kvserver: omit AbortSpan checks for buffered writes transactions

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

Author: arulajmani

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/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",

pkg/kv/kvserver/kvserverbase/knobs.go

@@ -9,7 +9,11 @@
 
 package kvserverbase
 
-import "time"
+import (
+	"time"
+
+	"github.com/cockroachdb/cockroach/pkg/util/uuid"
+)
 
 // BatchEvalTestingKnobs contains testing helpers that are used during batch evaluation.
 type BatchEvalTestingKnobs struct {
@@ -60,6 +64,10 @@ type BatchEvalTestingKnobs struct {
 
 	// CommitTriggerError is called at commit triggers to simulate errors.
 	CommitTriggerError func() error
+
+	// BeforeAbortSpanCheck is called before a request checks the abort span with
+	// the request's txn ID.
+	BeforeAbortSpanCheck func(id uuid.UUID)
 }
 
 // IntentResolverTestingKnobs contains testing helpers that are used during

pkg/kv/kvserver/replica_evaluate.go

@@ -240,16 +240,28 @@ func evaluateBatch(
 		// transactions on reads). Note that 1PC transactions have had their
 		// transaction field cleared by this point so we do not execute this
 		// check in that case.
+		//
+		// TODO(arul): this check assumes lock == Intent, which isn't true any
+		// longer. We could optimize this by making a distinction between locks
+		// acquired and previous writes performed.
 		if baHeader.Txn.IsLocking() {
 			// We don't check the abort span for a couple of special requests:
 			// - if the request is asking to abort the transaction, then don't check the
 			// AbortSpan; we don't want the request to be rejected if the transaction
 			// has already been aborted.
 			// - heartbeats don't check the abort span. If the txn is aborted, they'll
 			// return an aborted proto in their otherwise successful response.
+			// - if the request belongs to a transaction that has buffered all
+			// preceding writes on the client, we don't rely on the AbortSpan to
+			// correctly uphold read-your-own-write semantics.
+			//
 			// TODO(nvanbenschoten): Let's remove heartbeats from this allowlist when
 			// we rationalize the TODO in txnHeartbeater.heartbeat.
-			if !ba.IsSingleAbortTxnRequest() && !ba.IsSingleHeartbeatTxnRequest() {
+			if !ba.IsSingleAbortTxnRequest() && !ba.IsSingleHeartbeatTxnRequest() &&
+				!ba.HasBufferedAllPrecedingWrites {
+				if fn := rec.EvalKnobs().BeforeAbortSpanCheck; fn != nil {
+					fn(ba.Txn.ID)
+				}
 				if pErr := checkIfTxnAborted(ctx, rec, readWriter, *baHeader.Txn); pErr != nil {
 					return nil, result.Result{}, pErr
 				}