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() error {
     137           1 :         opEnd := i.opStartf("Close()")
     138           1 :         err := i.iter.Close()
     139           1 :         if err != nil {
     140           0 :                 opEnd(err)
     141           1 :         } else {
     142           1 :                 opEnd()
     143           1 :         }
     144           1 :         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 : }