feat(spanner): support multiplexed sessions for ReadWriteStmtBasedTransaction (#11852)

This PR adds support for multiplexed sessions in ReadWriteStmtBasedTransaction

Author: harshachinta

spanner/transaction.go

@@ -1910,16 +1910,20 @@ func NewReadWriteStmtBasedTransaction(ctx context.Context, c *Client) (*ReadWrit
 // NewReadWriteStmtBasedTransactionWithOptions is a configurable version of
 // NewReadWriteStmtBasedTransaction.
 func NewReadWriteStmtBasedTransactionWithOptions(ctx context.Context, c *Client, options TransactionOptions) (*ReadWriteStmtBasedTransaction, error) {
-	return newReadWriteStmtBasedTransactionWithSessionHandle(ctx, c, options, nil)
+	return newReadWriteStmtBasedTransactionWithSessionHandle(ctx, c, options, nil, nil)
 }
 
-func newReadWriteStmtBasedTransactionWithSessionHandle(ctx context.Context, c *Client, options TransactionOptions, sh *sessionHandle) (*ReadWriteStmtBasedTransaction, error) {
+func newReadWriteStmtBasedTransactionWithSessionHandle(ctx context.Context, c *Client, options TransactionOptions, sh *sessionHandle, previousTransactionID transactionID) (*ReadWriteStmtBasedTransaction, error) {
 	var (
 		err error
 		t   *ReadWriteStmtBasedTransaction
 	)
 	if sh == nil {
-		sh, err = c.idleSessions.take(ctx)
+		if c.idleSessions.isMultiplexedSessionForRWEnabled() {
+			sh, err = c.idleSessions.takeMultiplexed(ctx)
+		} else {
+			sh, err = c.idleSessions.take(ctx)
+		}
 		if err != nil {
 			// If session retrieval fails, just fail the transaction.
 			return nil, err
@@ -1931,6 +1935,12 @@ func newReadWriteStmtBasedTransactionWithSessionHandle(ctx context.Context, c *C
 		},
 		client: c,
 	}
+	if previousTransactionID != nil {
+		// The previousTx field is updated with the most recent transaction ID. This is needed for multiplexed sessions
+		// to increase the priority of the new transaction during retry attempt.
+		// This assignment is ignored for regular sessions.
+		t.previousTx = previousTransactionID
+	}
 	t.txReadOnly.sp = c.idleSessions
 	t.txReadOnly.sh = sh
 	t.txReadOnly.txReadEnv = t
@@ -1958,6 +1968,9 @@ func newReadWriteStmtBasedTransactionWithSessionHandle(ctx context.Context, c *C
 		}
 		return nil, err
 	}
+	if isUnimplementedErrorForMultiplexedRW(err) {
+		c.idleSessions.disableMultiplexedSessionForRW()
+	}
 	return t, err
 }
 
@@ -2003,8 +2016,17 @@ func (t *ReadWriteStmtBasedTransaction) ResetForRetry(ctx context.Context) (*Rea
 	if t.state != txAborted {
 		return nil, fmt.Errorf("ResetForRetry should only be called on an active transaction that was aborted by Spanner")
 	}
+
+	var previousTransactionID transactionID
+	if t.tx != nil {
+		// Track the current transactionId that is ABORTED.
+		previousTransactionID = t.tx
+	} else {
+		// In case the current transactionId is nil, then look at the previousTx.
+		previousTransactionID = t.previousTx
+	}
 	// Create a new transaction that re-uses the current session if it is available.
-	return newReadWriteStmtBasedTransactionWithSessionHandle(ctx, t.client, t.options, t.sh)
+	return newReadWriteStmtBasedTransactionWithSessionHandle(ctx, t.client, t.options, t.sh, previousTransactionID)
 }
 
 // writeOnlyTransaction provides the most efficient way of doing write-only

spanner/transaction_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package spanner
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
@@ -1199,6 +1200,165 @@ func TestReadWriteStmtBasedTransactionWithOptions(t *testing.T) {
 	}
 }
 
+// Verify that requests in a ReadWriteStmtBasedTransaction uses multiplexed sessions when enabled
+func TestReadWriteStmtBasedTransaction_UsesMultiplexedSession(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	cfg := SessionPoolConfig{
+		// Avoid regular session creation
+		MinOpened: 0,
+		MaxOpened: 0,
+		// Enabled multiplexed sessions
+		enableMultiplexSession:        true,
+		enableMultiplexedSessionForRW: true,
+	}
+	server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{
+		SessionPoolConfig: cfg,
+	})
+	defer teardown()
+	server.TestSpanner.PutExecutionTime(MethodCommitTransaction,
+		SimulatedExecutionTime{
+			Errors: []error{status.Errorf(codes.Aborted, "Transaction aborted")},
+		})
+
+	txn, err := NewReadWriteStmtBasedTransaction(ctx, client)
+	if err != nil {
+		t.Fatalf("got an error: %v", err)
+	}
+	_, err = txn.Commit(ctx)
+	if status.Code(err) != codes.Aborted || !strings.Contains(err.Error(), "Transaction aborted") {
+		t.Fatalf("got an incorrect error: %v", err)
+	}
+
+	requests := drainRequestsFromServer(server.TestSpanner)
+	if err := compareRequests([]interface{}{
+		&sppb.CreateSessionRequest{},
+		&sppb.BeginTransactionRequest{},
+		&sppb.CommitRequest{}}, requests); err != nil {
+		t.Fatal(err)
+	}
+	for _, req := range requests {
+		if c, ok := req.(*sppb.CommitRequest); ok {
+			if !strings.Contains(c.GetSession(), "multiplexed") {
+				t.Errorf("Expected session to be multiplexed")
+			}
+		}
+		if b, ok := req.(*sppb.BeginTransactionRequest); ok {
+			if !strings.Contains(b.GetSession(), "multiplexed") {
+				t.Errorf("Expected session to be multiplexed")
+			}
+		}
+	}
+}
+
+// Verify that in ReadWriteStmtBasedTransaction when a transaction using a multiplexed session fails with ABORTED then during retry the previousTransactionID is passed
+func TestReadWriteStmtBasedTransaction_UsesPreviousTransactionIDForMultiplexedSession_OnAbort(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	cfg := SessionPoolConfig{
+		// Avoid regular session creation
+		MinOpened: 0,
+		MaxOpened: 0,
+		// Enabled multiplexed sessions
+		enableMultiplexSession:        true,
+		enableMultiplexedSessionForRW: true,
+	}
+	server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{
+		SessionPoolConfig: cfg,
+	})
+	defer teardown()
+	server.TestSpanner.PutExecutionTime(MethodCommitTransaction,
+		SimulatedExecutionTime{
+			Errors: []error{status.Errorf(codes.Aborted, "Transaction aborted")},
+		})
+
+	txn, err := NewReadWriteStmtBasedTransaction(ctx, client)
+	if err != nil {
+		t.Fatalf("got an error: %v", err)
+	}
+	_, err = txn.Commit(ctx)
+	if status.Code(err) != codes.Aborted || !strings.Contains(err.Error(), "Transaction aborted") {
+		t.Fatalf("got an incorrect error: %v", err)
+	}
+	// Fetch the transaction ID of the first attempt
+	previousTransactionID := txn.tx
+	txn, err = txn.ResetForRetry(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, err = txn.Commit(ctx)
+	if err != nil {
+		t.Fatalf("expect commit to succeed but failed with error: %v", err)
+	}
+	requests := drainRequestsFromServer(server.TestSpanner)
+	beginTransactionRequestCount := 0
+	for _, req := range requests {
+		if b, ok := req.(*sppb.BeginTransactionRequest); ok {
+			beginTransactionRequestCount++
+			// The first BeginTransactionRequest will not have any previousTransactionID.
+			// Check if the second BeginTransactionRequest sets the previousTransactionID to correct value.
+			if beginTransactionRequestCount == 2 {
+				if !strings.Contains(b.GetSession(), "multiplexed") {
+					t.Errorf("Expected session to be multiplexed")
+				}
+				opts := b.Options.GetReadWrite()
+				if opts == nil {
+					t.Fatal("missing ReadWrite options")
+				}
+				if !bytes.Equal(opts.MultiplexedSessionPreviousTransactionId, previousTransactionID) {
+					t.Errorf("BeginTransactionRequest during retry: got prev txID %v, want %v",
+						opts.MultiplexedSessionPreviousTransactionId, previousTransactionID)
+				}
+			}
+		}
+	}
+}
+
+// Verify that in ReadWriteStmtBasedTransaction, commit request has precommit token set when using multiplexed sessions
+func TestReadWriteStmtBasedTransaction_SetsPrecommitToken(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	cfg := SessionPoolConfig{
+		// Avoid regular session creation
+		MinOpened: 0,
+		MaxOpened: 0,
+		// Enabled multiplexed sessions
+		enableMultiplexSession:        true,
+		enableMultiplexedSessionForRW: true,
+	}
+	server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{
+		SessionPoolConfig: cfg,
+	})
+	defer teardown()
+
+	txn, err := NewReadWriteStmtBasedTransaction(ctx, client)
+	if err != nil {
+		t.Fatalf("got an error: %v", err)
+	}
+	c, err := txn.Update(ctx, Statement{SQL: UpdateBarSetFoo})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if g, w := c, int64(UpdateBarSetFooRowCount); g != w {
+		t.Fatalf("update count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	_, err = txn.Commit(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	requests := drainRequestsFromServer(server.TestSpanner)
+	for _, req := range requests {
+		if commitReq, ok := req.(*sppb.CommitRequest); ok {
+			if commitReq.GetPrecommitToken() == nil || !strings.Contains(string(commitReq.GetPrecommitToken().PrecommitToken), "ResultSetPrecommitToken") {
+				t.Errorf("Expected precommit token 'ResultSetPrecommitToken', got %v", commitReq.GetPrecommitToken())
+			}
+			if !strings.Contains(commitReq.GetSession(), "multiplexed") {
+				t.Errorf("Expected session to be multiplexed")
+			}
+		}
+	}
+}
+
 func TestBatchDML_StatementBased_WithMultipleDML(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()