Merge #153514

153514: kvcoord: use separate flush for EndTxn with Prepare=true r=yuzefovich a=stevendanna

Because of #153513 it is not safe to send a batch with writes and an
EndTxn where Prepare=true as it might be immediately committed as a 1PC
transaction.

I've forgone a release note here because as far as I can tell
XA-transactions are not yet a documented feature.

Informs #153513
Fixes #144252

Release note: None

Co-authored-by: Steven Danna <[email protected]>

Author: craig[bot]

pkg/ccl/multitenantccl/tenantcapabilitiesccl/testdata/can_prepare_txns

@@ -3,14 +3,15 @@ SELECT * FROM [SHOW TENANT [10] WITH CAPABILITIES] WHERE capability_name = 'can_
 ----
 10 cluster-10 ready external can_prepare_txns false
 
-# TODO(#144252): remove this.
+
 exec-sql-tenant
-SET kv_transaction_buffered_writes_enabled = false
+CREATE TABLE t(a INT PRIMARY KEY)
 ----
 ok
 
-exec-sql-tenant
-CREATE TABLE t(a INT PRIMARY KEY)
+# Disable buffered writes so that the failing batch has the same requests.
+exec-privileged-op-tenant
+SET kv_transaction_buffered_writes_enabled = false
 ----
 ok
 
@@ -28,8 +29,12 @@ ok
 exec-privileged-op-tenant
 PREPARE TRANSACTION 'txn1'
 ----
-pq: ba: QueryIntent [/Tenant/10/Table/104/1/1/0], EndTxn(commit) [/Tenant/10/Table/104/1/1/0], [txn: ‹×›], [can-forward-ts] RPC error: rpc error: code = Unauthenticated desc = client tenant does not have capability "can_prepare_txns" (*kvpb.EndTxnRequest)
+pq: ba: QueryIntent [/Tenant/10/Table/104/1/1/0], EndTxn(prepare) [/Tenant/10/Table/104/1/1/0], [txn: ‹×›], [can-forward-ts] RPC error: rpc error: code = Unauthenticated desc = client tenant does not have capability "can_prepare_txns" (*kvpb.EndTxnRequest)
 
+exec-privileged-op-tenant
+RESET kv_transaction_buffered_writes_enabled
+----
+ok
 
 # Grant the capability.
 update-capabilities
@@ -57,13 +62,18 @@ ROLLBACK PREPARED 'txn2'
 ----
 ok
 
-
 # Revoke the capability using REVOKE syntax.
 update-capabilities
 ALTER TENANT [10] REVOKE CAPABILITY can_prepare_txns
 ----
 ok
 
+# Disable buffered writes so that the failing batch has the same requests.
+exec-privileged-op-tenant
+SET kv_transaction_buffered_writes_enabled = false
+----
+ok
+
 # Prepared transactions should no longer work.
 exec-privileged-op-tenant
 BEGIN
@@ -78,8 +88,12 @@ ok
 exec-privileged-op-tenant
 PREPARE TRANSACTION 'txn3'
 ----
-pq: ba: QueryIntent [/Tenant/10/Table/104/1/1/0], EndTxn(commit) [/Tenant/10/Table/104/1/1/0], [txn: ‹×›], [can-forward-ts] RPC error: rpc error: code = Unauthenticated desc = client tenant does not have capability "can_prepare_txns" (*kvpb.EndTxnRequest)
+pq: ba: QueryIntent [/Tenant/10/Table/104/1/1/0], EndTxn(prepare) [/Tenant/10/Table/104/1/1/0], [txn: ‹×›], [can-forward-ts] RPC error: rpc error: code = Unauthenticated desc = client tenant does not have capability "can_prepare_txns" (*kvpb.EndTxnRequest)
 
+exec-privileged-op-tenant
+RESET kv_transaction_buffered_writes_enabled
+----
+ok
 
 # However, transactions that have not acquired locks are able to be prepared,
 # since they don't actually prepare a transaction record in the KV layer. This

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -1762,7 +1762,7 @@ func (twb *txnWriteBuffer) flushBufferAndSendBatch(
 		return twb.wrapped.SendLocked(ctx, ba) // nothing to flush
 	}
 
-	_, hasEndTxn := ba.GetArg(kvpb.EndTxn)
+	endTxnArg, hasEndTxn := ba.GetArg(kvpb.EndTxn)
 	if !hasEndTxn {
 		// We're flushing the buffer even though the batch doesn't contain an EndTxn
 		// request. That means we buffered some writes and decided to disable write
@@ -1771,7 +1771,7 @@ func (twb *txnWriteBuffer) flushBufferAndSendBatch(
 	}
 
 	midTxnFlush := !hasEndTxn
-	splitBatchRequired := separateBatchIsNeeded(ba)
+	splitBatchRequired := separateBatchIsNeeded(ba, endTxnArg)
 
 	// Flush all buffered writes by pre-pending them to the requests being sent
 	// in the batch.
@@ -1861,7 +1861,16 @@ func requireAllFlushedRequestsProcessed(responses []kvpb.ResponseUnion) error {
 //
 // NB: If you are updating this function, you need to update
 // clearBatchRequestOptions as well.
-func separateBatchIsNeeded(ba *kvpb.BatchRequest) bool {
+func separateBatchIsNeeded(ba *kvpb.BatchRequest, optEndTxn kvpb.Request) bool {
+	if optEndTxn != nil {
+		// TODO(#153513): This should really be fixed server-side. Currently, if we
+		// send an EndTxn with Prepare, it will be erroneously evaluated as a 1PC
+		// commit.
+		if optEndTxn.(*kvpb.EndTxnRequest).Prepare {
+			return true
+		}
+	}
+
 	return ba.MightStopEarly() ||
 		ba.ReadConsistency != 0 ||
 		ba.WaitPolicy != 0 ||

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer_test.go

@@ -2186,7 +2186,7 @@ func TestTxnWriteBufferSplitsBatchesWhenNecessary(t *testing.T) {
 			br.Txn = ba.Txn
 			if requestCount == 0 {
 				// The buffer flush should not have any problematic settings.
-				require.False(t, separateBatchIsNeeded(ba))
+				require.False(t, separateBatchIsNeeded(ba, nil))
 			}
 			requestCount++
 			return br, nil
@@ -4090,7 +4090,7 @@ needs to be accounted for in the following functions:
 				}
 			}
 			req := &kvpb.BatchRequest{Header: header}
-			require.True(t, separateBatchIsNeeded(req),
+			require.True(t, separateBatchIsNeeded(req, nil),
 				"non-zero value for %s not handled in separateBatchIsNeeded", fieldName)
 
 			clearBatchRequestOptions(req)

pkg/kv/kvpb/api.go

@@ -248,7 +248,9 @@ var _ SafeFormatterRequest = (*EndTxnRequest)(nil)
 func (etr *EndTxnRequest) SafeFormat(s redact.SafePrinter, _ rune) {
 	s.Printf("%s(", etr.Method())
 	if etr.Commit {
-		if etr.IsParallelCommit() {
+		if etr.Prepare {
+			s.Printf("prepare")
+		} else if etr.IsParallelCommit() {
 			s.Printf("parallel commit")
 		} else {
 			s.Printf("commit")

pkg/sql/logictest/testdata/logic_test/two_phase_commit

@@ -1,9 +1,5 @@
 # LogicTest: !3node-tenant-default-configs !local-prepared
 
-# TODO(#144252): remove this.
-statement ok
-SET kv_transaction_buffered_writes_enabled = false;
-
 query T
 SHOW max_prepared_transactions
 ----