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 : "fmt"
9 :
10 : "github.com/cockroachdb/pebble/internal/base"
11 : "github.com/cockroachdb/pebble/internal/treeprinter"
12 : )
13 :
14 : // WrapFn is the prototype for a function that wraps a FragmentIterator.
15 : type WrapFn func(in FragmentIterator) FragmentIterator
16 :
17 : // InjectLogging wraps all iterators in a stack with logging iterators,
18 : // producing log messages showing each operation and its result.
19 0 : func InjectLogging(iter FragmentIterator, logger base.Logger) FragmentIterator {
20 0 : // All iterators in the stack will use the same logging state.
21 0 : state := &loggingState{
22 0 : log: logger,
23 0 : }
24 0 : var wrap WrapFn
25 0 : wrap = func(in FragmentIterator) FragmentIterator {
26 0 : if in == nil {
27 0 : return nil
28 0 : }
29 : // Recursively wrap all descendants.
30 0 : in.WrapChildren(wrap)
31 0 : return newLoggingIter(state, in)
32 : }
33 0 : return wrap(iter)
34 : }
35 :
36 0 : func newLoggingIter(state *loggingState, iter FragmentIterator) FragmentIterator {
37 0 : return &loggingIter{
38 0 : iter: iter,
39 0 : state: state,
40 0 : context: fmt.Sprintf("%T(%p):", iter, iter),
41 0 : }
42 0 : }
43 :
44 : // loggingIter is a pass-through FragmentIterator wrapper which performs checks
45 : // on what the wrapped iterator returns.
46 : type loggingIter struct {
47 : iter FragmentIterator
48 : state *loggingState
49 : context string
50 : }
51 :
52 : // loggingState is shared by all iterators in a stack.
53 : type loggingState struct {
54 : node treeprinter.Node
55 : log base.Logger
56 : }
57 :
58 0 : func (i *loggingIter) opStartf(format string, args ...any) func(results ...any) {
59 0 : savedNode := i.state.node
60 0 :
61 0 : n := i.state.node
62 0 : topLevelOp := false
63 0 : if n == (treeprinter.Node{}) {
64 0 : n = treeprinter.New()
65 0 : topLevelOp = true
66 0 : }
67 0 : op := fmt.Sprintf(format, args...)
68 0 :
69 0 : child := n.Childf("%s %s", i.context, op)
70 0 : i.state.node = child
71 0 :
72 0 : return func(results ...any) {
73 0 : if len(results) > 0 {
74 0 : child.Childf("%s", fmt.Sprint(results...))
75 0 : }
76 0 : if topLevelOp {
77 0 : for _, row := range n.FormattedRows() {
78 0 : i.state.log.Infof("%s\n", row)
79 0 : }
80 : }
81 0 : i.state.node = savedNode
82 : }
83 : }
84 :
85 : var _ FragmentIterator = (*loggingIter)(nil)
86 :
87 : // SeekGE implements FragmentIterator.
88 0 : func (i *loggingIter) SeekGE(key []byte) (*Span, error) {
89 0 : opEnd := i.opStartf("SeekGE(%q)", key)
90 0 : span, err := i.iter.SeekGE(key)
91 0 : opEnd(span, err)
92 0 : return span, err
93 0 : }
94 :
95 : // SeekLT implements FragmentIterator.
96 0 : func (i *loggingIter) SeekLT(key []byte) (*Span, error) {
97 0 : opEnd := i.opStartf("SeekLT(%q)", key)
98 0 : span, err := i.iter.SeekLT(key)
99 0 : opEnd(span, err)
100 0 : return span, err
101 0 : }
102 :
103 : // First implements FragmentIterator.
104 0 : func (i *loggingIter) First() (*Span, error) {
105 0 : opEnd := i.opStartf("First()")
106 0 : span, err := i.iter.First()
107 0 : opEnd(span, err)
108 0 : return span, err
109 0 : }
110 :
111 : // Last implements FragmentIterator.
112 0 : func (i *loggingIter) Last() (*Span, error) {
113 0 : opEnd := i.opStartf("Last()")
114 0 : span, err := i.iter.Last()
115 0 : opEnd(span, err)
116 0 : return span, err
117 0 : }
118 :
119 : // Next implements FragmentIterator.
120 0 : func (i *loggingIter) Next() (*Span, error) {
121 0 : opEnd := i.opStartf("Next()")
122 0 : span, err := i.iter.Next()
123 0 : opEnd(span, err)
124 0 : return span, err
125 0 : }
126 :
127 : // Prev implements FragmentIterator.
128 0 : func (i *loggingIter) Prev() (*Span, error) {
129 0 : opEnd := i.opStartf("Prev()")
130 0 : span, err := i.iter.Prev()
131 0 : opEnd(span, err)
132 0 : return span, err
133 0 : }
134 :
135 : // Close implements FragmentIterator.
136 0 : func (i *loggingIter) Close() error {
137 0 : opEnd := i.opStartf("Close()")
138 0 : err := i.iter.Close()
139 0 : if err != nil {
140 0 : opEnd(err)
141 0 : } else {
142 0 : opEnd()
143 0 : }
144 0 : return err
145 : }
146 :
147 : // WrapChildren implements FragmentIterator.
148 0 : func (i *loggingIter) WrapChildren(wrap WrapFn) {
149 0 : i.iter = wrap(i.iter)
150 0 : }
|