diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/MapBackedProvider.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/MapBackedProvider.java index 01126b10e7..d97941d88e 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/MapBackedProvider.java +++ b/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 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 dynamicClasses() { return createOrGetDynamicClasses(); } diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/context/MapBackedProviderTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/context/MapBackedProviderTest.groovy index f7794c05d0..168ec2c66e 100644 --- a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/context/MapBackedProviderTest.groovy +++ b/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")) - .or(named("datadog.trace.agent.test.context.BadClassToRemap")) - .transform(new AgentBuilder.Transformer() { + .type(named("datadog.trace.agent.test.context.ClassToRemap") + .or(named("datadog.trace.agent.test.context.BadClassToRemap"))) + builder = contextProvider.instrumentationTransformer(builder) + builder = builder + .with(new AgentBuilder.Listener() { @Override - DynamicType.Builder transform(DynamicType.Builder builder, TypeDescription typeDescription, ClassLoader classLoader, JavaModule module) { - 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) { } + 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")