Adjust the ScopedValue logic to correctly report the current scheduler in case of inherited scoped values

franz1981 · franz1981 · commit 1c19c0b45fe5 · 2025-11-18T16:00:33.000+01:00
diff --git a/core/src/main/java/io/netty/loom/EventLoopScheduler.java b/core/src/main/java/io/netty/loom/EventLoopScheduler.java
@@ -25,15 +25,32 @@ public EventLoopScheduler get() {
         }
     }
 
+    /**
+     * PerCarrier Read-Sub-Pollers are not created through the vThreadFactory and although the NettyScheduler
+     * can dynamically assign them to this scheduler, while running, won't see this scheduler as the current one.
+     * This means that if they try to unpark some virtual thread which belong to this scheduler,
+     * they can still wakeup the carrier thread. It's more a performance defect, but won't affect correctness.
+     */
+    public record SchedulerContext(long vThreadId, SchedulerRef scheduler) {
+
+        /**
+         * Get the assigned scheduler for the current thread, or null if this thread is not assigned to any scheduler.
+         */
+        @Override
+        public SchedulerRef scheduler() {
+            return (Thread.currentThread().threadId() == vThreadId) ? scheduler : null;
+        }
+    }
+
    private static final long MAX_WAIT_TASKS_NS = TimeUnit.HOURS.toNanos(1);
    // These are the soft-guaranteed yield times for the event loop whilst Thread.yield() is called.
    // Based on the status of the event loop (resuming from blocking or non-blocking, controlled by the running flag)
    // a different limit is applied.
    private static final long RUNNING_YIELD_US = TimeUnit.MICROSECONDS.toNanos(Integer.getInteger("io.netty.loom.running.yield.us", 1));
    private static final long IDLE_YIELD_US = TimeUnit.MICROSECONDS.toNanos(Integer.getInteger("io.netty.loom.idle.yield.us", 1));
    // This is required to allow sub-pollers to run on the correct scheduler
-   private static final ScopedValue<SchedulerRef> CURRENT_SCHEDULER = ScopedValue.newInstance();
-   private static final SchedulerRef EMPTY_REF = new SchedulerRef(null);
+   private static final ScopedValue<SchedulerContext> CURRENT_SCHEDULER = ScopedValue.newInstance();
+   private static final SchedulerContext EMPTY_SCHEDULER_CONTEXT = new SchedulerContext(-1, null);
    private final MpscUnboundedStream<Runnable> runQueue;
    private final ManualIoEventLoop ioEventLoop;
    private final Thread eventLoopThread;
@@ -42,17 +59,20 @@ public EventLoopScheduler get() {
    private volatile Runnable eventLoopContinuatioToRun;
    private final ThreadFactory vThreadFactory;
    private final AtomicBoolean eventLoopIsRunning;
-   private final SchedulerRef publishedReference;
+   private final SchedulerRef sharedRef;
 
    public EventLoopScheduler(IoEventLoopGroup parent, ThreadFactory threadFactory, IoHandlerFactory ioHandlerFactory, int resumedContinuationsExpectedCount) {
-      publishedReference = new SchedulerRef(this);
+      sharedRef = new SchedulerRef(this);
       eventLoopIsRunning = new AtomicBoolean(false);
       runQueue = new MpscUnboundedStream<>(resumedContinuationsExpectedCount);
       carrierThread = threadFactory.newThread(this::virtualThreadSchedulerLoop);
       var rawVTFactory = Thread.ofVirtual().factory();
       vThreadFactory = runnable ->
               NettyScheduler.assignUnstarted(rawVTFactory.newThread(
-                      () -> ScopedValue.where(CURRENT_SCHEDULER, publishedReference).run(runnable)), publishedReference);
+                      () ->   // this can be inherited by any thread created in the context of this virtual thread
+                              // but only the original virtual thread will have the correct scheduler context
+                              ScopedValue.where(CURRENT_SCHEDULER, new SchedulerContext(Thread.currentThread().threadId(), sharedRef)).run(runnable)
+              ), sharedRef);
       eventLoopThread = vThreadFactory.newThread(() -> FastThreadLocalThread.runWithFastThreadLocal(this::nettyEventLoop));
       ioEventLoop = new ManualIoEventLoop(parent, eventLoopThread,
             ioExecutor -> new AwakeAwareIoHandler(eventLoopIsRunning, ioHandlerFactory.newHandler(ioExecutor)));
@@ -149,7 +169,7 @@ private void virtualThreadSchedulerLoop() {
          runEventLoopContinuation();
       }
       // the event loop should be fully terminated
