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 0 : func InjectLogging(iter FragmentIterator, logger base.Logger) FragmentIterator {
21 0 : // All iterators in the stack will use the same logging state.
22 0 : state := &loggingState{
23 0 : log: logger,
24 0 : }
25 0 : var wrap WrapFn
26 0 : wrap = func(in FragmentIterator) FragmentIterator {
27 0 : if in == nil {
28 0 : return nil
29 0 : }
30 : // Recursively wrap all descendants.
31 0 : in.WrapChildren(wrap)
32 0 : return newLoggingIter(state, in)
33 : }
34 0 : return wrap(iter)
35 : }
36 :
37 0 : func newLoggingIter(state *loggingState, iter FragmentIterator) FragmentIterator {
38 0 : return &loggingIter{
39 0 : iter: iter,
40 0 : state: state,
41 0 : context: fmt.Sprintf("%T(%p):", iter, iter),
42 0 : }
43 0 : }
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 0 : func (i *loggingIter) opStartf(format string, args ...any) func(results ...any) {
60 0 : savedNode := i.state.node
61 0 :
62 0 : n := i.state.node
63 0 : topLevelOp := false
64 0 : if n == (treeprinter.Node{}) {
65 0 : n = treeprinter.New()
66 0 : topLevelOp = true
67 0 : }
68 0 : op := fmt.Sprintf(format, args...)
69 0 :
70 0 : child := n.Childf("%s %s", i.context, op)
71 0 : i.state.node = child
72 0 :
73 0 : return func(results ...any) {
74 0 : if len(results) > 0 {
75 0 : child.Childf("%s", fmt.Sprint(results...))
76 0 : }
77 0 : if topLevelOp {
78 0 : for _, row := range n.FormattedRows() {
79 0 : i.state.log.Infof("%s\n", row)
80 0 : }
81 : }
82 0 : i.state.node = savedNode
83 : }
84 : }
85 :
86 : var _ FragmentIterator = (*loggingIter)(nil)
87 :
88 : // SeekGE implements FragmentIterator.
89 0 : func (i *loggingIter) SeekGE(key []byte) (*Span, error) {
90 0 : opEnd := i.opStartf("SeekGE(%q)", key)
91 0 : span, err := i.iter.SeekGE(key)
92 0 : opEnd(span, err)
93 0 : return span, err
94 0 : }
95 :
96 : // SeekLT implements FragmentIterator.
97 0 : func (i *loggingIter) SeekLT(key []byte) (*Span, error) {
98 0 : opEnd := i.opStartf("SeekLT(%q)", key)
99 0 : span, err := i.iter.SeekLT(key)
100 0 : opEnd(span, err)
101 0 : return span, err
102 0 : }
103 :
104 : // First implements FragmentIterator.
105 0 : func (i *loggingIter) First() (*Span, error) {
106 0 : opEnd := i.opStartf("First()")
107 0 : span, err := i.iter.First()
108 0 : opEnd(span, err)
109 0 : return span, err
110 0 : }
111 :
112 : // Last implements FragmentIterator.
113 0 : func (i *loggingIter) Last() (*Span, error) {
114 0 : opEnd := i.opStartf("Last()")
115 0 : span, err := i.iter.Last()
116 0 : opEnd(span, err)
117 0 : return span, err
118 0 : }
119 :
120 : // Next implements FragmentIterator.
121 0 : func (i *loggingIter) Next() (*Span, error) {
122 0 : opEnd := i.opStartf("Next()")
123 0 : span, err := i.iter.Next()
124 0 : opEnd(span, err)
125 0 : return span, err
126 0 : }
127 :
128 : // Prev implements FragmentIterator.
129 0 : func (i *loggingIter) Prev() (*Span, error) {
130 0 : opEnd := i.opStartf("Prev()")
131 0 : span, err := i.iter.Prev()
132 0 : opEnd(span, err)
133 0 : return span, err
134 0 : }
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 0 : func (i *loggingIter) Close() {
143 0 : opEnd := i.opStartf("Close()")
144 0 : i.iter.Close()
145 0 : opEnd()
146 0 : }
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 : }
|
