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 : "github.com/cockroachdb/pebble/internal/treeprinter"
15 : )
16 :
17 : // OpKind indicates the type of iterator operation being performed.
18 : type OpKind int8
19 :
20 : const (
21 : // OpSeekGE indicates a SeekGE internal iterator operation.
22 : OpSeekGE OpKind = iota
23 : // OpSeekPrefixGE indicates a SeekPrefixGE internal iterator operation.
24 : OpSeekPrefixGE
25 : // OpSeekLT indicates a SeekLT internal iterator operation.
26 : OpSeekLT
27 : // OpFirst indicates a First internal iterator operation.
28 : OpFirst
29 : // OpLast indicates a Last internal iterator operation.
30 : OpLast
31 : // OpNext indicates a Next internal iterator operation.
32 : OpNext
33 : // OpNextPrefix indicates a NextPrefix internal iterator operation.
34 : OpNextPrefix
35 : // OpPrev indicates a Prev internal iterator operation.
36 : OpPrev
37 : // OpClose indicates a Close internal iterator operation.
38 : OpClose
39 : numOpKinds
40 : )
41 :
42 : var opNames = [numOpKinds]string{
43 : OpSeekGE: "OpSeekGE",
44 : OpSeekPrefixGE: "OpSeekPrefixGE",
45 : OpSeekLT: "OpSeekLT",
46 : OpFirst: "OpFirst",
47 : OpLast: "OpLast",
48 : OpNext: "OpNext",
49 : OpNextPrefix: "OpNextPrefix",
50 : OpPrev: "OpPrev",
51 : OpClose: "OpClose",
52 : }
53 :
54 : // OpKind implements Predicate.
55 : var _ Predicate = OpKind(0)
56 :
57 : // String imlements fmt.Stringer.
58 1 : func (o OpKind) String() string { return opNames[o] }
59 :
60 : // Evaluate implements Predicate.
61 1 : func (o OpKind) Evaluate(pctx *ProbeContext) bool { return pctx.Op.Kind == o }
62 :
63 : // Op describes an individual iterator operation being performed.
64 : type Op struct {
65 : Kind OpKind
66 : SeekKey []byte
67 : // Return is initialized with the return result of the underlying iterator.
68 : // Probes may mutate them.
69 : Return struct {
70 : KV *base.InternalKV
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.InternalKV {
140 1 : p.probeCtx.Op = preProbeOp
141 1 : if preProbeOp.Return.KV == 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.KV
147 : }
148 :
149 1 : func (p *probeIterator) SeekGE(key []byte, flags base.SeekGEFlags) *base.InternalKV {
150 1 : op := Op{
151 1 : Kind: OpSeekGE,
152 1 : SeekKey: key,
153 1 : }
154 1 : if p.iter != nil {
155 1 : op.Return.KV = p.iter.SeekGE(key, flags)
156 1 : }
157 1 : return p.handleOp(op)
158 : }
159 :
160 1 : func (p *probeIterator) SeekPrefixGE(prefix, key []byte, flags base.SeekGEFlags) *base.InternalKV {
161 1 : op := Op{
162 1 : Kind: OpSeekPrefixGE,
163 1 : SeekKey: key,
164 1 : }
165 1 : if p.iter != nil {
166 1 : op.Return.KV = p.iter.SeekPrefixGE(prefix, key, flags)
167 1 : }
168 1 : return p.handleOp(op)
169 : }
170 :
171 1 : func (p *probeIterator) SeekLT(key []byte, flags base.SeekLTFlags) *base.InternalKV {
172 1 : op := Op{
173 1 : Kind: OpSeekLT,
174 1 : SeekKey: key,
175 1 : }
176 1 : if p.iter != nil {
177 1 : op.Return.KV = p.iter.SeekLT(key, flags)
178 1 : }
179 1 : return p.handleOp(op)
180 : }
181 :
182 1 : func (p *probeIterator) First() *base.InternalKV {
183 1 : op := Op{Kind: OpFirst}
184 1 : if p.iter != nil {
185 1 : op.Return.KV = p.iter.First()
186 1 : }
187 1 : return p.handleOp(op)
188 : }
189 :
190 1 : func (p *probeIterator) Last() *base.InternalKV {
191 1 : op := Op{Kind: OpLast}
192 1 : if p.iter != nil {
193 1 : op.Return.KV = p.iter.Last()
194 1 : }
195 1 : return p.handleOp(op)
196 : }
197 :
198 1 : func (p *probeIterator) Next() *base.InternalKV {
199 1 : op := Op{Kind: OpNext}
200 1 : if p.iter != nil {
201 1 : op.Return.KV = p.iter.Next()
202 1 : }
203 1 : return p.handleOp(op)
204 : }
205 :
206 1 : func (p *probeIterator) NextPrefix(succKey []byte) *base.InternalKV {
207 1 : op := Op{Kind: OpNextPrefix, SeekKey: succKey}
208 1 : if p.iter != nil {
209 1 : op.Return.KV = p.iter.NextPrefix(succKey)
210 1 : }
211 1 : return p.handleOp(op)
212 : }
213 :
214 1 : func (p *probeIterator) Prev() *base.InternalKV {
215 1 : op := Op{Kind: OpPrev}
216 1 : if p.iter != nil {
217 1 : op.Return.KV = p.iter.Prev()
218 1 : }
219 1 : return p.handleOp(op)
220 : }
221 :
222 1 : func (p *probeIterator) Error() error {
223 1 : return p.probeCtx.Op.Return.Err
224 1 : }
225 :
226 1 : func (p *probeIterator) Close() error {
227 1 : op := Op{Kind: OpClose}
228 1 : if p.iter != nil {
229 1 : op.Return.Err = p.iter.Close()
230 1 : }
231 :
232 : // NB: Can't use handleOp because a close returns its error value directly
233 : // (and does not return a KV pair). We don't want to call iter.Error()
234 : // again, but rather use the error directly returned by iter.Close().
235 1 : p.probeCtx.Op = op
236 1 : p.probe.Probe(&p.probeCtx)
237 1 : return p.Error()
238 : }
239 :
240 0 : func (p *probeIterator) SetBounds(lower, upper []byte) {
241 0 : if p.iter != nil {
242 0 : p.iter.SetBounds(lower, upper)
243 0 : }
244 : }
245 :
246 0 : func (p *probeIterator) SetContext(ctx context.Context) {
247 0 : if p.iter != nil {
248 0 : p.iter.SetContext(ctx)
249 0 : }
250 : }
251 :
252 : // DebugTree is part of the InternalIterator interface.
253 0 : func (p *probeIterator) DebugTree(tp treeprinter.Node) {
254 0 : n := tp.Childf("%T(%p)", p, p)
255 0 : if p.iter != nil {
256 0 : p.iter.DebugTree(n)
257 0 : }
258 : }
259 :
260 0 : func (p *probeIterator) String() string {
261 0 : if p.iter != nil {
262 0 : return fmt.Sprintf("probeIterator(%q)", p.iter.String())
263 0 : }
264 0 : return "probeIterator(nil)"
265 : }
|
