github.com/cockroachdb/pebble@v1.1.1-0.20240513155919-3622ade60459/internal/keyspan/iter_test.go (about)

     1  // Copyright 2018 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  	"bytes"
     9  	"fmt"
    10  	"strings"
    11  	"testing"
    12  
    13  	"github.com/cockroachdb/datadriven"
    14  	"github.com/cockroachdb/pebble/internal/base"
    15  )
    16  
    17  func runFragmentIteratorCmd(iter FragmentIterator, input string, extraInfo func() string) string {
    18  	var b bytes.Buffer
    19  	for _, line := range strings.Split(input, "\n") {
    20  		parts := strings.Fields(line)
    21  		if len(parts) == 0 {
    22  			continue
    23  		}
    24  		var span *Span
    25  		switch parts[0] {
    26  		case "seek-ge":
    27  			if len(parts) != 2 {
    28  				return "seek-ge <key>\n"
    29  			}
    30  			span = iter.SeekGE([]byte(strings.TrimSpace(parts[1])))
    31  		case "seek-lt":
    32  			if len(parts) != 2 {
    33  				return "seek-lt <key>\n"
    34  			}
    35  			span = iter.SeekLT([]byte(strings.TrimSpace(parts[1])))
    36  		case "first":
    37  			span = iter.First()
    38  		case "last":
    39  			span = iter.Last()
    40  		case "next":
    41  			span = iter.Next()
    42  		case "prev":
    43  			span = iter.Prev()
    44  		default:
    45  			return fmt.Sprintf("unknown op: %s", parts[0])
    46  		}
    47  		if span != nil {
    48  			fmt.Fprintf(&b, "%s", span)
    49  			if extraInfo != nil {
    50  				fmt.Fprintf(&b, " (%s)", extraInfo())
    51  			}
    52  			b.WriteByte('\n')
    53  		} else if err := iter.Error(); err != nil {
    54  			fmt.Fprintf(&b, "err=%v\n", err)
    55  		} else {
    56  			fmt.Fprintf(&b, ".\n")
    57  		}
    58  	}
    59  	return b.String()
    60  }
    61  
    62  func TestIter(t *testing.T) {
    63  	var spans []Span
    64  	datadriven.RunTest(t, "testdata/iter", func(t *testing.T, d *datadriven.TestData) string {
    65  		switch d.Cmd {
    66  		case "define":
    67  			spans = nil
    68  			for _, line := range strings.Split(d.Input, "\n") {
    69  				spans = append(spans, ParseSpan(line))
    70  			}
    71  			return ""
    72  
    73  		case "iter":
    74  			iter := NewIter(base.DefaultComparer.Compare, spans)
    75  			defer iter.Close()
    76  			return runFragmentIteratorCmd(iter, d.Input, nil)
    77  		default:
    78  			return fmt.Sprintf("unknown command: %s", d.Cmd)
    79  		}
    80  	})
    81  }
    82  
    83  // invalidatingIter wraps a FragmentIterator and implements FragmentIterator
    84  // itself. Spans surfaced by the inner iterator are copied to buffers that are
    85  // zeroed by sbubsequent iterator positioning calls. This is intended to help
    86  // surface bugs in improper lifetime expectations of Spans.
    87  type invalidatingIter struct {
    88  	iter FragmentIterator
    89  	bufs [][]byte
    90  	keys []Key
    91  	span Span
    92  }
    93  
    94  // invalidatingIter implements FragmentIterator.
    95  var _ FragmentIterator = (*invalidatingIter)(nil)
    96  
    97  func (i *invalidatingIter) invalidate(s *Span) *Span {
    98  	// Zero the entirety of the byte bufs and the keys slice.
    99  	for j := range i.bufs {
   100  		for k := range i.bufs[j] {
   101  			i.bufs[j][k] = 0x00
   102  		}
   103  		i.bufs[j] = nil
   104  	}
   105  	for j := range i.keys {
   106  		i.keys[j] = Key{}
   107  	}
   108  	if s == nil {
   109  		return nil
   110  	}
   111  
   112  	// Copy all of the span's slices into slices owned by the invalidating iter
   113  	// that we can invalidate on a subsequent positioning method.
   114  	i.bufs = i.bufs[:0]
   115  	i.keys = i.keys[:0]
   116  	i.span = Span{
   117  		Start: i.saveBytes(s.Start),
   118  		End:   i.saveBytes(s.End),
   119  	}
   120  	for j := range s.Keys {
   121  		i.keys = append(i.keys, Key{
   122  			Trailer: s.Keys[j].Trailer,
   123  			Suffix:  i.saveBytes(s.Keys[j].Suffix),
   124  			Value:   i.saveBytes(s.Keys[j].Value),
   125  		})
   126  	}
   127  	i.span.Keys = i.keys
   128  	return &i.span
   129  }
   130  
   131  func (i *invalidatingIter) saveBytes(b []byte) []byte {
   132  	if b == nil {
   133  		return nil
   134  	}
   135  	saved := append([]byte(nil), b...)
   136  	i.bufs = append(i.bufs, saved)
   137  	return saved
   138  }
   139  
   140  func (i *invalidatingIter) SeekGE(key []byte) *Span { return i.invalidate(i.iter.SeekGE(key)) }
   141  func (i *invalidatingIter) SeekLT(key []byte) *Span { return i.invalidate(i.iter.SeekLT(key)) }
   142  func (i *invalidatingIter) First() *Span            { return i.invalidate(i.iter.First()) }
   143  func (i *invalidatingIter) Last() *Span             { return i.invalidate(i.iter.Last()) }
   144  func (i *invalidatingIter) Next() *Span             { return i.invalidate(i.iter.Next()) }
   145  func (i *invalidatingIter) Prev() *Span             { return i.invalidate(i.iter.Prev()) }
   146  func (i *invalidatingIter) Close() error            { return i.iter.Close() }
   147  func (i *invalidatingIter) Error() error            { return i.iter.Error() }