Improve vThreadFactory method

Author: franz1981

core/src/main/java/io/netty/loom/VirtualMultithreadIoEventLoopGroup.java

@@ -1,25 +1,21 @@
 package io.netty.loom;
 
+import java.util.ArrayList;
 import java.util.IdentityHashMap;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.atomic.AtomicLong;
 
 import io.netty.channel.IoEventLoop;
 import io.netty.channel.IoHandlerFactory;
 import io.netty.channel.MultiThreadIoEventLoopGroup;
-import io.netty.util.concurrent.FastThreadLocal;
 
 public class VirtualMultithreadIoEventLoopGroup extends MultiThreadIoEventLoopGroup {
 
    private static final int RESUMED_CONTINUATIONS_EXPECTED_COUNT = Integer.getInteger("io.netty.loom.resumed.continuations", 1024);
+   private ArrayList<EventLoopScheduler> eventLoopSchedulers;
+   private AtomicLong nextScheduler;
    private ThreadFactory threadFactory;
-   private IdentityHashMap<Thread, EventLoopScheduler> schedulers;
-   private final FastThreadLocal<EventLoopScheduler> v_thread_factory = new FastThreadLocal<>() {
-      @Override
-      protected EventLoopScheduler initialValue() {
-         return schedulers.get(Thread.currentThread());
-      }
-   };
 
    public VirtualMultithreadIoEventLoopGroup(int nThreads, IoHandlerFactory ioHandlerFactory) {
       super(nThreads, (Executor) command -> {
@@ -30,21 +26,46 @@ public VirtualMultithreadIoEventLoopGroup(int nThreads, IoHandlerFactory ioHandl
       }
    }
 
+    /**
+     * Return a {@link ThreadFactory} that creates virtual threads tied to an
+     * {@link EventLoopScheduler} of this group.
+     *
+     * <p>If the current thread has an associated {@link EventLoopScheduler} whose
+     * {@link io.netty.channel.IoEventLoop#parent()} is this group, that scheduler's
+     * {@code virtualThreadFactory()} is returned so newly created virtual threads
+     * are associated with the current event loop.</p>
+     *
+     * <p>Otherwise a randomly assigned scheduler from this group is used and its
+     * {@code virtualThreadFactory()} is returned.</p>
+     *
+     * @return a {@link ThreadFactory} producing virtual threads backed by an
+     * {@link EventLoopScheduler} of this group
+     **/
    public ThreadFactory vThreadFactory() {
-      return v_thread_factory.get().virtualThreadFactory();
+      var schedulerRef = EventLoopScheduler.currentThreadSchedulerContext().scheduler();
+      if (schedulerRef != null) {
+          var scheduler = schedulerRef.get();
+          if (scheduler != null && scheduler.ioEventLoop().parent() == this) {
+              return scheduler.virtualThreadFactory();
+          }
+      }
+      // assign a random one
+      int schedulerIndex = (int) (nextScheduler.getAndIncrement() % executorCount());
+      return eventLoopSchedulers.get(schedulerIndex).virtualThreadFactory();
    }
 
    @Override
    protected IoEventLoop newChild(Executor executor, IoHandlerFactory ioHandlerFactory,
                                   @SuppressWarnings("unused") Object... args) {
-      if (schedulers == null) {
-         schedulers = new IdentityHashMap<>();
+      if (eventLoopSchedulers == null) {
+         eventLoopSchedulers = new ArrayList<>(executorCount());
+         nextScheduler = new AtomicLong();
       }
       if (threadFactory == null) {
          threadFactory = newDefaultThreadFactory();
       }
       var customScheduler = new EventLoopScheduler(this, threadFactory, ioHandlerFactory, RESUMED_CONTINUATIONS_EXPECTED_COUNT);
-      schedulers.put(customScheduler.eventLoopThread(), customScheduler);
+      eventLoopSchedulers.add(customScheduler);
       return customScheduler.ioEventLoop();
    }