fix: handle thrown exceptions in runAsyncTransaction callback (#671)

Author: Brian Chen

google-cloud-firestore/src/main/java/com/google/cloud/firestore/TransactionRunner.java

@@ -145,28 +145,35 @@ public void run() {
    * Invokes the user callback on the user callback executor and returns the user-provided result.
    */
   private SettableApiFuture<T> invokeUserCallback() {
-    final SettableApiFuture<T> callbackResult = SettableApiFuture.create();
+    final SettableApiFuture<T> returnedResult = SettableApiFuture.create();
+
     userCallbackExecutor.execute(
         new Runnable() {
           @Override
           public void run() {
+            ApiFuture<T> userCallbackResult;
+            try {
+              userCallbackResult = userCallback.updateCallback(transaction);
+            } catch (Exception e) {
+              userCallbackResult = ApiFutures.immediateFailedFuture(e);
+            }
             ApiFutures.addCallback(
-                userCallback.updateCallback(transaction),
+                userCallbackResult,
                 new ApiFutureCallback<T>() {
                   @Override
                   public void onFailure(Throwable t) {
-                    callbackResult.setException(t);
+                    returnedResult.setException(t);
                   }
 
                   @Override
                   public void onSuccess(T result) {
-                    callbackResult.set(result);
+                    returnedResult.set(result);
                   }
                 },
-                MoreExecutors.directExecutor());
+                firestoreExecutor);
           }
         });
-    return callbackResult;
+    return returnedResult;
   }
 
   /** A callback that invokes the BeginTransaction callback. */

google-cloud-firestore/src/test/java/com/google/cloud/firestore/TransactionTest.java

@@ -253,7 +253,7 @@ public String updateCallback(Transaction transaction) throws Exception {
   }
 
   @Test
-  public void rollbackOnCallbackErrorAsync() {
+  public void rollbackOnCallbackApiFutureErrorAsync() {
     doReturn(beginResponse())
         .doReturn(rollbackResponse())
         .when(firestoreMock)
@@ -339,6 +339,31 @@ public ApiFuture<String> updateCallback(Transaction transaction) {
     assertEquals(1, requests.size());
   }
 
+  @Test
+  public void noRollbackOnThrownExceptionAsync() {
+    doReturn(beginResponse())
+        .doReturn(rollbackResponse())
+        .when(firestoreMock)
+        .sendRequest(requestCapture.capture(), Matchers.<UnaryCallable<Message, Message>>any());
+
+    ApiFuture<String> transaction =
+        firestoreMock.runAsyncTransaction(
+            new Transaction.AsyncFunction<String>() {
+              @Override
+              public ApiFuture<String> updateCallback(Transaction transaction) {
+                throw new RuntimeException("User exception");
+              }
+            },
+            options);
+
+    try {
+      transaction.get();
+      fail();
+    } catch (Exception e) {
+      assertTrue(e.getMessage().endsWith("User exception"));
+    }
+  }
+
   @Test
   public void limitsRetriesWithFailure() {
     ResponseStubber responseStubber =

google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITSystemTest.java

@@ -688,6 +688,24 @@ public String updateCallback(Transaction transaction) {
     }
   }
 
+  @Test
+  public void asyncTxFailsWithUserError() throws Exception {
+    try {
+      firestore
+          .runAsyncTransaction(
+              new Transaction.AsyncFunction<String>() {
+                @Override
+                public ApiFuture<String> updateCallback(Transaction transaction) {
+                  throw new RuntimeException("User exception");
+                }
+              })
+          .get();
+      fail();
+    } catch (Exception e) {
+      assertTrue(e.getMessage().endsWith("User exception"));
+    }
+  }
+
   @Test
   public void doesNotRetryTransactionsWithFailedPreconditions() {
     final DocumentReference documentReference = randomColl.document();