Add async API (2)

driver-core/src/main/com/mongodb/assertions/Assertions.java

@@ -92,7 +92,7 @@ public static <T> Iterable<T> notNullElements(final String name, final Iterable<
     public static <T> T notNull(final String name, final T value, final SingleResultCallback<?> callback) {
         if (value == null) {
             IllegalArgumentException exception = new IllegalArgumentException(name + " can not be null");
-            callback.onResult(null, exception);
+            callback.completeExceptionally(exception);
             throw exception;
         }
         return value;
@@ -122,7 +122,7 @@ public static void isTrue(final String name, final boolean condition) {
     public static void isTrue(final String name, final boolean condition, final SingleResultCallback<?> callback) {
         if (!condition) {
             IllegalStateException exception = new IllegalStateException("state should be: " + name);
-            callback.onResult(null, exception);
+            callback.completeExceptionally(exception);
             throw exception;
         }
     }

driver-core/src/main/com/mongodb/connection/AsyncCompletionHandler.java

@@ -16,6 +16,7 @@
 
 package com.mongodb.connection;
 
+import com.mongodb.internal.async.SingleResultCallback;
 import com.mongodb.lang.Nullable;
 
 /**
@@ -38,4 +39,17 @@ public interface AsyncCompletionHandler<T> {
      * @param t the exception that describes the failure
      */
     void failed(Throwable t);
+
+    /**
+     * @return this handler as a callback
+     */
+    default SingleResultCallback<T> asCallback() {
+        return (r, t) -> {
+            if (t != null) {
+                failed(t);
+            } else {
+                completed(r);
+            }
+        };
+    }
 }

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

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2008-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.mongodb.internal.async;
+
+/**
+ * See tests for usage (AsyncFunctionsTest).
+ * <p>
+ * This class is not part of the public API and may be removed or changed at any time
+ */
+@FunctionalInterface
+public interface AsyncConsumer<T> extends AsyncFunction<T, Void> {
+}

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

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2008-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.mongodb.internal.async;
+
+/**
+ * See tests for usage (AsyncFunctionsTest).
+ * <p>
+ * This class is not part of the public API and may be removed or changed at any time
+ */
+@FunctionalInterface
+public interface AsyncFunction<T, R> {
+    /**
+     * This should not be called externally, but should be implemented as a
+     * lambda. To "finish" an async chain, use one of the "finish" methods.
+     */
+    void unsafeFinish(T value, SingleResultCallback<R> callback);
+}

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

@@ -0,0 +1,158 @@
+/*
+ * Copyright 2008-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.mongodb.internal.async;
+
+import com.mongodb.internal.async.function.RetryState;
+import com.mongodb.internal.async.function.RetryingAsyncCallbackSupplier;
+
+import java.util.function.Predicate;
+import java.util.function.Supplier;
+
+/**
+ * See tests for usage (AsyncFunctionsTest).
+ * <p>
+ * This class is not part of the public API and may be removed or changed at any time
+ */
+@FunctionalInterface
+public interface AsyncRunnable extends AsyncSupplier<Void>, AsyncConsumer<Void> {
+
+    static AsyncRunnable beginAsync() {
+        return (c) -> c.complete(c);
+    }
+
+    /**
+     * Must be invoked at end of async chain
+     * @param runnable the sync code to invoke (under non-exceptional flow)
+     *                 prior to the callback
+     * @param callback the callback provided by the method the chain is used in
+     */
+    default void thenRunAndFinish(final Runnable runnable, final SingleResultCallback<Void> callback) {
+        this.finish((r, e) -> {
+            if (e != null) {
+                callback.completeExceptionally(e);
+                return;
+            }
+            try {
+                runnable.run();
+            } catch (Throwable t) {
+                callback.completeExceptionally(t);
+                return;
+            }
+            callback.complete(callback);
+        });
+    }
+
+    /**
+     * See {@link #thenRunAndFinish(Runnable, SingleResultCallback)}, but the runnable
+     * will always be executed, including on the exceptional path.
+     * @param runnable the runnable
+     * @param callback the callback
+     */
+    default void thenAlwaysRunAndFinish(final Runnable runnable, final SingleResultCallback<Void> callback) {
+        this.finish((r, e) -> {
+            try {
+                runnable.run();
+            } catch (Throwable t) {
+                if (e != null) {
+                    t.addSuppressed(e);
+                }
+                callback.completeExceptionally(t);
+                return;
+            }
+            callback.completeExceptionally(e);
+        });
+    }
+
+    /**
+     * @param runnable The async runnable to run after this runnable
+     * @return the composition of this runnable and the runnable, a runnable
+     */
+    default AsyncRunnable thenRun(final AsyncRunnable runnable) {
+        return (c) -> {
+            this.unsafeFinish((r, e) -> {
+                if (e == null) {
+                    runnable.unsafeFinish(c);
+                } else {
+                    c.completeExceptionally(e);
+                }
+            });
+        };
+    }
+
+    /**
+     * @param condition the condition to check
+     * @param runnable The async runnable to run after this runnable,
+     *                 if and only if the condition is met
+     * @return the composition of this runnable and the runnable, a runnable
+     */
+    default AsyncRunnable thenRunIf(final Supplier<Boolean> condition, final AsyncRunnable runnable) {
+        return (callback) -> {
+            this.unsafeFinish((r, e) -> {
+                if (e != null) {
+                    callback.completeExceptionally(e);
+                    return;
+                }
+                boolean matched;
+                try {
+                    matched = condition.get();
+                } catch (Throwable t) {
+                    callback.completeExceptionally(t);
+                    return;
+                }
+                if (matched) {
+                    runnable.unsafeFinish(callback);
+                } else {
+                    callback.complete(callback);
+                }
+            });
+        };
+    }
+
+    /**
+     * @param supplier The supplier to supply using after this runnable
+     * @return the composition of this runnable and the supplier, a supplier
+     * @param <R> The return type of the resulting supplier
+     */
+    default <R> AsyncSupplier<R> thenSupply(final AsyncSupplier<R> supplier) {
+        return (c) -> {
+            this.unsafeFinish((r, e) -> {
+                if (e == null) {
+                    supplier.unsafeFinish(c);
+                } else {
+                    c.completeExceptionally(e);
+                }
+            });
+        };
+    }
+
+    /**
+     * @param runnable    the runnable to loop
+     * @param shouldRetry condition under which to retry
+     * @return the composition of this, and the looping branch
+     * @see RetryingAsyncCallbackSupplier
+     */
+    default AsyncRunnable thenRunRetryingWhile(
+            final AsyncRunnable runnable, final Predicate<Throwable> shouldRetry) {
+        return thenRun(callback -> {
+            new RetryingAsyncCallbackSupplier<Void>(
+                    new RetryState(),
+                    (rs, lastAttemptFailure) -> shouldRetry.test(lastAttemptFailure),
+                    cb -> runnable.finish(cb) // finish is required here, to handle exceptions
+            ).get(callback);
+        });
+    }
+}

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

@@ -0,0 +1,137 @@
+/*
+ * Copyright 2008-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.mongodb.internal.async;
+
+import java.util.function.Predicate;
+
+
+/**
+ * See tests for usage (AsyncFunctionsTest).
+ * <p>
+ * This class is not part of the public API and may be removed or changed at any time
+ */
+@FunctionalInterface
+public interface AsyncSupplier<T> extends AsyncFunction<Void, T> {
+    /**
+     * This should not be called externally to this API. It should be
+     * implemented as a lambda. To "finish" an async chain, use one of
+     * the "finish" methods.
+     *
+     * @see #finish(SingleResultCallback)
+     */
+    void unsafeFinish(SingleResultCallback<T> callback);
+
+    /**
+     * This is the async variant of a supplier's get method.
+     * This method must only be used when this AsyncSupplier corresponds
+     * to a {@link java.util.function.Supplier} (and is therefore being
+     * used within an async chain method lambda).
+     * @param callback the callback
+     */
+    default void getAsync(final SingleResultCallback<T> callback) {
+        unsafeFinish(callback);
+    }
+
+    @Override
+    default void unsafeFinish(final Void value, final SingleResultCallback<T> callback) {
+        unsafeFinish(callback);
+    }
+
+    /**
+     * Must be invoked at end of async chain.
+     * @param callback the callback provided by the method the chain is used in
+     */
+    default void finish(final SingleResultCallback<T> callback) {
+        final boolean[] callbackInvoked = {false};
+        try {
+            this.unsafeFinish((v, e) -> {
+                callbackInvoked[0] = true;
+                callback.onResult(v, e);
+            });
+        } catch (Throwable t) {
+            if (callbackInvoked[0]) {
+                throw t;
+            } else {
+                callback.completeExceptionally(t);
+            }
+        }
+    }
+
+    /**
+     * @param function The async function to run after this supplier
+     * @return the composition of this supplier and the function, a supplier
+     * @param <R> The return type of the resulting supplier
+     */
+    default <R> AsyncSupplier<R> thenApply(final AsyncFunction<T, R> function) {
+        return (c) -> {
+            this.unsafeFinish((v, e) -> {
+                if (e == null) {
+                    function.unsafeFinish(v, c);
+                } else {
+                    c.completeExceptionally(e);
+                }
+            });
+        };
+    }
+
+
+    /**
+     * @param consumer The async consumer to run after this supplier
+     * @return the composition of this supplier and the consumer, a runnable
+     */
+    default AsyncRunnable thenConsume(final AsyncConsumer<T> consumer) {
+        return (c) -> {
+            this.unsafeFinish((v, e) -> {
+                if (e == null) {
+                    consumer.unsafeFinish(v, c);
+                } else {
+                    c.completeExceptionally(e);
+                }
+            });
+        };
+    }
+
+    /**
+     * @param errorCheck A check, comparable to a catch-if/otherwise-rethrow
+     * @param errorFunction   The branch to execute if the error matches
+     * @return The composition of this, and the conditional branch
+     */
+    default AsyncSupplier<T> onErrorIf(
+            final Predicate<Throwable> errorCheck,
+            final AsyncFunction<Throwable, T> errorFunction) {
+        return (callback) -> this.finish((r, e) -> {
+            if (e == null) {
+                callback.complete(r);
+                return;
+            }
+            boolean errorMatched;
+            try {
+                errorMatched = errorCheck.test(e);
+            } catch (Throwable t) {
+                t.addSuppressed(e);
+                callback.completeExceptionally(t);
+                return;
+            }
+            if (errorMatched) {
+                errorFunction.unsafeFinish(e, callback);
+            } else {
+                callback.completeExceptionally(e);
+            }
+        });
+    }
+
+}