From e004b4a354402707953e22a47912629cbd0538e4 Mon Sep 17 00:00:00 2001
From: Norman Maurer <norman_maurer@apple.com>
Date: Thu, 21 Dec 2017 20:35:32 +0100
Subject: [PATCH] Ensure ObjectCleaner will also be used when
 FastThreadLocal.set is used.

Motivation:

e329ca1 introduced the user of ObjectCleaner in FastThreadLocal but we missed the case to register our cleaner task if FastThreadLocal.set was called only.

Modifications:

- Use ObjectCleaner also when FastThreadLocal.set is used.
- Add test case.

Result:

ObjectCleaner is always used.
---
 .../util/concurrent/FastThreadLocal.java      | 16 ++++++++--
 .../util/concurrent/FastThreadLocalTest.java  | 29 ++++++++++++++-----
 2 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/common/src/main/java/io/netty/util/concurrent/FastThreadLocal.java b/common/src/main/java/io/netty/util/concurrent/FastThreadLocal.java
index 2881a356f9..aa5ab85e23 100644
--- a/common/src/main/java/io/netty/util/concurrent/FastThreadLocal.java
+++ b/common/src/main/java/io/netty/util/concurrent/FastThreadLocal.java
@@ -134,13 +134,18 @@ public class FastThreadLocal<V> {
      */
     @SuppressWarnings("unchecked")
     public final V get() {
-        final InternalThreadLocalMap threadLocalMap = InternalThreadLocalMap.get();
+        InternalThreadLocalMap threadLocalMap = InternalThreadLocalMap.get();
         Object v = threadLocalMap.indexedVariable(index);
         if (v != InternalThreadLocalMap.UNSET) {
             return (V) v;
         }
 
         V value = initialize(threadLocalMap);
+        registerCleaner(threadLocalMap);
+        return value;
+    }
+
+    private void registerCleaner(final InternalThreadLocalMap threadLocalMap) {
         Thread current = Thread.currentThread();
         if (!FastThreadLocalThread.willCleanupFastThreadLocals(current)) {
             // We will need to ensure we will trigger remove(InternalThreadLocalMap) so everything will be released
@@ -155,7 +160,6 @@ public class FastThreadLocal<V> {
                 }
             });
         }
-        return value;
     }
 
     /**
@@ -190,7 +194,13 @@ public class FastThreadLocal<V> {
      */
     public final void set(V value) {
         if (value != InternalThreadLocalMap.UNSET) {
-            set(InternalThreadLocalMap.get(), value);
+            InternalThreadLocalMap threadLocalMap = InternalThreadLocalMap.get();
+            boolean alreadySet = threadLocalMap.isIndexedVariableSet(index);
+            set(threadLocalMap, value);
+
+            if (!alreadySet) {
+                registerCleaner(threadLocalMap);
+            }
         } else {
             remove();
         }
diff --git a/common/src/test/java/io/netty/util/concurrent/FastThreadLocalTest.java b/common/src/test/java/io/netty/util/concurrent/FastThreadLocalTest.java
index b7ac6bdc34..d74d70c5b9 100644
--- a/common/src/test/java/io/netty/util/concurrent/FastThreadLocalTest.java
+++ b/common/src/test/java/io/netty/util/concurrent/FastThreadLocalTest.java
@@ -77,16 +77,26 @@ public class FastThreadLocalTest {
     }
 
     @Test(timeout = 4000)
-    public void testOnRemoveCalledForFastThreadLocal() throws Exception {
-        testOnRemoveCalled(true);
+    public void testOnRemoveCalledForFastThreadLocalGet() throws Exception {
+        testOnRemoveCalled(true, true);
     }
 
     @Test(timeout = 4000)
-    public void testOnRemoveCalledForNonFastThreadLocal() throws Exception {
-        testOnRemoveCalled(false);
+    public void testOnRemoveCalledForNonFastThreadLocalGet() throws Exception {
+        testOnRemoveCalled(false, true);
     }
 
-    private static void testOnRemoveCalled(boolean fastThreadLocal) throws Exception {
+    @Test(timeout = 4000)
+    public void testOnRemoveCalledForFastThreadLocalSet() throws Exception {
+        testOnRemoveCalled(true, false);
+    }
+
+    @Test(timeout = 4000)
+    public void testOnRemoveCalledForNonFastThreadLocalSet() throws Exception {
+        testOnRemoveCalled(false, false);
+    }
+
+    private static void testOnRemoveCalled(boolean fastThreadLocal, final boolean callGet) throws Exception {
 
         final TestFastThreadLocal threadLocal = new TestFastThreadLocal();
         final TestFastThreadLocal threadLocal2 = new TestFastThreadLocal();
@@ -94,8 +104,13 @@ public class FastThreadLocalTest {
         Runnable runnable = new Runnable() {
             @Override
             public void run() {
-                assertEquals(Thread.currentThread().getName(), threadLocal.get());
-                assertEquals(Thread.currentThread().getName(), threadLocal2.get());
+                if (callGet) {
+                    assertEquals(Thread.currentThread().getName(), threadLocal.get());
+                    assertEquals(Thread.currentThread().getName(), threadLocal2.get());
+                } else {
+                    threadLocal.set(Thread.currentThread().getName());
+                    threadLocal2.set(Thread.currentThread().getName());
+                }
             }
         };
         Thread thread = fastThreadLocal ? new FastThreadLocalThread(runnable) : new Thread(runnable);