diff --git a/pkg/kv/kvserver/kvstorage/BUILD.bazel b/pkg/kv/kvserver/kvstorage/BUILD.bazel index fa9afff01c20..3a696ce06a8f 100644 --- a/pkg/kv/kvserver/kvstorage/BUILD.bazel +++ b/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", diff --git a/pkg/kv/kvserver/kvstorage/destroy_test.go b/pkg/kv/kvserver/kvstorage/destroy_test.go index 7f5d2ff94e42..a849ec966392 100644 --- a/pkg/kv/kvserver/kvstorage/destroy_test.go +++ b/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}, diff --git a/pkg/kv/kvserver/kvstorage/storage.go b/pkg/kv/kvserver/kvstorage/storage.go index 77447d08de21..23c4f43bb0a8 100644 --- a/pkg/kv/kvserver/kvstorage/storage.go +++ b/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 +} diff --git a/pkg/kv/kvserver/kvstorage/storage_test.go b/pkg/kv/kvserver/kvstorage/storage_test.go new file mode 100644 index 000000000000..20a0a896f4f3 --- /dev/null +++ b/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) + }) + } +} diff --git a/pkg/kv/kvserver/replica_raftstorage.go b/pkg/kv/kvserver/replica_raftstorage.go index ce7b031259c9..0836cdde7d1d 100644 --- a/pkg/kv/kvserver/replica_raftstorage.go +++ b/pkg/kv/kvserver/replica_raftstorage.go @@ -879,128 +879,3 @@ func (r *Replica) destroyInfoRaftMuLocked() kvstorage.DestroyReplicaInfo { Keys: r.shMu.state.Desc.RSpan(), } } - -// 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 -} diff --git a/pkg/kv/kvserver/replica_read.go b/pkg/kv/kvserver/replica_read.go index 387132b3acd0..d6e529ddd04c 100644 --- a/pkg/kv/kvserver/replica_read.go +++ b/pkg/kv/kvserver/replica_read.go @@ -94,7 +94,6 @@ func (r *Replica) executeReadOnlyBatch( } if util.RaceEnabled { spans := g.LatchSpans() - spans.AddForbiddenMatcher(validateIsStateEngineSpan) rw = spanset.NewReadWriterAt(rw, spans, ba.Timestamp) } defer rw.Close() diff --git a/pkg/kv/kvserver/replica_test.go b/pkg/kv/kvserver/replica_test.go index ba6d9ff1b9b5..33e438c976d6 100644 --- a/pkg/kv/kvserver/replica_test.go +++ b/pkg/kv/kvserver/replica_test.go @@ -15201,153 +15201,3 @@ func TestLeaderlessWatcherInit(t *testing.T) { t.Fatalf("expected LeaderlessWatcher channel to be closed") } } - -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) - }) - } -} diff --git a/pkg/kv/kvserver/replica_write.go b/pkg/kv/kvserver/replica_write.go index 49de987b2a8b..bebf515a412d 100644 --- a/pkg/kv/kvserver/replica_write.go +++ b/pkg/kv/kvserver/replica_write.go @@ -815,7 +815,6 @@ func (r *Replica) newBatchedEngine(g *concurrency.Guard) (storage.Batch, *storag // any write latch would be declared at. But because of this, we don't // assert on access timestamps using spanset.NewBatchAt. spans := g.LatchSpans() - spans.AddForbiddenMatcher(validateIsStateEngineSpan) batch = spanset.NewBatch(batch, spans) } diff --git a/pkg/kv/kvserver/spanset/BUILD.bazel b/pkg/kv/kvserver/spanset/BUILD.bazel index 7cd201a681ba..c0d92753749d 100644 --- a/pkg/kv/kvserver/spanset/BUILD.bazel +++ b/pkg/kv/kvserver/spanset/BUILD.bazel @@ -4,6 +4,7 @@ go_library( name = "spanset", srcs = [ "batch.go", + "engine.go", "merge.go", "spanset.go", ], @@ -13,6 +14,7 @@ go_library( "//pkg/keys", "//pkg/roachpb", "//pkg/storage", + "//pkg/storage/enginepb", "//pkg/storage/fs", "//pkg/util/debugutil", "//pkg/util/hlc", @@ -20,6 +22,7 @@ go_library( "//pkg/util/protoutil", "@com_github_cockroachdb_errors//:errors", "@com_github_cockroachdb_pebble//:pebble", + "@com_github_cockroachdb_pebble//metrics", "@com_github_cockroachdb_pebble//rangekey", ], ) @@ -29,6 +32,7 @@ go_test( size = "small", srcs = [ "batch_test.go", + "engine_test.go", "merge_test.go", "spanset_test.go", ], @@ -42,6 +46,7 @@ go_test( "//pkg/testutils", "//pkg/util/hlc", "//pkg/util/leaktest", + "@com_github_cockroachdb_errors//:errors", "@com_github_stretchr_testify//require", ], ) diff --git a/pkg/kv/kvserver/spanset/batch.go b/pkg/kv/kvserver/spanset/batch.go index 565c20953a35..5b099fc2e207 100644 --- a/pkg/kv/kvserver/spanset/batch.go +++ b/pkg/kv/kvserver/spanset/batch.go @@ -964,6 +964,9 @@ func NewBatchAt(b storage.Batch, spans *SpanSet, ts hlc.Timestamp) storage.Batch // DisableReaderAssertions unwraps any storage.Reader implementations that may // assert access against a given SpanSet. +// TODO(ibrahim): Eventually we want to eliminate all the users of these generic +// disable functions, and use the specific disable functions +// (DisableUndeclaredSpanAssertions or DisableForbiddenSpanAssertions). func DisableReaderAssertions(reader storage.Reader) storage.Reader { switch v := reader.(type) { case ReadWriter: @@ -975,10 +978,32 @@ func DisableReaderAssertions(reader storage.Reader) storage.Reader { } } +// DisableWriterAssertions unwraps any storage.Writer implementations that may +// assert access against a given SpanSet. +// TODO(ibrahim): Eventually we want to eliminate all the users of these generic +// disable functions, and use the specific disable functions +// (DisableUndeclaredSpanAssertions or DisableForbiddenSpanAssertions). +func DisableWriterAssertions(writer storage.Writer) storage.Writer { + switch v := writer.(type) { + case spanSetWriter: + return DisableWriterAssertions(v.w) + case *spanSetWriteBatch: + return DisableWriterAssertions(v.spanSetWriter.w) + case *spanSetBatch: + return DisableWriterAssertions(v.spanSetWriter.w) + default: + return writer + } +} + // DisableReadWriterAssertions unwraps any storage.ReadWriter implementations // that may assert access against a given SpanSet. +// TODO(ibrahim): Eventually we want to eliminate all the users of these generic +// disable functions, and use the specific disable functions +// (DisableUndeclaredSpanAssertions or DisableForbiddenSpanAssertions). func DisableReadWriterAssertions(rw storage.ReadWriter) storage.ReadWriter { switch v := rw.(type) { + // TODO(ibrahim): fix the case where one case is a pointer and the other is not. case ReadWriter: return DisableReadWriterAssertions(v.spanSetWriter.w.(storage.ReadWriter)) case *spanSetBatch: @@ -995,9 +1020,18 @@ func DisableReadWriterAssertions(rw storage.ReadWriter) storage.ReadWriter { func DisableUndeclaredSpanAssertions(rw storage.ReadWriter) storage.ReadWriter { switch v := rw.(type) { case *spanSetBatch: - newSnapSetBatch := v.shallowCopy() - newSnapSetBatch.spans.DisableUndeclaredAccessAssertions() - return newSnapSetBatch + newSpanSetBatch := v.shallowCopy() + newSpanSetBatch.spans.DisableUndeclaredAccessAssertions() + + // Recursively disable on the underlying batch in case there are + // nested spanSetBatches. + newSpanSetBatch.b = DisableUndeclaredSpanAssertions(v.b).(storage.Batch) + // Update the reader and writer to point to the recursively processed batch. + newSpanSetBatch.spanSetReader.r = newSpanSetBatch.b + newSpanSetBatch.spanSetWriteBatch.spanSetWriter.w = newSpanSetBatch.b + newSpanSetBatch.spanSetWriteBatch.wb = newSpanSetBatch.b + return newSpanSetBatch + default: return rw } @@ -1011,10 +1045,27 @@ func DisableUndeclaredSpanAssertions(rw storage.ReadWriter) storage.ReadWriter { // function. func DisableForbiddenSpanAssertions(rw storage.ReadWriter) storage.ReadWriter { switch v := rw.(type) { + // TODO(ibrahim): We eventually want to remove OpLoggerBatch from this switch + // case. + case *storage.OpLoggerBatch: + // OpLoggerBatch embeds a storage.Batch. Recursively process the inner + // batch and wrap it back in an OpLoggerBatch to preserve the chain. + innerBatch := DisableForbiddenSpanAssertions(v.Batch).(storage.Batch) + return v.ShallowCopyWithBatch(innerBatch) + case *spanSetBatch: - newSnapSetBatch := v.shallowCopy() - newSnapSetBatch.spans.DisableForbiddenSpansAssertions() - return newSnapSetBatch + newSpanSetBatch := v.shallowCopy() + newSpanSetBatch.spans.DisableForbiddenSpansAssertions() + + // Recursively disable on the underlying batch in case there are + // nested spanSetBatches. + newSpanSetBatch.b = DisableForbiddenSpanAssertions(v.b).(storage.Batch) + // Update the reader and writer to point to the recursively processed batch. + newSpanSetBatch.spanSetReader.r = newSpanSetBatch.b + newSpanSetBatch.spanSetWriteBatch.spanSetWriter.w = newSpanSetBatch.b + newSpanSetBatch.spanSetWriteBatch.wb = newSpanSetBatch.b + return newSpanSetBatch + default: return rw } diff --git a/pkg/kv/kvserver/spanset/engine.go b/pkg/kv/kvserver/spanset/engine.go new file mode 100644 index 000000000000..0458de5b9e0a --- /dev/null +++ b/pkg/kv/kvserver/spanset/engine.go @@ -0,0 +1,295 @@ +// Copyright 2025 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. + +package spanset + +import ( + "context" + + "github.com/cockroachdb/cockroach/pkg/roachpb" + "github.com/cockroachdb/cockroach/pkg/storage" + "github.com/cockroachdb/cockroach/pkg/storage/enginepb" + "github.com/cockroachdb/cockroach/pkg/storage/fs" + "github.com/cockroachdb/cockroach/pkg/util/hlc" + "github.com/cockroachdb/pebble" + "github.com/cockroachdb/pebble/metrics" +) + +// spanSetEngine wraps an Engine and asserts that it does not access spans that +// do not belong to it. +type spanSetEngine struct { + ReadWriter + e storage.Engine + spans *SpanSet +} + +var _ storage.Engine = &spanSetEngine{} + +// NewEngine creates a new spanSetEngine wrapper. +func NewEngine(engine storage.Engine, forbiddenMatchers ...func(TrickySpan) error) storage.Engine { + spans := &SpanSet{} + // For engines, we disable undeclared access assertions as we only care about + // preventing access to spans that don't belong to this engine. + spans.DisableUndeclaredAccessAssertions() + // Add the forbidden matchers that assert against access to disallowed keys. + for _, matcher := range forbiddenMatchers { + spans.AddForbiddenMatcher(matcher) + } + return &spanSetEngine{ + ReadWriter: makeSpanSetReadWriter(engine, spans), + e: engine, + spans: spans, + } +} + +// NewBatch implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) NewBatch() storage.Batch { + return NewBatch(s.e.NewBatch(), s.spans) +} + +// NewBatch implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) NewReader(durability storage.DurabilityRequirement) storage.Reader { + return NewReader(s.e.NewReader(durability), s.spans, hlc.Timestamp{}) +} + +// NewReadOnly implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) NewReadOnly(durability storage.DurabilityRequirement) storage.ReadWriter { + return NewReadWriterAt(s.e.NewReadOnly(durability), s.spans, hlc.Timestamp{}) +} + +// NewUnindexedBatch implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) NewUnindexedBatch() storage.Batch { + return NewBatch(s.e.NewUnindexedBatch(), s.spans) +} + +// Excise implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) Excise(ctx context.Context, span roachpb.Span) error { + if err := s.spans.CheckAllowed(SpanReadWrite, TrickySpan{Key: span.Key, EndKey: span.EndKey}); err != nil { + return err + } + return s.e.Excise(ctx, span) +} + +// Download implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) Download(ctx context.Context, span roachpb.Span, copy bool) error { + if err := s.spans.CheckAllowed(SpanReadWrite, TrickySpan{Key: span.Key, EndKey: span.EndKey}); err != nil { + return err + } + return s.e.Download(ctx, span, copy) +} + +// CreateCheckpoint implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) CreateCheckpoint(dir string, spans []roachpb.Span) error { + for _, span := range spans { + if err := s.spans.CheckAllowed(SpanReadOnly, TrickySpan{Key: span.Key, EndKey: span.EndKey}); err != nil { + return err + } + } + return s.e.CreateCheckpoint(dir, spans) +} + +// GetTableMetrics implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) GetTableMetrics( + start, end roachpb.Key, +) ([]enginepb.SSTableMetricsInfo, error) { + if err := s.spans.CheckAllowed(SpanReadOnly, TrickySpan{Key: start, EndKey: end}); err != nil { + return nil, err + } + return s.e.GetTableMetrics(start, end) +} + +// CompactRange implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) CompactRange(ctx context.Context, start, end roachpb.Key) error { + if err := s.spans.CheckAllowed(SpanReadWrite, TrickySpan{Key: start, EndKey: end}); err != nil { + return err + } + return s.e.CompactRange(ctx, start, end) +} + +// ApproximateDiskBytes implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) ApproximateDiskBytes( + from, to roachpb.Key, +) (total, remote, external uint64, _ error) { + if err := s.spans.CheckAllowed(SpanReadOnly, TrickySpan{Key: from, EndKey: to}); err != nil { + return 0, 0, 0, err + } + return s.e.ApproximateDiskBytes(from, to) +} + +// NewSnapshot implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) NewSnapshot(keyRanges ...roachpb.Span) storage.Reader { + for _, span := range keyRanges { + if err := s.spans.CheckAllowed(SpanReadOnly, TrickySpan{Key: span.Key, EndKey: span.EndKey}); err != nil { + panic(err) + } + } + return NewReader(s.e.NewSnapshot(keyRanges...), s.spans, hlc.Timestamp{}) +} + +// NewWriteBatch implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) NewWriteBatch() storage.WriteBatch { + wb := s.e.NewWriteBatch() + return &spanSetWriteBatch{ + spanSetWriter: spanSetWriter{w: wb, spans: s.spans, spansOnly: true}, + wb: wb, + } +} + +// Attrs implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) Attrs() roachpb.Attributes { + return s.e.Attrs() +} + +// Capacity implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) Capacity() (roachpb.StoreCapacity, error) { + return s.e.Capacity() +} + +// Properties implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) Properties() roachpb.StoreProperties { + return s.e.Properties() +} + +// ProfileSeparatedValueRetrievals implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) ProfileSeparatedValueRetrievals( + ctx context.Context, +) (*metrics.ValueRetrievalProfile, error) { + return s.e.ProfileSeparatedValueRetrievals(ctx) +} + +// Compact implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) Compact(ctx context.Context) error { + return s.e.Compact(ctx) +} + +// Env implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) Env() *fs.Env { + return s.e.Env() +} + +// Flush implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) Flush() error { + return s.e.Flush() +} + +// GetMetrics implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) GetMetrics() storage.Metrics { + return s.e.GetMetrics() +} + +// GetEncryptionRegistries implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) GetEncryptionRegistries() (*fs.EncryptionRegistries, error) { + return s.e.GetEncryptionRegistries() +} + +// GetEnvStats implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) GetEnvStats() (*fs.EnvStats, error) { + return s.e.GetEnvStats() +} + +// GetAuxiliaryDir implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) GetAuxiliaryDir() string { + return s.e.GetAuxiliaryDir() +} + +// IngestLocalFiles implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) IngestLocalFiles(ctx context.Context, paths []string) error { + return s.e.IngestLocalFiles(ctx, paths) +} + +// IngestLocalFilesWithStats implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) IngestLocalFilesWithStats( + ctx context.Context, paths []string, +) (pebble.IngestOperationStats, error) { + return s.e.IngestLocalFilesWithStats(ctx, paths) +} + +// IngestAndExciseFiles implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) IngestAndExciseFiles( + ctx context.Context, + paths []string, + shared []pebble.SharedSSTMeta, + external []pebble.ExternalFile, + exciseSpan roachpb.Span, +) (pebble.IngestOperationStats, error) { + if err := s.spans.CheckAllowed(SpanReadWrite, TrickySpan{Key: exciseSpan.Key, EndKey: exciseSpan.EndKey}); err != nil { + return pebble.IngestOperationStats{}, err + } + return s.e.IngestAndExciseFiles(ctx, paths, shared, external, exciseSpan) +} + +// IngestExternalFiles implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) IngestExternalFiles( + ctx context.Context, external []pebble.ExternalFile, +) (pebble.IngestOperationStats, error) { + return s.e.IngestExternalFiles(ctx, external) +} + +// IngestLocalFilesToWriter implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) IngestLocalFilesToWriter( + ctx context.Context, paths []string, clearedSpans []roachpb.Span, writer storage.Writer, +) error { + return s.e.IngestLocalFilesToWriter(ctx, paths, clearedSpans, writer) +} + +// ScanStorageInternalKeys implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) ScanStorageInternalKeys( + start, end roachpb.Key, megabytesPerSecond int64, +) ([]enginepb.StorageInternalKeysMetrics, error) { + if err := s.spans.CheckAllowed(SpanReadWrite, TrickySpan{Key: start, EndKey: end}); err != nil { + return nil, err + } + return s.e.ScanStorageInternalKeys(start, end, megabytesPerSecond) +} + +// RegisterFlushCompletedCallback implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) RegisterFlushCompletedCallback(cb func()) { + s.e.RegisterFlushCompletedCallback(cb) +} + +// MinVersion implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) MinVersion() roachpb.Version { + return s.e.MinVersion() +} + +// SetMinVersion implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) SetMinVersion(version roachpb.Version) error { + return s.e.SetMinVersion(version) +} + +// SetCompactionConcurrency implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) SetCompactionConcurrency(n uint64) { + s.e.SetCompactionConcurrency(n) +} + +// SetStoreID implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) SetStoreID(ctx context.Context, storeID int32) error { + return s.e.SetStoreID(ctx, storeID) +} + +// GetStoreID implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) GetStoreID() (int32, error) { + return s.e.GetStoreID() +} + +// RegisterDiskSlowCallback implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) RegisterDiskSlowCallback(cb func(info pebble.DiskSlowInfo)) { + s.e.RegisterDiskSlowCallback(cb) +} + +// RegisterLowDiskSpaceCallback implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) RegisterLowDiskSpaceCallback(cb func(info pebble.LowDiskSpaceInfo)) { + s.e.RegisterLowDiskSpaceCallback(cb) +} + +// GetPebbleOptions implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) GetPebbleOptions() *pebble.Options { + return s.e.GetPebbleOptions() +} + +// GetDiskUnhealthy implements the storage.EngineWithoutRW interface. +func (s *spanSetEngine) GetDiskUnhealthy() bool { + return s.e.GetDiskUnhealthy() +} diff --git a/pkg/kv/kvserver/spanset/engine_test.go b/pkg/kv/kvserver/spanset/engine_test.go new file mode 100644 index 000000000000..bfb039280848 --- /dev/null +++ b/pkg/kv/kvserver/spanset/engine_test.go @@ -0,0 +1,67 @@ +// Copyright 2025 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. + +package spanset + +import ( + "context" + "testing" + + "github.com/cockroachdb/cockroach/pkg/roachpb" + "github.com/cockroachdb/cockroach/pkg/storage" + "github.com/cockroachdb/cockroach/pkg/util/leaktest" + "github.com/cockroachdb/errors" + "github.com/stretchr/testify/require" +) + +// TestSpanSetEngine tests that the spanSetEngine correctly enforces forbidden +// spans when performing engine operations. +func TestSpanSetEngine(t *testing.T) { + defer leaktest.AfterTest(t)() + + // Create an in-memory engine. + engine := storage.NewDefaultInMemForTesting() + defer engine.Close() + + allowedKey := roachpb.Key("allowed") + forbiddenKey := roachpb.Key("forbidden") + fm := func(span TrickySpan) error { + if Overlaps(roachpb.Span{ + Key: forbiddenKey, + }, span) { + return errors.Errorf("forbidden access: %s", span) + } + return nil + } + + // Wrap the engine with spanSetEngine. + wrappedEngine := NewEngine(engine, fm) + + // Create a batch from the wrapped engine. + batch := wrappedEngine.NewBatch() + defer batch.Close() + + // Writing to the declared span should succeed. + err := batch.PutUnversioned(allowedKey, []byte("value")) + require.NoError(t, err) + + // Writing to an undeclared span should fail. + err = batch.PutUnversioned(forbiddenKey, []byte("value")) + require.Error(t, err) + + // Run a direct engine call to Excise on a non-declared span; it should fail. + err = wrappedEngine.Excise(context.Background(), roachpb.Span{ + Key: forbiddenKey, + EndKey: forbiddenKey.Next(), + }) + require.Error(t, err) + + // Run a direct engine call to Excise on a declared span; it should succeed. + err = wrappedEngine.Excise(context.Background(), roachpb.Span{ + Key: allowedKey, + EndKey: allowedKey.Next(), + }) + require.NoError(t, err) +} diff --git a/pkg/storage/mvcc_logical_ops.go b/pkg/storage/mvcc_logical_ops.go index b7c8f8c9162c..10e24483625b 100644 --- a/pkg/storage/mvcc_logical_ops.go +++ b/pkg/storage/mvcc_logical_ops.go @@ -170,6 +170,14 @@ func (ol *OpLoggerBatch) LogicalOps() []enginepb.MVCCLogicalOp { return ol.ops } +// ShallowCopyWithBatch returns a shallow copy of the OpLoggerBatch with a +// different underlying batch. The ops slice is shared with the original. +func (ol *OpLoggerBatch) ShallowCopyWithBatch(b Batch) *OpLoggerBatch { + cpy := *ol + cpy.Batch = b + return &cpy +} + // DisableOpLogger disables op logging for the given read/writer. func DisableOpLogger(rw ReadWriter) ReadWriter { return &noOpLogger{ReadWriter: rw}