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.

Author: iskettaneh

pkg/kv/kvserver/batcheval/cmd_add_sstable.go

@@ -337,7 +337,7 @@ func EvalAddSSTable(
 	// addition, and instead just use this key-only iterator. If a caller actually
 	// needs to know what data is there, it must issue its own real Scan.
 	if args.ReturnFollowingLikelyNonEmptySpanStart {
-		existingIter, err := spanset.DisableReaderAssertions(readWriter).NewMVCCIterator(
+		existingIter, err := spanset.DisableLatchAssertions(readWriter).NewMVCCIterator(
 			ctx,
 			storage.MVCCKeyIterKind, // don't care if it is committed or not, just that it isn't empty.
 			storage.IterOptions{

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,59 @@ func DisableReadWriterAssertions(rw storage.ReadWriter) storage.ReadWriter {
 	}
 }
 
+func disableAssertionHelper(
+	rw storage.ReadWriter, disableAssertion func(spans *SpanSet),
+) storage.ReadWriter {
+	switch v := rw.(type) {
+	case *spanSetBatch:
+		// Create a SpanSet wrapper that will be used to disable a specific
+		// assertion ephemerally.
+		spans := &SpanSet{
+			spans:                  v.ReadWriter.spanSetReader.spans.spans,
+			forbiddenSpansMatchers: v.ReadWriter.spanSetReader.spans.forbiddenSpansMatchers,
+			allowForbidden:         v.ReadWriter.spanSetReader.spans.allowForbidden,
+			allowUndeclared:        v.ReadWriter.spanSetReader.spans.allowUndeclared,
+		}
+
+		// Call the function that will disable some assertion.
+		disableAssertion(spans)
+
+		// Create a new spanSetBatch with the modified span sets.
+		return &spanSetBatch{
+			ReadWriter: ReadWriter{
+				spanSetReader: spanSetReader{r: v.b, spans: spans, spansOnly: v.spansOnly, ts: v.ts},
+				spanSetWriter: spanSetWriter{w: v.b, spans: spans, spansOnly: v.spansOnly, ts: v.ts},
+			},
+			b:         v.b,
+			spans:     spans,
+			spansOnly: v.spansOnly,
+			ts:        v.ts,
+		}
+	default:
+		return rw
+	}
+}
+
+// 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 wrapper with the
+// latch assertion disabled.
+func DisableLatchAssertions(rw storage.ReadWriter) storage.ReadWriter {
+	return disableAssertionHelper(rw, func(spans *SpanSet) {
+		spans.DisableUndeclaredAccessAssertions()
+	})
+}
+
+// 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 wrapper with the forbidden span assertion disabled.
+func DisableForbiddenSpanAssertionsOnBatch(rw storage.ReadWriter) storage.ReadWriter {
+	return disableAssertionHelper(rw, func(spans *SpanSet) {
+		spans.DisableForbiddenSpansAssertions()
+	})
+}
+
 // 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,9 +349,19 @@ func (s *SpanSet) CheckAllowedAt(
 func (s *SpanSet) checkAllowed(
 	access SpanAccess, span roachpb.Span, check func(SpanAccess, Span) bool,
 ) error {
+	// 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)
+			}
+		}
+	}
+
+	// Unless explicitly disabled, check if we access any undeclared spans.
 	if s.allowUndeclared {
-		// If the request has specified that undeclared spans are allowed, do
-		// nothing.
 		return nil
 	}
 
@@ -352,6 +380,8 @@ func (s *SpanSet) checkAllowed(
 	}
 
 	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 +426,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
+}