spanset: assert that batches don't access store local and unreplicated RangeID local keys

This commit adds the following test-only assertions:
1) Generated batches don't touch store-local keys.
2) Generated batches don't touch unreplicated RangeID local keys.

We disable the check in exactly 3 locations we know that we currently
touch those keys.

Author: iskettaneh

pkg/keys/constants.go

@@ -163,6 +163,7 @@ var (
 	//
 	// LocalStorePrefix is the prefix identifying per-store data.
 	LocalStorePrefix = makeKey(LocalPrefix, roachpb.Key("s"))
+	LocalStoreMax    = roachpb.Key(LocalStorePrefix).PrefixEnd()
 	// localStoreClusterVersionSuffix stores the cluster-wide version
 	// information for this store, updated any time the operator
 	// updates the minimum cluster version.

pkg/keys/keys.go

@@ -308,6 +308,21 @@ func DecodeRangeIDKey(
 	return roachpb.RangeID(rangeInt), infix, suffix, b, nil
 }
 
+// DecodeRangeIDPrefix parses a local range ID prefix into range ID.
+func DecodeRangeIDPrefix(key roachpb.Key) (roachpb.RangeID, error) {
+	if !bytes.HasPrefix(key, LocalRangeIDPrefix) {
+		return 0, errors.Errorf("key %s does not have %s prefix", key, LocalRangeIDPrefix)
+	}
+	// Cut the prefix, the Range ID, and the infix specifier.
+	b := key[len(LocalRangeIDPrefix):]
+	_, rangeInt, err := encoding.DecodeUvarintAscending(b)
+	if err != nil {
+		return 0, err
+	}
+
+	return roachpb.RangeID(rangeInt), nil
+}
+
 // AbortSpanKey returns a range-local key by Range ID for an
 // AbortSpan entry, with detail specified by encoding the
 // supplied transaction ID.

pkg/kv/kvserver/batcheval/cmd_end_transaction.go

@@ -1340,8 +1340,10 @@ func splitTriggerHelper(
 	if err != nil {
 		return enginepb.MVCCStats{}, result.Result{}, errors.Wrap(err, "unable to fetch last replica GC timestamp")
 	}
+
 	if err := storage.MVCCPutProto(
-		ctx, batch, keys.RangeLastReplicaGCTimestampKey(split.RightDesc.RangeID), hlc.Timestamp{},
+		ctx, spanset.DisableForbiddenSpanAssertionsOnBatch(batch),
+		keys.RangeLastReplicaGCTimestampKey(split.RightDesc.RangeID), hlc.Timestamp{},
 		&replicaGCTS, storage.MVCCWriteOptions{Category: fs.BatchEvalReadCategory}); err != nil {
 		return enginepb.MVCCStats{}, result.Result{}, errors.Wrap(err, "unable to copy last replica GC timestamp")
 	}
@@ -1541,7 +1543,8 @@ func splitTriggerHelper(
 		// as all replicas will be responsible for writing it locally before
 		// applying the split.
 		if !rec.ClusterSettings().Version.IsActive(ctx, clusterversion.V25_4_WriteInitialTruncStateBeforeSplitApplication) {
-			if err := kvstorage.WriteInitialTruncState(ctx, batch, split.RightDesc.RangeID); err != nil {
+			if err := kvstorage.WriteInitialTruncState(ctx,
+				spanset.DisableForbiddenSpanAssertionsOnBatch(batch), split.RightDesc.RangeID); err != nil {
 				return enginepb.MVCCStats{}, result.Result{}, errors.Wrap(err, "unable to write initial Replica state")
 			}
 		}

pkg/kv/kvserver/batcheval/cmd_truncate_log.go

@@ -118,7 +118,9 @@ func TruncateLog(
 	// are not tracked in the raft log delta. The delta will be adjusted below
 	// raft.
 	// We can pass zero as nowNanos because we're only interested in SysBytes.
-	ms, err := storage.ComputeStats(ctx, readWriter, start, end, 0 /* nowNanos */)
+	// TODO(#157895): Use the log engine here instead of the state machine engine.
+	ms, err := storage.ComputeStats(ctx,
+		spanset.DisableForbiddenSpanAssertionsOnBatch(readWriter), start, end, 0 /* nowNanos */)
 	if err != nil {
 		return result.Result{}, errors.Wrap(err, "while computing stats of Raft log freed by truncation")
 	}

pkg/kv/kvserver/replica_test.go

@@ -15197,3 +15197,110 @@ func TestLeaderlessWatcherInit(t *testing.T) {
 		t.Fatalf("expected LeaderlessWatcher channel to be closed")
 	}
 }
+
+// TestOverlapsUnreplicatedRangeIDLocalKeys verifies that the function
+// overlapsUnreplicatedRangeIDLocalKeys() successfully catches any overlap with
+// unreplicated rangeID local keys.
+func TestOverlapsUnreplicatedRangeIDLocalKeys(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	testCases := []struct {
+		name      string
+		span      roachpb.Span
+		expectErr bool
+	}{
+		// Span before local RangeID.
+		{span: roachpb.Span{Key: roachpb.KeyMin, EndKey: keys.LocalRangeIDPrefix.AsRawKey()},
+			expectErr: false},
+		// Span after local RangeID.
+		{span: roachpb.Span{Key: keys.LocalRangePrefix, EndKey: roachpb.KeyMax},
+			expectErr: false},
+		// Span equals the full local RangeID span.
+		{span: roachpb.Span{Key: keys.LocalRangeIDPrefix.AsRawKey(),
+			EndKey: keys.LocalRangeIDPrefix.AsRawKey().PrefixEnd()},
+			expectErr: true},
+		// Span contains the full local RangeID span.
+		{span: roachpb.Span{Key: roachpb.KeyMin, EndKey: roachpb.KeyMax},
+			expectErr: true},
+		// Span partially overlaps the full RangeID span.
+		{span: roachpb.Span{Key: roachpb.KeyMin, EndKey: keys.RaftTruncatedStateKey(1)},
+			expectErr: true},
+		{span: roachpb.Span{Key: keys.RaftTruncatedStateKey(1), EndKey: roachpb.KeyMax},
+			expectErr: true},
+		// Span inside unreplicated local RangeID span.
+		{span: roachpb.Span{Key: keys.RangeTombstoneKey(1), EndKey: keys.RaftTruncatedStateKey(1)},
+			expectErr: true},
+		// Span inside replicated local RangeID span.
+		{span: roachpb.Span{Key: keys.RangeForceFlushKey(1), EndKey: keys.RangeLeaseKey(1)},
+			expectErr: false},
+		// Point span before local RangeID span.
+		{span: roachpb.Span{Key: keys.LocalRangeIDPrefix.AsRawKey().Prevish(1)}, expectErr: false},
+		// Point span after local RangeID span.
+		{span: roachpb.Span{Key: keys.LocalRangeIDPrefix.AsRawKey().PrefixEnd()}, expectErr: false},
+		// Point span inside local RangeID span.
+		{span: roachpb.Span{Key: keys.RaftTruncatedStateKey(1)}, expectErr: true},
+		// Point span with nil startKey before local RangeID span.
+		{span: roachpb.Span{EndKey: keys.LocalRangeIDPrefix.AsRawKey().Prevish(1)}, expectErr: false},
+		// Point span with nil startKey after local RangeID span.
+		{span: roachpb.Span{EndKey: keys.LocalRangeIDPrefix.AsRawKey().PrefixEnd().Next()}, expectErr: false},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			err := overlapsUnreplicatedRangeIDLocalKeys(spanset.TrickySpan(tc.span))
+			if tc.expectErr {
+				require.Errorf(t, err, "expected error for span %s", tc.span)
+			} else {
+				require.NoErrorf(t, err, "expected no error for span %s", tc.span)
+			}
+		})
+	}
+}
+
+// TestOverlapsStoreLocalKeys verifies that the function
+// overlapsStoreLocalKeys() successfully catches any overlap with
+// store local keys.
+func TestOverlapsStoreLocalKeys(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	testCases := []struct {
+		name      string
+		span      roachpb.Span
+		expectErr bool
+	}{
+		// Span before store local span.
+		{span: roachpb.Span{Key: roachpb.KeyMin, EndKey: keys.LocalStorePrefix}, expectErr: false},
+		// Span after store local span.
+		{span: roachpb.Span{Key: keys.LocalStoreMax, EndKey: roachpb.KeyMax}, expectErr: false},
+		// Span Contains store local span.
+		{span: roachpb.Span{Key: roachpb.KeyMin, EndKey: roachpb.KeyMax}, expectErr: true},
+		// Span partially overlaps with store local span.
+		{span: roachpb.Span{Key: roachpb.KeyMin, EndKey: keys.StoreIdentKey()}, expectErr: true},
+		{span: roachpb.Span{Key: keys.StoreIdentKey(), EndKey: roachpb.KeyMax}, expectErr: true},
+		// Span overlaps with store local keys.
+		{span: roachpb.Span{Key: keys.StoreGossipKey(), EndKey: keys.StoreIdentKey()}, expectErr: true},
+		// Point span before store local span.
+		{span: roachpb.Span{Key: roachpb.Key(keys.LocalStorePrefix).Prevish(1)}, expectErr: false},
+		// Point span after store local span.
+		{span: roachpb.Span{Key: keys.LocalStoreMax}, expectErr: false},
+		// Point span inside store local span.
+		{span: roachpb.Span{Key: keys.StoreIdentKey()}, expectErr: true},
+		// Point span with nil startKey before store local span.
+		{span: roachpb.Span{EndKey: roachpb.Key(keys.LocalStorePrefix).Prevish(1)}, expectErr: false},
+		// Point span with nil startKey after store local span.
+		{span: roachpb.Span{EndKey: keys.LocalStoreMax.Next()}, expectErr: false},
+		// Point span with nil startKey inside store local span.
+		{span: roachpb.Span{EndKey: keys.LocalStoreMax}, expectErr: true},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			err := overlapsStoreLocalKeys(spanset.TrickySpan(tc.span))
+			if tc.expectErr {
+				require.Errorf(t, err, "expected error for span %s", tc.span)
+			} else {
+				require.NoErrorf(t, err, "expected no error for span %s", tc.span)
+			}
+		})
+	}
+}

