Fix memory leak when ThreadLocal are used inside tested code (#79)

Add support for an annotation that does extra garbage collection after each test to try and avoid memory leaks.

Author: LudoP

gwtmockito/src/main/java/com/google/gwtmockito/GwtMockitoTestRunner.java

@@ -376,6 +376,9 @@ public void testStarted(Description description) throws Exception {
       super.run(wrapperNotifier);
     } finally {
       Thread.currentThread().setContextClassLoader(originalClassLoader);
+      if (unitTestClass.isAnnotationPresent(WithExperimentalGarbageCollection.class)) {
+        ThreadLocalCleaner.cleanUpThreadLocalValues(gwtMockitoClassLoader);
+      }
     }
   }
 

gwtmockito/src/main/java/com/google/gwtmockito/ThreadLocalCleaner.java

@@ -0,0 +1,57 @@
+package com.google.gwtmockito;
+
+import java.lang.ref.WeakReference;
+import java.lang.reflect.Field;
+
+/**
+ * Utility classes to modify the private ThreadLocalMap attribute from the java.lang.Thread
+ */
+final class ThreadLocalCleaner {
+
+    private ThreadLocalCleaner() {
+        // Nothing
+    }
+
+    /**
+     * Remove ThreadLocalMap entries of objects loaded through the given classLoader as they prevent garbage collection
+     * of this classLoader while the thread that holds the ThreadLocalMap is alive
+     */
+    public static void cleanUpThreadLocalValues(ClassLoader classLoader) {
+        try {
+            Object threadLocalMap = getPrivateAttribute(Thread.class, "threadLocals", Thread.currentThread());
+            WeakReference[] table = (WeakReference[]) getPrivateAttribute("table", threadLocalMap);
+            int length = table.length;
+            for (int i = 0; i < length; i++) {
+                WeakReference mapEntry = table[i];
+                if (mapEntry != null && mapEntry.get() != null) {
+                    ClassLoader mapEntryKeyClassLoader = mapEntry.get().getClass().getClassLoader();
+                    Field mapEntryValueField = getPrivateAttributeAccessibleField(mapEntry.getClass(), "value");
+                    ClassLoader mapEntryValueClassLoader = mapEntryValueField.get(mapEntry).getClass().getClassLoader();
+                    if (mapEntryKeyClassLoader == classLoader || mapEntryValueClassLoader == classLoader) {
+                        mapEntry.clear();
+                        mapEntryValueField.set(mapEntry, null);
+                        // The ThreadLocalMap is able to expunge the remaining stale entries, no need to remove it from the map
+                    }
+                }
+            }
+        } catch (IllegalAccessException | NoSuchFieldException e) {
+            throw new AssertionError("Unable to access expected class fields for cleaning ThreadLocal values", e);
+        }
+    }
+
+    private static Object getPrivateAttribute(String attributeName, Object holder) throws NoSuchFieldException, IllegalAccessException {
+        return getPrivateAttribute(holder.getClass(), attributeName, holder);
+    }
+
+    private static Object getPrivateAttribute(Class attributeClass, String attributeName, Object holder) throws NoSuchFieldException, IllegalAccessException {
+        Field field = getPrivateAttributeAccessibleField(attributeClass, attributeName);
+        return field.get(holder);
+    }
+
+    private static Field getPrivateAttributeAccessibleField(Class attributeClass, String attributeName) throws NoSuchFieldException {
+        Field field = attributeClass.getDeclaredField(attributeName);
+        field.setAccessible(true);
+        return field;
+    }
+
+}

gwtmockito/src/main/java/com/google/gwtmockito/WithClassesToStub.java

