Merge #159847

159847: kvserver: add spanSetWriteBatch to spanset package r=iskettaneh a=iskettaneh

This PR adds spanSetWriteBatch to the spanset package and asserts on accesses. This will be used later by an engine wrapper to assert on the engine accesses.

References: #158281

Release notes: None

Co-authored-by: iskettaneh <[email protected]>

Author: craig[bot]

pkg/kv/kvserver/spanset/batch.go

@@ -16,6 +16,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/storage/fs"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
+	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/pebble"
 	"github.com/cockroachdb/pebble/rangekey"
 )
@@ -764,83 +765,155 @@ func NewReader(r storage.Reader, spans *SpanSet, ts hlc.Timestamp) storage.Reade
 	return spanSetReader{r: r, spans: spans, ts: ts}
 }
 
+// NewReadWriter returns a storage.ReadWriter that asserts access of the
+// underlying ReadWriter against the given SpanSet.
+//
+// NewReadWriter clones and does not retain the provided span set.
+func NewReadWriter(rw storage.ReadWriter, spans *SpanSet) storage.ReadWriter {
+	return makeSpanSetReadWriter(rw, spans)
+}
+
 // NewReadWriterAt returns a storage.ReadWriter that asserts access of the
 // underlying ReadWriter against the given SpanSet at a given timestamp.
-// If zero timestamp is provided, accesses are considered non-MVCC.
 //
 // NewReadWriterAt clones and does not retain the provided span set.
 func NewReadWriterAt(rw storage.ReadWriter, spans *SpanSet, ts hlc.Timestamp) storage.ReadWriter {
 	return makeSpanSetReadWriterAt(rw, spans, ts)
 }
 
