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

@@ -766,7 +766,9 @@ func NewReadWriterAt(rw storage.ReadWriter, spans *SpanSet, ts hlc.Timestamp) st
 
 type spanSetBatch struct {
 	ReadWriter
-	b     storage.Batch
+	b storage.Batch
+	// TODO(ibrahim): The fields spans, spansOnly, and ts don't seem to be used.
+	// Consider removing them and performing the necessary clean ups.
 	spans *SpanSet
 
 	spansOnly bool
@@ -838,6 +840,22 @@ func (s spanSetBatch) ClearRawEncodedRange(start, end []byte) error {
 	return s.b.ClearRawEncodedRange(start, end)
 }
 
+// clone returns a shallow copy of the spanSetBatch. The returned batch shares
+// the same underlying storage.Batch but has its own spanSetBatch wrapper with
+// a new copy of the SpanSet that uses a shallow copy of the underlying spans.
+func (s spanSetBatch) clone() *spanSetBatch {
+	return &spanSetBatch{
+		ReadWriter: ReadWriter{
+			spanSetReader: spanSetReader{r: s.b, spans: s.spanSetReader.spans.ShallowCopy(), spansOnly: s.spansOnly, ts: s.ts},
+			spanSetWriter: spanSetWriter{w: s.b, spans: s.spanSetWriter.spans.ShallowCopy(), spansOnly: s.spansOnly, ts: s.ts},
+		},
+		b:         s.b,
+		spans:     s.spans.ShallowCopy(),
+		spansOnly: s.spansOnly,
+		ts:        s.ts,
+	}
+}
+
 // NewBatch returns a storage.Batch that asserts access of the underlying
 // Batch against the given SpanSet. We only consider span boundaries, associated
 // timestamps are not considered.
@@ -888,6 +906,40 @@ 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 wrapper with the
+// latch assertion disabled.
+func DisableLatchAssertions(rw storage.ReadWriter) storage.ReadWriter {
+	switch v := rw.(type) {
+	case *spanSetBatch:
+		newSnapSetBatch := v.clone()
+		newSnapSetBatch.spanSetReader.spans.DisableUndeclaredAccessAssertions()
+		newSnapSetBatch.spanSetWriter.spans.DisableUndeclaredAccessAssertions()
+		return newSnapSetBatch
+
+	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 wrapper with the forbidden span assertion disabled.
+func DisableForbiddenSpanAssertionsOnBatch(rw storage.ReadWriter) storage.ReadWriter {
+	switch v := rw.(type) {
+	case *spanSetBatch:
+		newSnapSetBatch := v.clone()
+		newSnapSetBatch.spanSetReader.spans.DisableForbiddenSpansAssertions()
+		newSnapSetBatch.spanSetWriter.spans.DisableForbiddenSpansAssertions()
+		return newSnapSetBatch
+
+	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,19 @@ 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
+}
+
+// ShallowCopy performs a shallow SpanSet copy.
+func (s *SpanSet) ShallowCopy() *SpanSet {
+	n := New()
+	n.spans = s.spans
+	n.forbiddenSpansMatchers = s.forbiddenSpansMatchers
+	n.allowUndeclared = s.allowUndeclared
+	n.allowForbidden = s.allowForbidden
 	return n
 }
 
@@ -201,14 +221,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 +359,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
 	}
 
@@ -396,3 +434,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
+}