chase fibers branch part 1

franz1981 · franz1981 · commit 7ebbeec57194 · 2025-11-18T16:00:33.000+01:00
diff --git a/core/src/main/java/io/netty/loom/GlobalDelegateThreadNettyScheduler.java b/core/src/main/java/io/netty/loom/GlobalDelegateThreadNettyScheduler.java
@@ -1,6 +1,8 @@
 package io.netty.loom;
 
+import java.lang.invoke.VarHandle;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * Global scheduler — provides execution for virtual threads that do not specify their own scheduler.
@@ -18,61 +20,53 @@
 
 public class GlobalDelegateThreadNettyScheduler implements Thread.VirtualThreadScheduler {
 
-   final Thread.VirtualThreadScheduler jdkBuildinScheduler;
+    private static GlobalDelegateThreadNettyScheduler INSTANCE;
 
-   final ConcurrentHashMap<Thread, Thread.VirtualThreadScheduler> internalSchedulerMappings = new ConcurrentHashMap<>();
+    final Thread.VirtualThreadScheduler jdkBuildinScheduler;
 
-   public GlobalDelegateThreadNettyScheduler(Thread.VirtualThreadScheduler jdkBuildinScheduler) {
-      this.jdkBuildinScheduler = jdkBuildinScheduler;
-   }
+    final ConcurrentHashMap<Thread, AtomicReference<VirtualThreadNettyScheduler>> unstartedThreads = new ConcurrentHashMap<>();
 
-   @Override
-   public void execute(Thread vthread, Runnable task) {
-      // When scheduling virtual threads, `GlobalDelegateThreadNettyScheduler::execute` will not be invoked concurrently for the same virtual thread,
-      // so it’s safe to replace `computeIfAbsent` with a `get + put` approach to reduce overhead.
-      Thread.VirtualThreadScheduler internalScheduler = internalSchedulerMappings.get(vthread);
-      if (internalScheduler == null) {
-         internalScheduler = determineScheduler();
-         // we remember ANY mapping, including the JDK built-in scheduler ones
-         internalSchedulerMappings.put(vthread, internalScheduler);
-      }
-      // regardless which scheduler is used, we need to remove the mapping once the vthread is terminated
-      internalScheduler.execute(vthread, () -> {
-         try {
-            task.run();
-         } finally {
-            if (vthread.getState() == Thread.State.TERMINATED) {
-               internalSchedulerMappings.remove(vthread);
+    public GlobalDelegateThreadNettyScheduler(Thread.VirtualThreadScheduler jdkBuildinScheduler) {
+        this.jdkBuildinScheduler = jdkBuildinScheduler;
+        INSTANCE = this;
+        VarHandle.storeStoreFence();
+    }
+
+    @Override
+    public void onStart(Thread.VirtualThreadTask virtualThreadTask) {
+        // TODO this is not great for 2 reasons:
+        // 1. we are doing a remove on a concurrent map even for v threads which are not really interesting to us
+        // 2. if a vThread will never start, it will leak here forever
+        var assignedSchedulerRef = unstartedThreads.remove(virtualThreadTask.thread());
+        if (assignedSchedulerRef == null) {
+            // TODO per carrier sub-pollers goes here, but we want them to inherit the scheduler from the caller context
+        }
+        var scheduler = assignedSchedulerRef != null ? assignedSchedulerRef.get() : null;
+        if (scheduler != null) {
+            // attach the assigned scheduler to the task
+            virtualThreadTask.attach(assignedSchedulerRef);
+            if (scheduler.execute(virtualThreadTask)) {
+                return;
             }
-         }
-      });
-   }
+            // the v thread has been rejected by its assigned scheduler, clean it up and fallback to JDK
+            virtualThreadTask.attach(null);
+        }
+        jdkBuildinScheduler.onStart(virtualThreadTask);
+    }
+
+    @Override
+    public void onContinue(Thread.VirtualThreadTask virtualThreadTask) {
 
-   private Thread.VirtualThreadScheduler determineScheduler() {
-      Thread callerThread = Thread.currentThread();
-      // platform thread
-      if (!callerThread.isVirtual()) {
-         return jdkBuildinScheduler;
-      }
-      VirtualThreadNettyScheduler current = VirtualThreadNettyScheduler.current();
-      // The current thread was spawned from a specific VirtualThreadNettyScheduler,
-      // so we continue using that scheduler.
-      if (current != null) {
-         return current;
-      }
-      Thread.VirtualThreadScheduler parentScheduler = internalSchedulerMappings.get(callerThread);
-      if (parentScheduler != null) {
-         return parentScheduler;
-      }
+    }
 
-      // if the caller thread is not known we assume it won't benefit from inheriting, so its children
-      // will use the JDK built-in scheduler
-      return jdkBuildinScheduler;
-   }
+    // just for benchmark
+    public Thread.VirtualThreadScheduler getJdkBuildinScheduler() {
+        return jdkBuildinScheduler;
+    }
 
-   // just for benchmark
-   public Thread.VirtualThreadScheduler getJdkBuildinScheduler() {
-      return jdkBuildinScheduler;
-   }
+    static Thread assignUnstarted(Thread unstarted, AtomicReference<VirtualThreadNettyScheduler> schedulerRef) {
+        INSTANCE.unstartedThreads.put(unstarted, schedulerRef);
+        return unstarted;
+    }
 }
 
diff --git a/core/src/main/java/io/netty/loom/VirtualThreadNettyScheduler.java b/core/src/main/java/io/netty/loom/VirtualThreadNettyScheduler.java
@@ -1,11 +1,10 @@
 package io.netty.loom;
 
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.Executor;
-import java.util.concurrent.RejectedExecutionException;
+import java.util.Queue;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.locks.LockSupport;
 
 import io.netty.channel.IoEventLoopGroup;
@@ -14,59 +13,45 @@
 import io.netty.util.concurrent.FastThreadLocalThread;
 import io.netty.util.internal.shaded.org.jctools.queues.MpscUnboundedArrayQueue;
 
-public class VirtualThreadNettyScheduler implements Thread.VirtualThreadScheduler {
+public class VirtualThreadNettyScheduler {
 
    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));
-   private static final ScopedValue<Thread.VirtualThreadScheduler> CURRENT_SCHEDULER = ScopedValue.newInstance();
-   private static final Thread.VirtualThreadScheduler EMPTY = new Thread.VirtualThreadScheduler() {
-      @Override
-      public void execute(Thread vthread, Runnable command) {
-         throw new RejectedExecutionException("VirtualThreadScheduler is empty");
-      }
-   };
-
-   private final MpscUnboundedArrayQueue<Runnable> externalContinuations;
+   // This is required to allow sub-pollers to run on the correct scheduler
+   private static final ScopedValue<AtomicReference<VirtualThreadNettyScheduler>> CURRENT_SCHEDULER = ScopedValue.newInstance();
+   private volatile MpscUnboundedArrayQueue<Runnable> submissionQueue;
+   private final MpscUnboundedArrayQueue<Runnable> runQueue;
    private final ManualIoEventLoop ioEventLoop;
    private final Thread eventLoopThread;
    private final Thread carrierThread;
    private volatile Thread parkedCarrierThread;
-   private volatile Runnable eventLoopContinuation;
-   private volatile boolean submittedEventLoopContinuation;
-   private final CountDownLatch eventLoopContinuationAvailable;
+   private volatile Runnable eventLoopContinuatioToRun;
    private final ThreadFactory vThreadFactory;
    private final AtomicBoolean running;
+   private final AtomicReference<VirtualThreadNettyScheduler> schedulerReference;
 
    public VirtualThreadNettyScheduler(IoEventLoopGroup parent, ThreadFactory threadFactory, IoHandlerFactory ioHandlerFactory, int resumedContinuationsExpectedCount) {
-      this.running = new AtomicBoolean(false);
-      this.externalContinuations = new MpscUnboundedArrayQueue<>(resumedContinuationsExpectedCount);
-      this.carrierThread = threadFactory.newThread(this::virtualThreadSchedulerLoop);
-      var builder = Thread.ofVirtual().scheduler(this);
-      ThreadFactory rawVTFactory = builder.factory();
-      this.vThreadFactory = (r) -> rawVTFactory.newThread(() -> ScopedValue.where(CURRENT_SCHEDULER, this).run(r));
-      this.eventLoopThread = vThreadFactory.newThread(
-            // this is enabling the adaptive allocator to use unshared magazines
-            () -> FastThreadLocalThread.runWithFastThreadLocal(this::nettyEventLoop));
-      this.ioEventLoop = new ManualIoEventLoop(parent, eventLoopThread,
+      schedulerReference = new AtomicReference<>(this);
+      running = new AtomicBoolean(false);
+      runQueue = new MpscUnboundedArrayQueue<>(resumedContinuationsExpectedCount);
+      submissionQueue = runQueue;
+      carrierThread = threadFactory.newThread(this::virtualThreadSchedulerLoop);
+      var rawVTFactory = Thread.ofVirtual().factory();
+      vThreadFactory = runnable ->
+              GlobalDelegateThreadNettyScheduler.assignUnstarted(rawVTFactory.newThread(
+                      () -> ScopedValue.where(CURRENT_SCHEDULER, schedulerReference).run(runnable)), schedulerReference);
+      eventLoopThread = vThreadFactory.newThread(() -> FastThreadLocalThread.runWithFastThreadLocal(this::nettyEventLoop));
+      ioEventLoop = new ManualIoEventLoop(parent, eventLoopThread,
             ioExecutor -> new AwakeAwareIoHandler(running, ioHandlerFactory.newHandler(ioExecutor)));
-      // we can start the carrier only after all the fields are initialized
-      eventLoopContinuationAvailable = new CountDownLatch(1);
       carrierThread.start();
-      // TODO we cannot make the virtual thread factory available until the event loop v thread is started:
-      //      we can save this if we can query the virtual thread associated with a continuation
-      try {
-         eventLoopContinuationAvailable.await();
-      } catch (InterruptedException e) {
-         throw new RuntimeException(e);
-      }
    }
 
    int externalContinuationsCount() {
-      return externalContinuations.size();
+      return runQueue.size();
    }
 
    public ThreadFactory virtualThreadFactory() {
@@ -111,29 +96,35 @@ private boolean runIO(boolean canBlock) {
                running.set(true);
             }
          } else {
-            running.set(false);
+            running.set(true);
          }
       }
       return ioEventLoop.runNow(RUNNING_YIELD_US) == 0;
    }
 
