Stop making hard references in HelperInjector

Also make awaitGC interruptable.

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java

@@ -1,14 +1,16 @@
 package datadog.trace.agent.tooling;
 
 import static datadog.trace.agent.tooling.ClassLoaderMatcher.BOOTSTRAP_CLASSLOADER;
+import static datadog.trace.bootstrap.WeakMap.Provider.newWeakMap;
 
+import datadog.trace.bootstrap.WeakMap;
 import java.io.File;
 import java.io.IOException;
+import java.security.SecureClassLoader;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.Map;
@@ -24,9 +26,13 @@ import net.bytebuddy.utility.JavaModule;
 /** Injects instrumentation helper classes into the user's classloader. */
 @Slf4j
 public class HelperInjector implements Transformer {
+  // Need this because we can't put null into the injectedClassLoaders map.
+  private static final ClassLoader BOOTSTRAP_CLASSLOADER_PLACEHOLDER =
+      new SecureClassLoader(null) {};
+
   private final Set<String> helperClassNames;
   private Map<TypeDescription, byte[]> helperMap = null;
-  private final Set<ClassLoader> injectedClassLoaders = new HashSet<>();
+  private final WeakMap<ClassLoader, Boolean> injectedClassLoaders = newWeakMap();
 
   /**
    * Construct HelperInjector.
@@ -80,15 +86,18 @@ public class HelperInjector implements Transformer {
   public DynamicType.Builder<?> transform(
       final DynamicType.Builder<?> builder,
       final TypeDescription typeDescription,
-      final ClassLoader classLoader,
+      ClassLoader classLoader,
       final JavaModule module) {
     if (helperClassNames.size() > 0) {
       synchronized (this) {
-        if (!injectedClassLoaders.contains(classLoader)) {
+        if (classLoader == BOOTSTRAP_CLASSLOADER) {
+          classLoader = BOOTSTRAP_CLASSLOADER_PLACEHOLDER;
+        }
+        if (!injectedClassLoaders.containsKey(classLoader)) {
           try {
             final Map<TypeDescription, byte[]> helperMap = getHelperMap();
             log.debug("Injecting classes onto classloader {} -> {}", classLoader, helperClassNames);
-            if (classLoader == BOOTSTRAP_CLASSLOADER) {
+            if (classLoader == BOOTSTRAP_CLASSLOADER_PLACEHOLDER) {
               final Map<TypeDescription, Class<?>> injected =
                   ClassInjector.UsingInstrumentation.of(
                           new File(System.getProperty("java.io.tmpdir")),
@@ -108,13 +117,13 @@ public class HelperInjector implements Transformer {
                     + ". Failed to inject helper classes into instance "
                     + classLoader
                     + " of type "
-                    + (classLoader == BOOTSTRAP_CLASSLOADER
+                    + (classLoader == BOOTSTRAP_CLASSLOADER_PLACEHOLDER
                         ? "<bootstrap>"
                         : classLoader.getClass().getName()),
                 e);
             throw new RuntimeException(e);
           }
-          injectedClassLoaders.add(classLoader);
+          injectedClassLoaders.put(classLoader, true);
         }
       }
     }

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/HelperInjectionTest.groovy

@@ -1,38 +1,54 @@
 package datadog.trace.agent.test
 
+
 import datadog.trace.agent.tooling.AgentInstaller
 import datadog.trace.agent.tooling.HelperInjector
 import datadog.trace.agent.tooling.Utils
 import net.bytebuddy.agent.ByteBuddyAgent
+import net.bytebuddy.description.type.TypeDescription
+import net.bytebuddy.dynamic.ClassFileLocator
+import net.bytebuddy.dynamic.loading.ClassInjector
 import spock.lang.Specification
+import spock.lang.Timeout
 
-import java.lang.reflect.Method
+import java.lang.ref.WeakReference
+import java.util.concurrent.atomic.AtomicReference
 
+import static datadog.trace.agent.test.utils.ClasspathUtils.isClassLoaded
 import static datadog.trace.agent.tooling.ClassLoaderMatcher.BOOTSTRAP_CLASSLOADER
+import static datadog.trace.util.gc.GCUtils.awaitGC
 
 class HelperInjectionTest extends Specification {
 
+  @Timeout(10)
   def "helpers injected to non-delegating classloader"() {
     setup:
     String helperClassName = HelperInjectionTest.getPackage().getName() + '.HelperClass'
     HelperInjector injector = new HelperInjector(helperClassName)
-    URLClassLoader emptyLoader = new URLClassLoader(new URL[0], (ClassLoader) null)
+    AtomicReference<URLClassLoader> emptyLoader = new AtomicReference<>(new URLClassLoader(new URL[0], (ClassLoader) null))
 
     when:
-    emptyLoader.loadClass(helperClassName)
+    emptyLoader.get().loadClass(helperClassName)
     then:
     thrown ClassNotFoundException
 
     when:
-    injector.transform(null, null, emptyLoader, null)
-    emptyLoader.loadClass(helperClassName)
+    injector.transform(null, null, emptyLoader.get(), null)
+    emptyLoader.get().loadClass(helperClassName)
     then:
-    isClassLoaded(helperClassName, emptyLoader)
+    isClassLoaded(helperClassName, emptyLoader.get())
     // injecting into emptyLoader should not load on agent's classloader
     !isClassLoaded(helperClassName, Utils.getAgentClassLoader())
 
-    cleanup:
-    emptyLoader?.close()
+    when: "references to emptyLoader are gone"
+    emptyLoader.get().close() // cleanup
+    def ref = new WeakReference(emptyLoader.get())
+    emptyLoader.set(null)
+
+    awaitGC(ref)
+
+    then: "HelperInjector doesn't prevent it from being collected"
+    null == ref.get()
   }
 
   def "helpers injected on bootstrap classloader"() {
@@ -55,16 +71,38 @@ class HelperInjectionTest extends Specification {
     helperClass.getClassLoader() == BOOTSTRAP_CLASSLOADER
   }
 
-  private static boolean isClassLoaded(String className, ClassLoader classLoader) {
-    final Method findLoadedClassMethod = ClassLoader.getDeclaredMethod("findLoadedClass", String)
-    try {
-      findLoadedClassMethod.setAccessible(true)
-      Class<?> loadedClass = (Class<?>) findLoadedClassMethod.invoke(classLoader, className)
-      return null != loadedClass && loadedClass.getClassLoader() == classLoader
-    } catch (Exception e) {
-      throw new IllegalStateException(e)
-    } finally {
-      findLoadedClassMethod.setAccessible(false)
-    }
+  def "check hard references on class injection"() {
+    setup:
+    String helperClassName = HelperInjectionTest.getPackage().getName() + '.HelperClass'
+
+    // Copied from HelperInjector:
+    final ClassFileLocator locator =
+      ClassFileLocator.ForClassLoader.of(Utils.getAgentClassLoader())
+    final byte[] classBytes = locator.locate(helperClassName).resolve()
+    final TypeDescription typeDesc =
+      new TypeDescription.Latent(
+        helperClassName, 0, null, Collections.<TypeDescription.Generic> emptyList())
+
+    AtomicReference<URLClassLoader> emptyLoader = new AtomicReference<>(new URLClassLoader(new URL[0], (ClassLoader) null))
+    AtomicReference<ClassInjector> injector = new AtomicReference<>(new ClassInjector.UsingReflection(emptyLoader.get()))
+    injector.get().inject([(typeDesc): classBytes])
+
+    when:
+    def injectorRef = new WeakReference(injector.get())
+    injector.set(null)
+
+    awaitGC(injectorRef)
+
+    then:
+    null == injectorRef.get()
+
+    when:
+    def loaderRef = new WeakReference(emptyLoader.get())
+    emptyLoader.set(null)
+
+    awaitGC(loaderRef)
+
+    then:
+    null == loaderRef.get()
   }
 }

dd-java-agent/src/test/groovy/datadog/trace/agent/integration/classloading/ClassLoadingTest.groovy

@@ -6,11 +6,13 @@ import datadog.trace.agent.test.IntegrationTestUtils
 import datadog.trace.api.Trace
 import datadog.trace.util.gc.GCUtils
 import spock.lang.Specification
+import spock.lang.Timeout
 
 import java.lang.ref.WeakReference
 
 import static datadog.trace.agent.test.IntegrationTestUtils.createJarWithClasses
 
+@Timeout(10)
 class ClassLoadingTest extends Specification {
 
   final URL[] classpath = [createJarWithClasses(ClassToInstrument, ClassToInstrumentChild, Trace)]

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/AgentTestRunner.java

@@ -178,13 +178,13 @@ public abstract class AgentTestRunner extends Specification {
 
   @AfterClass
   public static synchronized void agentCleanup() {
-    assert INSTRUMENTATION_ERROR_COUNT.get() == 0
-        : INSTRUMENTATION_ERROR_COUNT.get() + " Instrumentation errors during test";
-
     if (null != activeTransformer) {
       INSTRUMENTATION.removeTransformer(activeTransformer);
       activeTransformer = null;
     }
+    // Cleanup before assertion.
+    assert INSTRUMENTATION_ERROR_COUNT.get() == 0
+        : INSTRUMENTATION_ERROR_COUNT.get() + " Instrumentation errors during test";
   }
 
   public static void assertTraces(

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/utils/ClasspathUtils.java

@@ -13,6 +13,7 @@ import java.io.File;
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.lang.reflect.Method;
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.net.URLClassLoader;
@@ -156,4 +157,24 @@ public class ClasspathUtils {
     }
     return new URLClassLoader(urls.build().toArray(new URL[0]), null);
   }
+
+  // Moved this to a java class because groovy was adding a hard ref to classLoader
+  public static boolean isClassLoaded(final String className, final ClassLoader classLoader) {
+    try {
+      final Method findLoadedClassMethod =
+          ClassLoader.class.getDeclaredMethod("findLoadedClass", String.class);
+      try {
+        findLoadedClassMethod.setAccessible(true);
+        final Class<?> loadedClass =
+            (Class<?>) findLoadedClassMethod.invoke(classLoader, className);
+        return null != loadedClass && loadedClass.getClassLoader() == classLoader;
+      } catch (final Exception e) {
+        throw new IllegalStateException(e);
+      } finally {
+        findLoadedClassMethod.setAccessible(false);
+      }
+    } catch (final NoSuchMethodException e) {
+      throw new IllegalStateException(e);
+    }
+  }
 }

dd-java-agent/testing/src/test/java/muzzle/MuzzleWeakReferenceTest.java

@@ -14,7 +14,7 @@ public class MuzzleWeakReferenceTest {
    *
    * Even returning a WeakReference<ClassLoader> is enough for spock to create a strong ref.
    */
-  public static boolean classLoaderRefIsGarbageCollected() {
+  public static boolean classLoaderRefIsGarbageCollected() throws InterruptedException {
     ClassLoader loader = new URLClassLoader(new URL[0], null);
     final WeakReference<ClassLoader> clRef = new WeakReference<>(loader);
     final Reference[] refs =

utils/gc-utils/src/main/java/datadog/trace/util/gc/GCUtils.java

@@ -4,15 +4,18 @@ import java.lang.ref.WeakReference;
 
 public abstract class GCUtils {
 
-  public static void awaitGC() {
+  public static void awaitGC() throws InterruptedException {
     Object obj = new Object();
     final WeakReference<Object> ref = new WeakReference<>(obj);
     obj = null;
     awaitGC(ref);
   }
 
-  public static void awaitGC(final WeakReference<?> ref) {
+  public static void awaitGC(final WeakReference<?> ref) throws InterruptedException {
     while (ref.get() != null) {
+      if (Thread.interrupted()) {
+        throw new InterruptedException();
+      }
       System.gc();
       System.runFinalization();
     }