Helper Injector Transformer
MODIFY
dd-java-agent/tooling/src/main/java/dd/trace/DDAdvice.java
dd-java-agent/tooling/src/main/java/dd/trace/DDAdvice.java
@@ -18,12 +18,21 @@ public class DDAdvice extends AgentBuilder.Transformer.ForAdvice {
Method getAgentClassloaderMethod = agentClass.getMethod("getAgentClassLoader");
AGENT_CLASSLOADER = (ClassLoader) getAgentClassloaderMethod.invoke(null);
} catch (Throwable t) {
- log.error("Unable to locate agent classloader. Falling back to System Classloader");
- AGENT_CLASSLOADER = ClassLoader.getSystemClassLoader();
+ AGENT_CLASSLOADER = DDAdvice.class.getClassLoader();
+ log.error(
+ "Failed to locate agent classloader. Falling back to "
+ + AGENT_CLASSLOADER
+ + " -- "
+ + t.getMessage());
}
AGENT_CLASS_LOCATOR = ClassFileLocator.ForClassLoader.of(AGENT_CLASSLOADER);
}
+ /** Return the classloader the datadog agent is running on. */
+ public static ClassLoader getAgentClassLoader() {
+ return AGENT_CLASSLOADER;
+ }
+
/**
* Create bytebuddy advice with default datadog settings.
*

ADD
dd-java-agent/tooling/src/main/java/dd/trace/HelperInjector.java
+package dd.trace;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import lombok.extern.slf4j.Slf4j;
+import net.bytebuddy.agent.builder.AgentBuilder.Transformer;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.dynamic.ClassFileLocator;
+import net.bytebuddy.dynamic.DynamicType;
+import net.bytebuddy.dynamic.loading.ClassInjector;
+import net.bytebuddy.utility.JavaModule;
+
+/** Injects instrumentation helper classes into the user's classloader. */
+@Slf4j
+public class HelperInjector implements Transformer {
+ private final Set<String> helperClassNames;
+ private Map<TypeDescription, byte[]> helperMap = null;
+
+ /**
+ * Construct HelperInjector.
+ *
+ * @param helperClassNames binary names of the helper classes to inject. These class names must be
+ * resolvable by the classloader returned by dd.trace.DDAdvice#getAgentClassLoader()
+ */
+ public HelperInjector(String... helperClassNames) {
+ this.helperClassNames = new HashSet<String>(Arrays.asList(helperClassNames));
+ }
+
+ private synchronized Map<TypeDescription, byte[]> getHelperMap() throws ClassNotFoundException {
+ if (helperMap == null) {
+ helperMap = new HashMap<TypeDescription, byte[]>(helperClassNames.size());
+ for (String helperName : helperClassNames) {
+ Class<?> helper = DDAdvice.getAgentClassLoader().loadClass(helperName);
+ helperMap.put(
+ new TypeDescription.ForLoadedType(helper),
+ ClassFileLocator.ForClassLoader.read(helper).resolve());
+ }
+ }
+ return helperMap;
+ }
+
+ @Override
+ public DynamicType.Builder<?> transform(
+ DynamicType.Builder<?> builder,
+ TypeDescription typeDescription,
+ ClassLoader classLoader,
+ JavaModule module) {
+ if (helperClassNames.size() > 0) {
+ try {
+ new ClassInjector.UsingReflection(classLoader).inject(getHelperMap());
+ } catch (ClassNotFoundException cnfe) {
+ log.error("Failed to inject helper classes into " + classLoader, cnfe);
+ throw new RuntimeException(cnfe);
+ }
+ }
+ return builder;
+ }
+}

MODIFY
dd-java-agent/tooling/src/test/groovy/dd/test/TestInstrumentation.java
dd-java-agent/tooling/src/test/groovy/dd/test/TestInstrumentation.java
@@ -13,6 +13,7 @@ public class TestInstrumentation implements Instrumenter {
public AgentBuilder instrument(AgentBuilder agentBuilder) {
return agentBuilder
.type(named(TestInstrumentation.class.getName() + "$ClassToInstrument"))
+ .transform(new HelperInjector(TestInstrumentation.class.getName() + "$HelperClass"))
.transform(
DDAdvice.create()
.advice(

Helper Injector Transformer