Expose event loop -> thread factory mapping

Author: franz1981

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

@@ -1,6 +1,6 @@
 package io.netty.loom;
 
-import java.util.ArrayList;
+import java.util.*;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
@@ -14,7 +14,8 @@ public class VirtualMultithreadIoEventLoopGroup extends MultiThreadIoEventLoopGr
 
 	private static final int RESUMED_CONTINUATIONS_EXPECTED_COUNT = Integer
 			.getInteger("io.netty.loom.resumed.continuations", 1024);
-	private ArrayList<EventLoopScheduler> eventLoopSchedulers;
+	private Map<IoEventLoop, EventLoopScheduler> eventSchedulerMappings;
+	private List<EventLoopScheduler> schedulers;
 	private AtomicLong nextScheduler;
 	private ThreadFactory threadFactory;
 
@@ -31,6 +32,19 @@ private static void validateNettyAvailability() {
 		}
 	}
 
+	/**
+	 * Return a {@link ThreadFactory} that creates virtual threads tied to the
+	 * specified {@link IoEventLoop}'s {@link EventLoopScheduler}. Returns
+	 * {@code null} if the provided event loop is not associated with this group.
+	 */
+	public ThreadFactory vThreadFactoryOf(IoEventLoop eventLoop) {
+		EventLoopScheduler scheduler = eventSchedulerMappings.get(eventLoop);
+		if (scheduler == null) {
+			return null;
+		}
+		return scheduler.virtualThreadFactory();
+	}
+
 	/**
 	 * Return a {@link ThreadFactory} that creates virtual threads tied to an
 	 * {@link EventLoopScheduler} of this group.
@@ -60,23 +74,25 @@ public ThreadFactory vThreadFactory() {
 		}
 		// assign a random one
 		int schedulerIndex = (int) (nextScheduler.getAndIncrement() % executorCount());
-		return eventLoopSchedulers.get(schedulerIndex).virtualThreadFactory();
+		return schedulers.get(schedulerIndex).virtualThreadFactory();
 	}
 
 	@Override
 	protected IoEventLoop newChild(Executor executor, IoHandlerFactory ioHandlerFactory,
 			@SuppressWarnings("unused") Object... args) {
 		validateNettyAvailability();
-		if (eventLoopSchedulers == null) {
-			eventLoopSchedulers = new ArrayList<>(executorCount());
+		if (eventSchedulerMappings == null) {
+			eventSchedulerMappings = new IdentityHashMap<>(executorCount());
+			schedulers = new ArrayList<>(executorCount());
 			nextScheduler = new AtomicLong();
 		}
 		if (threadFactory == null) {
 			threadFactory = newDefaultThreadFactory();
 		}
 		var customScheduler = new EventLoopScheduler(this, threadFactory, ioHandlerFactory,
 				RESUMED_CONTINUATIONS_EXPECTED_COUNT);
-		eventLoopSchedulers.add(customScheduler);
+		eventSchedulerMappings.put(customScheduler.ioEventLoop(), customScheduler);
+		schedulers.add(customScheduler);
 		return customScheduler.ioEventLoop();
 	}
 

core/src/test/java/io/netty/loom/VirtualMultithreadIoEventLoopGroupTest.java

@@ -20,22 +20,15 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Stream;
 
-import io.netty.channel.Channel;
-import io.netty.channel.ChannelInboundHandlerAdapter;
-import io.netty.channel.ChannelInitializer;
-import io.netty.channel.IoHandle;
-import io.netty.channel.IoHandler;
-import io.netty.channel.IoHandlerContext;
-import io.netty.channel.IoHandlerFactory;
-import io.netty.channel.IoRegistration;
-import io.netty.channel.epoll.Epoll;
+import io.netty.channel.*;
 import io.netty.channel.local.LocalIoHandler;
 import io.netty.channel.epoll.EpollIoHandler;
 import io.netty.channel.epoll.EpollServerSocketChannel;
 
 import io.netty.channel.uring.IoUring;
 import io.netty.channel.uring.IoUringIoHandler;
 import io.netty.channel.uring.IoUringServerSocketChannel;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
 
@@ -475,7 +468,7 @@ void schedulerIsNotLeakingIfItsThreadFactoryOutliveIt(Transport transport)
 	@ParameterizedTest(name = "{index} => transport={0}")
 	@MethodSource("transportsAllowLocal")
 	void virtualThreadCanMakeProgressEvenIfEventLoopIsClosed(Transport transport)
-			throws InterruptedException, ExecutionException, TimeoutException, BrokenBarrierException {
+			throws InterruptedException, ExecutionException, BrokenBarrierException {
 		var group = new VirtualMultithreadIoEventLoopGroup(1, transport.handlerFactory());
 		final var barrier = new CyclicBarrier(2);
 		final var vThreadFactory = group.submit(group::vThreadFactory).get();
@@ -736,4 +729,23 @@ void testShutdownSchedulerOnLongBlockingIO(Transport transport)
 			}
 		}
 	}
+
+	@Test
+	void vThreadFactoryMappingShouldReturnNullIfNoneIsFound() throws InterruptedException, ExecutionException {
+		try (var otherGroup = new VirtualMultithreadIoEventLoopGroup(1, LocalIoHandler.newFactory());
+				var group = new VirtualMultithreadIoEventLoopGroup(1, NioIoHandler.newFactory())) {
+			var otherEventLoop = otherGroup.next();
+			assertNull(group.vThreadFactoryOf(otherEventLoop));
+		}
+	}
+
+	@Test
+	void vThreadFactoryMappingShouldReturnTheRightOne() throws InterruptedException, ExecutionException {
+		try (var group = new VirtualMultithreadIoEventLoopGroup(1, NioIoHandler.newFactory())) {
+			var eventLoopScheduler = group
+					.submit(() -> EventLoopScheduler.currentThreadSchedulerContext().scheduler().get()).get();
+			assertSame(eventLoopScheduler.virtualThreadFactory(),
+					group.vThreadFactoryOf(eventLoopScheduler.ioEventLoop()));
+		}
+	}
 }