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