fix: include mutations from original attempt in retry (#463)

Mutations that had been buffered during a transaction would not be included
in internal retries if the original transaction was aborted.

Author: olavloite

aborted_transactions_test.go

@@ -60,6 +60,51 @@ func TestCommitAborted(t *testing.T) {
 	}
 }
 
+func TestCommitWithMutationsAborted(t *testing.T) {
+	t.Parallel()
+
+	db, server, teardown := setupTestDBConnectionWithParams(t, "minSessions=1;maxSessions=1")
+	defer teardown()
+	ctx := context.Background()
+
+	conn, err := db.Conn(ctx)
+	if err != nil {
+		t.Fatalf("failed to open connection: %v", err)
+	}
+	defer func() { _ = conn.Close() }()
+
+	tx, err := conn.BeginTx(ctx, &sql.TxOptions{})
+	if err != nil {
+		t.Fatalf("begin failed: %v", err)
+	}
+	if err := conn.Raw(func(driverConn interface{}) error {
+		spannerConn, _ := driverConn.(SpannerConn)
+		mutation := spanner.Insert("foo", []string{}, []interface{}{})
+		return spannerConn.BufferWrite([]*spanner.Mutation{mutation})
+	}); err != nil {
+		t.Fatalf("failed to buffer mutations: %v", err)
+	}
+	// Abort the transaction on the first commit attempt.
+	server.TestSpanner.PutExecutionTime(testutil.MethodCommitTransaction, testutil.SimulatedExecutionTime{
+		Errors: []error{status.Error(codes.Aborted, "Aborted")},
+	})
+	err = tx.Commit()
+	if err != nil {
+		t.Fatalf("commit failed: %v", err)
+	}
+	reqs := drainRequestsFromServer(server.TestSpanner)
+	commitReqs := requestsOfType(reqs, reflect.TypeOf(&sppb.CommitRequest{}))
+	if g, w := len(commitReqs), 2; g != w {
+		t.Fatalf("commit request count mismatch\nGot: %v\nWant: %v", g, w)
+	}
+	for _, req := range commitReqs {
+		commitReq := req.(*sppb.CommitRequest)
+		if g, w := len(commitReq.Mutations), 1; g != w {
+			t.Fatalf("mutation count mismatch\n Got: %v\nWant: %v", g, w)
+		}
+	}
+}
+
 func TestCommitAbortedWithInternalRetriesDisabled(t *testing.T) {
 	t.Parallel()
 

transaction.go

@@ -233,6 +233,10 @@ type readWriteTransaction struct {
 	// transaction so far. These statements will be replayed on a new read write
 	// transaction if the initial attempt is aborted.
 	statements []retriableStatement
+
+	// mutations contains the buffered mutations of this transaction. These are
+	// added to the next transaction if the transaction executes an internal retry.
+	mutations []*spanner.Mutation
 }
 
 // retriableStatement is the interface that is used to keep track of statements
@@ -364,6 +368,10 @@ func (tx *readWriteTransaction) retry(ctx context.Context) (err error) {
 		tx.logger.Log(ctx, LevelNotice, "failed to reset transaction")
 		return err
 	}
+	// Re-apply the mutations from the previous transaction.
+	if err := tx.rwTx.BufferWrite(tx.mutations); err != nil {
+		return err
+	}
 	for _, stmt := range tx.statements {
 		tx.logger.Log(ctx, slog.LevelDebug, "retrying statement", "stmt", stmt)
 		err = stmt.retry(ctx, tx.rwTx)
@@ -600,6 +608,7 @@ func (tx *readWriteTransaction) runDmlBatch(ctx context.Context) (*result, error
 }
 
 func (tx *readWriteTransaction) BufferWrite(ms []*spanner.Mutation) error {
+	tx.mutations = append(tx.mutations, ms...)
 	return tx.rwTx.BufferWrite(ms)
 }