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

@@ -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,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,224 @@ 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,
+		},
+		{
+			name: "point end key at keys.MakeRangeIDUnreplicatedPrefix(0) with nil start",
+			span: roachpb.Span{Key: nil,
+				EndKey: keys.MakeRangeIDUnreplicatedPrefix(0)},
+			expectErr: false,
+		},
+		{
+			name: "point end key at keys.LocalRangePrefix with nil start",
+			span: roachpb.Span{Key: nil,
+				EndKey: keys.LocalRangePrefix},
+			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)
+			}
+		})
+	}
+}