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,123 @@ 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)
+
+	s := func(start, end roachpb.Key) roachpb.Span {
+		return roachpb.Span{Key: start, EndKey: end}
+	}
+
+	testCases := []struct {
+		span  roachpb.Span
+		notOk bool
+	}{
+		// Full spans not overlapping with unreplicated local RangeID spans.
+		{span: s(roachpb.KeyMin, keys.LocalRangeIDPrefix.AsRawKey())},
+		{span: s(keys.RangeForceFlushKey(1), keys.RangeLeaseKey(1))},
+		{span: s(keys.LocalRangeIDPrefix.AsRawKey().PrefixEnd(), roachpb.KeyMax)},
+
+		// Full spans overlapping with unreplicated local RangeID spans.
+		{span: s(roachpb.KeyMin, keys.RaftTruncatedStateKey(1)), notOk: true},
+		{span: s(keys.LocalRangeIDPrefix.AsRawKey(), keys.LocalRangeIDPrefix.AsRawKey().PrefixEnd()),
+			notOk: true},
+		{span: s(keys.RaftTruncatedStateKey(1), keys.RaftTruncatedStateKey(2)), notOk: true},
+		{span: s(keys.RaftTruncatedStateKey(1), roachpb.KeyMax), notOk: true},
+
+		// Point spans not overlapping with unreplicated local RangeID spans.
+		{span: s(roachpb.KeyMin, nil)},
+		{span: s(keys.LocalRangeIDPrefix.AsRawKey().Prevish(1), nil)},
+		{span: s(keys.RangeForceFlushKey(1), nil)},
+		{span: s(keys.LocalRangeIDPrefix.AsRawKey().PrefixEnd(), nil)},
+		{span: s(roachpb.KeyMax, nil)},
+
+		// Point spans overlapping with unreplicated local RangeID spans.
+		{span: s(keys.RangeTombstoneKey(1), nil), notOk: true},
+		{span: s(keys.RaftTruncatedStateKey(1), nil), notOk: true},
+		{span: s(keys.RaftTruncatedStateKey(2), nil), notOk: true},
+
+		// Tricky spans not overlapping with unreplicated local RangeID spans.
+		{span: s(nil, keys.LocalRangeIDPrefix.AsRawKey())},
+		{span: s(nil, keys.RangeForceFlushKey(1))},
+		{span: s(nil, keys.LocalRangeIDPrefix.AsRawKey().PrefixEnd().Next())},
+
+		// Tricky spans overlapping with unreplicated local RangeID spans.
+		{span: s(nil, keys.RangeTombstoneKey(1).Next()), notOk: true},
+		{span: s(nil, keys.RaftTruncatedStateKey(1).Next()), notOk: true},
+		{span: s(nil, keys.RaftTruncatedStateKey(2).Next()), notOk: true},
+	}
+
+	for _, tc := range testCases {
+		t.Run("", func(t *testing.T) {
+			err := overlapsUnreplicatedRangeIDLocalKeys(spanset.TrickySpan(tc.span))
+			if tc.notOk {
+				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)
+
+	s := func(start, end roachpb.Key) roachpb.Span {
+		return roachpb.Span{Key: start, EndKey: end}
+	}
+
+	testCases := []struct {
+		span  roachpb.Span
+		notOK bool
+	}{
+		// Full spans not overlapping with Store-local span.
+		{span: s(roachpb.KeyMin, keys.LocalStorePrefix)},
+		{span: s(keys.LocalStoreMax, roachpb.KeyMax)},
+
+		// Full spans overlapping with Store-local span.
+		{span: s(roachpb.KeyMin, roachpb.Key(keys.LocalStorePrefix).Next()), notOK: true},
+		{span: s(keys.LocalStorePrefix, keys.LocalStoreMax), notOK: true},
+		{span: s(keys.StoreGossipKey(), keys.StoreIdentKey()), notOK: true},
+		{span: s(keys.LocalStoreMax.Prevish(1), roachpb.KeyMax), notOK: true},
+
+		// Point spans not overlapping with Store-local span.
+		{span: s(roachpb.KeyMin, nil)},
+		{span: s(roachpb.Key(keys.LocalStorePrefix).Prevish(1), nil)},
+		{span: s(keys.LocalStoreMax.Next(), nil)},
+		{span: s(roachpb.KeyMax, nil)},
+
+		// Point spans overlapping with Store-local span.
+		{span: s(keys.LocalStorePrefix, nil), notOK: true},
+		{span: s(keys.StoreGossipKey(), nil), notOK: true},
+		{span: s(keys.LocalStoreMax.Prevish(1), nil), notOK: true},
+
+		// Tricky spans with nil StartKey not overlapping with Store-local span.
+		{span: s(nil, keys.LocalStorePrefix)},
+		{span: s(nil, keys.LocalStoreMax.Next())},
+
+		// Tricky spans with nil StartKey overlapping with Store-local span.
+		{span: s(nil, roachpb.Key(keys.LocalStorePrefix).Next()), notOK: true},
+		{span: s(nil, keys.StoreGossipKey()), notOK: true},
+		{span: s(nil, keys.LocalStoreMax), notOK: true},
+	}
+
+	for _, tc := range testCases {
+		t.Run("", func(t *testing.T) {
+			err := overlapsStoreLocalKeys(spanset.TrickySpan(tc.span))
+			if tc.notOK {
+				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,66 @@ 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 {
+	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
+	}
+
+	// At this point, we know that we overlap with fullRangeIDLocalSpans. If we
+	// are not completely within fullRangeIDLocalSpans, return an error as we
+	// collide with at least one unreplicated RangeIDLocal key.
+	if !spanset.Contains(fullRangeIDLocalSpans, span) {
+		return errors.Errorf("overlapping an unreplicated rangeID key")
+	}
+
+	// If the span in inside fullRangeIDLocalSpans, we expect that both start and
+	// end keys should be in the same rangeID.
+	rangeIDKey := span.Key
+	if rangeIDKey == nil {
+		rangeIDKey = span.EndKey
+	}
+
+	rangeID, err := keys.DecodeRangeIDPrefix(rangeIDKey)
+	if err != nil {
+		return errors.NewAssertionErrorWithWrappedErrf(err,
+			"could not decode range ID for span: %s", span)
+	}
+
+	if spanset.Overlaps(roachpb.Span{
+		Key:    keys.MakeRangeIDUnreplicatedPrefix(rangeID),
+		EndKey: keys.MakeRangeIDUnreplicatedPrefix(rangeID).PrefixEnd(),
+	}, 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
+}

pkg/kv/kvserver/spanset/batch.go

@@ -923,6 +923,8 @@ func DisableLatchAssertions(rw storage.ReadWriter) storage.ReadWriter {
 // forbidden span assertions disabled. It does not modify the original batch.
 // The returned batch shares the same underlying storage.Batch but has its own
 // SpanSet wrapper with the forbidden span assertion disabled.
+// TODO(ibrahim): We eventually want to eliminate all the users of this
+// function.
 func DisableForbiddenSpanAssertionsOnBatch(rw storage.ReadWriter) storage.ReadWriter {
 	switch v := rw.(type) {
 	case *spanSetBatch: