Check jax-rs AsyncResponse for span before starting new one
This will prevent the request from being broken due to replacing the unfinished span.
MODIFY
dd-java-agent/instrumentation/dropwizard/src/test/groovy/DropwizardAsyncTest.groovy
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) {
+ Response query_param(@QueryParam("some") String param, @Suspended final AsyncResponse asyncResponse) {
+ executor.execute {
controller(QUERY_PARAM) {
- Response.status(QUERY_PARAM.status).entity("some=$param".toString()).build()
+ asyncResponse.resume(Response.status(QUERY_PARAM.status).entity("some=$param".toString()).build())
+ }
}
}

MODIFY
dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/JavaExecutorInstrumentation.java
dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/JavaExecutorInstrumentation.java
@@ -4,6 +4,7 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScop
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 final class JavaExecutorInstrumentation extends AbstractExecutorInstrumen
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)),

MODIFY
dd-java-agent/instrumentation/jax-rs-annotations-2/src/main/java/datadog/trace/instrumentation/jaxrs2/JaxRsAnnotationsInstrumentation.java
dd-java-agent/instrumentation/jax-rs-annotations-2/src/main/java/datadog/trace/instrumentation/jaxrs2/JaxRsAnnotationsInstrumentation.java
@@ -15,6 +15,7 @@ import static net.bytebuddy.matcher.ElementMatchers.named;
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 final class JaxRsAnnotationsInstrumentation extends Instrumenter.Default
@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,6 +115,11 @@ public final class JaxRsAnnotationsInstrumentation extends Instrumenter.Default
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 final class JaxRsAnnotationsInstrumentation extends Instrumenter.Default
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 final class JaxRsAnnotationsInstrumentation extends Instrumenter.Default
return;
}
- AsyncResponse asyncResponse = null;
- for (final Object arg : args) {
- if (arg instanceof AsyncResponse) {
- asyncResponse = (AsyncResponse) arg;
- break;
- }
+ if (asyncResponse == null || !asyncResponse.isSuspended()) {
+ if (asyncResponse != null && !asyncResponse.isSuspended()) {
+ // Clear span from the asyncResponse
+ InstrumentationContext.get(AsyncResponse.class, AgentSpan.class).put(asyncResponse, null);
}
- if (asyncResponse != null && asyncResponse.isSuspended()) {
- InstrumentationContext.get(AsyncResponse.class, AgentSpan.class).put(asyncResponse, span);
- } else {
DECORATE.beforeFinish(span);
span.finish();
}

Check jax-rs AsyncResponse for span before starting new one

This will prevent the request from being broken due to replacing the unfinished span.