storage: unsafe key mangling for MVCCIterator

It can uncover bugs in code that misuses unsafe
keys. It uncovered the bug fixed in
https://github.com/cockroachdb/cockroach/pull/60460/files#diff-84108c53fd1e766ef8802b87b394981d3497d87c40d86e084f2ed77bba0ca71a

Release note: None

pkg/storage/intent_interleaving_iter.go

@@ -13,6 +13,7 @@ package storage
 import (
 	"bytes"
 	"fmt"
+	"math/rand"
 	"sync"
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
@@ -865,3 +866,72 @@ func newMVCCIteratorByCloningEngineIter(iter EngineIterator, opts IterOptions) M
 	}
 	return it
 }
+
+// unsageMVCCIterator is used in RaceEnabled test builds to randomly inject
+// changes to unsafe keys retrieved from MVCCIterators.
+type unsafeMVCCIterator struct {
+	MVCCIterator
+	keyBuf        []byte
+	rawKeyBuf     []byte
+	rawMVCCKeyBuf []byte
+}
+
+func wrapInUnsafeIter(iter MVCCIterator) MVCCIterator {
+	return &unsafeMVCCIterator{MVCCIterator: iter}
+}
+
+var _ MVCCIterator = &unsafeMVCCIterator{}
+
+func (i *unsafeMVCCIterator) SeekGE(key MVCCKey) {
+	i.mangleBufs()
+	i.MVCCIterator.SeekGE(key)
+}
+
+func (i *unsafeMVCCIterator) Next() {
+	i.mangleBufs()
+	i.MVCCIterator.Next()
+}
+
+func (i *unsafeMVCCIterator) NextKey() {
+	i.mangleBufs()
+	i.MVCCIterator.NextKey()
+}
+
+func (i *unsafeMVCCIterator) SeekLT(key MVCCKey) {
+	i.mangleBufs()
+	i.MVCCIterator.SeekLT(key)
+}
+
+func (i *unsafeMVCCIterator) Prev() {
+	i.mangleBufs()
+	i.MVCCIterator.Prev()
+}
+
+func (i *unsafeMVCCIterator) UnsafeKey() MVCCKey {
+	rv := i.MVCCIterator.UnsafeKey()
+	i.keyBuf = append(i.keyBuf[:0], rv.Key...)
+	rv.Key = i.keyBuf
+	return rv
+}
+
+func (i *unsafeMVCCIterator) UnsafeRawKey() []byte {
+	rv := i.MVCCIterator.UnsafeRawKey()
+	i.rawKeyBuf = append(i.rawKeyBuf[:0], rv...)
+	return i.rawKeyBuf
+}
+
+func (i *unsafeMVCCIterator) UnsafeRawMVCCKey() []byte {
+	rv := i.MVCCIterator.UnsafeRawMVCCKey()
+	i.rawMVCCKeyBuf = append(i.rawMVCCKeyBuf[:0], rv...)
+	return i.rawMVCCKeyBuf
+}
+
+func (i *unsafeMVCCIterator) mangleBufs() {
+	if rand.Intn(2) == 0 {
+		for _, b := range [3][]byte{i.keyBuf, i.rawKeyBuf, i.rawMVCCKeyBuf} {
+			for i := range b {
+				b[i] = 0
+			}
+		}
+	}
+}

pkg/storage/intent_interleaving_iter_test.go

@@ -312,7 +312,7 @@ func TestIntentInterleavingIter(t *testing.T) {
 				if d.HasArg("prefix") {
 					d.ScanArgs(t, "prefix", &opts.Prefix)
 				}
-				iter := newIntentInterleavingIterator(eng, opts)
+				iter := wrapInUnsafeIter(newIntentInterleavingIterator(eng, opts))
 				var b strings.Builder
 				defer iter.Close()
 				// pos is the original <file>:<lineno> prefix computed by

pkg/storage/pebble.go

@@ -29,6 +29,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/util"
 	"github.com/cockroachdb/cockroach/pkg/util/envutil"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -720,12 +721,18 @@ func (p *Pebble) NewMVCCIterator(iterKind MVCCIterKind, opts IterOptions) MVCCIt
 		// Doing defer r.Free() does not inline.
 		iter := r.NewMVCCIterator(iterKind, opts)
 		r.Free()
+		if util.RaceEnabled {
+			iter = wrapInUnsafeIter(iter)
+		}
 		return iter
 	}
