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