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

Author: mcculls

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");