From 205bf2db165152d0833f3704aca5c0047cf0a5ed Mon Sep 17 00:00:00 2001
From: Tomas Hofman <thofman@redhat.com>
Date: Wed, 27 Sep 2023 12:45:31 +0200
Subject: [PATCH] UNDERTOW-2309 Prevent memory leak in DefaultByteBufferPool

---
 .../server/DefaultByteBufferPool.java         | 30 ++++++++++++++++++-
 1 file changed, 29 insertions(+), 1 deletion(-)

diff --git a/core/src/main/java/io/undertow/server/DefaultByteBufferPool.java b/core/src/main/java/io/undertow/server/DefaultByteBufferPool.java
index f74a381bf6..519dd67d2d 100644
--- a/core/src/main/java/io/undertow/server/DefaultByteBufferPool.java
+++ b/core/src/main/java/io/undertow/server/DefaultByteBufferPool.java
@@ -26,7 +26,9 @@
 import java.nio.ByteBuffer;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
 
@@ -39,7 +41,7 @@
  */
 public class DefaultByteBufferPool implements ByteBufferPool {
 
-    private final ThreadLocal<ThreadLocalData> threadLocalCache = new ThreadLocal<>();
+    private final ThreadLocalCache threadLocalCache = new ThreadLocalCache();
     // Access requires synchronization on the threadLocalDataList instance
     private final List<WeakReference<ThreadLocalData>> threadLocalDataList = new ArrayList<>();
     private final ConcurrentLinkedQueue<ByteBuffer> queue = new ConcurrentLinkedQueue<>();
@@ -228,6 +230,7 @@ public void close() {
                     local.buffers.clear();
                 }
                 ref.clear();
+                threadLocalCache.remove(local);
             }
             threadLocalDataList.clear();
         }
@@ -332,4 +335,29 @@ 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.
+    private static class ThreadLocalCache {
+
+        Map<Thread, ThreadLocalData> localsByThread = new HashMap<>();
+
+        ThreadLocalData get() {
+            return localsByThread.get(Thread.currentThread());
+        }
+
+        void set(ThreadLocalData threadLocalData) {
+            localsByThread.put(Thread.currentThread(), threadLocalData);
+        }
+
+        void remove(ThreadLocalData threadLocalData) {
+            // Find the entry containing given data instance and remove it from the map.
+            for (Map.Entry<Thread, ThreadLocalData> entry: localsByThread.entrySet()) {
+                if (threadLocalData.equals(entry.getValue())) {
+                    localsByThread.remove(entry.getKey(), entry.getValue());
+                    break;
+                }
+            }
+        }
+    }
+
 }