Line data Source code
1 : // Copyright 2024 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 : "cmp"
9 : "fmt"
10 : "math/rand/v2"
11 : "slices"
12 :
13 : "github.com/cockroachdb/pebble"
14 : "github.com/cockroachdb/pebble/internal/testkeys"
15 : "github.com/cockroachdb/pebble/sstable"
16 : "github.com/cockroachdb/pebble/sstable/colblk"
17 : )
18 :
19 : var TestkeysKeyFormat = KeyFormat{
20 : Comparer: testkeys.Comparer,
21 2 : KeySchema: func() *colblk.KeySchema {
22 2 : kf := colblk.DefaultKeySchema(testkeys.Comparer, 16 /* bundle size */)
23 2 : return &kf
24 2 : }(),
25 2 : FormatKey: func(k UserKey) string { return string(k) },
26 2 : FormatKeySuffix: func(s UserKeySuffix) string { return string(s) },
27 1 : NewGenerator: func(km *keyManager, rng *rand.Rand, cfg OpConfig) KeyGenerator {
28 1 : return &testkeyKeyGenerator{
29 1 : keyManager: km,
30 1 : rng: rng,
31 1 : cfg: cfg,
32 1 : }
33 1 : },
34 2 : NewSuffixFilterMask: func() pebble.BlockPropertyFilterMask {
35 2 : return sstable.NewTestKeysMaskingFilter()
36 2 : },
37 2 : NewSuffixBlockPropertyFilter: func(filterMin, filterMax []byte) sstable.BlockPropertyFilter {
38 2 : low, err := testkeys.ParseSuffix(filterMin)
39 2 : if err != nil {
40 0 : panic(err)
41 : }
42 2 : high, err := testkeys.ParseSuffix(filterMax)
43 2 : if err != nil {
44 0 : panic(err)
45 : }
46 2 : return sstable.NewTestKeysBlockPropertyFilter(uint64(low), uint64(high))
47 : },
48 : }
49 :
50 : type testkeyKeyGenerator struct {
51 : keyManager *keyManager
52 : rng *rand.Rand
53 : cfg OpConfig
54 : }
55 :
56 : // RecordPrecedingKey may be invoked before generating keys to inform the key
57 : // generator of a key that was previously generated and used within a related
58 : // test context.
59 1 : func (kg *testkeyKeyGenerator) RecordPrecedingKey(key []byte) {
60 1 : // If the key has a suffix that's larger than the current max suffix,
61 1 : // ratchet up the maximum of the distribution of suffixes.
62 1 : if i := testkeys.Comparer.Split(key); i < len(key) {
63 1 : s, err := testkeys.ParseSuffix(key[i:])
64 1 : if err != nil {
65 0 : panic(err)
66 : }
67 1 : if uint64(s) > kg.cfg.writeSuffixDist.Max() {
68 1 : diff := uint64(s) - kg.cfg.writeSuffixDist.Max()
69 1 : kg.cfg.writeSuffixDist.IncMax(diff)
70 1 : }
71 : }
72 : }
73 :
74 : // RandKey returns a random key (either a previously known key, or a new key).
75 1 : func (kg *testkeyKeyGenerator) RandKey(newKeyProbability float64) []byte {
76 1 : return kg.randKey(newKeyProbability, nil /* bounds */)
77 1 : }
78 :
79 : // RandKeyInRange returns a random key (either a previously known key, or a new
80 : // key) in the given key range.
81 : func (kg *testkeyKeyGenerator) RandKeyInRange(
82 : newKeyProbability float64, kr pebble.KeyRange,
83 1 : ) []byte {
84 1 : return kg.randKey(newKeyProbability, &kr)
85 1 : }
86 :
87 : // RandPrefix returns a random prefix key (a key with no suffix).
88 1 : func (kg *testkeyKeyGenerator) RandPrefix(newPrefix float64) []byte {
89 1 : prefixes := kg.keyManager.prefixes()
90 1 : if len(prefixes) > 0 && kg.rng.Float64() > newPrefix {
91 1 : return pickOneUniform(kg.rng, prefixes)
92 1 : }
93 :
94 : // Use a new prefix.
95 1 : for {
96 1 : prefix := kg.generateKeyWithSuffix(4, 12, 0)
97 1 : if !kg.keyManager.prefixExists(prefix) {
98 1 : if !kg.keyManager.addNewKey(prefix) {
99 0 : panic("key must not exist if prefix doesn't exist")
100 : }
101 1 : return prefix
102 : }
103 : }
104 : }
105 :
106 : // UniqueKeys takes a key-generating function and uses it to generate n unique
107 : // keys, returning them in sorted order.
108 1 : func (kg *testkeyKeyGenerator) UniqueKeys(n int, genFn func() []byte) [][]byte {
109 1 : keys := make([][]byte, n)
110 1 : used := make(map[string]struct{}, n)
111 1 : for i := range keys {
112 1 : for attempts := 0; ; attempts++ {
113 1 : keys[i] = genFn()
114 1 : if _, exists := used[string(keys[i])]; !exists {
115 1 : break
116 : }
117 1 : if attempts > 100000 {
118 0 : panic("could not generate unique key")
119 : }
120 : }
121 1 : used[string(keys[i])] = struct{}{}
122 : }
123 1 : slices.SortFunc(keys, kg.cmp)
124 1 : return keys
125 : }
126 :
127 : // SkewedSuffix generates a random suffix according to the configuration's
128 : // suffix distribution. It takes a probability 0 ≤ p ≤ 1.0 indicating the
129 : // probability with which the generator should increase the max suffix generated
130 : // by the generator.
131 : //
132 : // May return a nil suffix, with the probability the configuration's suffix
133 : // distribution assigns to the zero suffix.
134 1 : func (kg *testkeyKeyGenerator) SkewedSuffix(incMaxProb float64) []byte {
135 1 : if suffix := kg.skewedSuffixInt(incMaxProb); suffix != 0 {
136 1 : return testkeys.Suffix(suffix)
137 1 : }
138 1 : return nil
139 : }
140 :
141 : // skewedSuffixInt is a helper of SkewedSuffix which returns the unencoded
142 : // suffix as an integer.
143 1 : func (kg *testkeyKeyGenerator) skewedSuffixInt(incMaxProb float64) int64 {
144 1 : if kg.rng.Float64() < incMaxProb {
145 1 : kg.cfg.writeSuffixDist.IncMax(1)
146 1 : }
147 1 : return int64(kg.cfg.writeSuffixDist.Uint64(kg.rng))
148 : }
149 :
150 : // IncMaxSuffix increases the max suffix range and returns the new maximum
151 : // suffix (which is guaranteed to be larger than any previously generated
152 : // suffix).
153 1 : func (kg *testkeyKeyGenerator) IncMaxSuffix() []byte {
154 1 : kg.cfg.writeSuffixDist.IncMax(1)
155 1 : return testkeys.Suffix(int64(kg.cfg.writeSuffixDist.Max()))
156 1 : }
157 :
158 1 : func (kg *testkeyKeyGenerator) SuffixRange() (low, high []byte) {
159 1 : a := kg.uniformSuffixInt()
160 1 : b := kg.uniformSuffixInt()
161 1 : if a > b {
162 1 : a, b = b, a
163 1 : } else if a == b {
164 1 : b++
165 1 : }
166 1 : return testkeys.Suffix(a), testkeys.Suffix(b)
167 : }
168 :
169 : // UniformSuffix returns a suffix in the same range as SkewedSuffix but with a
170 : // uniform distribution. This is used during reads to better exercise reading a
171 : // mix of older and newer keys. The suffix can be empty.
172 : //
173 : // May return a nil suffix.
174 1 : func (kg *testkeyKeyGenerator) UniformSuffix() []byte {
175 1 : if suffix := kg.uniformSuffixInt(); suffix != 0 {
176 1 : return testkeys.Suffix(suffix)
177 1 : }
178 1 : return nil
179 : }
180 :
181 : // uniformSuffixInt is a helper of UniformSuffix which returns the unencoded
182 : // suffix as an integer.
183 1 : func (kg *testkeyKeyGenerator) uniformSuffixInt() int64 {
184 1 : maxVal := kg.cfg.writeSuffixDist.Max()
185 1 : return kg.rng.Int64N(int64(maxVal))
186 1 : }
187 :
188 : // randKey returns a random key (either a previously known key or a new key).
189 : //
190 : // If bounds is not nil, the key will be inside the bounds.
191 1 : func (kg *testkeyKeyGenerator) randKey(newKeyProbability float64, bounds *pebble.KeyRange) []byte {
192 1 : var knownKeys [][]byte
193 1 : if bounds == nil {
194 1 : knownKeys = kg.keyManager.knownKeys()
195 1 : } else {
196 1 : if kg.cmp(bounds.Start, bounds.End) >= 0 {
197 0 : panic(fmt.Sprintf("invalid bounds [%q, %q)", bounds.Start, bounds.End))
198 : }
199 1 : knownKeys = kg.keyManager.knownKeysInRange(*bounds)
200 : }
201 1 : switch {
202 1 : case len(knownKeys) > 0 && kg.rng.Float64() > newKeyProbability:
203 1 : // Use an existing user key.
204 1 : return pickOneUniform(kg.rng, knownKeys)
205 :
206 1 : case len(knownKeys) > 0 && kg.rng.Float64() > kg.cfg.newPrefix:
207 1 : // Use an existing prefix but a new suffix, producing a new user key.
208 1 : prefixes := kg.keyManager.prefixes()
209 1 :
210 1 : // If we're constrained to a key range, find which existing prefixes
211 1 : // fall within that key range.
212 1 : if bounds != nil {
213 1 : s, _ := slices.BinarySearchFunc(prefixes, bounds.Start, kg.cmp)
214 1 : e, _ := slices.BinarySearchFunc(prefixes, bounds.End, kg.cmp)
215 1 : prefixes = prefixes[s:e]
216 1 : }
217 :
218 1 : if len(prefixes) > 0 {
219 1 : for {
220 1 : // Pick a prefix on each iteration in case most or all suffixes are
221 1 : // already in use for any individual prefix.
222 1 : p := kg.rng.IntN(len(prefixes))
223 1 : suffix := int64(kg.cfg.writeSuffixDist.Uint64(kg.rng))
224 1 :
225 1 : var key []byte
226 1 : if suffix > 0 {
227 1 : key = resizeBuffer(key, len(prefixes[p]), testkeys.SuffixLen(suffix))
228 1 : n := copy(key, prefixes[p])
229 1 : testkeys.WriteSuffix(key[n:], suffix)
230 1 : } else {
231 1 : key = resizeBuffer(key, len(prefixes[p]), 0)
232 1 : copy(key, prefixes[p])
233 1 : }
234 :
235 1 : if bounds == nil || (kg.cmp(key, bounds.Start) >= 0 && kg.cmp(key, bounds.End) < 0) {
236 1 : if kg.keyManager.addNewKey(key) {
237 1 : return key
238 1 : }
239 : }
240 :
241 : // If the generated key already existed, or the generated key
242 : // fell outside the provided bounds, increase the suffix
243 : // distribution and loop.
244 1 : kg.cfg.writeSuffixDist.IncMax(1)
245 : }
246 : }
247 : // Otherwise fall through to generating a new prefix.
248 : }
249 :
250 1 : if bounds == nil {
251 1 : suffix := kg.skewedSuffixInt(0.01)
252 1 : for {
253 1 : key := kg.generateKeyWithSuffix(4, 12, suffix)
254 1 : if !kg.keyManager.prefixExists(kg.prefix(key)) {
255 1 : if !kg.keyManager.addNewKey(key) {
256 0 : panic("key must not exist if prefix doesn't exist")
257 : }
258 1 : return key
259 : }
260 : }
261 : }
262 : // We need to generate a key between the bounds.
263 1 : startPrefix, startSuffix := kg.parseKey(bounds.Start)
264 1 : endPrefix, endSuffix := kg.parseKey(bounds.End)
265 1 : var prefix []byte
266 1 : var suffix int64
267 1 : if kg.equal(startPrefix, endPrefix) {
268 0 : prefix = startPrefix
269 0 : // Bounds have the same prefix, generate a suffix in-between.
270 0 : if cmpSuffix(startSuffix, endSuffix) >= 0 {
271 0 : panic(fmt.Sprintf("invalid bounds [%q, %q)", bounds.Start, bounds.End))
272 : }
273 0 : suffix = kg.skewedSuffixInt(0.01)
274 0 : for i := 0; !(cmpSuffix(startSuffix, suffix) <= 0 && cmpSuffix(suffix, endSuffix) < 0); i++ {
275 0 : if i > 10 {
276 0 : // This value is always >= startSuffix and < endSuffix.
277 0 : suffix = (startSuffix + endSuffix) / 2
278 0 : break
279 : }
280 : // The suffix we want must exist in the current suffix range, we don't
281 : // want to keep increasing it here.
282 0 : suffix = kg.skewedSuffixInt(0)
283 : }
284 1 : } else {
285 1 : prefix = testkeys.RandomPrefixInRange(startPrefix, endPrefix, kg.rng)
286 1 : suffix = kg.skewedSuffixInt(0.01)
287 1 : if kg.equal(prefix, startPrefix) {
288 1 : // We can't use a suffix which sorts before startSuffix.
289 1 : for i := 0; cmpSuffix(suffix, startSuffix) < 0; i++ {
290 0 : if i > 10 {
291 0 : suffix = startSuffix
292 0 : break
293 : }
294 0 : suffix = kg.skewedSuffixInt(0)
295 : }
296 : }
297 : }
298 1 : key := slices.Clip(prefix)
299 1 : if suffix != 0 {
300 1 : key = append(key, testkeys.Suffix(suffix)...)
301 1 : }
302 1 : if kg.cmp(key, bounds.Start) < 0 || kg.cmp(key, bounds.End) >= 0 {
303 0 : panic(fmt.Sprintf("invalid randKey %q; bounds: [%q, %q) %v %v",
304 0 : key, bounds.Start, bounds.End, kg.cmp(key, bounds.Start), kg.cmp(key, bounds.End)))
305 : }
306 : // We might (rarely) produce an existing key here, that's ok.
307 1 : kg.keyManager.addNewKey(key)
308 1 : return key
309 : }
310 :
311 : // generateKeyWithSuffix generates a key with a random prefix and the given
312 : // suffix. If the given suffix is 0, the key will not have a suffix.
313 : func (kg *testkeyKeyGenerator) generateKeyWithSuffix(
314 : minPrefixLen, maxPrefixLen int, suffix int64,
315 1 : ) []byte {
316 1 : prefix := randBytes(kg.rng, minPrefixLen, maxPrefixLen)
317 1 : if suffix == 0 {
318 1 : return prefix
319 1 : }
320 1 : return append(prefix, testkeys.Suffix(suffix)...)
321 : }
322 :
323 : // cmpSuffix compares two suffixes, where suffix 0 means there is no suffix.
324 1 : func cmpSuffix(s1, s2 int64) int {
325 1 : switch {
326 0 : case s1 == s2:
327 0 : return 0
328 0 : case s1 == 0:
329 0 : return -1
330 1 : case s2 == 0:
331 1 : return +1
332 0 : default:
333 0 : return cmp.Compare(s2, s1)
334 : }
335 : }
336 :
337 1 : func (kg *testkeyKeyGenerator) cmp(a, b []byte) int {
338 1 : return kg.keyManager.kf.Comparer.Compare(a, b)
339 1 : }
340 :
341 1 : func (kg *testkeyKeyGenerator) equal(a, b []byte) bool {
342 1 : return kg.keyManager.kf.Comparer.Equal(a, b)
343 1 : }
344 :
345 1 : func (kg *testkeyKeyGenerator) split(a []byte) int {
346 1 : return kg.keyManager.kf.Comparer.Split(a)
347 1 : }
348 :
349 1 : func (kg *testkeyKeyGenerator) prefix(a []byte) []byte {
350 1 : n := kg.split(a)
351 1 : return a[:n:n]
352 1 : }
353 :
354 1 : func (kg *testkeyKeyGenerator) parseKey(k []byte) (prefix []byte, suffix int64) {
355 1 : n := kg.split(k)
356 1 : if n == len(k) {
357 1 : return k, 0
358 1 : }
359 1 : suffix, err := testkeys.ParseSuffix(k[n:])
360 1 : if err != nil {
361 0 : panic(fmt.Sprintf("error parsing suffix for key %q", k))
362 : }
363 1 : return k[:n:n], suffix
364 : }
365 :
366 1 : func randBytes(rng *rand.Rand, minLen, maxLen int) []byte {
367 1 : n := minLen + rng.IntN(maxLen-minLen+1)
368 1 : if n == 0 {
369 1 : return nil
370 1 : }
371 : // NB: The actual random values are not particularly important. We only use
372 : // lowercase letters because that makes visual determination of ordering
373 : // easier, rather than having to remember the lexicographic ordering of
374 : // uppercase vs lowercase, or letters vs numbers vs punctuation.
375 1 : const letters = "abcdefghijklmnopqrstuvwxyz"
376 1 : const lettersLen = uint64(len(letters))
377 1 : const lettersCharsPerRand = 12 // floor(log(math.MaxUint64)/log(lettersLen))
378 1 :
379 1 : var r uint64
380 1 : var q int
381 1 : buf := make([]byte, n)
382 1 : for i := range buf {
383 1 : if q == 0 {
384 1 : r = rng.Uint64()
385 1 : q = lettersCharsPerRand
386 1 : }
387 1 : buf[i] = letters[r%lettersLen]
388 1 : r = r / lettersLen
389 1 : q--
390 : }
391 1 : return buf
392 : }
393 :
394 1 : func pickOneUniform[S ~[]E, E any](rng *rand.Rand, x S) E {
395 1 : return x[rng.IntN(len(x))]
396 1 : }
|
