Add throwTranslatedWriteException, refactoring, async helper

Author: katcharov

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

@@ -178,6 +178,42 @@ default AsyncRunnable thenRun(final AsyncRunnable runnable) {
         };
     }
 
+    /**
+     * The error check checks if the exception is an instance of the provided class.
+     * @see #thenRunTryCatchAsyncBlocks(AsyncRunnable, java.util.function.Predicate, AsyncFunction)
+     */
+    default <T extends Throwable> AsyncRunnable thenRunTryCatchAsyncBlocks(
+            final AsyncRunnable runnable,
+            final Class<T> exceptionClass,
+            final AsyncFunction<Throwable, Void> errorFunction) {
+        return thenRunTryCatchAsyncBlocks(runnable, e -> exceptionClass.isInstance(e), errorFunction);
+    }
+
+    /**
+     * Convenience method corresponding to a try-catch block in sync code.
+     * This MUST be used to properly handle cases where there is code above
+     * the block, whose errors must not be caught by an ensuing
+     * {@link #onErrorIf(java.util.function.Predicate, AsyncFunction)}.
+     *
+     * @param runnable corresponds to the contents of the try block
+     * @param errorCheck for matching on an error (or, a more complex condition)
+     * @param errorFunction corresponds to the contents of the catch block
+     * @return the composition of this runnable, a runnable that runs the
+     * provided runnable, followed by (composed with) the error function, which
+     * is conditional on there being an exception meeting the error check.
+     */
+    default AsyncRunnable thenRunTryCatchAsyncBlocks(
+            final AsyncRunnable runnable,
+            final Predicate<Throwable> errorCheck,
+            final AsyncFunction<Throwable, Void> errorFunction) {
+        return this.thenRun(c -> {
+            beginAsync()
+                    .thenRun(runnable)
+                    .onErrorIf(errorCheck, errorFunction)
+                    .finish(c);
+        });
+    }
+
     /**
      * @param condition the condition to check
      * @param runnable The async runnable to run after this runnable,

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

@@ -73,7 +73,6 @@
 import static com.mongodb.assertions.Assertions.isTrue;
 import static com.mongodb.assertions.Assertions.notNull;
 import static com.mongodb.internal.async.AsyncRunnable.beginAsync;
-import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback;
 import static com.mongodb.internal.connection.Authenticator.shouldAuthenticate;
 import static com.mongodb.internal.connection.CommandHelper.HELLO;
 import static com.mongodb.internal.connection.CommandHelper.LEGACY_HELLO;
@@ -633,19 +632,34 @@ private <T> T getCommandResult(final Decoder<T> decoder, final ResponseBuffers r
     @Override
     public void sendMessage(final List<ByteBuf> byteBuffers, final int lastRequestId) {
         notNull("stream is open", stream);
-
         if (isClosed()) {
             throw new MongoSocketClosedException("Cannot write to a closed stream", getServerAddress());
         }
-
         try {
             stream.write(byteBuffers);
         } catch (Exception e) {
             close();
-            throw translateWriteException(e);
+            throwTranslatedWriteException(e);
         }
     }
 
+    @Override
+    public void sendMessageAsync(final List<ByteBuf> byteBuffers, final int lastRequestId,
+            final SingleResultCallback<Void> callback) {
+        beginAsync().thenRun((c) -> {
+            notNull("stream is open", stream);
+            if (isClosed()) {
+                throw new MongoSocketClosedException("Cannot write to a closed stream", getServerAddress());
+            }
+            c.complete(c);
+        }).thenRunTryCatchAsyncBlocks(c -> {
+            stream.writeAsync(byteBuffers, c.asHandler());
+        }, Exception.class, (e, c) -> {
+            close();
+            throwTranslatedWriteException(e);
+        }).finish(callback);
+    }
+
     @Override
     public ResponseBuffers receiveMessage(final int responseTo) {
         assertNotNull(stream);
@@ -665,39 +679,6 @@ private ResponseBuffers receiveMessageWithAdditionalTimeout(final int additional
         }
     }
 
-    @Override
-    public void sendMessageAsync(final List<ByteBuf> byteBuffers, final int lastRequestId,
-            final SingleResultCallback<Void> callback) {
-        assertNotNull(stream);
-
-        if (isClosed()) {
-            callback.onResult(null, new MongoSocketClosedException("Can not read from a closed socket", getServerAddress()));
-            return;
-        }
-
-        writeAsync(byteBuffers, errorHandlingCallback(callback, LOGGER));
-    }
-
-    private void writeAsync(final List<ByteBuf> byteBuffers, final SingleResultCallback<Void> callback) {
-        try {
-            stream.writeAsync(byteBuffers, new AsyncCompletionHandler<Void>() {
-                @Override
-                public void completed(@Nullable final Void v) {
-                    callback.onResult(null, null);
-                }
-
-                @Override
-                public void failed(final Throwable t) {
-                    close();
-                    callback.onResult(null, translateWriteException(t));
-                }
-            });
-        } catch (Throwable t) {
-            close();
-            callback.onResult(null, t);
-        }
-    }
-
     @Override
     public void receiveMessageAsync(final int responseTo, final SingleResultCallback<ResponseBuffers> callback) {
         assertNotNull(stream);
@@ -762,6 +743,10 @@ private void updateSessionContext(final SessionContext sessionContext, final Res
         }
     }
 
+    private void throwTranslatedWriteException(final Throwable e) {
+        throw translateWriteException(e);
+    }
+
     private MongoException translateWriteException(final Throwable e) {
         if (e instanceof MongoException) {
             return (MongoException) e;

driver-core/src/test/unit/com/mongodb/internal/async/AsyncFunctionsTest.java

@@ -393,8 +393,8 @@ void testTryCatch() {
         // chain of 2 in try.
         // WARNING: "onErrorIf" will consider everything in
         // the preceding chain to be part of the try.
-        // Use nested async chains to define the beginning
-        // of the "try".
+        // Use nested async chains, or convenience methods,
+        // to define the beginning of the try.
         assertBehavesSameVariations(5,
                 () -> {
                     try {
@@ -491,6 +491,56 @@ void testTryCatch() {
                 });
     }
 
+    @Test
+    void testTryCatchHelper() {
+        assertBehavesSameVariations(4,
+                () -> {
+                    plain(0);
+                    try {
+                        sync(1);
+                    } catch (Throwable t) {
+                        plain(2);
+                        throw t;
+                    }
+                },
+                (callback) -> {
+                    beginAsync().thenRun(c -> {
+                        plain(0);
+                        c.complete(c);
+                    }).thenRunTryCatchAsyncBlocks(c -> {
+                        async(1, c);
+                    }, Throwable.class, (t, c) -> {
+                        plain(2);
+                        c.completeExceptionally(t);
+                    }).finish(callback);
+                });
+
+        assertBehavesSameVariations(5,
+                () -> {
+                    plain(0);
+                    try {
+                        sync(1);
+                    } catch (Throwable t) {
+                        plain(2);
+                        throw t;
+                    }
+                    sync(4);
+                },
+                (callback) -> {
+                    beginAsync().thenRun(c -> {
+                        plain(0);
+                        c.complete(c);
+                    }).thenRunTryCatchAsyncBlocks(c -> {
+                        async(1, c);
+                    }, Throwable.class, (t, c) -> {
+                        plain(2);
+                        c.completeExceptionally(t);
+                    }).thenRun(c -> {
+                        async(4, c);
+                    }).finish(callback);
+                });
+    }
+
     @Test
     void testTryCatchWithVariables() {
         // using supply etc.