Skip to content

Commit

Permalink
crdbtest: move crdb rowlbk benchmarks
Browse files Browse the repository at this point in the history 
Move the remaining crdb rowblk benchmarks to crdb and unexport test
key generation functionality.
  • Loading branch information
@RaduBerinde
RaduBerinde committed Oct 22, 2024
1 parent bcf1585 commit d9f7191
Show file tree
Hide file tree
Showing 5 changed files with 281 additions and 268 deletions.
118 changes: 0 additions & 118 deletions internal/crdbtest/crdb.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,14 +12,10 @@ import (
"encoding/binary"
"fmt"
"io"
"math/rand/v2"
"slices"
"strings"
"time"
"unsafe"

"github.com/cockroachdb/crlib/crbytes"
"github.com/cockroachdb/crlib/crstrings"
"github.com/cockroachdb/errors"
"github.com/cockroachdb/pebble/internal/base"
"github.com/cockroachdb/pebble/internal/invariants"
Expand Down Expand Up @@ -387,120 +383,6 @@ func getKeyPartFromEngineKey(engineKey []byte) (key []byte, ok bool) {
return engineKey[:keyPartEnd], true
}

// KeyConfig configures the shape of the random keys generated.
type KeyConfig struct {
PrefixAlphabetLen int // Number of bytes in the alphabet used for the prefix.
PrefixLenShared int // Number of bytes shared by all key prefixes.
PrefixLen int // Number of bytes in the prefix.
AvgKeysPerPrefix int // Average number of keys (with varying suffixes) per prefix.
BaseWallTime uint64 // Smallest MVCC WallTime.
PercentLogical int // Percent of keys with non-zero MVCC logical time.
}

func (cfg KeyConfig) String() string {
return fmt.Sprintf(
"AlphaLen=%d,Prefix=%d,Shared=%d,KeysPerPrefix=%d%s",
cfg.PrefixAlphabetLen, cfg.PrefixLen, cfg.PrefixLenShared,
cfg.AvgKeysPerPrefix,
crstrings.If(cfg.PercentLogical != 0, fmt.Sprintf(",Logical=%d", cfg.PercentLogical)),
)
}

// RandomKVs constructs count random KVs with the provided parameters.
func RandomKVs(rng *rand.Rand, count int, cfg KeyConfig, valueLen int) (keys, vals [][]byte) {
g := makeCockroachKeyGen(rng, cfg)
sharedPrefix := make([]byte, cfg.PrefixLenShared)
for i := 0; i < len(sharedPrefix); i++ {
sharedPrefix[i] = byte(rng.IntN(cfg.PrefixAlphabetLen) + 'a')
}

keys = make([][]byte, 0, count)
vals = make([][]byte, 0, count)
for len(keys) < count {
prefix := g.randPrefix(sharedPrefix)
// We use the exponential distribution so that we occasionally have many
// suffixes
n := int(rng.ExpFloat64() * float64(cfg.AvgKeysPerPrefix))
n = max(n, 1)
for i := 0; i < n && len(keys) < count; i++ {
wallTime, logicalTime := g.randTimestamp()
k := makeKey(prefix, wallTime, logicalTime)
v := make([]byte, valueLen)
for j := range v {
v[j] = byte(rng.Uint32())
}
keys = append(keys, k)
vals = append(vals, v)
}
}
slices.SortFunc(keys, Compare)
return keys, vals
}

func makeKey(prefix []byte, wallTime uint64, logicalTime uint32) []byte {
k := make([]byte, 0, len(prefix)+MaxSuffixLen)
k = append(k, prefix...)
return EncodeTimestamp(k, wallTime, logicalTime)
}

// RandomQueryKeys returns a slice of count random query keys. Each key has a
// random prefix uniformly chosen from the distinct prefixes in writtenKeys and
// a random timestamp.
//
// Note that by setting baseWallTime to be large enough, we can simulate a query
// pattern that always retrieves the latest version of any prefix.
func RandomQueryKeys(
rng *rand.Rand, count int, writtenKeys [][]byte, baseWallTime uint64,
) [][]byte {
// Gather prefixes.
prefixes := make([][]byte, len(writtenKeys))
for i, k := range writtenKeys {
prefixes[i] = k[:Split(k)-1]
}
slices.SortFunc(prefixes, bytes.Compare)
prefixes = slices.CompactFunc(prefixes, bytes.Equal)
result := make([][]byte, count)
for i := range result {
prefix := prefixes[rng.IntN(len(prefixes))]
wallTime := baseWallTime + rng.Uint64N(uint64(time.Hour))
var logicalTime uint32
if rng.IntN(10) == 0 {
logicalTime = rng.Uint32()
}
result[i] = makeKey(prefix, wallTime, logicalTime)
}
return result
}

type cockroachKeyGen struct {
rng *rand.Rand
cfg KeyConfig
}

func makeCockroachKeyGen(rng *rand.Rand, cfg KeyConfig) cockroachKeyGen {
return cockroachKeyGen{
rng: rng,
cfg: cfg,
}
}

func (g *cockroachKeyGen) randPrefix(blockPrefix []byte) []byte {
prefix := make([]byte, 0, g.cfg.PrefixLen+MaxSuffixLen)
prefix = append(prefix, blockPrefix...)
for len(prefix) < g.cfg.PrefixLen {
prefix = append(prefix, byte(g.rng.IntN(g.cfg.PrefixAlphabetLen)+'a'))
}
return prefix
}

func (g *cockroachKeyGen) randTimestamp() (wallTime uint64, logicalTime uint32) {
wallTime = g.cfg.BaseWallTime + g.rng.Uint64N(uint64(time.Hour))
if g.cfg.PercentLogical > 0 && g.rng.IntN(100) < g.cfg.PercentLogical {
logicalTime = g.rng.Uint32()
}
return wallTime, logicalTime
}

const (
cockroachColRoachKey int = iota
cockroachColMVCCWallTime
Expand Down
26 changes: 13 additions & 13 deletions internal/crdbtest/crdb_bench_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ func BenchmarkRandSeekInSST(b *testing.B) {
version: sstable.TableFormatPebblev5,
},
}
keyCfg := KeyConfig{
keyCfg := keyGenConfig{
PrefixAlphabetLen: 26,
PrefixLen: 12,
PrefixLenShared: 4,
Expand Down Expand Up @@ -83,11 +83,11 @@ func benchmarkRandSeekInSST(
b *testing.B,
rng *rand.Rand,
numKeys int,
keyCfg KeyConfig,
keyCfg keyGenConfig,
valueLen int,
writerOpts sstable.WriterOptions,
) {
keys, values := RandomKVs(rng, numKeys, keyCfg, valueLen)
keys, values := randomKVs(rng, numKeys, keyCfg, valueLen)
obj := &objstorage.MemObj{}
w := sstable.NewWriter(obj, writerOpts)
for i := range keys {
Expand Down Expand Up @@ -127,7 +127,7 @@ func benchmarkRandSeekInSST(
require.NoError(b, iter.Close())

const numQueryKeys = 65536
queryKeys := RandomQueryKeys(rng, numQueryKeys, keys, keyCfg.BaseWallTime)
queryKeys := randomQueryKeys(rng, numQueryKeys, keys, keyCfg.BaseWallTime)
b.ResetTimer()
for i := 0; i < b.N; i++ {
key := queryKeys[i%numQueryKeys]
Expand All @@ -152,7 +152,7 @@ func BenchmarkCockroachDataBlockWriter(b *testing.B) {
for _, prefixLen := range []int{8, 32, 128} {
lenShared := int(float64(prefixLen) * lenSharedPct)
for _, valueLen := range []int{8, 128, 1024} {
keyConfig := KeyConfig{
keyConfig := keyGenConfig{
PrefixAlphabetLen: alphaLen,
PrefixLen: prefixLen,
PrefixLenShared: lenShared,
Expand All @@ -169,7 +169,7 @@ func BenchmarkCockroachDataBlockWriter(b *testing.B) {
}
}

func benchmarkCockroachDataBlockWriter(b *testing.B, keyConfig KeyConfig, valueLen int) {
func benchmarkCockroachDataBlockWriter(b *testing.B, keyConfig keyGenConfig, valueLen int) {
const targetBlockSize = 32 << 10
seed := uint64(time.Now().UnixNano())
rng := rand.New(rand.NewPCG(0, seed))
Expand Down Expand Up @@ -201,7 +201,7 @@ func BenchmarkCockroachDataBlockIterFull(b *testing.B) {
for _, percentLogical := range []int{0, 50} {
for _, valueLen := range []int{8, 128, 1024} {
cfg := benchConfig{
KeyConfig: KeyConfig{
keyGenConfig: keyGenConfig{
PrefixAlphabetLen: alphaLen,
PrefixLen: prefixLen,
PrefixLenShared: lenShared,
Expand All @@ -223,7 +223,7 @@ func BenchmarkCockroachDataBlockIterFull(b *testing.B) {

var shortBenchConfigs = []benchConfig{
{
KeyConfig: KeyConfig{
keyGenConfig: keyGenConfig{
PrefixAlphabetLen: 8,
PrefixLen: 8,
PrefixLenShared: 4,
Expand All @@ -233,7 +233,7 @@ var shortBenchConfigs = []benchConfig{
ValueLen: 8,
},
{
KeyConfig: KeyConfig{
keyGenConfig: keyGenConfig{
PrefixAlphabetLen: 8,
PrefixLen: 128,
PrefixLenShared: 64,
Expand Down Expand Up @@ -297,12 +297,12 @@ func BenchmarkCockroachDataBlockIterTransforms(b *testing.B) {
}

type benchConfig struct {
KeyConfig
keyGenConfig
ValueLen int
}

func (cfg benchConfig) String() string {
return fmt.Sprintf("%s,ValueLen=%d", cfg.KeyConfig, cfg.ValueLen)
return fmt.Sprintf("%s,ValueLen=%d", cfg.keyGenConfig, cfg.ValueLen)
}

func benchmarkCockroachDataBlockIter(
Expand All @@ -313,7 +313,7 @@ func benchmarkCockroachDataBlockIter(
rng := rand.New(rand.NewPCG(0, seed))
cfg.BaseWallTime = seed

serializedBlock, keys, _ := generateDataBlock(rng, targetBlockSize, cfg.KeyConfig, cfg.ValueLen)
serializedBlock, keys, _ := generateDataBlock(rng, targetBlockSize, cfg.keyGenConfig, cfg.ValueLen)

var decoder colblk.DataBlockDecoder
var it colblk.DataBlockIter
Expand Down Expand Up @@ -353,7 +353,7 @@ func benchmarkCockroachDataBlockIter(
if queryLatest {
baseWallTime += 24 * uint64(time.Hour)
}
queryKeys := RandomQueryKeys(rng, numQueryKeys, keys, baseWallTime)
queryKeys := randomQueryKeys(rng, numQueryKeys, keys, baseWallTime)
b.ResetTimer()
for i := 0; i < b.N; i++ {
k := queryKeys[i%numQueryKeys]
Expand Down
Loading

0 comments on commit d9f7191

Please sign in to comment.