diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java index cadd63757398..a1ed225c71cc 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java @@ -464,7 +464,7 @@ public ApiFuture closeAsync() { if (isBatchActive()) { abortBatch(); } - if (isTransactionStarted()) { + if (isTransactionStarted() || isCurrentTransactionAborted()) { try { futures.add(rollbackAsync()); } catch (Exception exception) { @@ -1142,6 +1142,18 @@ private boolean hasTransactionalChanges() { return internalIsTransactionStarted() || this.connectionState.hasTransactionalChanges(); } + /** + * Returns true if the current transaction has been aborted and could not be retried. Ending such + * a transaction must be delegated to the underlying unit of work so that committing returns the + * abort error instead of silently succeeding, and rolling back transitions the state to + * ROLLED_BACK, cancels any scheduled keep-alive ping, and ends the OpenTelemetry span. + */ + private boolean isCurrentTransactionAborted() { + return internalIsInTransaction() + && this.currentUnitOfWork != null + && this.currentUnitOfWork.getState() == UnitOfWorkState.ABORTED; + } + @Override public Timestamp getReadTimestamp() { ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG); @@ -1387,7 +1399,7 @@ private ApiFuture endCurrentTransactionAsync( statementTag == null, "Statement tags are not supported for COMMIT or ROLLBACK"); ApiFuture res; try { - if (hasTransactionalChanges()) { + if (hasTransactionalChanges() || isCurrentTransactionAborted()) { res = endTransactionMethod.endAsync( callType, getCurrentUnitOfWorkOrStartNewUnitOfWork(parsedStatement)); diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java index 246c203c8674..594e8c6f6dde 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java @@ -473,7 +473,8 @@ private void checkValidStateAndMarkStarted() { } private boolean shouldPing() { - return isActive() + return getState().isActive() + && this.abortedException == null && keepAliveLock != null && keepTransactionAlive && !timedOutOrCancelled @@ -601,7 +602,9 @@ public boolean isActive() { } void checkAborted() { - if (this.state == UnitOfWorkState.ABORTED && this.abortedException != null) { + // Only check abortedException and not the state, as the ABORTED state can be overwritten by a + // commit that is called while (or after) an asynchronous statement aborted the transaction. + if (this.abortedException != null) { if (this.abortedException instanceof AbortedDueToConcurrentModificationException) { throw SpannerExceptionFactory.newAbortedDueToConcurrentModificationException( (AbortedDueToConcurrentModificationException) this.abortedException); @@ -1243,8 +1246,7 @@ private void handleAborted(AbortedException aborted) { } catch (Throwable t) { // ignore } - this.state = UnitOfWorkState.ABORTED; - this.abortedException = e; + markAborted(e); throw e; } catch (AbortedException abortedExceptionDuringRetry) { // Retry aborted, do another retry of the transaction. @@ -1269,8 +1271,7 @@ private void handleAborted(AbortedException aborted) { // ignore } // Set transaction state to aborted as the retry failed. - this.state = UnitOfWorkState.ABORTED; - this.abortedException = aborted; + markAborted(aborted); // Re-throw underlying exception. throw e; } @@ -1282,12 +1283,18 @@ private void handleAborted(AbortedException aborted) { // ignore } // Internal retry is not enabled. - this.state = UnitOfWorkState.ABORTED; - this.abortedException = aborted; + markAborted(aborted); throw aborted; } } + private void markAborted(AbortedException aborted) { + if (this.state == UnitOfWorkState.STARTED) { + this.state = UnitOfWorkState.ABORTED; + } + this.abortedException = aborted; + } + private void throwAbortWithRetryAttemptsExceeded() throws SpannerException { invokeTransactionRetryListenersOnFinish(RetryResult.RETRY_ABORTED_AND_MAX_ATTEMPTS_EXCEEDED); logger.fine( @@ -1302,11 +1309,10 @@ private void throwAbortWithRetryAttemptsExceeded() throws SpannerException { } catch (Throwable t) { // ignore } - this.state = UnitOfWorkState.ABORTED; - this.abortedException = + markAborted( (AbortedException) SpannerExceptionFactory.newSpannerException( - ErrorCode.ABORTED, MAX_INTERNAL_RETRIES_EXCEEDED); + ErrorCode.ABORTED, MAX_INTERNAL_RETRIES_EXCEEDED)); throw this.abortedException; } @@ -1327,7 +1333,9 @@ private void invokeTransactionRetryListenersOnFinish(RetryResult result) { @Override public Void call() { try { - if (state != UnitOfWorkState.ABORTED && rolledBackToSavepointException == null) { + if (state != UnitOfWorkState.ABORTED + && abortedException == null + && rolledBackToSavepointException == null) { // Make sure the transaction has actually started before we try to rollback. get(txContextFuture); txManager.rollback(); @@ -1356,10 +1364,11 @@ private ApiFuture rollbackAsync(CallType callType, boolean updateStatusAnd state == UnitOfWorkState.STARTED || state == UnitOfWorkState.ABORTED, "This transaction has status " + state.name()); cancelScheduledKeepAlivePing(); + boolean isAborted = (state == UnitOfWorkState.ABORTED) || (this.abortedException != null); if (updateStatusAndEndSpan) { state = UnitOfWorkState.ROLLED_BACK; } - if (txContextFuture != null && state != UnitOfWorkState.ABORTED) { + if (txContextFuture != null && !isAborted) { ApiFuture result = executeStatementAsync( callType, ROLLBACK_STATEMENT, rollbackCallable, SpannerGrpc.getRollbackMethod()); diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionAsyncApiAbortedTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionAsyncApiAbortedTest.java index 1f4378ecd86b..93512fe4b9dc 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionAsyncApiAbortedTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionAsyncApiAbortedTest.java @@ -18,6 +18,11 @@ import static com.google.cloud.spanner.SpannerApiFutures.get; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; import static org.junit.Assert.fail; import com.google.api.core.ApiFuture; @@ -35,6 +40,7 @@ import com.google.cloud.spanner.connection.ConnectionOptions.Builder; import com.google.cloud.spanner.connection.ITAbstractSpannerTest.ITConnection; import com.google.cloud.spanner.connection.StatementExecutor.StatementExecutorType; +import com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState; import com.google.common.collect.Collections2; import com.google.common.collect.ImmutableList; import com.google.common.collect.Lists; @@ -70,6 +76,7 @@ private static final class QueryResult { private static final class RetryCounter implements TransactionRetryListener { final CountDownLatch latch; + final CountDownLatch finishedLatch; int retryCount = 0; RetryCounter() { @@ -77,7 +84,12 @@ private static final class RetryCounter implements TransactionRetryListener { } RetryCounter(int countDown) { - latch = new CountDownLatch(countDown); + this(countDown, 0); + } + + RetryCounter(int startCountDown, int finishCountDown) { + latch = new CountDownLatch(startCountDown); + finishedLatch = new CountDownLatch(finishCountDown); } @Override @@ -88,7 +100,9 @@ public void retryStarting(Timestamp transactionStarted, long transactionId, int @Override public void retryFinished( - Timestamp transactionStarted, long transactionId, int retryAttempt, RetryResult result) {} + Timestamp transactionStarted, long transactionId, int retryAttempt, RetryResult result) { + finishedLatch.countDown(); + } } private static final ExecutorService singleThreadedExecutor = Executors.newSingleThreadExecutor(); @@ -531,6 +545,112 @@ public void testBlindUpdateAborted_WithConcurrentModification() { } } + @Test + public void testCommitAfterUpdateAborted_WithConcurrentModification() { + Statement update1 = Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=100"); + mockSpanner.putStatementResult(StatementResult.update(update1, 100)); + + RetryCounter counter = new RetryCounter(); + try (Connection connection = createConnection(counter)) { + get(connection.executeUpdateAsync(update1)); + mockSpanner.putStatementResult(StatementResult.update(update1, 200)); + + mockSpanner.abortNextStatement(); + ApiFuture insert = connection.executeUpdateAsync(INSERT_STATEMENT); + // Wait until the failed retry has marked the transaction as aborted before calling commit. + assertThrows(AbortedDueToConcurrentModificationException.class, () -> get(insert)); + assertEquals(1, counter.retryCount); + + // Committing a transaction that has been aborted must fail and must not silently succeed. + assertThrows( + AbortedDueToConcurrentModificationException.class, () -> get(connection.commitAsync())); + assertEquals(1, counter.retryCount); + } + } + + @Test + public void testBlindUpdateAborted_ConcurrentCommit_WithConcurrentModification() + throws InterruptedException { + Statement update1 = Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=100"); + for (int i = 0; i < 10; i++) { + mockSpanner.putStatementResult(StatementResult.update(update1, 100)); + + RetryCounter counter = new RetryCounter(0, 1); + try (Connection connection = createConnection(counter)) { + get(connection.executeUpdateAsync(update1)); + mockSpanner.putStatementResult(StatementResult.update(update1, 200)); + + mockSpanner.abortNextStatement(); + // Execute a blind (shoot-and-forget) update without waiting for its future. + connection.executeUpdateAsync(INSERT_STATEMENT); + // Trigger commitAsync concurrently as the internal retry finishes (while the blind update + // is still rolling back and marking the transaction as aborted on the executor thread). + assertThat(counter.finishedLatch.await(10L, TimeUnit.SECONDS)).isTrue(); + assertThrows( + AbortedDueToConcurrentModificationException.class, () -> get(connection.commitAsync())); + assertEquals(1, counter.retryCount); + } + } + } + + @Test + public void testRollbackAfterTransactionAborted_UpdatesStateToRolledBack() { + Statement update1 = Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=100"); + mockSpanner.putStatementResult(StatementResult.update(update1, 100)); + + RetryCounter counter = new RetryCounter(); + try (Connection connection = createConnection(counter)) { + get(connection.executeUpdateAsync(update1)); + mockSpanner.putStatementResult(StatementResult.update(update1, 200)); + + mockSpanner.abortNextStatement(); + ApiFuture insert = connection.executeUpdateAsync(INSERT_STATEMENT); + assertThrows(AbortedDueToConcurrentModificationException.class, () -> get(insert)); + assertEquals(1, counter.retryCount); + + // Verify the transaction is currently in the ABORTED state. + ConnectionImpl connectionImpl = (ConnectionImpl) connection; + ReadWriteTransaction transaction = + (ReadWriteTransaction) connectionImpl.getCurrentUnitOfWorkOrStartNewUnitOfWork(); + assertEquals(UnitOfWorkState.ABORTED, transaction.getState()); + + // Rollback the aborted transaction. + get(connection.rollbackAsync()); + + assertEquals(UnitOfWorkState.ROLLED_BACK, transaction.getState()); + } + } + + @Test + public void testRollbackAfterTransactionAborted_CancelsKeepAlivePing() { + Statement update1 = Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=100"); + mockSpanner.putStatementResult(StatementResult.update(update1, 100)); + + RetryCounter counter = new RetryCounter(); + try (Connection connection = createConnection(counter)) { + connection.setKeepTransactionAlive(true); + get(connection.executeUpdateAsync(update1)); + mockSpanner.putStatementResult(StatementResult.update(update1, 200)); + + ConnectionImpl connectionImpl = (ConnectionImpl) connection; + ReadWriteTransaction transaction = + (ReadWriteTransaction) connectionImpl.getCurrentUnitOfWorkOrStartNewUnitOfWork(); + assertNotNull(transaction.getKeepAliveFuture()); + assertFalse(transaction.getKeepAliveFuture().isCancelled()); + + mockSpanner.abortNextStatement(); + ApiFuture insert = connection.executeUpdateAsync(INSERT_STATEMENT); + assertThrows(AbortedDueToConcurrentModificationException.class, () -> get(insert)); + assertEquals(1, counter.retryCount); + assertEquals(UnitOfWorkState.ABORTED, transaction.getState()); + + // Rollback the aborted transaction. + get(connection.rollbackAsync()); + + assertNull(transaction.getKeepAliveFuture()); + } + } + @Test public void testMultipleBlindUpdatesAborted_WithConcurrentModification() { Statement update1 = Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=100");