Introduce marker mechanism for eagerly initializing helpers, and use …

…it to register AsyncResultExtensions (#8028)
DataDog · Nov 28, 2024 · a86c3ca · a86c3ca
1 parent 600378f
commit a86c3ca
Show file tree

Hide file tree

Showing 12 changed files with 55 additions and 54 deletions.
diff --git a/...n/java/datadog/trace/bootstrap/instrumentation/java/concurrent/AsyncResultExtensions.java b/...n/java/datadog/trace/bootstrap/instrumentation/java/concurrent/AsyncResultExtensions.java
@@ -2,6 +2,7 @@
 
 import static java.util.Collections.singletonList;
 
+import datadog.trace.api.Platform;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
@@ -22,6 +23,15 @@ public final class AsyncResultExtensions {
    */
   public static void register(AsyncResultExtension extension) {
     if (extension != null) {
+      if (Platform.isNativeImageBuilder()
+          && extension
+              .getClass()
+              .getClassLoader()
+              .getClass()
+              .getName()
+              .endsWith("ThrowawayClassLoader")) {
+        return; // spring-native expects this to be thrown away, not persisted
+      }
       EXTENSIONS.add(extension);
     }
   }

diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java
@@ -2,6 +2,7 @@
 
 import static datadog.trace.bootstrap.AgentClassLoading.INJECTING_HELPERS;
 
+import datadog.trace.bootstrap.instrumentation.api.EagerHelper;
 import java.io.IOException;
 import java.lang.ref.WeakReference;
 import java.security.CodeSource;
@@ -124,6 +125,18 @@ public DynamicType.Builder<?> transform(
             final JavaModule javaModule = JavaModule.ofType(classes.values().iterator().next());
             helperModules.add(new WeakReference<>(javaModule.unwrap()));
           }
+
+          // forcibly initialize any eager helpers
+          for (Class<?> clazz : classes.values()) {
+            if (EagerHelper.class.isAssignableFrom(clazz)) {
+              try {
+                clazz.getMethod("init").invoke(null);
+              } catch (Throwable e) {
+                log.debug("Problem initializing {}", clazz, e);
+              }
+            }
+          }
+
         } catch (final Exception e) {
           if (log.isErrorEnabled()) {
             log.error(

diff --git a/...og/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java b/...og/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java
@@ -96,7 +96,7 @@ public static void onEnter(@Advice.Argument(value = 0, readOnly = false) String[
               + "datadog.trace.bootstrap.benchmark.StaticEventLogger:build_time,"
               + "datadog.trace.bootstrap.blocking.BlockingExceptionHandler:build_time,"
               + "datadog.trace.bootstrap.InstrumentationErrors:build_time,"
-              + "datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtensions:rerun,"
+              + "datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtensions:build_time,"
               + "datadog.trace.bootstrap.instrumentation.java.concurrent.ConcurrentState:build_time,"
               + "datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter:build_time,"
               + "datadog.trace.bootstrap.instrumentation.java.concurrent.QueueTimeHelper:build_time,"
@@ -105,6 +105,10 @@ public static void onEnter(@Advice.Argument(value = 0, readOnly = false) String[
               + "datadog.trace.bootstrap.instrumentation.jfr.exceptions.ExceptionSampleEvent:build_time,"
               + "datadog.trace.bootstrap.instrumentation.jfr.backpressure.BackpressureSampleEvent:build_time,"
               + "datadog.trace.bootstrap.instrumentation.jfr.directallocation.DirectAllocationTotalEvent:build_time,"
+              + "datadog.trace.instrumentation.guava10.GuavaAsyncResultExtension:build_time,"
+              + "datadog.trace.instrumentation.reactivestreams.ReactiveStreamsAsyncResultExtension:build_time,"
+              + "datadog.trace.instrumentation.reactor.core.ReactorAsyncResultExtension:build_time,"
+              + "datadog.trace.instrumentation.rxjava2.RxJavaAsyncResultExtension:build_time,"
               + "datadog.trace.logging.LoggingSettingsDescription:build_time,"
               + "datadog.trace.logging.simplelogger.SLCompatFactory:build_time,"
               + "datadog.trace.logging.LogReporter:build_time,"

diff --git a/...ava-10/src/main/java/datadog/trace/instrumentation/guava10/GuavaAsyncResultExtension.java b/...ava-10/src/main/java/datadog/trace/instrumentation/guava10/GuavaAsyncResultExtension.java
@@ -2,12 +2,13 @@
 
 import com.google.common.util.concurrent.ListenableFuture;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.EagerHelper;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtension;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtensions;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.ExecutionException;
 
-public class GuavaAsyncResultExtension implements AsyncResultExtension {
+public class GuavaAsyncResultExtension implements AsyncResultExtension, EagerHelper {
   static {
     AsyncResultExtensions.register(new GuavaAsyncResultExtension());
   }
@@ -19,7 +20,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) {

diff --git a/.../src/main/java/datadog/trace/instrumentation/guava10/ListenableFutureInstrumentation.java b/.../src/main/java/datadog/trace/instrumentation/guava10/ListenableFutureInstrumentation.java
@@ -3,7 +3,6 @@
 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;
@@ -47,20 +46,11 @@ public Map<String, String> contextStore() {
 
   @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(

diff --git a/...src/main/java/datadog/trace/instrumentation/reactivestreams/PublisherInstrumentation.java b/...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;
@@ -69,7 +68,6 @@ public String[] helperClassNames() {
 
   @Override
   public void methodAdvice(MethodTransformer transformer) {
-    transformer.applyAdvice(isConstructor(), this.getClass().getName() + "$PublisherAdvice");
     transformer.applyAdvice(
         isMethod()
             .and(not(isStatic()))
@@ -79,13 +77,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(

diff --git a/...va/datadog/trace/instrumentation/reactivestreams/ReactiveStreamsAsyncResultExtension.java b/...va/datadog/trace/instrumentation/reactivestreams/ReactiveStreamsAsyncResultExtension.java
@@ -1,13 +1,14 @@
 package datadog.trace.instrumentation.reactivestreams;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.EagerHelper;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtension;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtensions;
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
 
-public class ReactiveStreamsAsyncResultExtension implements AsyncResultExtension {
+public class ReactiveStreamsAsyncResultExtension implements AsyncResultExtension, EagerHelper {
   static {
     AsyncResultExtensions.register(new ReactiveStreamsAsyncResultExtension());
   }
@@ -19,7 +20,7 @@ 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) {

diff --git a/...ain/java/datadog/trace/instrumentation/reactor/core/BlockingPublisherInstrumentation.java b/...ain/java/datadog/trace/instrumentation/reactor/core/BlockingPublisherInstrumentation.java
@@ -4,7 +4,6 @@
 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;
@@ -40,13 +39,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 +54,12 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
     return hasSuperType(namedOneOf("reactor.core.publisher.Mono", "reactor.core.publisher.Flux"));
   }
 
+  @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 +77,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();
-    }
-  }
 }
diff --git a/...src/main/java/datadog/trace/instrumentation/reactor/core/ReactorAsyncResultExtension.java b/...src/main/java/datadog/trace/instrumentation/reactor/core/ReactorAsyncResultExtension.java
@@ -1,12 +1,13 @@
 package datadog.trace.instrumentation.reactor.core;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.EagerHelper;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtension;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtensions;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
-public class ReactorAsyncResultExtension implements AsyncResultExtension {
+public class ReactorAsyncResultExtension implements AsyncResultExtension, EagerHelper {
   static {
     AsyncResultExtensions.register(new ReactorAsyncResultExtension());
   }
@@ -18,7 +19,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) {

diff --git a/...ava-2/src/main/java/datadog/trace/instrumentation/rxjava2/RxJavaAsyncResultExtension.java b/...ava-2/src/main/java/datadog/trace/instrumentation/rxjava2/RxJavaAsyncResultExtension.java
@@ -1,6 +1,7 @@
 package datadog.trace.instrumentation.rxjava2;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.EagerHelper;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtension;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtensions;
 import io.reactivex.Completable;
@@ -9,7 +10,7 @@
 import io.reactivex.Observable;
 import io.reactivex.Single;
 
-public class RxJavaAsyncResultExtension implements AsyncResultExtension {
+public class RxJavaAsyncResultExtension implements AsyncResultExtension, EagerHelper {
   static {
     AsyncResultExtensions.register(new RxJavaAsyncResultExtension());
   }
@@ -21,7 +22,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) {

diff --git a/...a-2/src/main/java/datadog/trace/instrumentation/rxjava2/RxJavaPluginsInstrumentation.java b/...a-2/src/main/java/datadog/trace/instrumentation/rxjava2/RxJavaPluginsInstrumentation.java
@@ -1,12 +1,9 @@
 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.api.InstrumenterConfig;
-import net.bytebuddy.asm.Advice;
 
 @AutoService(InstrumenterModule.class)
 public class RxJavaPluginsInstrumentation extends InstrumenterModule.Tracing
@@ -36,13 +33,6 @@ public String[] helperClassNames() {
 
   @Override
   public void methodAdvice(MethodTransformer transformer) {
-    transformer.applyAdvice(isMethod(), getClass().getName() + "$RxJavaPluginsAdvice");
-  }
-
-  public static class RxJavaPluginsAdvice {
-    @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void init() {
-      RxJavaAsyncResultExtension.initialize();
-    }
+    // no-op
   }
 }
diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/EagerHelper.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/EagerHelper.java
@@ -0,0 +1,8 @@
+package datadog.trace.bootstrap.instrumentation.api;
+
+/**
+ * Marker interface used to identify helpers that should be eagerly initialized when injected.
+ *
+ * <p>Eager helpers must declare a public static "init" method that takes no arguments.
+ */
+public interface EagerHelper {}