kv: add HasBufferedAllPreceedingWrites flag for BatchRequests

This patch adds a new flag to indicate that a transaction has
buffered all preceeding writes on the client. It also sets it
in the txnWriteBuffer.

In the future, we'll use this information to omit AbortSpan checks
on the server.

Epic: none

Release note: None

Author: arulajmani

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -146,26 +146,26 @@ type txnWriteBuffer struct {
 	// and disable write buffering going forward out of an abundance of caution.
 	// This is opted into by SQL.
 	//
-	// As a result, we have a nice invariant: if write buffering is enabled, then
-	// all writes performed by the transaction are buffered in memory. We can
-	// never have the case where a part of the write set is buffered, and the
-	// other part is replicated.
+	// As a result, we have a nice invariant: if write buffering is enabled,
+	// then all writes performed by the transaction are buffered in memory. We
+	// can never have the case where a part of the write set is buffered, and
+	// the other part is replicated.
 	//
-	// In the future, the invariant above allows us to omit checking the AbortSpan
-	// for transactions that have buffered writes enabled. The AbortSpan is used
-	// to ensure we don't violate read-your-own-write semantics for transactions
-	// that have been aborted by a conflicting transaction. As read-your-own-write
-	// semantics are upheld by the client, not the server, for transactions that
-	// use buffered writes, we can skip the AbortSpan check on the server.
+	// The invariant above allows us to omit checking the AbortSpan for
+	// transactions that have buffered writes enabled. The AbortSpan is used to
+	// ensure we don't violate read-your-own-write semantics for transactions
+	// that have been aborted by a conflicting transaction. As
+	// read-your-own-write semantics are upheld by the client, not the server,
+	// for transactions that use buffered writes, we can skip the AbortSpan
+	// check on the server.
 	//
 	// We currently track this via two state variables: `enabled` and `flushed`.
 	// Writes are only buffered if enabled && !flushed.
 	//
-	// `enabled` tracks whether buffering has been enabled/disabled externally via
-	// txn.SetBufferedWritesEnabled or because we are operating on a leaf
+	// `enabled` tracks whether buffering has been enabled/disabled externally
+	// via txn.SetBufferedWritesEnabled or because we are operating on a leaf
 	// transaction.
 	enabled bool
-	//
 	// `flushed` tracks whether the buffer has been previously flushed.
 	flushed bool
 
@@ -210,6 +210,15 @@ func (twb *txnWriteBuffer) SendLocked(
 
 	if !twb.shouldBuffer() {
 		return twb.wrapped.SendLocked(ctx, ba)
+	} else {
+		// If we're here, write buffering is enabled, and all writes until now
+		// have been buffered. Set the flag to indicate this.
+		//
+		// NB: We don't need a version check here (for v25.3) because this is only
+		// used by the server to optimize away the AbortSpan check. Even if we set
+		// this field, and the server is on a previous version, the worst that can
+		// happen is we'll perform this check, which is harmless.
+		ba.HasBufferedAllPrecedingWrites = true
 	}
 
 	if etArg, ok := ba.GetArg(kvpb.EndTxn); ok {

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer_test.go

@@ -2056,6 +2056,190 @@ func TestTxnWriteBufferBatchRequestValidation(t *testing.T) {
 	}
 }
 
+// TestTxnWriteBufferHasBufferedAllPrecedingWrites verifies that the
+// txnWriteBuffer correctly sets the HasBufferedAllPrecedingWrites flag.
+func TestTxnWriteBufferHasBufferedAllPrecedingWrites(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	txn := makeTxnProto()
+	txn.Sequence = 1
+	keyA, keyB, keyC := roachpb.Key("a"), roachpb.Key("b"), roachpb.Key("c")
+
+	for _, tc := range []struct {
+		name                             string
+		setup                            func(*txnWriteBuffer)
+		ba                               func(ba *kvpb.BatchRequest)
+		mockSend                         func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error)
+		expHasBufferedAllPrecedingWrites bool
+	}{
+		{
+			name: "batch with two Get requests",
+			ba: func(ba *kvpb.BatchRequest) {
+				getA := &kvpb.GetRequest{RequestHeader: kvpb.RequestHeader{Key: keyA, Sequence: txn.Sequence}}
+				getB := &kvpb.GetRequest{RequestHeader: kvpb.RequestHeader{Key: keyB, Sequence: txn.Sequence}}
+				ba.Add(getA, getB)
+			},
+			mockSend: func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+				require.Len(t, ba.Requests, 2)
+				require.IsType(t, &kvpb.GetRequest{}, ba.Requests[0].GetInner())
+				require.IsType(t, &kvpb.GetRequest{}, ba.Requests[1].GetInner())
+
+				require.True(t, ba.HasBufferedAllPrecedingWrites)
+
+				br := ba.CreateReply()
+				br.Txn = ba.Txn
+				return br, nil
+			},
+			expHasBufferedAllPrecedingWrites: true,
+		},
+		{
+			name: "batch with one Put and one Get request",
+			ba: func(ba *kvpb.BatchRequest) {
+				putA := putArgs(keyA, "valA", txn.Sequence)
+				getB := &kvpb.GetRequest{RequestHeader: kvpb.RequestHeader{Key: keyB, Sequence: txn.Sequence}}
+				ba.Add(putA, getB)
+			},
+			mockSend: func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+				require.Len(t, ba.Requests, 1)
+				require.IsType(t, &kvpb.GetRequest{}, ba.Requests[0].GetInner())
+
+				require.True(t, ba.HasBufferedAllPrecedingWrites)
+
+				br := ba.CreateReply()
+				br.Txn = ba.Txn
+				return br, nil
+			},
+			expHasBufferedAllPrecedingWrites: true,
+		},
+		{
+			name: "batch with one Put, one Get, and one Delete request",
+			ba: func(ba *kvpb.BatchRequest) {
+				putA := putArgs(keyA, "valA", txn.Sequence)
+				getB := &kvpb.GetRequest{RequestHeader: kvpb.RequestHeader{Key: keyB, Sequence: txn.Sequence}}
+				delC := delArgs(keyC, txn.Sequence)
+
+				ba.Add(putA, getB, delC)
+			},
+			mockSend: func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+				require.Len(t, ba.Requests, 1)
+				require.IsType(t, &kvpb.GetRequest{}, ba.Requests[0].GetInner())
+
+				require.True(t, ba.HasBufferedAllPrecedingWrites)
+
+				br := ba.CreateReply()
+				br.Txn = ba.Txn
+				return br, nil
+			},
+			expHasBufferedAllPrecedingWrites: true,
+		},
+		{
+			name: "batch with one DeleteRange and one Get request",
+			ba: func(ba *kvpb.BatchRequest) {
+				delRange := delRangeArgs(keyA, keyB, txn.Sequence)
+				getB := &kvpb.GetRequest{RequestHeader: kvpb.RequestHeader{Key: keyB, Sequence: txn.Sequence}}
+
+				ba.Add(delRange, getB)
+			},
+			mockSend: func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+				require.Len(t, ba.Requests, 2)
+				require.IsType(t, &kvpb.DeleteRangeRequest{}, ba.Requests[0].GetInner())
+				require.IsType(t, &kvpb.GetRequest{}, ba.Requests[1].GetInner())
+
+				require.True(t, ba.HasBufferedAllPrecedingWrites)
+
+				br := ba.CreateReply()
+				br.Txn = ba.Txn
+				return br, nil
+			},
+			expHasBufferedAllPrecedingWrites: false,
+		},
+		{
+			name: "flushed due to size limit",
+			setup: func(twb *txnWriteBuffer) {
+				bufferedWritesMaxBufferSize.Override(context.Background(), &twb.st.SV, 1)
+			},
+			ba: func(ba *kvpb.BatchRequest) {
+				putA := putArgs(keyA, "valA", txn.Sequence)
+
+				ba.Add(putA)
+			},
+			mockSend: func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+				require.Len(t, ba.Requests, 1)
+				require.IsType(t, &kvpb.PutRequest{}, ba.Requests[0].GetInner())
+
+				require.True(t, ba.HasBufferedAllPrecedingWrites)
+
+				br := ba.CreateReply()
+				br.Txn = ba.Txn
+				return br, nil
+			},
+			expHasBufferedAllPrecedingWrites: false,
+		},
+		{
+			name: "write buffering disabled",
+			setup: func(twb *txnWriteBuffer) {
+				twb.setEnabled(false)
+			},
+			ba: func(ba *kvpb.BatchRequest) {
+				putA := putArgs(keyA, "valA", txn.Sequence)
+
+				ba.Add(putA)
+			},
+			mockSend: func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+				require.Len(t, ba.Requests, 1)
+				require.IsType(t, &kvpb.PutRequest{}, ba.Requests[0].GetInner())
+
+				// NB: Should never be set if write buffering is disabled
+				require.False(t, ba.HasBufferedAllPrecedingWrites)
+
+				br := ba.CreateReply()
+				br.Txn = ba.Txn
+				return br, nil
+			},
+			expHasBufferedAllPrecedingWrites: false,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx := context.Background()
+			st := cluster.MakeTestingClusterSettings()
+			twb, mockSender := makeMockTxnWriteBuffer(st)
+
+			if tc.setup != nil {
+				tc.setup(&twb)
+			}
+
+			ba := &kvpb.BatchRequest{}
+			tc.ba(ba)
+			mockSender.MockSend(tc.mockSend)
+
+			br, pErr := twb.SendLocked(ctx, ba)
+			require.Nil(t, pErr)
+			require.NotNil(t, br)
+
+			// Go to commit the transaction and ensure HasBufferedAllPrecedingWrites
+			// is set correctly.
+			mockSender.MockSend(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+				require.Equal(t, tc.expHasBufferedAllPrecedingWrites, ba.HasBufferedAllPrecedingWrites)
+
+				br = ba.CreateReply()
+				br.Txn = ba.Txn
+				return br, nil
+			})
+
+			ba = &kvpb.BatchRequest{}
+			ba.Header = kvpb.Header{Txn: &txn}
+			ba.Add(&kvpb.EndTxnRequest{Commit: true})
+
+			br, pErr = twb.SendLocked(ctx, ba)
+			require.Nil(t, pErr)
+			require.NotNil(t, br)
+			require.Len(t, br.Responses, 1)
+			require.IsType(t, &kvpb.EndTxnResponse{}, br.Responses[0].GetInner())
+		})
+	}
+}
+
 // BenchmarkTxnWriteBuffer benchmarks the txnWriteBuffer. The test sets up a
 // transaction with an existing buffer and runs a single batch through
 // SendLocked and flushBufferAndSendBatch. The test varies the state of the

