Move ASM visiting method into context-store impl
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
@@ -18,7 +18,6 @@ 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.matcher.ElementMatcher;
import net.bytebuddy.utility.JavaModule;
@@ -101,8 +100,9 @@ public interface Instrumenter {
.and(new MuzzleMatcher())
.transform(DDTransformers.defaultTransformers());
agentBuilder = injectHelperClasses(agentBuilder);
- agentBuilder = applyContextStoreTransform(agentBuilder);
+ agentBuilder = contextProvider.instrumentationTransformer(agentBuilder);
agentBuilder = applyInstrumentationTransformers(agentBuilder);
+ agentBuilder = contextProvider.additionalInstrumentation(agentBuilder);
return agentBuilder;
}
@@ -128,26 +128,6 @@ public interface Instrumenter {
return agentBuilder;
}
- private AgentBuilder.Identified.Extendable applyContextStoreTransform(
- AgentBuilder.Identified.Extendable agentBuilder) {
- if (contextStore().size() > 0) {
- agentBuilder =
- agentBuilder.transform(
- new AgentBuilder.Transformer() {
- @Override
- public DynamicType.Builder<?> transform(
- DynamicType.Builder<?> builder,
- TypeDescription typeDescription,
- ClassLoader classLoader,
- JavaModule module) {
- return builder.visit(contextProvider.getInstrumentationVisitor());
- }
- });
- agentBuilder = agentBuilder.transform(new HelperInjector(contextProvider.dynamicClasses()));
- }
- return agentBuilder;
- }
-
/** Matches classes for which instrumentation is not muzzled. */
private class MuzzleMatcher implements AgentBuilder.RawMatcher {
@Override

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/InstrumentationContextProvider.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/InstrumentationContextProvider.java
package datadog.trace.agent.tooling.context;
-import java.util.Map;
import net.bytebuddy.agent.builder.AgentBuilder;
-import net.bytebuddy.asm.AsmVisitorWrapper;
public interface InstrumentationContextProvider {
- /** @return An AsmVisitorWrapper to run on the instrumentation's bytecode. */
- AsmVisitorWrapper getInstrumentationVisitor();
-
/**
- * @return A map of dynamic-class-name -> dynamic-class-bytes. These classes will be injected into
- * the runtime classloader.
+ * Hook to provide an agent builder after advice is applied to target class. Used to implement
+ * context-store lookup.
*/
- Map<String, byte[]> dynamicClasses();
+ AgentBuilder.Identified.Extendable instrumentationTransformer(
+ AgentBuilder.Identified.Extendable builder);
- /** Hook for the context impl to define additional instrumentation if needed. */
- AgentBuilder additionalInstrumentation(AgentBuilder builder);
+ /** Hook to define additional instrumentation. Run at instrumentation advice is hooked up. */
+ AgentBuilder.Identified.Extendable additionalInstrumentation(
+ AgentBuilder.Identified.Extendable builder);
}

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 datadog.trace.agent.tooling.HelperInjector;
import datadog.trace.agent.tooling.Instrumenter;
import datadog.trace.agent.tooling.Utils;
import datadog.trace.bootstrap.InstrumentationContext;
@@ -18,6 +19,7 @@ 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.dynamic.DynamicType;
import net.bytebuddy.implementation.Implementation;
import net.bytebuddy.jar.asm.ClassVisitor;
import net.bytebuddy.jar.asm.Label;
@@ -25,6 +27,7 @@ import net.bytebuddy.jar.asm.MethodVisitor;
import net.bytebuddy.jar.asm.Opcodes;
import net.bytebuddy.jar.asm.Type;
import net.bytebuddy.pool.TypePool;
+import net.bytebuddy.utility.JavaModule;
/**
* InstrumentationContextProvider which stores context in a global map.
@@ -77,8 +80,27 @@ public class MapBackedProvider implements InstrumentationContextProvider {
this.instrumenter = instrumenter;
}
+ public AgentBuilder.Identified.Extendable instrumentationTransformer(
+ AgentBuilder.Identified.Extendable builder) {
+ if (instrumenter.contextStore().size() > 0) {
+ builder =
+ builder.transform(
+ new AgentBuilder.Transformer() {
@Override
- public AsmVisitorWrapper getInstrumentationVisitor() {
+ public DynamicType.Builder<?> transform(
+ DynamicType.Builder<?> builder,
+ TypeDescription typeDescription,
+ ClassLoader classLoader,
+ JavaModule module) {
+ return builder.visit(getInstrumentationVisitor());
+ }
+ });
+ builder = builder.transform(new HelperInjector(dynamicClasses()));
+ }
+ return builder;
+ }
+
+ private AsmVisitorWrapper getInstrumentationVisitor() {
return new AsmVisitorWrapper() {
@Override
public int mergeWriter(int flags) {
@@ -217,13 +239,13 @@ public class MapBackedProvider implements InstrumentationContextProvider {
};
}
- @Override
- public Map<String, byte[]> dynamicClasses() {
+ private Map<String, byte[]> dynamicClasses() {
return createOrGetDynamicClasses();
}
@Override
- public AgentBuilder additionalInstrumentation(AgentBuilder builder) {
+ public AgentBuilder.Identified.Extendable additionalInstrumentation(
+ AgentBuilder.Identified.Extendable builder) {
return builder;
}

Move ASM visiting method into context-store impl