Install AsyncResult extensions during class initialization

Author: mcculls

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

@@ -0,0 +1,84 @@
+package datadog.trace.agent.tooling;
+
+import datadog.trace.util.Strings;
+import net.bytebuddy.asm.AsmVisitorWrapper;
+import net.bytebuddy.description.field.FieldDescription;
+import net.bytebuddy.description.field.FieldList;
+import net.bytebuddy.description.method.MethodList;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.implementation.Implementation;
+import net.bytebuddy.jar.asm.ClassVisitor;
+import net.bytebuddy.jar.asm.MethodVisitor;
+import net.bytebuddy.jar.asm.Opcodes;
+import net.bytebuddy.pool.TypePool;
+
+public final class PreloadingVisitor implements AsmVisitorWrapper {
+ private final String internalPreloadName;
+
+ public PreloadingVisitor(String preloadClass) {
+ this.internalPreloadName = Strings.getInternalName(preloadClass);
+ }
+
+ @Override
+ public int mergeWriter(int flags) {
+ return flags;
+ }
+
+ @Override
+ public int mergeReader(int flags) {
+ return flags;
+ }
+
+ @Override
+ public ClassVisitor wrap(
+ final TypeDescription instrumentedType,
+ final ClassVisitor classVisitor,
+ final Implementation.Context implementationContext,
+ final TypePool typePool,
+ final FieldList<FieldDescription.InDefinedShape> fields,
+ final MethodList<?> methods,
+ final int writerFlags,
+ final int readerFlags) {
+ return new ClassVisitor(Opcodes.ASM8, classVisitor) {
+
+ private static final String TYPE_INITIALIZER_NAME = "<clinit>";
+ private static final String VOID_METHOD_DESCRIPTOR = "()V";
+
+ private boolean hasTypeInitializer;
+
+ @Override
+ public MethodVisitor visitMethod(
+ int access, String name, String descriptor, String signature, String[] exceptions) {
+ MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
+ if (TYPE_INITIALIZER_NAME.equals(name)) {
+ hasTypeInitializer = true;
+ return new MethodVisitor(Opcodes.ASM8, mv) {
+ @Override
+ public void visitCode() {
+ visitMethodInsn(
+ Opcodes.INVOKESTATIC, internalPreloadName, "init", VOID_METHOD_DESCRIPTOR, false);
+ super.visitCode();
+ }
+ };
+ }
+ return mv;
+ }
+
+ @Override
+ public void visitEnd() {
+ if (!hasTypeInitializer) {
+ MethodVisitor mv =
+ visitMethod(
+ Opcodes.ACC_STATIC, TYPE_INITIALIZER_NAME, VOID_METHOD_DESCRIPTOR, null, null);
+ mv.visitCode();
+ mv.visitMethodInsn(
+ Opcodes.INVOKESTATIC, internalPreloadName, "init", VOID_METHOD_DESCRIPTOR, false);
+ mv.visitInsn(Opcodes.RETURN);
+ mv.visitMaxs(0, 0);
+ mv.visitEnd();
+ }
+ super.visitEnd();
+ }
+ };
+ }
+}

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

