Line data Source code
1 : // Copyright 2022 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 : "context"
9 :
10 : "github.com/cockroachdb/pebble/internal/base"
11 : "github.com/cockroachdb/pebble/internal/treeprinter"
12 : )
13 :
14 : // TODO(jackson): Consider removing this type and adding bounds enforcement
15 : // directly to the MergingIter. This type is probably too lightweight to warrant
16 : // its own type, but for now we implement it separately for expediency.
17 :
18 : // boundedIterPos records the position of the BoundedIter relative to the
19 : // underlying iterator's position. It's used to avoid Next/Prev-ing the iterator
20 : // if there can't possibly be another span within bounds, because the current
21 : // span overlaps the bound.
22 : //
23 : // Imagine bounds [a,c) and an iterator that seeks to a span [b,d). The span
24 : // [b,d) overlaps some portion of the iterator bounds, so the iterator must
25 : // return it. If the iterator is subsequently Nexted, Next can tell that the
26 : // iterator is exhausted without advancing the underlying iterator because the
27 : // current span's end bound of d is ≥ the upper bound of c. In this case, the
28 : // bounded iterator returns nil and records i.pos as posAtUpperLimit to remember
29 : // that the underlying iterator position does not match the current BoundedIter
30 : // position.
31 : type boundedIterPos int8
32 :
33 : const (
34 : posAtLowerLimit boundedIterPos = -1
35 : posAtIterSpan boundedIterPos = 0
36 : posAtUpperLimit boundedIterPos = +1
37 : )
38 :
39 : // BoundedIter implements FragmentIterator and enforces bounds.
40 : //
41 : // Like the point InternalIterator interface, the bounded iterator's forward
42 : // positioning routines (SeekGE, First, and Next) only check the upper bound.
43 : // The reverse positioning routines (SeekLT, Last, and Prev) only check the
44 : // lower bound. It is up to the caller to ensure that the forward positioning
45 : // routines respect the lower bound and the reverse positioning routines respect
46 : // the upper bound (i.e. calling SeekGE instead of First if there is a lower
47 : // bound, and SeekLT instead of Last if there is an upper bound).
48 : //
49 : // When the hasPrefix parameter indicates that the iterator is in prefix
50 : // iteration mode, BoundedIter elides any spans that do not overlap with the
51 : // prefix's keyspace. In prefix iteration mode, reverse iteration is disallowed,
52 : // except for an initial SeekLT with a seek key greater than or equal to the
53 : // prefix. In prefix iteration mode, the first seek must position the iterator
54 : // at or immediately before the first fragment covering a key greater than or
55 : // equal to the prefix.
56 : type BoundedIter struct {
57 : iter FragmentIterator
58 : iterSpan *Span
59 : cmp base.Compare
60 : split base.Split
61 : lower []byte
62 : upper []byte
63 : hasPrefix *bool
64 : prefix *[]byte
65 : pos boundedIterPos
66 : }
67 :
68 : // Init initializes the bounded iterator.
69 : //
70 : // In addition to the iterator bounds, Init takes pointers to a boolean
71 : // indicating whether the iterator is in prefix iteration mode and the prefix
72 : // key if it is. This is used to exclude spans that are outside the iteration
73 : // prefix.
74 : //
75 : // hasPrefix and prefix are allowed to be nil, however if hasPrefix != nil,
76 : // prefix must also not be nil.
77 : func (i *BoundedIter) Init(
78 : cmp base.Compare,
79 : split base.Split,
80 : iter FragmentIterator,
81 : lower, upper []byte,
82 : hasPrefix *bool,
83 : prefix *[]byte,
84 2 : ) {
85 2 : *i = BoundedIter{
86 2 : iter: iter,
87 2 : cmp: cmp,
88 2 : split: split,
89 2 : lower: lower,
90 2 : upper: upper,
91 2 : hasPrefix: hasPrefix,
92 2 : prefix: prefix,
93 2 : }
94 2 : }
95 :
96 : var _ FragmentIterator = (*BoundedIter)(nil)
97 :
98 : // Seek calls.
99 : //
100 : // Seek calls check iterator bounds in the direction of the seek. Additionally,
101 : // if the iterator is in prefix iteration mode, seek calls check both start and
102 : // end bounds against the prefix's bounds. We check both bounds for defense in
103 : // depth. This optimization has been a source of various bugs due to various
104 : // other prefix iteration optimizations that can result in seek keys that don't
105 : // respect the prefix bounds.
106 :
107 : // SeekGE implements FragmentIterator.
108 2 : func (i *BoundedIter) SeekGE(key []byte) (*Span, error) {
109 2 : s, err := i.iter.SeekGE(key)
110 2 : s, err = i.checkPrefixSpanStart(s, err)
111 2 : s, err = i.checkPrefixSpanEnd(s, err)
112 2 : return i.checkForwardBound(s, err)
113 2 : }
114 :
115 : // SeekLT implements FragmentIterator.
116 2 : func (i *BoundedIter) SeekLT(key []byte) (*Span, error) {
117 2 : s, err := i.iter.SeekLT(key)
118 2 : s, err = i.checkPrefixSpanStart(s, err)
119 2 : s, err = i.checkPrefixSpanEnd(s, err)
120 2 : return i.checkBackwardBound(s, err)
121 2 : }
122 :
123 : // First implements FragmentIterator.
124 2 : func (i *BoundedIter) First() (*Span, error) {
125 2 : s, err := i.iter.First()
126 2 : s, err = i.checkPrefixSpanStart(s, err)
127 2 : return i.checkForwardBound(s, err)
128 2 : }
129 :
130 : // Last implements FragmentIterator.
131 2 : func (i *BoundedIter) Last() (*Span, error) {
132 2 : s, err := i.iter.Last()
133 2 : s, err = i.checkPrefixSpanEnd(s, err)
134 2 : return i.checkBackwardBound(s, err)
135 2 : }
136 :
137 : // Next implements FragmentIterator.
138 2 : func (i *BoundedIter) Next() (*Span, error) {
139 2 : switch i.pos {
140 2 : case posAtLowerLimit:
141 2 : // The BoundedIter had previously returned nil, because it knew from
142 2 : // i.iterSpan's bounds that there was no previous span. To Next, we only
143 2 : // need to return the current iter span and reset i.pos to reflect that
144 2 : // we're no longer positioned at the limit.
145 2 : i.pos = posAtIterSpan
146 2 : return i.iterSpan, nil
147 2 : case posAtIterSpan:
148 2 : // If the span at the underlying iterator position extends to or beyond the
149 2 : // upper bound, we can avoid advancing because the next span is necessarily
150 2 : // out of bounds.
151 2 : if i.iterSpan != nil && i.upper != nil && i.cmp(i.iterSpan.End, i.upper) >= 0 {
152 2 : i.pos = posAtUpperLimit
153 2 : return nil, nil
154 2 : }
155 : // Similarly, if the span extends to the next prefix and we're in prefix
156 : // iteration mode, we can avoid advancing.
157 2 : if i.iterSpan != nil && i.hasPrefix != nil && *i.hasPrefix {
158 2 : ei := i.split(i.iterSpan.End)
159 2 : if i.cmp(i.iterSpan.End[:ei], *i.prefix) > 0 {
160 2 : i.pos = posAtUpperLimit
161 2 : return nil, nil
162 2 : }
163 : }
164 2 : return i.checkForwardBound(i.checkPrefixSpanStart(i.iter.Next()))
165 1 : case posAtUpperLimit:
166 1 : // Already exhausted.
167 1 : return nil, nil
168 0 : default:
169 0 : panic("unreachable")
170 : }
171 : }
172 :
173 : // Prev implements FragmentIterator.
174 2 : func (i *BoundedIter) Prev() (*Span, error) {
175 2 : switch i.pos {
176 1 : case posAtLowerLimit:
177 1 : // Already exhausted.
178 1 : return nil, nil
179 2 : case posAtIterSpan:
180 2 : // If the span at the underlying iterator position extends to or beyond
181 2 : // the lower bound, we can avoid advancing because the previous span is
182 2 : // necessarily out of bounds.
183 2 : if i.iterSpan != nil && i.lower != nil && i.cmp(i.iterSpan.Start, i.lower) <= 0 {
184 2 : i.pos = posAtLowerLimit
185 2 : return nil, nil
186 2 : }
187 : // Similarly, if the span extends to or beyond the current prefix and
188 : // we're in prefix iteration mode, we can avoid advancing.
189 2 : if i.iterSpan != nil && i.hasPrefix != nil && *i.hasPrefix {
190 2 : si := i.split(i.iterSpan.Start)
191 2 : if i.cmp(i.iterSpan.Start[:si], *i.prefix) < 0 {
192 2 : i.pos = posAtLowerLimit
193 2 : return nil, nil
194 2 : }
195 : }
196 2 : return i.checkBackwardBound(i.checkPrefixSpanEnd(i.iter.Prev()))
197 2 : case posAtUpperLimit:
198 2 : // The BoundedIter had previously returned nil, because it knew from
199 2 : // i.iterSpan's bounds that there was no next span. To Prev, we only
200 2 : // need to return the current iter span and reset i.pos to reflect that
201 2 : // we're no longer positioned at the limit.
202 2 : i.pos = posAtIterSpan
203 2 : return i.iterSpan, nil
204 0 : default:
205 0 : panic("unreachable")
206 : }
207 : }
208 :
209 : // Close implements FragmentIterator.
210 2 : func (i *BoundedIter) Close() {
211 2 : i.iter.Close()
212 2 : }
213 :
214 : // SetContext is part of the FragmentIterator interface.
215 0 : func (i *BoundedIter) SetContext(ctx context.Context) {
216 0 : i.iter.SetContext(ctx)
217 0 : }
218 :
219 : // SetBounds modifies the FragmentIterator's bounds.
220 2 : func (i *BoundedIter) SetBounds(lower, upper []byte) {
221 2 : i.lower, i.upper = lower, upper
222 2 : }
223 :
224 2 : func (i *BoundedIter) checkPrefixSpanStart(span *Span, err error) (*Span, error) {
225 2 : // Compare to the prefix's bounds, if in prefix iteration mode.
226 2 : if span != nil && i.hasPrefix != nil && *i.hasPrefix {
227 2 : si := i.split(span.Start)
228 2 : if i.cmp(span.Start[:si], *i.prefix) > 0 {
229 2 : // This span starts at a prefix that sorts after our current prefix.
230 2 : span = nil
231 2 : }
232 : }
233 2 : return span, err
234 : }
235 :
236 : // checkForwardBound enforces the upper bound, returning nil if the provided
237 : // span is wholly outside the upper bound. It also updates i.pos and i.iterSpan
238 : // to reflect the new iterator position.
239 2 : func (i *BoundedIter) checkForwardBound(span *Span, err error) (*Span, error) {
240 2 : // Compare to the upper bound.
241 2 : if span != nil && i.upper != nil && i.cmp(span.Start, i.upper) >= 0 {
242 2 : span = nil
243 2 : }
244 2 : i.iterSpan = span
245 2 : if i.pos != posAtIterSpan {
246 2 : i.pos = posAtIterSpan
247 2 : }
248 2 : return span, err
249 : }
250 :
251 2 : func (i *BoundedIter) checkPrefixSpanEnd(span *Span, err error) (*Span, error) {
252 2 : // Compare to the prefix's bounds, if in prefix iteration mode.
253 2 : if span != nil && i.hasPrefix != nil && *i.hasPrefix && i.cmp(span.End, *i.prefix) <= 0 {
254 2 : // This span ends before the current prefix.
255 2 : span = nil
256 2 : }
257 2 : return span, err
258 : }
259 :
260 : // checkBackward enforces the lower bound, returning nil if the provided span is
261 : // wholly outside the lower bound. It also updates i.pos and i.iterSpan to
262 : // reflect the new iterator position.
263 2 : func (i *BoundedIter) checkBackwardBound(span *Span, err error) (*Span, error) {
264 2 : // Compare to the lower bound.
265 2 : if span != nil && i.lower != nil && i.cmp(span.End, i.lower) <= 0 {
266 2 : span = nil
267 2 : }
268 2 : i.iterSpan = span
269 2 : if i.pos != posAtIterSpan {
270 2 : i.pos = posAtIterSpan
271 2 : }
272 2 : return span, err
273 : }
274 :
275 : // WrapChildren implements FragmentIterator.
276 0 : func (i *BoundedIter) WrapChildren(wrap WrapFn) {
277 0 : i.iter = wrap(i.iter)
278 0 : }
279 :
280 : // DebugTree is part of the FragmentIterator interface.
281 0 : func (i *BoundedIter) DebugTree(tp treeprinter.Node) {
282 0 : n := tp.Childf("%T(%p)", i, i)
283 0 : if i.iter != nil {
284 0 : i.iter.DebugTree(n)
285 0 : }
286 : }
|
