Repository: cockroachdb/pebble (go)

Issue:
metamorphic: use invalidating iterators in some runs
In some unit tests we use 'invalidating iterators' to ensure we're not relying on key/buffer stability where it's not guaranteed. Some runs of the metamorphic test could configure iterators with similar mechanics.

One possible explanation of cockroachlabs/support#2583 is corruption during compaction iteration. If there is a memory lifetime issue, use of an invalidating iterator in the metamorphic tests may help shake it out.

Pull request description:
db: use invalidating iterators with low probability under invariants
In invariants builds, with 10% probability wrap various internal iterators with
an invalidating iterator. This helps ensure we're not accidentally depending on
key byte slice stability where we shouldn't.

Close #2896.

Patch:
diff --git a/compaction.go b/compaction.go
index 945b8bf1..c3b42f53 100644
--- a/compaction.go
+++ b/compaction.go
@@ -18,6 +18,7 @@ import (
 
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/pebble/internal/base"
+	"github.com/cockroachdb/pebble/internal/invalidating"
 	"github.com/cockroachdb/pebble/internal/keyspan"
 	"github.com/cockroachdb/pebble/internal/manifest"
 	"github.com/cockroachdb/pebble/internal/private"
@@ -2847,6 +2848,7 @@ func (d *DB) runCompaction(
 		return nil, pendingOutputs, stats, err
 	}
 	c.allowedZeroSeqNum = c.allowZeroSeqNum()
+	iiter = invalidating.MaybeWrapIfInvariants(iiter)
 	iter := newCompactionIter(c.cmp, c.equal, c.formatKey, d.merge, iiter, snapshots,
 		&c.rangeDelFrag, &c.rangeKeyFrag, c.allowedZeroSeqNum, c.elideTombstone,
 		c.elideRangeTombstone, d.FormatMajorVersion())
diff --git a/db.go b/db.go
index 79a465bf..8794bc16 100644
--- a/db.go
+++ b/db.go
@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/pebble/internal/arenaskl"
 	"github.com/cockroachdb/pebble/internal/base"
+	"github.com/cockroachdb/pebble/internal/invalidating"
 	"github.com/cockroachdb/pebble/internal/invariants"
 	"github.com/cockroachdb/pebble/internal/keyspan"
 	"github.com/cockroachdb/pebble/internal/manifest"
@@ -1149,6 +1150,7 @@ func finishInitializingIter(ctx context.Context, buf *iterAlloc) *Iterator {
 				},
 			}
 			dbi.iter = &dbi.lazyCombinedIter
+			dbi.iter = invalidating.MaybeWrapIfInvariants(dbi.iter)
 		} else {
 			dbi.lazyCombinedIter.combinedIterState = combinedIterState{
 				initialized: true,
@@ -1426,7 +1428,8 @@ func (i *Iterator) constructPointIter(
 		li.initRangeDel(&mlevels[mlevelsIndex].rangeDelIter)
 		li.initBoundaryContext(&mlevels[mlevelsIndex].levelIterBoundaryContext)
 		li.initCombinedIterState(&i.lazyCombinedIter.combinedIterState)
-		mlevels[mlevelsIndex].iter = li
+		mlevels[mlevelsIndex].levelIter = li
+		mlevels[mlevelsIndex].iter = invalidating.MaybeWrapIfInvariants(li)
 
 		levelsIndex++
 		mlevelsIndex++
@@ -1452,7 +1455,7 @@ func (i *Iterator) constructPointIter(
 	buf.merging.snapshot = i.seqNum
 	buf.merging.batchSnapshot = i.batchSeqNum
 	buf.merging.combinedIterState = &i.lazyCombinedIter.combinedIterState
-	i.pointIter = &buf.merging
+	i.pointIter = invalidating.MaybeWrapIfInvariants(&buf.merging)
 	i.merging = &buf.merging
 }
 
diff --git a/internal/invalidating/iter.go b/internal/invalidating/iter.go
new file mode 100644
index 00000000..48909ec9
--- /dev/null
+++ b/internal/invalidating/iter.go
@@ -0,0 +1,162 @@
+// Copyright 2023 The LevelDB-Go and Pebble Authors. All rights reserved. Use
+// of this source code is governed by a BSD-style license that can be found in
+// the LICENSE file.
+
+package invalidating
+
+import (
+	"github.com/cockroachdb/pebble/internal/base"
+	"github.com/cockroachdb/pebble/internal/fastrand"
+	"github.com/cockroachdb/pebble/internal/invariants"
+)
+
+// MaybeWrapIfInvariants wraps some iterators with an invalidating iterator.
+// MaybeWrapIfInvariants does nothing in non-invariant builds.
+func MaybeWrapIfInvariants(iter base.InternalIterator) base.InternalIterator {
+	if invariants.Enabled {
+		if fastrand.Uint32n(10) == 1 {
+			return NewIter(iter)
+		}
+	}
+	return iter
+}
+
+// iter tests unsafe key/value slice reuse by modifying the last
+// returned key/value to all 1s.
+type iter struct {
+	iter        base.InternalIterator
+	lastKey     *base.InternalKey
+	lastValue   base.LazyValue
+	ignoreKinds [base.InternalKeyKindMax + 1]bool
+	err         error
+}
+
+// Option configures the behavior of an invalidating iterator.
+type Option interface {
+	apply(*iter)
+}
+
+type funcOpt func(*iter)
+
+func (f funcOpt) apply(i *iter) { f(i) }
+
+// IgnoreKinds constructs an Option that configures an invalidating iterator to
+// skip trashing k/v pairs with the provided key kinds. Some iterators provided
+// key stability guarantees for specific key kinds.
+func IgnoreKinds(kinds ...base.InternalKeyKind) Option {
+	return funcOpt(func(i *iter) {
+		for _, kind := range kinds {
+			i.ignoreKinds[kind] = true
+		}
+	})
+}
+
+// NewIter constructs a new invalidating iterator that wraps the provided
+// iterator, trashing buffers for previously returned keys.
+func NewIter(originalIterator base.InternalIterator, opts ...Option) base.InternalIterator {
+	i := &iter{iter: originalIterator}
+	for _, opt := range opts {
+		opt.apply(i)
+	}
+	return i
+}
+
+func (i *iter) update(
+	key *base.InternalKey, value base.LazyValue,
+) (*base.InternalKey, base.LazyValue) {
+	i.trashLastKV()
+	if key == nil {
+		i.lastKey = nil
+		i.lastValue = base.LazyValue{}
+		return nil, base.LazyValue{}
+	}
+
+	i.lastKey = &base.InternalKey{}
+	*i.lastKey = key.Clone()
+	i.lastValue = base.LazyValue{
+		ValueOrHandle: append(make([]byte, 0, len(value.ValueOrHandle)), value.ValueOrHandle...),
+	}
+	if value.Fetcher != nil {
+		fetcher := new(base.LazyFetcher)
+		*fetcher = *value.Fetcher
+		i.lastValue.Fetcher = fetcher
+	}
+	return i.lastKey, i.lastValue
+}
+
+func (i *iter) trashLastKV() {
+	if i.lastKey == nil {
+		return
+	}
+	if i.ignoreKinds[i.lastKey.Kind()] {
+		return
+	}
+
+	if i.lastKey != nil {
+		for j := range i.lastKey.UserKey {
+			i.lastKey.UserKey[j] = 0xff
+		}
+		i.lastKey.Trailer = 0xffffffffffffffff
+	}
+	for j := range i.lastValue.ValueOrHandle {
+		i.lastValue.ValueOrHandle[j] = 0xff
+	}
+	if i.lastValue.Fetcher != nil {
+		// Not all the LazyFetcher fields are visible, so we zero out the last
+		// value's Fetcher struct entirely.
+		*i.lastValue.Fetcher = base.LazyFetcher{}
+	}
+}
+
+func (i *iter) SeekGE(key []byte, flags base.SeekGEFlags) (*base.InternalKey, base.LazyValue) {
+	return i.update(i.iter.SeekGE(key, flags))
+}
+
+func (i *iter) SeekPrefixGE(
+	prefix, key []byte, flags base.SeekGEFlags,
+) (*base.InternalKey, base.LazyValue) {
+	return i.update(i.iter.SeekPrefixGE(prefix, key, flags))
+}
+
+func (i *iter) SeekLT(key []byte, flags base.SeekLTFlags) (*base.InternalKey, base.LazyValue) {
+	return i.update(i.iter.SeekLT(key, flags))
+}
+
+func (i *iter) First() (*base.InternalKey, base.LazyValue) {
+	return i.update(i.iter.First())
+}
+
+func (i *iter) Last() (*base.InternalKey, base.LazyValue) {
+	return i.update(i.iter.Last())
+}
+
+func (i *iter) Next() (*base.InternalKey, base.LazyValue) {
+	return i.update(i.iter.Next())
+}
+
+func (i *iter) Prev() (*base.InternalKey, base.LazyValue) {
+	return i.update(i.iter.Prev())
+}
+
+func (i *iter) NextPrefix(succKey []byte) (*base.InternalKey, base.LazyValue) {
+	return i.update(i.iter.NextPrefix(succKey))
+}
+
+func (i *iter) Error() error {
+	if err := i.iter.Error(); err != nil {
+		return err
+	}
+	return i.err
+}
+
+func (i *iter) Close() error {
+	return i.iter.Close()
+}
+
+func (i *iter) SetBounds(lower, upper []byte) {
+	i.iter.SetBounds(lower, upper)
+}
+
+func (i *iter) String() string {
+	return i.iter.String()
+}
diff --git a/iterator.go b/iterator.go
index 13b07f56..dfaa6461 100644
--- a/iterator.go
+++ b/iterator.go
@@ -780,39 +780,41 @@ func (i *Iterator) maybeSampleRead() {
 func (i *Iterator) sampleRead() {
 	var topFile *manifest.FileMetadata
 	topLevel, numOverlappingLevels := numLevels, 0
-	if mi, ok := i.iter.(*mergingIter); ok {
-		if len(mi.levels) > 1 {
-			mi.ForEachLevelIter(func(li *levelIter) bool {
-				l := manifest.LevelToInt(li.level)
-				if f := li.iterFile; f != nil {
-					var containsKey bool
-					if i.pos == iterPosNext || i.pos == iterPosCurForward ||
-						i.pos == iterPosCurForwardPaused {
-						containsKey = i.cmp(f.SmallestPointKey.UserKey, i.key) <= 0
-					} else if i.pos == iterPosPrev || i.pos == iterPosCurReverse ||
-						i.pos == iterPosCurReversePaused {
-						containsKey = i.cmp(f.LargestPointKey.UserKey, i.key) >= 0
-					}
-					// Do nothing if the current key is not contained in f's
-					// bounds. We could seek the LevelIterator at this level
-					// to find the right file, but the performance impacts of
-					// doing that are significant enough to negate the benefits
-					// of read sampling in the first place. See the discussion
-					// at:
-					// https://github.com/cockroachdb/pebble/pull/1041#issuecomment-763226492
-					if containsKey {
-						numOverlappingLevels++
-						if numOverlappingLevels >= 2 {
-							// Terminate the loop early if at least 2 overlapping levels are found.
-							return true
-						}
-						topLevel = l
-						topFile = f
+	mi := i.merging
+	if mi == nil {
+		return
+	}
+	if len(mi.levels) > 1 {
+		mi.ForEachLevelIter(func(li *levelIter) bool {
+			l := manifest.LevelToInt(li.level)
+			if f := li.iterFile; f != nil {
+				var containsKey bool
+				if i.pos == iterPosNext || i.pos == iterPosCurForward ||
+					i.pos == iterPosCurForwardPaused {
+					containsKey = i.cmp(f.SmallestPointKey.UserKey, i.key) <= 0
+				} else if i.pos == iterPosPrev || i.pos == iterPosCurReverse ||
+					i.pos == iterPosCurReversePaused {
+					containsKey = i.cmp(f.LargestPointKey.UserKey, i.key) >= 0
+				}
+				// Do nothing if the current key is not contained in f's
+				// bounds. We could seek the LevelIterator at this level
+				// to find the right file, but the performance impacts of
+				// doing that are significant enough to negate the benefits
+				// of read sampling in the first place. See the discussion
+				// at:
+				// https://github.com/cockroachdb/pebble/pull/1041#issuecomment-763226492
+				if containsKey {
+					numOverlappingLevels++
+					if numOverlappingLevels >= 2 {
+						// Terminate the loop early if at least 2 overlapping levels are found.
+						return true
 					}
+					topLevel = l
+					topFile = f
 				}
-				return false
-			})
-		}
+			}
+			return false
+		})
 	}
 	if topFile == nil || topLevel >= numLevels {
 		return
diff --git a/merging_iter.go b/merging_iter.go
index 4e2426d8..a0766869 100644
--- a/merging_iter.go
+++ b/merging_iter.go
@@ -27,6 +27,10 @@ type mergingIterLevel struct {
 	// iterKey and iterValue cache the current key and value iter are pointed at.
 	iterKey   *InternalKey
 	iterValue base.LazyValue
+	// levelIter is non-nil if this level's iter is ultimately backed by a
+	// *levelIter. The handle in iter may have wrapped the levelIter with
+	// intermediary internalIterator implementations.
+	levelIter *levelIter
 
 	// levelIterBoundaryContext's fields are set when using levelIter, in order
 	// to surface sstable boundary keys and file-level context. See levelIter
@@ -1371,9 +1375,9 @@ func (m *mergingIter) DebugString() string {
 }
 
 func (m *mergingIter) ForEachLevelIter(fn func(li *levelIter) bool) {
-	for _, iter := range m.levels {
-		if li, ok := iter.iter.(*levelIter); ok {
-			if done := fn(li); done {
+	for _, ml := range m.levels {
+		if ml.levelIter != nil {
+			if done := fn(ml.levelIter); done {
 				break
 			}
 		}


Test patch:
diff --git a/compaction_iter_test.go b/compaction_iter_test.go
index 4282e521..115b816c 100644
--- a/compaction_iter_test.go
+++ b/compaction_iter_test.go
@@ -16,6 +16,7 @@ import (
 
 	"github.com/cockroachdb/datadriven"
 	"github.com/cockroachdb/pebble/internal/base"
+	"github.com/cockroachdb/pebble/internal/invalidating"
 	"github.com/cockroachdb/pebble/internal/keyspan"
 	"github.com/cockroachdb/pebble/internal/rangekey"
 	"github.com/cockroachdb/pebble/internal/testkeys"
@@ -110,8 +111,7 @@ func TestCompactionIter(t *testing.T) {
 			fi,
 			keyspan.NewIter(base.DefaultComparer.Compare, rangeKeys),
 			keyspan.InterleavingIterOpts{})
-		iter := newInvalidatingIter(interleavingIter)
-		iter.ignoreKind(InternalKeyKindRangeDelete)
+		iter := invalidating.NewIter(interleavingIter, invalidating.IgnoreKinds(InternalKeyKindRangeDelete))
 		if merge == nil {
 			merge = func(key, value []byte) (base.ValueMerger, error) {
 				m := &debugMerger{}
diff --git a/iterator_test.go b/iterator_test.go
index df8d05d5..4ce2f811 100644
--- a/iterator_test.go
+++ b/iterator_test.go
@@ -21,6 +21,7 @@ import (
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/pebble/internal/base"
 	"github.com/cockroachdb/pebble/internal/bytealloc"
+	"github.com/cockroachdb/pebble/internal/invalidating"
 	"github.com/cockroachdb/pebble/internal/keyspan"
 	"github.com/cockroachdb/pebble/internal/manifest"
 	"github.com/cockroachdb/pebble/internal/testkeys"
@@ -224,123 +225,6 @@ func (f *fakeIter) SetBounds(lower, upper []byte) {
 	f.upper = upper
 }
 
-// invalidatingIter tests unsafe key/value slice reuse by modifying the last
-// returned key/value to all 1s.
-type invalidatingIter struct {
-	iter        internalIterator
-	lastKey     *InternalKey
-	lastValue   []byte
-	ignoreKinds [base.InternalKeyKindMax + 1]bool
-	err         error
-}
-
-func newInvalidatingIter(iter internalIterator) *invalidatingIter {
-	return &invalidatingIter{iter: iter}
-}
-
-func (i *invalidatingIter) ignoreKind(kind base.InternalKeyKind) {
-	i.ignoreKinds[kind] = true
-}
-
-func (i *invalidatingIter) update(
-	key *InternalKey, value base.LazyValue,
-) (*InternalKey, base.LazyValue) {
-	i.zeroLast()
-
-	v, _, err := value.Value(nil)
-	if err != nil {
-		i.err = err
-		key = nil
-	}
-	if key == nil {
-		i.lastKey = nil
-		i.lastValue = nil
-		return nil, LazyValue{}
-	}
-
-	i.lastKey = &InternalKey{}
-	*i.lastKey = key.Clone()
-	i.lastValue = make([]byte, len(v))
-	copy(i.lastValue, v)
-	return i.lastKey, base.MakeInPlaceValue(i.lastValue)
-}
-
-func (i *invalidatingIter) zeroLast() {
-	if i.lastKey == nil {
-		return
-	}
-	if i.ignoreKinds[i.lastKey.Kind()] {
-		return
-	}
-
-	if i.lastKey != nil {
-		for j := range i.lastKey.UserKey {
-			i.lastKey.UserKey[j] = 0xff
-		}
-		i.lastKey.Trailer = 0
-	}
-	for j := range i.lastValue {
-		i.lastValue[j] = 0xff
-	}
-}
-
-func (i *invalidatingIter) SeekGE(
-	key []byte, flags base.SeekGEFlags,
-) (*InternalKey, base.LazyValue) {
-	return i.update(i.iter.SeekGE(key, flags))
-}
-
-func (i *invalidatingIter) SeekPrefixGE(
-	prefix, key []byte, flags base.SeekGEFlags,
-) (*base.InternalKey, base.LazyValue) {
-	return i.update(i.iter.SeekPrefixGE(prefix, key, flags))
-}
-
-func (i *invalidatingIter) SeekLT(
-	key []byte, flags base.SeekLTFlags,
-) (*InternalKey, base.LazyValue) {
-	return i.update(i.iter.SeekLT(key, flags))
-}
-
-func (i *invalidatingIter) First() (*InternalKey, base.LazyValue) {
-	return i.update(i.iter.First())
-}
-
-func (i *invalidatingIter) Last() (*InternalKey, base.LazyValue) {
-	return i.update(i.iter.Last())
-}
-
-func (i *invalidatingIter) Next() (*InternalKey, base.LazyValue) {
-	return i.update(i.iter.Next())
-}
-
-func (i *invalidatingIter) Prev() (*InternalKey, base.LazyValue) {
-	return i.update(i.iter.Prev())
-}
-
-func (i *invalidatingIter) NextPrefix(succKey []byte) (*InternalKey, base.LazyValue) {
-	return i.update(i.iter.NextPrefix(succKey))
-}
-
-func (i *invalidatingIter) Error() error {
-	if err := i.iter.Error(); err != nil {
-		return err
-	}
-	return i.err
-}
-
-func (i *invalidatingIter) Close() error {
-	return i.iter.Close()
-}
-
-func (i *invalidatingIter) SetBounds(lower, upper []byte) {
-	i.iter.SetBounds(lower, upper)
-}
-
-func (i *invalidatingIter) String() string {
-	return i.iter.String()
-}
-
 // testIterator tests creating a combined iterator from a number of sub-
 // iterators. newFunc is a constructor function. splitFunc returns a random
 // split of the testKeyValuePairs slice such that walking a combined iterator
@@ -394,7 +278,7 @@ func testIterator(
 	}
 	for _, tc := range testCases {
 		var b bytes.Buffer
-		iter := newInvalidatingIter(newFunc(tc.iters...))
+		iter := invalidating.NewIter(newFunc(tc.iters...))
 		for key, _ := iter.First(); key != nil; key, _ = iter.Next() {
 			fmt.Fprintf(&b, "<%s:%d>", key.UserKey, key.SeqNum())
 		}
@@ -418,7 +302,7 @@ func testIterator(
 		for i, split := range splits {
 			iters[i] = newFakeIterator(nil, split...)
 		}
-		iter := newInternalIterAdapter(newInvalidatingIter(newFunc(iters...)))
+		iter := newInternalIterAdapter(invalidating.NewIter(newFunc(iters...)))
 		iter.First()
 
 		j := 0
@@ -525,7 +409,7 @@ func TestIterator(t *testing.T) {
 		iter.snapshot = seqNum
 		// NB: This Iterator cannot be cloned since it is not constructed
 		// with a readState. It suffices for this test.
-		it.iter = newInvalidatingIter(iter)
+		it.iter = invalidating.NewIter(iter)
 		return it
 	}
 
@@ -1131,7 +1015,7 @@ func TestIteratorSeekOptErrors(t *testing.T) {
 			keys:  keys,
 			vals:  vals,
 		}
-		errorIter = errorSeekIter{internalIterator: newInvalidatingIter(iter)}
+		errorIter = errorSeekIter{internalIterator: invalidating.NewIter(iter)}
 		// NB: This Iterator cannot be cloned since it is not constructed
 		// with a readState. It suffices for this test.
 		return &Iterator{