Skip to content

Commit e888b17

Browse files
tylerbensontylerbenson
authored
Merge pull request #582 from DataDog/tyler/servlet-dispatch
 Fix servlet async dispatch
2 parents 7b6c25b + e4a1240 commit e888b17

13 files changed

Lines changed: 703 additions & 395 deletions

‎dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/AbstractServlet3Instrumentation.java‎

Lines changed: 0 additions & 19 deletions
This file was deleted.

‎dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/AsyncContextInstrumentation.java‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
package datadog.trace.instrumentation.servlet3;
2+
3+
import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
4+
import static datadog.trace.instrumentation.servlet3.Servlet3Advice.SERVLET_SPAN;
5+
import static net.bytebuddy.matcher.ElementMatchers.isInterface;
6+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
7+
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
8+
import static net.bytebuddy.matcher.ElementMatchers.named;
9+
import static net.bytebuddy.matcher.ElementMatchers.not;
10+
11+
import com.google.auto.service.AutoService;
12+
import datadog.trace.agent.tooling.Instrumenter;
13+
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
14+
import io.opentracing.Span;
15+
import io.opentracing.propagation.Format;
16+
import io.opentracing.util.GlobalTracer;
17+
import java.util.Collections;
18+
import java.util.Map;
19+
import javax.servlet.AsyncContext;
20+
import javax.servlet.ServletRequest;
21+
import javax.servlet.http.HttpServletRequest;
22+
import net.bytebuddy.asm.Advice;
23+
import net.bytebuddy.description.type.TypeDescription;
24+
import net.bytebuddy.matcher.ElementMatcher;
25+
26+
@AutoService(Instrumenter.class)
27+
public final class AsyncContextInstrumentation extends Instrumenter.Default {
28+
29+
public AsyncContextInstrumentation() {
30+
super("servlet", "servlet-3");
31+
}
32+
33+
@Override
34+
public String[] helperClassNames() {
35+
return new String[] {"datadog.trace.instrumentation.servlet3.HttpServletRequestInjectAdapter"};
36+
}
37+
38+
@Override
39+
public ElementMatcher<TypeDescription> typeMatcher() {
40+
return not(isInterface()).and(safeHasSuperType(named("javax.servlet.AsyncContext")));
41+
}
42+
43+
@Override
44+
public Map<? extends ElementMatcher, String> transformers() {
45+
return Collections.singletonMap(
46+
isMethod().and(isPublic()).and(named("dispatch")), DispatchAdvice.class.getName());
47+
}
48+
49+
/**
50+
* When a request is dispatched, we want to close out the existing request and replace the
51+
* propagation info in the headers with the closed span.
52+
*/
53+
public static class DispatchAdvice {
54+
55+
@Advice.OnMethodEnter(suppress = Throwable.class)
56+
public static boolean enter(
57+
@Advice.This final AsyncContext context, @Advice.AllArguments final Object[] args) {
58+
final int depth = CallDepthThreadLocalMap.incrementCallDepth(AsyncContext.class);
59+
if (depth > 0) {
60+
return false;
61+
}
62+
63+
final ServletRequest request = context.getRequest();
64+
final Object spanAttr = request.getAttribute(SERVLET_SPAN);
65+
if (spanAttr instanceof Span) {
66+
request.removeAttribute(SERVLET_SPAN);
67+
final Span span = (Span) spanAttr;
68+
// Override propagation headers by injecting attributes with new values.
69+
if (request instanceof HttpServletRequest) {
70+
GlobalTracer.get()
71+
.inject(
72+
span.context(),
73+
Format.Builtin.TEXT_MAP,
74+
new HttpServletRequestInjectAdapter((HttpServletRequest) request));
75+
}
76+
final String path;
77+
if (args.length == 1 && args[0] instanceof String) {
78+
path = (String) args[0];
79+
} else if (args.length == 2 && args[1] instanceof String) {
80+
path = (String) args[1];
81+
} else {
82+
path = "true";
83+
}
84+
span.setTag("servlet.dispatch", path);
85+
span.finish();
86+
}
87+
return true;
88+
}
89+
90+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
91+
public static void exit(@Advice.Enter final boolean topLevel) {
92+
if (topLevel) {
93+
CallDepthThreadLocalMap.reset(AsyncContext.class);
94+
}
95+
}
96+
}
97+
}

‎dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/FilterChain3Instrumentation.java‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,19 @@
1515
import net.bytebuddy.matcher.ElementMatcher;
1616

