chore(spanner): track previous transaction id incase of retry during aborted 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
@@ -77,7 +78,9 @@ public TransactionContextFutureImpl beginAsync() {
 
   private ApiFuture<TransactionContext> internalBeginAsync(boolean firstAttempt) {
     txnState = TransactionState.STARTED;
-    txn = session.newTransaction(options);
+    ByteString previousAbortedTransactionId =
+        !firstAttempt && session.getIsMultiplexed() ? txn.transactionId : null;
+    txn = session.newTransaction(options, previousAbortedTransactionId);
     if (firstAttempt) {
       session.setActive(this);
     }

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 previousAbortedTransactionId) {
     TransactionOptions.Builder transactionOptions = TransactionOptions.newBuilder();
     if (options.withExcludeTxnFromChangeStreams() == Boolean.TRUE) {
       transactionOptions.setExcludeTxnFromChangeStreams(true);
@@ -78,6 +79,11 @@ static TransactionOptions createReadWriteTransactionOptions(Options options) {
     if (options.withOptimisticLock() == Boolean.TRUE) {
       readWrite.setReadLockMode(TransactionOptions.ReadWrite.ReadLockMode.OPTIMISTIC);
     }
+    if (previousAbortedTransactionId != null
+        && previousAbortedTransactionId != com.google.protobuf.ByteString.EMPTY) {
+      // TODO(sriharshach): uncomment this when multiplexed session R/W proto is published
+      // readWrite.setMultiplexedSessionPreviousTransactionId(previousAbortedTransactionId);
+    }
     transactionOptions.setReadWrite(readWrite);
     return transactionOptions.build();
   }
@@ -427,13 +433,17 @@ public void close() {
   }
 
   ApiFuture<ByteString> beginTransactionAsync(
-      Options transactionOptions, boolean routeToLeader, Map<SpannerRpc.Option, ?> channelHint) {
+      Options transactionOptions,
+      boolean routeToLeader,
+      Map<SpannerRpc.Option, ?> channelHint,
+      ByteString previousAbortedTransactionId) {
     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, previousAbortedTransactionId))
             .build();
     final ApiFuture<Transaction> requestFuture;
     try (IScope ignore = tracer.withSpan(span)) {
@@ -469,11 +479,12 @@ ApiFuture<ByteString> beginTransactionAsync(
     return res;
   }
 
-  TransactionContextImpl newTransaction(Options options) {
+  TransactionContextImpl newTransaction(Options options, ByteString previousAbortedTransactionId) {
     return TransactionContextImpl.newBuilder()
         .setSession(this)
         .setOptions(options)
         .setTransactionId(null)
+        .setPreviousAbortedTransactionId(previousAbortedTransactionId)
         .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, null);
       session.setActive(this);
       txnState = TransactionState.STARTED;
       return txn;
@@ -102,7 +103,10 @@ public TransactionContext resetForRetry() {
     }
     try (IScope s = tracer.withSpan(span)) {
       boolean useInlinedBegin = txn.transactionId != null;
-      txn = session.newTransaction(options);
+
+      ByteString previousAbortedTransactionId =
+          session.getIsMultiplexed() ? txn.transactionId : null;
+      txn = session.newTransaction(options, previousAbortedTransactionId);
       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 previousAbortedTransactionId;
       private Options options;
       private boolean trackTransactionStarter;
 
@@ -118,6 +121,11 @@ Builder setTrackTransactionStarter(boolean trackTransactionStarter) {
         return self();
       }
 
+      Builder setPreviousAbortedTransactionId(ByteString previousAbortedTransactionId) {
+        this.previousAbortedTransactionId = previousAbortedTransactionId;
+        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 previousAbortedTransactionId;
+
     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.previousAbortedTransactionId = builder.previousAbortedTransactionId;
     }
 
     @Override
@@ -283,7 +294,11 @@ 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(),
+              previousAbortedTransactionId);
       fut.addListener(
           () -> {
             try {
@@ -558,7 +573,9 @@ TransactionSelector getTransactionSelector() {
           }
           if (tx == null) {
             return TransactionSelector.newBuilder()
-                .setBegin(SessionImpl.createReadWriteTransactionOptions(options))
+                .setBegin(
+                    SessionImpl.createReadWriteTransactionOptions(
+                        options, previousAbortedTransactionId))
                 .build();
           } else {
             // Wait for the transaction to come available. The tx.get() call will fail with an
@@ -1079,7 +1096,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, null);
     this.tracer = session.getTracer();
   }
 
@@ -1117,7 +1134,10 @@ 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);
+
+            ByteString previousAbortedTransactionId =
+                session.getIsMultiplexed() ? txn.transactionId : null;
+            txn = session.newTransaction(options, previousAbortedTransactionId);
           }
           checkState(
               isValid, "TransactionRunner has been invalidated by a new operation on the session");