Line data Source code
1 : // Copyright 2025 The LevelDB-Go and Pebble Authors. All rights reserved. Use
2 : // of this source code is governed by a BSD-style license that can be found in
3 : // the LICENSE file.
4 :
5 : package metamorphic
6 :
7 : import (
8 : "fmt"
9 : "math/rand/v2"
10 : "slices"
11 :
12 : "github.com/cockroachdb/pebble"
13 : "github.com/cockroachdb/pebble/cockroachkvs"
14 : "github.com/cockroachdb/pebble/internal/testkeys"
15 : "github.com/cockroachdb/pebble/sstable"
16 : "github.com/cockroachdb/pebble/sstable/colblk"
17 : )
18 :
19 : // CockroachKeyFormat provides a KeyFormat implementation that uses
20 : // CockroachDB's key encoding (as defined in the cockroachkvs package).
21 : var CockroachKeyFormat = KeyFormat{
22 : Name: "cockroachkvs",
23 : Comparer: &cockroachkvs.Comparer,
24 2 : KeySchema: func() *colblk.KeySchema { return &cockroachkvs.KeySchema }(),
25 : BlockPropertyCollectors: cockroachkvs.BlockPropertyCollectors,
26 2 : FormatKey: func(k UserKey) string {
27 2 : if len(k) == 0 {
28 2 : return ""
29 2 : }
30 2 : return fmt.Sprint(cockroachkvs.FormatKey(k))
31 : },
32 2 : FormatKeySuffix: func(s UserKeySuffix) string {
33 2 : if len(s) == 0 {
34 2 : return ""
35 2 : }
36 2 : return fmt.Sprint(cockroachkvs.FormatKeySuffix(s))
37 : },
38 1 : ParseFormattedKey: func(formattedKey string) UserKey {
39 1 : return UserKey(cockroachkvs.ParseFormattedKey(formattedKey))
40 1 : },
41 1 : ParseFormattedKeySuffix: func(formattedKeySuffix string) UserKeySuffix {
42 1 : return UserKeySuffix(cockroachkvs.ParseFormattedKeySuffix(formattedKeySuffix))
43 1 : },
44 1 : NewGenerator: func(km *keyManager, rng *rand.Rand, cfg OpConfig) KeyGenerator {
45 1 : return &cockroachKeyGenerator{
46 1 : keyManager: km,
47 1 : rng: rng,
48 1 : cfg: cfg,
49 1 : // TODO(jackson): Vary maxLogical.
50 1 : suffixSpace: cockroachSuffixKeyspace{maxLogical: 2},
51 1 : }
52 1 : },
53 1 : NewSuffixFilterMask: func() pebble.BlockPropertyFilterMask {
54 1 : return &cockroachkvs.MVCCWallTimeIntervalRangeKeyMask{}
55 1 : },
56 1 : NewSuffixBlockPropertyFilter: func(minSuffix, maxSuffix []byte) sstable.BlockPropertyFilter {
57 1 : minWallTime, _, err := cockroachkvs.DecodeMVCCTimestampSuffix(maxSuffix)
58 1 : if err != nil {
59 0 : panic(err)
60 : }
61 1 : maxWallTime, _, err := cockroachkvs.DecodeMVCCTimestampSuffix(minSuffix)
62 1 : if err != nil {
63 0 : panic(err)
64 : }
65 1 : return cockroachkvs.NewMVCCTimeIntervalFilter(minWallTime, maxWallTime)
66 : },
67 : }
68 :
69 : type cockroachKeyGenerator struct {
70 : keyManager *keyManager
71 : rng *rand.Rand
72 : cfg OpConfig
73 : suffixSpace cockroachSuffixKeyspace
74 : }
75 :
76 : // RecordPrecedingKey may be invoked before generating keys to inform the key
77 : // generator of a key that was previously generated and used within a related
78 : // test context.
79 0 : func (kg *cockroachKeyGenerator) RecordPrecedingKey(key []byte) {
80 0 : // If the key has a suffix that's larger than the current max suffix,
81 0 : // ratchet up the maximum of the distribution of suffixes.
82 0 : if i := cockroachkvs.Comparer.Split(key); i < len(key) {
83 0 : suffixIdx := kg.suffixSpace.ToSuffixIndex(key[i:])
84 0 : if suffixIdx > suffixIndex(kg.cfg.writeSuffixDist.Max()) {
85 0 : diff := uint64(suffixIdx) - kg.cfg.writeSuffixDist.Max()
86 0 : kg.cfg.writeSuffixDist.IncMax(diff)
87 0 : }
88 : }
89 : }
90 :
91 : // ExtendPrefix extends the given prefix key with additional bytes,
92 : // returning a new prefix that sorts after the given prefix.
93 0 : func (kg *cockroachKeyGenerator) ExtendPrefix(prefix []byte) []byte {
94 0 : // Copy prefix and strip the delimiter byte.
95 0 : p := append(make([]byte, 0, len(prefix)+3), prefix[:len(prefix)-1]...)
96 0 : p = append(p, randBytes(kg.rng, 1, 3)...)
97 0 : p = append(p, 0x00) // Delimiter byte
98 0 : return p
99 0 : }
100 :
101 : // RandKey returns a random key (either a previously known key, or a new key).
102 1 : func (kg *cockroachKeyGenerator) RandKey(newKeyProbability float64) []byte {
103 1 : return kg.randKey(newKeyProbability, nil /* bounds */)
104 1 : }
105 :
106 : // RandKeyInRange returns a random key (either a previously known key, or a new
107 : // key) in the given key range.
108 : func (kg *cockroachKeyGenerator) RandKeyInRange(
109 : newKeyProbability float64, kr pebble.KeyRange,
110 1 : ) []byte {
111 1 : return kg.randKey(newKeyProbability, &kr)
112 1 : }
113 :
114 : // RandPrefix returns a random prefix key (a key with no suffix).
115 1 : func (kg *cockroachKeyGenerator) RandPrefix(newPrefix float64) []byte {
116 1 : prefixes := kg.keyManager.prefixes()
117 1 : if len(prefixes) > 0 && kg.rng.Float64() > newPrefix {
118 1 : return pickOneUniform(kg.rng, prefixes)
119 1 : }
120 :
121 : // Use a new prefix.
122 1 : for {
123 1 : prefix := kg.generateKeyWithSuffix(4, 12, 0)
124 1 : if !kg.keyManager.prefixExists(prefix) {
125 1 : if !kg.keyManager.addNewKey(prefix) {
126 0 : panic("key must not exist if prefix doesn't exist")
127 : }
128 1 : return prefix
129 : }
130 : }
131 : }
132 :
133 : // SkewedSuffix generates a random suffix according to the configuration's
134 : // suffix distribution. It takes a probability 0 ≤ p ≤ 1.0 indicating the
135 : // probability with which the generator should increase the max suffix generated
136 : // by the generator.
137 : //
138 : // May return a nil suffix, with the probability the configuration's suffix
139 : // distribution assigns to the zero suffix.
140 1 : func (kg *cockroachKeyGenerator) SkewedSuffix(incMaxProb float64) []byte {
141 1 : if suffixIdx := kg.skewedSuffixInt(incMaxProb); suffixIdx != 0 {
142 1 : return kg.suffixSpace.ToMaterializedSuffix(suffixIdx)
143 1 : }
144 0 : return nil
145 : }
146 :
147 : // skewedSuffixInt is a helper of SkewedSuffix which returns the unencoded
148 : // suffix as an integer.
149 1 : func (kg *cockroachKeyGenerator) skewedSuffixInt(incMaxProb float64) suffixIndex {
150 1 : if kg.rng.Float64() < incMaxProb {
151 1 : kg.cfg.writeSuffixDist.IncMax(1)
152 1 : }
153 1 : return suffixIndex(kg.cfg.writeSuffixDist.Uint64(kg.rng))
154 : }
155 :
156 : // IncMaxSuffix increases the max suffix range and returns the new maximum
157 : // suffix (which is guaranteed to be larger than any previously generated
158 : // suffix).
159 0 : func (kg *cockroachKeyGenerator) IncMaxSuffix() []byte {
160 0 : kg.cfg.writeSuffixDist.IncMax(1)
161 0 : s := suffixIndex(kg.cfg.writeSuffixDist.Max())
162 0 : return kg.suffixSpace.ToMaterializedSuffix(s)
163 0 : }
164 :
165 : // SuffixRange generates a new uniformly random range of suffixes (low, high]
166 : // such that high is guaranteed to be strictly greater (as defined by
167 : // ComparePointSuffixes) than low.
168 : //
169 : // The high suffix may be nil, in which case the suffix range represents all
170 : // suffixes ≥ low.
171 1 : func (kg *cockroachKeyGenerator) SuffixRange() (low, high []byte) {
172 1 : a := kg.uniformSuffixInt()
173 1 : b := kg.uniformSuffixInt()
174 1 : if a < b {
175 1 : a, b = b, a
176 1 : } else if a == b {
177 0 : a++
178 0 : }
179 1 : return kg.suffixSpace.ToMaterializedSuffix(a), kg.suffixSpace.ToMaterializedSuffix(b)
180 : }
181 :
182 : // UniformSuffix returns a suffix in the same range as SkewedSuffix but with a
183 : // uniform distribution. This is used during reads to better exercise reading a
184 : // mix of older and newer keys. The suffix can be empty.
185 : //
186 : // May return a nil suffix.
187 1 : func (kg *cockroachKeyGenerator) UniformSuffix() []byte {
188 1 : if suffix := kg.uniformSuffixInt(); suffix != 0 {
189 1 : return kg.suffixSpace.ToMaterializedSuffix(suffix)
190 1 : }
191 0 : return nil
192 : }
193 :
194 : // uniformSuffixInt is a helper of UniformSuffix which returns the suffix
195 : // index.
196 1 : func (kg *cockroachKeyGenerator) uniformSuffixInt() suffixIndex {
197 1 : maxVal := kg.cfg.writeSuffixDist.Max()
198 1 : return suffixIndex(kg.rng.Int64N(int64(maxVal)))
199 1 : }
200 :
201 : // randKey returns a random key (either a previously known key or a new key).
202 : //
203 : // If bounds is not nil, the key will be inside the bounds.
204 : func (kg *cockroachKeyGenerator) randKey(
205 : newKeyProbability float64, bounds *pebble.KeyRange,
206 1 : ) []byte {
207 1 : var knownKeys [][]byte
208 1 : if bounds == nil {
209 1 : knownKeys = kg.keyManager.knownKeys()
210 1 : } else {
211 1 : if cockroachkvs.Compare(bounds.Start, bounds.End) >= 0 {
212 0 : panic(fmt.Sprintf("invalid bounds [%q, %q)", bounds.Start, bounds.End))
213 : }
214 1 : knownKeys = kg.keyManager.knownKeysInRange(*bounds)
215 : }
216 1 : switch {
217 1 : case len(knownKeys) > 0 && kg.rng.Float64() > newKeyProbability:
218 1 : // Use an existing user key.
219 1 : return pickOneUniform(kg.rng, knownKeys)
220 :
221 1 : case len(knownKeys) > 0 && kg.rng.Float64() > kg.cfg.newPrefix:
222 1 : // Use an existing prefix but a new suffix, producing a new user key.
223 1 : prefixes := kg.keyManager.prefixes()
224 1 :
225 1 : // If we're constrained to a key range, find which existing prefixes
226 1 : // fall within that key range.
227 1 : if bounds != nil {
228 1 : s, _ := slices.BinarySearchFunc(prefixes, bounds.Start, cockroachkvs.Compare)
229 1 : e, _ := slices.BinarySearchFunc(prefixes, bounds.End, cockroachkvs.Compare)
230 1 : prefixes = prefixes[s:e]
231 1 : }
232 :
233 1 : if len(prefixes) > 0 {
234 1 : for {
235 1 : // Pick a prefix on each iteration in case most or all suffixes are
236 1 : // already in use for any individual prefix.
237 1 : p := kg.rng.IntN(len(prefixes))
238 1 : suffix := suffixIndex(kg.cfg.writeSuffixDist.Uint64(kg.rng))
239 1 :
240 1 : var key []byte
241 1 : if suffix > 0 {
242 1 : key = append(append(key, prefixes[p]...), kg.suffixSpace.ToMaterializedSuffix(suffix)...)
243 1 : } else {
244 1 : key = append(key, prefixes[p]...)
245 1 : }
246 1 : if bounds == nil || (cockroachkvs.Compare(key, bounds.Start) >= 0 &&
247 1 : cockroachkvs.Compare(key, bounds.End) < 0) {
248 1 : if kg.keyManager.addNewKey(key) {
249 1 : return key
250 1 : }
251 : }
252 :
253 : // If the generated key already existed, or the generated key
254 : // fell outside the provided bounds, increase the suffix
255 : // distribution and loop.
256 1 : kg.cfg.writeSuffixDist.IncMax(1)
257 : }
258 : }
259 : // Otherwise fall through to generating a new prefix.
260 : }
261 :
262 1 : if bounds == nil {
263 1 : suffixIdx := kg.skewedSuffixInt(0.01)
264 1 : for {
265 1 : key := kg.generateKeyWithSuffix(4, 12, suffixIdx)
266 1 : if !kg.keyManager.prefixExists(kg.keyManager.kf.Comparer.Split.Prefix(key)) {
267 1 : if !kg.keyManager.addNewKey(key) {
268 0 : panic("key must not exist if prefix doesn't exist")
269 : }
270 1 : return key
271 : }
272 : }
273 : }
274 : // We need to generate a key between the bounds.
275 1 : startPrefix, startSuffixIdx := kg.suffixSpace.Split(bounds.Start)
276 1 : endPrefix, endSuffixIdx := kg.suffixSpace.Split(bounds.End)
277 1 :
278 1 : var prefix []byte
279 1 : var suffixIdx suffixIndex
280 1 : if cockroachkvs.Equal(startPrefix, endPrefix) {
281 0 : prefix = startPrefix
282 0 : // Bounds have the same prefix, generate a suffix in-between.
283 0 : if startSuffixIdx <= endSuffixIdx {
284 0 : panic(fmt.Sprintf("invalid bounds [%q, %q)", bounds.Start, bounds.End))
285 : }
286 0 : suffixIdx = kg.skewedSuffixInt(0.01)
287 0 : for i := 0; !(startSuffixIdx >= suffixIdx && endSuffixIdx < suffixIdx); i++ {
288 0 : if i > 10 {
289 0 : // This value is always >= startSuffix and < endSuffix.
290 0 : suffixIdx = (startSuffixIdx + endSuffixIdx) / 2
291 0 : break
292 : }
293 : // The suffix we want must exist in the current suffix range, we don't
294 : // want to keep increasing it here.
295 0 : suffixIdx = kg.skewedSuffixInt(0)
296 : }
297 1 : } else {
298 1 : prefix = append(testkeys.RandomPrefixInRange(
299 1 : startPrefix[:len(startPrefix)-1], // Strip the delimiter byte.
300 1 : endPrefix[:len(endPrefix)-1], // Strip the delimiter byte.
301 1 : kg.rng,
302 1 : ), 0x00) // Add back the delimiter byte.
303 1 : suffixIdx = kg.skewedSuffixInt(0.01)
304 1 : if cockroachkvs.Equal(prefix, startPrefix) {
305 0 : // We can't use a suffix which sorts before startSuffix.
306 0 : for i := 0; suffixIdx > startSuffixIdx; i++ {
307 0 : if i > 10 {
308 0 : suffixIdx = startSuffixIdx
309 0 : break
310 : }
311 0 : suffixIdx = kg.skewedSuffixInt(0)
312 : }
313 : }
314 : }
315 1 : key := slices.Clip(prefix)
316 1 : if suffixIdx != 0 {
317 1 : key = append(key, kg.suffixSpace.ToMaterializedSuffix(suffixIdx)...)
318 1 : }
319 1 : if cockroachkvs.Compare(key, bounds.Start) < 0 || cockroachkvs.Compare(key, bounds.End) >= 0 {
320 0 : panic(fmt.Sprintf("invalid randKey %q; bounds: [%q, %q) %v %v",
321 0 : key, bounds.Start, bounds.End,
322 0 : cockroachkvs.Compare(key, bounds.Start),
323 0 : cockroachkvs.Compare(key, bounds.End)))
324 : }
325 : // We might (rarely) produce an existing key here, that's ok.
326 1 : kg.keyManager.addNewKey(key)
327 1 : return key
328 : }
329 :
330 : // generateKeyWithSuffix generates a key with a random prefix and the suffix
331 : // corresponding to the provided suffix index. If the given suffix index is 0,
332 : // the key will not have a suffix.
333 : func (kg *cockroachKeyGenerator) generateKeyWithSuffix(
334 : minPrefixLen, maxPrefixLen int, suffixIdx suffixIndex,
335 1 : ) []byte {
336 1 : prefix := randCockroachPrefix(kg.rng, minPrefixLen, maxPrefixLen)
337 1 : if suffixIdx == 0 {
338 1 : return prefix
339 1 : }
340 1 : return append(prefix, kg.suffixSpace.ToMaterializedSuffix(suffixIdx)...)
341 : }
342 :
343 1 : func randCockroachPrefix(rng *rand.Rand, minLen, maxLen int) []byte {
344 1 : n := minLen + rng.IntN(maxLen-minLen+1)
345 1 : if n == 0 {
346 0 : return nil
347 0 : }
348 : // NB: The actual random values are not particularly important. We only use
349 : // lowercase letters because that makes visual determination of ordering
350 : // easier, rather than having to remember the lexicographic ordering of
351 : // uppercase vs lowercase, or letters vs numbers vs punctuation.
352 1 : const letters = "abcdefghijklmnopqrstuvwxyz"
353 1 : const lettersLen = uint64(len(letters))
354 1 : const lettersCharsPerRand = 12 // floor(log(math.MaxUint64)/log(lettersLen))
355 1 :
356 1 : var r uint64
357 1 : var q int
358 1 : buf := make([]byte, n+1)
359 1 : for i := 0; i < n; i++ {
360 1 : if q == 0 {
361 1 : r = rng.Uint64()
362 1 : q = lettersCharsPerRand
363 1 : }
364 1 : buf[i] = letters[r%lettersLen]
365 1 : r = r / lettersLen
366 1 : q--
367 : }
368 1 : buf[n] = 0x00 // Delimiter byte
369 1 : return buf
370 : }
371 :
372 : // A suffixIndex represents a unique suffix. The suffixIndex exists within a
373 : // one-dimensional space of int64s, but is remapped into a two-dimensional space
374 : // of MVCC timestamps of (WallTime, Logical) tuples.
375 : type suffixIndex int64
376 :
377 : // cockroackSuffixKeyspace defines the mapping between a one-dimensional
378 : // suffixIndex and the suffix it represents within the two-dimensional
379 : // (wallTime, logical) space.
380 : //
381 : // The mapping is configued by the value of maxLogical, with maxLogical+1
382 : // possible logical timestamps at each wall time.
383 : //
384 : // TODO(jackson): Update to disallow non-zero logical timestamps when the wall
385 : // time is zero if we begin to prohibit it.
386 : //
387 : // +------------------------------------------------------+
388 : // | suffix | maxLogical |
389 : // | index | 0 | 1 | 2 | 3 | 4 |
390 : // +------------------------------------------------------+
391 : // | 0 | (0,0) | (0,0) | (0,0) | (0,0) | (0,0) |
392 : // | 1 | (1,0) | (0,1) | (0,1) | (0,1) | (0,1) |
393 : // | 2 | (2,0) | (1,0) | (0,2) | (0,2) | (0,2) |
394 : // | 3 | (3,0) | (1,1) | (1,0) | (0,3) | (0,3) |
395 : // | 4 | (4,0) | (2,0) | (1,1) | (1,0) | (0,4) |
396 : // | 5 | (5,0) | (2,1) | (1,2) | (1,1) | (1,0) |
397 : // | 6 | (6,0) | (3,0) | (2,0) | (1,2) | (1,1) |
398 : // | 7 | (7,0) | (3,1) | (2,1) | (1,3) | (1,2) |
399 : // +------------------------------------------------------+
400 : type cockroachSuffixKeyspace struct {
401 : maxLogical int64
402 : }
403 :
404 1 : func (ks cockroachSuffixKeyspace) ToMaterializedSuffix(s suffixIndex) []byte {
405 1 : // There are maxLogical+1 possible logical timestamps at each wall time.
406 1 : wallTime := int64(s) / (ks.maxLogical + 1)
407 1 : logical := int64(s) % (ks.maxLogical + 1)
408 1 : return cockroachkvs.NewTimestampSuffix(uint64(wallTime), uint32(logical))
409 1 : }
410 :
411 1 : func (ks cockroachSuffixKeyspace) ToSuffixIndex(suffix []byte) suffixIndex {
412 1 : wallTime, logical, err := cockroachkvs.DecodeMVCCTimestampSuffix(suffix)
413 1 : if err != nil {
414 0 : panic(err)
415 : }
416 1 : return suffixIndex(int64(wallTime)*(ks.maxLogical+1) + int64(logical))
417 : }
418 :
419 1 : func (ks cockroachSuffixKeyspace) Split(key []byte) (prefix []byte, suffixIdx suffixIndex) {
420 1 : i := cockroachkvs.Split(key)
421 1 : return key[:i], ks.ToSuffixIndex(key[i:])
422 1 : }
|