pkg/kv/kvserver/replica_write.go

@@ -10,6 +10,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/batcheval"
@@ -24,7 +25,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/storage"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
-	"github.com/cockroachdb/cockroach/pkg/util"
+	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/log/logcrash"
@@ -808,14 +809,18 @@ func (r *Replica) newBatchedEngine(g *concurrency.Guard) (storage.Batch, *storag
 		opLogger = storage.NewOpLoggerBatch(batch)
 		batch = opLogger
 	}
-	if util.RaceEnabled {
+	if buildutil.CrdbTestBuild {
 		// During writes we may encounter a versioned value newer than the request
 		// timestamp, and may have to retry at a higher timestamp. This is still
 		// safe as we're only ever writing at timestamps higher than the timestamp
 		// any write latch would be declared at. But because of this, we don't
 		// assert on access timestamps using spanset.NewBatchAt.
-		batch = spanset.NewBatch(batch, g.LatchSpans())
+		spans := g.LatchSpans()
+		spans.AddForbiddenMatcher(overlapsUnreplicatedRangeIDLocalKeys)
+		spans.AddForbiddenMatcher(overlapsStoreLocalKeys)
+		batch = spanset.NewBatch(batch, spans)
 	}
+
 	return batch, opLogger
 }
 
@@ -884,3 +889,76 @@ func releaseMVCCStats(ms *enginepb.MVCCStats) {
 	ms.Reset()
 	mvccStatsPool.Put(ms)
 }