1717
@AutoService(Instrumenter.class)
18-
public final class FilterChain3Instrumentation extends AbstractServlet3Instrumentation {
18+
public final class FilterChain3Instrumentation extends Instrumenter.Default {
19+
public FilterChain3Instrumentation() {
20+
super("servlet", "servlet-3");
21+
}
22+
23+
@Override
24+
public String[] helperClassNames() {
25+
return new String[] {
26+
"datadog.trace.instrumentation.servlet3.HttpServletRequestExtractAdapter",
27+
"datadog.trace.instrumentation.servlet3.HttpServletRequestExtractAdapter$MultivaluedMapFlatIterator",
28+
"datadog.trace.instrumentation.servlet3.TagSettingAsyncListener"
29+
};
30+
}
1931

2032
@Override
2133
public ElementMatcher<TypeDescription> typeMatcher() {

‎dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServlet3Instrumentation.java‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,19 @@
1515
import net.bytebuddy.matcher.ElementMatcher;
1616

1717
@AutoService(Instrumenter.class)
18-
public final class HttpServlet3Instrumentation extends AbstractServlet3Instrumentation {
18+
public final class HttpServlet3Instrumentation extends Instrumenter.Default {
19+
public HttpServlet3Instrumentation() {
20+
super("servlet", "servlet-3");
21+
}
22+
23+
@Override
24+
public String[] helperClassNames() {
25+
return new String[] {
26+
"datadog.trace.instrumentation.servlet3.HttpServletRequestExtractAdapter",
27+
"datadog.trace.instrumentation.servlet3.HttpServletRequestExtractAdapter$MultivaluedMapFlatIterator",
28+
"datadog.trace.instrumentation.servlet3.TagSettingAsyncListener"
29+
};
30+
}
1931

2032
@Override
2133
public ElementMatcher<TypeDescription> typeMatcher() {

‎dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletRequestExtractAdapter.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import io.opentracing.propagation.TextMap;
44
import java.util.AbstractMap;
55
import java.util.ArrayList;
6+
import java.util.Collections;
67
import java.util.Enumeration;
78
import java.util.HashMap;
89
import java.util.Iterator;
@@ -52,6 +53,20 @@ protected Map<String, List<String>> servletHeadersToMultiMap(
5253
headersResult.put(headerName, valuesList);
5354
}
5455

56+
/*
57+
* Read from the attributes and override the headers.
58+
* This is used by HttpServletRequestInjectAdapter when a request is async-dispatched.
59+
*/
60+
final Enumeration<String> attributeNamesIt = httpServletRequest.getAttributeNames();
61+
while (attributeNamesIt.hasMoreElements()) {
62+
final String attributeName = attributeNamesIt.nextElement();
63+
64+
final Object valuesIt = httpServletRequest.getAttribute(attributeName);
65+
if (valuesIt instanceof String) {
66+
headersResult.put(attributeName, Collections.singletonList((String) valuesIt));
67+
}
68+
}
69+
5570
return headersResult;
5671
}
5772

‎dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletRequestInjectAdapter.java‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
package datadog.trace.instrumentation.servlet3;
2+
3+
import io.opentracing.propagation.TextMap;
4+
import java.util.Iterator;
5+
import java.util.Map;
6+
import javax.servlet.http.HttpServletRequest;
7+
8+
/** Inject into request attributes since the request headers can't be modified. */
9+
public class HttpServletRequestInjectAdapter implements TextMap {
10+
11+
private final HttpServletRequest httpServletRequest;
12+
13+
public HttpServletRequestInjectAdapter(final HttpServletRequest httpServletRequest) {
14+
this.httpServletRequest = httpServletRequest;
15+
}
16+
17+
@Override
18+
public Iterator<Map.Entry<String, String>> iterator() {
19+
throw new UnsupportedOperationException(
20+
"This class should be used only with Tracer.extract()!");
21+
}
22+
23+
@Override
24+
public void put(final String key, final String value) {
25+
httpServletRequest.setAttribute(key, value);
26+
}
27+
}

‎dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java‎

Lines changed: 23 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,13 @@
2121
import net.bytebuddy.asm.Advice;
2222

2323
public class Servlet3Advice {
24+
public static final String SERVLET_SPAN = "datadog.servlet.span";
2425

2526
@Advice.OnMethodEnter(suppress = Throwable.class)
2627
public static Scope startSpan(
2728
@Advice.This final Object servlet, @Advice.Argument(0) final ServletRequest req) {
28-
if (GlobalTracer.get().activeSpan() != null || !(req instanceof HttpServletRequest)) {
29+
final Object spanAttr = req.getAttribute(SERVLET_SPAN);
30+
if (!(req instanceof HttpServletRequest) || spanAttr != null) {
2931
// Tracing might already be applied by the FilterChain. If so ignore this.
3032
return null;
3133
}
@@ -53,6 +55,8 @@ public static Scope startSpan(
5355
if (scope instanceof TraceScope) {
5456
((TraceScope) scope).setAsyncPropagation(true);
5557
}
58+
59+
req.setAttribute(SERVLET_SPAN, scope.span());
5660
return scope;
5761
}
5862

@@ -63,13 +67,11 @@ public static void stopSpan(
6367
@Advice.Enter final Scope scope,
6468
@Advice.Thrown final Throwable throwable) {
6569
// Set user.principal regardless of who created this span.
66-
final Span currentSpan = GlobalTracer.get().activeSpan();
67-
if (currentSpan != null) {
68-
if (request instanceof HttpServletRequest) {
69-
final Principal principal = ((HttpServletRequest) request).getUserPrincipal();
70-
if (principal != null) {
71-
currentSpan.setTag(DDTags.USER_NAME, principal.getName());
72-
}
70+
final Object spanAttr = request.getAttribute(SERVLET_SPAN);
71+
if (spanAttr instanceof Span && request instanceof HttpServletRequest) {
72+
final Principal principal = ((HttpServletRequest) request).getUserPrincipal();
73+
if (principal != null) {
74+
((Span) spanAttr).setTag(DDTags.USER_NAME, principal.getName());
7375
}
7476
}
7577

@@ -90,19 +92,23 @@ public static void stopSpan(
9092
((TraceScope) scope).setAsyncPropagation(false);
9193
}
9294
scope.close();
95+
req.removeAttribute(SERVLET_SPAN);
9396
span.finish(); // Finish the span manually since finishSpanOnClose was false
94-
} else if (req.isAsyncStarted()) {
95-
final AtomicBoolean activated = new AtomicBoolean(false);
96-
// what if async is already finished? This would not be called
97-
req.getAsyncContext().addListener(new TagSettingAsyncListener(activated, span));
98-
scope.close();
9997
} else {
100-
Tags.HTTP_STATUS.set(span, resp.getStatus());
101-
if (scope instanceof TraceScope) {
102-
((TraceScope) scope).setAsyncPropagation(false);
98+
final AtomicBoolean activated = new AtomicBoolean(false);
99+
if (req.isAsyncStarted()) {
100+
req.getAsyncContext().addListener(new TagSettingAsyncListener(activated, span));
101+
}
102+
// Check again in case the request finished before adding the listener.
103+
if (!req.isAsyncStarted() && activated.compareAndSet(false, true)) {
104+
Tags.HTTP_STATUS.set(span, resp.getStatus());
105+
if (scope instanceof TraceScope) {
106+
((TraceScope) scope).setAsyncPropagation(false);
107+
}
108+
req.removeAttribute(SERVLET_SPAN);
109+
span.finish(); // Finish the span manually since finishSpanOnClose was false
103110
}
104111
scope.close();
105-
span.finish(); // Finish the span manually since finishSpanOnClose was false
106112
}
107113
}
108114
}

‎dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/TagSettingAsyncListener.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import java.io.IOException;
1111
import java.util.Collections;
1212
import java.util.concurrent.atomic.AtomicBoolean;
13+
import javax.servlet.AsyncContext;
1314
import javax.servlet.AsyncEvent;
1415
import javax.servlet.AsyncListener;
1516
import javax.servlet.http.HttpServletResponse;
@@ -72,6 +73,12 @@ public void onError(final AsyncEvent event) throws IOException {
7273
}
7374
}
7475

76+
/** Re-attach listener for dispatch. */
7577
@Override
76-
public void onStartAsync(final AsyncEvent event) throws IOException {}
78+
public void onStartAsync(final AsyncEvent event) {
79+
final AsyncContext eventAsyncContext = event.getAsyncContext();
80+
if (eventAsyncContext != null) {
81+
eventAsyncContext.addListener(this, event.getSuppliedRequest(), event.getSuppliedResponse());
82+
}
83+
}
7784
}

0 commit comments

Comments
 (0)