Introduce marker mechanism for eagerly initializing helpers, and use it to register AsyncResultExtensions

Author: mcculls

dd-java-agent/agent-bootstrap/src/main/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);
  }
  }

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(

dd-java-agent/instrumentation/graal/native-image/src/main/java/datadog/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,"

dd-java-agent/instrumentation/guava-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) {

dd-java-agent/instrumentation/guava-10/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(

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;
@@ -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(

dd-java-agent/instrumentation/reactive-streams/src/main/java/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) {

dd-java-agent/instrumentation/reactor-core-3.1/src/main/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();
- }
- }
 }

dd-java-agent/instrumentation/reactor-core-3.1/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) {

dd-java-agent/instrumentation/rxjava-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) {