Move postProcessor option to InstrumenterModule, so it can be shared across the grouped instrumentations

Author: mcculls

dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java

@@ -9,7 +9,6 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
 import static net.bytebuddy.matcher.ElementMatchers.not;
 
-import datadog.trace.agent.tooling.Instrumenter.WithPostProcessor;
 import datadog.trace.agent.tooling.bytebuddy.ExceptionHandlers;
 import datadog.trace.agent.tooling.context.FieldBackedContextInjector;
 import datadog.trace.agent.tooling.context.FieldBackedContextMatcher;
@@ -73,17 +72,12 @@ public final class CombiningTransformerBuilder
   private Map<String, String> contextStore;
   private AgentBuilder.Transformer contextRequestRewriter;
   private HelperTransformer helperTransformer;
+  private Advice.PostProcessor.Factory postProcessor;
   private MuzzleCheck muzzle;
 
   // temporary buffer for collecting advice; reset for each instrumenter
   private final List<AgentBuilder.Transformer> advice = new ArrayList<>();
 
-  /**
-   * Post processor to be applied to instrumenter advices if they implement {@link
-   * WithPostProcessor}
-   */
-  private Advice.PostProcessor.Factory postProcessor;
-
   public CombiningTransformerBuilder(
       AgentBuilder agentBuilder, InstrumenterIndex instrumenterIndex) {
     this.agentBuilder = agentBuilder;
@@ -134,6 +128,8 @@ private void prepareInstrumentation(InstrumenterModule module, int instrumentati
                 module.useAgentCodeSource(), module.getClass().getSimpleName(), helperClassNames)
             : null;
 
+    postProcessor = module.postProcessor();
+
     muzzle = new MuzzleCheck(module, instrumentationId);
   }
 
@@ -208,9 +204,6 @@ private void buildTypeMatcher(Instrumenter member, int transformationId) {
 
   private void buildTypeAdvice(Instrumenter member, int transformationId) {
 
-    postProcessor =
-        member instanceof WithPostProcessor ? ((WithPostProcessor) member).postProcessor() : null;
-
     if (null != helperTransformer) {
       advice.add(helperTransformer);
     }

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java

@@ -5,7 +5,6 @@
 
 import java.security.ProtectionDomain;
 import java.util.Collection;
-import net.bytebuddy.asm.Advice;
 import net.bytebuddy.asm.AsmVisitorWrapper;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
@@ -82,11 +81,6 @@ interface WithTypeStructure {
     ElementMatcher<TypeDescription> structureMatcher();
   }
 
-  /** Instrumentation that wants to apply additional structure checks after type matching. */
-  interface WithPostProcessor {
-    Advice.PostProcessor.Factory postProcessor();
-  }
-
   /** Instrumentation that provides advice which affects the whole type. */
   interface HasTypeAdvice extends Instrumenter {
     /**

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java

@@ -156,6 +156,11 @@ public ElementMatcher<? super MethodDescription> methodIgnoreMatcher() {
     return isSynthetic();
   }
 
+  /** Override this to post-process the operand stack of any transformed methods. */
+  public Advice.PostProcessor.Factory postProcessor() {
+    return null;
+  }
+
   /**
    * Context stores to define for this instrumentation. Are added to matching class loaders.
    *
@@ -235,7 +240,7 @@ public boolean isApplicable(Set<TargetSystem> enabledSystems) {
 
   /** Parent class for all IAST related instrumentations */
   @SuppressForbidden
-  public abstract static class Iast extends InstrumenterModule implements WithPostProcessor {
+  public abstract static class Iast extends InstrumenterModule {
     public Iast(String instrumentationName, String... additionalNames) {
       super(instrumentationName, additionalNames);
     }

dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/TemplateAndMatrixVariablesInstrumentation.java

@@ -36,9 +36,7 @@
 /** Obtain template and matrix variables for RequestMappingInfoHandlerMapping. */
 @AutoService(InstrumenterModule.class)
 public class TemplateAndMatrixVariablesInstrumentation extends InstrumenterModule
-    implements Instrumenter.ForSingleType,
-        Instrumenter.HasMethodAdvice,
-        Instrumenter.WithPostProcessor {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
 
   private Advice.PostProcessor.Factory postProcessorFactory;
 

dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/TemplateVariablesUrlHandlerInstrumentation.java

@@ -35,9 +35,7 @@
 /** Obtain template and matrix variables for AbstractUrlHandlerMapping */
 @AutoService(InstrumenterModule.class)
 public class TemplateVariablesUrlHandlerInstrumentation extends InstrumenterModule
-    implements Instrumenter.ForSingleType,
-        Instrumenter.HasMethodAdvice,
-        Instrumenter.WithPostProcessor {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
 
   private Advice.PostProcessor.Factory postProcessorFactory;
 

dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java/datadog/trace/instrumentation/springweb6/TemplateAndMatrixVariablesInstrumentation.java

@@ -18,9 +18,7 @@
 /** Obtain template and matrix variables for RequestMappingInfoHandlerMapping. */
 @AutoService(InstrumenterModule.class)
 public class TemplateAndMatrixVariablesInstrumentation extends InstrumenterModule
-    implements Instrumenter.ForSingleType,
-        Instrumenter.HasMethodAdvice,
-        Instrumenter.WithPostProcessor {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
 
   private Advice.PostProcessor.Factory postProcessorFactory;
 

dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java/datadog/trace/instrumentation/springweb6/TemplateVariablesUrlHandlerInstrumentation.java

@@ -18,9 +18,7 @@
 /** Obtain template and matrix variables for AbstractUrlHandlerMapping */
 @AutoService(InstrumenterModule.class)
 public class TemplateVariablesUrlHandlerInstrumentation extends InstrumenterModule
-    implements Instrumenter.ForSingleType,
-        Instrumenter.HasMethodAdvice,
-        Instrumenter.WithPostProcessor {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
 
   private Advice.PostProcessor.Factory postProcessorFactory;
 

dd-java-agent/instrumentation/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/RouteImplInstrumentation.java

@@ -18,9 +18,7 @@
 
 @AutoService(InstrumenterModule.class)
 public class RouteImplInstrumentation extends InstrumenterModule
-    implements Instrumenter.ForKnownTypes,
-        Instrumenter.HasMethodAdvice,
-        Instrumenter.WithPostProcessor {
+    implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
 
   private Advice.PostProcessor.Factory postProcessorFactory;
 

dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RouteImplInstrumentation.java

@@ -16,9 +16,7 @@
 
 @AutoService(InstrumenterModule.class)
 public class RouteImplInstrumentation extends InstrumenterModule
-    implements Instrumenter.ForKnownTypes,
-        Instrumenter.HasMethodAdvice,
-        Instrumenter.WithPostProcessor {
+    implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
 
   private Advice.PostProcessor.Factory postProcessorFactory;