improve the ThreadLocal approach

stIncMale · stIncMale · commit 1284d59b129b · 2026-02-16T05:05:18.000-07:00
diff --git a/driver-core/src/main/com/mongodb/internal/async/function/AsyncCallbackLoop.java b/driver-core/src/main/com/mongodb/internal/async/function/AsyncCallbackLoop.java
@@ -21,8 +21,6 @@
 
 import java.util.function.Supplier;
 
-import static com.mongodb.assertions.Assertions.assertTrue;
-
 /**
  * A decorator that implements automatic repeating of an {@link AsyncCallbackRunnable}.
  * {@link AsyncCallbackLoop} may execute the original asynchronous function multiple times sequentially,
@@ -53,66 +51,103 @@ public AsyncCallbackLoop(final LoopState state, final AsyncCallbackRunnable body
 
     @Override
     public void run(final SingleResultCallback<Void> callback) {
-        body.initiateIteration(false, new ReusableLoopCallback(callback));
+        body.run(false, new ReusableLoopCallback(callback));
     }
 
     private static final class Body {
-        private final AsyncCallbackRunnable wrapped;
+        private final AsyncCallbackRunnable body;
         private final LoopState state;
-        private final ThreadLocal<Boolean> iterationIsExecutingSynchronously;
-        private final ThreadLocal<Status> status;
+        private final ThreadLocal<SameThreadDetectionStatus> sameThreadDetector;
 
-        private enum Status {
-            ITERATION_INITIATED,
-            LAST_ITERATION_COMPLETED,
-            ANOTHER_ITERATION_NEEDED
+        private enum SameThreadDetectionStatus {
+            NEGATIVE,
+            PROBING,
+            POSITIVE
         }
 
         private Body(final LoopState state, final AsyncCallbackRunnable body) {
-            this.wrapped = body;
+            this.body = body;
             this.state = state;
-            iterationIsExecutingSynchronously = ThreadLocal.withInitial(() -> false);
-            status = ThreadLocal.withInitial(() -> Status.ITERATION_INITIATED);
+            sameThreadDetector = ThreadLocal.withInitial(() -> SameThreadDetectionStatus.NEGATIVE);
         }
 
         /**
-         * Invoking this method initiates a new iteration of the loop. An iteration may be executed either
-         * synchronously or asynchronously with the execution of this method:
-         *
+         * Initiates a new iteration of the loop by invoking
+         * {@link #body}{@code .}{@link AsyncCallbackRunnable#run(SingleResultCallback) run}.
+         * The initiated iteration may be executed either synchronously or asynchronously with the method that initiated it:
          * <ul>
-         *     <li>synchronous execution: iteration completes before (in the happens-before order) the method completes;</li>
-         *     <li>asynchronous execution: the aforementioned relation does not exist.</li>
+         *     <li>synchronous execution—completion of the initiated iteration happens-before the method completion;</li>
+         *     <li>asynchronous execution—the aforementioned relation does not exist.</li>
          * </ul>
          *
+         * <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)}.
+         * 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)
+         * regardless of them being executed synchronously or asynchronously with the method that initiated them.
+         *
+         * <p>Initiating any but the {@linkplain LoopState#isFirstIteration() first} iteration is done using trampolining,
+         * which allows us to do it iteratively rather than recursively, if iterations are executed synchronously,
+         * and ensures stack usage does not increase with the number of iterations.
+         *
          * @return {@code true} iff it is known that another iteration must be initiated.
-         * Such information is available to this method only if the iteration it initiated has completed synchronously.
+         * This information is used only for trampolining, and is available only if the iteration executed synchronously.
+         *
+         * <p>It is impossible to detect whether an iteration is executed synchronously.
+         * It is, however, possible to detect whether an iteration is executed in the same thread as the method that initiated it,
+         * and we use it as proxy indicator of synchronous execution. Unfortunately, this means we do not support and behave incorrectly
+         * if an iteration is executed synchronously but in a thread different from the one in which the method that
+         * initiated the iteration was invoked.
+         *
+         * <p>The above limitation should not be a problem in practice:
+         * <ul>
+         *     <li>the only way to execute an iteration synchronously but in a different thread is to block the thread that
+         *     initiated the iteration by waiting for completion of the iteration by that other thread;</li>
+         *     <li>blocking a thread is forbidden in asynchronous code, and we do not do it;</li>
+         *     <li>therefore, we would not have an iteration that is executed synchronously but in a different thread.</li>
+         * </ul>
          */
