feat: support begin with AbortedException for manager interface

google-cloud-spanner/clirr-ignored-differences.xml

@@ -996,4 +996,16 @@
     <className>com/google/cloud/spanner/DatabaseClient</className>
     <method>com.google.cloud.spanner.Statement$StatementFactory getStatementFactory()</method>
   </difference>
+
+  <!-- Add method begin() with AbortedException in TransactionManager and AsyncTransactionManager -->
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/AsyncTransactionManager</className>
+    <method>com.google.cloud.spanner.AsyncTransactionManager$TransactionContextFuture beginAsync(com.google.cloud.spanner.AbortedException)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/TransactionManager</className>
+    <method>com.google.cloud.spanner.TransactionContext begin(com.google.cloud.spanner.AbortedException)</method>
+  </difference>
 </differences>

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

@@ -17,6 +17,7 @@
 package com.google.cloud.spanner;
 
 import com.google.api.gax.rpc.ApiException;
+import com.google.protobuf.ByteString;
 import javax.annotation.Nullable;
 
 /**
@@ -32,6 +33,8 @@ public class AbortedException extends SpannerException {
    */
   private static final boolean IS_RETRYABLE = false;
 
+  private ByteString transactionID;
+
   /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */
   AbortedException(
       DoNotConstructDirectly token, @Nullable String message, @Nullable Throwable cause) {
@@ -46,6 +49,9 @@ public class AbortedException extends SpannerException {
       @Nullable ApiException apiException,
       @Nullable XGoogSpannerRequestId reqId) {
     super(token, ErrorCode.ABORTED, IS_RETRYABLE, message, cause, apiException, reqId);
+    if (cause instanceof AbortedException) {
+      this.transactionID = ((AbortedException) cause).getTransactionID();
+    }
   }
 
   /**
@@ -55,4 +61,12 @@ public class AbortedException extends SpannerException {
   public boolean isEmulatorOnlySupportsOneTransactionException() {
     return getMessage().endsWith("The emulator only supports one transaction at a time.");
   }
+
+  void setTransactionID(ByteString transactionID) {
+    this.transactionID = transactionID;
+  }
+
+  ByteString getTransactionID() {
+    return this.transactionID;
+  }
 }

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

@@ -170,6 +170,21 @@ interface AsyncTransactionFunction<I, O> {
    */
   TransactionContextFuture beginAsync();
 
+  /**
+   * Initializes a new read-write transaction that is a retry of a previously aborted transaction.
+   * This method must be called before performing any operations, and it can only be invoked once
+   * per transaction lifecycle.
+   *
+   * <p>This method should only be used when multiplexed sessions are enabled to create a retry for
+   * a previously aborted transaction. This method can be used instead of {@link
+   * #resetForRetryAsync()} to create a retry. Using this method or {@link #resetForRetryAsync()}
+   * will have the same effect. You must pass in the {@link AbortedException} from the previous
+   * attempt to preserve the transaction's priority.
+   *
+   * <p>For regular sessions, this behaves the same as {@link #beginAsync()}.
+   */
+  TransactionContextFuture beginAsync(AbortedException exception);
+
   /**
    * 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,27 @@ 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 exception) {
+    Preconditions.checkState(txn == null, "begin can only be called once");
+    Preconditions.checkNotNull(exception, "AbortedException from the previous attempt is required");
+    ByteString abortedTransactionId =
+        exception.getTransactionID() != null ? exception.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 +200,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 exception) {
+    return getAsyncTransactionManager().beginAsync(exception);
+  }
+
   @Override
   public ApiFuture<Void> rollbackAsync() {
     return getAsyncTransactionManager().rollbackAsync();

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

@@ -49,6 +49,11 @@ public TransactionContext begin() {
     return getTransactionManager().begin();
   }
 
+  @Override
+  public TransactionContext begin(AbortedException exception) {
+    return getTransactionManager().begin(exception);
+  }
+
   @Override
   public void commit() {
     getTransactionManager().commit();

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

@@ -900,6 +900,13 @@ public TransactionContext begin() {
       return internalBegin();
     }
 
+    @Override
+    public TransactionContext begin(AbortedException exception) {
+      // For regular sessions, the input exception is ignored and the behavior is equivalent to
+      // calling {@link #begin()}.
+      return begin();
+    }
+
     private TransactionContext internalBegin() {
       TransactionContext res = new SessionPoolTransactionContext(this, delegate.begin());
       session.get().markUsed();

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

@@ -163,6 +163,13 @@ public void onSuccess(TransactionContext result) {
     return new TransactionContextFutureImpl(this, delegateTxnFuture);
   }
 
+  @Override
+  public TransactionContextFuture beginAsync(AbortedException exception) {
+    // For regular sessions, the input exception is ignored and the behavior is equivalent to
+    // calling {@link #beginAsync()}.
+    return beginAsync();
+  }
+
   @Override
   public void onError(Throwable t) {
     if (t instanceof AbortedException) {

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

@@ -61,6 +61,21 @@ enum TransactionState {
    */
   TransactionContext begin();
 
+  /**
+   * Initializes a new read-write transaction that is a retry of a previously aborted transaction.
+   * This method must be called before performing any operations, and it can only be invoked once
+   * per transaction lifecycle.
+   *
+   * <p>This method should only be used when multiplexed sessions are enabled to create a retry for
+   * a previously aborted transaction. This method can be used instead of {@link #resetForRetry()}
+   * to create a retry. Using this method or {@link #resetForRetry()} will have the same effect. You
+   * must pass in the {@link AbortedException} from the previous attempt to preserve the
+   * transaction's priority.
+   *
+   * <p>For regular sessions, this behaves the same as {@link #begin()}.
+   */
+  TransactionContext begin(AbortedException exception);
+
   /**
    * Commits the currently active transaction. If the transaction was already aborted, then this
    * would throw an {@link AbortedException}.

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

@@ -53,8 +53,21 @@ public void setSpan(ISpan span) {
   @Override
   public TransactionContext begin() {
     Preconditions.checkState(txn == null, "begin can only be called once");
+    return begin(ByteString.EMPTY);
+  }
+
+  @Override
+  public TransactionContext begin(AbortedException exception) {
+    Preconditions.checkState(txn == null, "begin can only be called once");
+    Preconditions.checkNotNull(exception, "AbortedException from the previous attempt is required");
+    ByteString previousAbortedTransactionID =
+        exception.getTransactionID() != null ? exception.getTransactionID() : ByteString.EMPTY;
+    return begin(previousAbortedTransactionID);
+  }
+
+  TransactionContext begin(ByteString previousTransactionId) {
     try (IScope s = tracer.withSpan(span)) {
-      txn = session.newTransaction(options, /* previousTransactionId = */ ByteString.EMPTY);
+      txn = session.newTransaction(options, previousTransactionId);
       session.setActive(this);
       txnState = TransactionState.STARTED;
       return txn;

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

@@ -793,6 +793,11 @@ public SpannerException onError(SpannerException e, boolean withBeginTransaction
         long delay = -1L;
         if (exceptionToThrow instanceof AbortedException) {
           delay = exceptionToThrow.getRetryDelayInMillis();
+          ((AbortedException) exceptionToThrow)
+              .setTransactionID(
+                  this.transactionId != null
+                      ? this.transactionId
+                      : this.getPreviousTransactionId());
         }
         if (delay == -1L) {
           txnLogger.log(