Move ASM visiting method into context-store impl

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

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/InstrumentationContextProvider.java

@@ -1,20 +1,17 @@
 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
+   * Hook to provide an agent builder after advice is applied to target class. Used to implement
-   *     the runtime classloader.
+   * 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. */
+  /** Hook to define additional instrumentation. Run at instrumentation advice is hooked up. */
-  AgentBuilder additionalInstrumentation(AgentBuilder builder);
+  AgentBuilder.Identified.Extendable additionalInstrumentation(
+      AgentBuilder.Identified.Extendable builder);
 }

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/MapBackedProvider.java

@@ -1,5 +1,6 @@
 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
+  private Map<String, byte[]> dynamicClasses() {
-  public Map<String, byte[]> dynamicClasses() {
     return createOrGetDynamicClasses();
   }
 
   @Override
-  public AgentBuilder additionalInstrumentation(AgentBuilder builder) {
+  public AgentBuilder.Identified.Extendable additionalInstrumentation(
+      AgentBuilder.Identified.Extendable builder) {
     return builder;
   }