test(spanner): unflake ConnectionAsyncApiAbortedTest

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java

@@ -464,7 +464,7 @@ public ApiFuture<Void> 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<Void> endCurrentTransactionAsync(
         statementTag == null, "Statement tags are not supported for COMMIT or ROLLBACK");
     ApiFuture<Void> res;
     try {
-      if (hasTransactionalChanges()) {
+      if (hasTransactionalChanges() || isCurrentTransactionAborted()) {
         res =
             endTransactionMethod.endAsync(
                 callType, getCurrentUnitOfWorkOrStartNewUnitOfWork(parsedStatement));

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<Void> 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<Void> result =
           executeStatementAsync(
               callType, ROLLBACK_STATEMENT, rollbackCallable, SpannerGrpc.getRollbackMethod());

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,14 +76,20 @@ private static final class QueryResult {
 
   private static final class RetryCounter implements TransactionRetryListener {
     final CountDownLatch latch;
+    final CountDownLatch finishedLatch;
     int retryCount = 0;
 
     RetryCounter() {
       this(0);
     }
 
     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<Long> 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<Long> 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<Long> 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");