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