release-22.2: storage: introduce test-build assertion iterator

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"
@@ -102,21 +103,37 @@ func (i *MVCCIterator) SeekLT(key storage.MVCCKey) {
 // Next is part of the storage.MVCCIterator interface.
 func (i *MVCCIterator) Next() {
 	i.i.Next()
-	i.checkAllowed(roachpb.Span{Key: i.UnsafeKey().Key}, false)
+	i.checkAllowedCurrPosForward(false)
 }
 
 // Prev is part of the storage.MVCCIterator interface.
 func (i *MVCCIterator) Prev() {
 	i.i.Prev()
-	i.checkAllowed(roachpb.Span{Key: i.UnsafeKey().Key}, false)
+	i.checkAllowedCurrPosForward(false)
 }
 
 // NextKey is part of the storage.MVCCIterator interface.
 func (i *MVCCIterator) NextKey() {
 	i.i.NextKey()
-	i.checkAllowed(roachpb.Span{Key: i.UnsafeKey().Key}, false)
+	i.checkAllowedCurrPosForward(false)
 }
 
+// checkAllowedCurrPosForward checks the span starting at the current iterator
+// position, if the current iterator position is valid.
+func (i *MVCCIterator) checkAllowedCurrPosForward(errIfDisallowed bool) {
+	i.invalid = false
+	i.err = nil
+	if ok, _ := i.i.Valid(); !ok {
+		// If the iterator is invalid after the operation, there's nothing to
+		// check. We allow uses of iterators to exceed the declared span bounds
+		// as long as the iterator itself is configured with proper boundaries.
+		return
+	}
+	i.checkAllowedValidPos(roachpb.Span{Key: i.UnsafeKey().Key}, errIfDisallowed)
+}
+
+// checkAllowed checks the provided span if the current iterator position is
+// valid.
 func (i *MVCCIterator) checkAllowed(span roachpb.Span, errIfDisallowed bool) {
 	i.invalid = false
 	i.err = nil
@@ -126,6 +143,10 @@ func (i *MVCCIterator) checkAllowed(span roachpb.Span, errIfDisallowed bool) {
 		// as long as the iterator itself is configured with proper boundaries.
 		return
 	}
+	i.checkAllowedValidPos(span, errIfDisallowed)
+}
+
+func (i *MVCCIterator) checkAllowedValidPos(span roachpb.Span, errIfDisallowed bool) {
 	var err error
 	if i.spansOnly {
 		err = i.spans.CheckAllowed(SpanReadOnly, span)
@@ -392,7 +413,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.