Be consistent when deciding when to do async propagation (#8555)

i.e. we should have a non-null, valid span and the async propagation flag should be set

(also remove some unused variables in the executor advice)

Author: mcculls

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/AdviceUtils.java

@@ -52,7 +52,7 @@ public static <T> void cancelTask(ContextStore<T, State> contextStore, final T t
 
   public static <T> void capture(ContextStore<T, State> contextStore, T task) {
     AgentSpan span = activeSpan();
-    if (span != null && isAsyncPropagationEnabled()) {
+    if (span != null && span.isValid() && isAsyncPropagationEnabled()) {
       State state = contextStore.get(task);
       if (null == state) {
         state = State.FACTORY.create();

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ExecutorInstrumentationUtils.java

@@ -1,12 +1,10 @@
 package datadog.trace.bootstrap.instrumentation.java.concurrent;
 
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.isAsyncPropagationEnabled;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType;
 
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-import java.util.concurrent.Executor;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -19,17 +17,17 @@ public final class ExecutorInstrumentationUtils {
    * Checks if given task should get state attached.
    *
    * @param task task object
-   * @param executor executor this task was scheduled on
+   * @param span active span
    * @return true iff given task object should be wrapped
    */
-  public static boolean shouldAttachStateToTask(final Object task, final Executor executor) {
+  public static boolean shouldAttachStateToTask(final Object task, final AgentSpan span) {
     if (task == null) {
       return false;
     }
     if (ExcludeFilter.exclude(ExcludeType.EXECUTOR, task)) {
       return false;
     }
-    return activeSpan() != null && isAsyncPropagationEnabled();
+    return span != null && span.isValid() && isAsyncPropagationEnabled();
   }
 
   /**
@@ -57,12 +55,10 @@ public static <T> State setupState(
   /**
    * Clean up after job submission method has exited.
    *
-   * @param executor the current executor
    * @param state task instrumentation state
    * @param throwable throwable that may have been thrown
    */
-  public static void cleanUpOnMethodExit(
-      final Executor executor, final State state, final Throwable throwable) {
+  public static void cleanUpOnMethodExit(final State state, final Throwable throwable) {
     if (null != state && null != throwable) {
       /*
       Note: this may potentially close somebody else's continuation if we didn't set it

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

@@ -54,14 +54,13 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class AddListenerAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static State addListenerEnter(
-        @Advice.Argument(value = 0, readOnly = false) Runnable task,
-        @Advice.Argument(1) final Executor executor) {
+        @Advice.Argument(value = 0, readOnly = false) Runnable task) {
       final AgentSpan span = activeSpan();
       if (null != span) {
         final Runnable newTask = RunnableWrapper.wrapIfNeeded(task);
         // It is important to check potentially wrapped task if we can instrument task in this
         // executor. Some executors do not support wrapped tasks.
-        if (ExecutorInstrumentationUtils.shouldAttachStateToTask(newTask, executor)) {
+        if (ExecutorInstrumentationUtils.shouldAttachStateToTask(newTask, span)) {
           task = newTask;
           final ContextStore<Runnable, State> contextStore =
               InstrumentationContext.get(Runnable.class, State.class);
@@ -73,10 +72,8 @@ public static State addListenerEnter(
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void addListenerExit(
-        @Advice.Argument(1) final Executor executor,
-        @Advice.Enter final State state,
-        @Advice.Thrown final Throwable throwable) {
-      ExecutorInstrumentationUtils.cleanUpOnMethodExit(executor, state, throwable);
+        @Advice.Enter final State state, @Advice.Thrown final Throwable throwable) {
+      ExecutorInstrumentationUtils.cleanUpOnMethodExit(state, throwable);
     }
 
     private static void muzzleCheck(final AbstractFuture<?> future) {

dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/executor/JavaExecutorInstrumentation.java

@@ -15,7 +15,6 @@
 import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
 import java.util.Collections;
 import java.util.Map;
-import java.util.concurrent.Executor;
 import java.util.concurrent.RunnableFuture;
 import net.bytebuddy.asm.Advice;
 
@@ -38,7 +37,6 @@ public static class SetExecuteRunnableStateAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static State enterJobSubmit(
-        @Advice.This final Executor executor,
         @Advice.Argument(value = 0, readOnly = false) Runnable task) {
       if (task instanceof RunnableFuture) {
         return null;
@@ -51,7 +49,7 @@ public static State enterJobSubmit(
         final Runnable newTask = RunnableWrapper.wrapIfNeeded(task);
         // It is important to check potentially wrapped task if we can instrument task in this
         // executor. Some executors do not support wrapped tasks.
-        if (ExecutorInstrumentationUtils.shouldAttachStateToTask(newTask, executor)) {
+        if (ExecutorInstrumentationUtils.shouldAttachStateToTask(newTask, span)) {
           task = newTask;
           final ContextStore<Runnable, State> contextStore =
               InstrumentationContext.get(Runnable.class, State.class);
@@ -63,10 +61,8 @@ public static State enterJobSubmit(
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exitJobSubmit(
-        @Advice.This final Executor executor,
-        @Advice.Enter final State state,
-        @Advice.Thrown final Throwable throwable) {
-      ExecutorInstrumentationUtils.cleanUpOnMethodExit(executor, state, throwable);
+        @Advice.Enter final State state, @Advice.Thrown final Throwable throwable) {
+      ExecutorInstrumentationUtils.cleanUpOnMethodExit(state, throwable);
     }
   }
 }

dd-java-agent/instrumentation/scala-promise/src/main/java/datadog/trace/instrumentation/scala/PromiseHelper.java

@@ -29,7 +29,7 @@ public class PromiseHelper {
    */
   public static AgentSpan getSpan() {
     final AgentSpan span = activeSpan();
-    if (null != span && isAsyncPropagationEnabled()) {
+    if (null != span && span.isValid() && isAsyncPropagationEnabled()) {
       return span;
     }
     return null;