Move HasMethodAdvice aspect to instrumentations, away from modules (#…

…8139)

* Load 'AbcModule' definitions from the instrumentation class-loader so they can be unloaded after startup
* Move postProcessor option to InstrumenterModule, so it can be shared across the grouped instrumentations

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/DatadogClassLoader.java

@@ -93,7 +93,9 @@ protected Enumeration<URL> findResources(String name) {
   @Override
   protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
     if (name.startsWith("datadog.trace.instrumentation.")
-        && (name.endsWith("$Muzzle") || name.endsWith("Instrumentation"))) {
+        && (name.endsWith("$Muzzle")
+            || name.endsWith("Instrumentation")
+            || name.endsWith("Module"))) {
       InstrumentationClassLoader cl;
       if (null == (cl = instrumentationClassLoader.get())) {
         synchronized (instrumentationClassLoaderLock) {

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-builder/src/test/groovy/datadog/trace/agent/test/DefaultInstrumenterForkedTest.groovy

@@ -141,8 +141,5 @@ class DefaultInstrumenterForkedTest extends DDSpecification {
     TestDefaultInstrumenter(String instrumentationName, String additionalName) {
       super(instrumentationName, [additionalName] as String[])
     }
-
-    @Override
-    void methodAdvice(MethodTransformer transformer) {}
   }
 }

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.
    *
@@ -191,7 +196,7 @@ protected final boolean isShortcutMatchingEnabled(boolean defaultToShortcut) {
   }
 
   /** Parent class for all tracing related instrumentations */
-  public abstract static class Tracing extends InstrumenterModule implements HasMethodAdvice {
+  public abstract static class Tracing extends InstrumenterModule {
     public Tracing(String instrumentationName, String... additionalNames) {
       super(instrumentationName, additionalNames);
     }
@@ -203,7 +208,7 @@ public boolean isApplicable(Set<TargetSystem> enabledSystems) {
   }
 
   /** Parent class for all profiling related instrumentations */
-  public abstract static class Profiling extends InstrumenterModule implements HasMethodAdvice {
+  public abstract static class Profiling extends InstrumenterModule {
     public Profiling(String instrumentationName, String... additionalNames) {
       super(instrumentationName, additionalNames);
     }
@@ -222,7 +227,7 @@ public boolean isEnabled() {
   }
 
   /** Parent class for all AppSec related instrumentations */
-  public abstract static class AppSec extends InstrumenterModule implements HasMethodAdvice {
+  public abstract static class AppSec extends InstrumenterModule {
     public AppSec(String instrumentationName, String... additionalNames) {
       super(instrumentationName, additionalNames);
     }
@@ -235,8 +240,7 @@ public boolean isApplicable(Set<TargetSystem> enabledSystems) {
 
   /** Parent class for all IAST related instrumentations */
   @SuppressForbidden
-  public abstract static class Iast extends InstrumenterModule
-      implements HasMethodAdvice, WithPostProcessor {
+  public abstract static class Iast extends InstrumenterModule {
     public Iast(String instrumentationName, String... additionalNames) {
       super(instrumentationName, additionalNames);
     }
@@ -291,7 +295,7 @@ protected boolean isOptOutEnabled() {
   }
 
   /** Parent class for all USM related instrumentations */
-  public abstract static class Usm extends InstrumenterModule implements HasMethodAdvice {
+  public abstract static class Usm extends InstrumenterModule {
     public Usm(String instrumentationName, String... additionalNames) {
       super(instrumentationName, additionalNames);
     }
@@ -303,7 +307,7 @@ public boolean isApplicable(Set<TargetSystem> enabledSystems) {
   }
 
   /** Parent class for all CI related instrumentations */
-  public abstract static class CiVisibility extends InstrumenterModule implements HasMethodAdvice {
+  public abstract static class CiVisibility extends InstrumenterModule {
     public CiVisibility(String instrumentationName, String... additionalNames) {
       super(instrumentationName, additionalNames);
     }

dd-java-agent/instrumentation/aerospike-4/src/main/java/datadog/trace/instrumentation/aerospike4/AerospikeClientInstrumentation.java

@@ -17,7 +17,7 @@
 
 @AutoService(InstrumenterModule.class)
 public final class AerospikeClientInstrumentation extends InstrumenterModule.Tracing
-    implements Instrumenter.ForSingleType {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
   public AerospikeClientInstrumentation() {
     super("aerospike");
   }

dd-java-agent/instrumentation/aerospike-4/src/main/java/datadog/trace/instrumentation/aerospike4/CommandInstrumentation.java

@@ -19,7 +19,7 @@
 
 @AutoService(InstrumenterModule.class)
 public final class CommandInstrumentation extends InstrumenterModule.Tracing
-    implements Instrumenter.ForSingleType {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
   public CommandInstrumentation() {
     super("aerospike");
   }

dd-java-agent/instrumentation/aerospike-4/src/main/java/datadog/trace/instrumentation/aerospike4/NioEventLoopInstrumentation.java

@@ -16,7 +16,7 @@
 
 @AutoService(InstrumenterModule.class)
 public final class NioEventLoopInstrumentation extends InstrumenterModule.Tracing
-    implements Instrumenter.ForSingleType {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
   public NioEventLoopInstrumentation() {
     super("aerospike", "java_concurrent");
   }

dd-java-agent/instrumentation/aerospike-4/src/main/java/datadog/trace/instrumentation/aerospike4/PartitionInstrumentation.java

@@ -20,7 +20,7 @@
 
 @AutoService(InstrumenterModule.class)
 public final class PartitionInstrumentation extends InstrumenterModule.Tracing
-    implements Instrumenter.ForSingleType {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
   public PartitionInstrumentation() {
     super("aerospike");
   }

dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java

@@ -21,7 +21,7 @@
 
 @AutoService(InstrumenterModule.class)
 public class AkkaActorCellInstrumentation extends InstrumenterModule.Tracing
-    implements Instrumenter.ForSingleType {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
 
   public AkkaActorCellInstrumentation() {
     super("akka_actor_receive", "akka_actor", "akka_concurrent", "java_concurrent");

dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaEnvelopeInstrumentation.java

@@ -15,7 +15,7 @@
 
 @AutoService(InstrumenterModule.class)
 public class AkkaEnvelopeInstrumentation extends InstrumenterModule.Tracing
-    implements Instrumenter.ForSingleType {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
 
   public AkkaEnvelopeInstrumentation() {
     super("akka_actor_send", "akka_actor", "akka_concurrent", "java_concurrent");

dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaForkJoinExecutorTaskInstrumentation.java

@@ -25,7 +25,7 @@
  */
 @AutoService(InstrumenterModule.class)
 public final class AkkaForkJoinExecutorTaskInstrumentation extends InstrumenterModule.Tracing
-    implements Instrumenter.ForSingleType {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
   public AkkaForkJoinExecutorTaskInstrumentation() {
     super("java_concurrent", "akka_concurrent");
   }

dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaForkJoinPoolInstrumentation.java

@@ -19,7 +19,7 @@
 
 @AutoService(InstrumenterModule.class)
 public final class AkkaForkJoinPoolInstrumentation extends InstrumenterModule.Tracing
-    implements Instrumenter.ForSingleType {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
 
   public AkkaForkJoinPoolInstrumentation() {
     super("java_concurrent", "akka_concurrent");

dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaForkJoinTaskInstrumentation.java

@@ -39,7 +39,7 @@
  */
 @AutoService(InstrumenterModule.class)
 public final class AkkaForkJoinTaskInstrumentation extends InstrumenterModule.Tracing
-    implements Instrumenter.ForTypeHierarchy, ExcludeFilterProvider {
+    implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice, ExcludeFilterProvider {
 
   public AkkaForkJoinTaskInstrumentation() {
     super("java_concurrent", "akka_concurrent");

dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java

@@ -22,7 +22,7 @@
 
 @AutoService(InstrumenterModule.class)
 public class AkkaMailboxInstrumentation extends InstrumenterModule.Tracing
-    implements Instrumenter.ForSingleType, ExcludeFilterProvider {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice, ExcludeFilterProvider {
 
   public AkkaMailboxInstrumentation() {
     super("akka_actor_mailbox", "akka_actor", "akka_concurrent", "java_concurrent");

dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaRoutedActorCellInstrumentation.java

@@ -19,7 +19,7 @@
 
 @AutoService(InstrumenterModule.class)
 public class AkkaRoutedActorCellInstrumentation extends InstrumenterModule.Tracing
-    implements Instrumenter.ForSingleType {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
 
   public AkkaRoutedActorCellInstrumentation() {
     super("akka_actor_send", "akka_actor", "akka_concurrent", "java_concurrent");

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttp2ServerInstrumentation.java

@@ -22,7 +22,7 @@
  */
 @AutoService(InstrumenterModule.class)
 public final class AkkaHttp2ServerInstrumentation extends InstrumenterModule.Tracing
-    implements Instrumenter.ForKnownTypes {
+    implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
   public AkkaHttp2ServerInstrumentation() {
     super("akka-http2", "akka-http", "akka-http-server");
   }

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpServerInstrumentation.java

@@ -51,7 +51,7 @@
  */
 @AutoService(InstrumenterModule.class)
 public final class AkkaHttpServerInstrumentation extends InstrumenterModule.Tracing
-    implements Instrumenter.ForSingleType {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
   public AkkaHttpServerInstrumentation() {
     super("akka-http", "akka-http-server");
   }

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpSingleRequestInstrumentation.java

@@ -24,7 +24,7 @@
 
 @AutoService(InstrumenterModule.class)
 public final class AkkaHttpSingleRequestInstrumentation extends InstrumenterModule.Tracing
-    implements Instrumenter.ForSingleType {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
   public AkkaHttpSingleRequestInstrumentation() {
     super("akka-http", "akka-http-client");
   }

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaPoolMasterActorInstrumentation.java

@@ -12,7 +12,7 @@
 
 @AutoService(InstrumenterModule.class)
 public final class AkkaPoolMasterActorInstrumentation extends InstrumenterModule.Tracing
-    implements Instrumenter.ForSingleType {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
   public AkkaPoolMasterActorInstrumentation() {
     super("akka-http", "akka-http-client");
   }

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DefaultExceptionHandlerInstrumentation.java

@@ -15,7 +15,7 @@
 
 @AutoService(InstrumenterModule.class)
 public class DefaultExceptionHandlerInstrumentation extends InstrumenterModule.AppSec
-    implements Instrumenter.ForSingleType {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
   public DefaultExceptionHandlerInstrumentation() {
     super("akka-http", "akka-http-server");
   }

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/appsec/Bug4304Instrumentation.java

@@ -25,7 +25,9 @@
 /** See https://github.com/akka/akka-http/issues/4304 */
 @AutoService(InstrumenterModule.class)
 public class Bug4304Instrumentation extends InstrumenterModule.AppSec
-    implements Instrumenter.ForTypeHierarchy, Instrumenter.WithTypeStructure {
+    implements Instrumenter.ForTypeHierarchy,
+        Instrumenter.WithTypeStructure,
+        Instrumenter.HasMethodAdvice {
   public Bug4304Instrumentation() {
     super("akka-http");
   }

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/appsec/ConfigProvideRemoteAddressHeaderInstrumentation.java

@@ -13,7 +13,7 @@
 
 @AutoService(InstrumenterModule.class)
 public class ConfigProvideRemoteAddressHeaderInstrumentation extends InstrumenterModule.AppSec
-    implements Instrumenter.ForSingleType {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
   public ConfigProvideRemoteAddressHeaderInstrumentation() {
     super("akka-http");
   }

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/appsec/FormDataToStrictInstrumentation.java

@@ -18,7 +18,9 @@
 /** @see akka.http.scaladsl.model.Multipart.FormData#toStrict(FiniteDuration, Materializer) */
 @AutoService(InstrumenterModule.class)
 public class FormDataToStrictInstrumentation extends InstrumenterModule.AppSec
-    implements Instrumenter.ForSingleType, ScalaListCollectorMuzzleReferences {
+    implements Instrumenter.ForSingleType,
+        Instrumenter.HasMethodAdvice,
+        ScalaListCollectorMuzzleReferences {
   public FormDataToStrictInstrumentation() {
     super("akka-http");
   }

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/appsec/JacksonUnmarshallerInstrumentation.java

@@ -16,7 +16,7 @@
 
 @AutoService(InstrumenterModule.class)
 public class JacksonUnmarshallerInstrumentation extends InstrumenterModule.AppSec
-    implements Instrumenter.ForSingleType {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
 
   public JacksonUnmarshallerInstrumentation() {
     super("akka-http");

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/appsec/MultipartUnmarshallersInstrumentation.java

@@ -14,7 +14,7 @@
 /** @see MultipartUnmarshallers */
 @AutoService(InstrumenterModule.class)
 public class MultipartUnmarshallersInstrumentation extends InstrumenterModule.AppSec
-    implements Instrumenter.ForKnownTypes {
+    implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
 
   private static final String TRAIT_NAME =
       "akka.http.scaladsl.unmarshalling.MultipartUnmarshallers";

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/appsec/PredefinedFromEntityUnmarshallersInstrumentation.java

@@ -20,7 +20,7 @@
  */
 @AutoService(InstrumenterModule.class)
 public class PredefinedFromEntityUnmarshallersInstrumentation extends InstrumenterModule.AppSec
-    implements Instrumenter.ForKnownTypes {
+    implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
 
   private static final String TRAIT_NAME =
       "akka.http.scaladsl.unmarshalling.PredefinedFromEntityUnmarshallers";

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/appsec/SprayUnmarshallerInstrumentation.java

@@ -14,7 +14,7 @@
 // TODO: move to separate module and have better support
 @AutoService(InstrumenterModule.class)
 public class SprayUnmarshallerInstrumentation extends InstrumenterModule.AppSec
-    implements Instrumenter.ForKnownTypes {
+    implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
 
   private static final String TRAIT_NAME =
       "akka.http.scaladsl.marshallers.sprayjson.SprayJsonSupport";

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/appsec/StrictFormCompanionInstrumentation.java

@@ -20,7 +20,7 @@
 /** @see akka.http.scaladsl.common.StrictForm$#unmarshaller(Unmarshaller, Unmarshaller) */
 @AutoService(InstrumenterModule.class)
 public class StrictFormCompanionInstrumentation extends InstrumenterModule.AppSec
-    implements Instrumenter.ForSingleType {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
   public StrictFormCompanionInstrumentation() {
     super("akka-http");
   }

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/CookieDirectivesInstrumentation.java

@@ -22,7 +22,7 @@
  */
 @AutoService(InstrumenterModule.class)
 public class CookieDirectivesInstrumentation extends InstrumenterModule.Iast
-    implements Instrumenter.ForKnownTypes {
+    implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
   public CookieDirectivesInstrumentation() {
     super("akka-http");
   }

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/CookieHeaderInstrumentation.java

@@ -33,7 +33,7 @@
  */
 @AutoService(InstrumenterModule.class)
 public class CookieHeaderInstrumentation extends InstrumenterModule.Iast
-    implements Instrumenter.ForSingleType {
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
   public CookieHeaderInstrumentation() {
     super("akka-http");
   }

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/ExtractDirectivesInstrumentation.java

@@ -34,7 +34,7 @@
  */
 @AutoService(InstrumenterModule.class)
 public class ExtractDirectivesInstrumentation extends InstrumenterModule.Iast
-    implements Instrumenter.ForKnownTypes {
+    implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
   public ExtractDirectivesInstrumentation() {
     super("akka-http");
   }