+
+// overlapsUnreplicatedRangeIDLocalKeys checks if the provided span overlaps
+// with any unreplicated rangeID local keys.
+// Note that we could receive the span with a nil startKey, which has a special
+// meaning that the span represents: [endKey.Prev(), endKey).
+func overlapsUnreplicatedRangeIDLocalKeys(span spanset.TrickySpan) error {
+	// getRangeIDLocalSpan is a helper function that returns the full
+	// unreplicated RangeID span for a given RangeID.
+	getRangeIDLocalSpan := func(rangeID roachpb.RangeID) roachpb.Span {
+		return roachpb.Span{Key: keys.MakeRangeIDUnreplicatedPrefix(rangeID),
+			EndKey: keys.MakeRangeIDUnreplicatedPrefix(rangeID).PrefixEnd()}
+	}
+
+	fullRangeIDLocalSpans := roachpb.Span{
+		Key:    keys.LocalRangeIDPrefix.AsRawKey(),
+		EndKey: keys.LocalRangeIDPrefix.AsRawKey().PrefixEnd(),
+	}
+
+	// If the provided span is completely outside the rangeID local spans for any
+	// rangeID, then there is no overlap with any rangeID local keys.
+	if !spanset.Overlaps(fullRangeIDLocalSpans, span) {
+		return nil
+	}
+
+	// If we are partially overlapping the full rangeID local span. We must be
+	// overlapping rangeID local keys.
+	if spanset.PartiallyOverlaps(fullRangeIDLocalSpans, span) {
+		return errors.Errorf("partially overlapping the unreplicated rangeID span")
+	}
+
+	// At this point, we know that either we are containing the
+	// fullRangeIDLocalSpans, or we are completely within fullRangeIDLocalSpans.
+	// Check to see if we fully contain fullRangeIDLocalSpans.
+	if !spanset.Contains(fullRangeIDLocalSpans, span) {
+		return errors.Errorf("completely containing the unreplicated rangeID span")
+	}
+
+	// We make an assumption here that whoever if the span in inside
+	// fullRangeIDLocalSpans, we expect that (a) The start key is not nil, and
+	// (b) both start and end keys should be in the same rangeID.
+	if span.Key == nil {
+		return errors.Errorf("unexpected nil start key")
+	}
+
+	rangeID, err := keys.DecodeRangeIDPrefix(span.Key)
+	if err != nil {
+		return errors.NewAssertionErrorWithWrappedErrf(err,
+			"could not decode range ID for key: %s", span.Key)
+	}
+
+	if spanset.Overlaps(getRangeIDLocalSpan(rangeID), span) {
+		return errors.Errorf("overlapping an unreplicated rangeID span")
+	}
+
+	return nil
+}
+
+// overlapsStoreLocalKeys returns an error if the provided span overlaps
+// with any store local keys.
+// Note that we could receive the span with a nil startKey, which has a special
+// meaning that the span represents: [endKey.Prev(), endKey).
+func overlapsStoreLocalKeys(span spanset.TrickySpan) error {
+	localStoreSpan := roachpb.Span{
+		Key:    keys.LocalStorePrefix,
+		EndKey: keys.LocalStoreMax,
+	}
+
+	if spanset.Overlaps(localStoreSpan, span) {
+		return errors.Errorf("overlaps with store local keys")
+	}
+
+	return nil
+}