Basic implementation outline and happy-path test
MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java
@@ -4,6 +4,7 @@ import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.failSafe;
import static datadog.trace.agent.tooling.Utils.getConfigEnabled;
import static net.bytebuddy.matcher.ElementMatchers.any;
+import datadog.trace.agent.tooling.context.InstrumentationContextProvider;
import datadog.trace.agent.tooling.muzzle.Reference;
import datadog.trace.agent.tooling.muzzle.ReferenceMatcher;
import java.security.ProtectionDomain;
@@ -16,6 +17,10 @@ import java.util.Set;
import lombok.extern.slf4j.Slf4j;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.dynamic.DynamicType;
+import net.bytebuddy.implementation.auxiliary.MethodCallProxy;
+import net.bytebuddy.implementation.bytecode.StackManipulation;
+import net.bytebuddy.implementation.bytecode.assign.Assigner;
import net.bytebuddy.matcher.ElementMatcher;
import net.bytebuddy.utility.JavaModule;
@@ -49,7 +54,10 @@ public interface Instrumenter {
@Slf4j
abstract class Default implements Instrumenter {
private final Set<String> instrumentationNames;
- private final String instrumentationPrimaryName;
+ // FIXME: make field private again
+ public final String instrumentationPrimaryName;
+ // TODO: create a generic instrumentaton pipeline attacher?
+ private final InstrumentationContextProvider contextProvider;
protected final boolean enabled;
protected final String packageName =
@@ -74,10 +82,11 @@ public interface Instrumenter {
}
}
enabled = anyEnabled;
+ contextProvider = InstrumentationContextProvider.Creator.contextProviderFor(this);
}
@Override
- public AgentBuilder instrument(final AgentBuilder parentAgentBuilder) {
+ public final AgentBuilder instrument(final AgentBuilder parentAgentBuilder) {
if (!enabled) {
log.debug("Instrumentation {} is disabled", this);
return parentAgentBuilder;
@@ -117,7 +126,14 @@ public interface Instrumenter {
new AgentBuilder.Transformer.ForAdvice()
.include(Utils.getAgentClassLoader())
.withExceptionHandler(ExceptionHandlers.defaultExceptionHandler())
- .advice(entry.getKey(), entry.getValue()));
+ .advice(entry.getKey(), entry.getValue()))
+ .transform(new AgentBuilder.Transformer() {
+ @Override
+ public DynamicType.Builder<?> transform(DynamicType.Builder<?> builder, TypeDescription typeDescription, ClassLoader classLoader, JavaModule module) {
+ return builder.visit(contextProvider.getInstrumentationVisitor());
+ }
+ })
+ ;
}
return agentBuilder;
}

ADD
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/InstrumentationContextProvider.java
+package datadog.trace.agent.tooling.context;
+
+import datadog.trace.agent.tooling.Instrumenter;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.AsmVisitorWrapper;
+
+import java.util.List;
+
+public interface InstrumentationContextProvider {
+
+ /**
+ * @return An AsmVisitorWrapper to run on the instrumentation's bytecode.
+ */
+ AsmVisitorWrapper getInstrumentationVisitor();
+
+ /**
+ * @return A list of classes in byte-array format. These classes will be injected into the runtime classloader.
+ */
+ List<byte[]> dynamicClasses();
+
+ AgentBuilder additionalInstrumentation(AgentBuilder builder);
+
+ // TODO: better place to put factory/creator
+ class Creator {
+ private Creator() {}
+
+ public static InstrumentationContextProvider contextProviderFor(Instrumenter.Default instrumenter) {
+ return new MapBackedProvider(instrumenter);
+ }
+ }
+}

ADD
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/MapBackedProvider.java
+package datadog.trace.agent.tooling.context;
+
+import datadog.trace.agent.tooling.Instrumenter;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.AsmVisitorWrapper;
+import net.bytebuddy.description.field.FieldDescription;
+import net.bytebuddy.description.field.FieldList;
+import net.bytebuddy.description.method.MethodList;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.implementation.Implementation;
+import net.bytebuddy.jar.asm.ClassVisitor;
+import net.bytebuddy.jar.asm.MethodVisitor;
+import net.bytebuddy.jar.asm.Opcodes;
+import net.bytebuddy.pool.TypePool;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+
+/**
+ * InstrumentationContextProvider which stores context in a global map.
+ */
+public class MapBackedProvider implements InstrumentationContextProvider {
+ private final Instrumenter.Default instrumenter;
+
+ public MapBackedProvider(Instrumenter.Default instrumenter) {
+ this.instrumenter = instrumenter;
+ }
+
+ @Override
+ public AsmVisitorWrapper getInstrumentationVisitor() {
+ return new AsmVisitorWrapper() {
+ @Override
+ public int mergeWriter(int flags) {
+ // FIXME: any writer flags needed?
+ return flags;
+ }
+
+ @Override
+ public int mergeReader(int flags) {
+ // FIXME: any reader flags needed?
+ return flags;
+ }
+
+ @Override
+ public ClassVisitor wrap(TypeDescription instrumentedType, ClassVisitor classVisitor, Implementation.Context implementationContext, TypePool typePool, FieldList<FieldDescription.InDefinedShape> fields, MethodList<?> methods, int writerFlags, int readerFlags) {
+ return new ClassVisitor(Opcodes.ASM7, classVisitor) {
+ @Override
+ public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
+ super.visit(version, access, name, signature, superName, interfaces);
+ }
+
+ @Override
+ public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
+ // TODO
+ System.out.println(instrumenter.instrumentationPrimaryName + " context visit method: " + name + descriptor);
+ for (Map.Entry<String, String> entry : instrumenter.contextStore().entrySet()) {
+ System.out.println(" -- remap: " + entry.getKey() + " -> " + entry.getValue());
+ }
+ return super.visitMethod(access, name, descriptor, signature, exceptions);
+ }
+ };
+ }
+ };
+ }
+
+ @Override
+ public List<byte[]> dynamicClasses() {
+ // TODO: create a map-holder class for each key in the context-store
+ return Collections.emptyList();
+ }
+
+ @Override
+ public AgentBuilder additionalInstrumentation(AgentBuilder builder) {
+ return builder;
+ }
+
+ /**
+ * Template class used to generate the class holding the global map.
+ */
+ private static final class MapHolder {
+ }
+}

