InstrumenterModule refactoring for java.concurrent module

dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/completablefuture/AsyncTaskInstrumentation.java

@@ -8,17 +8,14 @@
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE;
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 
-import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.ExcludeFilterProvider;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
 import java.util.Arrays;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.EnumMap;
 import java.util.Map;
 import java.util.concurrent.ForkJoinTask;
@@ -30,8 +27,7 @@
  * either context. This double instrumentation otherwise leads to excess scope creation and
  * duplicate checkpoint emission.
  */
-@AutoService(InstrumenterModule.class)
-public final class AsyncTaskInstrumentation extends InstrumenterModule.Tracing
+public final class AsyncTaskInstrumentation
     implements Instrumenter.ForBootstrap,
         Instrumenter.ForKnownTypes,
         Instrumenter.HasMethodAdvice,
@@ -42,20 +38,11 @@ public final class AsyncTaskInstrumentation extends InstrumenterModule.Tracing
     "java.util.concurrent.CompletableFuture$AsyncRun",
   };
 
-  public AsyncTaskInstrumentation() {
-    super("java_completablefuture", "java_concurrent");
-  }
-
   @Override
   public String[] knownMatchingTypes() {
     return CLASS_NAMES;
   }
 
-  @Override
-  public Map<String, String> contextStore() {
-    return Collections.singletonMap("java.util.concurrent.ForkJoinTask", State.class.getName());
-  }
-
   @Override
   public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(isConstructor(), getClass().getName() + "$Construct");

dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/completablefuture/CompletableFutureModule.java

@@ -0,0 +1,37 @@
+package datadog.trace.instrumentation.java.completablefuture;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.instrumentation.java.concurrent.ConcurrentState;
+import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/** Groups the instrumentations for CompletableFuture and related classes. */
+@AutoService(InstrumenterModule.class)
+public final class CompletableFutureModule extends InstrumenterModule.Tracing {
+
+  public CompletableFutureModule() {
+    super("java_completablefuture", "java_concurrent");
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    Map<String, String> stores = new HashMap<>();
+    stores.put("java.util.concurrent.ForkJoinTask", State.class.getName());
+    stores.put(
+        "java.util.concurrent.CompletableFuture$UniCompletion", ConcurrentState.class.getName());
+    return stores;
+  }
+
+  @Override
+  public List<Instrumenter> typeInstrumentations() {
+    return Arrays.asList(
+        new AsyncTaskInstrumentation(),
+        new CompletableFutureUniCompletionInstrumentation(),
+        new CompletableFutureUniCompletionSubclassInstrumentation());
+  }
+}

dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/completablefuture/CompletableFutureUniCompletionInstrumentation.java

@@ -3,18 +3,14 @@
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.EXECUTOR;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.FORK_JOIN_TASK;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE;
-import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 
-import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.ExcludeFilterProvider;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.ConcurrentState;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType;
 import java.util.Arrays;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.EnumMap;
 import java.util.List;
 import java.util.Map;
@@ -30,7 +26,7 @@
  * f.thenApplyAsync({ r -> ... })
  * }</pre>
  *
