Stop making hard references in HelperInjector
Also make awaitGC interruptable.
MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java
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);
}
}
}

MODIFY
dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/HelperInjectionTest.groovy
dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/HelperInjectionTest.groovy
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()
}
}

MODIFY
dd-java-agent/src/test/groovy/datadog/trace/agent/integration/classloading/ClassLoadingTest.groovy
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)]

MODIFY
dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/AgentTestRunner.java
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(

MODIFY
dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/utils/ClasspathUtils.java
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);
+ }
+ }
}

MODIFY
dd-java-agent/testing/src/test/java/muzzle/MuzzleWeakReferenceTest.java
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 =

MODIFY
utils/gc-utils/src/main/java/datadog/trace/util/gc/GCUtils.java
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();
}

Stop making hard references in HelperInjector

Also make awaitGC interruptable.