Replace ThreadLocal with a hashtable to support reliable cleanup (#588)

Co-authored-by: Scott M Stark <[email protected]>

core/impl-base/src/main/java/org/jboss/arquillian/core/impl/ManagerImpl.java

@@ -31,6 +31,7 @@
 import org.jboss.arquillian.core.api.threading.ExecutorService;
 import org.jboss.arquillian.core.impl.context.ApplicationContextImpl;
 import org.jboss.arquillian.core.impl.threading.ThreadedExecutorService;
+import org.jboss.arquillian.core.spi.ArquillianThreadLocal;
 import org.jboss.arquillian.core.spi.EventContext;
 import org.jboss.arquillian.core.spi.EventPoint;
 import org.jboss.arquillian.core.spi.Extension;
@@ -66,8 +67,8 @@ public class ManagerImpl implements Manager {
      * the higher level event will get the exception and re-fire it on the bus. We need to keep track of which exceptions
      * has been handled in the call chain so we can re-throw without re-firing on a higher level.
      */
-    private ThreadLocal<Set<Class<? extends Throwable>>> handledThrowables =
-        new ThreadLocal<Set<Class<? extends Throwable>>>() {
+    private ArquillianThreadLocal<Set<Class<? extends Throwable>>> handledThrowables =
+        new ArquillianThreadLocal<Set<Class<? extends Throwable>>>() {
             @Override
             protected Set<Class<? extends Throwable>> initialValue() {
                 return new HashSet<Class<? extends Throwable>>();
@@ -279,6 +280,8 @@ public void shutdown() {
             runtimeLogger.clear();
 
             handledThrowables.remove();
+            //Force cleanup:
+            handledThrowables.clear();
         }
         if (shutdownException != null) {
             UncheckedThrow.throwUnchecked(shutdownException);

core/spi/src/main/java/org/jboss/arquillian/core/spi/ArquillianThreadLocal.java

@@ -0,0 +1,59 @@
+package org.jboss.arquillian.core.spi;
+
+import java.util.Hashtable;
+
+/**
+ * Mapping for ThreadId to a value. Same as "ThreadLocal", but with simpler cleanup.
+ *
+ */
+public class ArquillianThreadLocal<T> {
+  private Hashtable<Long, T> table = new Hashtable<Long, T>();
+
+  protected T initialValue() {
+    return null;
+  }
+
+  /**
+   * Returns the value in the current thread's copy of this
+   * thread-local variable.  If the variable has no value for the
+   * current thread, it is first initialized to the value returned
+   * by an invocation of the {@link #initialValue} method.
+   *
+   * @return the current thread's value of this thread-local
+   */
+  public T get() {
+    Thread t = Thread.currentThread();
+    long threadId = t.getId();
+
+    if (table.containsKey(threadId)) {
+        return table.get(threadId);
+    }
+    else {
+      T value = initialValue();
+      table.put(threadId, value);
+      return value;
+    }
+  }
+
+  /**
+   * Removes the current thread's value for this thread-local
+   * variable.
+   *
+   */
+   public void remove() {
+     Thread t = Thread.currentThread();
+     long threadId = t.getId();
+
+     if (table.containsKey(threadId)) {
+       table.remove(threadId);
+     }
+   }
+
+   /**
+    * Clears the cache
+    */
+   public void clear() {
+     table.clear();
+   }
+}
+

core/spi/src/main/java/org/jboss/arquillian/core/spi/context/AbstractContext.java

@@ -20,6 +20,8 @@
 import java.util.Stack;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.logging.Logger;
+
+import org.jboss.arquillian.core.spi.ArquillianThreadLocal;
 import org.jboss.arquillian.core.spi.Validate;
 
 /**
@@ -33,7 +35,7 @@ public abstract class AbstractContext<T> implements Context, IdBoundContext<T> {
 
     private ConcurrentHashMap<T, ObjectStore> stores;
 
-    private ThreadLocal<Stack<StoreHolder<T>>> activeStore = new ThreadLocal<Stack<StoreHolder<T>>>() {
+    private ArquillianThreadLocal<Stack<StoreHolder<T>>> activeStore = new ArquillianThreadLocal<Stack<StoreHolder<T>>>() {
 
         @Override
         protected Stack<StoreHolder<T>> initialValue() {
@@ -105,6 +107,8 @@ public void clearAll() {
                 deactivateAll();
             }
             activeStore.remove();
+            //Force cleanup:
+            activeStore.clear();
             for (Map.Entry<T, ObjectStore> entry : stores.entrySet()) {
                 entry.getValue().clear();
             }