Merge #144799

144799: kvserver: omit AbortSpan checks for buffered writes transactions  r=stevendanna,tbg a=arulajmani

Transactions that use buffered writes do not rely on the AbortSpan
to correctly uphold read-your-own-writes semantics. As such, we can
omit AbortSpan checks for transactions that have buffered all writes,
from preceding batches, on the client.

Fixes #140593

Release note: None

Co-authored-by: Arul Ajmani <[email protected]>

Author: craig[bot]

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/kvclient/kvcoord/txn_test.go

@@ -38,10 +38,20 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/randutil"
+	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
+	"github.com/cockroachdb/cockroach/pkg/util/uuid"
 	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/require"
 )
 
+func checkGetResults(t *testing.T, expected map[string][]byte, results ...kv.Result) {
+	for _, result := range results {
+		require.Equal(t, 1, len(result.Rows))
+		require.Equal(t, expected[string(result.Rows[0].Key)], result.Rows[0].ValueBytes())
+	}
+	require.Len(t, expected, len(results))
+}
+
 // TestTxnDBBasics verifies that a simple transaction can be run and
 // either committed or aborted. On commit, mutations are visible; on
 // abort, mutations are never visible. During the txn, verify that
@@ -2371,3 +2381,82 @@ func TestLeafTransactionAdmissionHeader(t *testing.T) {
 	}
 	require.Equal(t, expectedLeafHeader, leafHeader)
 }
+
+// TestTxnBufferedWritesOmitAbortSpanChecks verifies that transactions that use
+// buffered writes do not check the AbortSpan, while still upholding
+// read-your-own-writes semantics.
+func TestTxnBufferedWritesOmitAbortSpanChecks(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+
+	var mu struct {
+		syncutil.Mutex
+		txnID uuid.UUID
+	}
+	s := createTestDBWithKnobs(t, &kvserver.StoreTestingKnobs{
+		EvalKnobs: kvserverbase.BatchEvalTestingKnobs{
+			BeforeAbortSpanCheck: func(id uuid.UUID) {
+				mu.Lock()
+				defer mu.Unlock()
+
+				if mu.txnID == id {
+					t.Fatal("transactions using buffered writes should not check the AbortSpan")
+				}
+			},
+		},
+	})
+	defer s.Stop()
+
+	value1 := []byte("value1")
+	valueConflict := []byte("conflict")
+
+	keyA := []byte("keyA")
+
+	txn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
+	txn.SetBufferedWritesEnabled(true)
+	mu.Lock()
+	mu.txnID = txn.ID()
+	mu.Unlock()
+
+	// Fix the transaction's commit timestamp.
+	_, err := txn.CommitTimestamp()
+	require.NoError(t, err)
+
+	// Put transactional value at keyA.
+	require.NoError(t, txn.Put(ctx, keyA, value1))
+
+	// Read what we just wrote.
+	b := txn.NewBatch()
+	b.Get(keyA)
+	require.NoError(t, txn.Run(ctx, b))
+	expected := map[string][]byte{
+		"keyA": value1,
+	}
+	checkGetResults(t, expected, b.Results...)
+
+	// Start another transaction that writes to keyA. This prevents us from
+	// committing at our original timestamp. Moreover, had we not been buffering
+	// our writes, this transaction would have resulted in aborting us and
+	// removing our intent.
+	err = s.DB.Txn(context.Background(), func(ctx context.Context, txn *kv.Txn) error {
+		require.NoError(t, txn.SetUserPriority(roachpb.MaxUserPriority))
+		return txn.Put(ctx, keyA, valueConflict)
+	})
+	require.NoError(t, err)
+
+	// Perform another read again. We should still see our previous write, not what
+	// the conflicting transaction wrote.
+	b = txn.NewBatch()
+	b.Get(keyA)
+	require.NoError(t, txn.Run(ctx, b))
+	expected = map[string][]byte{
+		"keyA": value1,
+	}
+	checkGetResults(t, expected, b.Results...)
+
+	// Try to commit the transaction. We should encounter a WriteTooOldError.
+	err = txn.Commit(ctx)
+	require.Error(t, err)
+	require.Regexp(t, "TransactionRetryWithProtoRefreshError: .*WriteTooOldError", err)
+}

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 {

pkg/kv/kvserver/kvserverbase/BUILD.bazel

@@ -28,6 +28,7 @@ go_library(
         "//pkg/util/quotapool",
         "//pkg/util/syncutil",
         "//pkg/util/timeutil",
+        "//pkg/util/uuid",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_cockroachdb_pebble//vfs",
         "@com_github_cockroachdb_redact//:redact",