Make Akka's ForkJoin pool instrumented class names configurable (#8766)

Author: nikita-tkachenko-datadog

dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaForkJoinExecutorTaskInstrumentation.java

@@ -11,6 +11,7 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
@@ -25,7 +26,9 @@
  */
 @AutoService(InstrumenterModule.class)
 public final class AkkaForkJoinExecutorTaskInstrumentation extends InstrumenterModule.Tracing
-    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+    implements Instrumenter.ForSingleType,
+        Instrumenter.ForConfiguredType,
+        Instrumenter.HasMethodAdvice {
   public AkkaForkJoinExecutorTaskInstrumentation() {
     super("java_concurrent", "akka_concurrent");
   }
@@ -40,6 +43,11 @@ public String instrumentedType() {
     return "akka.dispatch.ForkJoinExecutorConfigurator$AkkaForkJoinTask";
   }
 
+  @Override
+  public String configuredMatchingType() {
+    return InstrumenterConfig.get().getAkkaForkJoinExecutorTaskName();
+  }
+
   @Override
   public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(

dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaForkJoinPoolInstrumentation.java

@@ -12,14 +12,17 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
 
 @AutoService(InstrumenterModule.class)
 public final class AkkaForkJoinPoolInstrumentation extends InstrumenterModule.Tracing
-    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+    implements Instrumenter.ForSingleType,
+        Instrumenter.ForConfiguredType,
+        Instrumenter.HasMethodAdvice {
 
   public AkkaForkJoinPoolInstrumentation() {
     super("java_concurrent", "akka_concurrent");
@@ -30,6 +33,11 @@ public String instrumentedType() {
     return "akka.dispatch.forkjoin.ForkJoinPool";
   }
 
+  @Override
+  public String configuredMatchingType() {
+    return InstrumenterConfig.get().getAkkaForkJoinPoolName();
+  }
+
   @Override
   public Map<String, String> contextStore() {
     return singletonMap("akka.dispatch.forkjoin.ForkJoinTask", State.class.getName());
@@ -43,6 +51,7 @@ public void methodAdvice(MethodTransformer transformer) {
   }
 
   public static final class ExternalPush {
+
     @Advice.OnMethodEnter
     public static <T> void externalPush(@Advice.Argument(0) ForkJoinTask<T> task) {
       if (!exclude(FORK_JOIN_TASK, task)) {

dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaForkJoinTaskInstrumentation.java

@@ -19,6 +19,7 @@
 import datadog.trace.agent.tooling.ExcludeFilterProvider;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
@@ -52,14 +53,26 @@ public Map<String, String> contextStore() {
 
   @Override
   public String hierarchyMarkerType() {
-    return "akka.dispatch.forkjoin.ForkJoinTask";
+    String akkaForkJoinTaskName = InstrumenterConfig.get().getAkkaForkJoinTaskName();
+    return akkaForkJoinTaskName != null && !akkaForkJoinTaskName.isEmpty()
+        ? null // bypass the hint if custom class is configured
+        : "akka.dispatch.forkjoin.ForkJoinTask";
   }
 
   @Override
   public ElementMatcher<TypeDescription> hierarchyMatcher() {
     return notExcludedByName(FORK_JOIN_TASK)
         .and(declaresMethod(namedOneOf("exec", "fork", "cancel")))
-        .and(extendsClass(named(hierarchyMarkerType())));
+        .and(isForkJoinTaskSubclass());
+  }
+
+  private ElementMatcher<TypeDescription> isForkJoinTaskSubclass() {
+    ElementMatcher.Junction<TypeDescription> forkJoinTaskSubclass =
+        extendsClass(named("akka.dispatch.forkjoin.ForkJoinTask"));
+    String akkaForkJoinTaskName = InstrumenterConfig.get().getAkkaForkJoinTaskName();
+    return akkaForkJoinTaskName != null && !akkaForkJoinTaskName.isEmpty()
+        ? forkJoinTaskSubclass.or(extendsClass(named(akkaForkJoinTaskName)))
+        : forkJoinTaskSubclass;
   }
 
   @Override

dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java

@@ -72,6 +72,11 @@ public final class TraceInstrumentationConfig {
 
   public static final String JDBC_CONNECTION_CLASS_NAME = "trace.jdbc.connection.class.name";
 
+  public static final String AKKA_FORK_JOIN_TASK_NAME = "trace.akka.fork.join.task.name";
+  public static final String AKKA_FORK_JOIN_EXECUTOR_TASK_NAME =
+      "trace.akka.fork.join.executor.task.name";
+  public static final String AKKA_FORK_JOIN_POOL_NAME = "trace.akka.fork.join.pool.name";
+
   public static final String EXPERIMENTATAL_JEE_SPLIT_BY_DEPLOYMENT =
       "trace.experimental.jee.split-by-deployment";
 

internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java

@@ -32,6 +32,9 @@
 import static datadog.trace.api.config.ProfilingConfig.PROFILING_DIRECT_ALLOCATION_ENABLED_DEFAULT;
 import static datadog.trace.api.config.ProfilingConfig.PROFILING_ENABLED;
 import static datadog.trace.api.config.ProfilingConfig.PROFILING_ENABLED_DEFAULT;
+import static datadog.trace.api.config.TraceInstrumentationConfig.AKKA_FORK_JOIN_EXECUTOR_TASK_NAME;
+import static datadog.trace.api.config.TraceInstrumentationConfig.AKKA_FORK_JOIN_POOL_NAME;
+import static datadog.trace.api.config.TraceInstrumentationConfig.AKKA_FORK_JOIN_TASK_NAME;
 import static datadog.trace.api.config.TraceInstrumentationConfig.AXIS_TRANSPORT_CLASS_NAME;
 import static datadog.trace.api.config.TraceInstrumentationConfig.CODE_ORIGIN_FOR_SPANS_ENABLED;
 import static datadog.trace.api.config.TraceInstrumentationConfig.EXPERIMENTAL_DEFER_INTEGRATIONS_UNTIL;
@@ -129,6 +132,10 @@ public class InstrumenterConfig {
   private final boolean websocketTracingEnabled;
   private final boolean pekkoSchedulerEnabled;
 
+  private final String akkaForkJoinTaskName;
+  private final String akkaForkJoinExecutorTaskName;
+  private final String akkaForkJoinPoolName;
+
   private final boolean directAllocationProfilingEnabled;
 
   private final String instrumentationConfigId;
@@ -228,6 +235,10 @@ private InstrumenterConfig() {
     httpURLConnectionClassName = configProvider.getString(HTTP_URL_CONNECTION_CLASS_NAME, "");
     axisTransportClassName = configProvider.getString(AXIS_TRANSPORT_CLASS_NAME, "");
 
+    akkaForkJoinTaskName = configProvider.getString(AKKA_FORK_JOIN_TASK_NAME, "");
+    akkaForkJoinExecutorTaskName = configProvider.getString(AKKA_FORK_JOIN_EXECUTOR_TASK_NAME, "");
+    akkaForkJoinPoolName = configProvider.getString(AKKA_FORK_JOIN_POOL_NAME, "");
+
     directAllocationProfilingEnabled =
         configProvider.getBoolean(
             PROFILING_DIRECT_ALLOCATION_ENABLED, PROFILING_DIRECT_ALLOCATION_ENABLED_DEFAULT);
@@ -404,6 +415,18 @@ public String getAxisTransportClassName() {
     return axisTransportClassName;
   }
 
+  public String getAkkaForkJoinTaskName() {
+    return akkaForkJoinTaskName;
+  }
+
+  public String getAkkaForkJoinExecutorTaskName() {
+    return akkaForkJoinExecutorTaskName;
+  }
+
+  public String getAkkaForkJoinPoolName() {
+    return akkaForkJoinPoolName;
+  }
+
   public boolean isDirectAllocationProfilingEnabled() {
     return directAllocationProfilingEnabled;
   }