chore(spanner): refactor code for lock order

Author: harshachinta

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

@@ -78,9 +78,17 @@ public TransactionContextFutureImpl beginAsync() {
 
   private ApiFuture<TransactionContext> internalBeginAsync(boolean firstAttempt) {
     txnState = TransactionState.STARTED;
-    ByteString previousAbortedTransactionId =
-        !firstAttempt && session.getIsMultiplexed() ? txn.transactionId : null;
-    txn = session.newTransaction(options, previousAbortedTransactionId);
+
+    // Determine the latest transactionId when using a multiplexed session.
+    ByteString multiplexedSessionPreviousTransactionId = null;
+    if (session.getIsMultiplexed() && !firstAttempt) {
+      // Use the current transactionId if available, otherwise fallback to the previous aborted
+      // transactionId.
+      multiplexedSessionPreviousTransactionId =
+          txn.transactionId != null ? txn.transactionId : txn.previousTransactionId;
+    }
+
+    txn = session.newTransaction(options, multiplexedSessionPreviousTransactionId);
     if (firstAttempt) {
       session.setActive(this);
     }

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

@@ -70,7 +70,7 @@ static void throwIfTransactionsPending() {
   }
 
   static TransactionOptions createReadWriteTransactionOptions(
-      Options options, ByteString previousAbortedTransactionId) {
+      Options options, ByteString previousTransactionId) {
     TransactionOptions.Builder transactionOptions = TransactionOptions.newBuilder();
     if (options.withExcludeTxnFromChangeStreams() == Boolean.TRUE) {
       transactionOptions.setExcludeTxnFromChangeStreams(true);
@@ -79,10 +79,10 @@ static TransactionOptions createReadWriteTransactionOptions(
     if (options.withOptimisticLock() == Boolean.TRUE) {
       readWrite.setReadLockMode(TransactionOptions.ReadWrite.ReadLockMode.OPTIMISTIC);
     }
-    if (previousAbortedTransactionId != null
-        && previousAbortedTransactionId != com.google.protobuf.ByteString.EMPTY) {
+    if (previousTransactionId != null
+        && previousTransactionId != com.google.protobuf.ByteString.EMPTY) {
       // TODO(sriharshach): uncomment this when multiplexed session R/W proto is published
-      // readWrite.setMultiplexedSessionPreviousTransactionId(previousAbortedTransactionId);
+      // readWrite.setMultiplexedSessionPreviousTransactionId(previousTransactionId);
     }
     transactionOptions.setReadWrite(readWrite);
     return transactionOptions.build();
@@ -436,14 +436,14 @@ ApiFuture<ByteString> beginTransactionAsync(
       Options transactionOptions,
       boolean routeToLeader,
       Map<SpannerRpc.Option, ?> channelHint,
-      ByteString previousAbortedTransactionId) {
+      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, previousAbortedTransactionId))
+                createReadWriteTransactionOptions(transactionOptions, previousTransactionId))
             .build();
     final ApiFuture<Transaction> requestFuture;
     try (IScope ignore = tracer.withSpan(span)) {
@@ -479,12 +479,12 @@ ApiFuture<ByteString> beginTransactionAsync(
     return res;
   }
 
-  TransactionContextImpl newTransaction(Options options, ByteString previousAbortedTransactionId) {
+  TransactionContextImpl newTransaction(Options options, ByteString previousTransactionId) {
     return TransactionContextImpl.newBuilder()
         .setSession(this)
         .setOptions(options)
         .setTransactionId(null)
-        .setPreviousAbortedTransactionId(previousAbortedTransactionId)
+        .setPreviousTransactionId(previousTransactionId)
         .setOptions(options)
         .setTrackTransactionStarter(spanner.getOptions().isTrackTransactionStarter())
         .setRpc(spanner.getRpc())

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

@@ -104,9 +104,15 @@ public TransactionContext resetForRetry() {
     try (IScope s = tracer.withSpan(span)) {
       boolean useInlinedBegin = txn.transactionId != null;
 
-      ByteString previousAbortedTransactionId =
-          session.getIsMultiplexed() ? txn.transactionId : null;
-      txn = session.newTransaction(options, previousAbortedTransactionId);
+      // Determine the latest transactionId when using a multiplexed session.
+      ByteString multiplexedSessionPreviousTransactionId = null;
+      if (session.getIsMultiplexed()) {
+        // Use the current transactionId if available, otherwise fallback to the previous aborted
+        // transactionId.
+        multiplexedSessionPreviousTransactionId =
+            txn.transactionId != null ? txn.transactionId : txn.previousTransactionId;
+      }
+      txn = session.newTransaction(options, multiplexedSessionPreviousTransactionId);
       if (!useInlinedBegin) {
         txn.ensureTxn();
       }

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

@@ -95,7 +95,7 @@ static class Builder extends AbstractReadContext.Builder<Builder, TransactionCon
       private ByteString transactionId;
       // This field is set only when the transaction is created during a retry and uses a
       // multiplexed session.
-      private ByteString previousAbortedTransactionId;
+      private ByteString previousTransactionId;
       private Options options;
       private boolean trackTransactionStarter;
 
@@ -121,8 +121,8 @@ Builder setTrackTransactionStarter(boolean trackTransactionStarter) {
         return self();
       }
 
-      Builder setPreviousAbortedTransactionId(ByteString previousAbortedTransactionId) {
-        this.previousAbortedTransactionId = previousAbortedTransactionId;
+      Builder setPreviousTransactionId(ByteString previousTransactionId) {
+        this.previousTransactionId = previousTransactionId;
         return self();
       }
 
@@ -209,7 +209,7 @@ public void removeListener(Runnable listener) {
 
     volatile ByteString transactionId;
 
-    ByteString previousAbortedTransactionId;
+    ByteString previousTransactionId;
 
     private CommitResponse commitResponse;
     private final Clock clock;
@@ -226,7 +226,7 @@ private TransactionContextImpl(Builder builder) {
       this.channelHint =
           getChannelHintOptions(
               session.getOptions(), ThreadLocalRandom.current().nextLong(Long.MAX_VALUE));
-      this.previousAbortedTransactionId = builder.previousAbortedTransactionId;
+      this.previousTransactionId = builder.previousTransactionId;
     }
 
     @Override
@@ -295,10 +295,7 @@ private void createTxnAsync(final SettableApiFuture<Void> res) {
       span.addAnnotation("Creating Transaction");
       final ApiFuture<ByteString> fut =
           session.beginTransactionAsync(
-              options,
-              isRouteToLeader(),
-              getTransactionChannelHint(),
-              previousAbortedTransactionId);
+              options, isRouteToLeader(), getTransactionChannelHint(), previousTransactionId);
       fut.addListener(
           () -> {
             try {
@@ -574,8 +571,7 @@ TransactionSelector getTransactionSelector() {
           if (tx == null) {
             return TransactionSelector.newBuilder()
                 .setBegin(
-                    SessionImpl.createReadWriteTransactionOptions(
-                        options, previousAbortedTransactionId))
+                    SessionImpl.createReadWriteTransactionOptions(options, previousTransactionId))
                 .build();
           } else {
             // Wait for the transaction to come available. The tx.get() call will fail with an
@@ -1135,9 +1131,16 @@ private <T> T runInternal(final TransactionCallable<T> txCallable) {
             // actually start a transaction.
             useInlinedBegin = txn.transactionId != null;
 
-            ByteString previousAbortedTransactionId =
-                session.getIsMultiplexed() ? txn.transactionId : null;
-            txn = session.newTransaction(options, previousAbortedTransactionId);
+            // Determine the latest transactionId when using a multiplexed session.
+            ByteString multiplexedSessionPreviousTransactionId = null;
+            if (session.getIsMultiplexed()) {
+              // Use the current transactionId if available, otherwise fallback to the previous
+              // transactionId.
+              multiplexedSessionPreviousTransactionId =
+                  txn.transactionId != null ? txn.transactionId : txn.previousTransactionId;
+            }
+
+            txn = session.newTransaction(options, multiplexedSessionPreviousTransactionId);
           }
           checkState(
               isValid, "TransactionRunner has been invalidated by a new operation on the session");