-      publishedReference.ref = null;
+      sharedRef.ref = null;
       runQueue.close();
       // StoreLoad barrier
       while (!runQueue.isEmpty()) {
@@ -202,11 +222,8 @@ public boolean execute(Thread.VirtualThreadTask task) {
       var currentThread = Thread.currentThread();
       if (currentThread != eventLoopThread) {
           if (currentThread != carrierThread) {
-              // this is checking for "local" submissions: it assumes that
-              // currentThreadEventLoopScheduler() is the currently assigned one, and up to date
-              // WARNING!!!!!
-              // work-stealing could break this assumption if we don't update the CURRENT_SCHEDULER scoped value accordingly
-              if (currentThreadEventLoopScheduler() != publishedReference) {
+              // this is checking for "local" submissions
+              if (currentThreadSchedulerContext().scheduler() != sharedRef) {
                   ioEventLoop.wakeup();
                   LockSupport.unpark(parkedCarrierThread);
               }
@@ -220,8 +237,7 @@ public boolean execute(Thread.VirtualThreadTask task) {
       return true;
    }
 
-   public static SchedulerRef currentThreadEventLoopScheduler() {
-      var ref = CURRENT_SCHEDULER.orElse(EMPTY_REF);
-      return ref == EMPTY_REF ? null : ref;
+   public static SchedulerContext currentThreadSchedulerContext() {
+      return CURRENT_SCHEDULER.orElse(EMPTY_SCHEDULER_CONTEXT);
    }
 }
diff --git a/core/src/main/java/io/netty/loom/NettyScheduler.java b/core/src/main/java/io/netty/loom/NettyScheduler.java
@@ -53,20 +53,17 @@ public void onStart(Thread.VirtualThreadTask virtualThreadTask) {
             if (currentThread.isVirtual()) {
                 // TODO https://github.com/openjdk/loom/blob/12ddf39bb59252a8274d8b937bd075b2a6dbc3f8/src/java.base/share/classes/java/lang/VirtualThread.java#L270C18-L270C33
                 //      in theory should be easy to provide a VirtualThreadTask::current method to avoid the ScopedValue lookup
-                var schedulerRef = EventLoopScheduler.currentThreadEventLoopScheduler();
-                // TODO per carrier sub-pollers goes here, but we want them to inherit the scheduler from the caller context
+                var ctx = EventLoopScheduler.currentThreadSchedulerContext();
+                var schedulerRef = ctx.scheduler();
+                // See https://github.com/openjdk/loom/blob/12ddf39bb59252a8274d8b937bd075b2a6dbc3f8/src/java.base/share/classes/sun/nio/ch/Poller.java#L723C48-L723C59
                 if (schedulerRef != null) {
-                    var scheduler = schedulerRef.get();
-                    // See https://github.com/openjdk/loom/blob/12ddf39bb59252a8274d8b937bd075b2a6dbc3f8/src/java.base/share/classes/sun/nio/ch/Poller.java#L723C48-L723C59
-                    if (scheduler != null) {
-                        if (virtualThreadTask.thread().getName().endsWith("-Read-Poller")) {
-                            // attach the assigned scheduler to the task
-                            virtualThreadTask.attach(schedulerRef);
-                            if (scheduler.execute(virtualThreadTask)) {
-                                return;
-                            }
-                            virtualThreadTask.attach(null);
+                    var runningScheduler = schedulerRef.get();
+                    if (runningScheduler != null && virtualThreadTask.thread().getName().endsWith("-Read-Poller")) {
+                        virtualThreadTask.attach(schedulerRef);
+                        if (runningScheduler.execute(virtualThreadTask)) {
+                            return;
                         }
+                        virtualThreadTask.attach(null);
                     }
                 }
             }
@@ -89,9 +86,9 @@ public void onStart(Thread.VirtualThreadTask virtualThreadTask) {
     public void onContinue(Thread.VirtualThreadTask virtualThreadTask) {
         var attachment = virtualThreadTask.attachment();
         if (attachment instanceof SchedulerRef ref) {
-            var scheduler = ref.get();
-            if (scheduler != null) {
-                if (scheduler.execute(virtualThreadTask)) {
+            var assignedScheduler = ref.get();
+            if (assignedScheduler != null) {
+                if (assignedScheduler.execute(virtualThreadTask)) {
                     return;
                 }
             }
diff --git a/core/src/main/java/io/netty/loom/VirtualMultithreadIoEventLoopGroup.java b/core/src/main/java/io/netty/loom/VirtualMultithreadIoEventLoopGroup.java
@@ -48,4 +48,8 @@ protected IoEventLoop newChild(Executor executor, IoHandlerFactory ioHandlerFact
       return customScheduler.ioEventLoop();
    }
 
+    @Override
+    public void close() {
+        shutdownGracefully();
+    }
 }
diff --git a/core/src/test/java/io/netty/loom/VirtualMultithreadIoEventLoopGroupTest.java b/core/src/test/java/io/netty/loom/VirtualMultithreadIoEventLoopGroupTest.java
@@ -1,37 +1,17 @@
 package io.netty.loom;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertFalse;
-import static org.junit.jupiter.api.Assertions.assertInstanceOf;
-import static org.junit.jupiter.api.Assertions.assertNotEquals;
-import static org.junit.jupiter.api.Assertions.assertNotSame;
-import static org.junit.jupiter.api.Assertions.assertNull;
-import static org.junit.jupiter.api.Assertions.assertSame;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static java.util.concurrent.StructuredTaskScope.Joiner.allSuccessfulOrThrow;
+import static org.junit.jupiter.api.Assertions.*;
 import static org.junit.jupiter.api.Assumptions.assumeTrue;
 
-import java.io.IOException;
 import java.net.InetSocketAddress;
-import java.net.ServerSocket;
-import java.net.Socket;
 import java.net.URI;
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
-import java.util.Objects;
-import java.util.concurrent.BrokenBarrierException;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.CyclicBarrier;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.ThreadFactory;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
+import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.locks.LockSupport;
 
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelInboundHandlerAdapter;
@@ -331,37 +311,43 @@ public void channelRead(io.netty.channel.ChannelHandlerContext ctx, Object msg)
       group.shutdownGracefully();
    }
 
-   @Test
-   void schedulerIsInherited() throws InterruptedException, ExecutionException {
-      var group = new VirtualMultithreadIoEventLoopGroup(1, LocalIoHandler.newFactory());
-      final Thread expectedCarrier = group.submit(() -> LoomSupport.getCarrierThread(Thread.currentThread())).get();
-      final CompletableFuture<Thread> vfactoryCarrier = new CompletableFuture<>();
-      group.execute(() -> {
-         group.vThreadFactory().newThread(() -> {
-            vfactoryCarrier.complete(LoomSupport.getCarrierThread(Thread.currentThread()));
-         }).start();
-      });
-      final CompletableFuture<Thread> inheritedCarrier = new CompletableFuture<>();
-      group.execute(() -> {
-         group.vThreadFactory().newThread(() -> {
-            Thread.ofVirtual().start(() -> {
-               inheritedCarrier.complete(LoomSupport.getCarrierThread(Thread.currentThread()));
-            });
-         }).start();
-      });
-      final CompletableFuture<Thread> inheritedVFactoryCarrier = new CompletableFuture<>();
-      group.execute(() -> {
-         group.vThreadFactory().newThread(() -> {
-            Thread.ofVirtual().factory().newThread(() -> {
-               inheritedVFactoryCarrier.complete(LoomSupport.getCarrierThread(Thread.currentThread()));
+    @Test
+    void schedulerIsNotInheritedWithThreadOfVirtual() throws InterruptedException, ExecutionException {
+        try (var group = new VirtualMultithreadIoEventLoopGroup(1, LocalIoHandler.newFactory())) {
+            final var expectedScheduler = group.submit(() -> EventLoopScheduler.currentThreadSchedulerContext().scheduler().get()).get();
+            assertNotNull(expectedScheduler);
+            final var vThreadFactory = group.submit(group::vThreadFactory).get();
+            var groupEventLoopScheduler = new CompletableFuture<EventLoopScheduler>();
+            var noEventLoopScheduler = new CompletableFuture<>();
+            vThreadFactory.newThread(() -> {
+                Thread.ofVirtual().start(() -> noEventLoopScheduler.complete(EventLoopScheduler.currentThreadSchedulerContext().scheduler()));
+                groupEventLoopScheduler.complete(EventLoopScheduler.currentThreadSchedulerContext().scheduler().get());
             }).start();
-         }).start();
-      });
-      assertEquals(expectedCarrier, vfactoryCarrier.get());
-      assertEquals(expectedCarrier, inheritedCarrier.get());
-      assertEquals(expectedCarrier, inheritedVFactoryCarrier.get());
-      group.shutdownGracefully();
-   }
+            assertEquals(expectedScheduler, groupEventLoopScheduler.get());
+            assertNull(noEventLoopScheduler.get());
+        }
+    }
+
+    @Test
+    void schedulerIsInheritedByForkedVTFromTheRightFactory() throws InterruptedException, ExecutionException {
+        try (var group = new VirtualMultithreadIoEventLoopGroup(1, LocalIoHandler.newFactory())) {
+            final var expectedEventLoopScheduler = group.submit(() -> EventLoopScheduler.currentThreadSchedulerContext().scheduler().get()).get();
+            assertNotNull(expectedEventLoopScheduler);
+            final var vThreadFactory = group.submit(group::vThreadFactory).get();
+            var forkInheritedScheduler = new CompletableFuture<EventLoopScheduler>();
+            vThreadFactory.newThread(() -> {
+                try (var scope = StructuredTaskScope.open(allSuccessfulOrThrow(),
+                        cf -> cf.withThreadFactory(vThreadFactory))) {
+                    var task = scope.fork(() -> EventLoopScheduler.currentThreadSchedulerContext().scheduler().get());
+                    scope.join();
+                    forkInheritedScheduler.complete(task.get());
+                } catch (InterruptedException e) {
+                    forkInheritedScheduler.completeExceptionally(e);
+                }
+            }).start();
+            assertEquals(expectedEventLoopScheduler, forkInheritedScheduler.get());
+        }
+    }
 
    @Test
    void eventLoopSchedulerCanMakeProgressIfTheEventLoopIsBlocked() throws BrokenBarrierException, InterruptedException, TimeoutException {
@@ -422,16 +408,20 @@ private static void spinWait(long nanos) {
         }
     }
 
-   /*
-
-   @Test
-   public void testContainsJustBuiltinPollers() throws InterruptedException {
-      assertContainsJustBuiltinPollers(builtinGlobalScheduler());
-   }
-
+    /*
    @Test
    public void testPlatformThreadSpawnsVirtualThreads() throws InterruptedException {
-      testPlatformThreadSpawnsVirtualThreads(builtinGlobalScheduler());
+       CompletableFuture<Thread.VirtualThreadScheduler> schedulerCompletableFuture = new CompletableFuture<>();
+       Thread.ofPlatform()
+               .start(() -> {
+                   Thread.ofVirtual()
+                           .start(() -> {
+                               Thread.VirtualThreadScheduler virtualThreadScheduler = NettyScheduler.internalSchedulerMappings.get(Thread.currentThread());
+                               schedulerCompletableFuture.complete(virtualThreadScheduler);
+                           });
+               });
+       Thread.VirtualThreadScheduler virtualThreadScheduler = schedulerCompletableFuture.join();
+       assertEquals(NettyScheduler.jdkBuildinScheduler, virtualThreadScheduler);
    }
 
    @Test
@@ -588,20 +578,6 @@ private static void assertContainsJustBuiltinPollers(NettyScheduler NettySchedul
                         entry.getValue() == NettyScheduler.getJdkBuildinScheduler()));
    }
 
-   private void testPlatformThreadSpawnsVirtualThreads(NettyScheduler NettyScheduler) {
-      CompletableFuture<Thread.VirtualThreadScheduler> schedulerCompletableFuture = new CompletableFuture<>();
-      Thread.ofPlatform()
-            .start(() -> {
-               Thread.ofVirtual()
-                     .start(() -> {
-                        Thread.VirtualThreadScheduler virtualThreadScheduler = NettyScheduler.internalSchedulerMappings.get(Thread.currentThread());
-                        schedulerCompletableFuture.complete(virtualThreadScheduler);
-                     });
-            });
-      Thread.VirtualThreadScheduler virtualThreadScheduler = schedulerCompletableFuture.join();
-      assertEquals(NettyScheduler.jdkBuildinScheduler, virtualThreadScheduler);
-   }
-
    private void testNettyThreadSpawnsVirtualThreads(NettyScheduler NettyScheduler) throws InterruptedException {
       record NettyVirtualThreadContext(Thread ioEventLoopThread,
                                        Thread.VirtualThreadScheduler vtScheduler,Thread.VirtualThreadScheduler vtSchedulerFromScopeValue,

Original file line number	Diff line number	Diff line change
`@@ -48,4 +48,8 @@ protected IoEventLoop newChild(Executor executor, IoHandlerFactory ioHandlerFact`
`48`	`48`	`return customScheduler.ioEventLoop();`
`49`	`49`	`}`
`50`	`50`
	`51`	`+ @Override`
	`52`	`+ public void close() {`
	`53`	`+ shutdownGracefully();`
	`54`	`+ }`
`51`	`55`	`}`