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

This PR introduces changes to support the lock order preservation protocol for multiplexed sessions in read/write transactions. According to this protocol, when a read/write transaction on a multiplexed session is retried, the transaction ID from the previous abort must be passed when creating a new transaction. This ensures that the retried transaction is recognized as older, rather than being treated as a new one.

The transaction context object is structured as follows,
```
txn {
    transactionId: The transaction ID.
    previousTransactionId: The transaction ID of the most recently failed transaction.
}
```

Author: harshachinta

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java

@@ -28,6 +28,7 @@
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
 import com.google.common.util.concurrent.MoreExecutors;
+import com.google.protobuf.ByteString;
 
 /** Implementation of {@link AsyncTransactionManager}. */
 final class AsyncTransactionManagerImpl
@@ -80,7 +81,19 @@ public TransactionContextFutureImpl beginAsync() {
 
   private ApiFuture<TransactionContext> internalBeginAsync(boolean firstAttempt) {
     txnState = TransactionState.STARTED;
-    txn = session.newTransaction(options);
+
+    // Determine the latest transactionId when using a multiplexed session.
+    ByteString multiplexedSessionPreviousTransactionId = ByteString.EMPTY;
+    if (txn != null && session.getIsMultiplexed() && !firstAttempt) {
+      // Use the current transactionId if available, otherwise fallback to the previous aborted
+      // transactionId.
+      multiplexedSessionPreviousTransactionId =
+          txn.transactionId != null ? txn.transactionId : txn.getPreviousTransactionId();
+    }
+
+    txn =
+        session.newTransaction(
+            options, /* previousTransactionId = */ multiplexedSessionPreviousTransactionId);
     if (firstAttempt) {
       session.setActive(this);
     }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java

@@ -59,7 +59,7 @@ class DatabaseClientImpl implements DatabaseClient {
         /* useMultiplexedSessionBlindWrite = */ false,
         /* multiplexedSessionDatabaseClient = */ null,
         tracer,
-        false);
+        /* useMultiplexedSessionForRW = */ false);
   }
 
   DatabaseClientImpl(

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java

@@ -69,7 +69,8 @@ static void throwIfTransactionsPending() {
     }
   }
 
