fix(datastore): detach rollback context from transaction cancellation (#14602)

bhshkh · web-flow · commit 0a8d10d18d7e · 2026-05-19T09:04:57.000-07:00
Fix a lock contention and high-latency issue in RunInTransaction caused
by failing to send Rollback commands when a context timer expires
or is canceled.

Fixes: [b/346474631](http://b/346474631)
diff --git a/datastore/transaction.go b/datastore/transaction.go
@@ -475,6 +475,21 @@ func (t *Transaction) Commit() (c *Commit, err error) {
 // Returns last attempt rollback error if rollback fails even after retries
 func (t *Transaction) rollbackWithRetry() error {
 	var rollbackErr error
+	// Use a detached context for rollback so that if the parent context is cancelled/timed out,
+	// the rollback RPCs can still be sent to the server to release locks.
+	// We set a timeout of 5 seconds on the rollback context.
+	rollbackCtx, cancel := context.WithTimeout(context.WithoutCancel(t.ctx), 5*time.Second)
+	defer cancel()
+
+	// Temporarily override the transaction context for the rollback operation.
+	// Since Rollback and rollbackWithRetry are only called once at the end of a transaction,
+	// this is clean and safe.
+	origCtx := t.ctx
+	t.ctx = rollbackCtx
+	defer func() {
+		t.ctx = origCtx
+	}()
+
 	retryer := gax.OnCodes(rollbackRetryCodes, txnBackoff)
 	for rollbackAttempt := 0; rollbackAttempt < maxIndividualReqTxnRetry; rollbackAttempt++ {
 		rollbackErr = t.Rollback()
diff --git a/datastore/transaction_test.go b/datastore/transaction_test.go
@@ -881,3 +881,30 @@ func TestRunInTransaction(t *testing.T) {
 		})
 	}
 }
+
+func TestRunInTransaction_ContextCanceled(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	client, srv, cleanup := newMock(t)
+	defer cleanup()
+
+	srv.addRPC(&pb.BeginTransactionRequest{ProjectId: mockProjectID}, &pb.BeginTransactionResponse{Transaction: []byte("tid")})
+	// We expect Rollback to be called even though context is canceled.
+	srv.addRPC(&pb.RollbackRequest{ProjectId: mockProjectID, Transaction: []byte("tid")}, &pb.RollbackResponse{})
+
+	f := func(tx *Transaction) error {
+		cancel() // Cancel the context inside the transaction function
+		return nil
+	}
+
+	_, err := client.RunInTransaction(ctx, f)
+	if err == nil {
+		t.Fatal("expected error, got nil")
+	}
+	if status.Code(err) != codes.Canceled && err != context.Canceled {
+		t.Errorf("expected canceled error, got %v", err)
+	}
+
+	if len(srv.reqItems) != 0 {
+		t.Errorf("mock server has %d remaining expected RPCs (Rollback likely not called)", len(srv.reqItems))
+	}
+}
