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