fix: session retry could cause infinite wait (#616)

A "Session not found" when using an AsyncTransactionManager could cause an
infinite wait for an ApiFuture that would never be done.

Fixes #605

Author: olavloite

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

@@ -98,6 +98,7 @@ private ApiFuture<TransactionContext> internalBeginAsync(boolean firstAttempt) {
         new ApiFutureCallback<Void>() {
           @Override
           public void onFailure(Throwable t) {
+            onError(t);
             res.setException(SpannerExceptionFactory.newSpannerException(t));
           }
 
@@ -130,6 +131,7 @@ public ApiFuture<Timestamp> commitAsync() {
     }
     ApiFuture<Timestamp> res = txn.commitAsync();
     txnState = TransactionState.COMMITTED;
+
     ApiFutures.addCallback(
         res,
         new ApiFutureCallback<Timestamp>() {
@@ -174,10 +176,6 @@ public ApiFuture<Void> apply(Empty input) throws Exception {
 
   @Override
   public TransactionContextFuture resetForRetryAsync() {
-    if (txn == null || (!txn.isAborted() && txnState != TransactionState.ABORTED)) {
-      throw new IllegalStateException(
-          "resetForRetry can only be called if the previous attempt aborted");
-    }
     return new TransactionContextFutureImpl(this, internalBeginAsync(false));
   }
 

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

@@ -38,6 +38,9 @@ class SessionPoolAsyncTransactionManager
   @GuardedBy("lock")
   private TransactionState txnState;
 
+  @GuardedBy("lock")
+  private AbortedException abortedException;
+
   private final SessionPool pool;
   private volatile PooledSessionFuture session;
   private volatile SettableApiFuture<AsyncTransactionManagerImpl> delegate;
@@ -159,6 +162,7 @@ public void onError(Throwable t) {
     if (t instanceof AbortedException) {
       synchronized (lock) {
         txnState = TransactionState.ABORTED;
+        abortedException = (AbortedException) t;
       }
     }
   }
@@ -167,9 +171,12 @@ public void onError(Throwable t) {
   public ApiFuture<Timestamp> commitAsync() {
     synchronized (lock) {
       Preconditions.checkState(
-          txnState == TransactionState.STARTED,
+          txnState == TransactionState.STARTED || txnState == TransactionState.ABORTED,
           "commit can only be invoked if the transaction is in progress. Current state: "
               + txnState);
+      if (txnState == TransactionState.ABORTED) {
+        return ApiFutures.immediateFailedFuture(abortedException);
+      }
       txnState = TransactionState.COMMITTED;
     }
     return ApiFutures.transformAsync(
@@ -186,6 +193,7 @@ public void onFailure(Throwable t) {
                     synchronized (lock) {
                       if (t instanceof AbortedException) {
                         txnState = TransactionState.ABORTED;
+                        abortedException = (AbortedException) t;
                       } else {
                         txnState = TransactionState.COMMIT_FAILED;
                       }

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

@@ -109,6 +109,7 @@ class AsyncTransactionStatementImpl<I, O> extends ForwardingApiFuture<O>
             @Override
             public void onFailure(Throwable t) {
               mgr.onError(t);
+              statementResult.setException(t);
               txnResult.setException(t);
             }