kvserver: integrate spanSetWriteBatch into SpansetBatch

This commit integrates spanSetWriteBatch into SpansetBatch,
which will allow a future PR to assert all engine accesses
by wrapping Reader, Writer, ReadWriter, Batch with the
spanset type that does the assertions.

Author: iskettaneh

pkg/kv/kvserver/spanset/batch.go

@@ -765,9 +765,16 @@ 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 {
@@ -877,7 +884,8 @@ func (s spanSetWriteBatch) CommitStats() storage.BatchCommitStats {
 }
 
 type spanSetBatch struct {
-	ReadWriter
+	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.
@@ -889,6 +897,12 @@ type spanSetBatch struct {
 
 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) {
@@ -902,83 +916,49 @@ func (s spanSetBatch) NewBatchOnlyMVCCIterator(
 	return NewIteratorAt(mvccIter, s.spans, s.ts), nil
 }
 
-func (s spanSetBatch) Commit(sync bool) error {
-	return s.b.Commit(sync)
-}
-
-func (s spanSetBatch) CommitNoSyncWait() error {
-	return s.b.CommitNoSyncWait()
-}
-
-func (s spanSetBatch) SyncWait() error {
-	return s.b.SyncWait()
-}
-
-func (s spanSetBatch) Empty() bool {
-	return s.b.Empty()
-}
-
-func (s spanSetBatch) Count() uint32 {
-	return s.b.Count()
-}
-
-func (s spanSetBatch) Len() int {
-	return s.b.Len()
-}
-
-func (s spanSetBatch) Repr() []byte {
-	return s.b.Repr()
-}
-
-func (s spanSetBatch) CommitStats() storage.BatchCommitStats {
-	return s.b.CommitStats()
-}
-
-func (s spanSetBatch) PutInternalRangeKey(start, end []byte, key rangekey.Key) error {
-	return s.b.PutInternalRangeKey(start, end, key)
-}
-
-func (s spanSetBatch) PutInternalPointKey(key *pebble.InternalKey, value []byte) error {
-	return s.b.PutInternalPointKey(key, value)
-}
-
-func (s spanSetBatch) ClearRawEncodedRange(start, end []byte) error {
-	return s.b.ClearRawEncodedRange(start, end)
-}
-
 // shallowCopy 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) 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,
 	}
 }
 
@@ -987,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
 	}
@@ -1000,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
 	}
@@ -1016,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
 	}
@@ -1035,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
 	}