Install AsyncResult extensions during class initialization

Author: mcculls

dd-java-agent/instrumentation/guava-10/src/main/java/datadog/trace/instrumentation/guava10/ListenableFutureInstrumentation.java

@@ -3,7 +3,7 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
 import static java.util.Collections.singletonMap;
-import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.isTypeInitializer;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
@@ -48,13 +48,14 @@ public Map<String, String> contextStore() {
   @Override
   public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(
-        isConstructor(), ListenableFutureInstrumentation.class.getName() + "$AbstractFutureAdvice");
+        isTypeInitializer(),
+        ListenableFutureInstrumentation.class.getName() + "$InstallExtensionAdvice");
     transformer.applyAdvice(
         named("addListener").and(takesArguments(Runnable.class, Executor.class)),
         ListenableFutureInstrumentation.class.getName() + "$AddListenerAdvice");
   }
 
-  public static class AbstractFutureAdvice {
+  public static class InstallExtensionAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void init() {
       GuavaAsyncResultExtension.initialize();

dd-java-agent/instrumentation/reactive-streams/src/main/java/datadog/trace/instrumentation/reactivestreams/PublisherInstrumentation.java

@@ -5,9 +5,9 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
-import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isStatic;
+import static net.bytebuddy.matcher.ElementMatchers.isTypeInitializer;
 import static net.bytebuddy.matcher.ElementMatchers.not;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
@@ -69,7 +69,8 @@ public String[] helperClassNames() {
 
   @Override
   public void methodAdvice(MethodTransformer transformer) {
-    transformer.applyAdvice(isConstructor(), this.getClass().getName() + "$PublisherAdvice");
+    transformer.applyAdvice(
+        isTypeInitializer(), this.getClass().getName() + "$InstallExtensionAdvice");
     transformer.applyAdvice(
         isMethod()
             .and(not(isStatic()))
@@ -79,7 +80,7 @@ public void methodAdvice(MethodTransformer transformer) {
         getClass().getName() + "$PublisherSubscribeAdvice");
   }
 
-  public static class PublisherAdvice {
+  public static class InstallExtensionAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void init() {
       ReactiveStreamsAsyncResultExtension.initialize();

dd-java-agent/instrumentation/reactor-core-3.1/src/main/java/datadog/trace/instrumentation/reactor/core/BlockingPublisherInstrumentation.java

@@ -4,8 +4,8 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.nameStartsWith;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
-import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isTypeInitializer;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
@@ -42,7 +42,7 @@ public String[] helperClassNames() {
 
   @Override
   public void methodAdvice(MethodTransformer transformer) {
-    transformer.applyAdvice(isConstructor(), getClass().getName() + "$AsyncExtensionInstallAdvice");
+    transformer.applyAdvice(isTypeInitializer(), getClass().getName() + "$InstallExtensionAdvice");
     transformer.applyAdvice(
         isMethod().and(nameStartsWith("block")), getClass().getName() + "$BlockingAdvice");
   }
@@ -62,6 +62,13 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
     return hasSuperType(namedOneOf("reactor.core.publisher.Mono", "reactor.core.publisher.Flux"));
   }
 
+  public static class InstallExtensionAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void init() {
+      ReactorAsyncResultExtension.initialize();
+    }
+  }
+
   public static class BlockingAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static AgentScope before(@Advice.This final Publisher self) {
@@ -79,11 +86,4 @@ public static void after(@Advice.Enter final AgentScope scope) {
       }
     }
   }
-
-  public static class AsyncExtensionInstallAdvice {
-    @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void init() {
-      ReactorAsyncResultExtension.initialize();
-    }
-  }
 }

dd-java-agent/instrumentation/rxjava-2/src/main/java/datadog/trace/instrumentation/rxjava2/RxJavaPluginsInstrumentation.java

@@ -1,6 +1,6 @@
 package datadog.trace.instrumentation.rxjava2;
 
-import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isTypeInitializer;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
@@ -36,10 +36,10 @@ public String[] helperClassNames() {
 
   @Override
   public void methodAdvice(MethodTransformer transformer) {
-    transformer.applyAdvice(isMethod(), getClass().getName() + "$RxJavaPluginsAdvice");
+    transformer.applyAdvice(isTypeInitializer(), getClass().getName() + "$InstallExtensionAdvice");
   }
 
-  public static class RxJavaPluginsAdvice {
+  public static class InstallExtensionAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void init() {
       RxJavaAsyncResultExtension.initialize();