kvserver: add overlapsStateEngineKeys

This commit adds the function overlapsStateEngineKeys that
checks if a specific span overlaps with any StateEngine keys.

This will be used later to enforce RaftEngine assertions and
ensure that we do not touch and StateEngine keys.

References: #158281

Release note: None

Author: iskettaneh

pkg/keys/keys.go

@@ -1142,18 +1142,18 @@ func MakeRangeIDPrefixBuf(rangeID roachpb.RangeID) RangeIDPrefixBuf {
 	return RangeIDPrefixBuf(MakeRangeIDPrefix(rangeID))
 }
 
-func (b RangeIDPrefixBuf) replicatedPrefix() roachpb.Key {
+func (b RangeIDPrefixBuf) ReplicatedPrefix() roachpb.Key {
 	return append(roachpb.Key(b), LocalRangeIDReplicatedInfix...)
 }
 
-func (b RangeIDPrefixBuf) unreplicatedPrefix() roachpb.Key {
+func (b RangeIDPrefixBuf) UnreplicatedPrefix() roachpb.Key {
 	return append(roachpb.Key(b), localRangeIDUnreplicatedInfix...)
 }
 
 // AbortSpanKey returns a range-local key by Range ID for an AbortSpan
 // entry, with detail specified by encoding the supplied transaction ID.
 func (b RangeIDPrefixBuf) AbortSpanKey(txnID uuid.UUID) roachpb.Key {
-	key := append(b.replicatedPrefix(), LocalAbortSpanSuffix...)
+	key := append(b.ReplicatedPrefix(), LocalAbortSpanSuffix...)
 	return encoding.EncodeBytesAscending(key, txnID.GetBytes())
 }
 
@@ -1162,67 +1162,67 @@ func (b RangeIDPrefixBuf) AbortSpanKey(txnID uuid.UUID) roachpb.Key {
 // specific transaction should serialize on latches. The per-transaction bit is
 // achieved by encoding the supplied transaction ID into the key.
 func (b RangeIDPrefixBuf) ReplicatedSharedLocksTransactionLatchingKey(txnID uuid.UUID) roachpb.Key {
-	key := append(b.replicatedPrefix(), LocalReplicatedSharedLocksTransactionLatchingKeySuffix...)
+	key := append(b.ReplicatedPrefix(), LocalReplicatedSharedLocksTransactionLatchingKeySuffix...)
 	return encoding.EncodeBytesAscending(key, txnID.GetBytes())
 }
 
 // RangeAppliedStateKey returns a system-local key for the range applied state key.
 // See comment on RangeAppliedStateKey function.
 func (b RangeIDPrefixBuf) RangeAppliedStateKey() roachpb.Key {
-	return append(b.replicatedPrefix(), LocalRangeAppliedStateSuffix...)
+	return append(b.ReplicatedPrefix(), LocalRangeAppliedStateSuffix...)
 }
 
 // RangeForceFlushKey returns a system-local key for the range force flush
 // key.
 func (b RangeIDPrefixBuf) RangeForceFlushKey() roachpb.Key {
-	return append(b.replicatedPrefix(), LocalRangeForceFlushSuffix...)
+	return append(b.ReplicatedPrefix(), LocalRangeForceFlushSuffix...)
 }
 
 // RangeLeaseKey returns a system-local key for a range lease.
 func (b RangeIDPrefixBuf) RangeLeaseKey() roachpb.Key {
-	return append(b.replicatedPrefix(), LocalRangeLeaseSuffix...)
+	return append(b.ReplicatedPrefix(), LocalRangeLeaseSuffix...)
 }
 
 // RangePriorReadSummaryKey returns a system-local key for a range's prior read
 // summary.
 func (b RangeIDPrefixBuf) RangePriorReadSummaryKey() roachpb.Key {
-	return append(b.replicatedPrefix(), LocalRangePriorReadSummarySuffix...)
+	return append(b.ReplicatedPrefix(), LocalRangePriorReadSummarySuffix...)
 }
 
 // RangeGCThresholdKey returns a system-local key for the GC threshold.
 func (b RangeIDPrefixBuf) RangeGCThresholdKey() roachpb.Key {
-	return append(b.replicatedPrefix(), LocalRangeGCThresholdSuffix...)
+	return append(b.ReplicatedPrefix(), LocalRangeGCThresholdSuffix...)
 }
 
 // RangeGCHintKey returns a range-local key for the GC hint data.
 func (b RangeIDPrefixBuf) RangeGCHintKey() roachpb.Key {
-	return append(b.replicatedPrefix(), LocalRangeGCHintSuffix...)
+	return append(b.ReplicatedPrefix(), LocalRangeGCHintSuffix...)
 }
 
 // RangeVersionKey returns a system-local key for the range version.
 func (b RangeIDPrefixBuf) RangeVersionKey() roachpb.Key {
-	return append(b.replicatedPrefix(), LocalRangeVersionSuffix...)
+	return append(b.ReplicatedPrefix(), LocalRangeVersionSuffix...)
 }
 
 // RangeTombstoneKey returns a system-local key for a range tombstone.
 func (b RangeIDPrefixBuf) RangeTombstoneKey() roachpb.Key {
-	return append(b.unreplicatedPrefix(), LocalRangeTombstoneSuffix...)
+	return append(b.UnreplicatedPrefix(), LocalRangeTombstoneSuffix...)
 }
 
 // RaftTruncatedStateKey returns a system-local key for a RaftTruncatedState.
 func (b RangeIDPrefixBuf) RaftTruncatedStateKey() roachpb.Key {
-	return append(b.unreplicatedPrefix(), LocalRaftTruncatedStateSuffix...)
+	return append(b.UnreplicatedPrefix(), LocalRaftTruncatedStateSuffix...)
 }
 
 // RaftHardStateKey returns a system-local key for a Raft HardState.
 func (b RangeIDPrefixBuf) RaftHardStateKey() roachpb.Key {
-	return append(b.unreplicatedPrefix(), LocalRaftHardStateSuffix...)
+	return append(b.UnreplicatedPrefix(), LocalRaftHardStateSuffix...)
 }
 
 // RaftLogPrefix returns the system-local prefix shared by all Entries
 // in a Raft log.
 func (b RangeIDPrefixBuf) RaftLogPrefix() roachpb.Key {
-	return append(b.unreplicatedPrefix(), LocalRaftLogSuffix...)
+	return append(b.UnreplicatedPrefix(), LocalRaftLogSuffix...)
 }
 
 // RaftLogKey returns a system-local key for a Raft log entry.
@@ -1232,17 +1232,17 @@ func (b RangeIDPrefixBuf) RaftLogKey(logIndex kvpb.RaftIndex) roachpb.Key {
 
 // RaftReplicaIDKey returns a system-local key for a RaftReplicaID.
 func (b RangeIDPrefixBuf) RaftReplicaIDKey() roachpb.Key {
-	return append(b.unreplicatedPrefix(), LocalRaftReplicaIDSuffix...)
+	return append(b.UnreplicatedPrefix(), LocalRaftReplicaIDSuffix...)
 }
 
 // RangeLastReplicaGCTimestampKey returns a range-local key for
 // the range's last replica GC timestamp.
 func (b RangeIDPrefixBuf) RangeLastReplicaGCTimestampKey() roachpb.Key {
-	return append(b.unreplicatedPrefix(), LocalRangeLastReplicaGCTimestampSuffix...)
+	return append(b.UnreplicatedPrefix(), LocalRangeLastReplicaGCTimestampSuffix...)
 }
 
 // MVCCRangeKeyGCKey returns a range local key protecting range
 // tombstone mvcc stats calculations during range tombstone GC.
 func (b RangeIDPrefixBuf) MVCCRangeKeyGCKey() roachpb.Key {
-	return append(b.unreplicatedPrefix(), LocalRangeMVCCRangeKeyGCLockSuffix...)
+	return append(b.UnreplicatedPrefix(), LocalRangeMVCCRangeKeyGCLockSuffix...)
 }

pkg/kv/kvserver/replica_raftstorage.go

@@ -880,18 +880,26 @@ func (r *Replica) destroyInfoRaftMuLocked() kvstorage.DestroyReplicaInfo {
 	}
 }
 
-// overlapsLocalRaftKeys returns an error if the provided span overlaps with any
-// unreplicated rangeID local raft keys.
+// 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 overlapsLocalRaftKeys(span spanset.TrickySpan) error {
-	fullRangeIDLocalSpans := roachpb.Span{
-		Key:    keys.LocalRangeIDPrefix.AsRawKey(),
-		EndKey: keys.LocalRangeIDPrefix.AsRawKey().PrefixEnd(),
+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
 	}
@@ -919,42 +927,79 @@ func overlapsLocalRaftKeys(span spanset.TrickySpan) error {
 
 	// 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:    keys.MakeRangeIDUnreplicatedPrefix(rangeID),
-		EndKey: keys.MakeRangeIDUnreplicatedPrefix(rangeID).PrefixEnd(),
+		Key:    rangeIDPrefixBuf.UnreplicatedPrefix(),
+		EndKey: rangeIDPrefixBuf.UnreplicatedPrefix().PrefixEnd(),
 	}, span) {
 		return nil
 	}
 
-	// RangeTombstoneKey and RaftReplicaIDKey are considered part of the
-	// StateEngine keys. We need to exclude them from the check below.
-	if span.Key != nil && roachpb.Span(span).Equal(roachpb.Span{
-		Key: keys.RangeTombstoneKey(rangeID),
+	// 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 span.Key != nil && roachpb.Span(span).Equal(roachpb.Span{
-		Key: keys.RaftReplicaIDKey(rangeID),
+	if roachpb.Span(span).Equal(roachpb.Span{
+		Key: rangeIDPrefixBuf.RaftReplicaIDKey(),
 	}) {
 		return nil
 	}
 
 	return errors.Errorf("overlapping an unreplicated rangeID span")
 }
 
-// overlapsStoreLocalKeys returns an error if the provided span overlaps
-// with any store local keys.
+// 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 overlapsStoreLocalKeys(span spanset.TrickySpan) error {
-	localStoreSpan := roachpb.Span{
+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
 	}
 
-	if spanset.Overlaps(localStoreSpan, span) {
-		return errors.Errorf("overlaps with store local keys")
+	// 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/replica_read.go

@@ -94,8 +94,7 @@ func (r *Replica) executeReadOnlyBatch(
 	}
 	if util.RaceEnabled {
 		spans := g.LatchSpans()
-		spans.AddForbiddenMatcher(overlapsLocalRaftKeys)
-		spans.AddForbiddenMatcher(overlapsStoreLocalKeys)
+		spans.AddForbiddenMatcher(validateIsStateEngineSpan)
 		rw = spanset.NewReadWriterAt(rw, spans, ba.Timestamp)
 	}
 	defer rw.Close()