- * The general idea is to capture the current scope at the time of creation of the {@code
+ * <p>The general idea is to capture the current scope at the time of creation of the {@code
  * UniCompletion} and then activate that scope around the processing that happens in the {@code
  * tryFire} method.
  *
@@ -45,8 +41,7 @@
  * continuation {@code ConcurrentContinuation}, have been introduced to deal with the benign race
  * taking place that decides which thread actually get to run the user code that was supplied.
  */
-@AutoService(InstrumenterModule.class)
-public class CompletableFutureUniCompletionInstrumentation extends InstrumenterModule.Tracing
+public class CompletableFutureUniCompletionInstrumentation
     implements Instrumenter.ForBootstrap,
         Instrumenter.ForSingleType,
         Instrumenter.HasMethodAdvice,
@@ -56,30 +51,18 @@ public class CompletableFutureUniCompletionInstrumentation extends InstrumenterM
   static final String UNI_COMPLETION = COMPLETABLE_FUTURE + "$UniCompletion";
   static final String ADVICE_BASE = JAVA_UTIL_CONCURRENT + ".CompletableFutureAdvice$";
 
-  public CompletableFutureUniCompletionInstrumentation() {
-    super("java_completablefuture", "java_concurrent");
-  }
-
   @Override
   public String instrumentedType() {
     return UNI_COMPLETION;
   }
 
-  @Override
-  public Map<String, String> contextStore() {
-    return singletonMap(UNI_COMPLETION, ConcurrentState.class.getName());
-  }
-
   @Override
   public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(isConstructor(), ADVICE_BASE + "UniConstructor");
   }
 
   @Override
   public Map<ExcludeType, ? extends Collection<String>> excludedClasses() {
-    if (!isEnabled()) {
-      return Collections.emptyMap();
-    }
     String[] classes = {
       // This is not a subclass of UniCompletion and doesn't have a dependent CompletableFuture
       "java.util.concurrent.CompletableFuture$Completion",

dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/completablefuture/CompletableFutureUniCompletionSubclassInstrumentation.java

@@ -4,29 +4,18 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.nameStartsWith;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.instrumentation.java.completablefuture.CompletableFutureUniCompletionInstrumentation.*;
-import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
-import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.agent.tooling.InstrumenterModule;
-import datadog.trace.bootstrap.instrumentation.java.concurrent.ConcurrentState;
-import java.util.Map;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
 /** Described in {@link CompletableFutureUniCompletionInstrumentation} */
-@AutoService(InstrumenterModule.class)
-public class CompletableFutureUniCompletionSubclassInstrumentation
-    extends InstrumenterModule.Tracing
+public final class CompletableFutureUniCompletionSubclassInstrumentation
     implements Instrumenter.ForBootstrap,
         Instrumenter.ForTypeHierarchy,
         Instrumenter.HasMethodAdvice {
 
-  public CompletableFutureUniCompletionSubclassInstrumentation() {
-    super("java_completablefuture", "java_concurrent");
-  }
-
   @Override
   public String hierarchyMarkerType() {
     return null; // bootstrap type
@@ -37,11 +26,6 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
     return nameStartsWith(COMPLETABLE_FUTURE).and(extendsClass(named(UNI_COMPLETION)));
   }
 
-  @Override
-  public Map<String, String> contextStore() {
-    return singletonMap(UNI_COMPLETION, ConcurrentState.class.getName());
-  }
-
   @Override
   public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(

dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/AsyncPropagatingDisableInstrumentation.java

@@ -7,13 +7,10 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.isAsyncPropagationEnabled;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.setAsyncPropagationEnabled;
-import static datadog.trace.instrumentation.java.concurrent.ConcurrentInstrumentationNames.EXECUTOR_INSTRUMENTATION_NAME;
 import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy;
 import static net.bytebuddy.matcher.ElementMatchers.isTypeInitializer;
 
-import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.agent.tooling.InstrumenterModule;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -23,14 +20,9 @@
  * it can cause the trace to never be reported. Add matchers below to disable async propagation
  * during this period.
  */
-@AutoService(InstrumenterModule.class)
-public final class AsyncPropagatingDisableInstrumentation extends InstrumenterModule.Tracing
+public final class AsyncPropagatingDisableInstrumentation
     implements Instrumenter.CanShortcutTypeMatching, Instrumenter.HasMethodAdvice {
 
-  public AsyncPropagatingDisableInstrumentation() {
-    super(EXECUTOR_INSTRUMENTATION_NAME);
-  }
-
   private static final ElementMatcher.Junction<TypeDescription> RX_WORKERS =
       nameStartsWith("rx.").and(extendsClass(named("rx.Scheduler$Worker")));
   private static final ElementMatcher<TypeDescription> NETTY_UNSAFE =

dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/ExecutorModule.java

@@ -0,0 +1,54 @@
+package datadog.trace.instrumentation.java.concurrent;
+
+import static datadog.trace.instrumentation.java.concurrent.ConcurrentInstrumentationNames.EXECUTOR_INSTRUMENTATION_NAME;
+import static datadog.trace.instrumentation.java.concurrent.executor.ThreadPoolExecutorInstrumentation.TPE;
+import static java.util.Collections.singleton;
+
+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.instrumentation.java.concurrent.State;
+import datadog.trace.instrumentation.java.concurrent.executor.JavaExecutorInstrumentation;
+import datadog.trace.instrumentation.java.concurrent.executor.NonStandardExecutorInstrumentation;
+import datadog.trace.instrumentation.java.concurrent.executor.RejectedExecutionHandlerInstrumentation;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.RunnableFuture;
+
+@AutoService(InstrumenterModule.class)
+public class ExecutorModule extends InstrumenterModule.Tracing {
+  public ExecutorModule() {
+    super(EXECUTOR_INSTRUMENTATION_NAME);
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    final Map<String, String> contextStore = new HashMap<>();
+    contextStore.put(Runnable.class.getName(), State.class.getName());
+    contextStore.put(TPE, Boolean.class.getName());
+    contextStore.put(RunnableFuture.class.getName(), State.class.getName());
+    return contextStore;
+  }
+
+  @Override
+  public List<Instrumenter> typeInstrumentations() {
+    final List<Instrumenter> instrumenters = new ArrayList<>(5);
+    final InstrumenterConfig config = InstrumenterConfig.get();
+    instrumenters.add(new AsyncPropagatingDisableInstrumentation());
+    if (InstrumenterConfig.get().isIntegrationEnabled(singleton("new-task-for"), true)) {
+      instrumenters.add(new WrapRunnableAsNewTaskInstrumentation());
+    }
+    instrumenters.add(new JavaExecutorInstrumentation());
+    if (config.isIntegrationEnabled(singleton(EXECUTOR_INSTRUMENTATION_NAME + ".other"), true)) {
+      instrumenters.add(new NonStandardExecutorInstrumentation());
+    }
+    if (config.isIntegrationEnabled(singleton("rejected-execution-handler"), true)) {
+      instrumenters.add(new RejectedExecutionHandlerInstrumentation());
+    }
+
+    return instrumenters;
+  }
+}

dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/ProfilingQueuingModule.java

@@ -0,0 +1,44 @@
+package datadog.trace.instrumentation.java.concurrent;
+
+import static datadog.trace.instrumentation.java.concurrent.ConcurrentInstrumentationNames.EXECUTOR_INSTRUMENTATION_NAME;
+import static datadog.trace.instrumentation.java.concurrent.ConcurrentInstrumentationNames.FORK_JOIN_POOL_INSTRUMENTATION_NAME;
+import static java.util.Collections.singleton;
+
+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.api.config.ProfilingConfig;
+import datadog.trace.bootstrap.config.provider.ConfigProvider;
+import datadog.trace.instrumentation.java.concurrent.forkjoin.JavaForkJoinWorkQueueInstrumentation;
+import java.util.ArrayList;
+import java.util.List;
+
+@AutoService(InstrumenterModule.class)
+public class ProfilingQueuingModule extends InstrumenterModule.Profiling {
+  public ProfilingQueuingModule() {
+    super(EXECUTOR_INSTRUMENTATION_NAME);
+  }
+
+  @Override
+  public boolean isEnabled() {
+    return super.isEnabled()
+        && ConfigProvider.getInstance()
+            .getBoolean(
+                ProfilingConfig.PROFILING_QUEUEING_TIME_ENABLED,
+                ProfilingConfig.PROFILING_QUEUEING_TIME_ENABLED_DEFAULT);
+  }
+
+  @Override
+  public List<Instrumenter> typeInstrumentations() {
+    final List<Instrumenter> instrumenters = new ArrayList<>(2);
+    final InstrumenterConfig config = InstrumenterConfig.get();
+    if (config.isIntegrationEnabled(singleton(FORK_JOIN_POOL_INSTRUMENTATION_NAME), true)) {
+      instrumenters.add(new JavaForkJoinWorkQueueInstrumentation());
+    }
+    if (config.isIntegrationEnabled(singleton("task-unwrapping"), true)) {
+      instrumenters.add(new TaskUnwrappingInstrumentation());
+    }
+    return instrumenters;
+  }
+}

dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/TaskUnwrappingInstrumentation.java

@@ -1,29 +1,10 @@
 package datadog.trace.instrumentation.java.concurrent;
 
-import static datadog.trace.instrumentation.java.concurrent.ConcurrentInstrumentationNames.EXECUTOR_INSTRUMENTATION_NAME;
-
-import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.agent.tooling.bytebuddy.profiling.UnwrappingVisitor;
-import datadog.trace.api.config.ProfilingConfig;
-import datadog.trace.bootstrap.config.provider.ConfigProvider;
 
-@AutoService(InstrumenterModule.class)
-public class TaskUnwrappingInstrumentation extends InstrumenterModule.Profiling
+public class TaskUnwrappingInstrumentation
     implements Instrumenter.ForKnownTypes, Instrumenter.HasTypeAdvice {
-  public TaskUnwrappingInstrumentation() {
-    super(EXECUTOR_INSTRUMENTATION_NAME, "task-unwrapping");
-  }
-
-  @Override
-  public boolean isEnabled() {
-    return super.isEnabled()
-        && ConfigProvider.getInstance()
-            .getBoolean(
-                ProfilingConfig.PROFILING_QUEUEING_TIME_ENABLED,
-                ProfilingConfig.PROFILING_QUEUEING_TIME_ENABLED_DEFAULT);
-  }
 
   private static final String[] TYPES_WITH_FIELDS = {
     "java.util.concurrent.FutureTask",

dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/WrapRunnableAsNewTaskInstrumentation.java

@@ -4,15 +4,12 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.exclude;
-import static datadog.trace.instrumentation.java.concurrent.ConcurrentInstrumentationNames.EXECUTOR_INSTRUMENTATION_NAME;
 import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.not;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
-import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.NewTaskForPlaceholder;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.Wrapper;
 import java.util.concurrent.AbstractExecutorService;
@@ -22,12 +19,8 @@
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.matcher.ElementMatcher.Junction;
 
-@AutoService(InstrumenterModule.class)
-public final class WrapRunnableAsNewTaskInstrumentation extends InstrumenterModule.Tracing
+public final class WrapRunnableAsNewTaskInstrumentation
     implements Instrumenter.ForBootstrap, Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
-  public WrapRunnableAsNewTaskInstrumentation() {
-    super(EXECUTOR_INSTRUMENTATION_NAME, "new-task-for");
-  }
 
   @Override
   public String[] knownMatchingTypes() {