chore(spanner): preserving lock order - R/W mux session (googleapis#11585)

* chore(spanner): preserving lock order - R/W mux session

* fix build and tests

* remove not used options

Author: rahul2393

spanner/client.go

@@ -1065,8 +1065,13 @@ func (c *Client) rwTransaction(ctx context.Context, f func(context.Context, *Rea
 				return ToSpannerError(err)
 			}
 		} else {
+			var previousTx transactionID
+			if t != nil {
+				previousTx = t.previousTx
+			}
 			t = &ReadWriteTransaction{
 				txReadyOrClosed: make(chan struct{}),
+				previousTx:      previousTx,
 			}
 			t.txReadOnly.sh = sh
 		}

spanner/client_test.go

@@ -129,7 +129,7 @@ func setupMockedTestServerWithConfigAndGCPMultiendpointPool(t *testing.T, config
 	if err != nil {
 		t.Fatal(err)
 	}
-	if isMultiplexEnabled {
+	if isMultiplexEnabled || config.enableMultiplexSession {
 		waitFor(t, func() error {
 			client.idleSessions.mu.Lock()
 			defer client.idleSessions.mu.Unlock()

spanner/transaction.go

@@ -1165,6 +1165,7 @@ type ReadWriteTransaction struct {
 	// ReadWriteTransaction. It is set only once in ReadWriteTransaction.begin()
 	// during the initialization of ReadWriteTransaction.
 	tx             transactionID
+	previousTx     transactionID
 	precommitToken *sppb.MultiplexedSessionPrecommitToken
 
 	// txReadyOrClosed is for broadcasting that transaction ID has been returned
@@ -1473,7 +1474,8 @@ func (t *ReadWriteTransaction) getTransactionSelector() *sppb.TransactionSelecto
 			Begin: &sppb.TransactionOptions{
 				Mode: &sppb.TransactionOptions_ReadWrite_{
 					ReadWrite: &sppb.TransactionOptions_ReadWrite{
-						ReadLockMode: t.txOpts.ReadLockMode,
+						ReadLockMode:                            t.txOpts.ReadLockMode,
+						MultiplexedSessionPreviousTransactionId: t.previousTx,
 					},
 				},
 				ExcludeTxnFromChangeStreams: t.txOpts.ExcludeTxnFromChangeStreams,
@@ -1535,18 +1537,19 @@ func (t *ReadWriteTransaction) setSessionEligibilityForLongRunning(sh *sessionHa
 	}
 }
 
-func beginTransaction(ctx context.Context, sid string, client spannerClient, opts TransactionOptions, mutationKey *sppb.Mutation) (transactionID, *sppb.MultiplexedSessionPrecommitToken, error) {
-	res, err := client.BeginTransaction(ctx, &sppb.BeginTransactionRequest{
-		Session: sid,
+func beginTransaction(ctx context.Context, opts transactionBeginOptions) (transactionID, *sppb.MultiplexedSessionPrecommitToken, error) {
+	res, err := opts.client.BeginTransaction(ctx, &sppb.BeginTransactionRequest{
+		Session: opts.sessionID,
 		Options: &sppb.TransactionOptions{
 			Mode: &sppb.TransactionOptions_ReadWrite_{
 				ReadWrite: &sppb.TransactionOptions_ReadWrite{
-					ReadLockMode: opts.ReadLockMode,
+					ReadLockMode:                            opts.txOptions.ReadLockMode,
+					MultiplexedSessionPreviousTransactionId: opts.previousTx,
 				},
 			},
-			ExcludeTxnFromChangeStreams: opts.ExcludeTxnFromChangeStreams,
+			ExcludeTxnFromChangeStreams: opts.txOptions.ExcludeTxnFromChangeStreams,
 		},
-		MutationKey: mutationKey,
+		MutationKey: opts.mutation,
 	})
 	if err != nil {
 		return nil, nil, err
@@ -1581,6 +1584,7 @@ func (t *ReadWriteTransaction) begin(ctx context.Context, mutation *sppb.Mutatio
 		return nil
 	}
 	sh := t.sh
+	previousTx := t.previousTx
 	t.mu.Unlock()
 
 	var (
@@ -1604,7 +1608,13 @@ func (t *ReadWriteTransaction) begin(ctx context.Context, mutation *sppb.Mutatio
 		if sh != nil {
 			sh.updateLastUseTime()
 		}
-		tx, precommitToken, err = beginTransaction(contextWithOutgoingMetadata(ctx, sh.getMetadata(), t.disableRouteToLeader), sh.getID(), sh.getClient(), t.txOpts, mutation)
+		tx, precommitToken, err = beginTransaction(contextWithOutgoingMetadata(ctx, sh.getMetadata(), t.disableRouteToLeader), transactionBeginOptions{
+			sessionID:  sh.getID(),
+			client:     sh.getClient(),
+			txOptions:  t.txOpts,
+			mutation:   mutation,
+			previousTx: previousTx,
+		})
 		if isSessionNotFoundError(err) {
 			sh.destroy()
 			// this should not happen with multiplexed session, but if it does, we should not retry with multiplexed session
@@ -1791,6 +1801,9 @@ func (t *ReadWriteTransaction) runInTransaction(ctx context.Context, f func(cont
 		errDuringCommit = err != nil
 	}
 	if err != nil {
+		if t.tx != nil {
+			t.previousTx = t.tx
+		}
 		if isAbortedErr(err) {
 			// Retry the transaction using the same session on ABORT error.
 			// Cloud Spanner will create the new transaction with the previous
@@ -2073,3 +2086,12 @@ func isAbortedErr(err error) bool {
 	}
 	return false
 }
+
+// transactionBeginOptions holds the parameters for beginning a transaction.
+type transactionBeginOptions struct {
+	sessionID  string
+	client     spannerClient
+	txOptions  TransactionOptions
+	previousTx transactionID
+	mutation   *sppb.Mutation
+}

spanner/transaction_test.go

@@ -596,6 +596,131 @@ func TestCommitWithMultiplexedSessionRetry(t *testing.T) {
 	}
 }
 
+func TestClient_ReadWriteTransaction_PreviousTransactionID(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	cfg := SessionPoolConfig{
+		MinOpened:                     1,
+		MaxOpened:                     1,
+		enableMultiplexSession:        true,
+		enableMultiplexedSessionForRW: true,
+	}
+	server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{
+		DisableNativeMetrics: true,
+		SessionPoolConfig:    cfg,
+	})
+	defer teardown()
+
+	// First ExecuteSql will fail with InvalidArgument to force explicit BeginTransaction
+	invalidSQL := "Update FOO Set BAR=1"
+	server.TestSpanner.PutStatementResult(
+		invalidSQL,
+		&StatementResult{
+			Type: StatementResultError,
+			Err:  status.Error(codes.InvalidArgument, "Invalid update"),
+		},
+	)
+
+	// First Commit will fail with Aborted to force transaction retry
+	server.TestSpanner.PutExecutionTime(MethodCommitTransaction,
+		SimulatedExecutionTime{
+			Errors: []error{status.Error(codes.Aborted, "Transaction aborted")},
+		})
+
+	var attempts int
+	expectedAttempts := 4
+	_, err := client.ReadWriteTransaction(ctx, func(ctx context.Context, tx *ReadWriteTransaction) error {
+		if attempts == 1 || attempts == 3 {
+			// Replace the aborted result with a real result to prevent the
+			// transaction from aborting indefinitely.
+			server.TestSpanner.PutStatementResult(
+				invalidSQL,
+				&StatementResult{
+					Type:        StatementResultUpdateCount,
+					UpdateCount: 3,
+				},
+			)
+		}
+		if attempts == 2 {
+			server.TestSpanner.PutStatementResult(
+				invalidSQL,
+				&StatementResult{
+					Type: StatementResultError,
+					Err:  status.Error(codes.InvalidArgument, "Invalid update"),
+				},
+			)
+		}
+		attempts++
+		// First attempt: Inline begin fails due to invalid update
+		// Second attempt: Explicit begin succeeds but commit aborts
+		// Third attempt: Inline begin fails, previousTx set to txn2
+		// Fourth attempt: Explicit begin succeeds with previousTx=txn2
+		_, err := tx.Update(ctx, NewStatement(invalidSQL))
+		if err != nil {
+			return err
+		}
+		return nil
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if attempts != expectedAttempts {
+		t.Fatalf("got %d attempts, want %d", attempts, expectedAttempts)
+	}
+
+	requests := drainRequestsFromServer(server.TestSpanner)
+	var txID2 []byte
+	var foundPrevTxID bool
+
+	// Verify the sequence of requests and transaction IDs
+	for i, req := range requests {
+		switch r := req.(type) {
+		case *sppb.ExecuteSqlRequest:
+			// First and third attempts use inline begin
+			if i == 1 || i == 5 {
+				if _, ok := r.Transaction.GetSelector().(*sppb.TransactionSelector_Begin); !ok {
+					t.Errorf("Request %d: got %T, want Begin", i, r.Transaction.GetSelector())
+				}
+			}
+			if txID2 == nil && r.Transaction.GetId() != nil {
+				txID2 = r.Transaction.GetId()
+			}
+		case *sppb.BeginTransactionRequest:
+			if i == 7 {
+				opts := r.Options.GetReadWrite()
+				if opts == nil {
+					t.Fatal("Request 7: missing ReadWrite options")
+				}
+				if !testEqual(opts.MultiplexedSessionPreviousTransactionId, txID2) {
+					t.Errorf("Request 7: got prev txID %v, want %v",
+						opts.MultiplexedSessionPreviousTransactionId, txID2)
+				}
+				foundPrevTxID = true
+			}
+		}
+	}
+
+	if !foundPrevTxID {
+		t.Error("Did not find BeginTransaction request with previous transaction ID")
+	}
+
+	// Verify the complete sequence of requests
+	wantRequests := []interface{}{
+		&sppb.BatchCreateSessionsRequest{},
+		&sppb.ExecuteSqlRequest{},       // Attempt 1: Inline begin fails
+		&sppb.BeginTransactionRequest{}, // Attempt 2: Explicit begin
+		&sppb.ExecuteSqlRequest{},       // Attempt 2: Update succeeds
+		&sppb.CommitRequest{},           // Attempt 2: Commit aborts
+		&sppb.ExecuteSqlRequest{},       // Attempt 3: Inline begin fails
+		&sppb.BeginTransactionRequest{}, // Attempt 4: Explicit begin with prev txID
+		&sppb.ExecuteSqlRequest{},       // Attempt 4: Update succeeds
+		&sppb.CommitRequest{},           // Attempt 4: Commit succeeds
+	}
+	if err := compareRequestsWithConfig(wantRequests, requests, &cfg); err != nil {
+		t.Fatal(err)
+	}
+}
+
 func TestMutationOnlyCaseAborted(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()
@@ -1229,6 +1354,10 @@ func shouldHaveReceived(server InMemSpannerServer, want []interface{}) ([]interf
 
 // Compares expected requests (want) with actual requests (got).
 func compareRequests(want []interface{}, got []interface{}) error {
+	return compareRequestsWithConfig(want, got, nil)
+}
+
+func compareRequestsWithConfig(want []interface{}, got []interface{}, config *SessionPoolConfig) error {
 	if reflect.TypeOf(want[0]) != reflect.TypeOf(&sppb.BatchCreateSessionsRequest{}) {
 		sessReq := 0
 		for i := 0; i < len(want); i++ {
@@ -1239,7 +1368,7 @@ func compareRequests(want []interface{}, got []interface{}) error {
 		}
 		want[0], want[sessReq] = want[sessReq], want[0]
 	}
-	if isMultiplexEnabled {
+	if isMultiplexEnabled || (config != nil && config.enableMultiplexSession) {
 		if reflect.TypeOf(want[0]) != reflect.TypeOf(&sppb.CreateSessionRequest{}) {
 			want = append([]interface{}{&sppb.CreateSessionRequest{}}, want...)
 		}