refactor the ThreadLocal approach

Author: stIncMale

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

@@ -51,7 +51,7 @@ public AsyncCallbackLoop(final LoopState state, final AsyncCallbackRunnable body
 
     @Override
     public void run(final SingleResultCallback<Void> callback) {
-        body.run(false, new ReusableLoopCallback(callback));
+        body.run(false, callback);
     }
 
     private static final class Body {
@@ -82,7 +82,7 @@ private Body(final LoopState state, final AsyncCallbackRunnable body) {
          *
          * <p>If another iteration is needed, it is initiated from the callback passed to
          * {@link #body}{@code .}{@link AsyncCallbackRunnable#run(SingleResultCallback) run}
-         * by invoking {@link #run(boolean, ReusableLoopCallback)}.
+         * by invoking {@link #run(boolean, SingleResultCallback)}.
          * Completing the initiated iteration is {@linkplain SingleResultCallback#onResult(Object, Throwable) invoking} the callback.
          * Thus, it is guaranteed that all iterations are executed sequentially with each other
          * (that is, completion of one iteration happens-before initiation of the next one)
@@ -109,7 +109,7 @@ private Body(final LoopState state, final AsyncCallbackRunnable body) {
          *     <li>therefore, we would not have an iteration that is executed synchronously but in a different thread.</li>
          * </ul>
          */
-        boolean run(final boolean trampolining, final ReusableLoopCallback callback) {
+        boolean run(final boolean trampolining, final SingleResultCallback<Void> afterLoopCallback) {
             // The `trampoliningResult` variable must be used only if the initiated iteration is executed synchronously with
             // the current method, which must be detected separately.
             //
@@ -122,7 +122,7 @@ boolean run(final boolean trampolining, final ReusableLoopCallback callback) {
             boolean[] trampoliningResult = {false};
             sameThreadDetector.set(SameThreadDetectionStatus.PROBING);
             body.run((r, t) -> {
-                if (callback.onResult(state, r, t)) {
+                if (completeIfNeeded(afterLoopCallback, r, t)) {
                     // If we are trampolining, then here we bounce up, trampolining completes and so is the whole loop;
                     // otherwise, the whole loop simply completes.
                     return;
@@ -139,9 +139,10 @@ boolean run(final boolean trampolining, final ReusableLoopCallback callback) {
                         sameThreadDetector.remove();
                     }
                 }
+                // trampolining
                 boolean anotherIterationNeeded;
-                do { // trampolining
-                    anotherIterationNeeded = run(true, callback);
+                do {
+                    anotherIterationNeeded = run(true, afterLoopCallback);
                 } while (anotherIterationNeeded);
             });
             try {
@@ -150,39 +151,28 @@ boolean run(final boolean trampolining, final ReusableLoopCallback callback) {
                 sameThreadDetector.remove();
             }
         }
-    }
-
-    /**
-     * This callback is allowed to be {@linkplain #onResult(LoopState, Void, Throwable) completed} more than once.
-     */
-    @NotThreadSafe
-    private static final class ReusableLoopCallback {
-        private final SingleResultCallback<Void> wrapped;
-
-        ReusableLoopCallback(final SingleResultCallback<Void> callback) {
-            wrapped = callback;
-        }
 
         /**
-         * @return {@code true} iff the {@linkplain ReusableLoopCallback#ReusableLoopCallback(SingleResultCallback) wrapped}
-         * {@link SingleResultCallback} is {@linkplain SingleResultCallback#onResult(Object, Throwable) completed}.
+         * @return {@code true} iff the {@code afterLoopCallback} was
+         * {@linkplain SingleResultCallback#onResult(Object, Throwable) completed}.
          */
-        public boolean onResult(final LoopState state, @Nullable final Void result, @Nullable final Throwable t) {
+        private boolean completeIfNeeded(final SingleResultCallback<Void> afterLoopCallback,
+                @Nullable final Void result, @Nullable final Throwable t) {
             if (t != null) {
-                wrapped.onResult(null, t);
+                afterLoopCallback.onResult(null, t);
                 return true;
             } else {
-                boolean continueLooping;
+                boolean anotherIterationNeeded;
                 try {
-                    continueLooping = state.advance();
+                    anotherIterationNeeded = state.advance();
                 } catch (Throwable e) {
-                    wrapped.onResult(null, e);
+                    afterLoopCallback.onResult(null, e);
                     return true;
                 }
-                if (continueLooping) {
+                if (anotherIterationNeeded) {
                     return false;
                 } else {
-                    wrapped.onResult(result, null);
+                    afterLoopCallback.onResult(result, null);
                     return true;
                 }
             }

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

@@ -32,6 +32,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeUnit;
 
 import static com.mongodb.internal.async.AsyncRunnable.beginAsync;
@@ -77,28 +78,32 @@ private enum IterationExecutionType {
         SYNC_SAME_THREAD,
         SYNC_DIFFERENT_THREAD,
         ASYNC,
+        MIXED_SYNC_SAME_AND_ASYNC
     }
 
     @ParameterizedTest()
     @CsvSource({
-            "10, 0, SYNC_SAME_THREAD, 0",
-//            "10, 0, SYNC_DIFFERENT_THREAD, 0",
-            "10, 0, ASYNC, 4",
-            "10, 4, ASYNC, 0"
+            "10, 0, SYNC_SAME_THREAD, 0, true",
+//            "10, 0, SYNC_DIFFERENT_THREAD, 0, true",
+            "10, 0, ASYNC, 4, true",
+            "10, 4, ASYNC, 0, true",
+            "1_000_000, 0, MIXED_SYNC_SAME_AND_ASYNC, 0, false",
     })
     void testThenRunDoWhileLoop(
             final int counterInitialValue,
             final int blockSyncPartOfIterationTotalSeconds,
             final IterationExecutionType executionType,
-            final int delayAsyncExecutionTotalSeconds) throws Exception {
+            final int delayAsyncExecutionTotalSeconds,
+            final boolean verbose) throws Exception {
         System.err.printf("baselineStackDepth=%d%n%n", Thread.currentThread().getStackTrace().length);
         Duration blockSyncPartOfIterationTotalDuration = Duration.ofSeconds(blockSyncPartOfIterationTotalSeconds);
         com.mongodb.assertions.Assertions.assertTrue(
                 executionType.equals(IterationExecutionType.ASYNC) || delayAsyncExecutionTotalSeconds == 0);
         Duration delayAsyncExecutionTotalDuration = Duration.ofSeconds(delayAsyncExecutionTotalSeconds);
         StartTime start = StartTime.now();
         CompletableFuture<Void> join = new CompletableFuture<>();
-        asyncLoop(new Counter(counterInitialValue), blockSyncPartOfIterationTotalDuration, executionType, delayAsyncExecutionTotalDuration,
+        asyncLoop(new Counter(counterInitialValue, verbose),
+                blockSyncPartOfIterationTotalDuration, executionType, delayAsyncExecutionTotalDuration, verbose,
             (r, t) -> {
                 System.err.printf("test callback completed callStackDepth=%s, r=%s, t=%s%n",
                         Thread.currentThread().getStackTrace().length, r, exceptionToString(t));
@@ -114,25 +119,31 @@ private static void asyncLoop(
             final Duration blockSyncPartOfIterationTotalDuration,
             final IterationExecutionType executionType,
             final Duration delayAsyncExecutionTotalDuration,
+            final boolean verbose,
             final SingleResultCallback<Void> callback) {
         beginAsync().thenRunDoWhileLoop(c -> {
             sleep(blockSyncPartOfIterationTotalDuration.dividedBy(counter.initial()));
             StartTime start = StartTime.now();
-            asyncPartOfIteration(counter, executionType, delayAsyncExecutionTotalDuration, c);
-            System.err.printf("\tasyncPartOfIteration returned in %s%n", start.elapsed());
+            asyncPartOfIteration(counter, executionType, delayAsyncExecutionTotalDuration, verbose, c);
+            if (verbose) {
+                System.err.printf("\tasyncPartOfIteration returned in %s%n", start.elapsed());
+            }
         }, () -> !counter.done()).finish(callback);
     }
 
     private static void asyncPartOfIteration(
             final Counter counter,
             final IterationExecutionType executionType,
             final Duration delayAsyncExecutionTotalDuration,
+            final boolean verbose,
             final SingleResultCallback<Void> callback) {
         Runnable asyncPartOfIteration = () -> {
             counter.countDown();
             StartTime start = StartTime.now();
             callback.complete(callback);
-            System.err.printf("\tasyncPartOfIteration callback.complete returned in %s%n", start.elapsed());
+            if (verbose) {
+                System.err.printf("\tasyncPartOfIteration callback.complete returned in %s%n", start.elapsed());
+            }
         };
         switch (executionType) {
             case SYNC_SAME_THREAD: {
@@ -150,6 +161,15 @@ private static void asyncPartOfIteration(
                         delayAsyncExecutionTotalDuration.dividedBy(counter.initial()).toNanos(), TimeUnit.NANOSECONDS);
                 break;
             }
+            case MIXED_SYNC_SAME_AND_ASYNC: {
+                if (ThreadLocalRandom.current().nextBoolean()) {
+                    asyncPartOfIteration.run();
+                } else {
+                    executor.schedule(asyncPartOfIteration,
+                            delayAsyncExecutionTotalDuration.dividedBy(counter.initial()).toNanos(), TimeUnit.NANOSECONDS);
+                }
+                break;
+            }
             default: {
                 com.mongodb.assertions.Assertions.fail(executionType.toString());
             }
@@ -159,10 +179,12 @@ private static void asyncPartOfIteration(
     private static final class Counter {
         private final int initial;
         private int current;
+        private final boolean verbose;
 
-        Counter(final int initial) {
+        Counter(final int initial, final boolean verbose) {
             this.initial = initial;
             this.current = initial;
+            this.verbose = verbose;
         }
 
         int initial() {
@@ -173,7 +195,10 @@ void countDown() {
             com.mongodb.assertions.Assertions.assertTrue(current > 0);
             int previous = current;
             int decremented = --current;
-            System.err.printf("counted %d->%d callStackDepth=%d %n", previous, decremented, Thread.currentThread().getStackTrace().length);
+            if (verbose || decremented % 100_000 == 0) {
+                System.err.printf("counted %d->%d callStackDepth=%d %n",
+                        previous, decremented, Thread.currentThread().getStackTrace().length);
+            }
         }
 
         boolean done() {