feat: support begin with AbortedException for manager interface (#3835)

Aborted transactions should ideally be retried in the same transaction manager instance because the client library will ensure to populate the "Prev txn attempt's txn id" on retry (for mux sessions) which helps preserve the lock order(priority) of the transaction in scenarios of high contention thats causing txn wounding. This is achieved by using `resetForRetry()` that automatically takes care of preserving the lock order.

But if the customer application retries aborted transactions on a new TransactionManager instance with out using the `resetForRetry()`, we lose the lock order of the previous attempt.

To address this scenario of preserving the lockorder across transaction manger instance for a retry of the same transaction
Current proposal as per client lib team recommendation:
1. Overload the [Begin](https://togithub.com/googleapis/java-spanner/blob/9940b66dabe519495a69dcc0020e2587c16f7285/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManager.java#L62) API in the TransactionManager to accept an AbortedException argument.
2. Client library will populate the aborted transaction's transaction ID in the [AbortedException](https://togithub.com/googleapis/java-spanner/blob/efb168015aee1c4cba16f666b89be19e6f04fdbe/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbortedException.java#L27) object which is propagated up the customer's application stack.
3. On the customer end, they will just need to pass in this AbortedException object to the "Begin" API when they retry the transaction on the new transaction manager instance.

Note: This will be a no-op for regular sessions.
b/407037309

Author: harshachinta

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;