Repository: cockroachdb/pebble
Pull request #3679 (2024-06-21 02:52:19)

sstable: wrap raw fragment iterators with checkers, fix issues found
#### keyspan: copy Suffix/Value in CopyFrom


#### tool: use DeepClone instead of ShallowClone

We were performing shallow clones of spans, which is not sufficient
when the spans come from a FragmentIterator.

#### sstable: wrap raw fragment iterators with checkers

This change wraps raw range del and range key iterators with
`assertIter` and `invalidatingIter` and we fix a few cases of using
stale values (especially Prefix/Value for range keys).

cockroach-teamcity: 

<!-- Reviewable:start -->
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/cockroachdb/pebble/3679)
<!-- Reviewable:end -->


Patch:
diff --git a/internal/keyspan/assert_iter.go b/internal/keyspan/assert_iter.go
index 6fb6b59b..53e04b49 100644
--- a/internal/keyspan/assert_iter.go
+++ b/internal/keyspan/assert_iter.go
@@ -9,6 +9,7 @@ import (
 
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/pebble/internal/base"
+	"github.com/cockroachdb/pebble/internal/fastrand"
 	"github.com/cockroachdb/pebble/internal/invariants"
 )
 
@@ -20,12 +21,21 @@ func Assert(iter FragmentIterator, cmp base.Compare) FragmentIterator {
 	}
 }
 
