cockroachkvs: export RandomKVs

Export the utility used to generate random CRDB KVs.

Author: RaduBerinde

cockroachkvs/cockroachkvs_bench_test.go

@@ -67,7 +67,7 @@ func BenchmarkRandSeekInSST(b *testing.B) {
 			version:  sstable.TableFormatPebblev6,
 		},
 	}
-	keyCfg := keyGenConfig{
+	keyCfg := KeyGenConfig{
 		PrefixAlphabetLen: 26,
 		RoachKeyLen:       12,
 		PrefixLenShared:   4,
@@ -95,11 +95,11 @@ func benchmarkRandSeekInSST(
 	b *testing.B,
 	rng *rand.Rand,
 	numKeys int,
-	keyCfg keyGenConfig,
+	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 {
@@ -168,12 +168,12 @@ func benchmarkRandSeekInSST(
 func BenchmarkCockroachDataColBlockWriter(b *testing.B) {
 	for _, cfg := range benchConfigs {
 		b.Run(cfg.String(), func(b *testing.B) {
-			benchmarkCockroachDataColBlockWriter(b, cfg.keyGenConfig, cfg.ValueLen)
+			benchmarkCockroachDataColBlockWriter(b, cfg.KeyGenConfig, cfg.ValueLen)
 		})
 	}
 }
 
-func benchmarkCockroachDataColBlockWriter(b *testing.B, keyConfig keyGenConfig, valueLen int) {
+func benchmarkCockroachDataColBlockWriter(b *testing.B, keyConfig KeyGenConfig, valueLen int) {
 	const targetBlockSize = 32 << 10
 	seed := uint64(time.Now().UnixNano())
 	rng := rand.New(rand.NewPCG(0, seed))
@@ -198,7 +198,7 @@ func benchmarkCockroachDataColBlockWriter(b *testing.B, keyConfig keyGenConfig,
 
 var benchConfigs = []benchConfig{
 	{
-		keyGenConfig: keyGenConfig{
+		KeyGenConfig: KeyGenConfig{
 			PrefixAlphabetLen: 8,
 			RoachKeyLen:       8,
 			PrefixLenShared:   4,
@@ -208,7 +208,7 @@ var benchConfigs = []benchConfig{
 		ValueLen: 8,
 	},
 	{
-		keyGenConfig: keyGenConfig{
+		KeyGenConfig: KeyGenConfig{
 			PrefixAlphabetLen: 8,
 			RoachKeyLen:       128,
 			PrefixLenShared:   64,
@@ -218,7 +218,7 @@ var benchConfigs = []benchConfig{
 		ValueLen: 128,
 	},
 	{
-		keyGenConfig: keyGenConfig{
+		KeyGenConfig: KeyGenConfig{
 			PrefixAlphabetLen: 26,
 			RoachKeyLen:       1024,
 			PrefixLenShared:   512,
@@ -282,12 +282,12 @@ func BenchmarkCockroachDataColBlockIterTransforms(b *testing.B) {
 }
 
 type benchConfig struct {
-	keyGenConfig
+	KeyGenConfig
 	ValueLen int
 }
 
 func (cfg benchConfig) String() string {
-	return fmt.Sprintf("%s,ValueLen=%d", cfg.keyGenConfig, cfg.ValueLen)
+	return fmt.Sprintf("%s,ValueLen=%d", cfg.KeyGenConfig, cfg.ValueLen)
 }
 
 func benchmarkCockroachDataColBlockIter(
@@ -298,7 +298,7 @@ func benchmarkCockroachDataColBlockIter(
 	rng := rand.New(rand.NewPCG(0, seed))
 	cfg.BaseWallTime = seed
 
-	serializedBlock, keys, _ := generateDataBlock(rng, targetBlockSize, cfg.keyGenConfig, cfg.ValueLen)
+	serializedBlock, keys, _ := generateDataBlock(rng, targetBlockSize, cfg.KeyGenConfig, cfg.ValueLen)
 
 	var decoder colblk.DataBlockDecoder
 	var it colblk.DataBlockIter

cockroachkvs/cockroachkvs_test.go

@@ -316,7 +316,7 @@ func TestRandKeys(t *testing.T) {
 		seed := uint64(1234)
 		count := 10
 		valueLen := 4
-		cfg := keyGenConfig{
+		cfg := KeyGenConfig{
 			PrefixAlphabetLen: 8,
 			PrefixLenShared:   4,
 			RoachKeyLen:       8,
@@ -344,7 +344,7 @@ func TestRandKeys(t *testing.T) {
 		switch d.Cmd {
 		case "rand-kvs":
 			var vals [][]byte
-			keys, vals = randomKVs(rng, count, cfg, valueLen)
+			keys, vals = RandomKVs(rng, count, cfg, valueLen)
 			for i, key := range keys {
 				fmt.Fprintf(&buf, "%s = %X\n", formatUserKey(key), vals[i])
 			}
@@ -369,7 +369,7 @@ func TestCockroachDataColBlock(t *testing.T) {
 	rng := rand.New(rand.NewPCG(0, seed))
 
 	for i := 0; i < 100; i++ {
-		keyCfg := keyGenConfig{
+		keyCfg := KeyGenConfig{
 			PrefixAlphabetLen: 2 + rng.IntN(25),
 			RoachKeyLen:       10 + int(rng.ExpFloat64()*10),
 			AvgKeysPerPrefix:  1 + int(rng.ExpFloat64()*[]float64{1, 10, 100}[rng.IntN(3)]),
@@ -400,7 +400,7 @@ func TestCockroachDataColBlock(t *testing.T) {
 	}
 }
 
-func testCockroachDataColBlock(t *testing.T, seed uint64, keyCfg keyGenConfig) {
+func testCockroachDataColBlock(t *testing.T, seed uint64, keyCfg KeyGenConfig) {
 	const targetBlockSize = 32 << 10
 	rng := rand.New(rand.NewPCG(0, seed))
 	valueLen := 1 + rng.IntN(300)
@@ -453,9 +453,9 @@ func testCockroachDataColBlock(t *testing.T, seed uint64, keyCfg keyGenConfig) {
 // parameters. Returns the serialized block data and the keys and values
 // written.
 func generateDataBlock(
-	rng *rand.Rand, targetBlockSize int, cfg keyGenConfig, valueLen int,
+	rng *rand.Rand, targetBlockSize int, cfg KeyGenConfig, valueLen int,
 ) (data []byte, keys [][]byte, values [][]byte) {
-	keys, values = randomKVs(rng, targetBlockSize/valueLen, cfg, valueLen)
+	keys, values = RandomKVs(rng, targetBlockSize/valueLen, cfg, valueLen)
 
 	var w colblk.DataBlockEncoder
 	w.Init(&KeySchema)
@@ -478,80 +478,6 @@ func (g getInternalValuer) GetInternalValueForPrefixAndValueHandle(
 	return g(handle)
 }
 
-// keyGenConfig configures the shape of the random keys generated.
-type keyGenConfig struct {
-	PrefixAlphabetLen  int    // Number of bytes in the alphabet used for the prefix.
-	PrefixLenShared    int    // Number of bytes shared by all key prefixes.
-	RoachKeyLen        int    // Number of bytes in the prefix (without the 0 sentinel byte).
-	AvgKeysPerPrefix   int    // Average number of keys (with varying suffixes) per prefix.
-	BaseWallTime       uint64 // Smallest MVCC WallTime.
-	PercentLogical     int    // Percent of MVCC keys with non-zero MVCC logical time.
-	PercentEmptySuffix int    // Percent of keys with empty suffix.
-	PercentLockSuffix  int    // Percent of keys with lock suffix.
-}
-
-func (cfg keyGenConfig) String() string {
-	return fmt.Sprintf(
-		"AlphaLen=%d,Prefix=%d,Shared=%d,KeysPerPrefix=%d%s",
-		cfg.PrefixAlphabetLen, cfg.RoachKeyLen, 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 keyGenConfig, 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)
-	for len(keys) < count {
-		roachKey := g.randRoachKey(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++ {
-			if cfg.PercentEmptySuffix+cfg.PercentLockSuffix > 0 {
-				if r := rng.IntN(100); r < cfg.PercentEmptySuffix+cfg.PercentLockSuffix {
-					k := append(roachKey, 0)
-					if r < cfg.PercentLockSuffix {
-						// Generate a lock key suffix.
-						for j := 0; j < engineKeyVersionLockTableLen; j++ {
-							k = append(k, byte(g.rng.IntN(g.cfg.PrefixAlphabetLen)+'a'))
-						}
-						k = append(k, engineKeyVersionLockTableLen+1)
-					}
-					keys = append(keys, k)
-					continue
-				}
-			}
-
-			wallTime, logicalTime := g.randTimestamp()
-			k := makeMVCCKey(roachKey, wallTime, logicalTime)
-			keys = append(keys, k)
-		}
-	}
-	slices.SortFunc(keys, Compare)
-	vals = make([][]byte, count)
-	for i := range vals {
-		v := make([]byte, valueLen)
-		for j := range v {
-			v[j] = byte(rng.Uint32())
-		}
-		vals[i] = v
-	}
-	return keys, vals
-}
-
-func makeMVCCKey(roachKey []byte, wallTime uint64, logicalTime uint32) []byte {
-	k := slices.Grow(slices.Clip(roachKey), MaxSuffixLen)
-	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.
@@ -581,35 +507,6 @@ func randomQueryKeys(
 	return result
 }
 
-type cockroachKeyGen struct {
-	rng *rand.Rand
-	cfg keyGenConfig
-}
-
-func makeCockroachKeyGen(rng *rand.Rand, cfg keyGenConfig) cockroachKeyGen {
-	return cockroachKeyGen{
-		rng: rng,
-		cfg: cfg,
-	}
-}
-
-func (g *cockroachKeyGen) randRoachKey(blockPrefix []byte) []byte {
-	roachKey := make([]byte, 0, g.cfg.RoachKeyLen+MaxSuffixLen)
-	roachKey = append(roachKey, blockPrefix...)
-	for len(roachKey) < g.cfg.RoachKeyLen {
-		roachKey = append(roachKey, byte(g.rng.IntN(g.cfg.PrefixAlphabetLen)+'a'))
-	}
-	return roachKey
-}
-
-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
-}
-
 // formatUserKey formats a user key in the format:
 //
 //	<roach-key> [ @ <version-hex> ]

cockroachkvs/rowblk_bench_test.go

@@ -18,16 +18,16 @@ import (
 func BenchmarkCockroachDataRowBlockWriter(b *testing.B) {
 	for _, cfg := range benchConfigs {
 		b.Run(cfg.String(), func(b *testing.B) {
-			benchmarkCockroachDataRowBlockWriter(b, cfg.keyGenConfig, cfg.ValueLen)
+			benchmarkCockroachDataRowBlockWriter(b, cfg.KeyGenConfig, cfg.ValueLen)
 		})
 	}
 }
 
-func benchmarkCockroachDataRowBlockWriter(b *testing.B, keyConfig keyGenConfig, valueLen int) {
+func benchmarkCockroachDataRowBlockWriter(b *testing.B, keyConfig KeyGenConfig, valueLen int) {
 	const targetBlockSize = 32 << 10
 	seed := uint64(time.Now().UnixNano())
 	rng := rand.New(rand.NewPCG(0, seed))
-	keys, values := randomKVs(rng, targetBlockSize/valueLen, keyConfig, valueLen)
+	keys, values := RandomKVs(rng, targetBlockSize/valueLen, keyConfig, valueLen)
 
 	var w rowblk.Writer
 	w.RestartInterval = 16
@@ -54,16 +54,16 @@ func benchmarkCockroachDataRowBlockWriter(b *testing.B, keyConfig keyGenConfig,
 func BenchmarkCockroachDataRowBlockIter(b *testing.B) {
 	for _, cfg := range benchConfigs {
 		b.Run(cfg.String(), func(b *testing.B) {
-			benchmarkCockroachDataRowBlockIter(b, cfg.keyGenConfig, cfg.ValueLen)
+			benchmarkCockroachDataRowBlockIter(b, cfg.KeyGenConfig, cfg.ValueLen)
 		})
 	}
 }
 
-func benchmarkCockroachDataRowBlockIter(b *testing.B, keyConfig keyGenConfig, valueLen int) {
+func benchmarkCockroachDataRowBlockIter(b *testing.B, keyConfig KeyGenConfig, valueLen int) {
 	const targetBlockSize = 32 << 10
 	seed := uint64(time.Now().UnixNano())
 	rng := rand.New(rand.NewPCG(0, seed))
-	keys, values := randomKVs(rng, targetBlockSize/valueLen, keyConfig, valueLen)
+	keys, values := RandomKVs(rng, targetBlockSize/valueLen, keyConfig, valueLen)
 
 	var w rowblk.Writer
 	w.RestartInterval = 16