test(spanner): unflake ConnectionAsyncApiAbortedTest

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

@@ -1142,6 +1142,17 @@ private boolean hasTransactionalChanges() {
     return internalIsTransactionStarted() || this.connectionState.hasTransactionalChanges();
   }
 
+  /**
+   * Returns true if the current transaction has been aborted and could not be retried. Committing
+   * such a transaction must return the abort error instead of silently succeeding. This can happen
+   * when a statement that was executed asynchronously aborts before commit is called.
+   */
+  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 +1398,8 @@ private ApiFuture<Void> endCurrentTransactionAsync(
         statementTag == null, "Statement tags are not supported for COMMIT or ROLLBACK");
     ApiFuture<Void> res;
     try {
-      if (hasTransactionalChanges()) {
+      if (hasTransactionalChanges()
+          || (endTransactionMethod == commit && isCurrentTransactionAborted())) {
         res =
             endTransactionMethod.endAsync(
                 callType, getCurrentUnitOfWorkOrStartNewUnitOfWork(parsedStatement));

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

@@ -601,7 +601,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);
@@ -1327,7 +1329,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 +1360,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,8 @@
 
 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.assertThrows;
 import static org.junit.Assert.fail;
 
 import com.google.api.core.ApiFuture;
@@ -531,6 +533,29 @@ public void testBlindUpdateAborted_WithConcurrentModification() {
     }
   }
 
+  @Test
+  public void testCommitAfterBlindUpdateAborted_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 testMultipleBlindUpdatesAborted_WithConcurrentModification() {
     Statement update1 = Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=100");