-type spanSetBatch struct {
-	ReadWriter
-	b storage.Batch
-	// TODO(ibrahim): The fields spans, spansOnly, and ts don't seem to be used.
-	// Consider removing or marking them as intended.
-	spans *SpanSet
+// spanSetWriteBatch wraps a storage.WriteBatch and adds span checking.
+type spanSetWriteBatch struct {
+	spanSetWriter
+	wb storage.WriteBatch
+}
 
-	spansOnly bool
-	ts        hlc.Timestamp
+var _ storage.WriteBatch = (*spanSetWriteBatch)(nil)
+
+// ClearRawEncodedRange implements storage.InternalWriter.
+func (s spanSetWriteBatch) ClearRawEncodedRange(start, end []byte) error {
+	// Decode the engine keys to check spans.
+	startKey, ok := storage.DecodeEngineKey(start)
+	if !ok {
+		return errors.Errorf("cannot decode start engine key")
+	}
+	endKey, ok := storage.DecodeEngineKey(end)
+	if !ok {
+		return errors.Errorf("cannot decode end engine key")
+	}
+	if err := s.spanSetWriter.checkAllowedRange(startKey.Key, endKey.Key); err != nil {
+		return err
+	}
+	return s.wb.ClearRawEncodedRange(start, end)
 }
 
-var _ storage.Batch = spanSetBatch{}
+// PutInternalRangeKey implements storage.InternalWriter.
+func (s spanSetWriteBatch) PutInternalRangeKey(start, end []byte, key rangekey.Key) error {
+	// Decode the engine keys to check spans.
+	startKey, ok := storage.DecodeEngineKey(start)
+	if !ok {
+		return errors.Errorf("cannot decode start engine key")
+	}
+	endKey, ok := storage.DecodeEngineKey(end)
+	if !ok {
+		return errors.Errorf("cannot decode end engine key")
+	}
+	if err := s.spanSetWriter.checkAllowedRange(startKey.Key, endKey.Key); err != nil {
+		return err
+	}
+	return s.wb.PutInternalRangeKey(start, end, key)
+}
 
-func (s spanSetBatch) NewBatchOnlyMVCCIterator(
-	ctx context.Context, opts storage.IterOptions,
-) (storage.MVCCIterator, error) {
-	mvccIter, err := s.b.NewBatchOnlyMVCCIterator(ctx, opts)
-	if err != nil {
-		return nil, err
+// PutInternalPointKey implements storage.InternalWriter.
+func (s spanSetWriteBatch) PutInternalPointKey(key *pebble.InternalKey, value []byte) error {
+	// Decode the internal key to check spans.
+	engKey, ok := storage.DecodeEngineKey(key.UserKey)
+	if !ok {
+		return errors.Errorf("cannot decode engine key")
 	}
-	if s.spansOnly {
-		return NewIterator(mvccIter, s.spans), nil
+
+	if err := s.spanSetWriter.checkAllowed(engKey.Key); err != nil {
+		return err
 	}
-	return NewIteratorAt(mvccIter, s.spans, s.ts), nil
+
+	return s.wb.PutInternalPointKey(key, value)
 }
 
-func (s spanSetBatch) Commit(sync bool) error {
-	return s.b.Commit(sync)
+// Close implements storage.WriteBatch.
+func (s spanSetWriteBatch) Close() {
+	s.wb.Close()
 }
 
-func (s spanSetBatch) CommitNoSyncWait() error {
-	return s.b.CommitNoSyncWait()
+// Commit implements storage.WriteBatch.
+func (s spanSetWriteBatch) Commit(sync bool) error {
+	return s.wb.Commit(sync)
 }
 
-func (s spanSetBatch) SyncWait() error {
-	return s.b.SyncWait()
+// CommitNoSyncWait implements storage.WriteBatch.
+func (s spanSetWriteBatch) CommitNoSyncWait() error {
+	return s.wb.CommitNoSyncWait()
 }
 
-func (s spanSetBatch) Empty() bool {
-	return s.b.Empty()
+// SyncWait implements storage.WriteBatch.
+func (s spanSetWriteBatch) SyncWait() error {
+	return s.wb.SyncWait()
 }
 
-func (s spanSetBatch) Count() uint32 {
-	return s.b.Count()
+// Empty implements storage.WriteBatch.
+func (s spanSetWriteBatch) Empty() bool {
+	return s.wb.Empty()
 }
 
-func (s spanSetBatch) Len() int {
-	return s.b.Len()
+// Count implements storage.WriteBatch.
+func (s spanSetWriteBatch) Count() uint32 {
+	return s.wb.Count()
 }
 
-func (s spanSetBatch) Repr() []byte {
-	return s.b.Repr()
+// Len implements storage.WriteBatch.
+func (s spanSetWriteBatch) Len() int {
+	return s.wb.Len()
 }
 
-func (s spanSetBatch) CommitStats() storage.BatchCommitStats {
-	return s.b.CommitStats()
+// Repr implements storage.WriteBatch.
+func (s spanSetWriteBatch) Repr() []byte {
+	return s.wb.Repr()
 }
 
-func (s spanSetBatch) PutInternalRangeKey(start, end []byte, key rangekey.Key) error {
-	return s.b.PutInternalRangeKey(start, end, key)
+// CommitStats implements storage.WriteBatch.
+func (s spanSetWriteBatch) CommitStats() storage.BatchCommitStats {
+	return s.wb.CommitStats()
 }
 
-func (s spanSetBatch) PutInternalPointKey(key *pebble.InternalKey, value []byte) error {
-	return s.b.PutInternalPointKey(key, value)
+type spanSetBatch struct {
+	spanSetReader
+	spanSetWriteBatch
+	b storage.Batch
+	// TODO(ibrahim): The fields spans, spansOnly, and ts don't seem to be used.
+	// Consider removing or marking them as intended.
+	spans *SpanSet
+
+	spansOnly bool
+	ts        hlc.Timestamp
 }
 
-func (s spanSetBatch) ClearRawEncodedRange(start, end []byte) error {
-	return s.b.ClearRawEncodedRange(start, end)
+var _ storage.Batch = spanSetBatch{}
+
+// Close implements storage.Batch (resolves ambiguity between Reader.Close
+// and WriteBatch.Close).
+func (s spanSetBatch) Close() {
+	s.wb.Close()
+}
+
+func (s spanSetBatch) NewBatchOnlyMVCCIterator(
+	ctx context.Context, opts storage.IterOptions,
+) (storage.MVCCIterator, error) {
+	mvccIter, err := s.b.NewBatchOnlyMVCCIterator(ctx, opts)
+	if err != nil {
+		return nil, err
+	}
+	if s.spansOnly {
+		return NewIterator(mvccIter, s.spans), nil
+	}
+	return NewIteratorAt(mvccIter, s.spans, s.ts), nil
 }
 
 // shallowCopy returns a shallow copy of the spanSetBatch. The returned batch
@@ -849,33 +922,43 @@ func (s spanSetBatch) ClearRawEncodedRange(start, end []byte) error {
 // spans.
 func (s spanSetBatch) shallowCopy() *spanSetBatch {
 	b := s
-	b.spanSetReader.spans = b.spanSetReader.spans.ShallowCopy()
-	b.spanSetWriter.spans = b.spanSetWriter.spans.ShallowCopy()
 	b.spans = b.spans.ShallowCopy()
+	b.spanSetReader.spans = b.spans
+	b.spanSetWriter.spans = b.spans
 	return &b
 }
 
 // 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.
 func NewBatch(b storage.Batch, spans *SpanSet) storage.Batch {
+	spans = addLockTableSpans(spans)
 	return &spanSetBatch{
-		ReadWriter: makeSpanSetReadWriter(b, spans),
-		b:          b,
-		spans:      spans,
-		spansOnly:  true,
+		spanSetReader: spanSetReader{r: b, spans: spans, spansOnly: true},
+		spanSetWriteBatch: spanSetWriteBatch{
+			spanSetWriter: spanSetWriter{w: b, spans: spans, spansOnly: true},
+			wb:            b,
+		},
+		b:         b,
+		spans:     spans,
+		spansOnly: true,
 	}
 }
 
 // NewBatchAt returns an storage.Batch that asserts access of the underlying
 // Batch against the given SpanSet at the given timestamp.
 // If the zero timestamp is used, all accesses are considered non-MVCC.
 func NewBatchAt(b storage.Batch, spans *SpanSet, ts hlc.Timestamp) storage.Batch {
+	spans = addLockTableSpans(spans)
 	return &spanSetBatch{
-		ReadWriter: makeSpanSetReadWriterAt(b, spans, ts),
-		b:          b,
-		spans:      spans,
-		ts:         ts,
+		spanSetReader: spanSetReader{r: b, spans: spans, ts: ts},
+		spanSetWriteBatch: spanSetWriteBatch{
+			spanSetWriter: spanSetWriter{w: b, spans: spans, ts: ts},
+			wb:            b,
+		},
+		b:     b,
+		spans: spans,
+		ts:    ts,
 	}
 }
 
@@ -884,9 +967,9 @@ func NewBatchAt(b storage.Batch, spans *SpanSet, ts hlc.Timestamp) storage.Batch
 func DisableReaderAssertions(reader storage.Reader) storage.Reader {
 	switch v := reader.(type) {
 	case ReadWriter:
-		return DisableReaderAssertions(v.r)
+		return DisableReaderAssertions(v.spanSetReader.r)
 	case *spanSetBatch:
-		return DisableReaderAssertions(v.r)
+		return DisableReaderAssertions(v.spanSetReader.r)
 	default:
 		return reader
 	}
@@ -897,9 +980,9 @@ func DisableReaderAssertions(reader storage.Reader) storage.Reader {
 func DisableReadWriterAssertions(rw storage.ReadWriter) storage.ReadWriter {
 	switch v := rw.(type) {
 	case ReadWriter:
-		return DisableReadWriterAssertions(v.w.(storage.ReadWriter))
+		return DisableReadWriterAssertions(v.spanSetWriter.w.(storage.ReadWriter))
 	case *spanSetBatch:
-		return DisableReadWriterAssertions(v.w.(storage.ReadWriter))
+		return DisableReadWriterAssertions(v.spanSetWriteBatch.spanSetWriter.w.(storage.ReadWriter))
 	default:
 		return rw
 	}
@@ -913,10 +996,8 @@ func DisableUndeclaredSpanAssertions(rw storage.ReadWriter) storage.ReadWriter {
 	switch v := rw.(type) {
 	case *spanSetBatch:
 		newSnapSetBatch := v.shallowCopy()
-		newSnapSetBatch.spanSetReader.spans.DisableUndeclaredAccessAssertions()
-		newSnapSetBatch.spanSetWriter.spans.DisableUndeclaredAccessAssertions()
+		newSnapSetBatch.spans.DisableUndeclaredAccessAssertions()
 		return newSnapSetBatch
-
 	default:
 		return rw
 	}
@@ -932,10 +1013,8 @@ func DisableForbiddenSpanAssertions(rw storage.ReadWriter) storage.ReadWriter {
 	switch v := rw.(type) {
 	case *spanSetBatch:
 		newSnapSetBatch := v.shallowCopy()
-		newSnapSetBatch.spanSetReader.spans.DisableForbiddenSpansAssertions()
-		newSnapSetBatch.spanSetWriter.spans.DisableForbiddenSpansAssertions()
+		newSnapSetBatch.spans.DisableForbiddenSpansAssertions()
 		return newSnapSetBatch
-
 	default:
 		return rw
 	}