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 : KV *base.InternalKV
70 : Err error
71 : }
72 : }
73 :
74 : // Probe defines an interface for probes that may inspect or mutate internal
75 : // iterator behavior.
76 : type Probe interface {
77 : // Probe inspects, and possibly manipulates, iterator operations' results.
78 : Probe(*ProbeContext)
79 : }
80 :
81 : // ProbeContext provides the context within which a Probe is run. It includes
82 : // information about the iterator operation in progress.
83 : type ProbeContext struct {
84 : Op
85 : ProbeState
86 : }
87 :
88 : // ProbeState holds state additional to the context of the operation that's
89 : // accessible to probes.
90 : type ProbeState struct {
91 : *base.Comparer
92 : Log io.Writer
93 : }
94 :
95 : // Attach takes an iterator, an initial state and a probe, returning an iterator
96 : // that will invoke the provided Probe on all internal iterator operations.
97 : func Attach(
98 : iter base.InternalIterator, initialState ProbeState, probes ...Probe,
99 1 : ) base.InternalIterator {
100 1 : for i := range probes {
101 1 : iter = &probeIterator{
102 1 : iter: iter,
103 1 : probe: probes[i],
104 1 : probeCtx: ProbeContext{
105 1 : ProbeState: initialState,
106 1 : },
107 1 : }
108 1 : }
109 1 : return iter
110 : }
111 :
112 : // MustParseProbes parses each DSL string as a separate probe, returning a slice
113 : // of parsed probes. Panics if any of the probes fail to parse.
114 1 : func MustParseProbes(parser *dsl.Parser[Probe], probeDSLs ...string) []Probe {
115 1 : probes := make([]Probe, len(probeDSLs))
116 1 : var err error
117 1 : for i := range probeDSLs {
118 1 : probes[i], err = parser.Parse(probeDSLs[i])
119 1 : if err != nil {
120 0 : panic(err)
121 : }
122 : }
123 1 : return probes
124 : }
125 :
126 : type probeIterator struct {
127 : iter base.InternalIterator
128 : probe Probe
129 : probeCtx ProbeContext
130 : }
131 :
132 : // Assert that errIterator implements the internal iterator interface.
133 : var _ base.InternalIterator = (*probeIterator)(nil)
134 :
135 : // handleOp takes an Op representing the iterator operation performed, and the
136 : // underlying iterator's return value. It populates `Return.Err` and invokes the
137 : // probe.
138 1 : func (p *probeIterator) handleOp(preProbeOp Op) *base.InternalKV {
139 1 : p.probeCtx.Op = preProbeOp
140 1 : if preProbeOp.Return.KV == nil && p.iter != nil {
141 1 : p.probeCtx.Op.Return.Err = p.iter.Error()
142 1 : }
143 :
144 1 : p.probe.Probe(&p.probeCtx)
145 1 : return p.probeCtx.Op.Return.KV
146 : }
147 :
148 1 : func (p *probeIterator) SeekGE(key []byte, flags base.SeekGEFlags) *base.InternalKV {
149 1 : op := Op{
150 1 : Kind: OpSeekGE,
151 1 : SeekKey: key,
152 1 : }
153 1 : if p.iter != nil {
154 1 : op.Return.KV = p.iter.SeekGE(key, flags)
155 1 : }
156 1 : return p.handleOp(op)
157 : }
158 :
159 1 : func (p *probeIterator) SeekPrefixGE(prefix, key []byte, flags base.SeekGEFlags) *base.InternalKV {
160 1 : op := Op{
161 1 : Kind: OpSeekPrefixGE,
162 1 : SeekKey: key,
163 1 : }
164 1 : if p.iter != nil {
165 1 : op.Return.KV = p.iter.SeekPrefixGE(prefix, key, flags)
166 1 : }
167 1 : return p.handleOp(op)
168 : }
169 :
170 1 : func (p *probeIterator) SeekLT(key []byte, flags base.SeekLTFlags) *base.InternalKV {
171 1 : op := Op{
172 1 : Kind: OpSeekLT,
173 1 : SeekKey: key,
174 1 : }
175 1 : if p.iter != nil {
176 1 : op.Return.KV = p.iter.SeekLT(key, flags)
177 1 : }
178 1 : return p.handleOp(op)
179 : }
180 :
181 1 : func (p *probeIterator) First() *base.InternalKV {
182 1 : op := Op{Kind: OpFirst}
183 1 : if p.iter != nil {
184 1 : op.Return.KV = p.iter.First()
185 1 : }
186 1 : return p.handleOp(op)
187 : }
188 :
189 1 : func (p *probeIterator) Last() *base.InternalKV {
190 1 : op := Op{Kind: OpLast}
191 1 : if p.iter != nil {
192 1 : op.Return.KV = p.iter.Last()
193 1 : }
194 1 : return p.handleOp(op)
195 : }
196 :
197 1 : func (p *probeIterator) Next() *base.InternalKV {
198 1 : op := Op{Kind: OpNext}
199 1 : if p.iter != nil {
200 1 : op.Return.KV = p.iter.Next()
201 1 : }
202 1 : return p.handleOp(op)
203 : }
204 :
205 1 : func (p *probeIterator) NextPrefix(succKey []byte) *base.InternalKV {
206 1 : op := Op{Kind: OpNextPrefix, SeekKey: succKey}
207 1 : if p.iter != nil {
208 1 : op.Return.KV = p.iter.NextPrefix(succKey)
209 1 : }
210 1 : return p.handleOp(op)
211 : }
212 :
213 1 : func (p *probeIterator) Prev() *base.InternalKV {
214 1 : op := Op{Kind: OpPrev}
215 1 : if p.iter != nil {
216 1 : op.Return.KV = p.iter.Prev()
217 1 : }
218 1 : return p.handleOp(op)
219 : }
220 :
221 1 : func (p *probeIterator) Error() error {
222 1 : return p.probeCtx.Op.Return.Err
223 1 : }
224 :
225 1 : func (p *probeIterator) Close() error {
226 1 : op := Op{Kind: OpClose}
227 1 : if p.iter != nil {
228 1 : op.Return.Err = p.iter.Close()
229 1 : }
230 :
231 : // NB: Can't use handleOp because a close returns its error value directly
232 : // (and does not return a KV pair). We don't want to call iter.Error()
233 : // again, but rather use the error directly returned by iter.Close().
234 1 : p.probeCtx.Op = op
235 1 : p.probe.Probe(&p.probeCtx)
236 1 : return p.Error()
237 : }
238 :
239 0 : func (p *probeIterator) SetBounds(lower, upper []byte) {
240 0 : if p.iter != nil {
241 0 : p.iter.SetBounds(lower, upper)
242 0 : }
243 : }
244 :
245 0 : func (p *probeIterator) SetContext(ctx context.Context) {
246 0 : if p.iter != nil {
247 0 : p.iter.SetContext(ctx)
248 0 : }
249 : }
250 :
251 0 : func (p *probeIterator) String() string {
252 0 : if p.iter != nil {
253 0 : return fmt.Sprintf("probeIterator(%q)", p.iter.String())
254 0 : }
255 0 : return "probeIterator(nil)"
256 : }
|
