Inject dynamic classes into topmost class loader
MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/MapBackedProvider.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/MapBackedProvider.java
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();
}

MODIFY
dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/context/MapBackedProviderTest.groovy
dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/context/MapBackedProviderTest.groovy
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,36 +21,38 @@ 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() {
- @Override
- DynamicType.Builder<?> transform(DynamicType.Builder<?> builder, TypeDescription typeDescription, ClassLoader classLoader, JavaModule module) {
- return builder.visit(contextProvider.getInstrumentationVisitor())
- }})
- .transform(new HelperInjector(contextProvider.dynamicClasses()))
+ .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
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) {}
+
@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) {}
})
@@ -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")

Inject dynamic classes into topmost class loader