Implement OIDC auth for async

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

@@ -17,7 +17,6 @@
 
 package com.mongodb.assertions;
 
-import com.mongodb.internal.async.SingleResultCallback;
 import com.mongodb.lang.Nullable;
 
 import java.util.Collection;
@@ -50,25 +49,6 @@ public static <T> T notNull(final String name, final T value) {
         return value;
     }
 
-    /**
-     * Throw IllegalArgumentException if the value is null.
-     *
-     * @param name  the parameter name
-     * @param value the value that should not be null
-     * @param callback  the callback that also is passed the exception if the value is null
-     * @param <T>   the value type
-     * @return the value
-     * @throws java.lang.IllegalArgumentException if value is null
-     */
-    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);
-            throw exception;
-        }
-        return value;
-    }
-
     /**
      * Throw IllegalStateException if the condition if false.
      *
@@ -82,22 +62,6 @@ public static void isTrue(final String name, final boolean condition) {
         }
     }
 
-    /**
-     * Throw IllegalStateException if the condition if false.
-     *
-     * @param name      the name of the state that is being checked
-     * @param condition the condition about the parameter to check
-     * @param callback  the callback that also is passed the exception if the condition is not true
-     * @throws java.lang.IllegalStateException if the condition is false
-     */
-    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);
-            throw exception;
-        }
-    }
-
     /**
      * Throw IllegalArgumentException if the condition if false.
      *

driver-core/src/main/com/mongodb/internal/Locks.java

@@ -17,6 +17,8 @@
 package com.mongodb.internal;
 
 import com.mongodb.MongoInterruptedException;
+import com.mongodb.internal.async.AsyncRunnable;
+import com.mongodb.internal.async.SingleResultCallback;
 
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.StampedLock;
@@ -33,7 +35,23 @@ public static void withLock(final Lock lock, final Runnable action) {
         });
     }
 
-    public static <V> V withLock(final StampedLock lock, final Supplier<V> supplier) {
+    public static void withLockAsync(final StampedLock lock, final AsyncRunnable runnable,
+            final SingleResultCallback<Void> callback) {
+        long stamp;
+        try {
+            stamp = lock.writeLockInterruptibly();
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            callback.onResult(null, new MongoInterruptedException("Interrupted waiting for lock", e));
+            return;
+        }
+
+        runnable.thenAlwaysRunAndFinish(() -> {
+            lock.unlockWrite(stamp);
+        }, callback);
+    }
+
+    public static void withLock(final StampedLock lock, final Runnable runnable) {
         long stamp;
         try {
             stamp = lock.writeLockInterruptibly();
@@ -42,7 +60,7 @@ public static <V> V withLock(final StampedLock lock, final Supplier<V> supplier)
             throw new MongoInterruptedException("Interrupted waiting for lock", e);
         }
         try {
-            return supplier.get();
+            runnable.run();
         } finally {
             lock.unlockWrite(stamp);
         }
@@ -55,15 +73,15 @@ public static <V> V withLock(final Lock lock, final Supplier<V> supplier) {
     public static <V, E extends Exception> V checkedWithLock(final Lock lock, final CheckedSupplier<V, E> supplier) throws E {
         try {
             lock.lockInterruptibly();
-            try {
-                return supplier.get();
-            } finally {
-                lock.unlock();
-            }
         } catch (InterruptedException e) {
             Thread.currentThread().interrupt();
             throw new MongoInterruptedException("Interrupted waiting for lock", e);
         }
+        try {
+            return supplier.get();
+        } finally {
+            lock.unlock();
+        }
     }
 
     private Locks() {

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

@@ -0,0 +1,188 @@
+/*
+ * 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;
+
+/**
+ * See AsyncRunnableTest for usage
+ */
+public interface AsyncRunnable {
+
+    static AsyncRunnable beginAsync() {
+        return (c) -> c.onResult(null, null);
+    }
+
+    void runUnsafe(SingleResultCallback<Void> callback); // NoResultCallback
+
+    /**
+     * Must be invoked at end of async chain. Wraps the lambda in an error
+     * handler.
+     * @param callback the callback provided by the method the chain is used in
+     */
+    default void finish(final SingleResultCallback<Void> callback) {
+        try {
+            this.runUnsafe((v, e) -> {
+                try {
+                    callback.onResult(v, e);
+                } catch (Throwable t) {
+                    throw new CallbackThrew("Unexpected Throwable thrown from callback: ", e);
+                }
+            });
+        } catch (CallbackThrew t) {
+            // ignore
+        } catch (Throwable t) {
+            callback.onResult(null, t);
+        }
+    }
+
+    /**
+     * 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.onResult(null, e);
+                return;
+            }
+            try {
+                runnable.run();
+            } catch (Throwable t) {
+                callback.onResult(null, t);
+                return;
+            }
+            callback.onResult(null, null);
+        });
+    }
+
+    /**
+     * 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) {
+                callback.onResult(null, t);
+                return;
+            }
+            callback.onResult(r, e);
+        });
+    }
+
+    /**
+     * @param runnable The async runnable to run after this one
+     * @return the composition of this and the runnable
+     */
+    default AsyncRunnable thenRun(final AsyncRunnable runnable) {
+        return (c) -> {
+            this.finish((r, e) -> {
+                if (e != null) {
+                    c.onResult(null, e);
+                    return;
+                }
+                try {
+                    runnable.finish(c);
+                } catch (Throwable t) {
+                    c.onResult(null, t);
+                }
+            });
+        };
+    }
+
+    /**
+     * @param supplier The supplier to supply using after this runnable.
+     * @return the composition of this runnable and the supplier
+     * @param <T> The return type of the supplier
+     */
+    default <T> AsyncSupplier<T> thenSupply(final AsyncSupplier<T> supplier) {
+        return (c) -> {
+            this.finish((r, e) -> {
+                if (e != null) {
+                    c.onResult(null, e);
+                    return;
+                }
+                try {
+                    supplier.finish(c);
+                } catch (Throwable t) {
+                    c.onResult(null, t);
+                }
+            });
+        };
+    }
+
+    /**
+     * @param errorCheck A check, comparable to a catch-if/otherwise-rethrow
+     * @param runnable   The branch to execute if the error matches
+     * @return The composition of this, and the conditional branch
+     */
+    default AsyncRunnable onErrorRunIf(
+            final Predicate<Throwable> errorCheck,
+            final AsyncRunnable runnable) {
+        return (callback) -> this.finish((r, e) -> {
+            if (e == null) {
+                callback.onResult(r, null);
+                return;
+            }
+            try {
+                boolean check = errorCheck.test(e);
+                if (check) {
+                    runnable.finish(callback);
+                    return;
+                }
+            } catch (Throwable t) {
+                callback.onResult(null, t);
+                return;
+            }
+            callback.onResult(r, 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 this.thenRun(callback -> {
+            new RetryingAsyncCallbackSupplier<Void>(
+                    new RetryState(),
+                    (rs, lastAttemptFailure) -> shouldRetry.test(lastAttemptFailure),
+                    cb -> runnable.finish(cb)
+            ).get(callback);
+        });
+    }
+
+    final class CallbackThrew extends AssertionError {
+        private static final long serialVersionUID = 875624357420415700L;
+
+        public CallbackThrew(final String s, final Throwable e) {
+            super(s, e);
+        }
+    }
+}

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

@@ -0,0 +1,78 @@
+/*
+ * 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;
+
+import static com.mongodb.internal.async.AsyncRunnable.CallbackThrew;
+
+/**
+ * See AsyncRunnableTest for usage
+ */
+public interface AsyncSupplier<T> {
+
+    void supplyUnsafe(SingleResultCallback<T> 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) {
+        try {
+            this.supplyUnsafe((v, e) -> {
+                try {
+                    callback.onResult(v, e);
+                } catch (Throwable t) {
+                    throw new CallbackThrew("Unexpected Throwable thrown from callback: ", e);
+                }
+            });
+        } catch (CallbackThrew t) {
+            // ignore
+        } catch (Throwable t) {
+            callback.onResult(null, t);
+        }
+    }
+
+    /**
+     * @see AsyncRunnable#onErrorRunIf(Predicate, AsyncRunnable).
+     *
+     * @param errorCheck A check, comparable to a catch-if/otherwise-rethrow
+     * @param supplier   The branch to execute if the error matches
+     * @return The composition of this, and the conditional branch
+     */
+    default AsyncSupplier<T> onErrorSupplyIf(
+            final Predicate<Throwable> errorCheck,
+            final AsyncSupplier<T> supplier) {
+        return (callback) -> this.finish((r, e) -> {
+            if (e == null) {
+                callback.onResult(r, null);
+                return;
+            }
+            try {
+                boolean check = errorCheck.test(e);
+                if (check) {
+                    supplier.finish(callback);
+                    return;
+                }
+            } catch (Throwable t) {
+                callback.onResult(null, t);
+                return;
+            }
+            callback.onResult(r, e);
+        });
+    }
+}