Repository: netty/netty
Pull request #15243: Add API to wrap virtual thread task for FastThreadLocal support [Port of #15222]
Author: login: yawkat
Created: 2025-05-22T16:11:37Z | Merged: 2025-05-23T04:57:01Z



Diff:
diff --git a/buffer/src/main/java/io/netty5/buffer/adapt/AdaptivePoolingAllocator.java b/buffer/src/main/java/io/netty5/buffer/adapt/AdaptivePoolingAllocator.java
index 9226f1e596f..258d5dbb7ec 100644
--- a/buffer/src/main/java/io/netty5/buffer/adapt/AdaptivePoolingAllocator.java
+++ b/buffer/src/main/java/io/netty5/buffer/adapt/AdaptivePoolingAllocator.java
@@ -140,7 +140,7 @@ protected Object initialValue() {
                         return NO_MAGAZINE;
                     }
                     Magazine mag = new Magazine(AdaptivePoolingAllocator.this, executor);
-                    if (FastThreadLocalThread.willCleanupFastThreadLocals(Thread.currentThread())) {
+                    if (FastThreadLocalThread.currentThreadWillCleanupFastThreadLocals()) {
                         // Only add it to the liveMagazines if we can guarantee that onRemoval(...) is called,
                         // as otherwise we might end up holding the reference forever.
                         liveMagazines.add(mag);
@@ -214,7 +214,7 @@ public Buffer allocate(int size) {
             int sizeBucket = AllocationStatistics.sizeBucket(size); // Compute outside of Magazine lock for better ILP.
             FastThreadLocal<Object> threadLocalMagazine = this.threadLocalMagazine;
             Thread currentThread = Thread.currentThread();
-            if (threadLocalMagazine != null && currentThread instanceof FastThreadLocalThread) {
+            if (threadLocalMagazine != null && FastThreadLocalThread.currentThreadHasFastThreadLocal()) {
                 Object mag = threadLocalMagazine.get();
                 if (mag != NO_MAGAZINE) {
                     return ((Magazine) mag).allocate(size, sizeBucket);
diff --git a/buffer/src/main/java/io/netty5/buffer/internal/CleanerPool.java b/buffer/src/main/java/io/netty5/buffer/internal/CleanerPool.java
index a8d9cc5b321..4f0dbb07a08 100644
--- a/buffer/src/main/java/io/netty5/buffer/internal/CleanerPool.java
+++ b/buffer/src/main/java/io/netty5/buffer/internal/CleanerPool.java
@@ -111,7 +111,7 @@ private static final class CleanerThreadLocal extends FastThreadLocal<Cleaner> {
 
         @Override
         protected Cleaner initialValue() {
-            if (!EVENT_LOOP_USE_POOL && Thread.currentThread() instanceof FastThreadLocalThread) {
+            if (!EVENT_LOOP_USE_POOL && FastThreadLocalThread.currentThreadHasFastThreadLocal()) {
                 // Allocate one dedicated cleaner for the caller event-loop thread
                 return Cleaner.create();
             }
diff --git a/buffer/src/main/java/io/netty5/buffer/pool/PooledBufferAllocator.java b/buffer/src/main/java/io/netty5/buffer/pool/PooledBufferAllocator.java
index 5304d601e88..d432ff5d383 100644
--- a/buffer/src/main/java/io/netty5/buffer/pool/PooledBufferAllocator.java
+++ b/buffer/src/main/java/io/netty5/buffer/pool/PooledBufferAllocator.java
@@ -437,7 +437,7 @@ protected synchronized PoolThreadCache initialValue() {
             final EventExecutor executor = ThreadExecutorMap.currentExecutor();
             if (useCacheForAllThreads ||
                 // If the current thread is a FastThreadLocalThread we will always use the cache
-                current instanceof FastThreadLocalThread ||
+                FastThreadLocalThread.currentThreadHasFastThreadLocal() ||
                 // The Thread is used by an EventExecutor, let's use the cache as the chances are good that we
                 // will allocate a lot!
                 executor != null) {
diff --git a/common/src/main/java/io/netty5/util/Recycler.java b/common/src/main/java/io/netty5/util/Recycler.java
index 1568d79ed0d..875c9c64f50 100644
--- a/common/src/main/java/io/netty5/util/Recycler.java
+++ b/common/src/main/java/io/netty5/util/Recycler.java
@@ -142,7 +142,9 @@ protected Recycler(int maxCapacityPerThread, int ratio, int chunkSize) {
 
     @SuppressWarnings("unchecked")
     public final T get() {
-        if (maxCapacityPerThread == 0 || Thread.currentThread().isVirtual()) {
+        if (maxCapacityPerThread == 0 ||
+                (Thread.currentThread().isVirtual() &&
+                        !FastThreadLocalThread.currentThreadHasFastThreadLocal())) {
             return newObject((Handle<T>) NOOP_HANDLE);
         }
         LocalPool<T> localPool = threadLocal.get();
@@ -165,7 +167,7 @@ public final T get() {
 
     @VisibleForTesting
     final int threadLocalSize() {
-        if (Thread.currentThread().isVirtual()) {
+        if (Thread.currentThread().isVirtual() && !FastThreadLocalThread.currentThreadHasFastThreadLocal()) {
             return 0;
         }
         LocalPool<T> localPool = threadLocal.getIfExists();
@@ -266,7 +268,8 @@ private static final class LocalPool<T> implements MessagePassingQueue.Consumer<
             this.chunkSize = chunkSize;
             batch = new ArrayDeque<DefaultHandle<T>>(chunkSize);
             Thread currentThread = Thread.currentThread();
-            owner = !BATCH_FAST_TL_ONLY || currentThread instanceof FastThreadLocalThread ? currentThread : null;
+            owner = !BATCH_FAST_TL_ONLY || FastThreadLocalThread.currentThreadHasFastThreadLocal()
+                    ? currentThread : null;
             if (BLOCKING_POOL) {
                 pooledHandles = new BlockingMessageQueue<>(maxCapacity);
             } else {
diff --git a/common/src/main/java/io/netty5/util/concurrent/FastThreadLocalThread.java b/common/src/main/java/io/netty5/util/concurrent/FastThreadLocalThread.java
index f1a8f9ce893..ddf2c93c16d 100644
--- a/common/src/main/java/io/netty5/util/concurrent/FastThreadLocalThread.java
+++ b/common/src/main/java/io/netty5/util/concurrent/FastThreadLocalThread.java
@@ -19,6 +19,9 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.Arrays;
+import java.util.concurrent.atomic.AtomicReference;
+
 /**
  * A special {@link Thread} that provides fast access to {@link FastThreadLocal} variables.
  */
@@ -26,6 +29,11 @@ public class FastThreadLocalThread extends Thread {
 
     private static final Logger logger = LoggerFactory.getLogger(FastThreadLocalThread.class);
 
+    /**
+     * Sorted array of thread IDs that are treated like {@link FastThreadLocalThread}.
+     */
+    private static final AtomicReference<long[]> fallbackThreads = new AtomicReference<>(null);
+
     // This will be set to true if we have a chance to wrap the Runnable.
     private final boolean cleanupFastThreadLocals;
 
@@ -96,21 +104,106 @@ final void setThreadLocalMap(InternalThreadLocalMap threadLocalMap) {
 
     /**
      * Returns {@code true} if {@link FastThreadLocal#removeAll()} will be called once {@link #run()} completes.
+     *
+     * @deprecated Use {@link FastThreadLocalThread#currentThreadWillCleanupFastThreadLocals()} instead
      */
-    @UnstableApi
+    @Deprecated
     public boolean willCleanupFastThreadLocals() {
         return cleanupFastThreadLocals;
     }
 
     /**
      * Returns {@code true} if {@link FastThreadLocal#removeAll()} will be called once {@link Thread#run()} completes.
+     *
+     * @deprecated Use {@link FastThreadLocalThread#currentThreadWillCleanupFastThreadLocals()} instead
      */
-    @UnstableApi
+    @Deprecated
     public static boolean willCleanupFastThreadLocals(Thread thread) {
         return thread instanceof FastThreadLocalThread &&
                 ((FastThreadLocalThread) thread).willCleanupFastThreadLocals();
     }
 
+    /**
+     * Returns {@code true} if {@link FastThreadLocal#removeAll()} will be called once {@link Thread#run()} completes.
+     */
+    public static boolean currentThreadWillCleanupFastThreadLocals() {
+        // intentionally doesn't accept a thread parameter to work with ScopedValue in the future
+        Thread currentThread = currentThread();
+        if (currentThread instanceof FastThreadLocalThread) {
+            return ((FastThreadLocalThread) currentThread).willCleanupFastThreadLocals();
+        }
+        return isFastThreadLocalVirtualThread();
+    }
+
+    /**
+     * Returns {@code true} if this thread supports {@link FastThreadLocal}.
+     */
+    public static boolean currentThreadHasFastThreadLocal() {
+        // intentionally doesn't accept a thread parameter to work with ScopedValue in the future
+        return currentThread() instanceof FastThreadLocalThread || isFastThreadLocalVirtualThread();
+    }
+
+    private static boolean isFastThreadLocalVirtualThread() {
+        long[] arr = fallbackThreads.get();
+        if (arr == null) {
+            return false;
+        }
+        return Arrays.binarySearch(arr, Thread.currentThread().threadId()) >= 0;
+    }
+
+    /**
+     * Run the given task with {@link FastThreadLocal} support. This call should wrap the runnable for any thread that
+     * is long-running enough to make treating it as a {@link FastThreadLocalThread} reasonable, but that can't
+     * actually extend this class (e.g. because it's a virtual thread). Netty will use optimizations for recyclers and
+     * allocators as if this was a {@link FastThreadLocalThread}.
+     * <p>This method will clean up any {@link FastThreadLocal}s at the end, and
+     * {@link #currentThreadWillCleanupFastThreadLocals()} will return {@code true}.
+     * <p>At the moment, {@link FastThreadLocal} uses normal {@link ThreadLocal} as the backing storage here, but in
+     * the future this may be replaced with scoped values, if semantics can be preserved and performance is good.
+     *
+     * @param runnable The task to run
+     */
+    public static void runWithFastThreadLocal(Runnable runnable) {
+        Thread current = currentThread();
+        if (current instanceof FastThreadLocalThread) {
+            throw new IllegalStateException("Caller is a real FastThreadLocalThread");
+        }
+        long id = current.threadId();
+        fallbackThreads.updateAndGet(arr -> {
+            if (arr == null) {
+                return new long[] { id };
+            }
+            int index = Arrays.binarySearch(arr, id);
+            if (index >= 0) {
+                throw new IllegalStateException("Reentrant call to run()");
+            }
+            index = ~index; // same as -(index + 1)
+            long[] next = new long[arr.length + 1];
+            System.arraycopy(arr, 0, next, 0, index);
+            next[index] = id;
+            System.arraycopy(arr, index, next, index + 1, arr.length - index);
+            return next;
+        });
+        try {
+            runnable.run();
+        } finally {
+            fallbackThreads.getAndUpdate(arr -> {
+                if (arr == null || (arr.length == 1 && arr[0] == id)) {
+                    return null;
+                }
+                int index = Arrays.binarySearch(arr, id);
+                if (index < 0) {
+                    return arr;
+                }
+                long[] next = new long[arr.length - 1];
+                System.arraycopy(arr, 0, next, 0, index);
+                System.arraycopy(arr, index + 1, next, index, arr.length - index - 1);
+                return next;
+            });
+            FastThreadLocal.removeAll();
+        }
+    }
+
     /**
      * Query whether this thread is allowed to perform blocking calls or not.
      * {@link FastThreadLocalThread}s are often used in event-loops, where blocking calls are forbidden in order to
diff --git a/common/src/test/java/io/netty5/util/concurrent/FastThreadLocalTest.java b/common/src/test/java/io/netty5/util/concurrent/FastThreadLocalTest.java
index eb890a08fc9..a455b02693c 100644
--- a/common/src/test/java/io/netty5/util/concurrent/FastThreadLocalTest.java
+++ b/common/src/test/java/io/netty5/util/concurrent/FastThreadLocalTest.java
@@ -23,6 +23,12 @@
 import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
 
 import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.FutureTask;
+import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -34,6 +40,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.nullValue;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -161,33 +168,100 @@ public void testMultipleSetRemove_multipleThreadLocal() throws Exception {
         assertEquals(0, mapRef.get().size());
     }
 
+    @Test
+    public void testWrappedProperties() {
+        assertFalse(FastThreadLocalThread.currentThreadWillCleanupFastThreadLocals());
+        assertFalse(FastThreadLocalThread.currentThreadHasFastThreadLocal());
+        FastThreadLocalThread.runWithFastThreadLocal(() -> {
+            assertTrue(FastThreadLocalThread.currentThreadWillCleanupFastThreadLocals());
+            assertTrue(FastThreadLocalThread.currentThreadHasFastThreadLocal());
+        });
+    }
+
+    @Test
+    public void testWrapMany() throws ExecutionException, InterruptedException {
+        class Worker implements Runnable {
+            final Semaphore semaphore = new Semaphore(0);
+            final FutureTask<?> task = new FutureTask<>(this, null);
+
+            @Override
+            public void run() {
+                assertFalse(FastThreadLocalThread.currentThreadWillCleanupFastThreadLocals());
+                assertFalse(FastThreadLocalThread.currentThreadHasFastThreadLocal());
+                semaphore.acquireUninterruptibly();
+                FastThreadLocalThread.runWithFastThreadLocal(() -> {
+                    assertTrue(FastThreadLocalThread.currentThreadWillCleanupFastThreadLocals());
+                    assertTrue(FastThreadLocalThread.currentThreadHasFastThreadLocal());
+                    semaphore.acquireUninterruptibly();
+                    assertTrue(FastThreadLocalThread.currentThreadWillCleanupFastThreadLocals());
+                    assertTrue(FastThreadLocalThread.currentThreadHasFastThreadLocal());
+                });
+                assertFalse(FastThreadLocalThread.currentThreadWillCleanupFastThreadLocals());
+                assertFalse(FastThreadLocalThread.currentThreadHasFastThreadLocal());
+            }
+        }
+
+        int n = 100;
+        List<Worker> workers = new ArrayList<>();
+        for (int i = 0; i < n; i++) {
+            Worker worker = new Worker();
+            workers.add(worker);
+        }
+        Collections.shuffle(workers);
+        for (int i = 0; i < workers.size(); i++) {
+            new Thread(workers.get(i).task, "worker-" + i).start();
+        }
+        for (int i = 0; i < 2; i++) {
+            Collections.shuffle(workers);
+            for (Worker worker : workers) {
+                worker.semaphore.release();
+            }
+        }
+        for (Worker worker : workers) {
+            worker.task.get();
+        }
+    }
+
     @Test
     @Timeout(value = 4000, unit = TimeUnit.MILLISECONDS)
     public void testOnRemoveCalledForFastThreadLocalGet() throws Exception {
-        testOnRemoveCalled(true, true);
+        testOnRemoveCalled(true, false, true);
     }
 
     @Disabled("onRemoval(...) not called with non FastThreadLocal")
     @Test
     @Timeout(value = 4000, unit = TimeUnit.MILLISECONDS)
     public void testOnRemoveCalledForNonFastThreadLocalGet() throws Exception {
-        testOnRemoveCalled(false, true);
+        testOnRemoveCalled(false, false, true);
     }
 
     @Test
     @Timeout(value = 4000, unit = TimeUnit.MILLISECONDS)
     public void testOnRemoveCalledForFastThreadLocalSet() throws Exception {
-        testOnRemoveCalled(true, false);
+        testOnRemoveCalled(true, false, false);
     }
 
     @Disabled("onRemoval(...) not called with non FastThreadLocal")
     @Test
     @Timeout(value = 4000, unit = TimeUnit.MILLISECONDS)
     public void testOnRemoveCalledForNonFastThreadLocalSet() throws Exception {
-        testOnRemoveCalled(false, false);
+        testOnRemoveCalled(false, false, false);
     }
 
-    private static void testOnRemoveCalled(boolean fastThreadLocal, final boolean callGet) throws Exception {
+    @Test
+    @Timeout(value = 4000, unit = TimeUnit.MILLISECONDS)
+    public void testOnRemoveCalledForWrappedGet() throws Exception {
+        testOnRemoveCalled(false, true, true);
+    }
+
+    @Test
+    @Timeout(value = 4000, unit = TimeUnit.MILLISECONDS)
+    public void testOnRemoveCalledForWrappedSet() throws Exception {
+        testOnRemoveCalled(false, true, false);
+    }
+
+    private static void testOnRemoveCalled(boolean fastThreadLocal, boolean wrap, final boolean callGet)
+            throws Exception {
 
         final TestFastThreadLocal threadLocal = new TestFastThreadLocal();
         final TestFastThreadLocal threadLocal2 = new TestFastThreadLocal();
@@ -201,6 +275,10 @@ private static void testOnRemoveCalled(boolean fastThreadLocal, final boolean ca
                 threadLocal2.set(Thread.currentThread().getName());
             }
         };
+        if (wrap) {
+            Runnable r = runnable;
+            runnable = () -> FastThreadLocalThread.runWithFastThreadLocal(r);
+        }
         Thread thread = fastThreadLocal ? new FastThreadLocalThread(runnable) : new Thread(runnable);
         thread.start();
         thread.join();


Reviews:
user: login: chrisvest
body: Yeah, the adaptive allocator in `main` has fallen behind the 4.2 one.

Spotted one issue.
state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2025-05-22T16:25:23Z
user: login: yawkat
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T16:30:54Z
user: login: chrisvest
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-22T16:31:28Z
user: login: yawkat
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T16:31:33Z

Review comments:
diff_hunk: @@ -140,7 +140,7 @@ protected Object initialValue() {
                         return NO_MAGAZINE;
                     }
                     Magazine mag = new Magazine(AdaptivePoolingAllocator.this, executor);
-                    if (FastThreadLocalThread.willCleanupFastThreadLocals(Thread.currentThread())) {
+                    if (!FastThreadLocalThread.currentThreadWillCleanupFastThreadLocals()) {
path: buffer/src/main/java/io/netty5/buffer/adapt/AdaptivePoolingAllocator.java
user: login: chrisvest
body: This boolean should not be inverted.
created_at: 2025-05-22T16:23:24Z
author_association: MEMBER
diff_hunk: @@ -140,7 +140,7 @@ protected Object initialValue() {
                         return NO_MAGAZINE;
                     }
                     Magazine mag = new Magazine(AdaptivePoolingAllocator.this, executor);
-                    if (FastThreadLocalThread.willCleanupFastThreadLocals(Thread.currentThread())) {
+                    if (!FastThreadLocalThread.currentThreadWillCleanupFastThreadLocals()) {
path: buffer/src/main/java/io/netty5/buffer/adapt/AdaptivePoolingAllocator.java
user: login: yawkat
body: ```suggestion
                    if (FastThreadLocalThread.currentThreadWillCleanupFastThreadLocals()) {
```
created_at: 2025-05-22T16:30:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -140,7 +140,7 @@ protected Object initialValue() {
                         return NO_MAGAZINE;
                     }
                     Magazine mag = new Magazine(AdaptivePoolingAllocator.this, executor);
-                    if (FastThreadLocalThread.willCleanupFastThreadLocals(Thread.currentThread())) {
+                    if (!FastThreadLocalThread.currentThreadWillCleanupFastThreadLocals()) {
path: buffer/src/main/java/io/netty5/buffer/adapt/AdaptivePoolingAllocator.java
user: login: yawkat
body: yep, I think this check was flipped in 4.2, so my dumb copy failed :D
created_at: 2025-05-22T16:31:32Z
author_association: CONTRIBUTOR

PR comments:
user: login: yawkat
created_at: 2025-05-22T16:12:12Z
body: @normanmaurer I think AdaptivePoolingAllocator is missing some ports of the latest changes
author_association: CONTRIBUTOR
user: login: normanmaurer
created_at: 2025-05-22T16:13:13Z
body: @yawkat it is a different implementation mostly... @chrisvest can you check ?
author_association: MEMBER
user: login: normanmaurer
created_at: 2025-05-23T04:57:10Z
body: @yawkat thanks
author_association: MEMBER