pkg/kv/kvpb/api.proto

@@ -2928,8 +2928,6 @@ message Header {
   // * A destination node older than 24.1 will not see this field.
   RangeInfo proxy_range_info = 34;
 
-  reserved 7, 10, 12, 14, 20;
-
   WriteOptions write_options = 35;
 
   // DeadlockTimeout specifies the amount of time that a request will wait on a
@@ -2945,7 +2943,27 @@ message Header {
   google.protobuf.Duration deadlock_timeout = 36 [(gogoproto.nullable) = false,
     (gogoproto.stdduration) = true];
 
-  // Next ID: 37
+  // HasBufferedAllPrecedingWrites, if set, indicates that the batch belongs to
+  // a transaction that has buffered all of its writes (from preceding batches)
+  // on the client. 
+  //
+  // The server may use this field to omit checking the AbortSpan. Transactions
+  // use the AbortSpan to check whether they've been aborted or not. If they
+  // have, any intents they may have previously written could be removed by
+  // concurrent transactions, which means a transaction may not have a guarantee
+  // to read its own writes. So, transactions eagerly check the AbortSpan to
+  // identify this case and eagerly return a TransactionAbortedError to the
+  // client, instead of breaking read-your-own-writes. However, transactions
+  // that have buffered all writes on the client uphold read-your-own-writes
+  // semantics by joining results from the KVServer against the write buffer;
+  // simply put, they do not rely on the AbortSpan to uphold
+  // read-your-own-writes. As such, they can eschew checking the AbortSpan on
+  // the server.
+  bool has_buffered_all_preceding_writes = 37;
+  
+  reserved 7, 10, 12, 14, 20;
+
+  // Next ID: 38
 }
 
 message WriteOptions {