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.

Fixes: #156537

Release note: None

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) (rangeID roachpb.RangeID, err 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

@@ -1328,8 +1328,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")
 	}
@@ -1529,7 +1531,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,8 @@ 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 */)
+	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,212 @@ 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
+	}{
+		{
+			name: "point start key before unreplicated rangeID local span",
+			span: roachpb.Span{Key: keys.MakeRangeIDUnreplicatedPrefix(1).
+				Prevish(roachpb.PrevishKeyLength), EndKey: nil},
+			expectErr: false,
+		},
+		{
+			name:      "point start key within unreplicated rangeID local span",
+			span:      roachpb.Span{Key: keys.RaftTruncatedStateKey(1), EndKey: nil},
+			expectErr: true,
+		},
+		{
+			name: "point start key at unreplicated rangeID local span end",
+			span: roachpb.Span{Key: keys.MakeRangeIDUnreplicatedPrefix(1).
+				PrefixEnd(), EndKey: nil},
+			expectErr: false,
+		},
+		{
+			name:      "point end key before unreplicated rangeID local span",
+			span:      roachpb.Span{Key: nil, EndKey: keys.MakeRangeIDUnreplicatedPrefix(1)},
+			expectErr: false,
+		},
+		{
+			name:      "point end key within unreplicated rangeID local span",
+			span:      roachpb.Span{Key: nil, EndKey: keys.RaftTruncatedStateKey(1)},
+			expectErr: true,
+		},
+		{
+			name: "point end key at unreplicated rangeID local span end",
+			span: roachpb.Span{Key: nil, EndKey: keys.MakeRangeIDUnreplicatedPrefix(1).
+				PrefixEnd()},
+			expectErr: true,
+		},
+		{
+			name: "point end key after unreplicated rangeID local span",
+			span: roachpb.Span{Key: nil, EndKey: keys.MakeRangeIDUnreplicatedPrefix(1).
+				PrefixEnd().Next()},
+			expectErr: false,
+		},
+		{
+			name:      "span fully before unreplicated rangeID local span",
+			span:      roachpb.Span{Key: roachpb.KeyMin, EndKey: keys.MakeRangeIDUnreplicatedPrefix(0)},
+			expectErr: false,
+		},
+		{
+			name:      "span fully after unreplicated rangeID local span",
+			span:      roachpb.Span{Key: keys.LocalRangePrefix, EndKey: roachpb.KeyMax},
+			expectErr: false,
+		},
+		{
+			name:      "span fully overlaps unreplicated rangeID local span",
+			span:      roachpb.Span{Key: roachpb.KeyMin, EndKey: roachpb.KeyMax},
+			expectErr: true,
+		},
+		{
+			name:      "span partially overlaps with start of unreplicated rangeID local span",
+			span:      roachpb.Span{Key: roachpb.KeyMin, EndKey: keys.RaftTruncatedStateKey(1)},
+			expectErr: true,
+		},
+		{
+			name:      "span partially overlaps with end of unreplicated rangeID local span",
+			span:      roachpb.Span{Key: keys.RaftTruncatedStateKey(1), EndKey: roachpb.KeyMax},
+			expectErr: true,
+		},
+		{
+			name: "span within with unreplicated rangeID local span",
+			span: roachpb.Span{Key: keys.RaftTruncatedStateKey(1),
+				EndKey: keys.RangeLastReplicaGCTimestampKey(1)},
+			expectErr: true,
+		},
+		{
+			name: "span overlaps with multiple unreplicated range local span",
+			span: roachpb.Span{Key: keys.RaftTruncatedStateKey(1),
+				EndKey: keys.RaftTruncatedStateKey(3)},
+			expectErr: true,
+		},
+		{
+			name: "span from end of unreplicated RangeID, to next replicated RangeID",
+			span: roachpb.Span{Key: keys.MakeRangeIDUnreplicatedPrefix(1).PrefixEnd(),
+				EndKey: keys.RangeVersionKey(2)},
+			expectErr: false,
+		},
+		{
+			name: "span from end of unreplicated RangeID, to next unreplicated RangeID",
+			span: roachpb.Span{Key: keys.MakeRangeIDUnreplicatedPrefix(1).PrefixEnd(),
+				EndKey: keys.RangeLastReplicaGCTimestampKey(2)},
+			expectErr: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			err := overlapsUnreplicatedRangeIDLocalKeys(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
+	}{
+		{
+			name: "point start key before store local span",
+			span: roachpb.Span{Key: roachpb.Key(keys.LocalStorePrefix).
+				Prevish(roachpb.PrevishKeyLength), EndKey: nil},
+			expectErr: false,
+		},
+		{
+			name:      "point start key at store local span start",
+			span:      roachpb.Span{Key: keys.LocalStorePrefix, EndKey: nil},
+			expectErr: true,
+		},
+		{
+			name:      "point start key within store local span",
+			span:      roachpb.Span{Key: keys.StoreIdentKey(), EndKey: nil},
+			expectErr: true,
+		},
+		{
+			name:      "point start key at store local span end",
+			span:      roachpb.Span{Key: keys.LocalStoreMax, EndKey: nil},
+			expectErr: false,
+		},
+		{
+			name:      "point end key before store local span",
+			span:      roachpb.Span{Key: nil, EndKey: keys.LocalStorePrefix},
+			expectErr: false,
+		},
+		{
+			name:      "point end key at store local span start",
+			span:      roachpb.Span{Key: nil, EndKey: roachpb.Key(keys.LocalStorePrefix).Next()},
+			expectErr: true,
+		},
+		{
+			name:      "point end key within store local span",
+			span:      roachpb.Span{Key: nil, EndKey: keys.StoreIdentKey()},
+			expectErr: true,
+		},
+		{
+			name:      "point end key at store local span end",
+			span:      roachpb.Span{Key: nil, EndKey: keys.LocalStoreMax},
+			expectErr: true,
+		},
+		{
+			name:      "span fully before store local span",
+			span:      roachpb.Span{Key: roachpb.KeyMin, EndKey: keys.LocalStorePrefix},
+			expectErr: false,
+		},
+		{
+			name:      "span fully after store local span",
+			span:      roachpb.Span{Key: keys.LocalStoreMax, EndKey: roachpb.KeyMax},
+			expectErr: false,
+		},
+		{
+			name:      "span fully overlaps store local span",
+			span:      roachpb.Span{Key: roachpb.KeyMin, EndKey: roachpb.KeyMax},
+			expectErr: true,
+		},
+		{
+			name:      "span partially overlaps with start of store local span",
+			span:      roachpb.Span{Key: roachpb.KeyMin, EndKey: keys.StoreIdentKey()},
+			expectErr: true,
+		},
+		{
+			name:      "span partially overlaps with end of store local span",
+			span:      roachpb.Span{Key: keys.StoreIdentKey(), EndKey: roachpb.KeyMax},
+			expectErr: true,
+		},
+		{
+			name:      "span within store local span",
+			span:      roachpb.Span{Key: keys.StoreGossipKey(), EndKey: keys.StoreIdentKey()},
+			expectErr: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			err := overlapsStoreLocalKeys(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"
@@ -814,8 +815,12 @@ func (r *Replica) newBatchedEngine(g *concurrency.Guard) (storage.Batch, *storag
 		// 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,102 @@ func releaseMVCCStats(ms *enginepb.MVCCStats) {
 	ms.Reset()
 	mvccStatsPool.Put(ms)
 }
+
+// overlapsUnreplicatedRangeIDLocalKeys checks if the provided span overlaps
+// with any unreplicated rangeID local keys.
+func overlapsUnreplicatedRangeIDLocalKeys(span roachpb.Span) 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.MakeRangeIDUnreplicatedPrefix(0),
+		EndKey: keys.LocalRangePrefix,
+	}
+
+	// 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
+	}
+
+	// Overlapping the full rangeID local span. We must be overlapping rangeID
+	// local keys.
+	if spanset.Contains(span, fullRangeIDLocalSpans) {
+		return errors.Errorf("overlapping an unreplicated rangeID span")
+	}
+
+	// At this point, we know that there is an overlap. If our span is not fully
+	// inside fullRangeIDLocalSpans, then we must be partially overlapping with
+	// RangeID local spans.
+	if !spanset.Contains(fullRangeIDLocalSpans, span) {
+		return errors.Errorf("partially overlapping an unreplicated rangeID span")
+	}
+
+	// Our span is fully inside fullRangeIDLocalSpans.
+	if span.Key == nil {
+		if span.EndKey.Compare(fullRangeIDLocalSpans.EndKey) == 0 {
+			return errors.Errorf("overlapping an unreplicated rangeID span")
+		}
+
+		rangeID, err := keys.DecodeRangeIDPrefix(span.EndKey)
+		if err != nil {
+			panic(err) // cannot happen
+		}
+
+		if spanset.Overlaps(getRangeIDLocalSpan(rangeID), span) {
+			return errors.Errorf("overlapping an unreplicated rangeID span")
+		}
+
+		return nil
+	}
+
+	rangeID, err := keys.DecodeRangeIDPrefix(span.Key)
+	if err != nil {
+		panic(err) // cannot happen
+	}
+
+	if spanset.Overlaps(getRangeIDLocalSpan(rangeID), span) {
+		return errors.Errorf("overlapping an unreplicated rangeID span")
+	}
+
+	// One last case to check. If the start key is just at the end of some RangeID
+	// local key, we need to check if the span is overlapping some RangeID local
+	// keys that are ahead of startKey (on the next RangeID).
+	if span.EndKey != nil {
+		endRangeID, err := keys.DecodeRangeIDPrefix(span.EndKey)
+		if err != nil {
+			panic(err) // cannot happen
+		}
+
+		if endRangeID-rangeID > 1 {
+			// If we are crossing more than one rangeID, we must have overlapped with
+			// some unreplicated local RangeID keys.
+			return errors.Errorf("overlapping an unreplicated rangeID span")
+		}
+
+		if spanset.Overlaps(getRangeIDLocalSpan(endRangeID), 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.
+func overlapsStoreLocalKeys(span roachpb.Span) 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
+}