Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Check jax-rs AsyncResponse for span before starting new one #1403

Merged
merged 3 commits into from
Apr 28, 2020
Merged

Check jax-rs AsyncResponse for span before starting new one #1403

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
2e8f695
Check jax-rs AsyncResponse for span before starting new one
tylerbenson Apr 27, 2020
90987a8
Merge branch 'master' into tyler/jax-rs-nested-async
tylerbenson Apr 28, 2020
1f69da6
Remove nested if.
tylerbenson Apr 28, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions dd-java-agent/instrumentation/dropwizard/src/test/groovy/DropwizardAsyncTest.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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())
}
}
}

Expand Down
3 changes: 2 additions & 1 deletion .../main/java/datadog/trace/instrumentation/java/concurrent/JavaExecutorInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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)),
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not immediately related, though I discovered this problem in the same investigation and it is tangentially related...

When an executor has multiple execute methods, one calling the other, previously they were both instrumented and we would see Failed to set continuation because another continuation is already set log messages. This is because the first call would create a continuation, set it on the task, then call the next execute method, which would do the same, but it would close the second one and keep the first. These log messages were a red herring for my investigation as it didn't actually impact the application reporting.

JavaExecutorInstrumentation.class.getName() + "$SetExecuteRunnableStateAdvice");
transformers.put(
named("execute").and(takesArgument(0, ForkJoinTask.class)),
Expand Down
47 changes: 36 additions & 11 deletions ...2/src/main/java/datadog/trace/instrumentation/jaxrs2/JaxRsAnnotationsInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
}

Comment on lines +87 to +108
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

And this doesn't need to be done for datadog.trace.instrumentation.jaxrs1.JaxRsAnnotationsInstrumentation?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

jaxrs1 doesn't support async. that is the main distinction between the two.

// Rename the parent span according to the path represented by these annotations.
final AgentSpan parent = activeSpan();

Expand All @@ -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);
Expand All @@ -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();
}
Expand Down