-        Status initiateIteration(final boolean trampolining, final ReusableLoopCallback callback) {
-            iterationIsExecutingSynchronously.set(true);
-            wrapped.run((r, t) -> {
-                boolean localIterationIsExecutingSynchronously = iterationIsExecutingSynchronously.get();
+        boolean run(final boolean trampolining, final ReusableLoopCallback callback) {
+            // The `trampoliningResult` variable must be used only if the initiated iteration is executed synchronously with
+            // the current method, which must be detected separately.
+            //
+            // It may be tempting to detect whether the iteration was executed synchronously by reading from the variable
+            // and observing a write that is part of the callback execution. However, if the iteration is executed asynchronously with
+            // the current method, then the aforementioned conflicting write and read actions are not ordered by
+            // the happens-before relation, the execution contains a data race and the read is allowed to observe the write.
+            // If such observation happens when the iteration is executed asynchronously, then we have a false positive.
+            // Furthermore, depending on the nature of the value read, it may not be trustworthy.
+            boolean[] trampoliningResult = {false};
+            sameThreadDetector.set(SameThreadDetectionStatus.PROBING);
+            body.run((r, t) -> {
                 if (callback.onResult(state, r, t)) {
-                    status.set(Status.LAST_ITERATION_COMPLETED);
+                    // If we are trampolining, then here we bounce up, trampolining completes and so is the whole loop;
+                    // otherwise, the whole loop simply completes.
                     return;
                 }
-                if (trampolining && localIterationIsExecutingSynchronously) {
-                    // bounce
-                    status.set(Status.ANOTHER_ITERATION_NEEDED);
-                    return;
+                if (trampolining) {
+                    boolean sameThread = sameThreadDetector.get().equals(SameThreadDetectionStatus.PROBING);
+                    if (sameThread) {
+                        // Bounce up if we are trampolining and the iteration was executed synchronously;
+                        // otherwise proceed to begin trampolining.
+                        sameThreadDetector.set(SameThreadDetectionStatus.POSITIVE);
+                        trampoliningResult[0] = true;
+                        return;
+                    } else {
+                        sameThreadDetector.remove();
+                    }
                 }
-                Status localStatus;
-                do {
-                    localStatus = initiateIteration(true, callback);
-                } while (localStatus.equals(Status.ANOTHER_ITERATION_NEEDED));
-                status.set(localStatus);
-
-                // VAKOTODO remove thread-locals if executed asynchronously
+                boolean anotherIterationNeeded;
+                do { // trampolining
+                    anotherIterationNeeded = run(true, callback);
+                } while (anotherIterationNeeded);
             });
             try {
-                return status.get();
+                return sameThreadDetector.get().equals(SameThreadDetectionStatus.POSITIVE) && trampoliningResult[0];
             } finally {
-                status.remove();
-                iterationIsExecutingSynchronously.remove();
+                sameThreadDetector.remove();
             }
         }
     }
diff --git a/driver-core/src/test/unit/com/mongodb/internal/async/VakoTest.java b/driver-core/src/test/unit/com/mongodb/internal/async/VakoTest.java
@@ -19,6 +19,8 @@
 import com.mongodb.internal.async.function.LoopState;
 import com.mongodb.internal.time.StartTime;
 import com.mongodb.lang.Nullable;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.CsvSource;
 
@@ -28,18 +30,32 @@
 import java.time.Duration;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ForkJoinPool;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 
 import static com.mongodb.internal.async.AsyncRunnable.beginAsync;
 
 class VakoTest {
+    private static ScheduledExecutorService executor;
+
+    @BeforeAll
+    static void beforeAll() {
+        executor = Executors.newScheduledThreadPool(2);
+    }
+
+    @AfterAll
+    static void afterAll() throws InterruptedException {
+        executor.shutdownNow();
+        com.mongodb.assertions.Assertions.assertTrue(executor.awaitTermination(1, TimeUnit.MINUTES));
+    }
+
     @ParameterizedTest
     @CsvSource({
             "10"
     })
     void asyncCallbackLoop(final int iterations) throws Exception {
-        System.err.printf("baselineStackDepth=%d%n", Thread.currentThread().getStackTrace().length);
+        System.err.printf("baselineStackDepth=%d%n%n", Thread.currentThread().getStackTrace().length);
         CompletableFuture<Void> join = new CompletableFuture<>();
         LoopState loopState = new LoopState();
         new AsyncCallbackLoop(loopState, c -> {
@@ -51,78 +67,92 @@ void asyncCallbackLoop(final int iterations) throws Exception {
         }).run((r, t) -> {
             System.err.printf("test callback completed callStackDepth=%d, r=%s, t=%s%n",
                     Thread.currentThread().getStackTrace().length, r, exceptionToString(t));
-            if (t != null) {
-                join.completeExceptionally(t);
-            } else {
-                join.complete(r);
-            }
+            complete(join, r, t);
         });
         join.get();
         System.err.printf("%n%nDONE%n%n");
     }
 
+    private enum IterationExecutionType {
+        SYNC_SAME_THREAD,
+        SYNC_DIFFERENT_THREAD,
+        ASYNC,
+    }
+
     @ParameterizedTest()
     @CsvSource({
-            "0, false, 0, 10",
-            "0, true, 4, 10",
-            "4, true, 0, 10"
+            "10, 0, SYNC_SAME_THREAD, 0",
+//            "10, 0, SYNC_DIFFERENT_THREAD, 0",
+            "10, 0, ASYNC, 4",
+            "10, 4, ASYNC, 0"
     })
     void testThenRunDoWhileLoop(
-            final int blockInAsyncMethodTotalSeconds,
-            final boolean asyncExecution,
-            final int delayAsyncExecutionTotalSeconds,
-            final int counterInitialValue) throws Exception {
-        Duration blockInAsyncMethodTotalDuration = Duration.ofSeconds(blockInAsyncMethodTotalSeconds);
-        com.mongodb.assertions.Assertions.assertTrue(asyncExecution || delayAsyncExecutionTotalSeconds == 0);
+            final int counterInitialValue,
+            final int blockSyncPartOfIterationTotalSeconds,
+            final IterationExecutionType executionType,
+            final int delayAsyncExecutionTotalSeconds) 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();
-        System.err.printf("baselineStackDepth=%d%n", Thread.currentThread().getStackTrace().length);
         CompletableFuture<Void> join = new CompletableFuture<>();
-        asyncMethod1(blockInAsyncMethodTotalDuration, asyncExecution, delayAsyncExecutionTotalDuration, new Counter(counterInitialValue),
+        asyncLoop(new Counter(counterInitialValue), blockSyncPartOfIterationTotalDuration, executionType, delayAsyncExecutionTotalDuration,
             (r, t) -> {
-                System.err.printf("TEST callback completed callStackDepth=%s, r=%s, t=%s%n",
+                System.err.printf("test callback completed callStackDepth=%s, r=%s, t=%s%n",
                         Thread.currentThread().getStackTrace().length, r, exceptionToString(t));
-                if (t != null) {
-                    join.completeExceptionally(t);
-                } else {
-                    join.complete(r);
-                }
+                complete(join, r, t);
         });
-        System.err.printf("asyncMethod1 returned in %s%n", start.elapsed());
+        System.err.printf("\tasyncLoop returned in %s%n", start.elapsed());
         join.get();
         System.err.printf("%n%nDONE%n%n");
     }
 
-    private static void asyncMethod1(
-            final Duration blockInAsyncMethodTotalDuration,
-            final boolean asyncExecution,
-            final Duration delayAsyncExecutionTotalDuration,
+    private static void asyncLoop(
             final Counter counter,
+            final Duration blockSyncPartOfIterationTotalDuration,
+            final IterationExecutionType executionType,
+            final Duration delayAsyncExecutionTotalDuration,
             final SingleResultCallback<Void> callback) {
         beginAsync().thenRunDoWhileLoop(c -> {
-            sleep(blockInAsyncMethodTotalDuration.dividedBy(counter.initial()));
+            sleep(blockSyncPartOfIterationTotalDuration.dividedBy(counter.initial()));
             StartTime start = StartTime.now();
-            asyncMethod2(asyncExecution, delayAsyncExecutionTotalDuration, counter, c);
-            System.err.printf("asyncMethod2 returned in %s%n", start.elapsed());
+            asyncPartOfIteration(counter, executionType, delayAsyncExecutionTotalDuration, c);
+            System.err.printf("\tasyncPartOfIteration returned in %s%n", start.elapsed());
         }, () -> !counter.done()).finish(callback);
     }
 
-    private static void asyncMethod2(
-            final boolean asyncExecution,
-            final Duration delayAsyncExecutionTotalDuration,
+    private static void asyncPartOfIteration(
             final Counter counter,
+            final IterationExecutionType executionType,
+            final Duration delayAsyncExecutionTotalDuration,
             final SingleResultCallback<Void> callback) {
-        Runnable action = () -> {
-            sleep(delayAsyncExecutionTotalDuration.dividedBy(counter.initial()));
+        Runnable asyncPartOfIteration = () -> {
             counter.countDown();
             StartTime start = StartTime.now();
             callback.complete(callback);
-            System.err.printf("asyncMethod2 callback.complete returned in %s%n", start.elapsed());
+            System.err.printf("\tasyncPartOfIteration callback.complete returned in %s%n", start.elapsed());
         };
-        if (asyncExecution) {
-            ForkJoinPool.commonPool().execute(action);
-        } else {
-            action.run();
+        switch (executionType) {
+            case SYNC_SAME_THREAD: {
+                asyncPartOfIteration.run();
+                break;
+            }
+            case SYNC_DIFFERENT_THREAD: {
+                Thread guaranteedDifferentThread = new Thread(asyncPartOfIteration);
+                guaranteedDifferentThread.start();
+                join(guaranteedDifferentThread);
+                break;
+            }
+            case ASYNC: {
+                executor.schedule(asyncPartOfIteration,
+                        delayAsyncExecutionTotalDuration.dividedBy(counter.initial()).toNanos(), TimeUnit.NANOSECONDS);
+                break;
+            }
+            default: {
+                com.mongodb.assertions.Assertions.fail(executionType.toString());
+            }
         }
     }
 
@@ -170,6 +200,23 @@ private static String exceptionToString(@Nullable final Throwable t) {
         }
     }
 
+    private static <T> void complete(final CompletableFuture<T> future, @Nullable final T result, @Nullable final Throwable t) {
+        if (t != null) {
+            future.completeExceptionally(t);
+        } else {
+            future.complete(result);
+        }
+    }
+
+    private static void join(final Thread thread) {
+        try {
+            thread.join();
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            throw new RuntimeException(e);
+        }
+    }
+
     private static void sleep(final Duration duration) {
         if (duration.isZero()) {
             return;
@@ -179,7 +226,6 @@ private static void sleep(final Duration duration) {
             long durationMsPartFromNsPart = TimeUnit.MILLISECONDS.convert(duration.getNano(), TimeUnit.NANOSECONDS);
             long sleepMs = TimeUnit.MILLISECONDS.convert(duration.getSeconds(), TimeUnit.SECONDS) + durationMsPartFromNsPart;
             int sleepNs = Math.toIntExact(durationNsPart - TimeUnit.NANOSECONDS.convert(durationMsPartFromNsPart, TimeUnit.MILLISECONDS));
-            System.err.printf("sleeping for %d ms %d ns%n", sleepMs, sleepNs);
             Thread.sleep(sleepMs, sleepNs);
         } catch (InterruptedException e) {
             Thread.currentThread().interrupt();
