Inject all map holder classes into the bootstrap
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
@@ -10,7 +10,6 @@ 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;
@@ -112,7 +111,8 @@ public class MapBackedProvider implements InstrumentationContextProvider {
return injector.transform(
builder,
typeDescription,
- findClassLoaderForMapHolderInjection(classLoader),
+ // dynamic map classes will always go to the bootstrap
+ BOOTSTRAP_CLASSLOADER,
module);
}
});
@@ -259,49 +259,6 @@ 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();
}

DELETE
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.Instrumenter
-import datadog.trace.agent.tooling.context.MapBackedProvider
-import net.bytebuddy.agent.ByteBuddyAgent
-import net.bytebuddy.agent.builder.AgentBuilder
-import net.bytebuddy.agent.builder.ResettableClassFileTransformer
-import net.bytebuddy.description.type.TypeDescription
-import net.bytebuddy.dynamic.DynamicType
-import net.bytebuddy.matcher.ElementMatcher
-import net.bytebuddy.utility.JavaModule
-import spock.lang.Shared
-import spock.lang.Specification
-
-import java.lang.ref.WeakReference
-import java.util.concurrent.Callable
-
-import static net.bytebuddy.matcher.ElementMatchers.named
-
-class MapBackedProviderTest extends Specification {
- @Shared
- ResettableClassFileTransformer transformer
- @Shared
- MapBackedProvider contextProvider
-
- def setupSpec() {
- 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")))
- 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) {}
- })
-
- ByteBuddyAgent.install()
- transformer = builder.installOn(ByteBuddyAgent.getInstrumentation())
- }
-
- def cleanupSpec() {
- transformer.reset(ByteBuddyAgent.getInstrumentation(), AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
- }
-
- def "correct api usage stores state in map"() {
- when:
- ClassToRemap instance1 = new ClassToRemap()
- ClassToRemap.mapObject(instance1)
- ClassToRemap.mapObject(instance1)
- ClassToRemap instance2 = new ClassToRemap()
- ClassToRemap.mapObject(instance2)
- ClassToRemap instance3 = new ClassToRemap()
-
- then:
- ClassToRemap.mapObject(instance1) == 3
- ClassToRemap.mapObject(instance2) == 2
- ClassToRemap.mapObject(instance3) == 1
-
- when:
- Runnable r1 = new TestRunnable()
- Runnable r2 = new TestRunnable()
- ClassToRemap.mapOtherObject(r1)
-
- then:
- ClassToRemap.mapOtherObject(r1) == 2
- ClassToRemap.mapOtherObject(r2) == 1
- }
-
- def "backing map should not create strong refs to user instances"() {
- when:
- ClassToRemap instance = new ClassToRemap()
- int count = ClassToRemap.mapObject(instance)
- WeakReference<ClassToRemap> instanceRef = new WeakReference(instance)
- instance = null
- TestUtils.awaitGC(instanceRef)
-
- then:
- instanceRef.get() == null
- count == 1
- }
-
- def "incorrect api usage fails at class load time"() {
- when:
- BadClassToRemap.getName()
-
- then:
- noExceptionThrown()
-
- when:
- BadClassToRemap.unmappedObjectError(new Callable() {
- @Override
- Object call() throws Exception {
- return null
- }
- })
-
- then:
- 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")
- }
-
- @Override
- ElementMatcher<? super TypeDescription> typeMatcher() {
- return named("not invoked")
- }
-
- @Override
- Map<? extends ElementMatcher, String> transformers() {
- return Collections.emptyMap()
- }
-
- @Override
- Map<String, String> contextStore() {
- final Map<String, String> contextStore = new HashMap<>(2)
- contextStore.put('datadog.trace.agent.test.context.ClassToRemap', 'datadog.trace.agent.test.context.ClassToRemap$State')
- contextStore.put('java.lang.Runnable', 'datadog.trace.agent.test.context.ClassToRemap$State')
- return contextStore
- }
- }
-
- static class TestRunnable implements Runnable {
- @Override
- void run() {}
- }
-}

DELETE
dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/test/context/BadClassToRemap.java
-package datadog.trace.agent.test.context;
-
-import datadog.trace.bootstrap.InstrumentationContext;
-import java.util.concurrent.Callable;
-
-/** A class which incorrectly uses the context api. */
-public class BadClassToRemap {
- /**
- * Callable does not have context class of type State. This should throw an exception during class
- * transform.
- */
- public static int unmappedObjectError(Callable callable) {
- final State state = InstrumentationContext.get(callable, Callable.class, State.class);
- return ++state.anInt;
- }
-
- public static class State {
- int anInt = 0;
- }
-}

DELETE
dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/test/context/ClassToRemap.java
-package datadog.trace.agent.test.context;
-
-import datadog.trace.bootstrap.InstrumentationContext;
-
-/** A class which correctly uses the context api. */
-public class ClassToRemap {
-
- /**
- * ClassToRemap has context class of type State.
- *
- * @return The value of State#anInt
- */
- public static int mapObject(final ClassToRemap classToRemap) {
- final State state = InstrumentationContext.get(classToRemap, ClassToRemap.class, State.class);
- return ++state.anInt;
- }
-
- /**
- * Runnable has context class of type State.
- *
- * @return The value of State#anInt
- */
- public static int mapOtherObject(final Runnable runnable) {
- State state = InstrumentationContext.get(runnable, Runnable.class, State.class);
- return ++state.anInt;
- }
-
- public static class State {
- public int anInt = 0;
- public Object anObject = new Object();
- }
-}