@@ -19,7 +19,7 @@ public class GuavaAsyncResultExtension implements AsyncResultExtension {
  * class initialization. This will ensure this extension will only be registered once under {@link
  * AsyncResultExtensions}.
  */
- public static void initialize() {}
+ public static void init() {}
 
  @Override
  public boolean supports(Class<?> result) {

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

@@ -3,13 +3,13 @@
 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.takesArguments;
 
 import com.google.auto.service.AutoService;
 import com.google.common.util.concurrent.AbstractFuture;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.agent.tooling.PreloadingVisitor;
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
@@ -22,7 +22,7 @@
 
 @AutoService(InstrumenterModule.class)
 public class ListenableFutureInstrumentation extends InstrumenterModule.Tracing
- implements Instrumenter.ForSingleType {
+ implements Instrumenter.ForSingleType, Instrumenter.HasTypeAdvice {
 
  public ListenableFutureInstrumentation() {
  super("guava");
@@ -45,22 +45,18 @@ public Map<String, String> contextStore() {
  return singletonMap(Runnable.class.getName(), State.class.getName());
  }
 
+ @Override
+ public void typeAdvice(TypeTransformer transformer) {
+ transformer.applyAdvice(new PreloadingVisitor(packageName + ".GuavaAsyncResultExtension"));
+ }
+
  @Override
  public void methodAdvice(MethodTransformer transformer) {
- transformer.applyAdvice(
- isConstructor(), ListenableFutureInstrumentation.class.getName() + "$AbstractFutureAdvice");
  transformer.applyAdvice(
  named("addListener").and(takesArguments(Runnable.class, Executor.class)),
  ListenableFutureInstrumentation.class.getName() + "$AddListenerAdvice");
  }
 
- public static class AbstractFutureAdvice {
- @Advice.OnMethodExit(suppress = Throwable.class)
- public static void init() {
- GuavaAsyncResultExtension.initialize();
- }
- }
-
  public static class AddListenerAdvice {
  @Advice.OnMethodEnter(suppress = Throwable.class)
  public static State addListenerEnter(

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

@@ -5,7 +5,6 @@
 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.not;
@@ -15,6 +14,7 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.agent.tooling.PreloadingVisitor;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -33,7 +33,7 @@
  */
 @AutoService(InstrumenterModule.class)
 public class PublisherInstrumentation extends InstrumenterModule.Tracing
- implements Instrumenter.ForTypeHierarchy {
+ implements Instrumenter.ForTypeHierarchy, Instrumenter.HasTypeAdvice {
 
  public PublisherInstrumentation() {
  super("reactive-streams", "reactive-streams-1");
@@ -67,9 +67,14 @@ public String[] helperClassNames() {
  };
  }
 
+ @Override
+ public void typeAdvice(TypeTransformer transformer) {
+ transformer.applyAdvice(
+ new PreloadingVisitor(packageName + ".ReactiveStreamsAsyncResultExtension"));
+ }
+
  @Override
  public void methodAdvice(MethodTransformer transformer) {
- transformer.applyAdvice(isConstructor(), this.getClass().getName() + "$PublisherAdvice");
  transformer.applyAdvice(
  isMethod()
  .and(not(isStatic()))
@@ -79,13 +84,6 @@ public void methodAdvice(MethodTransformer transformer) {
  getClass().getName() + "$PublisherSubscribeAdvice");
  }
 
- public static class PublisherAdvice {
- @Advice.OnMethodExit(suppress = Throwable.class)
- public static void init() {
- ReactiveStreamsAsyncResultExtension.initialize();
- }
- }
-
  public static class PublisherSubscribeAdvice {
  @Advice.OnMethodEnter(suppress = Throwable.class)
  public static AgentScope onSubscribe(

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

@@ -19,12 +19,11 @@ public class ReactiveStreamsAsyncResultExtension implements AsyncResultExtension
  * class initialization. This will ensure this extension will only be registered once under {@link
  * AsyncResultExtensions}.
  */
- public static void initialize() {}
+ public static void init() {}
 
  @Override
  public boolean supports(Class<?> result) {
- boolean ret = result == Publisher.class;
- return ret;
+ return result == Publisher.class;
  }
 
  @Override

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

@@ -4,12 +4,12 @@
 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 com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.agent.tooling.PreloadingVisitor;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -28,7 +28,7 @@
  */
 @AutoService(InstrumenterModule.class)
 public class BlockingPublisherInstrumentation extends InstrumenterModule.Tracing
- implements Instrumenter.ForTypeHierarchy {
+ implements Instrumenter.ForTypeHierarchy, Instrumenter.HasTypeAdvice {
  public BlockingPublisherInstrumentation() {
  super("reactor-core");
  }
@@ -40,13 +40,6 @@ public String[] helperClassNames() {
  };
  }
 
- @Override
- public void methodAdvice(MethodTransformer transformer) {
- transformer.applyAdvice(isConstructor(), getClass().getName() + "$AsyncExtensionInstallAdvice");
- transformer.applyAdvice(
- isMethod().and(nameStartsWith("block")), getClass().getName() + "$BlockingAdvice");
- }
-
  @Override
  public Map<String, String> contextStore() {
  return Collections.singletonMap("org.reactivestreams.Publisher", AgentSpan.class.getName());
@@ -62,6 +55,17 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
  return hasSuperType(namedOneOf("reactor.core.publisher.Mono", "reactor.core.publisher.Flux"));
  }
 
+ @Override
+ public void typeAdvice(TypeTransformer transformer) {
+ transformer.applyAdvice(new PreloadingVisitor(packageName + ".ReactorAsyncResultExtension"));
+ }
+
+ @Override
+ public void methodAdvice(MethodTransformer transformer) {
+ transformer.applyAdvice(
+ isMethod().and(nameStartsWith("block")), getClass().getName() + "$BlockingAdvice");
+ }
+
  public static class BlockingAdvice {
  @Advice.OnMethodEnter(suppress = Throwable.class)
  public static AgentScope before(@Advice.This final Publisher self) {
@@ -79,11 +83,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/reactor-core-3.1/src/main/java/datadog/trace/instrumentation/reactor/core/ReactorAsyncResultExtension.java

@@ -18,7 +18,7 @@ public class ReactorAsyncResultExtension implements AsyncResultExtension {
  * class initialization. This will ensure this extension will only be registered once under {@link
  * AsyncResultExtensions}.
  */
- public static void initialize() {}
+ public static void init() {}
 
  @Override
  public boolean supports(Class<?> result) {

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

@@ -21,7 +21,7 @@ public class RxJavaAsyncResultExtension implements AsyncResultExtension {
  * class initialization. This will ensure this extension will only be registered once under {@link
  * AsyncResultExtensions}.
  */
- public static void initialize() {}
+ public static void init() {}
 
  @Override
  public boolean supports(Class<?> result) {

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

@@ -1,16 +1,14 @@
 package datadog.trace.instrumentation.rxjava2;
 
-import static net.bytebuddy.matcher.ElementMatchers.isMethod;
-
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.agent.tooling.PreloadingVisitor;
 import datadog.trace.api.InstrumenterConfig;
-import net.bytebuddy.asm.Advice;
 
 @AutoService(InstrumenterModule.class)
 public class RxJavaPluginsInstrumentation extends InstrumenterModule.Tracing
- implements Instrumenter.ForSingleType {
+ implements Instrumenter.ForSingleType, Instrumenter.HasTypeAdvice {
 
  public RxJavaPluginsInstrumentation() {
  super("rxjava");
@@ -35,14 +33,12 @@ public String[] helperClassNames() {
  }
 
  @Override
- public void methodAdvice(MethodTransformer transformer) {
- transformer.applyAdvice(isMethod(), getClass().getName() + "$RxJavaPluginsAdvice");
+ public void typeAdvice(TypeTransformer transformer) {
+ transformer.applyAdvice(new PreloadingVisitor(packageName + ".RxJavaAsyncResultExtension"));
  }
 
- public static class RxJavaPluginsAdvice {
- @Advice.OnMethodExit(suppress = Throwable.class)
- public static void init() {
- RxJavaAsyncResultExtension.initialize();
- }
+ @Override
+ public void methodAdvice(MethodTransformer transformer) {
+ // no-op
  }
 }