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 : "bytes"
9 : "fmt"
10 : "go/token"
11 : "io"
12 : "reflect"
13 : "strconv"
14 : "strings"
15 :
16 : "github.com/cockroachdb/errors"
17 : "github.com/cockroachdb/pebble/internal/dsl"
18 : "github.com/cockroachdb/pebble/internal/treeprinter"
19 : )
20 :
21 : // This file contains testing facilities for Spans and FragmentIterators. It's
22 : // defined here so that it may be used by the keyspan package to test its
23 : // various FragmentIterator implementations.
24 : //
25 : // TODO(jackson): Move keyspan.{Span,Key,FragmentIterator} into internal/base,
26 : // and then move the testing facilities to an independent package, eg
27 : // internal/itertest. Alternatively, make all tests that use it use the
28 : // keyspan_test package, which can then import a separate itertest package.
29 :
30 : // probe defines an interface for probes that may inspect or mutate internal
31 : // span iterator behavior.
32 : type probe interface {
33 : // probe inspects, and possibly manipulates, iterator operations' results.
34 : probe(*probeContext)
35 : }
36 :
37 0 : func parseProbes(probeDSLs ...string) []probe {
38 0 : probes := make([]probe, len(probeDSLs))
39 0 : var err error
40 0 : for i := range probeDSLs {
41 0 : probes[i], err = probeParser.Parse(probeDSLs[i])
42 0 : if err != nil {
43 0 : panic(err)
44 : }
45 : }
46 0 : return probes
47 : }
48 :
49 0 : func attachProbes(iter FragmentIterator, pctx probeContext, probes ...probe) FragmentIterator {
50 0 : if pctx.log == nil {
51 0 : pctx.log = io.Discard
52 0 : }
53 0 : for i := range probes {
54 0 : iter = &probeIterator{
55 0 : iter: iter,
56 0 : probe: probes[i],
57 0 : probeCtx: pctx,
58 0 : }
59 0 : }
60 0 : return iter
61 : }
62 :
63 : // ParseAndAttachProbes parses DSL probes and attaches them to an iterator.
64 : func ParseAndAttachProbes(
65 : iter FragmentIterator, log io.Writer, probeDSLs ...string,
66 0 : ) FragmentIterator {
67 0 : pctx := probeContext{log: log}
68 0 : return attachProbes(iter, pctx, parseProbes(probeDSLs...)...)
69 0 : }
70 :
71 : // probeContext provides the context within which a probe is run. It includes
72 : // information about the iterator operation in progress.
73 : type probeContext struct {
74 : op
75 : log io.Writer
76 : }
77 :
78 : type op struct {
79 : Kind OpKind
80 : SeekKey []byte
81 : Span *Span
82 : Err error
83 : }
84 :
85 : // ErrInjected is an error artificially injected for testing.
86 : var ErrInjected = &errorProbe{name: "ErrInjected", err: errors.New("injected error")}
87 :
88 1 : var probeParser = func() *dsl.Parser[probe] {
89 1 : valuerParser := dsl.NewParser[valuer]()
90 1 : valuerParser.DefineConstant("StartKey", func() valuer { return startKey{} })
91 1 : valuerParser.DefineFunc("Bytes",
92 1 : func(p *dsl.Parser[valuer], s *dsl.Scanner) valuer {
93 0 : v := bytesConstant{bytes: []byte(s.ConsumeString())}
94 0 : s.Consume(token.RPAREN)
95 0 : return v
96 0 : })
97 :
98 1 : predicateParser := dsl.NewPredicateParser[*probeContext]()
99 1 : predicateParser.DefineFunc("Equal",
100 1 : func(p *dsl.Parser[dsl.Predicate[*probeContext]], s *dsl.Scanner) dsl.Predicate[*probeContext] {
101 0 : eq := equal{
102 0 : valuerParser.ParseFromPos(s, s.Scan()),
103 0 : valuerParser.ParseFromPos(s, s.Scan()),
104 0 : }
105 0 : s.Consume(token.RPAREN)
106 0 : return eq
107 0 : })
108 1 : for i, name := range opNames {
109 1 : opKind := OpKind(i)
110 1 : predicateParser.DefineConstant(name, func() dsl.Predicate[*probeContext] {
111 0 : // An OpKind implements dsl.Predicate[*probeContext].
112 0 : return opKind
113 0 : })
114 : }
115 1 : probeParser := dsl.NewParser[probe]()
116 1 : probeParser.DefineConstant("ErrInjected", func() probe { return ErrInjected })
117 1 : probeParser.DefineConstant("noop", func() probe { return noop{} })
118 1 : probeParser.DefineFunc("If",
119 1 : func(p *dsl.Parser[probe], s *dsl.Scanner) probe {
120 0 : probe := ifProbe{
121 0 : predicateParser.ParseFromPos(s, s.Scan()),
122 0 : probeParser.ParseFromPos(s, s.Scan()),
123 0 : probeParser.ParseFromPos(s, s.Scan()),
124 0 : }
125 0 : s.Consume(token.RPAREN)
126 0 : return probe
127 0 : })
128 1 : probeParser.DefineFunc("Return",
129 1 : func(p *dsl.Parser[probe], s *dsl.Scanner) (ret probe) {
130 0 : switch tok := s.Scan(); tok.Kind {
131 0 : case token.STRING:
132 0 : str, err := strconv.Unquote(tok.Lit)
133 0 : if err != nil {
134 0 : panic(err)
135 : }
136 0 : span := ParseSpan(str)
137 0 : ret = returnSpan{s: &span}
138 0 : case token.IDENT:
139 0 : switch tok.Lit {
140 0 : case "nil":
141 0 : ret = returnSpan{s: nil}
142 0 : default:
143 0 : panic(errors.Newf("unrecognized return value %q", tok.Lit))
144 : }
145 : }
146 0 : s.Consume(token.RPAREN)
147 0 : return ret
148 : })
149 1 : probeParser.DefineFunc("Log",
150 1 : func(p *dsl.Parser[probe], s *dsl.Scanner) (ret probe) {
151 0 : ret = loggingProbe{prefix: s.ConsumeString()}
152 0 : s.Consume(token.RPAREN)
153 0 : return ret
154 0 : })
155 1 : return probeParser
156 : }()
157 :
158 : // probe implementations
159 :
160 : type errorProbe struct {
161 : name string
162 : err error
163 : }
164 :
165 0 : func (p *errorProbe) String() string { return p.name }
166 0 : func (p *errorProbe) Error() error { return p.err }
167 0 : func (p *errorProbe) probe(pctx *probeContext) {
168 0 : pctx.op.Err = p.err
169 0 : pctx.op.Span = nil
170 0 : }
171 :
172 : // ifProbe is a conditional probe. If its predicate evaluates to true, it probes
173 : // using its Then probe. If its predicate evalutes to false, it probes using its
174 : // Else probe.
175 : type ifProbe struct {
176 : Predicate dsl.Predicate[*probeContext]
177 : Then probe
178 : Else probe
179 : }
180 :
181 0 : func (p ifProbe) String() string { return fmt.Sprintf("(If %s %s %s)", p.Predicate, p.Then, p.Else) }
182 0 : func (p ifProbe) probe(pctx *probeContext) {
183 0 : if p.Predicate.Evaluate(pctx) {
184 0 : p.Then.probe(pctx)
185 0 : } else {
186 0 : p.Else.probe(pctx)
187 0 : }
188 : }
189 :
190 : type returnSpan struct {
191 : s *Span
192 : }
193 :
194 0 : func (p returnSpan) String() string {
195 0 : if p.s == nil {
196 0 : return "(Return nil)"
197 0 : }
198 0 : return fmt.Sprintf("(Return %q)", p.s.String())
199 : }
200 :
201 0 : func (p returnSpan) probe(pctx *probeContext) {
202 0 : pctx.op.Span = p.s
203 0 : pctx.op.Err = nil
204 0 : }
205 :
206 : type noop struct{}
207 :
208 0 : func (noop) String() string { return "Noop" }
209 0 : func (noop) probe(pctx *probeContext) {}
210 :
211 : type loggingProbe struct {
212 : prefix string
213 : }
214 :
215 0 : func (lp loggingProbe) String() string { return fmt.Sprintf("(Log %q)", lp.prefix) }
216 0 : func (lp loggingProbe) probe(pctx *probeContext) {
217 0 : opStr := strings.TrimPrefix(pctx.op.Kind.String(), "Op")
218 0 : fmt.Fprintf(pctx.log, "%s%s(", lp.prefix, opStr)
219 0 : if pctx.op.SeekKey != nil {
220 0 : fmt.Fprintf(pctx.log, "%q", pctx.op.SeekKey)
221 0 : }
222 0 : fmt.Fprint(pctx.log, ") = ")
223 0 : if pctx.op.Span == nil {
224 0 : fmt.Fprint(pctx.log, "nil")
225 0 : if pctx.op.Err != nil {
226 0 : fmt.Fprintf(pctx.log, " <err=%q>", pctx.op.Err)
227 0 : }
228 0 : } else {
229 0 : fmt.Fprint(pctx.log, pctx.op.Span.String())
230 0 : }
231 0 : fmt.Fprintln(pctx.log)
232 : }
233 :
234 : // dsl.Predicate[*probeContext] implementations.
235 :
236 : type equal struct {
237 : a, b valuer
238 : }
239 :
240 0 : func (e equal) String() string { return fmt.Sprintf("(Equal %s %s)", e.a, e.b) }
241 0 : func (e equal) Evaluate(pctx *probeContext) bool {
242 0 : return reflect.DeepEqual(e.a.value(pctx), e.b.value(pctx))
243 0 : }
244 :
245 : // OpKind indicates the type of iterator operation being performed.
246 : type OpKind int8
247 :
248 : // OpKind values.
249 : const (
250 : OpSeekGE OpKind = iota
251 : OpSeekLT
252 : OpFirst
253 : OpLast
254 : OpNext
255 : OpPrev
256 : OpClose
257 : numOpKinds
258 : )
259 :
260 0 : func (o OpKind) String() string { return opNames[o] }
261 :
262 : // Evaluate implements dsl.Predicate.
263 0 : func (o OpKind) Evaluate(pctx *probeContext) bool { return pctx.op.Kind == o }
264 :
265 : var opNames = [numOpKinds]string{
266 : OpSeekGE: "OpSeekGE",
267 : OpSeekLT: "OpSeekLT",
268 : OpFirst: "OpFirst",
269 : OpLast: "OpLast",
270 : OpNext: "OpNext",
271 : OpPrev: "OpPrev",
272 : OpClose: "OpClose",
273 : }
274 :
275 : // valuer implementations
276 :
277 : type valuer interface {
278 : fmt.Stringer
279 : value(pctx *probeContext) any
280 : }
281 :
282 : type bytesConstant struct {
283 : bytes []byte
284 : }
285 :
286 0 : func (b bytesConstant) String() string { return fmt.Sprintf("%q", string(b.bytes)) }
287 0 : func (b bytesConstant) value(pctx *probeContext) any { return b.bytes }
288 :
289 : type startKey struct{}
290 :
291 0 : func (s startKey) String() string { return "StartKey" }
292 0 : func (s startKey) value(pctx *probeContext) any {
293 0 : if pctx.op.Span == nil {
294 0 : return nil
295 0 : }
296 0 : return pctx.op.Span.Start
297 : }
298 :
299 : type probeIterator struct {
300 : iter FragmentIterator
301 : probe probe
302 : probeCtx probeContext
303 : }
304 :
305 : // Assert that probeIterator implements the fragment iterator interface.
306 : var _ FragmentIterator = (*probeIterator)(nil)
307 :
308 0 : func (p *probeIterator) handleOp(preProbeOp op) (*Span, error) {
309 0 : p.probeCtx.op = preProbeOp
310 0 : p.probe.probe(&p.probeCtx)
311 0 : return p.probeCtx.op.Span, p.probeCtx.op.Err
312 0 : }
313 :
314 0 : func (p *probeIterator) SeekGE(key []byte) (*Span, error) {
315 0 : op := op{
316 0 : Kind: OpSeekGE,
317 0 : SeekKey: key,
318 0 : }
319 0 : if p.iter != nil {
320 0 : op.Span, op.Err = p.iter.SeekGE(key)
321 0 : }
322 0 : return p.handleOp(op)
323 : }
324 :
325 0 : func (p *probeIterator) SeekLT(key []byte) (*Span, error) {
326 0 : op := op{
327 0 : Kind: OpSeekLT,
328 0 : SeekKey: key,
329 0 : }
330 0 : if p.iter != nil {
331 0 : op.Span, op.Err = p.iter.SeekLT(key)
332 0 : }
333 0 : return p.handleOp(op)
334 : }
335 :
336 0 : func (p *probeIterator) First() (*Span, error) {
337 0 : op := op{Kind: OpFirst}
338 0 : if p.iter != nil {
339 0 : op.Span, op.Err = p.iter.First()
340 0 : }
341 0 : return p.handleOp(op)
342 : }
343 :
344 0 : func (p *probeIterator) Last() (*Span, error) {
345 0 : op := op{Kind: OpLast}
346 0 : if p.iter != nil {
347 0 : op.Span, op.Err = p.iter.Last()
348 0 : }
349 0 : return p.handleOp(op)
350 : }
351 :
352 0 : func (p *probeIterator) Next() (*Span, error) {
353 0 : op := op{Kind: OpNext}
354 0 : if p.iter != nil {
355 0 : op.Span, op.Err = p.iter.Next()
356 0 : }
357 0 : return p.handleOp(op)
358 : }
359 :
360 0 : func (p *probeIterator) Prev() (*Span, error) {
361 0 : op := op{Kind: OpPrev}
362 0 : if p.iter != nil {
363 0 : op.Span, op.Err = p.iter.Prev()
364 0 : }
365 0 : return p.handleOp(op)
366 : }
367 :
368 0 : func (p *probeIterator) Close() {
369 0 : op := op{Kind: OpClose}
370 0 : if p.iter != nil {
371 0 : p.iter.Close()
372 0 : }
373 0 : _, _ = p.handleOp(op)
374 : }
375 :
376 0 : func (p *probeIterator) WrapChildren(wrap WrapFn) {
377 0 : p.iter = wrap(p.iter)
378 0 : }
379 :
380 : // DebugTree is part of the FragmentIterator interface.
381 0 : func (p *probeIterator) DebugTree(tp treeprinter.Node) {
382 0 : n := tp.Childf("%T(%p)", p, p)
383 0 : if p.iter != nil {
384 0 : p.iter.DebugTree(n)
385 0 : }
386 : }
387 :
388 : // RunIterCmd evaluates a datadriven command controlling an internal
389 : // keyspan.FragmentIterator, writing the results of the iterator operations to
390 : // the provided writer.
391 0 : func RunIterCmd(tdInput string, iter FragmentIterator, w io.Writer) {
392 0 : lines := strings.Split(strings.TrimSpace(tdInput), "\n")
393 0 : for i, line := range lines {
394 0 : if i > 0 {
395 0 : fmt.Fprintln(w)
396 0 : }
397 0 : line = strings.TrimSpace(line)
398 0 : i := strings.IndexByte(line, '#')
399 0 : iterCmd := line
400 0 : if i > 0 {
401 0 : iterCmd = string(line[:i])
402 0 : }
403 0 : runIterOp(w, iter, iterCmd)
404 : }
405 : }
406 :
407 : var iterDelim = map[rune]bool{',': true, ' ': true, '(': true, ')': true, '"': true}
408 :
409 0 : func runIterOp(w io.Writer, it FragmentIterator, op string) {
410 0 : fields := strings.FieldsFunc(op, func(r rune) bool { return iterDelim[r] })
411 0 : var s *Span
412 0 : var err error
413 0 : switch strings.ToLower(fields[0]) {
414 0 : case "first":
415 0 : s, err = it.First()
416 0 : case "last":
417 0 : s, err = it.Last()
418 0 : case "seekge", "seek-ge":
419 0 : if len(fields) == 1 {
420 0 : panic(fmt.Sprintf("unable to parse iter op %q", op))
421 : }
422 0 : s, err = it.SeekGE([]byte(fields[1]))
423 0 : case "seeklt", "seek-lt":
424 0 : if len(fields) == 1 {
425 0 : panic(fmt.Sprintf("unable to parse iter op %q", op))
426 : }
427 0 : s, err = it.SeekLT([]byte(fields[1]))
428 0 : case "next":
429 0 : s, err = it.Next()
430 0 : case "prev":
431 0 : s, err = it.Prev()
432 0 : default:
433 0 : panic(fmt.Sprintf("unrecognized iter op %q", fields[0]))
434 : }
435 0 : switch {
436 0 : case err != nil:
437 0 : fmt.Fprintf(w, "<nil> err=<%s>", err)
438 0 : case s == nil:
439 0 : fmt.Fprint(w, "<nil>")
440 0 : default:
441 0 : fmt.Fprint(w, s)
442 : }
443 : }
444 :
445 : // RunFragmentIteratorCmd runs a command on an iterator; intended for testing.
446 0 : func RunFragmentIteratorCmd(iter FragmentIterator, input string, extraInfo func() string) string {
447 0 : var b bytes.Buffer
448 0 : for _, line := range strings.Split(input, "\n") {
449 0 : parts := strings.Fields(line)
450 0 : if len(parts) == 0 {
451 0 : continue
452 : }
453 0 : var span *Span
454 0 : var err error
455 0 : switch parts[0] {
456 0 : case "seek-ge":
457 0 : if len(parts) != 2 {
458 0 : return "seek-ge <key>\n"
459 0 : }
460 0 : span, err = iter.SeekGE([]byte(strings.TrimSpace(parts[1])))
461 0 : case "seek-lt":
462 0 : if len(parts) != 2 {
463 0 : return "seek-lt <key>\n"
464 0 : }
465 0 : span, err = iter.SeekLT([]byte(strings.TrimSpace(parts[1])))
466 0 : case "first":
467 0 : span, err = iter.First()
468 0 : case "last":
469 0 : span, err = iter.Last()
470 0 : case "next":
471 0 : span, err = iter.Next()
472 0 : case "prev":
473 0 : span, err = iter.Prev()
474 0 : default:
475 0 : return fmt.Sprintf("unknown op: %s", parts[0])
476 : }
477 0 : switch {
478 0 : case err != nil:
479 0 : fmt.Fprintf(&b, "err=%v\n", err)
480 0 : case span == nil:
481 0 : fmt.Fprintf(&b, ".\n")
482 0 : default:
483 0 : fmt.Fprintf(&b, "%s", span)
484 0 : if extraInfo != nil {
485 0 : fmt.Fprintf(&b, " (%s)", extraInfo())
486 0 : }
487 0 : b.WriteByte('\n')
488 : }
489 : }
490 0 : return b.String()
491 : }
492 :
493 : // NewInvalidatingIter wraps a FragmentIterator; spans surfaced by the inner
494 : // iterator are copied to buffers that are zeroed by subsequent iterator
495 : // positioning calls. This is intended to help surface bugs in improper lifetime
496 : // expectations of Spans.
497 1 : func NewInvalidatingIter(iter FragmentIterator) FragmentIterator {
498 1 : return &invalidatingIter{
499 1 : iter: iter,
500 1 : }
501 1 : }
502 :
503 : type invalidatingIter struct {
504 : iter FragmentIterator
505 : bufs [][]byte
506 : keys []Key
507 : span Span
508 : }
509 :
510 : // invalidatingIter implements FragmentIterator.
511 : var _ FragmentIterator = (*invalidatingIter)(nil)
512 :
513 1 : func (i *invalidatingIter) invalidate(s *Span, err error) (*Span, error) {
514 1 : // Mangle the entirety of the byte bufs and the keys slice.
515 1 : for j := range i.bufs {
516 1 : for k := range i.bufs[j] {
517 1 : i.bufs[j][k] = 0xff
518 1 : }
519 1 : i.bufs[j] = nil
520 : }
521 1 : for j := range i.keys {
522 1 : i.keys[j] = Key{}
523 1 : }
524 1 : if s == nil {
525 1 : return nil, err
526 1 : }
527 :
528 : // Copy all of the span's slices into slices owned by the invalidating iter
529 : // that we can invalidate on a subsequent positioning method.
530 1 : i.bufs = i.bufs[:0]
531 1 : i.keys = i.keys[:0]
532 1 : i.span = Span{
533 1 : Start: i.saveBytes(s.Start),
534 1 : End: i.saveBytes(s.End),
535 1 : KeysOrder: s.KeysOrder,
536 1 : }
537 1 : for j := range s.Keys {
538 1 : i.keys = append(i.keys, Key{
539 1 : Trailer: s.Keys[j].Trailer,
540 1 : Suffix: i.saveBytes(s.Keys[j].Suffix),
541 1 : Value: i.saveBytes(s.Keys[j].Value),
542 1 : })
543 1 : }
544 1 : i.span.Keys = i.keys
545 1 : return &i.span, err
546 : }
547 :
548 1 : func (i *invalidatingIter) saveBytes(b []byte) []byte {
549 1 : if b == nil {
550 1 : return nil
551 1 : }
552 1 : saved := append([]byte(nil), b...)
553 1 : i.bufs = append(i.bufs, saved)
554 1 : return saved
555 : }
556 :
557 1 : func (i *invalidatingIter) SeekGE(key []byte) (*Span, error) { return i.invalidate(i.iter.SeekGE(key)) }
558 1 : func (i *invalidatingIter) SeekLT(key []byte) (*Span, error) { return i.invalidate(i.iter.SeekLT(key)) }
559 1 : func (i *invalidatingIter) First() (*Span, error) { return i.invalidate(i.iter.First()) }
560 1 : func (i *invalidatingIter) Last() (*Span, error) { return i.invalidate(i.iter.Last()) }
561 1 : func (i *invalidatingIter) Next() (*Span, error) { return i.invalidate(i.iter.Next()) }
562 1 : func (i *invalidatingIter) Prev() (*Span, error) { return i.invalidate(i.iter.Prev()) }
563 :
564 1 : func (i *invalidatingIter) Close() {
565 1 : _, _ = i.invalidate(nil, nil)
566 1 : i.iter.Close()
567 1 : }
568 :
569 0 : func (i *invalidatingIter) WrapChildren(wrap WrapFn) {
570 0 : i.iter = wrap(i.iter)
571 0 : }
572 :
573 : // DebugTree is part of the FragmentIterator interface.
574 0 : func (i *invalidatingIter) DebugTree(tp treeprinter.Node) {
575 0 : n := tp.Childf("%T(%p)", i, i)
576 0 : if i.iter != nil {
577 0 : i.iter.DebugTree(n)
578 0 : }
579 : }
|
