diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java index 7dda260102..d4bf7f4a47 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java @@ -98,6 +98,7 @@ private ApiFuture internalBeginAsync(boolean firstAttempt) { new ApiFutureCallback() { @Override public void onFailure(Throwable t) { + onError(t); res.setException(SpannerExceptionFactory.newSpannerException(t)); } @@ -130,6 +131,7 @@ public ApiFuture commitAsync() { } ApiFuture res = txn.commitAsync(); txnState = TransactionState.COMMITTED; + ApiFutures.addCallback( res, new ApiFutureCallback() { @@ -174,10 +176,6 @@ public ApiFuture 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)); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java index 515286fb11..9dc947332f 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java +++ b/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 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 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; } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java index be21a947d1..c9d9c04d64 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java @@ -109,6 +109,7 @@ class AsyncTransactionStatementImpl extends ForwardingApiFuture @Override public void onFailure(Throwable t) { mgr.onError(t); + statementResult.setException(t); txnResult.setException(t); }