MODIFY
dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/ExceptionHandlerTest.groovy
dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/ExceptionHandlerTest.groovy
@@ -7,6 +7,7 @@ import datadog.trace.agent.tooling.ExceptionHandlers
import datadog.trace.bootstrap.ExceptionLogger
import net.bytebuddy.agent.ByteBuddyAgent
import net.bytebuddy.agent.builder.AgentBuilder
+import net.bytebuddy.agent.builder.ResettableClassFileTransformer
import net.bytebuddy.dynamic.ClassFileLocator
import org.slf4j.LoggerFactory
import spock.lang.Shared
@@ -18,6 +19,8 @@ import static net.bytebuddy.matcher.ElementMatchers.named
class ExceptionHandlerTest extends Specification {
@Shared
ListAppender testAppender = new ListAppender()
+ @Shared
+ ResettableClassFileTransformer transformer
def setupSpec() {
AgentBuilder builder = new AgentBuilder.Default()
@@ -40,7 +43,7 @@ class ExceptionHandlerTest extends Specification {
BadAdvice.NoOpAdvice.getName()))
ByteBuddyAgent.install()
- builder.installOn(ByteBuddyAgent.getInstrumentation())
+ transformer = builder.installOn(ByteBuddyAgent.getInstrumentation())
final Logger logger = (Logger) LoggerFactory.getLogger(ExceptionLogger)
testAppender.setContext(logger.getLoggerContext())
@@ -50,6 +53,7 @@ class ExceptionHandlerTest extends Specification {
def cleanupSpec() {
testAppender.stop()
+ transformer.reset(ByteBuddyAgent.getInstrumentation(), AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
}
def "exception handler invoked"() {

ADD
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.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 static net.bytebuddy.matcher.ElementMatchers.named
+
+class MapBackedProviderTest extends Specification {
+ @Shared
+ ResettableClassFileTransformer transformer
+
+ def setupSpec() {
+ final MapBackedProvider contextProvider = new MapBackedProvider(new TestInstrumenter())
+
+ AgentBuilder builder = new AgentBuilder.Default()
+ .disableClassFormatChanges()
+ .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
+ .type(named("datadog.trace.agent.test.context.ClassToRemap"))
+ .transform(new AgentBuilder.Transformer() {
+ @Override
+ DynamicType.Builder<?> transform(DynamicType.Builder<?> builder, TypeDescription typeDescription, ClassLoader classLoader, JavaModule module) {
+ return builder.visit(contextProvider.getInstrumentationVisitor())
+ }
+ })
+
+ ByteBuddyAgent.install()
+ transformer = builder.installOn(ByteBuddyAgent.getInstrumentation())
+ }
+
+ def cleanupSpec() {
+ transformer.reset(ByteBuddyAgent.getInstrumentation(), AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
+ }
+
+ def "correct api usage stores state in map"() {
+ setup:
+ ClassToRemap instance1 = new ClassToRemap()
+ ClassToRemap.mapObject(instance1)
+ ClassToRemap.mapObject(instance1)
+ ClassToRemap instance2 = new ClassToRemap()
+ ClassToRemap.mapObject(instance1)
+ ClassToRemap instance3 = new ClassToRemap()
+
+ expect:
+ ClassToRemap.mapObject(instance1) == 3
+ ClassToRemap.mapObject(instance2) == 2
+ ClassToRemap.mapObject(instance3) == 1
+ }
+
+
+ // TODO
+ def "incorrect api usage fails instrumentation"() {
+ setup:
+ println "loaded class: " + BadClassToRemap.getName()
+ expect:
+ 1 == 1
+ // Instrumentation remapper does not apply.
+ }
+
+ 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
+ }
+ }
+}

ADD
dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/test/context/BadClassToRemap.java
+package datadog.trace.agent.test.context;
+
+import datadog.trace.agent.tooling.context.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, State.class);
+ return ++state.anInt;
+ }
+
+ public static class State {
+ int anInt = 0;
+ }
+}

ADD
dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/test/context/ClassToRemap.java
+package datadog.trace.agent.test.context;
+
+import datadog.trace.agent.tooling.context.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, 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) {
+ final State state = InstrumentationContext.get(runnable, State.class);
+ return ++state.anInt;
+ }
+
+ public static class State {
+ int anInt = 0;
+ }
+}

ADD
dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/test/context/UserClass.java
+package datadog.trace.agent.test.context;
+
+public class UserClass {
+}

Basic implementation outline and happy-path test