kvserver/spanset: introduce forbidden spans matchers

This commit introduces the concept of forbidden spans, which will
allow us to assert our batches don't modify Raft's engine keys.

References: #156537

Release note: None

Author: iskettaneh

pkg/kv/kvserver/batcheval/cmd_delete.go

@@ -58,7 +58,7 @@ func Delete(
 	// If requested, replace point tombstones with range tombstones.
 	if cArgs.EvalCtx.EvalKnobs().UseRangeTombstonesForPointDeletes && h.Txn == nil {
 		if err := storage.ReplacePointTombstonesWithRangeTombstones(
-			ctx, spanset.DisableReadWriterAssertions(readWriter),
+			ctx, spanset.DisableLatchAssertions(readWriter),
 			cArgs.Stats, args.Key, args.EndKey); err != nil {
 			return result.Result{}, err
 		}

pkg/kv/kvserver/batcheval/cmd_delete_range.go

@@ -302,7 +302,7 @@ func deleteRangeTransactional(
 	// If requested, replace point tombstones with range tombstones.
 	if cArgs.EvalCtx.EvalKnobs().UseRangeTombstonesForPointDeletes && h.Txn == nil {
 		if err := storage.ReplacePointTombstonesWithRangeTombstones(
-			ctx, spanset.DisableReadWriterAssertions(readWriter),
+			ctx, spanset.DisableLatchAssertions(readWriter),
 			cArgs.Stats, args.Key, args.EndKey); err != nil {
 			return result.Result{}, err
 		}

pkg/kv/kvserver/batcheval/cmd_end_transaction.go

@@ -719,7 +719,7 @@ func resolveLocalLocksWithPagination(
 				// If requested, replace point tombstones with range tombstones.
 				if ok && evalCtx.EvalKnobs().UseRangeTombstonesForPointDeletes {
 					if err := storage.ReplacePointTombstonesWithRangeTombstones(
-						ctx, spanset.DisableReadWriterAssertions(readWriter),
+						ctx, spanset.DisableLatchAssertions(readWriter),
 						ms, update.Key, update.EndKey); err != nil {
 						return 0, 0, 0, errors.Wrapf(err,
 							"replacing point tombstones with range tombstones for write intent at %s on end transaction [%s]",
@@ -757,7 +757,7 @@ func resolveLocalLocksWithPagination(
 			// If requested, replace point tombstones with range tombstones.
 			if evalCtx.EvalKnobs().UseRangeTombstonesForPointDeletes {
 				if err := storage.ReplacePointTombstonesWithRangeTombstones(
-					ctx, spanset.DisableReadWriterAssertions(readWriter),
+					ctx, spanset.DisableLatchAssertions(readWriter),
 					ms, update.Key, update.EndKey); err != nil {
 					return 0, 0, 0, errors.Wrapf(err,
 						"replacing point tombstones with range tombstones for write intent range at %s on end transaction [%s]",

pkg/kv/kvserver/batcheval/cmd_resolve_intent.go

@@ -129,7 +129,7 @@ func ResolveIntent(
 	// If requested, replace point tombstones with range tombstones.
 	if ok && cArgs.EvalCtx.EvalKnobs().UseRangeTombstonesForPointDeletes {
 		if err := storage.ReplacePointTombstonesWithRangeTombstones(ctx,
-			spanset.DisableReadWriterAssertions(readWriter), ms, update.Key, update.EndKey); err != nil {
+			spanset.DisableLatchAssertions(readWriter), ms, update.Key, update.EndKey); err != nil {
 			return result.Result{}, err
 		}
 	}

pkg/kv/kvserver/batcheval/cmd_resolve_intent_range.go

@@ -99,7 +99,7 @@ func ResolveIntentRange(
 	// If requested, replace point tombstones with range tombstones.
 	if cArgs.EvalCtx.EvalKnobs().UseRangeTombstonesForPointDeletes {
 		if err := storage.ReplacePointTombstonesWithRangeTombstones(ctx,
-			spanset.DisableReadWriterAssertions(readWriter), ms, args.Key, args.EndKey); err != nil {
+			spanset.DisableLatchAssertions(readWriter), ms, args.Key, args.EndKey); err != nil {
 			return result.Result{}, err
 		}
 	}

pkg/kv/kvserver/spanset/batch.go

@@ -888,6 +888,64 @@ func DisableReadWriterAssertions(rw storage.ReadWriter) storage.ReadWriter {
 	}
 }
 
+// DisableLatchAssertions returns a new batch wrapper with latch assertions
+// disabled. It does not modify the original batch. The returned batch shares
+// the same underlying storage.Batch but has its own SpanSet copies with
+// assertions disabled.
+func DisableLatchAssertions(rw storage.ReadWriter) storage.ReadWriter {
+	switch v := rw.(type) {
+	case *spanSetBatch:
+		// Clone the span sets and disable assertions on the copies
+		readSpans := v.ReadWriter.spanSetReader.spans.Copy()
+		writeSpans := v.ReadWriter.spanSetWriter.spans.Copy()
+		readSpans.DisableUndeclaredAccessAssertions()
+		writeSpans.DisableUndeclaredAccessAssertions()
+
+		// Create a new spanSetBatch with the modified span sets.
+		return &spanSetBatch{
+			ReadWriter: ReadWriter{
+				spanSetReader: spanSetReader{r: v.b, spans: readSpans, spansOnly: v.spansOnly, ts: v.ts},
+				spanSetWriter: spanSetWriter{w: v.b, spans: writeSpans, spansOnly: v.spansOnly, ts: v.ts},
+			},
+			b:         v.b,
+			spans:     readSpans, // Use the read spans as the main spans
+			spansOnly: v.spansOnly,
+			ts:        v.ts,
+		}
+	default:
+		return rw
+	}
+}
+
+// DisableForbiddenSpanAssertionsOnBatch returns a new batch wrapper with
+// forbidden span assertions disabled. It does not modify the original batch.
+// The returned batch shares the same underlying storage.Batch but has its own
+// SpanSet copies with assertions disabled.
+func DisableForbiddenSpanAssertionsOnBatch(rw storage.ReadWriter) storage.ReadWriter {
+	switch v := rw.(type) {
+	case *spanSetBatch:
+		// Clone the span sets and disable assertions on the copies
+		readSpans := v.ReadWriter.spanSetReader.spans.Copy()
+		writeSpans := v.ReadWriter.spanSetWriter.spans.Copy()
+		readSpans.DisableForbiddenSpansAssertions()
+		writeSpans.DisableForbiddenSpansAssertions()
+
+		// Create a new spanSetBatch with the modified span sets.
+		return &spanSetBatch{
+			ReadWriter: ReadWriter{
+				spanSetReader: spanSetReader{r: v.b, spans: readSpans, spansOnly: v.spansOnly, ts: v.ts},
+				spanSetWriter: spanSetWriter{w: v.b, spans: writeSpans, spansOnly: v.spansOnly, ts: v.ts},
+			},
+			b:         v.b,
+			spans:     readSpans, // Use the read spans as the main spans
+			spansOnly: v.spansOnly,
+			ts:        v.ts,
+		}
+	default:
+		return rw
+	}
+}
+
 // addLockTableSpans adds corresponding lock table spans for the declared
 // spans. This is to implicitly allow raw access to separated intents in the
 // lock table for any declared keys. Explicitly declaring lock table spans is

pkg/kv/kvserver/spanset/spanset.go

@@ -83,8 +83,14 @@ type Span struct {
 // The Span slice for a particular access and scope contains non-overlapping
 // spans in increasing key order after calls to SortAndDedup.
 type SpanSet struct {
-	spans           [NumSpanAccess][NumSpanScope][]Span
-	allowUndeclared bool
+	spans [NumSpanAccess][NumSpanScope][]Span
+	// forbiddenSpansMatchers are functions that return an error if the given span
+	// shouldn't be accessed (forbidden). This allows for complex pattern matching
+	// like forbidding specific keys across all range IDs without enumerating them
+	// explicitly.
+	forbiddenSpansMatchers []func(roachpb.Span) error
+	allowUndeclared        bool
+	allowForbidden         bool
 }
 
 var spanSetPool = sync.Pool{
@@ -113,6 +119,8 @@ func (s *SpanSet) Release() {
 			s.spans[sa][ss] = recycle
 		}
 	}
+	s.forbiddenSpansMatchers = nil
+	s.allowForbidden = false
 	s.allowUndeclared = false
 	spanSetPool.Put(s)
 }
@@ -155,7 +163,9 @@ func (s *SpanSet) Copy() *SpanSet {
 			n.spans[sa][ss] = append(n.spans[sa][ss], s.spans[sa][ss]...)
 		}
 	}
+	n.forbiddenSpansMatchers = append(n.forbiddenSpansMatchers, s.forbiddenSpansMatchers...)
 	n.allowUndeclared = s.allowUndeclared
+	n.allowForbidden = s.allowForbidden
 	return n
 }
 
@@ -201,14 +211,22 @@ func (s *SpanSet) AddMVCC(access SpanAccess, span roachpb.Span, timestamp hlc.Ti
 	s.spans[access][scope] = append(s.spans[access][scope], Span{Span: span, Timestamp: timestamp})
 }
 
+// AddForbiddenMatcher adds a forbidden span matcher. The matcher is a function
+// that is called for each span access to check if it should be forbidden.
+func (s *SpanSet) AddForbiddenMatcher(matcher func(roachpb.Span) error) {
+	s.forbiddenSpansMatchers = append(s.forbiddenSpansMatchers, matcher)
+}
+
 // Merge merges all spans in s2 into s. s2 is not modified.
 func (s *SpanSet) Merge(s2 *SpanSet) {
 	for sa := SpanAccess(0); sa < NumSpanAccess; sa++ {
 		for ss := SpanScope(0); ss < NumSpanScope; ss++ {
 			s.spans[sa][ss] = append(s.spans[sa][ss], s2.spans[sa][ss]...)
 		}
 	}
+	s.forbiddenSpansMatchers = append(s.forbiddenSpansMatchers, s2.forbiddenSpansMatchers...)
 	s.allowUndeclared = s2.allowUndeclared
+	s.allowForbidden = s2.allowForbidden
 	s.SortAndDedup()
 }
 
@@ -331,27 +349,37 @@ func (s *SpanSet) CheckAllowedAt(
 func (s *SpanSet) checkAllowed(
 	access SpanAccess, span roachpb.Span, check func(SpanAccess, Span) bool,
 ) error {
-	if s.allowUndeclared {
-		// If the request has specified that undeclared spans are allowed, do
-		// nothing.
-		return nil
+	// Unless explicitly disabled, check if we access any forbidden spans.
+	if !s.allowForbidden {
+		// Check if the span is forbidden.
+		for _, matcher := range s.forbiddenSpansMatchers {
+			if err := matcher(span); err != nil {
+				return errors.Errorf("cannot %s span %s: matches forbidden pattern",
+					access, span)
+			}
+		}
 	}
 
-	scope := SpanGlobal
-	if (span.Key != nil && keys.IsLocal(span.Key)) ||
-		(span.EndKey != nil && keys.IsLocal(span.EndKey)) {
-		scope = SpanLocal
-	}
+	// Unless explicitly disabled, check if we access any undeclared spans.
+	if !s.allowUndeclared {
+		scope := SpanGlobal
+		if (span.Key != nil && keys.IsLocal(span.Key)) ||
+			(span.EndKey != nil && keys.IsLocal(span.EndKey)) {
+			scope = SpanLocal
+		}
 
-	for ac := access; ac < NumSpanAccess; ac++ {
-		for _, cur := range s.spans[ac][scope] {
-			if contains(cur.Span, span) && check(ac, cur) {
-				return nil
+		for ac := access; ac < NumSpanAccess; ac++ {
+			for _, cur := range s.spans[ac][scope] {
+				if contains(cur.Span, span) && check(ac, cur) {
+					return nil
+				}
 			}
 		}
+
+		return errors.Errorf("cannot %s undeclared span %s\ndeclared:\n%s\nstack:\n%s", access, span, s, debugutil.Stack())
 	}
 
-	return errors.Errorf("cannot %s undeclared span %s\ndeclared:\n%s\nstack:\n%s", access, span, s, debugutil.Stack())
+	return nil
 }
 
 // contains returns whether s1 contains s2. Unlike Span.Contains, this function
@@ -396,3 +424,8 @@ func (s *SpanSet) Validate() error {
 func (s *SpanSet) DisableUndeclaredAccessAssertions() {
 	s.allowUndeclared = true
 }
+
+// DisableForbiddenAssertions disables forbidden spans assertions.
+func (s *SpanSet) DisableForbiddenSpansAssertions() {
+	s.allowForbidden = true
+}