fix(firestore): Correct ReadWrite transaction retries (googleapis#12893)

Fixes: 

- [b/443038452](b/443038452)
- [b/291261375](b/291261375)
- [b/248162068](b/248162068)

Retry logic has been changed as per
[go/transaction-retry-matrix](http://go/transaction-retry-matrix) and
[go/transaction-retries](go/transaction-retries)

Author: bhshkh

firestore/transaction.go

@@ -102,6 +102,21 @@ var (
 
 type transactionInProgressKey struct{}
 
+// isAborted checks if an error from a transaction operation
+// indicates that the entire transaction should be retried.
+func isAborted(err error) bool {
+	s, ok := status.FromError(err)
+	if !ok {
+		return false // Not a gRPC error
+	}
+	switch s.Code() {
+	case codes.Aborted:
+		return true
+	default:
+		return false
+	}
+}
+
 // RunTransaction runs f in a transaction. f should use the transaction it is given
 // for all Firestore operations. For any operation requiring a context, f should use
 // the context it is passed, not the first argument to RunTransaction.
@@ -162,41 +177,42 @@ func (c *Client) RunTransaction(ctx context.Context, f func(context.Context, *Tr
 			return err
 		}
 		t.id = res.Transaction
+
 		err = f(context.WithValue(ctx, transactionInProgressKey{}, 1), t)
 		// Read after write can only be checked client-side, so we make sure to check
 		// even if the user does not.
 		if err == nil && t.readAfterWrite {
 			err = errReadAfterWrite
 		}
-		if err != nil {
-			t.rollback()
-			// Prefer f's returned error to rollback error.
-			return err
-		}
-		t.ctx = trace.StartSpan(t.ctx, "cloud.google.com/go/firestore.Client.Commit")
-		commitResponse, err = t.c.c.Commit(t.ctx, &pb.CommitRequest{
-			Database:    t.c.path(),
-			Writes:      t.writes,
-			Transaction: t.id,
-		})
-		trace.EndSpan(t.ctx, err)
 
-		// on success, handle the commit response
 		if err == nil {
-			for _, opt := range opts {
-				opt.handleCommitResponse(commitResponse)
+			t.ctx = trace.StartSpan(t.ctx, "cloud.google.com/go/firestore.Client.Commit")
+			commitResponse, err = t.c.c.Commit(t.ctx, &pb.CommitRequest{
+				Database:    t.c.path(),
+				Writes:      t.writes,
+				Transaction: t.id,
+			})
+			trace.EndSpan(t.ctx, err)
+
+			// on success, handle the commit response
+			if err == nil {
+				for _, opt := range opts {
+					opt.handleCommitResponse(commitResponse)
+				}
+				return nil
 			}
 		}
 
-		// If a read-write transaction returns Aborted, retry.
-		// On success or other failures, return here.
-		if t.readOnly || status.Code(err) != codes.Aborted {
-			// According to the Firestore team, we should not roll back here
-			// if err != nil. But spanner does.
-			// See https://code.googlesource.com/gocloud/+/master/spanner/transaction.go#740.
+		// At this point, `err` is non-nil. It came from `f` or `Commit`.
+		t.rollback()
+
+		// If not a retryable error, or if read-only, return now.
+		// (We've already rolled back).
+		if t.readOnly || !isAborted(err) {
 			return err
 		}
 
+		// It's a retryable error, so continue to backoff and retry logic.
 		if txOpts == nil {
 			// txOpts can only be nil if is the first retry of a read-write transaction.
 			// (It is only set here and in the body of "if t.readOnly" above.)
@@ -207,6 +223,9 @@ func (c *Client) RunTransaction(ctx context.Context, f func(context.Context, *Tr
 					ReadWrite: &pb.TransactionOptions_ReadWrite{RetryTransaction: t.id},
 				},
 			}
+		} else if rw := txOpts.GetReadWrite(); rw != nil {
+			// Update transaction ID for read-write retries.
+			rw.RetryTransaction = t.id
 		}
 		// Use exponential backoff to avoid contention with other running
 		// transactions.
@@ -218,11 +237,7 @@ func (c *Client) RunTransaction(ctx context.Context, f func(context.Context, *Tr
 		// Reset state for the next attempt.
 		t.writes = nil
 	}
-	// If we run out of retries, return the last error we saw (which should
-	// be the Aborted from Commit, or a context error).
-	if err != nil {
-		t.rollback()
-	}
+
 	return err
 }
 

firestore/transaction_test.go

@@ -140,18 +140,23 @@ func TestRunTransaction(t *testing.T) {
 	srv.reset()
 	srv.addRPC(beginReq, beginRes)
 	srv.addRPC(commitReq, status.Errorf(codes.Aborted, ""))
-	srv.addRPC(
-		&pb.BeginTransactionRequest{
-			Database: db,
-			Options: &pb.TransactionOptions{
-				Mode: &pb.TransactionOptions_ReadWrite_{
-					ReadWrite: &pb.TransactionOptions_ReadWrite{RetryTransaction: tid},
-				},
+	rollbackReq := &pb.RollbackRequest{Database: db, Transaction: tid}
+	srv.addRPC(rollbackReq, &emptypb.Empty{})
+	tid2 := []byte{2} // Use a new transaction ID for the response to the retry BeginTransaction
+	beginReqRetry := &pb.BeginTransactionRequest{
+		Database: db,
+		Options: &pb.TransactionOptions{
+			Mode: &pb.TransactionOptions_ReadWrite_{
+				ReadWrite: &pb.TransactionOptions_ReadWrite{RetryTransaction: tid}, // Retries with the previous transaction's ID
 			},
 		},
-		beginRes,
-	)
-	srv.addRPC(commitReq, &pb.CommitResponse{CommitTime: aTimestamp})
+	}
+	beginRes2 := &pb.BeginTransactionResponse{Transaction: tid2} // New transaction ID for the successful attempt
+	srv.addRPC(beginReqRetry, beginRes2)
+
+	// Attempt 2: Commit succeeds with the new transaction ID (tid2).
+	commitReq2 := &pb.CommitRequest{Database: db, Transaction: tid2}
+	srv.addRPC(commitReq2, &pb.CommitResponse{CommitTime: aTimestamp})
 	err = c.RunTransaction(ctx, func(_ context.Context, tx *Transaction) error { return nil })
 	if err != nil {
 		t.Fatal(err)
@@ -167,7 +172,17 @@ func TestTransactionErrors(t *testing.T) {
 	var (
 		tid        = []byte{1}
 		unknownErr = status.Errorf(codes.Unknown, "so sad")
-		beginReq   = &pb.BeginTransactionRequest{
+		abortedErr = status.Errorf(codes.Aborted, "not so sad.retryable")
+
+		beginRetryReq = &pb.BeginTransactionRequest{
+			Database: db,
+			Options: &pb.TransactionOptions{
+				Mode: &pb.TransactionOptions_ReadWrite_{
+					ReadWrite: &pb.TransactionOptions_ReadWrite{RetryTransaction: tid},
+				},
+			},
+		}
+		beginReq = &pb.BeginTransactionRequest{
 			Database: db,
 		}
 		beginRes = &pb.BeginTransactionResponse{Transaction: tid}
@@ -180,8 +195,19 @@ func TestTransactionErrors(t *testing.T) {
 			Documents:           []string{db + "/documents/C/a"},
 			ConsistencySelector: &pb.BatchGetDocumentsRequest_Transaction{Transaction: tid},
 		}
+		getRes = []interface{}{
+			&pb.BatchGetDocumentsResponse{
+				Result: &pb.BatchGetDocumentsResponse_Found{Found: &pb.Document{
+					Name:       "projects/projectID/databases/(default)/documents/C/a",
+					CreateTime: aTimestamp,
+					UpdateTime: aTimestamp2,
+				}},
+				ReadTime: aTimestamp2,
+			},
+		}
 		rollbackReq = &pb.RollbackRequest{Database: db, Transaction: tid}
 		commitReq   = &pb.CommitRequest{Database: db, Transaction: tid}
+		commitRes   = &pb.CommitResponse{CommitTime: aTimestamp}
 	)
 
 	t.Run("BeginTransaction has a permanent error", func(t *testing.T) {
@@ -239,6 +265,8 @@ func TestTransactionErrors(t *testing.T) {
 			},
 		})
 		srv.addRPC(commitReq, unknownErr)
+		srv.addRPC(rollbackReq, &emptypb.Empty{})
+
 		err := c.RunTransaction(ctx, get)
 		if status.Code(err) != codes.Unknown {
 			t.Errorf("got <%v>, want Unknown", err)
@@ -356,22 +384,38 @@ func TestTransactionErrors(t *testing.T) {
 	})
 
 	t.Run("Too many retries", func(t *testing.T) {
+		// Use tid = 1 for the first attempt.
+		// Use tid = 2 for the second attempt.
+		tid1 := []byte{1}
+		tid2 := []byte{2}
+		beginRes2 := &pb.BeginTransactionResponse{Transaction: tid2}
+
 		srv.reset()
-		srv.addRPC(beginReq, beginRes)
-		srv.addRPC(commitReq, status.Errorf(codes.Aborted, ""))
-		srv.addRPC(
-			&pb.BeginTransactionRequest{
-				Database: db,
-				Options: &pb.TransactionOptions{
-					Mode: &pb.TransactionOptions_ReadWrite_{
-						ReadWrite: &pb.TransactionOptions_ReadWrite{RetryTransaction: tid},
-					},
+
+		// Attempt 1 (Fails)
+		srv.addRPC(beginReq, beginRes)                          // 1. BeginTransaction (tid1)
+		srv.addRPC(commitReq, status.Errorf(codes.Aborted, "")) // 2. Commit (tid1) fails (Aborted)
+		srv.addRPC(rollbackReq, &emptypb.Empty{})               // 3. Rollback (tid1)
+
+		// Attempt 2 (Fails)
+		beginReqRetry := &pb.BeginTransactionRequest{
+			Database: db,
+			Options: &pb.TransactionOptions{
+				Mode: &pb.TransactionOptions_ReadWrite_{
+					ReadWrite: &pb.TransactionOptions_ReadWrite{RetryTransaction: tid1}, // Retries with previous ID (tid1)
 				},
 			},
-			beginRes,
-		)
-		srv.addRPC(commitReq, status.Errorf(codes.Aborted, ""))
-		srv.addRPC(rollbackReq, &emptypb.Empty{})
+		}
+		// The retry BeginTransaction should return a new ID (tid2)
+		srv.addRPC(beginReqRetry, beginRes2) // 4. BeginTransaction (tid2)
+
+		commitReq2 := &pb.CommitRequest{Database: db, Transaction: tid2} // New commit request with tid2
+		srv.addRPC(commitReq2, status.Errorf(codes.Aborted, ""))         // 5. Commit (tid2) fails (Aborted)
+
+		// Final Rollback on Aborted error when MaxAttempts is reached
+		rollbackReq2 := &pb.RollbackRequest{Database: db, Transaction: tid2}
+		srv.addRPC(rollbackReq2, &emptypb.Empty{}) // 6. Rollback (tid2)
+
 		err := c.RunTransaction(ctx, func(context.Context, *Transaction) error { return nil },
 			MaxAttempts(2))
 		if status.Code(err) != codes.Aborted {
@@ -397,6 +441,22 @@ func TestTransactionErrors(t *testing.T) {
 		}
 	})
 
+	t.Run("Get has a retryable error", func(t *testing.T) {
+		srv.reset()
+		srv.addRPC(beginReq, beginRes)
+		srv.addRPC(getReq, abortedErr)
+		srv.addRPC(rollbackReq, &emptypb.Empty{})
+		srv.addRPC(beginRetryReq, beginRes)
+		srv.addRPC(getReq, getRes)
+		srv.addRPC(commitReq, commitRes)
+		err := c.RunTransaction(ctx, get)
+		if err != nil {
+			t.Errorf("got <%v>, want nil", err)
+		}
+		if !srv.isEmpty() {
+			t.Errorf("Expected %+v requests but not received. srv.reqItems: %+v", len(srv.reqItems), srv.reqItems)
+		}
+	})
 }
 
 func TestTransactionGetAll(t *testing.T) {
@@ -439,6 +499,7 @@ func TestRunTransaction_Retries(t *testing.T) {
 	const db = "projects/projectID/databases/(default)"
 	tid := []byte{1}
 
+	// Attempt 1: Begin
 	srv.addRPC(
 		&pb.BeginTransactionRequest{Database: db},
 		&pb.BeginTransactionResponse{Transaction: tid},
@@ -455,6 +516,7 @@ func TestRunTransaction_Retries(t *testing.T) {
 		Fields: map[string]*pb.Value{"count": intval(7)},
 	}
 
+	// Attempt 1: Commit (Fails)
 	srv.addRPC(
 		&pb.CommitRequest{
 			Database:    db,
@@ -470,6 +532,10 @@ func TestRunTransaction_Retries(t *testing.T) {
 		status.Errorf(codes.Aborted, "something failed! please retry me!"),
 	)
 
+	rollbackReq := &pb.RollbackRequest{Database: db, Transaction: tid}
+	srv.addRPC(rollbackReq, &emptypb.Empty{})
+
+	// Attempt 2: Begin (Retry)
 	srv.addRPC(
 		&pb.BeginTransactionRequest{
 			Database: db,