-  static TransactionOptions createReadWriteTransactionOptions(Options options) {
+  static TransactionOptions createReadWriteTransactionOptions(
+      Options options, ByteString previousTransactionId) {
     TransactionOptions.Builder transactionOptions = TransactionOptions.newBuilder();
     if (options.withExcludeTxnFromChangeStreams() == Boolean.TRUE) {
       transactionOptions.setExcludeTxnFromChangeStreams(true);
@@ -78,6 +79,10 @@ static TransactionOptions createReadWriteTransactionOptions(Options options) {
     if (options.withOptimisticLock() == Boolean.TRUE) {
       readWrite.setReadLockMode(TransactionOptions.ReadWrite.ReadLockMode.OPTIMISTIC);
     }
+    if (previousTransactionId != null
+        && previousTransactionId != com.google.protobuf.ByteString.EMPTY) {
+      readWrite.setMultiplexedSessionPreviousTransactionId(previousTransactionId);
+    }
     transactionOptions.setReadWrite(readWrite);
     return transactionOptions.build();
   }
@@ -427,13 +432,17 @@ public void close() {
   }
 
   ApiFuture<ByteString> beginTransactionAsync(
-      Options transactionOptions, boolean routeToLeader, Map<SpannerRpc.Option, ?> channelHint) {
+      Options transactionOptions,
+      boolean routeToLeader,
+      Map<SpannerRpc.Option, ?> channelHint,
+      ByteString previousTransactionId) {
     final SettableApiFuture<ByteString> res = SettableApiFuture.create();
     final ISpan span = tracer.spanBuilder(SpannerImpl.BEGIN_TRANSACTION);
     final BeginTransactionRequest request =
         BeginTransactionRequest.newBuilder()
             .setSession(getName())
-            .setOptions(createReadWriteTransactionOptions(transactionOptions))
+            .setOptions(
+                createReadWriteTransactionOptions(transactionOptions, previousTransactionId))
             .build();
     final ApiFuture<Transaction> requestFuture;
     try (IScope ignore = tracer.withSpan(span)) {
@@ -469,11 +478,12 @@ ApiFuture<ByteString> beginTransactionAsync(
     return res;
   }
 
-  TransactionContextImpl newTransaction(Options options) {
+  TransactionContextImpl newTransaction(Options options, ByteString previousTransactionId) {
     return TransactionContextImpl.newBuilder()
         .setSession(this)
         .setOptions(options)
         .setTransactionId(null)
+        .setPreviousTransactionId(previousTransactionId)
         .setOptions(options)
         .setTrackTransactionStarter(spanner.getOptions().isTrackTransactionStarter())
         .setRpc(spanner.getRpc())

google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java

@@ -20,6 +20,7 @@
 import com.google.cloud.spanner.Options.TransactionOption;
 import com.google.cloud.spanner.SessionImpl.SessionTransaction;
 import com.google.common.base.Preconditions;
+import com.google.protobuf.ByteString;
 
 /** Implementation of {@link TransactionManager}. */
 final class TransactionManagerImpl implements TransactionManager, SessionTransaction {
@@ -53,7 +54,7 @@ public void setSpan(ISpan span) {
   public TransactionContext begin() {
     Preconditions.checkState(txn == null, "begin can only be called once");
     try (IScope s = tracer.withSpan(span)) {
-      txn = session.newTransaction(options);
+      txn = session.newTransaction(options, /* previousTransactionId = */ ByteString.EMPTY);
       session.setActive(this);
       txnState = TransactionState.STARTED;
       return txn;
@@ -102,7 +103,18 @@ public TransactionContext resetForRetry() {
     }
     try (IScope s = tracer.withSpan(span)) {
       boolean useInlinedBegin = txn.transactionId != null;
-      txn = session.newTransaction(options);
+
+      // Determine the latest transactionId when using a multiplexed session.
+      ByteString multiplexedSessionPreviousTransactionId = ByteString.EMPTY;
+      if (session.getIsMultiplexed()) {
+        // Use the current transactionId if available, otherwise fallback to the previous aborted
+        // transactionId.
+        multiplexedSessionPreviousTransactionId =
+            txn.transactionId != null ? txn.transactionId : txn.getPreviousTransactionId();
+      }
+      txn =
+          session.newTransaction(
+              options, /* previousTransactionId = */ multiplexedSessionPreviousTransactionId);
       if (!useInlinedBegin) {
         txn.ensureTxn();
       }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java

@@ -93,6 +93,9 @@ static class Builder extends AbstractReadContext.Builder<Builder, TransactionCon
 
       private Clock clock = new Clock();
       private ByteString transactionId;
+      // This field is set only when the transaction is created during a retry and uses a
+      // multiplexed session.
+      private ByteString previousTransactionId;
       private Options options;
       private boolean trackTransactionStarter;
 
@@ -118,6 +121,11 @@ Builder setTrackTransactionStarter(boolean trackTransactionStarter) {
         return self();
       }
 
+      Builder setPreviousTransactionId(ByteString previousTransactionId) {
+        this.previousTransactionId = previousTransactionId;
+        return self();
+      }
+
       @Override
       TransactionContextImpl build() {
         Preconditions.checkState(this.options != null, "Options must be set");
@@ -201,6 +209,8 @@ public void removeListener(Runnable listener) {
 
     volatile ByteString transactionId;
 
+    final ByteString previousTransactionId;
+
     private CommitResponse commitResponse;
     private final Clock clock;
 
@@ -216,6 +226,7 @@ private TransactionContextImpl(Builder builder) {
       this.channelHint =
           getChannelHintOptions(
               session.getOptions(), ThreadLocalRandom.current().nextLong(Long.MAX_VALUE));
+      this.previousTransactionId = builder.previousTransactionId;
     }
 
     @Override
@@ -246,6 +257,10 @@ private void decreaseAsyncOperations() {
       }
     }
 
+    ByteString getPreviousTransactionId() {
+      return this.previousTransactionId;
+    }
+
     @Override
     public void close() {
       // Only mark the context as closed, but do not end the tracer span, as that is done by the
@@ -283,7 +298,8 @@ ApiFuture<Void> ensureTxnAsync() {
     private void createTxnAsync(final SettableApiFuture<Void> res) {
       span.addAnnotation("Creating Transaction");
       final ApiFuture<ByteString> fut =
-          session.beginTransactionAsync(options, isRouteToLeader(), getTransactionChannelHint());
+          session.beginTransactionAsync(
+              options, isRouteToLeader(), getTransactionChannelHint(), getPreviousTransactionId());
       fut.addListener(
           () -> {
             try {
@@ -558,7 +574,9 @@ TransactionSelector getTransactionSelector() {
           }
           if (tx == null) {
             return TransactionSelector.newBuilder()
-                .setBegin(SessionImpl.createReadWriteTransactionOptions(options))
+                .setBegin(
+                    SessionImpl.createReadWriteTransactionOptions(
+                        options, getPreviousTransactionId()))
                 .build();
           } else {
             // Wait for the transaction to come available. The tx.get() call will fail with an
@@ -1079,7 +1097,7 @@ public TransactionRunner allowNestedTransaction() {
   TransactionRunnerImpl(SessionImpl session, TransactionOption... options) {
     this.session = session;
     this.options = Options.fromTransactionOptions(options);
-    this.txn = session.newTransaction(this.options);
+    this.txn = session.newTransaction(this.options, /* previousTransactionId = */ ByteString.EMPTY);
     this.tracer = session.getTracer();
   }
 
@@ -1118,7 +1136,19 @@ private <T> T runInternal(final TransactionCallable<T> txCallable) {
             // Do not inline the BeginTransaction during a retry if the initial attempt did not
             // actually start a transaction.
             useInlinedBegin = txn.transactionId != null;
-            txn = session.newTransaction(options);
+
+            // Determine the latest transactionId when using a multiplexed session.
+            ByteString multiplexedSessionPreviousTransactionId = ByteString.EMPTY;
+            if (session.getIsMultiplexed()) {
+              // Use the current transactionId if available, otherwise fallback to the previous
+              // transactionId.
+              multiplexedSessionPreviousTransactionId =
+                  txn.transactionId != null ? txn.transactionId : txn.getPreviousTransactionId();
+            }
+
+            txn =
+                session.newTransaction(
+                    options, /* previousTransactionId = */ multiplexedSessionPreviousTransactionId);
           }
           checkState(
               isValid, "TransactionRunner has been invalidated by a new operation on the session");

google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerImplTest.java

@@ -16,12 +16,18 @@
 
 package com.google.cloud.spanner;
 
+import static org.junit.Assert.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.clearInvocations;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import com.google.api.core.ApiFutures;
 import com.google.cloud.Timestamp;
+import com.google.protobuf.ByteString;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Scope;
 import org.junit.Test;
@@ -42,7 +48,7 @@ public void testCommitReturnsCommitStats() {
     when(oTspan.makeCurrent()).thenReturn(mock(Scope.class));
     try (AsyncTransactionManagerImpl manager =
         new AsyncTransactionManagerImpl(session, span, Options.commitStats())) {
-      when(session.newTransaction(Options.fromTransactionOptions(Options.commitStats())))
+      when(session.newTransaction(eq(Options.fromTransactionOptions(Options.commitStats())), any()))
           .thenReturn(transaction);
       when(transaction.ensureTxnAsync()).thenReturn(ApiFutures.immediateFuture(null));
       Timestamp commitTimestamp = Timestamp.ofTimeMicroseconds(1);
@@ -54,4 +60,67 @@ public void testCommitReturnsCommitStats() {
       verify(transaction).commitAsync();
     }
   }
+
+  @Test
+  public void testRetryUsesPreviousTransactionIdOnMultiplexedSession() {
+    // Set up mock transaction IDs
+    final ByteString mockTransactionId = ByteString.copyFromUtf8("mockTransactionId");
+    final ByteString mockPreviousTransactionId =
+        ByteString.copyFromUtf8("mockPreviousTransactionId");
+
+    Span oTspan = mock(Span.class);
+    ISpan span = new OpenTelemetrySpan(oTspan);
+    when(oTspan.makeCurrent()).thenReturn(mock(Scope.class));
+    // Mark the session as multiplexed.
+    when(session.getIsMultiplexed()).thenReturn(true);
+
+    // Initialize a mock transaction with transactionId = null, previousTransactionId = null.
+    transaction = mock(TransactionRunnerImpl.TransactionContextImpl.class);
+    when(transaction.ensureTxnAsync()).thenReturn(ApiFutures.immediateFuture(null));
+    when(session.newTransaction(eq(Options.fromTransactionOptions(Options.commitStats())), any()))
+        .thenReturn(transaction);
+
+    // Simulate an ABORTED error being thrown when `commitAsync()` is called.
+    doThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.ABORTED, ""))
+        .when(transaction)
+        .commitAsync();
+
+    try (AsyncTransactionManagerImpl manager =
+        new AsyncTransactionManagerImpl(session, span, Options.commitStats())) {
+      manager.beginAsync();
+
+      // Verify that for the first transaction attempt, the `previousTransactionId` is
+      // ByteString.EMPTY.
+      // This is because no transaction has been previously aborted at this point.
+      verify(session)
+          .newTransaction(Options.fromTransactionOptions(Options.commitStats()), ByteString.EMPTY);
+      assertThrows(AbortedException.class, manager::commitAsync);
+      clearInvocations(session);
+
+      // Mock the transaction object to contain transactionID=null and
+      // previousTransactionId=mockPreviousTransactionId
+      when(transaction.getPreviousTransactionId()).thenReturn(mockPreviousTransactionId);
+      manager.resetForRetryAsync();
+      // Verify that in the first retry attempt, the `previousTransactionId`
+      // (mockPreviousTransactionId) is passed to the new transaction.
+      // This allows Spanner to retry the transaction using the ID of the aborted transaction.
+      verify(session)
+          .newTransaction(
+              Options.fromTransactionOptions(Options.commitStats()), mockPreviousTransactionId);
+      assertThrows(AbortedException.class, manager::commitAsync);
+      clearInvocations(session);
+
+      // Mock the transaction object to contain transactionID=mockTransactionId and
+      // previousTransactionId=mockPreviousTransactionId and transactionID = null
+      transaction.transactionId = mockTransactionId;
+      manager.resetForRetryAsync();
+      // Verify that the latest `transactionId` (mockTransactionId) is used in the retry.
+      // This ensures the retry logic is working as expected with the latest transaction ID.
+      verify(session)
+          .newTransaction(Options.fromTransactionOptions(Options.commitStats()), mockTransactionId);
+
+      when(transaction.rollbackAsync()).thenReturn(ApiFutures.immediateFuture(null));
+      manager.closeAsync();
+    }
+  }
 }