-	iter := newPebbleIterator(p.db, nil, opts)
+	iter := MVCCIterator(newPebbleIterator(p.db, nil, opts))
 	if iter == nil {
 		panic("couldn't create a new iterator")
 	}
+	if util.RaceEnabled {
+		iter = wrapInUnsafeIter(iter)
+	}
 	return iter
 }
 
@@ -1339,12 +1346,19 @@ func (p *pebbleReadOnly) NewMVCCIterator(iterKind MVCCIterKind, opts IterOptions
 		// Doing defer r.Free() does not inline.
 		iter := r.NewMVCCIterator(iterKind, opts)
 		r.Free()
+		if util.RaceEnabled {
+			iter = wrapInUnsafeIter(iter)
+		}
 		return iter
 	}
 
 	if !opts.MinTimestampHint.IsEmpty() {
 		// MVCCIterators that specify timestamp bounds cannot be cached.
-		return newPebbleIterator(p.parent.db, nil, opts)
+		iter := MVCCIterator(newPebbleIterator(p.parent.db, nil, opts))
+		if util.RaceEnabled {
+			iter = wrapInUnsafeIter(iter)
+		}
+		return iter
 	}
 
 	iter := &p.normalIter
@@ -1369,7 +1383,11 @@ func (p *pebbleReadOnly) NewMVCCIterator(iterKind MVCCIterKind, opts IterOptions
 	}
 
 	iter.inuse = true
-	return iter
+	var rv MVCCIterator = iter
+	if util.RaceEnabled {
+		rv = wrapInUnsafeIter(rv)
+	}
+	return rv
 }
 
 // NewEngineIterator implements the Engine interface.
@@ -1582,9 +1600,16 @@ func (p *pebbleSnapshot) NewMVCCIterator(iterKind MVCCIterKind, opts IterOptions
 		// Doing defer r.Free() does not inline.
 		iter := r.NewMVCCIterator(iterKind, opts)
 		r.Free()
+		if util.RaceEnabled {
+			iter = wrapInUnsafeIter(iter)
+		}
 		return iter
 	}
-	return newPebbleIterator(p.snapshot, nil, opts)
+	iter := MVCCIterator(newPebbleIterator(p.snapshot, nil, opts))
+	if util.RaceEnabled {
+		iter = wrapInUnsafeIter(iter)
+	}
+	return iter
 }
 
 // NewEngineIterator implements the Reader interface.

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/util"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
@@ -206,12 +207,19 @@ func (p *pebbleBatch) NewMVCCIterator(iterKind MVCCIterKind, opts IterOptions) M
 		// Doing defer r.Free() does not inline.
 		iter := r.NewMVCCIterator(iterKind, opts)
 		r.Free()
+		if util.RaceEnabled {
+			iter = wrapInUnsafeIter(iter)
+		}
 		return iter
 	}
 
 	if !opts.MinTimestampHint.IsEmpty() {
 		// MVCCIterators that specify timestamp bounds cannot be cached.
-		return newPebbleIterator(p.batch, nil, opts)
+		iter := MVCCIterator(newPebbleIterator(p.batch, nil, opts))
+		if util.RaceEnabled {
+			iter = wrapInUnsafeIter(iter)
+		}
+		return iter
 	}
 
 	iter := &p.normalIter
@@ -239,7 +247,11 @@ func (p *pebbleBatch) NewMVCCIterator(iterKind MVCCIterKind, opts IterOptions) M
 	}
 
 	iter.inuse = true
-	return iter
+	var rv MVCCIterator = iter
+	if util.RaceEnabled {
+		rv = wrapInUnsafeIter(rv)
+	}
+	return rv
 }
 
 // NewEngineIterator implements the Batch interface.