kvcoord: fix EndTxn elision race with in-flight requests

`TxnCoordSender` has a fast-path that avoids sending `EndTxn` requests
if the transaction has not taken out any locks. However, it relied on
`txnPipeliner` to keep track of locks, and it only registers these when
responses arrive. Since `EndTxn(commit=false)` requests can be sent
asynchronously (e.g. by the client disconnecting), these can arrive
while the write request is still in flight, causing `TxnCoordSender` to
incorrectly believe the transaction hasn't taken out any locks and elide
the `EndTxn` request. This would cause the written intents to be left
behind.

This patch changes `TxnCoordSender` to keep track of any locking
requests that have been sent, and only elides `EndTxn` requests when no
locking requests have been seen.

`txnCommitter` also elides `EndTxn` requests, but this is done below the
`txnFinalizeBarrier` which serializes the requests, so this is safe.

Release note (bug fix): Fixed a race condition where transaction
cleanup would not be triggered if the transaction was rolled back while
all locking requests were still in flight.

Author: erikgrinaker

pkg/kv/kvclient/kvcoord/txn_coord_sender.go

@@ -112,6 +112,13 @@ type TxnCoordSender struct {
 		// (a retryable TransactionAbortedError in case of the async abort).
 		closed bool
 
+		// sentLocking is set once this transaction has sent a locking request. It
+		// is used to elide EndTxn requests via finalizeNonLockingTxnLocked. We
+		// can't rely on txnPipeliner for this, since it doesn't track locks until
+		// after responses have been received, and an EndTxn(commit=false) could
+		// arrive before that happens.
+		sentLocking bool
+
 		// txn is the Transaction proto attached to all the requests and updated on
 		// all the responses.
 		txn roachpb.Transaction
@@ -485,7 +492,7 @@ func (tc *TxnCoordSender) Send(
 		return nil, pErr
 	}
 
-	if ba.IsSingleEndTxnRequest() && !tc.interceptorAlloc.txnPipeliner.hasAcquiredLocks() {
+	if ba.IsSingleEndTxnRequest() && !tc.mu.sentLocking {
 		return nil, tc.finalizeNonLockingTxnLocked(ctx, ba)
 	}
 
@@ -516,6 +523,10 @@ func (tc *TxnCoordSender) Send(
 		return nil, nil
 	}
 
+	if !tc.mu.sentLocking && ba.IsLocking() {
+		tc.mu.sentLocking = true
+	}
+
 	// Clone the Txn's Proto so that future modifications can be made without
 	// worrying about synchronization.
 	ba.Txn = tc.mu.txn.Clone()

pkg/kv/kvserver/intent_resolver_integration_test.go

@@ -401,6 +401,12 @@ func TestReliableIntentCleanup(t *testing.T) {
 		testTxnExecute := func(t *testing.T, spec testTxnSpec, txn *kv.Txn) (roachpb.Key, error) {
 			txnKey := genKey(spec.singleRange)
 
+			if spec.finalize == "cancelAsync" && spec.abort != "" {
+				// This would require coordinating the abort, cancel, and put
+				// goroutines. Doesn't seem worth the added complexity.
+				require.Fail(t, "Can't combine finalize=cancelAsync and abort")
+			}
+
 			// If requested, spin off txn aborter goroutines, returning errors
 			// (if any) via abortErrC.
 			//
@@ -461,6 +467,33 @@ func TestReliableIntentCleanup(t *testing.T) {
 				require.Fail(t, "invalid abort type", "abort=%v", spec.abort)
 			}
 
+			// If requested, cancel the context while the put is in flight and
+			// roll back the txn with the cancelled context (which is what the
+			// SQL session would do). This is different from finalize=cancel in
+			// that it's a regression test for EndTxn elision:
+			// https://github.com/cockroachdb/cockroach/issues/65587
+			cancelAsyncErrC := make(chan error, 1)
+			if spec.finalize == "cancelAsync" {
+				readyC := blockPut(t, txnKey)
+				require.NoError(t, tc.Stopper().RunAsyncTask(ctx, "cancel", func(ctx context.Context) {
+					unblockC := <-readyC
+					defer close(unblockC)
+					defer close(cancelAsyncErrC)
+
+					// Roll back with a new, cancelled context and wait to
+					// unblock Put. This is to provoke EndTxn elision by making
+					// sure the txn abort executes before the Put request. We
+					// don't cancel the main context, since that would cause the
+					// Put to error immediately, racing with the txn abort.
+					rollbackCtx, rollbackCancel := context.WithCancel(ctx)
+					rollbackCancel()
+					if err := txn.Rollback(rollbackCtx); err != nil && !errors.Is(err, context.Canceled) {
+						cancelAsyncErrC <- err
+					}
+					time.Sleep(100 * time.Millisecond)
+				}))
+			}
+
 			// Write numKeys KV pairs in batches of batchSize as a single txn.
 			const batchSize = 10000
 			batch := txn.NewBatch()
@@ -475,6 +508,11 @@ func TestReliableIntentCleanup(t *testing.T) {
 						return nil, err
 					}
 					batch = txn.NewBatch()
+					// If we did an async cancel, submitting another Put violates
+					// the txn protocol (i.e. the client has already gone away).
+					if spec.finalize == "cancelAsync" {
+						break
+					}
 				}
 			}
 
@@ -499,6 +537,10 @@ func TestReliableIntentCleanup(t *testing.T) {
 				if err := txn.Rollback(rollbackCtx); err != nil && !errors.Is(err, context.Canceled) {
 					return nil, err
 				}
+			case "cancelAsync":
+				if err := <-cancelAsyncErrC; err != nil {
+					return nil, err
+				}
 			default:
 				require.Fail(t, "invalid finalize type", "finalize=%v", spec.finalize)
 			}
@@ -577,9 +619,14 @@ func TestReliableIntentCleanup(t *testing.T) {
 						})
 						return
 					}
-					finalize := []interface{}{"commit", "rollback", "cancel"}
+					finalize := []interface{}{"commit", "rollback", "cancel", "cancelAsync"}
 					testutils.RunValues(t, "finalize", finalize, func(t *testing.T, finalize interface{}) {
 						abort := []interface{}{"no", "push", "heartbeat"}
+						if finalize == "cancelAsync" {
+							// Async cancel can't currently run together with an
+							// abort, since we'd have to coordinate the tasks.
+							abort = []interface{}{"no"}
+						}
 						testutils.RunValues(t, "abort", abort, func(t *testing.T, abort interface{}) {
 							if abort.(string) == "no" {
 								abort = "" // "no" just makes the test output better