Skip to content

Commit 04771b6

Browse files
amarzialiamarziali
authored
🍒 9333 - Support async servlet for RUM injection (#9343)
* Support async servlet for RUM injection (cherry picked from commit 736260f) * Move isRumEnabled to instrumenter config (cherry picked from commit 54d6a7e)
1 parent 8b4cecd commit 04771b6

File tree

26 files changed

+534
-37
lines changed

26 files changed

+534
-37
lines changed

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeOutputStream.java

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -172,7 +172,6 @@ private void drain() throws IOException {
172172

173173
public void commit() throws IOException {
174174
if (filter || wasDraining) {
175-
filter = false;
176175
drain();
177176
}
178177
}
@@ -190,4 +189,8 @@ public void close() throws IOException {
190189
downstream.close();
191190
}
192191
}
192+
193+
public void setFilter(boolean filter) {
194+
this.filter = filter;
195+
}
193196
}

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeWriter.java

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -173,7 +173,6 @@ private void drain() throws IOException {
173173

174174
public void commit() throws IOException {
175175
if (filter || wasDraining) {
176-
filter = false;
177176
drain();
178177
}
179178
}
@@ -191,4 +190,8 @@ public void close() throws IOException {
191190
downstream.close();
192191
}
193192
}
193+
194+
public void setFilter(boolean filter) {
195+
this.filter = filter;
196+
}
194197
}

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

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
package datadog.trace.instrumentation.servlet3;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
4+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
6+
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_RUM_INJECTED;
7+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
8+
9+
import com.google.auto.service.AutoService;
10+
import datadog.trace.agent.tooling.Instrumenter;
11+
import datadog.trace.agent.tooling.InstrumenterModule;
12+
import datadog.trace.api.InstrumenterConfig;
13+
import javax.servlet.AsyncContext;
14+
import net.bytebuddy.asm.Advice;
15+
import net.bytebuddy.description.type.TypeDescription;
16+
import net.bytebuddy.matcher.ElementMatcher;
17+
18+
@AutoService(InstrumenterModule.class)
19+
public class RumAsyncContextInstrumentation extends InstrumenterModule.Tracing
20+
implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
21+
22+
public RumAsyncContextInstrumentation() {
23+
super("servlet", "servlet-3", "servlet-3-async-context");
24+
}
25+
26+
@Override
27+
public String hierarchyMarkerType() {
28+
return "javax.servlet.AsyncContext";
29+
}
30+
31+
@Override
32+
public String[] helperClassNames() {
33+
return new String[] {
34+
packageName + ".RumHttpServletResponseWrapper", packageName + ".WrappedServletOutputStream",
35+
};
36+
}
37+
38+
@Override
39+
public ElementMatcher<TypeDescription> hierarchyMatcher() {
40+
return implementsInterface(named(hierarchyMarkerType()));
41+
}
42+
43+
@Override
44+
public boolean isEnabled() {
45+
return super.isEnabled() && InstrumenterConfig.get().isRumEnabled();
46+
}
47+
48+
@Override
49+
public void methodAdvice(MethodTransformer transformer) {
50+
transformer.applyAdvice(
51+
isMethod().and(namedOneOf("complete", "dispatch")), getClass().getName() + "$CommitAdvice");
52+
}
53+
54+
public static class CommitAdvice {
55+
@Advice.OnMethodEnter(suppress = Throwable.class)
56+
public static void commitRumBuffer(@Advice.This final AsyncContext asyncContext) {
57+
final Object maybeRumWrappedResponse =
58+
asyncContext.getRequest().getAttribute(DD_RUM_INJECTED);
59+
if (maybeRumWrappedResponse instanceof RumHttpServletResponseWrapper) {
60+
((RumHttpServletResponseWrapper) maybeRumWrappedResponse).commit();
61+
}
62+
}
63+
}
64+
}

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

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
package datadog.trace.instrumentation.servlet3;
2+
3+
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_RUM_INJECTED;
4+
5+
import javax.servlet.AsyncContext;
6+
import javax.servlet.ServletRequest;
7+
import javax.servlet.ServletResponse;
8+
import javax.servlet.http.HttpServletRequest;
9+
import javax.servlet.http.HttpServletRequestWrapper;
10+
import javax.servlet.http.HttpServletResponse;
11+
12+
public class RumHttpServletRequestWrapper extends HttpServletRequestWrapper {
13+
14+
private final HttpServletResponse response;
15+
16+
public RumHttpServletRequestWrapper(
17+
final HttpServletRequest request, final HttpServletResponse response) {
18+
super(request);
19+
this.response = response;
20+
}
21+
22+
@Override
23+
public AsyncContext startAsync() throws IllegalStateException {
24+
// need to hide this method otherwise we cannot control the wrapped response used asynchronously
25+
return startAsync(getRequest(), response);
26+
}
27+
28+
@Override
29+
public AsyncContext startAsync(ServletRequest servletRequest, ServletResponse servletResponse)
30+
throws IllegalStateException {
31+
// deactivate the previous wrapper
32+
final Object maybeRumWrappedResponse = (servletRequest.getAttribute(DD_RUM_INJECTED));
33+
if (maybeRumWrappedResponse instanceof RumHttpServletResponseWrapper) {
34+
((RumHttpServletResponseWrapper) maybeRumWrappedResponse).commit();
35+
((RumHttpServletResponseWrapper) maybeRumWrappedResponse).stopFiltering();
36+
}
37+
ServletResponse actualResponse = servletResponse;
38+
// rewrap it
39+
if (servletResponse instanceof HttpServletResponse) {
40+
actualResponse = new RumHttpServletResponseWrapper((HttpServletResponse) servletResponse);
41+
servletRequest.setAttribute(DD_RUM_INJECTED, actualResponse);
42+
}
43+
return super.startAsync(servletRequest, actualResponse);
44+
}
45+
}

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,7 @@ public void setContentType(String type) {
131131
}
132132
if (!shouldInject) {
133133
commit();
134+
stopFiltering();
134135
}
135136
super.setContentType(type);
136137
}
@@ -149,4 +150,14 @@ public void commit() {
149150
}
150151
}
151152
}
153+
154+
public void stopFiltering() {
155+
shouldInject = false;
156+
if (wrappedPipeWriter != null) {
157+
wrappedPipeWriter.setFilter(false);
158+
}
159+
if (outputStream != null) {
160+
outputStream.setFilter(false);
161+
}
162+
}
152163
}

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

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -46,11 +46,16 @@ public static boolean onEnter(
4646
final HttpServletRequest httpServletRequest = (HttpServletRequest) request;
4747
HttpServletResponse httpServletResponse = (HttpServletResponse) response;
4848

49-
if (RumInjector.get().isEnabled() && httpServletRequest.getAttribute(DD_RUM_INJECTED) == null) {
50-
httpServletRequest.setAttribute(DD_RUM_INJECTED, Boolean.TRUE);
51-
rumServletWrapper = new RumHttpServletResponseWrapper(httpServletResponse);
52-
httpServletResponse = rumServletWrapper;
53-
response = httpServletResponse;
49+
if (RumInjector.get().isEnabled()) {
50+
final Object maybeRumWrapper = httpServletRequest.getAttribute(DD_RUM_INJECTED);
51+
if (maybeRumWrapper instanceof RumHttpServletResponseWrapper) {
52+
rumServletWrapper = (RumHttpServletResponseWrapper) maybeRumWrapper;
53+
} else {
54+
rumServletWrapper = new RumHttpServletResponseWrapper((HttpServletResponse) response);
55+
httpServletRequest.setAttribute(DD_RUM_INJECTED, rumServletWrapper);
56+
response = rumServletWrapper;
57+
request = new RumHttpServletRequestWrapper(httpServletRequest, rumServletWrapper);
58+
}
5459
}
5560

5661
Object dispatchSpan = request.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE);

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ public String[] helperClassNames() {
5252
packageName + ".Servlet3Decorator",
5353
packageName + ".ServletRequestURIAdapter",
5454
packageName + ".FinishAsyncDispatchListener",
55+
packageName + ".RumHttpServletRequestWrapper",
5556
packageName + ".RumHttpServletResponseWrapper",
5657
packageName + ".WrappedServletOutputStream",
5758
"datadog.trace.instrumentation.servlet.ServletBlockingHelper",

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,4 +86,8 @@ public void setWriteListener(WriteListener writeListener) {
8686
public void commit() throws IOException {
8787
filtered.commit();
8888
}
89+
90+
public void setFilter(boolean filter) {
91+
filtered.setFilter(filter);
92+
}
8993
}

‎dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaServletInstrumentation.java

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,9 @@ public String hierarchyMarkerType() {
4242
@Override
4343
public String[] helperClassNames() {
4444
return new String[] {
45-
packageName + ".RumHttpServletResponseWrapper", packageName + ".WrappedServletOutputStream",
45+
packageName + ".RumHttpServletRequestWrapper",
46+
packageName + ".RumHttpServletResponseWrapper",
47+
packageName + ".WrappedServletOutputStream",
4648
};
4749
}
4850

@@ -67,7 +69,7 @@ public void methodAdvice(MethodTransformer transformer) {
6769
public static class JakartaServletAdvice {
6870
@Advice.OnMethodEnter(suppress = Throwable.class)
6971
public static AgentSpan before(
70-
@Advice.Argument(0) final ServletRequest request,
72+
@Advice.Argument(value = 0, readOnly = false) ServletRequest request,
7173
@Advice.Argument(value = 1, readOnly = false) ServletResponse response,
7274
@Advice.Local("rumServletWrapper") RumHttpServletResponseWrapper rumServletWrapper) {
7375
if (!(request instanceof HttpServletRequest)) {
@@ -77,11 +79,16 @@ public static AgentSpan before(
7779
if (response instanceof HttpServletResponse) {
7880
final HttpServletRequest httpServletRequest = (HttpServletRequest) request;
7981

80-
if (RumInjector.get().isEnabled()
81-
&& httpServletRequest.getAttribute(DD_RUM_INJECTED) == null) {
82-
httpServletRequest.setAttribute(DD_RUM_INJECTED, Boolean.TRUE);
83-
rumServletWrapper = new RumHttpServletResponseWrapper((HttpServletResponse) response);
84-
response = rumServletWrapper;
82+
if (RumInjector.get().isEnabled()) {
83+
final Object maybeRumWrapper = httpServletRequest.getAttribute(DD_RUM_INJECTED);
84+
if (maybeRumWrapper instanceof RumHttpServletResponseWrapper) {
85+
rumServletWrapper = (RumHttpServletResponseWrapper) maybeRumWrapper;
86+
} else {
87+
rumServletWrapper = new RumHttpServletResponseWrapper((HttpServletResponse) response);
88+
httpServletRequest.setAttribute(DD_RUM_INJECTED, rumServletWrapper);
89+
response = rumServletWrapper;
90+
request = new RumHttpServletRequestWrapper(httpServletRequest, rumServletWrapper);
91+
}
8592
}
8693
}
8794

‎dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/RumAsyncContextInstrumentation.java

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
package datadog.trace.instrumentation.servlet5;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
4+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
6+
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_RUM_INJECTED;
7+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
8+
9+
import com.google.auto.service.AutoService;
10+
import datadog.trace.agent.tooling.Instrumenter;
11+
import datadog.trace.agent.tooling.InstrumenterModule;
12+
import datadog.trace.api.InstrumenterConfig;
13+
import jakarta.servlet.AsyncContext;
14+
import net.bytebuddy.asm.Advice;
15+
import net.bytebuddy.description.type.TypeDescription;
16+
import net.bytebuddy.matcher.ElementMatcher;
17+
18+
@AutoService(InstrumenterModule.class)
19+
public class RumAsyncContextInstrumentation extends InstrumenterModule.Tracing
20+
implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
21+
22+
public RumAsyncContextInstrumentation() {
23+
super("servlet", "servlet-5", "servlet-5-async-context");
24+
}
25+
26+
@Override
27+
public String hierarchyMarkerType() {
28+
return "jakarta.servlet.AsyncContext";
29+
}
30+
31+
@Override
32+
public String[] helperClassNames() {
33+
return new String[] {
34+
packageName + ".RumHttpServletResponseWrapper", packageName + ".WrappedServletOutputStream",
35+
};
36+
}
37+
38+
@Override
39+
public ElementMatcher<TypeDescription> hierarchyMatcher() {
40+
return implementsInterface(named(hierarchyMarkerType()));
41+
}
42+
43+
@Override
44+
public boolean isEnabled() {
45+
return super.isEnabled() && InstrumenterConfig.get().isRumEnabled();
46+
}
47+
48+
@Override
49+
public void methodAdvice(MethodTransformer transformer) {
50+
transformer.applyAdvice(
51+
isMethod().and(namedOneOf("complete", "dispatch")), getClass().getName() + "$CommitAdvice");
52+
}
53+
54+
public static class CommitAdvice {
55+
@Advice.OnMethodEnter(suppress = Throwable.class)
56+
public static void commitRumBuffer(@Advice.This final AsyncContext asyncContext) {
57+
final Object maybeRumWrappedResponse =
58+
asyncContext.getRequest().getAttribute(DD_RUM_INJECTED);
59+
if (maybeRumWrappedResponse instanceof RumHttpServletResponseWrapper) {
60+
((RumHttpServletResponseWrapper) maybeRumWrappedResponse).commit();
61+
}
62+
}
63+
}
64+
}

0 commit comments

Comments
 (0)