Inject dynamic classes into topmost class loader

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/MapBackedProvider.java

@@ -1,5 +1,7 @@
 package datadog.trace.agent.tooling.context;
 
+import static datadog.trace.agent.tooling.ClassLoaderMatcher.BOOTSTRAP_CLASSLOADER;
+
 import datadog.trace.agent.tooling.HelperInjector;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.Utils;
@@ -8,6 +10,7 @@ import datadog.trace.bootstrap.WeakMap;
 import java.lang.reflect.Method;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.atomic.AtomicReference;
 import jdk.internal.org.objectweb.asm.ClassWriter;
@@ -95,7 +98,24 @@ public class MapBackedProvider implements InstrumentationContextProvider {
                   return builder.visit(getInstrumentationVisitor());
                 }
               });
-      builder = builder.transform(new HelperInjector(dynamicClasses()));
+      builder =
+          builder.transform(
+              new AgentBuilder.Transformer() {
+                final HelperInjector injector = new HelperInjector(dynamicClasses());
+
+                @Override
+                public DynamicType.Builder<?> transform(
+                    DynamicType.Builder<?> builder,
+                    TypeDescription typeDescription,
+                    ClassLoader classLoader,
+                    JavaModule module) {
+                  return injector.transform(
+                      builder,
+                      typeDescription,
+                      findClassLoaderForMapHolderInjection(classLoader),
+                      module);
+                }
+              });
     }
     return builder;
   }
@@ -239,6 +259,49 @@ public class MapBackedProvider implements InstrumentationContextProvider {
     };
   }
 
+  /**
+   * Find the topmost classloader in the classloader hierarchy (beginning from startingLoader) which
+   * can load all user classes defined in the instrumentation context store.
+   */
+  private ClassLoader findClassLoaderForMapHolderInjection(final ClassLoader startingLoader) {
+    final Set<String> userClassNames = instrumenter.contextStore().keySet();
+    ClassLoader lastGoodLoader = startingLoader;
+    ClassLoader searchLoader = startingLoader;
+    // search up the classloader hierarchy for a classloader which can see all user classes
+    while (true) {
+      log.debug(
+          "{}: Searching for classloader to inject dynamic context map: {}",
+          instrumenter,
+          searchLoader);
+      if (searchLoader == BOOTSTRAP_CLASSLOADER) {
+        searchLoader = Utils.getBootstrapProxy();
+      }
+      boolean allClassesFound = true;
+      for (final String userClassName : userClassNames) {
+        if (searchLoader.getResource(Utils.getResourceName(userClassName)) == null) {
+          allClassesFound = false;
+          break;
+        }
+      }
+      if (!allClassesFound) {
+        // current searchLoader can't see resources. Use last good classloader.
+        break;
+      }
+      lastGoodLoader = searchLoader;
+      if (searchLoader == Utils.getBootstrapProxy()) {
+        // all classloaders can see resources. Use bootstrap
+        searchLoader = BOOTSTRAP_CLASSLOADER;
+        lastGoodLoader = BOOTSTRAP_CLASSLOADER;
+        break;
+      }
+      searchLoader = searchLoader.getParent();
+    }
+
+    log.debug(
+        "{}: Found classloader to inject dynamic context map: {}", instrumenter, lastGoodLoader);
+    return lastGoodLoader;
+  }
+
   private Map<String, byte[]> dynamicClasses() {
     return createOrGetDynamicClasses();
   }

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

@@ -1,7 +1,6 @@
 package datadog.trace.agent.test.context
 
 import datadog.trace.agent.test.TestUtils
-import datadog.trace.agent.tooling.HelperInjector
 import datadog.trace.agent.tooling.Instrumenter
 import datadog.trace.agent.tooling.context.MapBackedProvider
 import net.bytebuddy.agent.ByteBuddyAgent
@@ -22,38 +21,40 @@ import static net.bytebuddy.matcher.ElementMatchers.named
 class MapBackedProviderTest extends Specification {
   @Shared
   ResettableClassFileTransformer transformer
+  @Shared
+  MapBackedProvider contextProvider
 
   def setupSpec() {
-    final MapBackedProvider contextProvider = new MapBackedProvider(new TestInstrumenter())
+    contextProvider = new MapBackedProvider(new TestInstrumenter())
 
     AgentBuilder builder = new AgentBuilder.Default()
       .disableClassFormatChanges()
       .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
-      .type(named("datadog.trace.agent.test.context.ClassToRemap"))
+      .type(named("datadog.trace.agent.test.context.ClassToRemap")
-      .or(named("datadog.trace.agent.test.context.BadClassToRemap"))
+      .or(named("datadog.trace.agent.test.context.BadClassToRemap")))
-      .transform(new AgentBuilder.Transformer() {
+    builder = contextProvider.instrumentationTransformer(builder)
+    builder = builder
+      .with(new AgentBuilder.Listener() {
       @Override
-      DynamicType.Builder<?> transform(DynamicType.Builder<?> builder, TypeDescription typeDescription, ClassLoader classLoader, JavaModule module) {
+      void onDiscovery(String typeName, ClassLoader classLoader, JavaModule module, boolean loaded) {}
-        return builder.visit(contextProvider.getInstrumentationVisitor())
+
-      }})
-    .transform(new HelperInjector(contextProvider.dynamicClasses()))
-    .with(new AgentBuilder.Listener() {
-      @Override
-      void onDiscovery(String typeName, ClassLoader classLoader, JavaModule module, boolean loaded) { }
       @Override
       void onTransformation(TypeDescription typeDescription, ClassLoader classLoader, JavaModule module, boolean loaded, DynamicType dynamicType) {
         assert !"datadog.trace.agent.test.context.BadClassToRemap".equals(typeDescription.getName())
       }
+
       @Override
-      void onIgnored(TypeDescription typeDescription, ClassLoader classLoader, JavaModule module, boolean loaded) { }
+      void onIgnored(TypeDescription typeDescription, ClassLoader classLoader, JavaModule module, boolean loaded) {}
+
       @Override
       void onError(String typeName, ClassLoader classLoader, JavaModule module, boolean loaded, Throwable throwable) {
         assert "datadog.trace.agent.test.context.BadClassToRemap".equals(typeName)
         System.err.println("Exception during test")
         throwable.printStackTrace()
       }
+
       @Override
-      void onComplete(String typeName, ClassLoader classLoader, JavaModule module, boolean loaded) { }
+      void onComplete(String typeName, ClassLoader classLoader, JavaModule module, boolean loaded) {}
     })
 
     ByteBuddyAgent.install()
@@ -120,6 +121,16 @@ class MapBackedProviderTest extends Specification {
     thrown RuntimeException
   }
 
+  def "inject dynamic classes into correct classloader"() {
+    when:
+    for (final String dynamicClassName : contextProvider.dynamicClasses().keySet()) {
+      assert getClass().getClassLoader().loadClass(dynamicClassName).getClassLoader() == getClass().getClassLoader()
+    }
+
+    then:
+    noExceptionThrown()
+  }
+
   static class TestInstrumenter extends Instrumenter.Default {
     TestInstrumenter() {
       super("test")