Merge pull request #2680 from DataDog/tyler/instrumenter

Complete transformer inversion refactor
DataDog · May 4, 2021 · d4a8983 · d4a8983
2 parents 08939d5 + d8d4c45
commit d4a8983
Show file tree

Hide file tree

Showing 264 changed files with 654 additions and 1,456 deletions.
diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java
@@ -294,19 +294,7 @@ public AgentBuilder.Transformer transformer() {
      * Instrumenters should register each advice transformation by calling {@link
      * AdviceTransformation#applyAdvice(ElementMatcher, String)} one or more times.
      */
-    // TODO: once everything is fully migrated, make this method abstract.
-    public void adviceTransformations(AdviceTransformation transformation) {
-      for (final Map.Entry<? extends ElementMatcher<? super MethodDescription>, String> entry :
-          transformers().entrySet()) {
-        transformation.applyAdvice(entry.getKey(), entry.getValue());
-      }
-    }
-
-    @Deprecated
-    protected Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-      throw new UnsupportedOperationException(
-          "Either transformers or adviceTransformations must be overridden.");
-    }
+    public abstract void adviceTransformations(AdviceTransformation transformation);
 
     /**
      * Context stores to define for this instrumentation. Are added to matching class loaders.

diff --git a/...ent-tooling/src/test/groovy/datadog/trace/agent/test/DefaultInstrumenterForkedTest.groovy b/...ent-tooling/src/test/groovy/datadog/trace/agent/test/DefaultInstrumenterForkedTest.groovy
@@ -168,8 +168,7 @@ class DefaultInstrumenterForkedTest extends DDSpecification {
     }
 
     @Override
-    Map<ElementMatcher, String> transformers() {
-      return Collections.emptyMap()
+    void adviceTransformations(AdviceTransformation transformation) {
     }
   }
 }
diff --git a/...-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/TestInstrumentationClasses.java b/...-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/TestInstrumentationClasses.java
@@ -1,11 +1,9 @@
 package datadog.trace.agent.tooling.muzzle;
 
-import static java.util.Collections.emptyMap;
 import static net.bytebuddy.matcher.ElementMatchers.none;
 
 import datadog.trace.agent.tooling.Instrumenter;
 import java.util.Map;
-import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -31,9 +29,7 @@ public ElementMatcher<? super TypeDescription> typeMatcher() {
     }
 
     @Override
-    public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-      return emptyMap();
-    }
+    public void adviceTransformations(AdviceTransformation transformation) {}
   }
 
   public static class EmptyInst extends BaseInst {

diff --git a/...rc/main/java/datadog/trace/instrumentation/aerospike4/AerospikeClientInstrumentation.java b/...rc/main/java/datadog/trace/instrumentation/aerospike4/AerospikeClientInstrumentation.java
@@ -13,10 +13,7 @@
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-import java.util.HashMap;
-import java.util.Map;
 import net.bytebuddy.asm.Advice;
-import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -39,19 +36,17 @@ public String[] helperClassNames() {
   }
 
   @Override
-  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-    final Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
-    transformers.put(
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
         isMethod()
             .and(isPublic())
             .and(takesArgument(0, nameStartsWith("com.aerospike.client.policy"))),
         getClass().getName() + "$TraceSyncRequestAdvice");
-    transformers.put(
+    transformation.applyAdvice(
         isMethod()
             .and(isPublic())
             .and(takesArgument(1, nameStartsWith("com.aerospike.client.listener"))),
         getClass().getName() + "$TraceAsyncRequestAdvice");
-    return transformers;
   }
 
   public static final class TraceSyncRequestAdvice {

diff --git a/...pike-4/src/main/java/datadog/trace/instrumentation/aerospike4/CommandInstrumentation.java b/...pike-4/src/main/java/datadog/trace/instrumentation/aerospike4/CommandInstrumentation.java
@@ -3,7 +3,6 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
 import static datadog.trace.instrumentation.aerospike4.AerospikeClientDecorator.DECORATE;
-import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.returns;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -15,9 +14,7 @@
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.api.DDSpanTypes;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-import java.util.Map;
 import net.bytebuddy.asm.Advice;
-import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -40,8 +37,8 @@ public String[] helperClassNames() {
   }
 
   @Override
-  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-    return singletonMap(
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
         isMethod()
             .and(named("getNode"))
             .and(takesArgument(0, named("com.aerospike.client.cluster.Cluster")))

diff --git a/...4/src/main/java/datadog/trace/instrumentation/aerospike4/NioEventLoopInstrumentation.java b/...4/src/main/java/datadog/trace/instrumentation/aerospike4/NioEventLoopInstrumentation.java
@@ -3,18 +3,15 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.exclude;
-import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
-import java.util.Map;
 import java.util.concurrent.FutureTask;
 import java.util.concurrent.RunnableFuture;
 import net.bytebuddy.asm.Advice;
-import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -30,8 +27,8 @@ public ElementMatcher<TypeDescription> typeMatcher() {
   }
 
   @Override
-  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-    return singletonMap(
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
         isMethod()
             .and(named("execute"))
             .and(takesArguments(1))

diff --git a/...ke-4/src/main/java/datadog/trace/instrumentation/aerospike4/PartitionInstrumentation.java b/...ke-4/src/main/java/datadog/trace/instrumentation/aerospike4/PartitionInstrumentation.java
@@ -4,7 +4,6 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
 import static datadog.trace.instrumentation.aerospike4.AerospikeClientDecorator.DECORATE;
-import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.returns;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -16,9 +15,7 @@
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.api.DDSpanTypes;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-import java.util.Map;
 import net.bytebuddy.asm.Advice;
-import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -41,8 +38,8 @@ public String[] helperClassNames() {
   }
 
   @Override
-  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-    return singletonMap(
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
         isMethod()
             .and(namedOneOf("getNodeRead", "getNodeWrite"))
             .and(takesArgument(0, named("com.aerospike.client.cluster.Cluster")))

diff --git a/...main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java b/...main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java
@@ -17,7 +17,6 @@
 import datadog.trace.context.TraceScope;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
-import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -39,8 +38,9 @@ public Map<String, String> contextStore() {
   }
 
   @Override
-  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-    return singletonMap(isMethod().and(named("invoke")), getClass().getName() + "$InvokeAdvice");
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
+        isMethod().and(named("invoke")), getClass().getName() + "$InvokeAdvice");
   }
 
   /**

diff --git a/.../main/java/datadog/trace/instrumentation/akka/concurrent/AkkaEnvelopeInstrumentation.java b/.../main/java/datadog/trace/instrumentation/akka/concurrent/AkkaEnvelopeInstrumentation.java
@@ -13,7 +13,6 @@
 import datadog.trace.context.TraceScope;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
-import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -35,8 +34,8 @@ public Map<String, String> contextStore() {
   }
 
   @Override
-  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-    return singletonMap(isConstructor(), getClass().getName() + "$ConstructAdvice");
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(isConstructor(), getClass().getName() + "$ConstructAdvice");
   }
 
   public static class ConstructAdvice {

diff --git a/...atadog/trace/instrumentation/akka/concurrent/AkkaForkJoinExecutorTaskInstrumentation.java b/...atadog/trace/instrumentation/akka/concurrent/AkkaForkJoinExecutorTaskInstrumentation.java
@@ -5,7 +5,6 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.endTaskScope;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.startTaskScope;
-import static java.util.Collections.unmodifiableMap;
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -16,10 +15,8 @@
 import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
 import datadog.trace.context.TraceScope;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
-import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -53,13 +50,11 @@ public ElementMatcher<? super TypeDescription> typeMatcher() {
   }
 
   @Override
-  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-    Map<ElementMatcher<MethodDescription>, String> transformers = new HashMap<>(4);
-    transformers.put(
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
         isConstructor().and(takesArgument(0, named(Runnable.class.getName()))),
         getClass().getName() + "$Construct");
-    transformers.put(isMethod().and(named("run")), getClass().getName() + "$Run");
-    return unmodifiableMap(transformers);
+    transformation.applyAdvice(isMethod().and(named("run")), getClass().getName() + "$Run");
   }
 
   public static final class Construct {

diff --git a/...n/java/datadog/trace/instrumentation/akka/concurrent/AkkaForkJoinPoolInstrumentation.java b/...n/java/datadog/trace/instrumentation/akka/concurrent/AkkaForkJoinPoolInstrumentation.java
@@ -13,10 +13,8 @@
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
 import datadog.trace.context.TraceScope;
-import java.util.Collections;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
-import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -38,8 +36,8 @@ public Map<String, String> contextStore() {
   }
 
   @Override
-  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-    return Collections.singletonMap(
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
         isMethod().and(namedOneOf("externalPush", "fullExternalPush")),
         getClass().getName() + "$ExternalPush");
   }

diff --git a/...n/java/datadog/trace/instrumentation/akka/concurrent/AkkaForkJoinTaskInstrumentation.java b/...n/java/datadog/trace/instrumentation/akka/concurrent/AkkaForkJoinTaskInstrumentation.java
@@ -20,10 +20,8 @@
 import datadog.trace.context.TraceScope;
 import java.util.Arrays;
 import java.util.Collection;
-import java.util.HashMap;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
-import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import net.bytebuddy.matcher.ElementMatchers;
@@ -65,12 +63,11 @@ public ElementMatcher<? super TypeDescription> typeMatcher() {
   }
 
   @Override
-  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-    Map<ElementMatcher<MethodDescription>, String> transformers = new HashMap<>(4);
-    transformers.put(isMethod().and(namedOneOf("doExec", "exec")), getClass().getName() + "$Exec");
-    transformers.put(isMethod().and(named("fork")), getClass().getName() + "$Fork");
-    transformers.put(isMethod().and(named("cancel")), getClass().getName() + "$Cancel");
-    return transformers;
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
+        isMethod().and(namedOneOf("doExec", "exec")), getClass().getName() + "$Exec");
+    transformation.applyAdvice(isMethod().and(named("fork")), getClass().getName() + "$Fork");
+    transformation.applyAdvice(isMethod().and(named("cancel")), getClass().getName() + "$Cancel");
   }
 
   public static final class Exec {

diff --git a/...c/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java b/...c/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java
@@ -4,7 +4,6 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
 import static java.util.Collections.singletonList;
-import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 
 import com.google.auto.service.AutoService;
@@ -20,7 +19,6 @@
 import java.util.List;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
-import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -38,8 +36,8 @@ public ElementMatcher<? super TypeDescription> typeMatcher() {
   }
 
   @Override
-  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-    return singletonMap(
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
         isMethod().and(named("run")), getClass().getName() + "$SuppressMailboxRunAdvice");
   }
 

diff --git a/...ava/datadog/trace/instrumentation/akka/concurrent/AkkaRoutedActorCellInstrumentation.java b/...ava/datadog/trace/instrumentation/akka/concurrent/AkkaRoutedActorCellInstrumentation.java
@@ -14,7 +14,6 @@
 import datadog.trace.context.TraceScope;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
-import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import net.bytebuddy.matcher.ElementMatchers;
@@ -37,8 +36,8 @@ public Map<String, String> contextStore() {
   }
 
   @Override
-  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-    return singletonMap(
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
         isMethod()
             .and(
                 named("sendMessage")

diff --git a/.../src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttp2ServerInstrumentation.java b/.../src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttp2ServerInstrumentation.java
@@ -9,11 +9,7 @@
 import akka.stream.Materializer;
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.Map;
 import net.bytebuddy.asm.Advice;
-import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import scala.Function1;
@@ -48,27 +44,25 @@ public String[] helperClassNames() {
   }
 
   @Override
-  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-    Map<ElementMatcher<MethodDescription>, String> transformers = new HashMap<>(4);
-    transformers.put(
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
         takesArguments(8)
             .and(named("bindAndHandleAsync"))
             .and(takesArgument(0, named("scala.Function1")))
             .and(takesArgument(7, named("akka.stream.Materializer"))),
         getClass().getName() + "$Http2BindAndHandleAsync8ArgAdvice");
-    transformers.put(
+    transformation.applyAdvice(
         takesArguments(7)
             .and(named("bindAndHandleAsync"))
             .and(takesArgument(0, named("scala.Function1")))
             .and(takesArgument(6, named("akka.stream.Materializer"))),
         getClass().getName() + "$Http2BindAndHandleAsync7ArgAdvice");
-    transformers.put(
+    transformation.applyAdvice(
         takesArguments(6)
             .and(named("bindAndHandleAsync"))
             .and(takesArgument(0, named("scala.Function1")))
             .and(takesArgument(5, named("akka.stream.Materializer"))),
         getClass().getName() + "$Http2BindAndHandleAsync6ArgAdvice");
-    return Collections.unmodifiableMap(transformers);
   }
 
   public static class Http2BindAndHandleAsync8ArgAdvice {