1// SPDX-License-Identifier: MIT
2
3package kv
4
5import (
6 "bytes"
7 "fmt"
8 "math/rand/v2"
9 "slices"
10 "sync"
11 "sync/atomic"
12 "testing"
13)
14
15// model is the reference implementation: a map plus a sort on demand.
16type model map[string]Item
17
18func (m model) clone() model {
19 c := make(model, len(m))
20 for k, v := range m {
21 c[k] = v
22 }
23 return c
24}
25
26// sorted returns the keys of m in ascending order.
27func (m model) sorted() []string {
28 ks := make([]string, 0, len(m))
29 for k := range m {
30 ks = append(ks, k)
31 }
32 slices.Sort(ks)
33 return ks
34}
35
36// within returns the part of the sorted keys in [start, end); nil end = unbounded.
37func within(sorted []string, start, end []byte) []string {
38 lo, _ := slices.BinarySearch(sorted, string(start))
39 hi := len(sorted)
40 if end != nil {
41 hi, _ = slices.BinarySearch(sorted, string(end))
42 }
43 if hi < lo {
44 return nil
45 }
46 return sorted[lo:hi]
47}
48
49type kvPair struct {
50 key string
51 item Item
52}
53
54func collect(seq func(func([]byte, Item) bool)) []kvPair {
55 var out []kvPair
56 seq(func(k []byte, it Item) bool {
57 out = append(out, kvPair{string(k), it})
58 return true
59 })
60 return out
61}
62
63func equalItem(a, b Item) bool {
64 return a.Version == b.Version && bytes.Equal(a.Value, b.Value)
65}
66
67// checkStructure verifies the B-tree invariants: ordering, fill bounds,
68// uniform leaf depth and the entry count.
69func checkStructure(t *testing.T, tr tree) {
70 t.Helper()
71 if tr.root == nil {
72 if tr.n != 0 {
73 t.Fatalf("empty root but n=%d", tr.n)
74 }
75 return
76 }
77 var count, leafDepth int
78 var prev []byte
79 var walk func(n *node, depth int, isRoot bool)
80 walk = func(n *node, depth int, isRoot bool) {
81 if len(n.entries) == 0 || len(n.entries) > maxEntries || (!isRoot && len(n.entries) < minEntries) {
82 t.Fatalf("node with %d entries at depth %d", len(n.entries), depth)
83 }
84 if n.leaf() {
85 if leafDepth == 0 {
86 leafDepth = depth
87 } else if leafDepth != depth {
88 t.Fatalf("leaves at depths %d and %d", leafDepth, depth)
89 }
90 } else if len(n.children) != len(n.entries)+1 {
91 t.Fatalf("%d entries but %d children", len(n.entries), len(n.children))
92 }
93 for i, e := range n.entries {
94 if !n.leaf() {
95 walk(n.children[i], depth+1, false)
96 }
97 if prev != nil && bytes.Compare(prev, e.key) >= 0 {
98 t.Fatalf("keys out of order: %x then %x", prev, e.key)
99 }
100 prev = e.key
101 count++
102 }
103 if !n.leaf() {
104 walk(n.children[len(n.entries)], depth+1, false)
105 }
106 }
107 walk(tr.root, 1, true)
108 if count != tr.n {
109 t.Fatalf("tree has %d entries, n=%d", count, tr.n)
110 }
111}
112
113// checkAgainst compares the tree with the model through get, full scan,
114// random ranges and random prefixes.
115func checkAgainst(t *testing.T, rng *rand.Rand, tr tree, m model, genKey func() []byte) {
116 t.Helper()
117 checkStructure(t, tr)
118 if tr.len() != len(m) {
119 t.Fatalf("len = %d, want %d", tr.len(), len(m))
120 }
121 s := newSnapshot(tr, 7)
122 if s.Len() != len(m) || s.Seq() != 7 {
123 t.Fatalf("snapshot Len/Seq = %d/%d", s.Len(), s.Seq())
124 }
125 sorted := m.sorted()
126 verify := func(what string, got []kvPair, start, end []byte) {
127 t.Helper()
128 want := within(sorted, start, end)
129 if len(got) != len(want) {
130 t.Fatalf("%s: got %d items, want %d", what, len(got), len(want))
131 }
132 for i, k := range want {
133 if got[i].key != k || !equalItem(got[i].item, m[k]) {
134 t.Fatalf("%s: item %d = %x, want %x", what, i, got[i].key, k)
135 }
136 }
137 }
138 verify("full scan", collect(s.Range(bg, nil, nil)), nil, nil)
139
140 for range 20 {
141 k := genKey()
142 got, ok := s.Get(bg, k)
143 want, wok := m[string(k)]
144 if ok != wok || (ok && !equalItem(got, want)) {
145 t.Fatalf("Get(%x) = %v,%v want %v,%v", k, got, ok, want, wok)
146 }
147 }
148 for k, want := range m {
149 if got, ok := s.Get(bg, []byte(k)); !ok || !equalItem(got, want) {
150 t.Fatalf("Get(%x) = %v,%v want %v", k, got, ok, want)
151 }
152 break // one arbitrary present key is enough per check
153 }
154 for range 10 {
155 a, b := genKey(), genKey()
156 if bytes.Compare(a, b) > 0 {
157 a, b = b, a
158 }
159 verify("range", collect(s.Range(bg, a, b)), a, b)
160 verify("range open end", collect(s.Range(bg, a, nil)), a, nil)
161 verify("range open start", collect(s.Range(bg, nil, b)), nil, b)
162 }
163 for range 10 {
164 p := genKey()
165 p = p[:rng.IntN(len(p)+1)] // may be empty
166 var want []string
167 for _, k := range sorted {
168 if bytes.HasPrefix([]byte(k), p) {
169 want = append(want, k)
170 }
171 }
172 got := collect(s.Prefix(bg, p))
173 if len(got) != len(want) {
174 t.Fatalf("prefix %x: got %d items, want %d", p, len(got), len(want))
175 }
176 for i, k := range want {
177 if got[i].key != k {
178 t.Fatalf("prefix %x: item %d = %x, want %x", p, i, got[i].key, k)
179 }
180 }
181 }
182 // Early break stops the walk.
183 seen := 0
184 for range s.Range(bg, nil, nil) {
185 seen++
186 if seen == 3 {
187 break
188 }
189 }
190 if seen != min(3, len(m)) {
191 t.Fatalf("early break saw %d items", seen)
192 }
193}
194
195type frozen struct {
196 tr tree
197 m model
198}
199
200func TestTreeModel(t *testing.T) {
201 seeds, ops := 50, 10000
202 if testing.Short() {
203 seeds, ops = 8, 2000
204 }
205 alphabet := []byte{0, 1, 2, 'a', 'b', 'c', 0x7F, 0x80, 0xFD, 0xFE, 0xFF, 'x', 'y', 'z', 0x10, 0x20}
206 for seed := range seeds {
207 t.Run(fmt.Sprintf("seed=%d", seed), func(t *testing.T) {
208 rng := rand.New(rand.NewPCG(uint64(seed), 0xe131))
209 alpha := alphabet[:2+seed%(len(alphabet)-1)]
210 genKey := func() []byte {
211 k := make([]byte, 1+rng.IntN(5))
212 for i := range k {
213 k[i] = alpha[rng.IntN(len(alpha))]
214 }
215 return k
216 }
217 var (
218 tr tree
219 m = model{}
220 version uint64
221 history []frozen
222 )
223 apply := func(del bool, k []byte, put func(k, v []byte, ver uint64), rm func(k []byte)) {
224 version++
225 if del {
226 rm(k)
227 delete(m, string(k))
228 return
229 }
230 v := make([]byte, rng.IntN(8))
231 for i := range v {
232 v[i] = byte(rng.IntN(256))
233 }
234 put(k, v, version)
235 m[string(k)] = Item{Value: v, Version: version}
236 }
237 for done, next := 0, 0; done < ops; {
238 if rng.IntN(5) == 0 {
239 // A group of ops through one builder.
240 b := tr.builder()
241 n := 1 + rng.IntN(40)
242 for range n {
243 k := genKey()
244 apply(rng.IntN(3) == 0, k, b.put, b.delete)
245 if got, ok := b.get(k); ok != (m[string(k)].Version != 0) || (ok && !equalItem(got, m[string(k)])) {
246 t.Fatalf("seed %d: builder.get mismatch", seed)
247 }
248 if rng.IntN(512) == 0 {
249 // Freeze mid-group and keep using the builder: the
250 // frozen tree must never see later builder writes.
251 history = append(history, frozen{b.freeze(), m.clone()})
252 }
253 }
254 if b.len() != len(m) {
255 t.Fatalf("seed %d: builder.len = %d, want %d", seed, b.len(), len(m))
256 }
257 tr = b.freeze()
258 done += n
259 } else {
260 k := genKey()
261 apply(rng.IntN(3) == 0, k, func(k, v []byte, ver uint64) { tr = tr.put(k, v, ver) }, func(k []byte) { tr = tr.delete(k) })
262 done++
263 }
264 if done >= next {
265 next = done + 500 + rng.IntN(500)
266 checkAgainst(t, rng, tr, m, genKey)
267 history = append(history, frozen{tr, m.clone()})
268 }
269 }
270 checkAgainst(t, rng, tr, m, genKey)
271 // Persistence: every older tree still equals its model.
272 for _, h := range history {
273 checkAgainst(t, rng, h.tr, h.m, genKey)
274 }
275 // Drain the tree in random order, with and without a builder.
276 keys := m.sorted()
277 rng.Shuffle(len(keys), func(i, j int) { keys[i], keys[j] = keys[j], keys[i] })
278 for i, k := range keys {
279 if i%2 == 0 {
280 tr = tr.delete([]byte(k))
281 } else {
282 b := tr.builder()
283 b.delete([]byte(k))
284 b.delete([]byte(k)) // absent: no-op
285 tr = b.freeze()
286 }
287 delete(m, k)
288 if i%97 == 0 {
289 checkStructure(t, tr)
290 }
291 }
292 checkAgainst(t, rng, tr, m, genKey)
293 if tr.root != nil || tr.len() != 0 {
294 t.Fatalf("seed %d: tree not empty after draining", seed)
295 }
296 // Draining must not have disturbed the older trees.
297 for _, h := range history {
298 checkAgainst(t, rng, h.tr, h.m, genKey)
299 }
300 })
301 }
302}
303
304// TestBuilderReuse checks that nodes of a frozen tree are never mutated by
305// later builder operations, and that nodes owned by the builder are reused.
306func TestBuilderReuse(t *testing.T) {
307 var base tree
308 for i := range 500 {
309 base = base.put(fmt.Appendf(nil, "k%04d", i), []byte("v"), 1)
310 }
311 b := base.builder()
312 b.put([]byte("k0100"), []byte("new"), 2)
313 first := b.root
314 b.put([]byte("k0101"), []byte("new"), 2)
315 if b.root != first {
316 t.Fatal("builder copied its own root twice")
317 }
318 frozenTree := b.freeze()
319 b.put([]byte("k0102"), []byte("again"), 3)
320 if b.root == frozenTree.root {
321 t.Fatal("builder mutated the root of a frozen tree")
322 }
323 if it, _ := frozenTree.get([]byte("k0102")); string(it.Value) != "v" {
324 t.Fatalf("frozen tree changed: %q", it.Value)
325 }
326 if it, _ := base.get([]byte("k0100")); string(it.Value) != "v" || it.Version != 1 {
327 t.Fatalf("base tree changed: %v", it)
328 }
329 checkStructure(t, base)
330 checkStructure(t, frozenTree)
331 checkStructure(t, b.freeze())
332}
333
334func TestPrefixEnd(t *testing.T) {
335 tests := []struct{ in, want []byte }{
336 {nil, nil},
337 {[]byte{}, nil},
338 {[]byte("a"), []byte("b")},
339 {[]byte("ab"), []byte("ac")},
340 {[]byte{'a', 0xFF}, []byte("b")},
341 {[]byte{'a', 0xFF, 0xFF}, []byte("b")},
342 {[]byte{0xFF}, nil},
343 {[]byte{0xFF, 0xFF}, nil},
344 }
345 for _, tt := range tests {
346 if got := prefixEnd(tt.in); !bytes.Equal(got, tt.want) || (got == nil) != (tt.want == nil) {
347 t.Errorf("prefixEnd(%x) = %x, want %x", tt.in, got, tt.want)
348 }
349 }
350}
351
352func TestPrefixAllFF(t *testing.T) {
353 var tr tree
354 for i, k := range [][]byte{{0xFE}, {0xFF}, {0xFF, 0x00}, {0xFF, 0xFF, 0xFF}} {
355 tr = tr.put(k, nil, uint64(i+1))
356 }
357 s := newSnapshot(tr, 4)
358 if got := collect(s.Prefix(bg, []byte{0xFF})); len(got) != 3 {
359 t.Fatalf("prefix FF: %d items, want 3", len(got))
360 }
361 if got := collect(s.Prefix(bg, nil)); len(got) != 4 {
362 t.Fatalf("empty prefix: %d items, want 4", len(got))
363 }
364}
365
366func benchTree(n int) (*Snapshot, [][]byte) {
367 keys := make([][]byte, n)
368 b := tree{}.builder()
369 for i := range keys {
370 keys[i] = fmt.Appendf(nil, "key/%08d", i)
371 b.put(keys[i], []byte("value"), uint64(i+1))
372 }
373 return newSnapshot(b.freeze(), uint64(n)), keys
374}
375
376func BenchmarkSnapshotGet(b *testing.B) {
377 s, keys := benchTree(100_000)
378 b.ReportAllocs()
379 b.ResetTimer()
380 for i := 0; b.Loop(); i++ {
381 if _, ok := s.Get(bg, keys[(i*7919)%len(keys)]); !ok {
382 b.Fatal("missing key")
383 }
384 }
385}
386
387func BenchmarkRange(b *testing.B) {
388 s, _ := benchTree(100_000)
389 b.ReportAllocs()
390 for b.Loop() {
391 n := 0
392 for range s.Range(bg, nil, nil) {
393 n++
394 }
395 if n != 100_000 {
396 b.Fatalf("saw %d items", n)
397 }
398 }
399}
400
401func BenchmarkTreePut(b *testing.B) {
402 _, keys := benchTree(100_000)
403 b.ReportAllocs()
404 var tr tree
405 for i := 0; b.Loop(); i++ {
406 tr = tr.put(keys[(i*7919)%len(keys)], nil, uint64(i))
407 }
408}
409
410func BenchmarkBuilderPut(b *testing.B) {
411 _, keys := benchTree(100_000)
412 b.ReportAllocs()
413 bl := tree{}.builder()
414 for i := 0; b.Loop(); i++ {
415 bl.put(keys[(i*7919)%len(keys)], nil, uint64(i))
416 if i%64 == 63 {
417 bl.freeze()
418 }
419 }
420}
421
422// TestTreeDeep builds a tree of several levels, then deletes ascending,
423// exercising merges and borrows on internal nodes.
424func TestTreeDeep(t *testing.T) {
425 const n = 20000
426 var tr tree
427 for i := range n {
428 tr = tr.put(fmt.Appendf(nil, "%06d", i), nil, 1)
429 }
430 checkStructure(t, tr)
431 depth := 0
432 for x := tr.root; x != nil; depth++ {
433 if x.leaf() {
434 depth++
435 break
436 }
437 x = x.children[0]
438 }
439 if depth < 3 {
440 t.Fatalf("depth = %d, want >= 3", depth)
441 }
442 full := tr
443 for i := range n {
444 tr = tr.delete(fmt.Appendf(nil, "%06d", i))
445 if i%251 == 0 {
446 checkStructure(t, tr)
447 }
448 }
449 if tr.len() != 0 || tr.root != nil {
450 t.Fatalf("tree not empty: len=%d", tr.len())
451 }
452 checkStructure(t, full)
453}
454
455// TestConcurrentReaders reads published snapshots without locks while one
456// builder keeps writing and freezing, so -race sees any shared mutation.
457func TestConcurrentReaders(t *testing.T) {
458 var cur atomic.Pointer[Snapshot]
459 cur.Store(newSnapshot(tree{}, 0))
460 stop := make(chan struct{})
461 var wg sync.WaitGroup
462 for range 8 {
463 wg.Go(func() {
464 for {
465 select {
466 case <-stop:
467 return
468 default:
469 }
470 s := cur.Load()
471 n := 0
472 var prev []byte
473 for k, it := range s.Range(bg, nil, nil) {
474 if prev != nil && bytes.Compare(prev, k) >= 0 {
475 t.Errorf("seq %d: keys out of order", s.Seq())
476 return
477 }
478 if it.Version > s.Seq() {
479 t.Errorf("seq %d: version %d from the future", s.Seq(), it.Version)
480 return
481 }
482 prev = k
483 n++
484 }
485 if n != s.Len() {
486 t.Errorf("seq %d: ranged %d keys, Len %d", s.Seq(), n, s.Len())
487 return
488 }
489 }
490 })
491 }
492 rng := rand.New(rand.NewPCG(1, 2))
493 b := cur.Load().t.builder()
494 for seq := uint64(1); seq <= 3000; seq++ {
495 for range 1 + rng.IntN(8) {
496 k := fmt.Appendf(nil, "k%04d", rng.IntN(2000))
497 if rng.IntN(3) == 0 {
498 b.delete(k)
499 } else {
500 b.put(k, []byte("v"), seq)
501 }
502 }
503 cur.Store(newSnapshot(b.freeze(), seq))
504 }
505 close(stop)
506 wg.Wait()
507}