Ensure helpers are injected on a classloader that has instrumented code.

Author: tylerbenson

dd-java-agent-ittests/src/test/java/com/datadoghq/agent/instrumentation/annotation/TraceAnnotationsTest.java

@@ -21,6 +21,9 @@ public class TraceAnnotationsTest {
 
   @Before
   public void beforeTest() throws Exception {
+    Class.forName("com.datadoghq.agent.InstrumentationRulesManager")
+        .getMethod("registerClassLoad")
+        .invoke(null);
     try {
       GlobalTracer.register(tracer);
     } catch (final Exception e) {

dd-java-agent/src/main/java/com/datadoghq/agent/TracingAgent.java

@@ -25,7 +25,11 @@
 
 import dd.trace.Instrumenter;
 import java.lang.instrument.Instrumentation;
+import java.lang.reflect.Method;
+import java.util.Collections;
 import java.util.ServiceLoader;
+import java.util.Set;
+import java.util.WeakHashMap;
 import lombok.extern.slf4j.Slf4j;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.description.type.TypeDescription;
@@ -103,6 +107,9 @@ public static void addByteBuddy(final Instrumentation inst) {
   @Slf4j
   static class Listener implements AgentBuilder.Listener {
 
+    private final Set<ClassLoader> initializedClassloaders =
+        Collections.newSetFromMap(new WeakHashMap<ClassLoader, Boolean>());
+
     @Override
     public void onError(
         final String typeName,
@@ -121,6 +128,26 @@ public void onTransformation(
         final boolean loaded,
         final DynamicType dynamicType) {
       log.debug("Transformed {}", typeDescription);
+
+      if (classLoader == null) {
+        return;
+      }
+      synchronized (classLoader) {
+        if (initializedClassloaders.contains(classLoader)) {
+          return;
+        }
+        initializedClassloaders.add(classLoader);
+
+        try {
+          final Class<?> rulesManager =
+              Class.forName("com.datadoghq.agent.InstrumentationRulesManager", true, classLoader);
+          final Method registerClassLoad =
+              rulesManager.getDeclaredMethod("registerClassLoad", Object.class);
+          registerClassLoad.invoke(null, classLoader);
+        } catch (final Throwable e) {
+          log.info("ClassLoad Registration for target " + classLoader, e);
+        }
+      }
     }
 
     @Override