[Backport] Fix exception propagation in Async API methods

driver-core/src/main/com/mongodb/internal/async/AsyncFunction.java

@@ -18,6 +18,8 @@
 
 import com.mongodb.lang.Nullable;
 
+import java.util.concurrent.atomic.AtomicBoolean;
+
 /**
  * See {@link AsyncRunnable}
  * <p>
@@ -33,4 +35,28 @@ public interface AsyncFunction<T, R> {
      * @param callback the callback
      */
     void unsafeFinish(T value, SingleResultCallback<R> callback);
+
+    /**
+     * Must be invoked at end of async chain or when executing a callback handler supplied by the caller.
+     *
+     * @param callback the callback provided by the method the chain is used in.
+     */
+    default void finish(final T value, final SingleResultCallback<R> callback) {
+        final AtomicBoolean callbackInvoked = new AtomicBoolean(false);
+        try {
+            this.unsafeFinish(value, (v, e) -> {
+                if (!callbackInvoked.compareAndSet(false, true)) {
+                    throw new AssertionError(String.format("Callback has been already completed. It could happen "
+                            + "if code throws an exception after invoking an async method. Value: %s", v), e);
+                }
+                callback.onResult(v, e);
+            });
+        } catch (Throwable t) {
+            if (!callbackInvoked.compareAndSet(false, true)) {
+                throw t;
+            } else {
+                callback.completeExceptionally(t);
+            }
+        }
+    }
 }

driver-core/src/main/com/mongodb/internal/async/AsyncRunnable.java

@@ -170,7 +170,9 @@ default AsyncRunnable thenRun(final AsyncRunnable runnable) {
         return (c) -> {
             this.unsafeFinish((r, e) -> {
                 if (e == null) {
-                    runnable.unsafeFinish(c);
+                    /* If 'runnable' is executed on a different thread from the one that executed the initial 'finish()',
+                     then invoking 'finish()' within 'runnable' will catch and propagate any exceptions to 'c' (the callback). */
+                    runnable.finish(c);
                 } else {
                     c.completeExceptionally(e);
                 }
@@ -199,7 +201,7 @@ default AsyncRunnable thenRunIf(final Supplier<Boolean> condition, final AsyncRu
                     return;
                 }
                 if (matched) {
-                    runnable.unsafeFinish(callback);
+                    runnable.finish(callback);
                 } else {
                     callback.complete(callback);
                 }
@@ -216,7 +218,7 @@ default <R> AsyncSupplier<R> thenSupply(final AsyncSupplier<R> supplier) {
         return (c) -> {
             this.unsafeFinish((r, e) -> {
                 if (e == null) {
-                    supplier.unsafeFinish(c);
+                    supplier.finish(c);
                 } else {
                     c.completeExceptionally(e);
                 }

driver-core/src/main/com/mongodb/internal/async/AsyncSupplier.java

@@ -18,6 +18,7 @@
 
 import com.mongodb.lang.Nullable;
 
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Predicate;
 
 
@@ -54,18 +55,25 @@ default void unsafeFinish(@Nullable final Void value, final SingleResultCallback
     }
 
     /**
-     * Must be invoked at end of async chain.
+     * Must be invoked at end of async chain or when executing a callback handler supplied by the caller.
+     *
+     * @see #thenApply(AsyncFunction)
+     * @see #thenConsume(AsyncConsumer)
+     * @see #onErrorIf(Predicate, AsyncFunction)
      * @param callback the callback provided by the method the chain is used in
      */
     default void finish(final SingleResultCallback<T> callback) {
-        final boolean[] callbackInvoked = {false};
+        final AtomicBoolean callbackInvoked = new AtomicBoolean(false);
         try {
             this.unsafeFinish((v, e) -> {
-                callbackInvoked[0] = true;
+                if (!callbackInvoked.compareAndSet(false, true)) {
+                    throw new AssertionError(String.format("Callback has been already completed. It could happen "
+                            + "if code throws an exception after invoking an async method. Value: %s", v), e);
+                }
                 callback.onResult(v, e);
             });
         } catch (Throwable t) {
-            if (callbackInvoked[0]) {
+            if (!callbackInvoked.compareAndSet(false, true)) {
                 throw t;
             } else {
                 callback.completeExceptionally(t);
@@ -80,9 +88,9 @@ default void finish(final SingleResultCallback<T> callback) {
      */
     default <R> AsyncSupplier<R> thenApply(final AsyncFunction<T, R> function) {
         return (c) -> {
-            this.unsafeFinish((v, e) -> {
+            this.finish((v, e) -> {
                 if (e == null) {
-                    function.unsafeFinish(v, c);
+                    function.finish(v, c);
                 } else {
                     c.completeExceptionally(e);
                 }
@@ -99,7 +107,7 @@ default AsyncRunnable thenConsume(final AsyncConsumer<T> consumer) {
         return (c) -> {
             this.unsafeFinish((v, e) -> {
                 if (e == null) {
-                    consumer.unsafeFinish(v, c);
+                    consumer.finish(v, c);
                 } else {
                     c.completeExceptionally(e);
                 }
@@ -131,7 +139,7 @@ default AsyncSupplier<T> onErrorIf(
                 return;
             }
             if (errorMatched) {
-                errorFunction.unsafeFinish(e, callback);
+                errorFunction.finish(e, callback);
             } else {
                 callback.completeExceptionally(e);
             }

driver-core/src/main/com/mongodb/internal/connection/InternalStreamConnection.java

@@ -595,6 +595,7 @@ private <T> void sendCommandMessageAsync(final int messageId, final Decoder<T> d
                         return;
                     }
                     assertNotNull(responseBuffers);
+                    T commandResult;
                     try {
                         updateSessionContext(sessionContext, responseBuffers);
                         boolean commandOk =
@@ -609,13 +610,14 @@ private <T> void sendCommandMessageAsync(final int messageId, final Decoder<T> d
                         }
                         commandEventSender.sendSucceededEvent(responseBuffers);
 
-                        T result1 = getCommandResult(decoder, responseBuffers, messageId);
-                        callback.onResult(result1, null);
+                        commandResult = getCommandResult(decoder, responseBuffers, messageId);
                     } catch (Throwable localThrowable) {
                         callback.onResult(null, localThrowable);
+                        return;
                     } finally {
                         responseBuffers.close();
                     }
+                    callback.onResult(commandResult, null);
                 }));
             }
         });

driver-core/src/main/com/mongodb/internal/connection/InternalStreamConnectionInitializer.java

@@ -98,7 +98,14 @@ public void startHandshakeAsync(final InternalConnection internalConnection,
                         callback.onResult(null, t instanceof MongoException ? mapHelloException((MongoException) t) : t);
                     } else {
                         setSpeculativeAuthenticateResponse(helloResult);
-                        callback.onResult(createInitializationDescription(helloResult, internalConnection, startTime), null);
+                        InternalConnectionInitializationDescription initializationDescription;
+                        try {
+                            initializationDescription = createInitializationDescription(helloResult, internalConnection, startTime);
+                        } catch (Throwable localThrowable) {
+                            callback.onResult(null, localThrowable);
+                            return;
+                        }
+                        callback.onResult(initializationDescription, null);
                     }
                 });
     }