chore(spanner): review comments

Author: harshachinta

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

@@ -171,13 +171,15 @@ interface AsyncTransactionFunction<I, O> {
   TransactionContextFuture beginAsync();
 
   /**
-   * Initializes a new read-write transaction. This method must be called before performing any
-   * operations, and it can only be invoked once per transaction lifecycle.
+   * 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 is especially useful in scenarios involving multiplexed sessions and when creating a
-   * new transaction for retry attempts. If {@link #resetForRetryAsync()} is not used, you can pass
-   * the {@link AbortedException} from a previous attempt here to preserve the transaction's
-   * priority.
+   * <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()}.
    */

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

@@ -82,6 +82,7 @@ public TransactionContextFutureImpl beginAsync() {
   @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));

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

@@ -62,12 +62,15 @@ enum TransactionState {
   TransactionContext begin();
 
   /**
-   * Initializes a new read-write transaction. This method must be called before performing any
-   * operations, and it can only be invoked once per transaction lifecycle.
+   * 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 is especially useful in scenarios involving multiplexed sessions and when creating a
-   * new transaction for retry attempts. If {@link #resetForRetry()} is not used, you can pass the
-   * {@link AbortedException} from a previous attempt here to preserve the transaction's priority.
+   * <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()}.
    */

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

@@ -59,14 +59,15 @@ public TransactionContext begin() {
   @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 = */ previousTransactionId);
+      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,7 +793,11 @@ public SpannerException onError(SpannerException e, boolean withBeginTransaction
         long delay = -1L;
         if (exceptionToThrow instanceof AbortedException) {
           delay = exceptionToThrow.getRetryDelayInMillis();
-          ((AbortedException) exceptionToThrow).setTransactionID(this.transactionId);
+          ((AbortedException) exceptionToThrow)
+              .setTransactionID(
+                  this.transactionId != null
+                      ? this.transactionId
+                      : this.getPreviousTransactionId());
         }
         if (delay == -1L) {
           txnLogger.log(