Reduce contention in DefaultByteBufferPool thread-local lookups

Author: matthew-bayer

benchmarks/src/main/java/io/undertow/benchmarks/DefaultByteBufferPoolBenchmark.java

@@ -0,0 +1,100 @@
+/*
+ * JBoss, Home of Professional Open Source.
+ * Copyright 2026 Red Hat, Inc., and individual contributors
+ * as indicated by the @author tags.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.undertow.benchmarks;
+
+import io.undertow.connector.PooledByteBuffer;
+import io.undertow.server.DefaultByteBufferPool;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Warmup;
+
+import java.util.concurrent.TimeUnit;
+
+@State(Scope.Benchmark)
+@Warmup(iterations = 3, time = 3)
+@Measurement(iterations = 5, time = 3)
+@Fork(3)
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(TimeUnit.SECONDS)
+public class DefaultByteBufferPoolBenchmark {
+
+    @Param({"false", "true"})
+    private boolean direct;
+
+    @Param({"0", "6", "12"})
+    private int cacheSize;
+
+    private DefaultByteBufferPool pool;
+
+    @Setup
+    public void setup() {
+        pool = new DefaultByteBufferPool(direct, 16 * 1024 - 20, 1024, cacheSize);
+    }
+
+    @TearDown
+    public void tearDown() {
+        pool.close();
+        pool.getArrayBackedPool().close();
+    }
+
+    @Benchmark
+    public int allocateAndRelease(Buffers state) {
+        int sum = 0;
+        try {
+            for (int i = 0; i < state.buffers.length; i++) {
+                PooledByteBuffer buffer = pool.allocate();
+                state.buffers[i] = buffer;
+                buffer.getBuffer().putInt(i);
+            }
+            for (PooledByteBuffer buffer : state.buffers) {
+                sum += buffer.getBuffer().getInt(0);
+            }
+            return sum;
+        } finally {
+            for (int i = 0; i < state.buffers.length; i++) {
+                if (state.buffers[i] != null) {
+                    state.buffers[i].close();
+                    state.buffers[i] = null;
+                }
+            }
+        }
+    }
+
+    @State(Scope.Thread)
+    public static class Buffers {
+
+        @Param({"1", "16"})
+        private int batchSize;
+
+        private PooledByteBuffer[] buffers;
+
+        @Setup
+        public void setup() {
+            buffers = new PooledByteBuffer[batchSize];
+        }
+    }
+}

core/src/main/java/io/undertow/server/DefaultByteBufferPool.java

@@ -42,7 +42,7 @@
 // TODO: move this somewhere more appropriate
 public class DefaultByteBufferPool implements ByteBufferPool {
 
-    private final ThreadLocalCache threadLocalCache = new ThreadLocalCache();
+    private final ThreadLocalCache threadLocalCache;
     // Access requires synchronization on the threadLocalDataList instance
     private final List<WeakReference<ThreadLocalData>> threadLocalDataList = new ArrayList<>();
     private final ConcurrentLinkedQueue<ByteBuffer> queue = new ConcurrentLinkedQueue<>();
@@ -85,6 +85,7 @@ public DefaultByteBufferPool(boolean direct, int bufferSize, int maximumPoolSize
         this.bufferSize = bufferSize;
         this.maximumPoolSize = maximumPoolSize;
         this.threadLocalCacheSize = threadLocalCacheSize;
+        this.threadLocalCache = new ThreadLocalCache(threadLocalCacheSize > 0);
         this.leakDectionPercent = leakDecetionPercent;
         if(direct) {
             arrayBackedPool = new DefaultByteBufferPool(false, bufferSize, maximumPoolSize, 0, leakDecetionPercent);
@@ -301,6 +302,7 @@ public String toString() {
     private class ThreadLocalData {
         final ArrayDeque<ByteBuffer> buffers = new ArrayDeque<>(threadLocalCacheSize);
         int allocationDepth = 0;
+        int cacheAccesses;
 
         @Override
         protected void finalize() throws Throwable {
@@ -340,18 +342,34 @@ protected void finalize() throws Throwable {
         }
     }
 
-    // This is used instead of Java ThreadLocal class. Unlike in the ThreadLocal class, the remove() method in this
-    // class can be called by a different thread than the one that initialized the data.
+    // The registry permits cleanup from another thread.
     private static class ThreadLocalCache {
 
+        private final ThreadLocal<WeakReference<ThreadLocalData>> localReference;
+
         final Map<Thread, ThreadLocalData> localsByThread = Collections.synchronizedMap(new WeakHashMap<>());
 
+        ThreadLocalCache(boolean cachingEnabled) {
+            localReference = cachingEnabled ? new ThreadLocal<>() : null;
+        }
+
         ThreadLocalData get() {
-            return localsByThread.get(Thread.currentThread());
+            WeakReference<ThreadLocalData> reference = localReference == null ? null : localReference.get();
+            ThreadLocalData local = reference == null ? null : reference.get();
+            if (local == null) {
+                return localsByThread.get(Thread.currentThread());
+            }
+            // WeakHashMap only releases stale values when accessed; keep retiring dead threads' caches.
+            if ((++local.cacheAccesses & 255) == 0) {
+                localsByThread.size();
+            }
+            return local;
         }
 
         void set(ThreadLocalData threadLocalData) {
             localsByThread.put(Thread.currentThread(), threadLocalData);
+            // A long-lived thread must not keep this pool alive through ThreadLocalData.
+            localReference.set(new WeakReference<>(threadLocalData));
         }
 
         void remove(ThreadLocalData threadLocalData) {