@@ -21,7 +21,7 @@
 import java.lang.annotation.Target;
 
 /**
- * Annotation allowing the the test to configure the set of classes that will be stubbed completely
+ * Annotation allowing the test to configure the set of classes that will be stubbed completely
  * when the test is executed. If a class is stubbed, all of its non-abstract methods (including
  * native and final methods) will be replaced with a no-op implementation that returns a dummy value
  * if needed. This stubbing takes place at the classloader level, so all instances of the class

gwtmockito/src/main/java/com/google/gwtmockito/WithExperimentalGarbageCollection.java

@@ -0,0 +1,23 @@
+package com.google.gwtmockito;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Annotation allowing the test to enable a clean up in the ThreadLocalMap at the end of the test
+ * in order to avoid a memory leak.
+ *
+ * <p> Without this annotation, if a test uses a ThreadLocal object load through GwtMockitoClassLoader
+ * and stored in the test runner thread context, the garbage collection of the GwtMockitoClassLoader
+ * will be prevented.
+ *
+ * <p> Note: the memory leak issue and this experimental garbage collection fix is dependent of the JVM
+ * (observed with the OpenJDK). It will raise an AssertionError on uncompatible JVM.
+ */
+@Target(ElementType.TYPE)
+@Retention(RetentionPolicy.RUNTIME)
+public @interface WithExperimentalGarbageCollection {
+
+}

gwtmockito/src/main/java/com/google/gwtmockito/WithPackagesToLoadViaStandardClassLoader.java

@@ -21,7 +21,7 @@
 import java.lang.annotation.Target;
 
 /**
- * Annotation allowing the the test to configure a list of package names that should always be
+ * Annotation allowing the test to configure a list of package names that should always be
  * loaded via the standard system classloader instead of through GwtMockito's custom classloader.
  * Any subpackages of these packages will also be loaded with the standard loader. If you're
  * getting "loader constraint violation" errors, try defining adding this annotation to your test

gwtmockito/src/test/java/com/google/gwtmockito/GwtMockitoMemoryLeakTest.java

@@ -0,0 +1,48 @@
+package com.google.gwtmockito;
+
+import com.google.gwtmockito.subpackage.ThreadLocalUsage;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runner.notification.RunNotifier;
+import org.junit.runners.JUnit4;
+import org.junit.runners.model.InitializationError;
+
+import java.lang.ref.WeakReference;
+
+import static junit.framework.TestCase.assertNull;
+
+public class GwtMockitoMemoryLeakTest {
+
+    @Test
+    public void shouldGarbageCollectGwtClassLoaderWhenThreadLocalIsUsed() throws InitializationError {
+        GwtMockitoTestRunner runner = new GwtMockitoTestRunner(TestWithThreadLocal.class);
+        runner.run(new RunNotifier());
+        WeakReference<GwtMockitoTestRunner> wkRunner = new WeakReference<>(runner);
+        // remove the reference to the runner to allow the garbage collector to collect it.
+        runner = null;
+        // Call the Garbage Collector to (try to) force the collection of the runner.
+        // This test is not perfect and may depend of the JVM Garbage Collector implementation as
+        // as return from "System.gc()" simply garantee that a "best effort to reclaim space from all discarded objects"
+        // has been done.
+        System.gc();
+        // Check if the garbage colllector has collected the runner instance through the WeakReference
+        assertNull("Expected to garbage collect the GwtMockitoTestRunner", wkRunner.get());
+    }
+
+    @RunWith(JUnit4.class)
+    public static class DummyTestClass {
+
+        @Test
+        public void dummy() {
+        }
+    }
+
+    @RunWith(JUnit4.class)
+    @WithExperimentalGarbageCollection
+    public static class TestWithThreadLocal {
+        @Test
+        public void dummy() {
+            ThreadLocalUsage threadLocalUsage = new ThreadLocalUsage();
+        }
+    }
+}

gwtmockito/src/test/java/com/google/gwtmockito/subpackage/ThreadLocalUsage.java

@@ -0,0 +1,5 @@
+package com.google.gwtmockito.subpackage;
+
+public class ThreadLocalUsage {
+    private static ThreadLocal<ThreadLocalUsage> MY_DUMMY_LOCAL = ThreadLocal.withInitial(() -> new ThreadLocalUsage());
+}