kvstorage: assert each engine only access its spans

pkg/kv/kvserver/kvstorage/BUILD.bazel

@@ -21,11 +21,13 @@ go_library(
         "//pkg/kv/kvserver/kvserverpb",
         "//pkg/kv/kvserver/logstore",
         "//pkg/kv/kvserver/rditer",
+        "//pkg/kv/kvserver/spanset",
         "//pkg/raft/raftpb",
         "//pkg/roachpb",
         "//pkg/storage",
         "//pkg/storage/enginepb",
         "//pkg/storage/fs",
+        "//pkg/util",
         "//pkg/util/buildutil",
         "//pkg/util/hlc",
         "//pkg/util/iterutil",
@@ -45,6 +47,7 @@ go_test(
         "init_test.go",
         "initial_test.go",
         "stateloader_test.go",
+        "storage_test.go",
     ],
     data = glob(["testdata/**"]),
     embed = [":kvstorage"],
@@ -56,6 +59,7 @@ go_test(
         "//pkg/kv/kvserver/kvserverpb",
         "//pkg/kv/kvserver/logstore",
         "//pkg/kv/kvserver/print",
+        "//pkg/kv/kvserver/spanset",
         "//pkg/raft/raftpb",
         "//pkg/roachpb",
         "//pkg/settings/cluster",

pkg/kv/kvserver/kvstorage/destroy_test.go

@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/logstore"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/print"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/spanset"
 	"github.com/cockroachdb/cockroach/pkg/raft/raftpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/storage"
@@ -139,7 +140,7 @@ func (r *replicaInfo) createStateMachine(ctx context.Context, t *testing.T, w st
 	// TODO(pav-kv): figure out whether LastReplicaGCTimestamp should be in the
 	// log or state engine.
 	require.NoError(t, storage.MVCCBlindPutProto(
-		ctx, w,
+		ctx, spanset.DisableWriterAssertions(w),
 		keys.RangeLastReplicaGCTimestampKey(r.id.RangeID),
 		hlc.Timestamp{}, /* timestamp */
 		&hlc.Timestamp{WallTime: 12345678},

pkg/kv/kvserver/kvstorage/storage.go

@@ -6,8 +6,13 @@
 package kvstorage
 
 import (
+	"github.com/cockroachdb/cockroach/pkg/keys"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/spanset"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/storage"
+	"github.com/cockroachdb/cockroach/pkg/util"
 	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
+	"github.com/cockroachdb/errors"
 )
 
 // The following are type aliases that help annotating various storage
@@ -121,8 +126,15 @@ type Engines struct {
 // MakeEngines creates an Engines handle in which both state machine and log
 // engine reside in the same physical engine.
 func MakeEngines(eng storage.Engine) Engines {
-	// TODO(#158281): in test builds, wrap the engines in a way that allows
-	// verifying that all accesses touch the correct subset of keys.
+	if util.RaceEnabled {
+		// Wrap the engines with span set engines to catch incorrect engine
+		// accesses.
+		return Engines{
+			stateEngine: spanset.NewEngine(eng, validateIsStateEngineSpan),
+			logEngine:   spanset.NewEngine(eng, validateIsRaftEngineSpan),
+			todoEngine:  eng,
+		}
+	}
 	return Engines{
 		stateEngine: eng,
 		todoEngine:  eng,
@@ -137,6 +149,16 @@ func MakeSeparatedEnginesForTesting(state, log storage.Engine) Engines {
 	if !buildutil.CrdbTestBuild {
 		panic("separated engines are not supported")
 	}
+	if util.RaceEnabled {
+		// Wrap the engines with span set engines to catch incorrect engine
+		// accesses.
+		return Engines{
+			stateEngine: spanset.NewEngine(state, validateIsStateEngineSpan),
+			todoEngine:  nil,
+			logEngine:   spanset.NewEngine(log, validateIsRaftEngineSpan),
+			separated:   true,
+		}
+	}
 	return Engines{
 		stateEngine: state,
 		todoEngine:  nil,
@@ -178,3 +200,128 @@ func (e *Engines) TODOEngine() storage.Engine {
 func (e *Engines) Separated() bool {
 	return e.separated
 }
+
+// validateIsStateEngineSpan asserts that the provided span only overlaps with
+// keys in the State engine and returns an error if not.
+// Note that we could receive the span with a nil startKey, which has a special
+// meaning that the span represents: [endKey.Prev(), endKey).
+func validateIsStateEngineSpan(span spanset.TrickySpan) error {
+	// If the provided span overlaps with local store span, it cannot be a
+	// StateEngine span because Store-local keys belong to the LogEngine.
+	if spanset.Overlaps(roachpb.Span{
+		Key:    keys.LocalStorePrefix,
+		EndKey: keys.LocalStoreMax,
+	}, span) {
+		return errors.Errorf("overlaps with store local keys")
+	}
+
+	// If the provided span is completely outside the rangeID local spans for any
+	// rangeID, then there is no overlap with any rangeID local keys.
+	fullRangeIDLocalSpans := roachpb.Span{
+		Key:    keys.LocalRangeIDPrefix.AsRawKey(),
+		EndKey: keys.LocalRangeIDPrefix.AsRawKey().PrefixEnd(),
+	}
+	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
+	// make an assumption that spans should respect the local RangeID tree
+	// structure, and that spans that partially overlaps with
+	// fullRangeIDLocalSpans don't make logical sense.
+	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 the span is inside RangeIDLocalSpans but outside RangeIDUnreplicated,
+	// it cannot overlap local raft keys.
+	rangeIDPrefixBuf := keys.MakeRangeIDPrefixBuf(rangeID)
+	if !spanset.Overlaps(roachpb.Span{
+		Key:    rangeIDPrefixBuf.UnreplicatedPrefix(),
+		EndKey: rangeIDPrefixBuf.UnreplicatedPrefix().PrefixEnd(),
+	}, span) {
+		return nil
+	}
+
+	// RangeTombstoneKey and RaftReplicaIDKey belong to the StateEngine, and can
+	// be accessed as point keys.
+	if roachpb.Span(span).Equal(roachpb.Span{
+		Key: rangeIDPrefixBuf.RangeTombstoneKey(),
+	}) {
+		return nil
+	}
+
+	if roachpb.Span(span).Equal(roachpb.Span{
+		Key: rangeIDPrefixBuf.RaftReplicaIDKey(),
+	}) {
+		return nil
+	}
+
+	return errors.Errorf("overlapping an unreplicated rangeID span")
+}
+
+// validateIsRaftEngineSpan asserts that the provided span only overlaps with
+// keys in the Raft engine and returns an error if not.
+// Note that we could receive the span with a nil startKey, which has a special
+// meaning that the span represents: [endKey.Prev(), endKey).
+func validateIsRaftEngineSpan(span spanset.TrickySpan) error {
+	// The LogEngine owns only Store-local and RangeID-local raft keys. A span
+	// inside Store-local is correct. If it's only partially inside, an error is
+	// returned below, as part of checking RangeID-local spans.
+	if spanset.Contains(roachpb.Span{
+		Key:    keys.LocalStorePrefix,
+		EndKey: keys.LocalStoreMax,
+	}, span) {
+		return nil
+	}
+
+	// At this point, the remaining possible LogEngine keys are inside
+	// LocalRangeID spans. If the span is not completely inside it, it must
+	// overlap with some StateEngine keys.
+	if !spanset.Contains(roachpb.Span{
+		Key:    keys.LocalRangeIDPrefix.AsRawKey(),
+		EndKey: keys.LocalRangeIDPrefix.AsRawKey().PrefixEnd(),
+	}, span) {
+		return errors.Errorf("overlaps with state engine keys")
+	}
+
+	// If the span in inside LocalRangeID, we assume 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)
+	}
+	rangeIDPrefixBuf := keys.MakeRangeIDPrefixBuf(rangeID)
+	if !spanset.Contains(roachpb.Span{
+		Key:    rangeIDPrefixBuf.UnreplicatedPrefix(),
+		EndKey: rangeIDPrefixBuf.UnreplicatedPrefix().PrefixEnd(),
+	}, span) {
+		return errors.Errorf("overlaps with state engine keys")
+	}
+	if spanset.Overlaps(roachpb.Span{Key: rangeIDPrefixBuf.RangeTombstoneKey()}, span) {
+		return errors.Errorf("overlaps with state engine keys")
+	}
+	if spanset.Overlaps(roachpb.Span{Key: rangeIDPrefixBuf.RaftReplicaIDKey()}, span) {
+		return errors.Errorf("overlaps with state engine keys")
+	}
+
+	return nil
+}

pkg/kv/kvserver/kvstorage/storage_test.go

@@ -0,0 +1,167 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package kvstorage
+
+import (
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/keys"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/spanset"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/stretchr/testify/require"
+)
+
+func TestValidateIsStateEngineSpan(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 state engine spans.
+		{span: s(roachpb.KeyMin, keys.LocalRangeIDPrefix.AsRawKey())},
+		{span: s(keys.RangeForceFlushKey(1), keys.RangeLeaseKey(1))},
+		{span: s(keys.LocalStoreMax, roachpb.KeyMax)},
+
+		// Full non-state engine spans.
+		{span: s(roachpb.KeyMin, keys.MakeRangeIDUnreplicatedPrefix(1)), notOk: true}, // partial overlap
+		{span: s(roachpb.KeyMin, roachpb.Key(keys.LocalStorePrefix).Next()), notOk: true},
+		{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.MakeRangeIDUnreplicatedPrefix(1).PrefixEnd(), roachpb.KeyMax), notOk: true}, // partial overlap
+		{span: s(keys.MakeRangeIDUnreplicatedPrefix(1),
+			keys.MakeRangeIDUnreplicatedPrefix(1).PrefixEnd()), notOk: true},
+		{span: s(keys.RangeTombstoneKey(1), keys.RaftTruncatedStateKey(1)), notOk: true},
+		{span: s(keys.RaftReplicaIDKey(1), keys.RaftTruncatedStateKey(1)), notOk: true},
+		{span: s(keys.RaftTruncatedStateKey(1), roachpb.KeyMax), 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 state engine spans.
+		{span: s(roachpb.KeyMin, nil)},
+		{span: s(keys.LocalRangeIDPrefix.AsRawKey().Prevish(1), nil)},
+		{span: s(keys.RangeForceFlushKey(1), nil)},
+		{span: s(keys.MakeRangeIDUnreplicatedPrefix(1).Prevish(1), nil)},
+		{span: s(keys.RangeTombstoneKey(1), nil)},
+		{span: s(keys.RaftReplicaIDKey(1), nil)},
+		{span: s(keys.MakeRangeIDUnreplicatedPrefix(1).PrefixEnd(), nil)},
+		{span: s(keys.LocalRangeIDPrefix.AsRawKey().PrefixEnd(), nil)},
+		{span: s(roachpb.Key(keys.LocalStorePrefix).Prevish(1), nil)},
+		{span: s(keys.LocalStoreMax, nil)},
+		{span: s(keys.LocalStoreMax.Next(), nil)},
+		{span: s(roachpb.KeyMax, nil)},
+
+		// Point non-state engine spans.
+		{span: s(keys.MakeRangeIDUnreplicatedPrefix(1), nil), notOk: true},
+		{span: s(keys.RaftTruncatedStateKey(1), nil), notOk: true},
+		{span: s(keys.RaftTruncatedStateKey(2), nil), notOk: true},
+		{span: s(keys.MakeRangeIDUnreplicatedPrefix(1).PrefixEnd().Prevish(1), nil), notOk: true},
+		{span: s(keys.LocalStorePrefix, nil), notOk: true},
+		{span: s(keys.StoreGossipKey(), nil), notOk: true},
+		{span: s(keys.LocalStoreMax.Prevish(1), nil), notOk: true},
+
+		// Tricky state engine spans.
+		{span: s(nil, keys.LocalRangeIDPrefix.AsRawKey())},
+		{span: s(nil, keys.MakeRangeIDUnreplicatedPrefix(1))},
+		{span: s(nil, keys.RangeForceFlushKey(1))},
+		{span: s(nil, keys.MakeRangeIDUnreplicatedPrefix(1).PrefixEnd().Next())},
+		{span: s(nil, keys.LocalRangeIDPrefix.AsRawKey().PrefixEnd().Next())},
+		{span: s(nil, keys.LocalStorePrefix)},
+		{span: s(nil, keys.LocalStoreMax.Next())},
+
+		// Tricky non-state engine spans.
+		{span: s(nil, keys.MakeRangeIDUnreplicatedPrefix(1).Next()), notOk: true},
+		{span: s(nil, keys.RaftTruncatedStateKey(1).Next()), notOk: true},
+		{span: s(nil, keys.RaftTruncatedStateKey(2).Next()), notOk: true},
+		{span: s(nil, keys.MakeRangeIDUnreplicatedPrefix(1).PrefixEnd()), notOk: true},
+		{span: s(nil, keys.LocalRangeIDPrefix.AsRawKey().PrefixEnd()), notOk: true}, // can't decode RangeID.
+		{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 := validateIsStateEngineSpan(spanset.TrickySpan(tc.span))
+			require.Equal(t, tc.notOk, err != nil, tc.span)
+		})
+	}
+}
+
+func TestValidateIsRaftEngineSpan(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 state engine spans.
+		{span: s(keys.RangeTombstoneKey(1).Next(), keys.RaftReplicaIDKey(1))},
+		{span: s(keys.RaftReplicaIDKey(1).Next(), keys.RangeLastReplicaGCTimestampKey(1).Next())},
+		{span: s(keys.LocalStorePrefix, keys.LocalStoreMax)},
+
+		// Full spans overlapping with state engine spans.
+		{span: s(keys.MinKey, keys.LocalStorePrefix), notOk: true},
+		{span: s(keys.RangeGCThresholdKey(1), keys.RangeVersionKey(1)), notOk: true},
+		{span: s(keys.RangeTombstoneKey(1), keys.RaftReplicaIDKey(1)), notOk: true},
+		{span: s(keys.RaftReplicaIDKey(1), keys.LocalRangeIDPrefix.AsRawKey().PrefixEnd()), notOk: true},
+		{span: s(keys.RangeGCThresholdKey(1), keys.RangeGCThresholdKey(2)), notOk: true},
+		{span: s(keys.LocalRangeIDPrefix.AsRawKey().PrefixEnd(), keys.LocalStorePrefix), notOk: true},
+		{span: s(keys.LocalStoreMax, keys.MaxKey), notOk: true},
+
+		// Point spans not overlapping with state engine spans.
+		{span: s(keys.RaftLogKey(1, 1), nil)},
+		{span: s(keys.RangeTombstoneKey(1).Next(), nil)},
+		{span: s(keys.RaftReplicaIDKey(1).Next(), nil)},
+		{span: s(keys.RaftTruncatedStateKey(1).Next(), nil)},
+		{span: s(keys.LocalStorePrefix, nil)},
+		{span: s(roachpb.Key(keys.LocalStorePrefix).Next(), nil)},
+		{span: s(keys.LocalStoreMax.Prevish(1), nil)},
+
+		// Point spans overlapping with state engine spans.
+		{span: s(keys.LocalRangeIDPrefix.AsRawKey(), nil), notOk: true}, // invalid start key
+		{span: s(keys.RangeLeaseKey(1), nil), notOk: true},
+		{span: s(keys.RangeTombstoneKey(1), nil), notOk: true},
+		{span: s(keys.RaftReplicaIDKey(1), nil), notOk: true},
+		{span: s(keys.RangeTombstoneKey(2), nil), notOk: true},
+		{span: s(keys.RaftReplicaIDKey(2), nil), notOk: true},
+		{span: s(keys.LocalRangeIDPrefix.AsRawKey().PrefixEnd(), nil), notOk: true},
+
+		// Tricky spans not overlapping with state engine spans.
+		{span: s(nil, keys.RaftLogKey(1, 1))},
+		{span: s(nil, keys.RangeTombstoneKey(1))},
+		{span: s(nil, keys.RaftReplicaIDKey(1))},
+		{span: s(nil, roachpb.Key(keys.LocalStorePrefix).PrefixEnd())},
+
+		// Tricky spans overlapping with state engine spans.
+		{span: s(nil, keys.LocalRangeIDPrefix.AsRawKey()), notOk: true},
+		{span: s(nil, keys.LocalStorePrefix), notOk: true},
+		{span: s(nil, keys.RangeLeaseKey(1)), notOk: true},
+		{span: s(nil, keys.RangeTombstoneKey(1).Next()), notOk: true},
+		{span: s(nil, keys.RaftReplicaIDKey(1).Next()), notOk: true},
+		{span: s(nil, roachpb.Key(keys.LocalStorePrefix).PrefixEnd().Next()), notOk: true},
+	}
+
+	for _, tc := range testCases {
+		t.Run("", func(t *testing.T) {
+			err := validateIsRaftEngineSpan(spanset.TrickySpan(tc.span))
+			require.Equal(t, tc.notOk, err != nil, tc.span)
+		})
+	}
+}