storage: introduce test-build assertion iterator

Backport cockroachdb#90859, cockroachdb#96685 and cockroachdb#97222 to mangle the buffers returned by Pebble
iterators to ensure MVCC and Cockroach code respects the Pebble iterator memory
lifetimes.

The RangeBounds() and RangeKeys() validity assertions are omitted, since 22.2
allowed these methods to be invoked on an invalid iterator or an iterator
positioned at a point key.

Release note: None
Release justification: Non-production code changes

build/bazelutil/check.sh

@@ -61,6 +61,8 @@ pkg/cmd/roachtest/BUILD.bazel
 EXISTING_CRDB_TEST_BUILD_CONSTRAINTS="
 pkg/util/buildutil/crdb_test_off.go://go:build !crdb_test || crdb_test_off
 pkg/util/buildutil/crdb_test_on.go://go:build crdb_test && !crdb_test_off
+pkg/storage/pebbleiter/crdb_test_off.go://go:build !crdb_test || crdb_test_off
+pkg/storage/pebbleiter/crdb_test_on.go://go:build crdb_test && !crdb_test_off
 "
 
 if [ -z "${COCKROACH_BAZEL_CHECK_FAST:-}" ]; then

pkg/BUILD.bazel

@@ -505,6 +505,7 @@ ALL_TESTS = [
     "//pkg/storage/enginepb:enginepb_test",
     "//pkg/storage/fs:fs_test",
     "//pkg/storage/metamorphic:metamorphic_test",
+    "//pkg/storage/pebbleiter:pebbleiter_test",
     "//pkg/storage:storage_test",
     "//pkg/testutils/docker:docker_test",
     "//pkg/testutils/floatcmp:floatcmp_test",
@@ -1805,6 +1806,8 @@ GO_TARGETS = [
     "//pkg/storage/fs:fs_test",
     "//pkg/storage/metamorphic:metamorphic",
     "//pkg/storage/metamorphic:metamorphic_test",
+    "//pkg/storage/pebbleiter:pebbleiter",
+    "//pkg/storage/pebbleiter:pebbleiter_test",
     "//pkg/storage:storage",
     "//pkg/storage:storage_test",
     "//pkg/streaming:streaming",
@@ -2861,6 +2864,7 @@ GET_X_DATA_TARGETS = [
     "//pkg/storage/enginepb:get_x_data",
     "//pkg/storage/fs:get_x_data",
     "//pkg/storage/metamorphic:get_x_data",
+    "//pkg/storage/pebbleiter:get_x_data",
     "//pkg/streaming:get_x_data",
     "//pkg/testutils:get_x_data",
     "//pkg/testutils/buildutil:get_x_data",

pkg/kv/kvserver/spanset/BUILD.bazel

@@ -14,6 +14,7 @@ go_library(
         "//pkg/keys",
         "//pkg/roachpb",
         "//pkg/storage",
+        "//pkg/storage/pebbleiter",
         "//pkg/util/hlc",
         "//pkg/util/log",
         "//pkg/util/protoutil",

pkg/kv/kvserver/spanset/batch.go

@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/storage"
+	"github.com/cockroachdb/cockroach/pkg/storage/pebbleiter"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
@@ -392,7 +393,7 @@ func (i *EngineIterator) UnsafeRawEngineKey() []byte {
 }
 
 // GetRawIter is part of the storage.EngineIterator interface.
-func (i *EngineIterator) GetRawIter() *pebble.Iterator {
+func (i *EngineIterator) GetRawIter() pebbleiter.Iterator {
 	return i.i.GetRawIter()
 }
 

pkg/storage/BUILD.bazel

@@ -61,6 +61,7 @@ go_library(
         "//pkg/settings/cluster",
         "//pkg/storage/enginepb",
         "//pkg/storage/fs",
+        "//pkg/storage/pebbleiter",
         "//pkg/util",
         "//pkg/util/admission",
         "//pkg/util/bufalloc",

pkg/storage/disk_map.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/diskmap"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/storage/pebbleiter"
 	"github.com/cockroachdb/cockroach/pkg/util/encoding"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/errors"
@@ -52,7 +53,7 @@ type pebbleMapBatchWriter struct {
 // pebbleMapIterator iterates over the keys of a pebbleMap in sorted order.
 type pebbleMapIterator struct {
 	allowDuplicates bool
-	iter            *pebble.Iterator
+	iter            pebbleiter.Iterator
 	// makeKey is a function that transforms a key into a byte slice with a prefix
 	// used to SeekGE() the underlying iterator.
 	makeKey func(k []byte) []byte
@@ -114,9 +115,9 @@ func (r *pebbleMap) makeKeyWithSequence(k []byte) []byte {
 func (r *pebbleMap) NewIterator() diskmap.SortedDiskMapIterator {
 	return &pebbleMapIterator{
 		allowDuplicates: r.allowDuplicates,
-		iter: r.store.NewIter(&pebble.IterOptions{
+		iter: pebbleiter.MaybeWrap(r.store.NewIter(&pebble.IterOptions{
 			UpperBound: roachpb.Key(r.prefix).PrefixEnd(),
-		}),
+		})),
 		makeKey: r.makeKey,
 		prefix:  r.prefix,
 	}

pkg/storage/engine.go

@@ -21,6 +21,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
 	"github.com/cockroachdb/cockroach/pkg/storage/fs"
+	"github.com/cockroachdb/cockroach/pkg/storage/pebbleiter"
 	"github.com/cockroachdb/cockroach/pkg/util/envutil"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/iterutil"
@@ -340,7 +341,7 @@ type EngineIterator interface {
 	Value() []byte
 	// GetRawIter is a low-level method only for use in the storage package,
 	// that returns the underlying pebble Iterator.
-	GetRawIter() *pebble.Iterator
+	GetRawIter() pebbleiter.Iterator
 	// SeekEngineKeyGEWithLimit is similar to SeekEngineKeyGE, but takes an
 	// additional exclusive upper limit parameter. The limit is semantically
 	// best-effort, and is an optimization to avoid O(n^2) iteration behavior in
@@ -1551,17 +1552,7 @@ func assertSimpleMVCCIteratorInvariants(iter SimpleMVCCIterator) error {
 					rangeKey, value.Value.RawBytes)
 			}
 		}
-
-	} else {
-		// Bounds and range keys must be empty.
-		if bounds := iter.RangeBounds(); !bounds.Equal(roachpb.Span{}) {
-			return errors.AssertionFailedf("hasRange=false but RangeBounds=%s", bounds)
-		}
-		if r := iter.RangeKeys(); !r.IsEmpty() || !r.Bounds.Equal(roachpb.Span{}) {
-			return errors.AssertionFailedf("hasRange=false but RangeKeys=%s", r)
-		}
 	}
-
 	return nil
 }
 

pkg/storage/engine_test.go

@@ -2171,10 +2171,6 @@ func TestEngineRangeKeysUnsupported(t *testing.T) {
 					hasPoint, hasRange := iter.HasPointAndRange()
 					require.True(t, hasPoint)
 					require.False(t, hasRange)
-					rangeBounds, err := iter.EngineRangeBounds()
-					require.NoError(t, err)
-					require.Empty(t, rangeBounds)
-					require.Empty(t, iter.EngineRangeKeys())
 
 					// Exhaust the iterator.
 					ok, err = iter.NextEngineKey()

pkg/storage/mvcc_history_metamorphic_iterator_test.go

@@ -172,7 +172,7 @@ func (m *metamorphicIterator) moveAround() {
 	}
 
 	hasPoint, _ := m.it.HasPointAndRange()
-	rangeKeys := m.it.RangeKeys().Clone()
+	rangeKeys := rangeKeysIfExist(m.it).Clone()
 	var rangeKeysIgnoringTime storage.MVCCRangeKeyStack
 	if iit != nil {
 		rangeKeysIgnoringTime = iit.RangeKeysIgnoringTime()
@@ -224,7 +224,7 @@ func (m *metamorphicIterator) moveAround() {
 					if m.it.UnsafeKey().Equal(cur) {
 						break // made it
 					}
-					printfln("step: %s %s [changed=%t]", m.it.UnsafeKey(), m.it.RangeKeys(), m.it.RangeKeyChanged())
+					printfln("step: %s %s [changed=%t]", m.it.UnsafeKey(), rangeKeysIfExist(m.it), m.it.RangeKeyChanged())
 					if iit != nil {
 						// If we're an incremental iterator with time bounds, and `cur` is not within bounds,
 						// would miss it if we used Next. So call NextIgnoringTime unconditionally.
@@ -247,7 +247,7 @@ func (m *metamorphicIterator) moveAround() {
 					valid, err := m.it.Valid()
 					require.Nil(m.t, err)
 					require.True(m.t, valid, "unable to recover original position following SeekLT")
-					printfln("rev-step: %s %s [changed=%t]", m.it.UnsafeKey(), m.it.RangeKeys(), m.it.RangeKeyChanged())
+					printfln("rev-step: %s %s [changed=%t]", m.it.UnsafeKey(), rangeKeysIfExist(m.it), m.it.RangeKeyChanged())
 					if m.it.UnsafeKey().Equal(cur) {
 						printfln("done")
 						break // made it
@@ -268,13 +268,13 @@ func (m *metamorphicIterator) moveAround() {
 			rangeKeysIgnoringTime2 = iit.RangeKeysIgnoringTime()
 		}
 		printfln("recovered position: %s hasPoint=%t, rangeKeys=%s, rangeKeysIgnoringTime=%s",
-			m.it.UnsafeKey(), hasPoint2, m.it.RangeKeys(), rangeKeysIgnoringTime2)
+			m.it.UnsafeKey(), hasPoint2, rangeKeysIfExist(m.it), rangeKeysIgnoringTime2)
 	}
 	// Back where we started and hopefully in an indistinguishable state.
 	// When the stack is empty, sometimes it's a nil slice and sometimes zero
 	// slice. A similar problem exists with MVCCRangeKeyVersion.Value. Sidestep
 	// them by comparing strings.
-	require.Equal(m.t, fmt.Sprint(rangeKeys), fmt.Sprint(m.it.RangeKeys()))
+	require.Equal(m.t, fmt.Sprint(rangeKeys), fmt.Sprint(rangeKeysIfExist(m.it)))
 	if iit != nil {
 		require.Equal(m.t, fmt.Sprint(rangeKeysIgnoringTime), fmt.Sprint(iit.RangeKeysIgnoringTime()))
 	}

pkg/storage/mvcc_history_test.go

@@ -1868,9 +1868,18 @@ func printIter(e *evalCtx) {
 	}
 }
 
+func rangeKeysIfExist(it storage.SimpleMVCCIterator) storage.MVCCRangeKeyStack {
+	if valid, err := it.Valid(); !valid || err != nil {
+		return storage.MVCCRangeKeyStack{}
+	} else if _, hasRange := it.HasPointAndRange(); !hasRange {
+		return storage.MVCCRangeKeyStack{}
+	}
+	return it.RangeKeys()
+}
+
 func checkAndUpdateRangeKeyChanged(e *evalCtx) bool {
 	rangeKeyChanged := e.iter.RangeKeyChanged()
-	rangeKeys := e.iter.RangeKeys()
+	rangeKeys := rangeKeysIfExist(e.iter)
 
 	if incrIter := e.tryMVCCIncrementalIter(); incrIter != nil {
 		// For MVCCIncrementalIterator, make sure RangeKeyChangedIgnoringTime() fires

pkg/storage/pebble.go

@@ -34,6 +34,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
 	"github.com/cockroachdb/cockroach/pkg/storage/fs"
+	"github.com/cockroachdb/cockroach/pkg/storage/pebbleiter"
 	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/cockroach/pkg/util/envutil"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
@@ -1912,7 +1913,7 @@ type pebbleReadOnly struct {
 	prefixEngineIter pebbleIterator
 	normalEngineIter pebbleIterator
 
-	iter       *pebble.Iterator
+	iter       pebbleiter.Iterator
 	iterUsed   bool // avoids cloning after PinEngineStateForIterators()
 	durability DurabilityRequirement
 	closed     bool
@@ -2086,7 +2087,7 @@ func (p *pebbleReadOnly) PinEngineStateForIterators() error {
 		if p.durability == GuaranteedDurability {
 			o = &pebble.IterOptions{OnlyReadGuaranteedDurable: true}
 		}
-		p.iter = p.parent.db.NewIter(o)
+		p.iter = pebbleiter.MaybeWrap(p.parent.db.NewIter(o))
 		// NB: p.iterUsed == false avoids cloning this in NewMVCCIterator(), since
 		// we've just created it.
 	}

pkg/storage/pebble_batch.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
+	"github.com/cockroachdb/cockroach/pkg/storage/pebbleiter"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/pebble"
@@ -45,7 +46,7 @@ type pebbleBatch struct {
 	prefixEngineIter pebbleIterator
 	normalEngineIter pebbleIterator
 
-	iter      *pebble.Iterator
+	iter      pebbleiter.Iterator
 	iterUsed  bool // avoids cloning after PinEngineStateForIterators()
 	writeOnly bool
 	closed    bool
@@ -243,9 +244,9 @@ func (p *pebbleBatch) SupportsRangeKeys() bool {
 func (p *pebbleBatch) PinEngineStateForIterators() error {
 	if p.iter == nil {
 		if p.batch.Indexed() {
-			p.iter = p.batch.NewIter(nil)
+			p.iter = pebbleiter.MaybeWrap(p.batch.NewIter(nil))
 		} else {
-			p.iter = p.db.NewIter(nil)
+			p.iter = pebbleiter.MaybeWrap(p.db.NewIter(nil))
 		}
 		// NB: p.iterUsed == false avoids cloning this in NewMVCCIterator(). We've
 		// just created it, so cloning it would just be overhead.

pkg/storage/pebble_iterator.go

@@ -17,6 +17,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/storage/pebbleiter"
 	"github.com/cockroachdb/cockroach/pkg/util"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
@@ -31,7 +32,7 @@ import (
 // should only be used in one of the two modes.
 type pebbleIterator struct {
 	// Underlying iterator for the DB.
-	iter    *pebble.Iterator
+	iter    pebbleiter.Iterator
 	options pebble.IterOptions
 	// Reusable buffer for MVCCKey or EngineKey encoding.
 	keyBuf []byte
@@ -91,27 +92,30 @@ func newPebbleIterator(
 	p := pebbleIterPool.Get().(*pebbleIterator)
 	p.reusable = false // defensive
 	p.init(nil, opts, durability, supportsRangeKeys)
-	p.iter = handle.NewIter(&p.options)
+	p.iter = pebbleiter.MaybeWrap(handle.NewIter(&p.options))
 	return p
 }
 
 // newPebbleIteratorByCloning creates a new Pebble iterator by cloning the given
 // iterator and reconfiguring it.
 func newPebbleIteratorByCloning(
-	iter *pebble.Iterator, opts IterOptions, durability DurabilityRequirement, supportsRangeKeys bool,
+	iter pebbleiter.Iterator,
+	opts IterOptions,
+	durability DurabilityRequirement,
+	supportsRangeKeys bool,
 ) *pebbleIterator {
-	var err error
 	p := pebbleIterPool.Get().(*pebbleIterator)
 	p.reusable = false // defensive
 	p.init(nil, opts, durability, supportsRangeKeys)
-	p.iter, err = iter.Clone(pebble.CloneOptions{
+	iter, err := iter.Clone(pebble.CloneOptions{
 		IterOptions:      &p.options,
 		RefreshBatchView: true,
 	})
 	if err != nil {
 		p.Close()
 		panic(err)
 	}
+	p.iter = iter
 	return p
 }
 
@@ -128,11 +132,12 @@ func newPebbleSSTIterator(
 		externalIterOpts = append(externalIterOpts, pebble.ExternalIterForwardOnly{})
 	}
 
-	var err error
-	if p.iter, err = pebble.NewExternalIter(DefaultPebbleOptions(), &p.options, files, externalIterOpts...); err != nil {
+	iter, err := pebble.NewExternalIter(DefaultPebbleOptions(), &p.options, files, externalIterOpts...)
+	if err != nil {
 		p.Close()
 		return nil, err
 	}
+	p.iter = pebbleiter.MaybeWrap(iter)
 	p.external = true
 	return p, nil
 }
@@ -141,7 +146,10 @@ func newPebbleSSTIterator(
 // reconfiguring the given iter. It is valid to pass a nil iter and then create
 // p.iter using p.options, to avoid redundant reconfiguration via SetOptions().
 func (p *pebbleIterator) init(
-	iter *pebble.Iterator, opts IterOptions, durability DurabilityRequirement, supportsRangeKeys bool,
+	iter pebbleiter.Iterator,
+	opts IterOptions,
+	durability DurabilityRequirement,
+	supportsRangeKeys bool,
 ) {
 	*p = pebbleIterator{
 		iter:               iter,
@@ -164,7 +172,7 @@ func (p *pebbleIterator) init(
 // 3. iter == nil: create a new iterator from handle.
 func (p *pebbleIterator) initReuseOrCreate(
 	handle pebble.Reader,
-	iter *pebble.Iterator,
+	iter pebbleiter.Iterator,
 	clone bool,
 	opts IterOptions,
 	durability DurabilityRequirement,
@@ -177,7 +185,7 @@ func (p *pebbleIterator) initReuseOrCreate(
 
 	p.init(nil, opts, durability, supportsRangeKeys)
 	if iter == nil {
-		p.iter = handle.NewIter(&p.options)
+		p.iter = pebbleiter.MaybeWrap(handle.NewIter(&p.options))
 	} else if clone {
 		var err error
 		p.iter, err = iter.Clone(pebble.CloneOptions{
@@ -917,7 +925,7 @@ func (p *pebbleIterator) IsPrefix() bool {
 }
 
 // GetRawIter is part of the EngineIterator interface.
-func (p *pebbleIterator) GetRawIter() *pebble.Iterator {
+func (p *pebbleIterator) GetRawIter() pebbleiter.Iterator {
 	return p.iter
 }