fix: fix aborted handling of batchUpdateAsync (#421)

Author: olavloite

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

@@ -540,7 +540,7 @@ public ApiFuture<long[]> batchUpdateAsync(Iterable<Statement> statements) {
         decreaseAsyncOperations();
         throw t;
       }
-      final ApiFuture<long[]> updateCounts =
+      ApiFuture<long[]> updateCounts =
           ApiFutures.transform(
               response,
               new ApiFunction<ExecuteBatchDmlResponse, long[]>() {
@@ -565,19 +565,24 @@ public long[] apply(ExecuteBatchDmlResponse input) {
                 }
               },
               MoreExecutors.directExecutor());
+      updateCounts =
+          ApiFutures.catching(
+              updateCounts,
+              Throwable.class,
+              new ApiFunction<Throwable, long[]>() {
+                @Override
+                public long[] apply(Throwable input) {
+                  SpannerException e = SpannerExceptionFactory.newSpannerException(input);
+                  onError(e);
+                  throw e;
+                }
+              },
+              MoreExecutors.directExecutor());
       updateCounts.addListener(
           new Runnable() {
             @Override
             public void run() {
-              try {
-                updateCounts.get();
-              } catch (ExecutionException e) {
-                onError(SpannerExceptionFactory.newSpannerException(e.getCause()));
-              } catch (InterruptedException e) {
-                onError(SpannerExceptionFactory.propagateInterrupt(e));
-              } finally {
-                decreaseAsyncOperations();
-              }
+              decreaseAsyncOperations();
             }
           },
           MoreExecutors.directExecutor());

google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java

@@ -734,6 +734,47 @@ public ApiFuture<long[]> apply(TransactionContext txn, Void input)
             CommitRequest.class);
   }
 
+  @Test
+  public void asyncTransactionManagerBatchUpdateAbortedBeforeFirstStatement() throws Exception {
+    final AtomicInteger attempt = new AtomicInteger();
+    try (AsyncTransactionManager mgr = clientWithEmptySessionPool().transactionManagerAsync()) {
+      TransactionContextFuture txn = mgr.beginAsync();
+      while (true) {
+        try {
+          txn.then(
+                  new AsyncTransactionFunction<Void, long[]>() {
+                    @Override
+                    public ApiFuture<long[]> apply(TransactionContext txn, Void input)
+                        throws Exception {
+                      if (attempt.incrementAndGet() == 1) {
+                        mockSpanner.abortTransaction(txn);
+                      }
+                      return txn.batchUpdateAsync(
+                          ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT));
+                    }
+                  },
+                  executor)
+              .commitAsync()
+              .get();
+          break;
+        } catch (AbortedException e) {
+          txn = mgr.resetForRetryAsync();
+        }
+      }
+    }
+    assertThat(attempt.get()).isEqualTo(2);
+    // There should only be 1 CommitRequest, as the first attempt should abort already after the
+    // ExecuteBatchDmlRequest.
+    assertThat(mockSpanner.getRequestTypes())
+        .containsExactly(
+            BatchCreateSessionsRequest.class,
+            BeginTransactionRequest.class,
+            ExecuteBatchDmlRequest.class,
+            BeginTransactionRequest.class,
+            ExecuteBatchDmlRequest.class,
+            CommitRequest.class);
+  }
+
   @Test
   public void asyncTransactionManagerWithBatchUpdateCommitAborted() throws Exception {
     try (AsyncTransactionManager mgr = clientWithEmptySessionPool().transactionManagerAsync()) {