+   private boolean runEventLoopContinuation() {
+      assert Thread.currentThread() == carrierThread;
+      var eventLoopContinuation = this.eventLoopContinuatioToRun;
+      if (eventLoopContinuation != null) {
+         this.eventLoopContinuatioToRun = null;
+         eventLoopContinuation.run();
+         return true;
+      }
+      return false;
+   }
+
    private void virtualThreadSchedulerLoop() {
       // start the event loop thread
       var eventLoop = this.ioEventLoop;
       eventLoopThread.start();
-      // we expect here the continuation to be set up already
-      var eventLoopContinuation = this.eventLoopContinuation;
-      assert eventLoopContinuation != null && eventLoopContinuationAvailable.getCount() == 0;
+      assert eventLoopContinuatioToRun != null;
       // we keep on running until the event loop is shutting-down
       while (!eventLoop.isTerminated()) {
          // if the event loop was idle, we apply a different limit to the yield time
          final boolean eventLoopRunning = running.get();
          final long yieldDurationNs = eventLoopRunning ? RUNNING_YIELD_US : IDLE_YIELD_US;
          int count = runExternalContinuations(yieldDurationNs);
-         if (submittedEventLoopContinuation) {
-            submittedEventLoopContinuation = false;
-            eventLoopContinuation.run();
-         } else if (count == 0) {
+         if (!runEventLoopContinuation() && count == 0) {
             // nothing to run, including the event loop: we can park
             parkedCarrierThread = carrierThread;
             if (canBlock()) {
@@ -142,24 +133,27 @@ private void virtualThreadSchedulerLoop() {
             parkedCarrierThread = null;
          }
       }
-      while (!canBlock()) {
-         // we still have continuations to run, let's run them
+      // make sure the event loop thread is fully terminated
+      // TODO verify that isAlive works!
+      while (eventLoopThread.isAlive()) {
          runExternalContinuations(RUNNING_YIELD_US);
-         if (submittedEventLoopContinuation) {
-            submittedEventLoopContinuation = false;
-            eventLoopContinuation.run();
-         }
+         runEventLoopContinuation();
+      }
+      schedulerReference.set(null);
+      submissionQueue = null;
+      // StoreLoad barrier
+      while (!runQueue.isEmpty()) {
+         runExternalContinuations(IDLE_YIELD_US);
       }
    }
 
    private boolean canBlock() {
-      return externalContinuations.isEmpty() && !submittedEventLoopContinuation;
+      return runQueue.isEmpty() && eventLoopContinuatioToRun == null;
    }
 
    private int runExternalContinuations(long deadlineNs) {
-      assert eventLoopContinuation != null;
       final long startDrainingNs = System.nanoTime();
-      var ready = this.externalContinuations;
+      var ready = this.runQueue;
       int runContinuations = 0;
       for (; ; ) {
          var continuation = ready.poll();
@@ -176,51 +170,48 @@ private int runExternalContinuations(long deadlineNs) {
       return runContinuations;
    }
 
-   @Override
-   public void execute(Thread vthread, Runnable command) {
-      // The default scheduler won't shut down, but Netty's event loop can!
-      Runnable eventLoopContinuation = this.eventLoopContinuation;
-      if (eventLoopContinuation == null) {
-         eventLoopContinuation = setEventLoopContinuation(command);
+   private boolean eventLoopContinuation(Thread.VirtualThreadTask task) {
+      if (eventLoopContinuatioToRun != null) {
+         assert task.thread() != eventLoopThread;
+         return false;
       }
-      if (eventLoopContinuation == command) {
-         submittedEventLoopContinuation = true;
-      } else {
-         if (ioEventLoop.isTerminated()) {
-            throw new RejectedExecutionException("event loop is shutting down");
-         }
-         externalContinuations.offer(command);
-      }
-      if (!inEventLoop(Thread.currentThread())) {
-         // TODO: if we have access to the scheduler brought by the continuation,
-         //       we could skip the wakeup if matches with this.
-         ioEventLoop.wakeup();
-         LockSupport.unpark(parkedCarrierThread);
-      } else if (eventLoopContinuation != command && !running.get()) {
-         // since the event loop was blocked, it's fine if we try to consume some continuations, if any
-         Thread.yield();
+      if (task.thread() == eventLoopThread) {
+         eventLoopContinuatioToRun = task;
+         return true;
       }
+      return false;
    }
 
-   private boolean inEventLoop(Thread thread) {
-       if (!thread.isVirtual()) {
-           return ioEventLoop.inEventLoop(thread);
-       }
-
-       return CURRENT_SCHEDULER.orElse(EMPTY) == this;
-   }
-
-   private Runnable setEventLoopContinuation(Runnable command) {
-      // this is the first command, we need to set the continuation
-      this.eventLoopContinuation = command;
-      // we need to notify the event loop that we have a continuation
-      eventLoopContinuationAvailable.countDown();
-      return command;
-   }
-
-   public static VirtualThreadNettyScheduler current() {
-       Thread.VirtualThreadScheduler virtualThreadScheduler = CURRENT_SCHEDULER.orElse(EMPTY);
-       return virtualThreadScheduler == EMPTY ? null : (VirtualThreadNettyScheduler) virtualThreadScheduler;
+   public boolean execute(Thread.VirtualThreadTask task) {
+      boolean isEventLoopContinuation = eventLoopContinuation(task);
+      if (!isEventLoopContinuation) {
+        var q = submissionQueue;
+        if (q == null) {
+            return false;
+        }
+        q.offer(task);
+        // double-check in case we are shutting down
+        if (submissionQueue == null) {
+            // TODO
+            // What if:
+            // 1) we added to the queue
+            // 2) we got paused for long time
+            // 3) the scheduler found the task and run it
+            // 4) we are resuming here and return false
+            // This is not ok, need to be fixed!
+            return false;
+        }
+      }
+      if (!ioEventLoop.inEventLoop(Thread.currentThread())) {
+          // TODO: if we have access to the scheduler brought by the continuation,
+          //       we could skip the wakeup if matches with this.
+          ioEventLoop.wakeup();
+          LockSupport.unpark(parkedCarrierThread);
+      } else if (!isEventLoopContinuation && !running.get()) {
+          assert eventLoopContinuatioToRun == null;
+          // since the event loop was blocked, it's fine if we try to consume some continuations, if any
+          Thread.yield();
+      }
+      return true;
    }
-
 }
