Line data Source code
1 : // Copyright 2023 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 keyspan
6 :
7 : import (
8 : "fmt"
9 :
10 : "github.com/cockroachdb/errors"
11 : "github.com/cockroachdb/pebble/internal/base"
12 : "github.com/cockroachdb/pebble/internal/invariants"
13 : )
14 :
15 : // Assert wraps an iterator and asserts that operations return sane results.
16 1 : func Assert(iter FragmentIterator, cmp base.Compare) FragmentIterator {
17 1 : return &assertIter{
18 1 : iter: iter,
19 1 : cmp: cmp,
20 1 : }
21 1 : }
22 :
23 : // MaybeAssert wraps an iterator and asserts that operations return sane
24 : // results if we are in testing mode.
25 1 : func MaybeAssert(iter FragmentIterator, cmp base.Compare) FragmentIterator {
26 1 : if invariants.Enabled && iter != nil {
27 1 : // Don't wrap an assertIter.
28 1 : if _, ok := iter.(*assertIter); !ok {
29 1 : return Assert(iter, cmp)
30 1 : }
31 : }
32 1 : return iter
33 : }
34 :
35 : // AssertUserKeyBounds wraps an iterator and asserts that all spans are within
36 : // the given bounds [lower, upper).
37 : func AssertUserKeyBounds(
38 : iter FragmentIterator, lower, upper []byte, cmp base.Compare,
39 0 : ) FragmentIterator {
40 0 : return AssertBounds(iter, base.MakeSearchKey(lower), upper, cmp)
41 0 : }
42 :
43 : // AssertBounds wraps an iterator and asserts that all spans are within the
44 : // given bounds [lower.UserKey, upper), and that all keys in a span that starts
45 : // exactly at lower.UserKey are >= lower.
46 : //
47 : // The asymmetry here is due to fragment spans having exclusive end user keys.
48 : func AssertBounds(
49 : iter FragmentIterator, lower base.InternalKey, upper []byte, cmp base.Compare,
50 1 : ) FragmentIterator {
51 1 : i := &assertIter{
52 1 : iter: iter,
53 1 : cmp: cmp,
54 1 : }
55 1 : i.checkBounds.enabled = true
56 1 : i.checkBounds.lower = lower
57 1 : i.checkBounds.upper = upper
58 1 : return i
59 1 : }
60 :
61 : // assertIter is a pass-through FragmentIterator wrapper which performs checks
62 : // on what the wrapped iterator returns.
63 : //
64 : // It verifies that results for various operations are sane, and it optionally
65 : // verifies that spans are within given bounds.
66 : type assertIter struct {
67 : iter FragmentIterator
68 : cmp base.Compare
69 : checkBounds struct {
70 : enabled bool
71 : lower base.InternalKey
72 : upper []byte
73 : }
74 : lastSpanStart []byte
75 : lastSpanEnd []byte
76 : }
77 :
78 : var _ FragmentIterator = (*assertIter)(nil)
79 :
80 0 : func (i *assertIter) panicf(format string, args ...interface{}) {
81 0 : str := fmt.Sprintf(format, args...)
82 0 : panic(errors.AssertionFailedf("%s; wraps %T", str, i.iter))
83 : }
84 :
85 1 : func (i *assertIter) check(span *Span) {
86 1 : i.lastSpanStart = i.lastSpanStart[:0]
87 1 : i.lastSpanEnd = i.lastSpanEnd[:0]
88 1 : if span == nil {
89 1 : return
90 1 : }
91 1 : if i.checkBounds.enabled {
92 1 : lower := i.checkBounds.lower
93 1 : switch startCmp := i.cmp(span.Start, lower.UserKey); {
94 0 : case startCmp < 0:
95 0 : i.panicf("lower bound %q violated by span %s", lower.UserKey, span)
96 1 : case startCmp == 0:
97 1 : // Note: trailers are in descending order.
98 1 : if len(span.Keys) > 0 && span.SmallestKey().Trailer > lower.Trailer {
99 0 : i.panicf("lower bound %s violated by key %s", lower, span.SmallestKey())
100 0 : }
101 : }
102 1 : if i.cmp(span.End, i.checkBounds.upper) > 0 {
103 0 : i.panicf("upper bound %q violated by span %s", i.checkBounds.upper, span)
104 0 : }
105 : }
106 : // Save the span to check Next/Prev operations.
107 1 : i.lastSpanStart = append(i.lastSpanStart, span.Start...)
108 1 : i.lastSpanEnd = append(i.lastSpanEnd, span.End...)
109 : }
110 :
111 : // SeekGE implements FragmentIterator.
112 1 : func (i *assertIter) SeekGE(key []byte) (*Span, error) {
113 1 : span, err := i.iter.SeekGE(key)
114 1 : if span != nil && i.cmp(span.End, key) <= 0 {
115 0 : i.panicf("incorrect SeekGE(%q) span %s", key, span)
116 0 : }
117 1 : i.check(span)
118 1 : return span, err
119 : }
120 :
121 : // SeekLT implements FragmentIterator.
122 1 : func (i *assertIter) SeekLT(key []byte) (*Span, error) {
123 1 : span, err := i.iter.SeekLT(key)
124 1 : if span != nil && i.cmp(span.Start, key) >= 0 {
125 0 : i.panicf("incorrect SeekLT(%q) span %s", key, span)
126 0 : }
127 1 : i.check(span)
128 1 : return span, err
129 : }
130 :
131 : // First implements FragmentIterator.
132 1 : func (i *assertIter) First() (*Span, error) {
133 1 : span, err := i.iter.First()
134 1 : i.check(span)
135 1 : return span, err
136 1 : }
137 :
138 : // Last implements FragmentIterator.
139 1 : func (i *assertIter) Last() (*Span, error) {
140 1 : span, err := i.iter.Last()
141 1 : i.check(span)
142 1 : return span, err
143 1 : }
144 :
145 : // Next implements FragmentIterator.
146 1 : func (i *assertIter) Next() (*Span, error) {
147 1 : span, err := i.iter.Next()
148 1 : if span != nil && len(i.lastSpanEnd) > 0 && i.cmp(i.lastSpanEnd, span.Start) > 0 {
149 0 : i.panicf("Next span %s not after last span end %q", span, i.lastSpanEnd)
150 0 : }
151 1 : i.check(span)
152 1 : return span, err
153 : }
154 :
155 : // Prev implements FragmentIterator.
156 1 : func (i *assertIter) Prev() (*Span, error) {
157 1 : span, err := i.iter.Prev()
158 1 : if span != nil && len(i.lastSpanStart) > 0 && i.cmp(i.lastSpanStart, span.End) < 0 {
159 0 : i.panicf("Prev span %s not before last span start %q", span, i.lastSpanStart)
160 0 : }
161 1 : i.check(span)
162 1 : return span, err
163 : }
164 :
165 : // Close implements FragmentIterator.
166 1 : func (i *assertIter) Close() error {
167 1 : return i.iter.Close()
168 1 : }
169 :
170 : // WrapChildren implements FragmentIterator.
171 0 : func (i *assertIter) WrapChildren(wrap WrapFn) {
172 0 : i.iter = wrap(i.iter)
173 0 : }
|