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 itertest
6 :
7 : import (
8 : "context"
9 : "fmt"
10 : "io"
11 :
12 : "github.com/cockroachdb/pebble/internal/base"
13 : "github.com/cockroachdb/pebble/internal/dsl"
14 : )
15 :
16 : // OpKind indicates the type of iterator operation being performed.
17 : type OpKind int8
18 :
19 : const (
20 : // OpSeekGE indicates a SeekGE internal iterator operation.
21 : OpSeekGE OpKind = iota
22 : // OpSeekPrefixGE indicates a SeekPrefixGE internal iterator operation.
23 : OpSeekPrefixGE
24 : // OpSeekLT indicates a SeekLT internal iterator operation.
25 : OpSeekLT
26 : // OpFirst indicates a First internal iterator operation.
27 : OpFirst
28 : // OpLast indicates a Last internal iterator operation.
29 : OpLast
30 : // OpNext indicates a Next internal iterator operation.
31 : OpNext
32 : // OpNextPrefix indicates a NextPrefix internal iterator operation.
33 : OpNextPrefix
34 : // OpPrev indicates a Prev internal iterator operation.
35 : OpPrev
36 : // OpClose indicates a Close internal iterator operation.
37 : OpClose
38 : numOpKinds
39 : )
40 :
41 : var opNames = [numOpKinds]string{
42 : OpSeekGE: "OpSeekGE",
43 : OpSeekPrefixGE: "OpSeekPrefixGE",
44 : OpSeekLT: "OpSeekLT",
45 : OpFirst: "OpFirst",
46 : OpLast: "OpLast",
47 : OpNext: "OpNext",
48 : OpNextPrefix: "OpNextPrefix",
49 : OpPrev: "OpPrev",
50 : OpClose: "OpClose",
51 : }
52 :
53 : // OpKind implements Predicate.
54 : var _ Predicate = OpKind(0)
55 :
56 : // String imlements fmt.Stringer.
57 1 : func (o OpKind) String() string { return opNames[o] }
58 :
59 : // Evaluate implements Predicate.
60 1 : func (o OpKind) Evaluate(pctx *ProbeContext) bool { return pctx.Op.Kind == o }
61 :
62 : // Op describes an individual iterator operation being performed.
63 : type Op struct {
64 : Kind OpKind
65 : SeekKey []byte
66 : // Return is initialized with the return result of the underlying iterator.
67 : // Probes may mutate them.
68 : Return struct {
69 : Key *base.InternalKey
70 : Value base.LazyValue
71 : Err error
72 : }
73 : }
74 :
75 : // Probe defines an interface for probes that may inspect or mutate internal
76 : // iterator behavior.
77 : type Probe interface {
78 : // Probe inspects, and possibly manipulates, iterator operations' results.
79 : Probe(*ProbeContext)
80 : }
81 :
82 : // ProbeContext provides the context within which a Probe is run. It includes
83 : // information about the iterator operation in progress.
84 : type ProbeContext struct {
85 : Op
86 : ProbeState
87 : }
88 :
89 : // ProbeState holds state additional to the context of the operation that's
90 : // accessible to probes.
91 : type ProbeState struct {
92 : *base.Comparer
93 : Log io.Writer
94 : }
95 :
96 : // Attach takes an iterator, an initial state and a probe, returning an iterator
97 : // that will invoke the provided Probe on all internal iterator operations.
98 : func Attach(
99 : iter base.InternalIterator, initialState ProbeState, probes ...Probe,
100 1 : ) base.InternalIterator {
101 1 : for i := range probes {
102 1 : iter = &probeIterator{
103 1 : iter: iter,
104 1 : probe: probes[i],
105 1 : probeCtx: ProbeContext{
106 1 : ProbeState: initialState,
107 1 : },
108 1 : }
109 1 : }
110 1 : return iter
111 : }
112 :
113 : // MustParseProbes parses each DSL string as a separate probe, returning a slice
114 : // of parsed probes. Panics if any of the probes fail to parse.
115 1 : func MustParseProbes(parser *dsl.Parser[Probe], probeDSLs ...string) []Probe {
116 1 : probes := make([]Probe, len(probeDSLs))
117 1 : var err error
118 1 : for i := range probeDSLs {
119 1 : probes[i], err = parser.Parse(probeDSLs[i])
120 1 : if err != nil {
121 0 : panic(err)
122 : }
123 : }
124 1 : return probes
125 : }
126 :
127 : type probeIterator struct {
128 : iter base.InternalIterator
129 : probe Probe
130 : probeCtx ProbeContext
131 : }
132 :
133 : // Assert that errIterator implements the internal iterator interface.
134 : var _ base.InternalIterator = (*probeIterator)(nil)
135 :
136 : // handleOp takes an Op representing the iterator operation performed, and the
137 : // underlying iterator's return value. It populates `Return.Err` and invokes the
138 : // probe.
139 1 : func (p *probeIterator) handleOp(preProbeOp Op) (*base.InternalKey, base.LazyValue) {
140 1 : p.probeCtx.Op = preProbeOp
141 1 : if preProbeOp.Return.Key == nil && p.iter != nil {
142 1 : p.probeCtx.Op.Return.Err = p.iter.Error()
143 1 : }
144 :
145 1 : p.probe.Probe(&p.probeCtx)
146 1 : return p.probeCtx.Op.Return.Key, p.probeCtx.Op.Return.Value
147 : }
148 :
149 : func (p *probeIterator) SeekGE(
150 : key []byte, flags base.SeekGEFlags,
151 1 : ) (*base.InternalKey, base.LazyValue) {
152 1 : op := Op{
153 1 : Kind: OpSeekGE,
154 1 : SeekKey: key,
155 1 : }
156 1 : if p.iter != nil {
157 1 : op.Return.Key, op.Return.Value = p.iter.SeekGE(key, flags)
158 1 : }
159 1 : return p.handleOp(op)
160 : }
161 :
162 : func (p *probeIterator) SeekPrefixGE(
163 : prefix, key []byte, flags base.SeekGEFlags,
164 1 : ) (*base.InternalKey, base.LazyValue) {
165 1 : op := Op{
166 1 : Kind: OpSeekPrefixGE,
167 1 : SeekKey: key,
168 1 : }
169 1 : if p.iter != nil {
170 1 : op.Return.Key, op.Return.Value = p.iter.SeekPrefixGE(prefix, key, flags)
171 1 : }
172 1 : return p.handleOp(op)
173 : }
174 :
175 : func (p *probeIterator) SeekLT(
176 : key []byte, flags base.SeekLTFlags,
177 1 : ) (*base.InternalKey, base.LazyValue) {
178 1 : op := Op{
179 1 : Kind: OpSeekLT,
180 1 : SeekKey: key,
181 1 : }
182 1 : if p.iter != nil {
183 1 : op.Return.Key, op.Return.Value = p.iter.SeekLT(key, flags)
184 1 : }
185 1 : return p.handleOp(op)
186 : }
187 :
188 1 : func (p *probeIterator) First() (*base.InternalKey, base.LazyValue) {
189 1 : op := Op{Kind: OpFirst}
190 1 : if p.iter != nil {
191 1 : op.Return.Key, op.Return.Value = p.iter.First()
192 1 : }
193 1 : return p.handleOp(op)
194 : }
195 :
196 1 : func (p *probeIterator) Last() (*base.InternalKey, base.LazyValue) {
197 1 : op := Op{Kind: OpLast}
198 1 : if p.iter != nil {
199 1 : op.Return.Key, op.Return.Value = p.iter.Last()
200 1 : }
201 1 : return p.handleOp(op)
202 : }
203 :
204 1 : func (p *probeIterator) Next() (*base.InternalKey, base.LazyValue) {
205 1 : op := Op{Kind: OpNext}
206 1 : if p.iter != nil {
207 1 : op.Return.Key, op.Return.Value = p.iter.Next()
208 1 : }
209 1 : return p.handleOp(op)
210 : }
211 :
212 1 : func (p *probeIterator) NextPrefix(succKey []byte) (*base.InternalKey, base.LazyValue) {
213 1 : op := Op{Kind: OpNextPrefix, SeekKey: succKey}
214 1 : if p.iter != nil {
215 1 : op.Return.Key, op.Return.Value = p.iter.NextPrefix(succKey)
216 1 : }
217 1 : return p.handleOp(op)
218 : }
219 :
220 1 : func (p *probeIterator) Prev() (*base.InternalKey, base.LazyValue) {
221 1 : op := Op{Kind: OpPrev}
222 1 : if p.iter != nil {
223 1 : op.Return.Key, op.Return.Value = p.iter.Prev()
224 1 : }
225 1 : return p.handleOp(op)
226 : }
227 :
228 1 : func (p *probeIterator) Error() error {
229 1 : return p.probeCtx.Op.Return.Err
230 1 : }
231 :
232 1 : func (p *probeIterator) Close() error {
233 1 : op := Op{Kind: OpClose}
234 1 : if p.iter != nil {
235 1 : op.Return.Err = p.iter.Close()
236 1 : }
237 :
238 : // NB: Can't use handleOp because a close returns its error value directly
239 : // (and does not return a KV pair). We don't want to call iter.Error()
240 : // again, but rather use the error directly returned by iter.Close().
241 1 : p.probeCtx.Op = op
242 1 : p.probe.Probe(&p.probeCtx)
243 1 : return p.Error()
244 : }
245 :
246 0 : func (p *probeIterator) SetBounds(lower, upper []byte) {
247 0 : if p.iter != nil {
248 0 : p.iter.SetBounds(lower, upper)
249 0 : }
250 : }
251 :
252 0 : func (p *probeIterator) SetContext(ctx context.Context) {
253 0 : if p.iter != nil {
254 0 : p.iter.SetContext(ctx)
255 0 : }
256 : }
257 :
258 0 : func (p *probeIterator) String() string {
259 0 : if p.iter != nil {
260 0 : return fmt.Sprintf("probeIterator(%q)", p.iter.String())
261 0 : }
262 0 : return "probeIterator(nil)"
263 : }
|
