Merge pull request #1403 from DataDog/tyler/jax-rs-nested-async

DataDog · Apr 28, 2020 · 03b6f57 · 03b6f57
2 parents 3808bbf + 1f69da6
commit 03b6f57
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 16 deletions.
diff --git a/dd-java-agent/instrumentation/dropwizard/src/test/groovy/DropwizardAsyncTest.groovy b/dd-java-agent/instrumentation/dropwizard/src/test/groovy/DropwizardAsyncTest.groovy
@@ -2,13 +2,13 @@ import io.dropwizard.Application
 import io.dropwizard.Configuration
 import io.dropwizard.setup.Bootstrap
 import io.dropwizard.setup.Environment
+
 import javax.ws.rs.GET
 import javax.ws.rs.Path
 import javax.ws.rs.QueryParam
 import javax.ws.rs.container.AsyncResponse
 import javax.ws.rs.container.Suspended
 import javax.ws.rs.core.Response
-
 import java.util.concurrent.Executors
 
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.ERROR
@@ -60,9 +60,11 @@ class DropwizardAsyncTest extends DropwizardTest {
 
     @GET
     @Path("query")
-    Response query_param(@QueryParam("some") String param) {
-      controller(QUERY_PARAM) {
-        Response.status(QUERY_PARAM.status).entity("some=$param".toString()).build()
+    Response query_param(@QueryParam("some") String param, @Suspended final AsyncResponse asyncResponse) {
+      executor.execute {
+        controller(QUERY_PARAM) {
+          asyncResponse.resume(Response.status(QUERY_PARAM.status).entity("some=$param".toString()).build())
+        }
       }
     }
 

diff --git a/.../main/java/datadog/trace/instrumentation/java/concurrent/JavaExecutorInstrumentation.java b/.../main/java/datadog/trace/instrumentation/java/concurrent/JavaExecutorInstrumentation.java
@@ -4,6 +4,7 @@
 import static net.bytebuddy.matcher.ElementMatchers.nameMatches;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
@@ -44,7 +45,7 @@ public Map<String, String> contextStore() {
   public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
     final Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
     transformers.put(
-        named("execute").and(takesArgument(0, Runnable.class)),
+        named("execute").and(takesArgument(0, Runnable.class)).and(takesArguments(1)),
         JavaExecutorInstrumentation.class.getName() + "$SetExecuteRunnableStateAdvice");
     transformers.put(
         named("execute").and(takesArgument(0, ForkJoinTask.class)),

diff --git a/...2/src/main/java/datadog/trace/instrumentation/jaxrs2/JaxRsAnnotationsInstrumentation.java b/...2/src/main/java/datadog/trace/instrumentation/jaxrs2/JaxRsAnnotationsInstrumentation.java
@@ -15,6 +15,7 @@
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -83,7 +84,28 @@ public static class JaxRsAnnotationsAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static AgentScope nameSpan(
-        @Advice.This final Object target, @Advice.Origin final Method method) {
+        @Advice.This final Object target,
+        @Advice.Origin final Method method,
+        @Advice.AllArguments final Object[] args,
+        @Advice.Local("asyncResponse") AsyncResponse asyncResponse) {
+      ContextStore<AsyncResponse, AgentSpan> contextStore = null;
+      for (final Object arg : args) {
+        if (arg instanceof AsyncResponse) {
+          asyncResponse = (AsyncResponse) arg;
+          contextStore = InstrumentationContext.get(AsyncResponse.class, AgentSpan.class);
+          if (contextStore.get(asyncResponse) != null) {
+            /**
+             * We are probably in a recursive call and don't want to start a new span because it
+             * would replace the existing span in the asyncResponse and cause it to never finish. We
+             * could work around this by using a list instead, but we likely don't want the extra
+             * span anyway.
+             */
+            return null;
+          }
+          break;
+        }
+      }
+
       // Rename the parent span according to the path represented by these annotations.
       final AgentSpan parent = activeSpan();
 
@@ -93,14 +115,22 @@ public static AgentScope nameSpan(
 
       final AgentScope scope = activateSpan(span, false);
       scope.setAsyncPropagation(true);
+
+      if (contextStore != null && asyncResponse != null) {
+        contextStore.put(asyncResponse, span);
+      }
+
       return scope;
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopSpan(
         @Advice.Enter final AgentScope scope,
         @Advice.Thrown final Throwable throwable,
-        @Advice.AllArguments final Object[] args) {
+        @Advice.Local("asyncResponse") final AsyncResponse asyncResponse) {
+      if (scope == null) {
+        return;
+      }
       final AgentSpan span = scope.span();
       if (throwable != null) {
         DECORATE.onError(span, throwable);
@@ -110,16 +140,11 @@ public static void stopSpan(
         return;
       }
 
-      AsyncResponse asyncResponse = null;
-      for (final Object arg : args) {
-        if (arg instanceof AsyncResponse) {
-          asyncResponse = (AsyncResponse) arg;
-          break;
-        }
+      if (asyncResponse != null && !asyncResponse.isSuspended()) {
+        // Clear span from the asyncResponse. Logically this should never happen. Added to be safe.
+        InstrumentationContext.get(AsyncResponse.class, AgentSpan.class).put(asyncResponse, null);
       }
-      if (asyncResponse != null && asyncResponse.isSuspended()) {
-        InstrumentationContext.get(AsyncResponse.class, AgentSpan.class).put(asyncResponse, span);
-      } else {
+      if (asyncResponse == null || !asyncResponse.isSuspended()) {
         DECORATE.beforeFinish(span);
         span.finish();
       }