kvcoord: add request validation to txnWriteBuffer

This adds early detection of unsupported request types and request
options. SendLocked now errors on any request that the txnWriteBuffer
doesn't explicitly support. We may consider flushing the buffer
instead to increase the set of transactions that can use this feature.

Fixes #143894

Release note: None

Author: stevendanna

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -20,6 +20,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
 	"github.com/cockroachdb/cockroach/pkg/storage/mvccencoding"
 	"github.com/cockroachdb/cockroach/pkg/storage/mvcceval"
+	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/cockroach/pkg/util/humanizeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/errors"
@@ -46,8 +47,8 @@ var bufferedWritesMaxBufferSize = settings.RegisterIntSetting(
 
 // txnWriteBuffer is a txnInterceptor that buffers transactional writes until
 // commit time. Moreover, it also decomposes read-write KV operations (e.g.
-// CPuts, InitPuts) into separate (locking) read and write operations, buffering
-// the latter until commit time.
+// CPuts) into separate (locking) read and write operations, buffering the
+// latter until commit time.
 //
 // Buffering writes until commit time has four main benefits:
 //
@@ -253,7 +254,16 @@ func (twb *txnWriteBuffer) SendLocked(
 		return twb.flushBufferAndSendBatch(ctx, ba)
 	}
 
-	transformedBa, ts := twb.applyTransformations(ctx, ba)
+	if err := twb.validateBatch(ba); err != nil {
+		// We could choose to twb.flushBufferAndSendBatch
+		// here. For now, we return an error.
+		return nil, kvpb.NewError(err)
+	}
+
+	transformedBa, ts, pErr := twb.applyTransformations(ctx, ba)
+	if pErr != nil {
+		return nil, pErr
+	}
 
 	if len(transformedBa.Requests) == 0 {
 		// Lower layers (the DistSender and the KVServer) do not expect/handle empty
@@ -278,6 +288,80 @@ func (twb *txnWriteBuffer) SendLocked(
 	return twb.mergeResponseWithTransformations(ctx, ts, br)
 }
 
+// validateBatch returns an error if the batch is unsupported
+// by the txnWriteBuffer.
+func (twb *txnWriteBuffer) validateBatch(ba *kvpb.BatchRequest) error {
+	if ba.WriteOptions != nil {
+		// OriginTimestamp and OriginID are currently only used by Logical Data
+		// Replication (LDR). These options are unsupported at the moment as we
+		// don't store the inbound batch options in the buffer.
+		if ba.WriteOptions.OriginTimestamp.IsSet() {
+			return errors.AssertionFailedf("transaction write buffer does not support batches with OriginTimestamp set")
+		}
+		if ba.WriteOptions.OriginID != 0 {
+			return errors.AssertionFailedf("transaction write buffer does not support batches with OriginID set")
+		}
+	}
+	return twb.validateRequests(ba)
+}
+
+// validateRequests returns an error if any of the requests in the batch
+// are unsupported by the txnWriteBuffer.
+func (twb *txnWriteBuffer) validateRequests(ba *kvpb.BatchRequest) error {
+	for _, ru := range ba.Requests {
+		req := ru.GetInner()
+		switch t := req.(type) {
+		case *kvpb.ConditionalPutRequest:
+			// Our client side ConditionalPutRequest evaluation does not know how to
+			// handle the origin timestamp setting. Doing so would require sending a
+			// GetRequest with RawMVCCValues set and parsing the MVCCValueHeader.
+			if t.OriginTimestamp.IsSet() {
+				return unsupportedOptionError(t.Method(), "OriginTimestamp")
+			}
+		case *kvpb.PutRequest:
+		case *kvpb.DeleteRequest:
+		case *kvpb.GetRequest:
+			// ReturnRawMVCCValues is unsupported because we don't know how to serve
+			// such reads from the write buffer currently.
+			if t.ReturnRawMVCCValues {
+				return unsupportedOptionError(t.Method(), "ReturnRawMVCCValue")
+			}
+		case *kvpb.ScanRequest:
+			// ReturnRawMVCCValues is unsupported because we don't know how to serve
+			// such reads from the write buffer currently.
+			if t.ReturnRawMVCCValues {
+				return unsupportedOptionError(t.Method(), "ReturnRawMVCCValue")
+			}
+			if t.ScanFormat == kvpb.COL_BATCH_RESPONSE {
+				return unsupportedOptionError(t.Method(), "COL_BATCH_RESPONSE scan format")
+			}
+		case *kvpb.ReverseScanRequest:
+			// ReturnRawMVCCValues is unsupported because we don't know how to serve
+			// such reads from the write buffer currently.
+			if t.ReturnRawMVCCValues {
+				return unsupportedOptionError(t.Method(), "ReturnRawMVCCValue")
+			}
+			if t.ScanFormat == kvpb.COL_BATCH_RESPONSE {
+				return unsupportedOptionError(t.Method(), "COL_BATCH_RESPONSE scan format")
+			}
+		default:
+			// All other requests are unsupported. Note that we assume EndTxn and
+			// DeleteRange requests were handled explicitly before this method was
+			// called.
+			return unsupportedMethodError(t.Method())
+		}
+	}
+	return nil
+}
+
+func unsupportedMethodError(m kvpb.Method) error {
+	return errors.AssertionFailedf("transaction write buffer does not support %s requests", m)
+}
+
+func unsupportedOptionError(m kvpb.Method, option string) error {
+	return errors.AssertionFailedf("transaction write buffer does not support %s requests with %s", m, option)
+}
+
 // estimateSize returns a conservative estimate by which the buffer will grow in
 // size if the writes from the supplied batch request are buffered.
 func (twb *txnWriteBuffer) estimateSize(ba *kvpb.BatchRequest) int64 {
@@ -539,7 +623,7 @@ func (twb *txnWriteBuffer) closeLocked() {}
 // TODO(arul): Augment this comment as these expand.
 func (twb *txnWriteBuffer) applyTransformations(
 	ctx context.Context, ba *kvpb.BatchRequest,
-) (*kvpb.BatchRequest, transformations) {
+) (*kvpb.BatchRequest, transformations, *kvpb.Error) {
 	baRemote := ba.ShallowCopy()
 	// TODO(arul): We could improve performance here by pre-allocating
 	// baRemote.Requests to the correct size by counting the number of Puts/Dels
@@ -699,9 +783,6 @@ func (twb *txnWriteBuffer) applyTransformations(
 				// We've constructed a response that we'll stitch together with the
 				// result on the response path; eschew sending the request to the KV
 				// layer.
-				//
-				// TODO(arul): if the ReturnRawMVCCValues flag is set, we'll need to
-				// flush the buffer.
 				continue
 			}
 			// Wasn't served locally; send the request to the KV layer.
@@ -736,10 +817,10 @@ func (twb *txnWriteBuffer) applyTransformations(
 			baRemote.Requests = append(baRemote.Requests, ru)
 
 		default:
-			baRemote.Requests = append(baRemote.Requests, ru)
+			return nil, nil, kvpb.NewError(unsupportedMethodError(t.Method()))
 		}
 	}
-	return baRemote, ts
+	return baRemote, ts, nil
 }
 
 // seekItemForSpan returns a bufferedWrite appropriate for use with a
@@ -1103,9 +1184,7 @@ func (t transformation) toResp(
 		ru.MustSetInner(reverseScanResp)
 
 	default:
-		// This is only possible once we start decomposing read-write requests into
-		// separate bits.
-		panic("unimplemented")
+		return ru, kvpb.NewError(unsupportedMethodError(req.Method()))
 	}
 
 	return ru, nil
@@ -1725,7 +1804,7 @@ func (m *respMerger) toReverseScanResp(
 
 // assertTrue panics with a message if the supplied condition isn't true.
 func assertTrue(cond bool, msg string) {
-	if !cond {
+	if !cond && buildutil.CrdbTestBuild {
 		panic(msg)
 	}
 }

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer_test.go

@@ -1851,3 +1851,139 @@ func TestTxnWriteBufferClearsBufferOnEpochBump(t *testing.T) {
 	require.Equal(t, 0, int(twb.bufferSize))
 	require.Equal(t, numCalled, mockSender.NumCalled())
 }
+
+// TestTxnWriteBufferBatchRequestValidation verifies that the txnWriteBuffer
+// rejects requests that it doesn't know how to support.
+func TestTxnWriteBufferBatchRequestValidation(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+	twb, mockSender := makeMockTxnWriteBuffer(cluster.MakeClusterSettings())
+
+	type testCase struct {
+		name string
+		ba   func() *kvpb.BatchRequest
+	}
+
+	txn := makeTxnProto()
+	txn.Sequence = 1
+	keyA, keyC := roachpb.Key("a"), roachpb.Key("c")
+
+	tests := []testCase{
+		{
+			name: "batch with OriginTimestamp",
+			ba: func() *kvpb.BatchRequest {
+				header := kvpb.Header{
+					Txn: &txn,
+					WriteOptions: &kvpb.WriteOptions{
+						OriginTimestamp: hlc.Timestamp{WallTime: 1},
+					},
+				}
+				return &kvpb.BatchRequest{Header: header}
+			},
+		},
+		{
+			name: "batch with OriginID",
+			ba: func() *kvpb.BatchRequest {
+				header := kvpb.Header{
+					Txn: &txn,
+					WriteOptions: &kvpb.WriteOptions{
+						OriginID: 1,
+					},
+				}
+				return &kvpb.BatchRequest{Header: header}
+			},
+		},
+		{
+			name: "batch with InitPut",
+			ba: func() *kvpb.BatchRequest {
+				b := &kvpb.BatchRequest{Header: kvpb.Header{Txn: &txn}}
+				b.Add(&kvpb.InitPutRequest{
+					RequestHeader: kvpb.RequestHeader{Key: keyA, Sequence: txn.Sequence},
+					Value:         roachpb.Value{},
+				})
+				return b
+			},
+		},
+		{
+			name: "batch with Increment",
+			ba: func() *kvpb.BatchRequest {
+				b := &kvpb.BatchRequest{Header: kvpb.Header{Txn: &txn}}
+				b.Add(&kvpb.IncrementRequest{
+					RequestHeader: kvpb.RequestHeader{Key: keyA, Sequence: txn.Sequence},
+				})
+				return b
+			},
+		},
+		{
+			name: "batch with ReturnRawMVCCValues Scan",
+			ba: func() *kvpb.BatchRequest {
+				b := &kvpb.BatchRequest{Header: kvpb.Header{Txn: &txn}}
+				r := &kvpb.ScanRequest{
+					ReturnRawMVCCValues: true,
+					RequestHeader:       kvpb.RequestHeader{Key: keyA, EndKey: keyC, Sequence: txn.Sequence},
+				}
+				b.Add(r)
+				return b
+			},
+		},
+		{
+			name: "batch with ReturnRawMVCCValues ReverseScan",
+			ba: func() *kvpb.BatchRequest {
+				b := &kvpb.BatchRequest{Header: kvpb.Header{Txn: &txn}}
+				r := &kvpb.ReverseScanRequest{
+					ReturnRawMVCCValues: true,
+					RequestHeader:       kvpb.RequestHeader{Key: keyA, EndKey: keyC, Sequence: txn.Sequence},
+				}
+				b.Add(r)
+				return b
+			},
+		},
+		{
+			name: "batch with ReturnRawMVCCValues Get",
+			ba: func() *kvpb.BatchRequest {
+				b := &kvpb.BatchRequest{Header: kvpb.Header{Txn: &txn}}
+				r := &kvpb.GetRequest{
+					ReturnRawMVCCValues: true,
+					RequestHeader:       kvpb.RequestHeader{Key: keyA, Sequence: txn.Sequence},
+				}
+				b.Add(r)
+				return b
+			},
+		},
+		{
+			name: "batch with COL_BATCH_RESPONSE Scan",
+			ba: func() *kvpb.BatchRequest {
+				b := &kvpb.BatchRequest{Header: kvpb.Header{Txn: &txn}}
+				r := &kvpb.ScanRequest{
+					ScanFormat:    kvpb.COL_BATCH_RESPONSE,
+					RequestHeader: kvpb.RequestHeader{Key: keyA, EndKey: keyC, Sequence: txn.Sequence},
+				}
+				b.Add(r)
+				return b
+			},
+		},
+		{
+			name: "batch with COL_BATCH_RESPONSE ReverseScan",
+			ba: func() *kvpb.BatchRequest {
+				b := &kvpb.BatchRequest{Header: kvpb.Header{Txn: &txn}}
+				r := &kvpb.ReverseScanRequest{
+					ScanFormat:    kvpb.COL_BATCH_RESPONSE,
+					RequestHeader: kvpb.RequestHeader{Key: keyA, EndKey: keyC, Sequence: txn.Sequence},
+				}
+				b.Add(r)
+				return b
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			numCalledBefore := mockSender.NumCalled()
+			_, pErr := twb.SendLocked(ctx, tc.ba())
+			require.NotNil(t, pErr)
+			require.Equal(t, numCalledBefore, mockSender.NumCalled())
+
+		})
+	}
+}