third attempt

Author: stIncMale

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

@@ -16,11 +16,17 @@
 package com.mongodb.internal.async.function;
 
 import com.mongodb.annotations.NotThreadSafe;
+import com.mongodb.assertions.Assertions;
+import com.mongodb.connection.AsyncCompletionHandler;
 import com.mongodb.internal.async.SingleResultCallback;
+import com.mongodb.internal.async.function.CallbackChain.Element;
 import com.mongodb.lang.Nullable;
 
 import java.util.function.Supplier;
 
+import static com.mongodb.assertions.Assertions.assertNotNull;
+import static com.mongodb.assertions.Assertions.assertNull;
+
 /**
  * A decorator that implements automatic repeating of an {@link AsyncCallbackRunnable}.
  * {@link AsyncCallbackLoop} may execute the original asynchronous function multiple times sequentially,
@@ -39,10 +45,8 @@
  */
 @NotThreadSafe
 public final class AsyncCallbackLoop implements AsyncCallbackRunnable {
-    @Nullable
-    private final CallbackChain chain;
     private final LoopState state;
-    private final AsyncCallbackRunnable body;
+    private final Body body;
 
     /**
      * @param state The {@link LoopState} to be deemed as initial for the purpose of the new {@link AsyncCallbackLoop}.
@@ -53,43 +57,67 @@ public AsyncCallbackLoop(final LoopState state, final AsyncCallbackRunnable body
     }
 
     public AsyncCallbackLoop(final boolean optimized, final LoopState state, final AsyncCallbackRunnable body) {
-        this.chain = optimized ? new CallbackChain() : null;
         this.state = state;
-        this.body = body;
+        this.body = new Body(optimized, body);
     }
 
     @Override
     public void run(final SingleResultCallback<Void> callback) {
         body.run(new LoopingCallback(callback));
     }
 
+    private final class Body {
+        private final AsyncCallbackRunnable wrapped;
+        @Nullable
+        private final CallbackChain chain;
+
+        private Body(final boolean optimized, final AsyncCallbackRunnable body) {
+            this.wrapped = body;
+            this.chain = optimized ? new CallbackChain() : null;
+        }
+
+        @Nullable
+        Element run(final LoopingCallback loopingCallback) {
+            Element[] mutableElement = new Element[1];
+            wrapped.run((r, t) -> {
+                Element nextCallbackToComplete = loopingCallback.onResult(r, t);
+                if (!CallbackChain.execute(chain, nextCallbackToComplete)) {
+                    mutableElement[0] = nextCallbackToComplete;
+                }
+            });
+            return mutableElement[0];
+        }
+    }
+
     /**
      * This callback is allowed to be completed more than once.
      */
     @NotThreadSafe
-    private class LoopingCallback implements SingleResultCallback<Void> {
+    private class LoopingCallback {
         private final SingleResultCallback<Void> wrapped;
 
         LoopingCallback(final SingleResultCallback<Void> callback) {
             wrapped = callback;
         }
 
-        @Override
-        public void onResult(@Nullable final Void result, @Nullable final Throwable t) {
+        @Nullable
+        public Element onResult(@Nullable final Void result, @Nullable final Throwable t) {
             if (t != null) {
                 wrapped.onResult(null, t);
+                return null;
             } else {
                 boolean continueLooping;
                 try {
                     continueLooping = state.advance();
                 } catch (Throwable e) {
                     wrapped.onResult(null, e);
-                    return;
+                    return null;
                 }
                 if (continueLooping) {
-                    CallbackChain.execute(chain, () -> body.run(this));
+                    return () -> body.run(this);
                 } else {
                     wrapped.onResult(result, null);
+                    return null;
                 }
             }
         }

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

@@ -15,52 +15,49 @@
  */
 package com.mongodb.internal.async.function;
 
+import com.mongodb.annotations.NotThreadSafe;
 import com.mongodb.lang.Nullable;
 
 import java.util.concurrent.atomic.AtomicReference;
 
-import static com.mongodb.assertions.Assertions.assertNotNull;
-import static com.mongodb.assertions.Assertions.assertNull;
-
+@NotThreadSafe
 final class CallbackChain {
-    @Nullable
-    private Runnable next;
-    private int runEnteredCounter;
-    private final AtomicReference<Thread> thread;
+    private int enteringCounter;
 
     CallbackChain() {
-        runEnteredCounter = 0;
-        thread = new AtomicReference<>();
+        enteringCounter = 0;
     }
 
-    static void execute(@Nullable final CallbackChain chain, final Runnable next) {
+    static boolean execute(@Nullable final CallbackChain chain, @Nullable final Element element) {
+        if (element == null) {
+            return false;
+        }
         if (chain != null) {
-            chain.execute(next);
+            return chain.execute(element);
         } else {
-            next.run();
+            element.execute();
+            return true;
         }
     }
 
-    // VAKOTODO figure out thread safety
-    private void execute(final Runnable next) {
-        assertNotNull(next);
-        assertNull(this.next);
-        this.next = next;
-
-//        if (!thread.compareAndSet(null, Thread.currentThread())) {
-//            assertTrue(Thread.currentThread() == thread.get());
-//        }
-        boolean recursive = ++runEnteredCounter > 1;
+    private boolean execute(final Element element) {
+        boolean reentered = ++enteringCounter > 1;
         try {
-            if (recursive) {
-                return;
+            if (reentered) {
+                return false;
             }
-            for (Runnable localNext = next; localNext != null; localNext = this.next) {
-                this.next = null;
-                localNext.run();
+            Element next = element.execute();
+            while (next != null) {
+                next = next.execute();
             }
         } finally {
-            runEnteredCounter--;
+            enteringCounter--;
         }
+        return true;
+    }
+
+    interface Element {
+        @Nullable
+        Element execute();
     }
 }