diff --git a/dd-java-agent/instrumentation/dropwizard/src/test/groovy/DropwizardAsyncTest.groovy b/dd-java-agent/instrumentation/dropwizard/src/test/groovy/DropwizardAsyncTest.groovy index f2d7e3b1920..2a1569655c9 100644 --- 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/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/JavaExecutorInstrumentation.java b/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/JavaExecutorInstrumentation.java index af4fe50b8ab..3edfdb23249 100644 --- a/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/JavaExecutorInstrumentation.java +++ b/dd-java-agent/instrumentation/java-concurrent/src/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 contextStore() { public Map, String> transformers() { final Map, 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/dd-java-agent/instrumentation/jax-rs-annotations-2/src/main/java/datadog/trace/instrumentation/jaxrs2/JaxRsAnnotationsInstrumentation.java b/dd-java-agent/instrumentation/jax-rs-annotations-2/src/main/java/datadog/trace/instrumentation/jaxrs2/JaxRsAnnotationsInstrumentation.java index 6098f4a7c16..52f757707b0 100644 --- a/dd-java-agent/instrumentation/jax-rs-annotations-2/src/main/java/datadog/trace/instrumentation/jaxrs2/JaxRsAnnotationsInstrumentation.java +++ b/dd-java-agent/instrumentation/jax-rs-annotations-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 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,6 +115,11 @@ public static AgentScope nameSpan( final AgentScope scope = activateSpan(span, false); scope.setAsyncPropagation(true); + + if (contextStore != null && asyncResponse != null) { + contextStore.put(asyncResponse, span); + } + return scope; } @@ -100,7 +127,10 @@ public static AgentScope nameSpan( 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(); }