kvstorage: assert each engine only access its spans

The commit adds engine assertions to assert that Log
engine only accesses log engine spans, and state engine
only accesses state engine spans.

Author: iskettaneh

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,18 @@ 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.
+		wrappedStateEngine := spanset.NewSpanSetEngine(eng, validateIsStateEngineSpan)
+		wrappedLogEngine := spanset.NewSpanSetEngine(eng, validateIsRaftEngineSpan)
+		return Engines{
+			stateEngine: wrappedStateEngine,
+			logEngine:   wrappedLogEngine,
+			todoEngine:  eng,
+		}
+	}
+
 	return Engines{
 		stateEngine: eng,
 		todoEngine:  eng,
@@ -137,6 +152,18 @@ 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.
+		wrappedStateEngine := spanset.NewSpanSetEngine(state, validateIsStateEngineSpan)
+		wrappedLogEngine := spanset.NewSpanSetEngine(log, validateIsRaftEngineSpan)
+		return Engines{
+			stateEngine: wrappedStateEngine,
+			todoEngine:  nil,
+			logEngine:   wrappedLogEngine,
+			separated:   true,
+		}
+	}
 	return Engines{
 		stateEngine: state,
 		todoEngine:  nil,
@@ -178,3 +205,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)
+		})
+	}
+}