From 529d256ae16a2fb699d78040c3cfd64c62123cbd Mon Sep 17 00:00:00 2001
From: Jackson Owens <jackson@cockroachlabs.com>
Date: Wed, 13 Sep 2023 17:21:16 -0400
Subject: [PATCH] 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.
---
 compaction.go                 |  2 ++
 db.go                         |  6 ++++--
 internal/invalidating/iter.go | 17 ++++++++++++++++-
 3 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/compaction.go b/compaction.go
index 945b8bf150..c3b42f53b1 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 381163842a..8794bc162b 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,8 +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++
@@ -1453,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
index 056384ef74..48909ec936 100644
--- a/internal/invalidating/iter.go
+++ b/internal/invalidating/iter.go
@@ -4,7 +4,22 @@
 
 package invalidating
 
-import "github.com/cockroachdb/pebble/internal/base"
+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.