ADD
dd-java-agent/testing/src/test/groovy/context/MapBackedProviderTest.groovy
+package context
+
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.agent.test.TestUtils
+import net.bytebuddy.utility.JavaModule
+
+import java.lang.ref.WeakReference
+
+class MapBackedProviderTest extends AgentTestRunner {
+
+ def setupSpec() {
+ assert new UserClass1().isInstrumented()
+ }
+
+ @Override
+ boolean onInstrumentationError(
+ final String typeName,
+ final ClassLoader classLoader,
+ final JavaModule module,
+ final boolean loaded,
+ final Throwable throwable) {
+ // UserClass2 asserts on incorrect api usage. Error expected.
+ return !(typeName.equals(UserClass2.getName()) && throwable.getMessage().startsWith("Incorrect Context Api Usage detected."))
+ }
+
+ def "correct api usage stores state in map"() {
+ when:
+ UserClass1 instance1 = new UserClass1()
+ UserClass1 instance2 = new UserClass1()
+ instance1.incrementContextCount()
+
+ then:
+ instance1.incrementContextCount() == 2
+ instance2.incrementContextCount() == 1
+ }
+
+ def "backing map should not create strong refs to user instances"() {
+ when:
+ UserClass1 instance = new UserClass1()
+ final int count = instance.incrementContextCount()
+ WeakReference<UserClass1> instanceRef = new WeakReference(instance)
+ instance = null
+ TestUtils.awaitGC(instanceRef)
+
+ then:
+ instanceRef.get() == null
+ count == 1
+ }
+
+ def "incorrect api usage fails at class load time"() {
+ expect:
+ !new UserClass2().isInstrumented()
+ }
+}

ADD
dd-java-agent/testing/src/test/java/context/ContextTestInstrumentation.java
+package context;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.InstrumentationContext;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public class ContextTestInstrumentation extends Instrumenter.Default {
+ public ContextTestInstrumentation() {
+ super("context-test-isntrumenter1");
+ }
+
+ @Override
+ public ElementMatcher<? super TypeDescription> typeMatcher() {
+ return named(getClass().getPackage().getName() + ".UserClass1")
+ .or(named(getClass().getPackage().getName() + ".UserClass2"));
+ }
+
+ @Override
+ public Map<? extends ElementMatcher, String> transformers() {
+ final Map<ElementMatcher, String> transformers = new HashMap<>(2);
+ transformers.put(named("isInstrumented"), MarkInstrumentedAdvice.class.getName());
+ transformers.put(named("incrementContextCount"), CorrectContextApiUsageAdvice.class.getName());
+ transformers.put(
+ named("incrementContextCountCountBroken"), IncorrectContextApiUsageAdvice.class.getName());
+ return transformers;
+ }
+
+ @Override
+ public String[] helperClassNames() {
+ return new String[] {getClass().getName() + "$UserClass1State"};
+ }
+
+ @Override
+ public Map<String, String> contextStore() {
+ return Collections.singletonMap(
+ getClass().getPackage().getName() + ".UserClass1",
+ getClass().getName() + "$UserClass1State");
+ }
+
+ public static class MarkInstrumentedAdvice {
+ @Advice.OnMethodExit
+ public static void markInstrumented(@Advice.Return(readOnly = false) boolean isInstrumented) {
+ isInstrumented = true;
+ }
+ }
+
+ public static class CorrectContextApiUsageAdvice {
+ @Advice.OnMethodExit
+ public static void storeAndIncrement(
+ @Advice.This Object thiz, @Advice.Return(readOnly = false) int contextCount) {
+ UserClass1State state =
+ InstrumentationContext.get((UserClass1) thiz, UserClass1.class, UserClass1State.class);
+ contextCount = ++state.count;
+ }
+ }
+
+ public static class IncorrectContextApiUsageAdvice {
+ @Advice.OnMethodExit
+ public static void storeAndIncrement(
+ @Advice.This Object thiz, @Advice.Return(readOnly = false) int contextCount) {
+ UserClass1State state = InstrumentationContext.get(thiz, Object.class, UserClass1State.class);
+ contextCount = ++state.count;
+ }
+ }
+
+ public static class UserClass1State {
+ int count = 0;
+ }
+}

ADD
dd-java-agent/testing/src/test/java/context/UserClass1.java
+package context;
+
+public class UserClass1 {
+ public boolean isInstrumented() {
+ // implementation replaced with test instrumentation
+ return false;
+ }
+
+ public int incrementContextCount() {
+ // implementation replaced with test instrumentation
+ return -1;
+ }
+}

ADD
dd-java-agent/testing/src/test/java/context/UserClass2.java
+package context;
+
+public class UserClass2 {
+ public boolean isInstrumented() {
+ return false;
+ }
+
+ public int incrementContextCountCountBroken() {
+ // instrumentation will not apply to this class because advice incorrectly uses context api
+ return -1;
+ }
+}

Inject all map holder classes into the bootstrap