add MethodInvocation wrapper for spanning @async invocations

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/matcher/AdditionalLibraryIgnoresMatcher.java

@@ -42,7 +42,8 @@ public boolean matches(final T target) {
     }
 
     if (name.startsWith("org.springframework.")) {
-      if (name.startsWith("org.springframework.aop.")
+      if ((name.startsWith("org.springframework.aop.")
+              && !(name.equals("org.springframework.aop.interceptor.AsyncExecutionInterceptor")))
           || name.startsWith("org.springframework.cache.")
           || name.startsWith("org.springframework.dao.")
           || name.startsWith("org.springframework.ejb.")

dd-java-agent/instrumentation/spring-scheduling-3.1/src/main/java/datadog/trace/instrumentation/springscheduling/SpannedMethodInvocation.java

@@ -0,0 +1,62 @@
+package datadog.trace.instrumentation.springscheduling;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.springscheduling.SpringSchedulingDecorator.DECORATE;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import java.lang.reflect.AccessibleObject;
+import java.lang.reflect.Method;
+import org.aopalliance.intercept.MethodInvocation;
+
+public class SpannedMethodInvocation implements MethodInvocation {
+
+  private final AgentSpan parent;
+  private final MethodInvocation delegate;
+
+  public SpannedMethodInvocation(AgentSpan parent, MethodInvocation delegate) {
+    this.parent = parent;
+    this.delegate = delegate;
+  }
+
+  @Override
+  public Method getMethod() {
+    return delegate.getMethod();
+  }
+
+  @Override
+  public Object[] getArguments() {
+    return delegate.getArguments();
+  }
+
+  @Override
+  public Object proceed() throws Throwable {
+    CharSequence spanName = DECORATE.spanNameForMethod(delegate.getMethod());
+    // TODO kill all APIs requiring String parameters with fire
+    final AgentSpan span =
+        parent == null
+            ? startSpan(spanName.toString())
+            : startSpan(spanName.toString(), parent.context());
+    try (AgentScope scope = activateSpan(span)) {
+      // question: is this necessary? What does it do?
+      // if the delegate does async work is everything OK because of this?
+      // if the delegate does async work, should I need to worry about it here?
+      scope.setAsyncPropagation(true);
+      Object result = delegate.proceed();
+      // question? Why can't this just be AutoCloseable? Dogma?
+      span.finish();
+      return result;
+    }
+  }
+
+  @Override
+  public Object getThis() {
+    return delegate.getThis();
+  }
+
+  @Override
+  public AccessibleObject getStaticPart() {
+    return delegate.getStaticPart();
+  }
+}

dd-java-agent/instrumentation/spring-scheduling-3.1/src/main/java/datadog/trace/instrumentation/springscheduling/SpringAsyncAdvice.java

@@ -0,0 +1,16 @@
+package datadog.trace.instrumentation.springscheduling;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
+
+import net.bytebuddy.asm.Advice;
+import org.aopalliance.intercept.MethodInvocation;
+
+public class SpringAsyncAdvice {
+
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static void scheduleAsync(
+      @Advice.Argument(value = 0, readOnly = false) MethodInvocation invocation) {
+    // wrap so that when the invocation is invoked, it can be wrapped with a span's start and stop
+    invocation = new SpannedMethodInvocation(activeSpan(), invocation);
+  }
+}

dd-java-agent/instrumentation/spring-scheduling-3.1/src/main/java/datadog/trace/instrumentation/springscheduling/SpringAsyncInstrumentation.java

@@ -0,0 +1,38 @@
+package datadog.trace.instrumentation.springscheduling;
+
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import java.util.Collections;
+import java.util.Map;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import net.bytebuddy.matcher.ElementMatchers;
+
+@AutoService(Instrumenter.class)
+public class SpringAsyncInstrumentation extends Instrumenter.Default {
+
+  public SpringAsyncInstrumentation() {
+    super("spring-async");
+  }
+
+  @Override
+  public ElementMatcher<? super TypeDescription> typeMatcher() {
+    return named("org.springframework.aop.interceptor.AsyncExecutionInterceptor");
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return Collections.singletonMap(
+        isMethod()
+            .and(
+                named("invoke")
+                    .and(
+                        ElementMatchers.takesArgument(
+                            0, named("org.aopalliance.intercept.MethodInvocation")))),
+        packageName + ".SpringAsyncAdvice");
+  }
+}

dd-java-agent/instrumentation/spring-scheduling-3.1/src/test/groovy/SpringAsyncTest.groovy

@@ -0,0 +1,79 @@
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.api.Trace
+import org.springframework.context.annotation.AnnotationConfigApplicationContext
+import spock.lang.Ignore
+
+import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
+
+class SpringAsyncTest extends AgentTestRunner {
+
+  @Ignore("this requires improved context propagation in the CompletableFuture instrumentation")
+  def "context propagated through @async annotation" () {
+    setup:
+    def context = new AnnotationConfigApplicationContext(AsyncTaskConfig)
+    AsyncTask asyncTask = context.getBean(AsyncTask)
+    when:
+    runUnderTrace("root") {
+      asyncTask.async()
+        .thenApply({
+          new Leaf(it)
+        })
+        .thenApplyAsync({it.leaf() })
+        .join()
+    }
+    then:
+    assertTraces(1) {
+      trace(0, 3) {
+        span(0) {
+          resourceName "root"
+          threadNameStartsWith "main"
+        }
+        span(1) {
+          resourceName "AsyncTask.async"
+          childOf span(0)
+        }
+        span(2) {
+          resourceName "Leaf.leaf"
+          childOf span(1)
+        }
+      }
+    }
+  }
+
+  def "context propagated through to @async task" () {
+    setup:
+    def context = new AnnotationConfigApplicationContext(AsyncTaskConfig)
+    AsyncTask asyncTask = context.getBean(AsyncTask)
+    when:
+    runUnderTrace("root") {
+      asyncTask.async().join()
+    }
+    then:
+    assertTraces(1) {
+      trace(0, 2) {
+        span(0) {
+          resourceName "root"
+          threadNameStartsWith "main"
+        }
+        span(1) {
+          resourceName "AsyncTask.async"
+          childOf span(0)
+        }
+      }
+    }
+  }
+
+  static class Leaf {
+
+    private int i
+
+    Leaf(int i) {
+      this.i = i
+    }
+
+    @Trace
+    int leaf() {
+      return i
+    }
+  }
+}

dd-java-agent/instrumentation/spring-scheduling-3.1/src/test/groovy/SpringSchedulingTest.groovy

@@ -1,12 +1,9 @@
 import datadog.trace.agent.test.AgentTestRunner
-import datadog.trace.api.Trace
 import datadog.trace.bootstrap.instrumentation.api.Tags
 import org.springframework.context.annotation.AnnotationConfigApplicationContext
 
 import java.util.concurrent.TimeUnit
 
-import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
-
 class SpringSchedulingTest extends AgentTestRunner {
 
   def "schedule trigger test according to cron expression"() {
@@ -86,45 +83,4 @@ class SpringSchedulingTest extends AgentTestRunner {
     cleanup:
     context.close()
   }
-
-  def "context propagated through @async annotation" () {
-    setup:
-    def context = new AnnotationConfigApplicationContext(AsyncTaskConfig)
-    AsyncTask asyncTask = context.getBean(AsyncTask)
-
-    when:
-    runUnderTrace("root") {
-      asyncTask.async()
-      .thenApply({
-        new Leaf().leaf(it)
-      }).join()
-    }
-    then:
-    assertTraces(1) {
-      trace(0, 3) {
-        span(0) {
-          resourceName "root"
-          threadNameStartsWith "main"
-        }
-        span(1) {
-          resourceName "AsyncTask.async"
-          childOf span(0)
-          threadNameStartsWith "SimpleAsyncTaskExecutor"
-        }
-        span(2) {
-          resourceName "Leaf.leaf"
-          childOf span(1)
-          threadNameStartsWith "SimpleAsyncTaskExecutor"
-        }
-      }
-    }
-
-  }
-
-  class Leaf {
-    @Trace
-    int leaf(int i) {
-      return i
-    }
-  }
 }