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 : "slices"
11 :
12 : "github.com/cockroachdb/pebble"
13 : "github.com/cockroachdb/pebble/internal/testkeys"
14 : "golang.org/x/exp/rand"
15 : )
16 :
17 : type keyGenerator struct {
18 : keyManager *keyManager
19 : rng *rand.Rand
20 : cfg OpConfig
21 : }
22 :
23 1 : func newKeyGenerator(km *keyManager, rng *rand.Rand, cfg OpConfig) *keyGenerator {
24 1 : return &keyGenerator{
25 1 : keyManager: km,
26 1 : rng: rng,
27 1 : cfg: cfg,
28 1 : }
29 1 : }
30 :
31 : // RandKey returns a random key (either a previously known key, or a new key).
32 1 : func (kg *keyGenerator) RandKey(newKeyProbability float64) []byte {
33 1 : return kg.randKey(newKeyProbability, nil /* bounds */)
34 1 : }
35 :
36 : // RandKeyInRange returns a random key (either a previously known key, or a new
37 : // key) in the given key range.
38 1 : func (kg *keyGenerator) RandKeyInRange(newKeyProbability float64, kr pebble.KeyRange) []byte {
39 1 : return kg.randKey(newKeyProbability, &kr)
40 1 : }
41 :
42 : // RandPrefix returns a random prefix key (a key with no suffix).
43 1 : func (kg *keyGenerator) RandPrefix(newPrefix float64) []byte {
44 1 : prefixes := kg.keyManager.prefixes()
45 1 : if len(prefixes) > 0 && kg.rng.Float64() > newPrefix {
46 1 : return pickOneUniform(kg.rng, prefixes)
47 1 : }
48 :
49 : // Use a new prefix.
50 1 : for {
51 1 : prefix := kg.generateKeyWithSuffix(4, 12, 0)
52 1 : if !kg.keyManager.prefixExists(prefix) {
53 1 : if !kg.keyManager.addNewKey(prefix) {
54 0 : panic("key must not exist if prefix doesn't exist")
55 : }
56 1 : 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 1 : func (kg *keyGenerator) UniqueKeys(n int, genFn func() []byte) [][]byte {
64 1 : keys := make([][]byte, n)
65 1 : used := make(map[string]struct{}, n)
66 1 : for i := range keys {
67 1 : for attempts := 0; ; attempts++ {
68 1 : keys[i] = genFn()
69 1 : if _, exists := used[string(keys[i])]; !exists {
70 1 : break
71 : }
72 1 : if attempts > 100000 {
73 0 : panic("could not generate unique key")
74 : }
75 : }
76 1 : used[string(keys[i])] = struct{}{}
77 : }
78 1 : slices.SortFunc(keys, kg.cmp)
79 1 : 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 1 : func (kg *keyGenerator) SkewedSuffix(incMaxProb float64) []byte {
90 1 : if suffix := kg.SkewedSuffixInt(incMaxProb); suffix != 0 {
91 1 : return testkeys.Suffix(suffix)
92 1 : }
93 1 : return nil
94 : }
95 :
96 : // SkewedSuffixInt is a variant of SkewedSuffix which returns the unencoded
97 : // suffix as an integer.
98 1 : func (kg *keyGenerator) SkewedSuffixInt(incMaxProb float64) int64 {
99 1 : if kg.rng.Float64() < incMaxProb {
100 1 : kg.cfg.writeSuffixDist.IncMax(1)
101 1 : }
102 1 : 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 1 : func (kg *keyGenerator) IncMaxSuffix() []byte {
109 1 : kg.cfg.writeSuffixDist.IncMax(1)
110 1 : return testkeys.Suffix(int64(kg.cfg.writeSuffixDist.Max()))
111 1 : }
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 1 : func (kg *keyGenerator) UniformSuffix() []byte {
119 1 : if suffix := kg.UniformSuffixInt(); suffix != 0 {
120 1 : return testkeys.Suffix(suffix)
121 1 : }
122 1 : return nil
123 : }
124 :
125 : // UniformSuffixInt is a variant of UniformSuffix which returns the unencoded
126 : // suffix as an integer.
127 1 : func (kg *keyGenerator) UniformSuffixInt() int64 {
128 1 : maxVal := kg.cfg.writeSuffixDist.Max()
129 1 : return kg.rng.Int63n(int64(maxVal))
130 1 : }
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 1 : func (kg *keyGenerator) randKey(newKeyProbability float64, bounds *pebble.KeyRange) []byte {
136 1 : var knownKeys [][]byte
137 1 : if bounds == nil {
138 1 : knownKeys = kg.keyManager.knownKeys()
139 1 : } else {
140 1 : if kg.cmp(bounds.Start, bounds.End) >= 0 {
141 0 : panic(fmt.Sprintf("invalid bounds [%q, %q)", bounds.Start, bounds.End))
142 : }
143 1 : knownKeys = kg.keyManager.knownKeysInRange(*bounds)
144 : }
145 1 : switch {
146 1 : case len(knownKeys) > 0 && kg.rng.Float64() > newKeyProbability:
147 1 : // Use an existing user key.
148 1 : return pickOneUniform(kg.rng, knownKeys)
149 :
150 1 : case len(knownKeys) > 0 && kg.rng.Float64() > kg.cfg.newPrefix:
151 1 : // Use an existing prefix but a new suffix, producing a new user key.
152 1 : prefixes := kg.keyManager.prefixes()
153 1 :
154 1 : // If we're constrained to a key range, find which existing prefixes
155 1 : // fall within that key range.
156 1 : if bounds != nil {
157 1 : s, _ := slices.BinarySearchFunc(prefixes, bounds.Start, kg.cmp)
158 1 : e, _ := slices.BinarySearchFunc(prefixes, bounds.End, kg.cmp)
159 1 : prefixes = prefixes[s:e]
160 1 : }
161 :
162 1 : if len(prefixes) > 0 {
163 1 : for {
164 1 : // Pick a prefix on each iteration in case most or all suffixes are
165 1 : // already in use for any individual prefix.
166 1 : p := kg.rng.Intn(len(prefixes))
167 1 : suffix := int64(kg.cfg.writeSuffixDist.Uint64(kg.rng))
168 1 :
169 1 : var key []byte
170 1 : if suffix > 0 {
171 1 : key = resizeBuffer(key, len(prefixes[p]), testkeys.SuffixLen(suffix))
172 1 : n := copy(key, prefixes[p])
173 1 : testkeys.WriteSuffix(key[n:], suffix)
174 1 : } else {
175 1 : key = resizeBuffer(key, len(prefixes[p]), 0)
176 1 : copy(key, prefixes[p])
177 1 : }
178 :
179 1 : if bounds == nil || (kg.cmp(key, bounds.Start) >= 0 && kg.cmp(key, bounds.End) < 0) {
180 1 : if kg.keyManager.addNewKey(key) {
181 1 : return key
182 1 : }
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 1 : kg.cfg.writeSuffixDist.IncMax(1)
189 : }
190 : }
191 : // Otherwise fall through to generating a new prefix.
192 : }
193 :
194 1 : if bounds == nil {
195 1 : suffix := kg.SkewedSuffixInt(0.01)
196 1 : for {
197 1 : key := kg.generateKeyWithSuffix(4, 12, suffix)
198 1 : if !kg.keyManager.prefixExists(kg.prefix(key)) {
199 1 : if !kg.keyManager.addNewKey(key) {
200 0 : panic("key must not exist if prefix doesn't exist")
201 : }
202 1 : return key
203 : }
204 : }
205 : }
206 : // We need to generate a key between the bounds.
207 1 : startPrefix, startSuffix := kg.parseKey(bounds.Start)
208 1 : endPrefix, endSuffix := kg.parseKey(bounds.End)
209 1 : var prefix []byte
210 1 : var suffix int64
211 1 : 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 1 : } else {
229 1 : prefix = testkeys.RandomPrefixInRange(startPrefix, endPrefix, kg.rng)
230 1 : suffix = kg.SkewedSuffixInt(0.01)
231 1 : 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 1 : key := slices.Clip(prefix)
243 1 : if suffix != 0 {
244 1 : key = append(key, testkeys.Suffix(suffix)...)
245 1 : }
246 1 : 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 1 : kg.keyManager.addNewKey(key)
251 1 : 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 1 : func (kg *keyGenerator) generateKeyWithSuffix(minPrefixLen, maxPrefixLen int, suffix int64) []byte {
257 1 : prefix := randBytes(kg.rng, minPrefixLen, maxPrefixLen)
258 1 : if suffix == 0 {
259 1 : return prefix
260 1 : }
261 1 : 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 1 : func (kg *keyGenerator) cmp(a, b []byte) int {
279 1 : return kg.keyManager.comparer.Compare(a, b)
280 1 : }
281 :
282 1 : func (kg *keyGenerator) equal(a, b []byte) bool {
283 1 : return kg.keyManager.comparer.Equal(a, b)
284 1 : }
285 :
286 1 : func (kg *keyGenerator) split(a []byte) int {
287 1 : return kg.keyManager.comparer.Split(a)
288 1 : }
289 :
290 1 : func (kg *keyGenerator) prefix(a []byte) []byte {
291 1 : n := kg.split(a)
292 1 : return a[:n:n]
293 1 : }
294 :
295 1 : func (kg *keyGenerator) parseKey(k []byte) (prefix []byte, suffix int64) {
296 1 : n := kg.split(k)
297 1 : if n == len(k) {
298 1 : return k, 0
299 1 : }
300 1 : suffix, err := testkeys.ParseSuffix(k[n:])
301 1 : if err != nil {
302 0 : panic(fmt.Sprintf("error parsing suffix for key %q", k))
303 : }
304 1 : return k[:n:n], suffix
305 : }
306 :
307 1 : func randBytes(rng *rand.Rand, minLen, maxLen int) []byte {
308 1 : n := minLen + rng.Intn(maxLen-minLen+1)
309 1 : if n == 0 {
310 1 : return nil
311 1 : }
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 1 : const letters = "abcdefghijklmnopqrstuvwxyz"
317 1 : const lettersLen = uint64(len(letters))
318 1 : const lettersCharsPerRand = 12 // floor(log(math.MaxUint64)/log(lettersLen))
319 1 :
320 1 : var r uint64
321 1 : var q int
322 1 : buf := make([]byte, n)
323 1 : for i := range buf {
324 1 : if q == 0 {
325 1 : r = rng.Uint64()
326 1 : q = lettersCharsPerRand
327 1 : }
328 1 : buf[i] = letters[r%lettersLen]
329 1 : r = r / lettersLen
330 1 : q--
331 : }
332 1 : return buf
333 : }
334 :
335 1 : func pickOneUniform[S ~[]E, E any](rng *rand.Rand, x S) E {
336 1 : return x[rng.Intn(len(x))]
337 1 : }
|