-// MaybeAssert wraps an iterator and asserts that operations return sane
-// results if we are in testing mode.
+// MaybeAssert potentially wraps an iterator and asserts that operations return
+// sane results if we are in testing mode.
 func MaybeAssert(iter FragmentIterator, cmp base.Compare) FragmentIterator {
-	if invariants.Enabled && iter != nil {
+	if invariants.Enabled && iter != nil && fastrand.Uint32()%4 > 0 {
 		// Don't wrap an assertIter.
-		if _, ok := iter.(*assertIter); !ok {
+		switch iter.(type) {
+		case *assertIter:
+			return iter
+		case *invalidatingIter:
+			return iter
+		default:
+			// Also wrap with an invalidating iterator some of the time.
+			if fastrand.Uint32()%2 > 0 {
+				iter = NewInvalidatingIter(iter)
+			}
 			return Assert(iter, cmp)
 		}
 	}
diff --git a/internal/keyspan/interleaving_iter.go b/internal/keyspan/interleaving_iter.go
index dda0da6d..883526ba 100644
--- a/internal/keyspan/interleaving_iter.go
+++ b/internal/keyspan/interleaving_iter.go
@@ -1117,9 +1117,13 @@ func (i *InterleavingIter) savePoint(kv *base.InternalKV) {
 //
 // Span will never return an invalid or empty span.
 func (i *InterleavingIter) Span() *Span {
+	if invariants.Enabled && i.pointIter == nil {
+		panic("Span() called after close")
+	}
 	if !i.withinSpan || len(i.span.Keys) == 0 {
 		return nil
-	} else if i.truncated {
+	}
+	if i.truncated {
 		return &i.truncatedSpan
 	}
 	return i.span
@@ -1153,7 +1157,9 @@ func (i *InterleavingIter) Error() error {
 // Close implements (base.InternalIterator).Close.
 func (i *InterleavingIter) Close() error {
 	err := i.pointIter.Close()
+	i.pointIter = nil
 	i.keyspanIter.Close()
+	i.keyspanIter = nil
 	return err
 }
 
diff --git a/internal/keyspan/iter.go b/internal/keyspan/iter.go
index f8bc722e..217221c6 100644
--- a/internal/keyspan/iter.go
+++ b/internal/keyspan/iter.go
@@ -57,7 +57,8 @@ type FragmentIterator interface {
 
 	// Close closes the iterator. It is not in general valid to call Close
 	// multiple times. Other methods should not be called after the iterator has
-	// been closed.
+	// been closed. Spans returned by a previous method should also not be used
+	// after the iterator has been closed.
 	Close()
 
 	// WrapChildren wraps any child iterators using the given function. The
diff --git a/internal/keyspan/span.go b/internal/keyspan/span.go
index b78cc44e..3aa07eef 100644
--- a/internal/keyspan/span.go
+++ b/internal/keyspan/span.go
@@ -7,6 +7,7 @@ package keyspan // import "github.com/cockroachdb/pebble/internal/keyspan"
 import (
 	"bytes"
 	"fmt"
+	"slices"
 	"sort"
 	"strconv"
 	"strings"
@@ -99,6 +100,28 @@ func (k Key) Equal(equal base.Equal, b Key) bool {
 		bytes.Equal(k.Value, b.Value)
 }
 
+// CopyFrom copies the contents of another key, retaining the Suffix and Value slices.
+func (k *Key) CopyFrom(other Key) {
+	k.Trailer = other.Trailer
+	k.Suffix = append(k.Suffix[:0], other.Suffix...)
+	k.Value = append(k.Value[:0], other.Value...)
+}
+
+// Clone creates a deep clone of the key, copying the Suffix and Value
+// slices.
+func (k Key) Clone() Key {
+	res := Key{
+		Trailer: k.Trailer,
+	}
+	if len(k.Suffix) > 0 {
+		res.Suffix = slices.Clone(k.Suffix)
+	}
+	if len(k.Value) > 0 {
+		res.Value = slices.Clone(k.Value)
+	}
+	return res
+}
+
 func (k Key) String() string {
 	var b strings.Builder
 	fmt.Fprintf(&b, "(#%d,%s", k.SeqNum(), k.Kind())
@@ -317,41 +340,17 @@ func (s *Span) VisibleAt(snapshot uint64) bool {
 	}
 }
 
-// ShallowClone returns the span with a Keys slice owned by the span itself.
-// None of the key byte slices are cloned (see Span.DeepClone).
-func (s *Span) ShallowClone() Span {
-	c := Span{
-		Start:     s.Start,
-		End:       s.End,
-		Keys:      make([]Key, len(s.Keys)),
-		KeysOrder: s.KeysOrder,
-	}
-	copy(c.Keys, s.Keys)
-	return c
-}
-
-// DeepClone clones the span, creating copies of all contained slices. DeepClone
-// is intended for non-production code paths like tests, the level checker, etc
-// because it is allocation heavy.
-func (s *Span) DeepClone() Span {
+// Clone clones the span, creating copies of all contained slices. Clone is
+// allocation heavy and should not be used in hot paths.
+func (s *Span) Clone() Span {
 	c := Span{
-		Start:     make([]byte, len(s.Start)),
-		End:       make([]byte, len(s.End)),
-		Keys:      make([]Key, len(s.Keys)),
+		Start:     slices.Clone(s.Start),
+		End:       slices.Clone(s.End),
 		KeysOrder: s.KeysOrder,
 	}
-	copy(c.Start, s.Start)
-	copy(c.End, s.End)
-	for i := range s.Keys {
-		c.Keys[i].Trailer = s.Keys[i].Trailer
-		if len(s.Keys[i].Suffix) > 0 {
-			c.Keys[i].Suffix = make([]byte, len(s.Keys[i].Suffix))
-			copy(c.Keys[i].Suffix, s.Keys[i].Suffix)
-		}
-		if len(s.Keys[i].Value) > 0 {
-			c.Keys[i].Value = make([]byte, len(s.Keys[i].Value))
-			copy(c.Keys[i].Value, s.Keys[i].Value)
-		}
+	c.Keys = make([]Key, len(s.Keys))
+	for i := range c.Keys {
+		c.Keys[i] = s.Keys[i].Clone()
 	}
 	return c
 }
@@ -407,12 +406,22 @@ func (s *Span) Reset() {
 	s.Keys = s.Keys[:0]
 }
 
-// CopyFrom copies the contents of the other span, retaining the slices
-// allocated in this span. Key.Suffix and Key.Value slices are not cloned.
+// CopyFrom deep-copies the contents of the other span, retaining the slices
+// allocated in this span.
 func (s *Span) CopyFrom(other *Span) {
 	s.Start = append(s.Start[:0], other.Start...)
 	s.End = append(s.End[:0], other.End...)
-	s.Keys = append(s.Keys[:0], other.Keys...)
+
+	// We want to preserve any existing Suffix/Value buffers.
+	if cap(s.Keys) >= len(other.Keys) {
+		s.Keys = s.Keys[:len(other.Keys)]
+	} else {
+		s.Keys = append(s.Keys[:cap(s.Keys)], make([]Key, len(other.Keys)-cap(s.Keys))...)
+	}
+	for i := range other.Keys {
+		s.Keys[i].CopyFrom(other.Keys[i])
+	}
+
 	s.KeysOrder = other.KeysOrder
 }
 
diff --git a/level_checker.go b/level_checker.go
index 12580672..1f24272c 100644
--- a/level_checker.go
+++ b/level_checker.go
@@ -459,7 +459,7 @@ func addTombstonesFromIter(
 		if t.Empty() {
 			continue
 		}
-		t = t.DeepClone()
+		t = t.Clone()
 		// This is mainly a test for rangeDelV2 formatted blocks which are expected to
 		// be ordered and fragmented on disk. But we anyways check for memtables,
 		// rangeDelV1 as well.
diff --git a/scan_internal.go b/scan_internal.go
index 8eb87bdd..fcec4ab8 100644
--- a/scan_internal.go
+++ b/scan_internal.go
@@ -778,7 +778,7 @@ func scanInternalImpl(
 				// call visitRangeKey.
 				keysCopy := make([]keyspan.Key, len(span.Keys))
 				for i := range span.Keys {
-					keysCopy[i] = span.Keys[i]
+					keysCopy[i].CopyFrom(span.Keys[i])
 					keysCopy[i].Trailer = base.MakeTrailer(0, span.Keys[i].Kind())
 				}
 				keyspan.SortKeysByTrailer(&keysCopy)
diff --git a/sstable/block_iter.go b/sstable/block_iter.go
index 6180ffdf..681494a3 100644
--- a/sstable/block_iter.go
+++ b/sstable/block_iter.go
@@ -8,13 +8,13 @@ import (
 	"bytes"
 	"context"
 	"encoding/binary"
+	"slices"
 	"unsafe"
 
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/pebble/internal/base"
 	"github.com/cockroachdb/pebble/internal/invariants"
 	"github.com/cockroachdb/pebble/internal/manual"
-	"golang.org/x/exp/slices"
 )
 
 // blockIter is an iterator over a single block of data.
diff --git a/sstable/reader.go b/sstable/reader.go
index a4fe0c75..f738c1cf 100644
--- a/sstable/reader.go
+++ b/sstable/reader.go
@@ -461,7 +461,7 @@ func (r *Reader) NewRawRangeDelIter(transforms IterTransforms) (keyspan.Fragment
 	if err := i.blockIter.initHandle(r.Compare, r.Split, h, transforms); err != nil {
 		return nil, err
 	}
-	return i, nil
+	return keyspan.MaybeAssert(i, r.Compare), nil
 }
 
 // NewRawRangeKeyIter returns an internal iterator for the contents of the
@@ -488,7 +488,7 @@ func (r *Reader) NewRawRangeKeyIter(transforms IterTransforms) (keyspan.Fragment
 	if err := i.blockIter.initHandle(r.Compare, r.Split, h, transforms); err != nil {
 		return nil, err
 	}
-	return i, nil
+	return keyspan.MaybeAssert(i, r.Compare), nil
 }
 
 func (r *Reader) readIndex(
diff --git a/tool/find.go b/tool/find.go
index a379efa8..666b3b6c 100644
--- a/tool/find.go
+++ b/tool/find.go
@@ -487,7 +487,7 @@ func (f *findT) searchTables(stdout io.Writer, searchKey []byte, refs []findRef)
 					if !t.Contains(r.Compare, searchKey) {
 						continue
 					}
-					tombstones = append(tombstones, t.ShallowClone())
+					tombstones = append(tombstones, t.Clone())
 				}
 				if err != nil {
 					return nil, err
diff --git a/tool/sstable.go b/tool/sstable.go
index 9ef22174..2fc88a53 100644
--- a/tool/sstable.go
+++ b/tool/sstable.go
@@ -407,7 +407,7 @@ func (s *sstableT) runScan(cmd *cobra.Command, args []string) {
 					// The range tombstone lies before the scan range.
 					continue
 				}
-				tombstones = append(tombstones, t.ShallowClone())
+				tombstones = append(tombstones, t.Clone())
 			}
 			if err != nil {
 				return nil, err


Test patch:
diff --git a/internal/keyspan/test_utils.go b/internal/keyspan/test_utils.go
index e9c1ef93..9a1040a4 100644
--- a/internal/keyspan/test_utils.go
+++ b/internal/keyspan/test_utils.go
@@ -502,10 +502,10 @@ type invalidatingIter struct {
 var _ FragmentIterator = (*invalidatingIter)(nil)
 
 func (i *invalidatingIter) invalidate(s *Span, err error) (*Span, error) {
-	// Zero the entirety of the byte bufs and the keys slice.
+	// Mangle the entirety of the byte bufs and the keys slice.
 	for j := range i.bufs {
 		for k := range i.bufs[j] {
-			i.bufs[j][k] = 0x00
+			i.bufs[j][k] = 0xff
 		}
 		i.bufs[j] = nil
 	}
@@ -521,8 +521,9 @@ func (i *invalidatingIter) invalidate(s *Span, err error) (*Span, error) {
 	i.bufs = i.bufs[:0]
 	i.keys = i.keys[:0]
 	i.span = Span{
-		Start: i.saveBytes(s.Start),
-		End:   i.saveBytes(s.End),
+		Start:     i.saveBytes(s.Start),
+		End:       i.saveBytes(s.End),
+		KeysOrder: s.KeysOrder,
 	}
 	for j := range s.Keys {
 		i.keys = append(i.keys, Key{
@@ -550,7 +551,12 @@ func (i *invalidatingIter) First() (*Span, error)            { return i.invalida
 func (i *invalidatingIter) Last() (*Span, error)             { return i.invalidate(i.iter.Last()) }
 func (i *invalidatingIter) Next() (*Span, error)             { return i.invalidate(i.iter.Next()) }
 func (i *invalidatingIter) Prev() (*Span, error)             { return i.invalidate(i.iter.Prev()) }
-func (i *invalidatingIter) Close()                           { i.iter.Close() }
+
+func (i *invalidatingIter) Close() {
+	_, _ = i.invalidate(nil, nil)
+	i.iter.Close()
+}
+
 func (i *invalidatingIter) WrapChildren(wrap WrapFn) {
 	i.iter = wrap(i.iter)
 }
diff --git a/internal/keyspan/truncate_test.go b/internal/keyspan/truncate_test.go
index 0253a088..c70c6057 100644
--- a/internal/keyspan/truncate_test.go
+++ b/internal/keyspan/truncate_test.go
@@ -58,7 +58,7 @@ func TestTruncate(t *testing.T) {
 			var s *Span
 			var err error
 			for s, err = tIter.First(); s != nil; s, err = tIter.Next() {
-				truncated = append(truncated, s.ShallowClone())
+				truncated = append(truncated, s.Clone())
 			}
 			require.NoError(t, err)
 			return formatAlphabeticSpans(truncated)
diff --git a/internal/lint/lint_test.go b/internal/lint/lint_test.go
index 1ff9b527..3386f788 100644
--- a/internal/lint/lint_test.go
+++ b/internal/lint/lint_test.go
@@ -180,8 +180,9 @@ func TestLint(t *testing.T) {
 
 		// Forbidden-import-pkg -> permitted-replacement-pkg
 		forbiddenImports := map[string]string{
-			"errors":     "github.com/cockroachdb/errors",
-			"pkg/errors": "github.com/cockroachdb/errors",
+			"errors":                  "github.com/cockroachdb/errors",
+			"pkg/errors":              "github.com/cockroachdb/errors",
+			"golang.org/x/exp/slices": "slices",
 		}
 
 		// grepBuf creates a grep string that matches any forbidden import pkgs.
diff --git a/sstable/test_utils.go b/sstable/test_utils.go
index f3a8fae4..0a161384 100644
--- a/sstable/test_utils.go
+++ b/sstable/test_utils.go
@@ -32,14 +32,14 @@ func ReadAll(
 	if rangeDelIter := testutils.CheckErr(reader.NewRawRangeDelIter(NoTransforms)); rangeDelIter != nil {
 		defer rangeDelIter.Close()
 		for s := testutils.CheckErr(rangeDelIter.First()); s != nil; s = testutils.CheckErr(rangeDelIter.Next()) {
-			rangeDels = append(rangeDels, s.DeepClone())
+			rangeDels = append(rangeDels, s.Clone())
 		}
 	}
 
 	if rangeKeyIter := testutils.CheckErr(reader.NewRawRangeKeyIter(NoTransforms)); rangeKeyIter != nil {
 		defer rangeKeyIter.Close()
 		for s := testutils.CheckErr(rangeKeyIter.First()); s != nil; s = testutils.CheckErr(rangeKeyIter.Next()) {
-			rangeKeys = append(rangeKeys, s.DeepClone())
+			rangeKeys = append(rangeKeys, s.Clone())
 		}
 	}
 	return points, rangeDels, rangeKeys