chore(spanner): support AsyncTransactionManager

Author: harshachinta

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

@@ -170,6 +170,8 @@ interface AsyncTransactionFunction<I, O> {
    */
   TransactionContextFuture beginAsync();
 
+  TransactionContextFuture beginAsync(AbortedException abortedException);
+
   /**
    * Rolls back the currently active transaction. In most cases there should be no need to call this
    * explicitly since {@link #close()} would automatically roll back any active transaction.

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

@@ -76,14 +76,24 @@ public ApiFuture<Void> closeAsync() {
   @Override
   public TransactionContextFutureImpl beginAsync() {
     Preconditions.checkState(txn == null, "begin can only be called once");
-    return new TransactionContextFutureImpl(this, internalBeginAsync(true));
+    return new TransactionContextFutureImpl(this, internalBeginAsync(true, ByteString.EMPTY));
   }
 
-  private ApiFuture<TransactionContext> internalBeginAsync(boolean firstAttempt) {
+  @Override
+  public TransactionContextFutureImpl beginAsync(AbortedException abortedException) {
+    Preconditions.checkState(txn == null, "begin can only be called once");
+    ByteString abortedTransactionId = abortedException.getTransactionID() != null ? abortedException.getTransactionID() : ByteString.EMPTY;
+    return new TransactionContextFutureImpl(this, internalBeginAsync(true, abortedTransactionId));
+  }
+
+  private ApiFuture<TransactionContext> internalBeginAsync(boolean firstAttempt, ByteString abortedTransactionID) {
     txnState = TransactionState.STARTED;
 
     // Determine the latest transactionId when using a multiplexed session.
     ByteString multiplexedSessionPreviousTransactionId = ByteString.EMPTY;
+    if (firstAttempt && session.getIsMultiplexed()) {
+      multiplexedSessionPreviousTransactionId = abortedTransactionID;
+    }
     if (txn != null && session.getIsMultiplexed() && !firstAttempt) {
       // Use the current transactionId if available, otherwise fallback to the previous aborted
       // transactionId.
@@ -187,7 +197,7 @@ public TransactionContextFuture resetForRetryAsync() {
       throw new IllegalStateException(
           "resetForRetry can only be called if the previous attempt aborted");
     }
-    return new TransactionContextFutureImpl(this, internalBeginAsync(false));
+    return new TransactionContextFutureImpl(this, internalBeginAsync(false, ByteString.EMPTY));
   }
 
   @Override

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

@@ -50,6 +50,11 @@ public TransactionContextFuture beginAsync() {
     return getAsyncTransactionManager().beginAsync();
   }
 
+  @Override
+  public TransactionContextFuture beginAsync(AbortedException abortedException) {
+    return getAsyncTransactionManager().beginAsync(abortedException);
+  }
+
   @Override
   public ApiFuture<Void> rollbackAsync() {
     return getAsyncTransactionManager().rollbackAsync();

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

@@ -163,6 +163,11 @@ public void onSuccess(TransactionContext result) {
     return new TransactionContextFutureImpl(this, delegateTxnFuture);
   }
 
+  @Override
+  public TransactionContextFuture beginAsync(AbortedException abortedException) {
+    return beginAsync();
+  }
+
   @Override
   public void onError(Throwable t) {
     if (t instanceof AbortedException) {

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java

@@ -46,6 +46,7 @@
 import com.google.api.gax.longrunning.OperationFuture;
 import com.google.cloud.NoCredentials;
 import com.google.cloud.Timestamp;
+import com.google.cloud.spanner.AbortedException;
 import com.google.cloud.spanner.BatchClient;
 import com.google.cloud.spanner.BatchReadOnlyTransaction;
 import com.google.cloud.spanner.BatchTransactionId;
@@ -120,6 +121,11 @@ public TransactionContext begin() {
       return txContext;
     }
 
+    @Override
+    public TransactionContext begin(AbortedException exception) {
+      return begin();
+    }
+
     @Override
     public void commit() {
       Timestamp commitTimestamp = Timestamp.now();