Skip to content

Commit 09b7c65

Browse files
amarzialimcculls
amarziali
and
mcculls
authored
Support CompletableFuture on spring webmvc controllers (#8659)
* Support CompletableFuture on spring webmvc controllers * Update dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/SpringWebHttpServerDecorator.java Co-authored-by: Stuart McCulloch <[email protected]> * Update dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/SpringWebHttpServerDecorator.java Co-authored-by: Stuart McCulloch <[email protected]> --------- Co-authored-by: Stuart McCulloch <[email protected]>
1 parent 089d1cc commit 09b7c65

File tree

16 files changed

+294
-60
lines changed

16 files changed

+294
-60
lines changed

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/AsyncResultExtensions.java

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -61,17 +61,17 @@ public Object apply(Object result, AgentSpan span) {
6161
}
6262
return null;
6363
}
64+
}
6465

65-
private <T> BiConsumer<T, Throwable> finishSpan(AgentSpan span) {
66-
return (o, throwable) -> {
67-
if (throwable != null) {
68-
span.addThrowable(
69-
throwable instanceof ExecutionException || throwable instanceof CompletionException
70-
? throwable.getCause()
71-
: throwable);
72-
}
73-
span.finish();
74-
};
75-
}
66+
public static <T> BiConsumer<T, Throwable> finishSpan(AgentSpan span) {
67+
return (o, throwable) -> {
68+
if (throwable != null) {
69+
span.addThrowable(
70+
throwable instanceof ExecutionException || throwable instanceof CompletionException
71+
? throwable.getCause()
72+
: throwable);
73+
}
74+
span.finish();
75+
};
7676
}
7777
}

‎dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -357,6 +357,7 @@
357357
0 org.springframework.util.StreamUtils
358358
# Included for IAST Spring mvc unvalidated redirect and xss vulnerability
359359
0 org.springframework.web.method.support.HandlerMethodReturnValueHandlerComposite
360+
0 org.springframework.web.method.support.InvocableHandlerMethod
360361
2 org.xml.*
361362
2 org.yaml.snakeyaml.*
362363
# Need for IAST sink

‎dd-java-agent/instrumentation/spring-webmvc-3.1/build.gradle

Lines changed: 0 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -3,27 +3,8 @@ muzzle {
33
group = 'org.springframework'
44
module = 'spring-webmvc'
55
versions = "[3.1.0.RELEASE,6)"
6-
skipVersions += [
7-
'1.2.1',
8-
'1.2.2',
9-
'1.2.3',
10-
'1.2.4'] // broken releases... missing dependencies
116
skipVersions += '3.2.1.RELEASE' // missing a required class. (bad release?)
127
extraDependency "javax.servlet:javax.servlet-api:3.0.1"
13-
// assertInverse = true
14-
}
15-
16-
// FIXME: webmvc depends on web, so we need a separate integration for spring-web specifically.
17-
fail {
18-
group = 'org.springframework'
19-
module = 'spring-web'
20-
versions = "[,]"
21-
skipVersions += [
22-
'1.2.1',
23-
'1.2.2',
24-
'1.2.3',
25-
'1.2.4'] // broken releases... missing dependencies
26-
extraDependency "javax.servlet:javax.servlet-api:3.0.1"
278
}
289

2910
pass {

‎dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/test/boot/AppConfig.groovy

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,13 +8,15 @@ import org.springframework.boot.web.server.WebServerFactoryCustomizer
88
import org.springframework.boot.web.servlet.filter.OrderedRequestContextFilter
99
import org.springframework.context.annotation.Bean
1010
import org.springframework.core.Ordered
11+
import org.springframework.scheduling.annotation.EnableAsync
1112
import org.springframework.web.filter.RequestContextFilter
1213
import org.springframework.web.servlet.config.annotation.PathMatchConfigurer
1314
import org.springframework.web.servlet.config.annotation.WebMvcConfigurerAdapter
1415
import org.springframework.web.util.UrlPathHelper
1516

1617
// Component scan defeats the purpose of configuring with specific classes
1718
@SpringBootApplication(scanBasePackages = "doesnotexist")
19+
@EnableAsync
1820
class AppConfig extends WebMvcConfigurerAdapter {
1921

2022
@Bean

‎dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/test/boot/SpringBootBasedTest.groovy

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import javax.servlet.http.HttpServletRequest
2727
import javax.servlet.http.HttpServletResponse
2828

2929
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.EXCEPTION
30+
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.FORWARDED
3031
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.LOGIN
3132
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.MATRIX_PARAM
3233
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.NOT_FOUND
@@ -290,7 +291,7 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
290291
}
291292

292293
boolean hasResponseSpan(ServerEndpoint endpoint) {
293-
return endpoint == REDIRECT || endpoint == NOT_FOUND || endpoint == LOGIN
294+
return endpoint == REDIRECT || endpoint == NOT_FOUND || endpoint == LOGIN || endpoint == FORWARDED
294295
}
295296

296297
@Override
@@ -341,6 +342,18 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
341342
defaultTags()
342343
}
343344
}
345+
} else if (endpoint == FORWARDED) {
346+
trace.span {
347+
serviceName expectedServiceName()
348+
operationName 'servlet.dispatch'
349+
resourceName 'servlet.dispatch'
350+
tags {
351+
"$Tags.COMPONENT" 'java-web-servlet-async-dispatcher'
352+
'servlet.context' "/$servletContext"
353+
'servlet.path' '/forwarded'
354+
defaultTags()
355+
}
356+
}
344357
} else {
345358
throw new UnsupportedOperationException("responseSpan not implemented for " + endpoint)
346359
}

‎dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/test/boot/TestController.groovy

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import org.springframework.web.bind.annotation.ResponseBody
1919
import org.springframework.web.servlet.view.RedirectView
2020

2121
import javax.servlet.http.HttpServletRequest
22+
import java.util.concurrent.CompletableFuture
2223

2324
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_JSON
2425
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_URLENCODED
@@ -48,9 +49,11 @@ class TestController {
4849

4950
@RequestMapping("/forwarded")
5051
@ResponseBody
51-
String forwarded(HttpServletRequest request) {
52-
HttpServerTest.controller(FORWARDED) {
53-
request.getHeader("x-forwarded-for")
52+
CompletableFuture<String> forwarded(HttpServletRequest request) {
53+
CompletableFuture.supplyAsync {
54+
HttpServerTest.controller(FORWARDED) {
55+
request.getHeader("x-forwarded-for")
56+
}
5457
}
5558
}
5659

‎dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java

Lines changed: 36 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
88
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
99
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE;
10+
import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_CONTINUE_SUFFIX;
11+
import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_PREFIX_KEY;
1012
import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DECORATE;
1113
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
1214
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
@@ -22,6 +24,7 @@
2224
import net.bytebuddy.asm.Advice;
2325
import net.bytebuddy.description.type.TypeDescription;
2426
import net.bytebuddy.matcher.ElementMatcher;
27+
import org.springframework.web.method.HandlerMethod;
2528

2629
@AutoService(InstrumenterModule.class)
2730
public final class HandlerAdapterInstrumentation extends InstrumenterModule.Tracing
@@ -64,7 +67,10 @@ public static class ControllerAdvice {
6467
@Advice.OnMethodEnter(suppress = Throwable.class)
6568
public static AgentScope nameResourceAndStartSpan(
6669
@Advice.Argument(0) final HttpServletRequest request,
67-
@Advice.Argument(2) final Object handler) {
70+
@Advice.Argument(2) final Object handler,
71+
@Advice.Local("handlerSpanKey") String handlerSpanKey) {
72+
handlerSpanKey = "";
73+
6874
// Name the parent span based on the matching pattern
6975
Object parentSpan = request.getAttribute(DD_SPAN_ATTRIBUTE);
7076
if (parentSpan instanceof AgentSpan) {
@@ -77,24 +83,47 @@ public static AgentScope nameResourceAndStartSpan(
7783

7884
// Now create a span for handler/controller execution.
7985

86+
final String handlerKey;
87+
if (handler instanceof HandlerMethod) {
88+
handlerKey = ((HandlerMethod) handler).getBean().getClass().getName();
89+
} else {
90+
handlerKey = handler.getClass().getName();
91+
}
92+
handlerSpanKey = DD_HANDLER_SPAN_PREFIX_KEY + handlerKey;
93+
final Object existingSpan = request.getAttribute(handlerSpanKey);
94+
if (existingSpan instanceof AgentSpan) {
95+
return activateSpan((AgentSpan) existingSpan);
96+
}
97+
8098
final AgentSpan span = startSpan(DECORATE.spanName()).setMeasured(true);
8199
DECORATE.afterStart(span);
82100
DECORATE.onHandle(span, handler);
83-
101+
request.setAttribute(handlerSpanKey, span);
84102
return activateSpan(span);
85103
}
86104

87105
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
88106
public static void stopSpan(
89-
@Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) {
107+
@Advice.Argument(0) final HttpServletRequest request,
108+
@Advice.Enter final AgentScope scope,
109+
@Advice.Thrown final Throwable throwable,
110+
@Advice.Local("handlerSpanKey") String handlerSpanKey) {
90111
if (scope == null) {
91112
return;
92113
}
93-
94-
DECORATE.onError(scope, throwable);
95-
DECORATE.beforeFinish(scope);
114+
boolean finish =
115+
!Boolean.TRUE.equals(
116+
request.getAttribute(handlerSpanKey + DD_HANDLER_SPAN_CONTINUE_SUFFIX));
117+
final AgentSpan span = scope.span();
96118
scope.close();
97-
scope.span().finish();
119+
if (throwable != null) {
120+
DECORATE.onError(span, throwable);
121+
finish = true;
122+
}
123+
if (finish) {
124+
DECORATE.beforeFinish(span);
125+
span.finish();
126+
}
98127
}
99128
}
100129
}

‎dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/InvocableHandlerMethodInstrumentation.java

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
package datadog.trace.instrumentation.springweb;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_CONTINUE_SUFFIX;
5+
import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_PREFIX_KEY;
6+
7+
import com.google.auto.service.AutoService;
8+
import datadog.trace.agent.tooling.Instrumenter;
9+
import datadog.trace.agent.tooling.InstrumenterModule;
10+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
11+
import datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtensions;
12+
import java.util.concurrent.CompletionStage;
13+
import net.bytebuddy.asm.Advice;
14+
import net.bytebuddy.implementation.bytecode.assign.Assigner;
15+
import org.springframework.web.context.request.NativeWebRequest;
16+
import org.springframework.web.context.request.ServletWebRequest;
17+
import org.springframework.web.method.support.InvocableHandlerMethod;
18+
19+
@AutoService(InstrumenterModule.class)
20+
public class InvocableHandlerMethodInstrumentation extends InstrumenterModule.Tracing
21+
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
22+
public InvocableHandlerMethodInstrumentation() {
23+
super("spring-web");
24+
}
25+
26+
@Override
27+
public String instrumentedType() {
28+
return "org.springframework.web.method.support.InvocableHandlerMethod";
29+
}
30+
31+
@Override
32+
public void methodAdvice(MethodTransformer transformer) {
33+
transformer.applyAdvice(
34+
named("invokeForRequest"), getClass().getName() + "$WrapContinuableResultAdvice");
35+
}
36+
37+
@Override
38+
public String[] helperClassNames() {
39+
return new String[] {
40+
packageName + ".SpringWebHttpServerDecorator", packageName + ".ServletRequestURIAdapter",
41+
};
42+
}
43+
44+
public static class WrapContinuableResultAdvice {
45+
@Advice.OnMethodExit(suppress = Throwable.class)
46+
public static void after(
47+
@Advice.Argument(value = 0, typing = Assigner.Typing.DYNAMIC)
48+
final NativeWebRequest nativeWebRequest,
49+
@Advice.Return(readOnly = false) Object result,
50+
@Advice.This final InvocableHandlerMethod self) {
51+
if (!(nativeWebRequest instanceof ServletWebRequest)
52+
|| !(result instanceof CompletionStage)) {
53+
return;
54+
}
55+
ServletWebRequest servletWebRequest = (ServletWebRequest) nativeWebRequest;
56+
final String handlerSpanKey =
57+
DD_HANDLER_SPAN_PREFIX_KEY + self.getBean().getClass().getName();
58+
59+
if (Boolean.TRUE.equals(
60+
servletWebRequest.getAttribute(
61+
handlerSpanKey + DD_HANDLER_SPAN_CONTINUE_SUFFIX, ServletWebRequest.SCOPE_REQUEST))) {
62+
return;
63+
}
64+
65+
Object span = servletWebRequest.getAttribute(handlerSpanKey, ServletWebRequest.SCOPE_REQUEST);
66+
if (!(span instanceof AgentSpan)) {
67+
return;
68+
}
69+
servletWebRequest.setAttribute(
70+
handlerSpanKey + DD_HANDLER_SPAN_CONTINUE_SUFFIX, true, ServletWebRequest.SCOPE_REQUEST);
71+
result =
72+
((CompletionStage<?>) result)
73+
.whenComplete(AsyncResultExtensions.finishSpan((AgentSpan) span));
74+
}
75+
}
76+
}

‎dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/SpringWebHttpServerDecorator.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@ public class SpringWebHttpServerDecorator
3030
new SpringWebHttpServerDecorator(UTF8BytesString.create("spring-web-controller"));
3131
public static final SpringWebHttpServerDecorator DECORATE_RENDER =
3232
new SpringWebHttpServerDecorator(UTF8BytesString.create("spring-webmvc"));
33+
public static final String DD_HANDLER_SPAN_PREFIX_KEY = "dd.handler.span.";
34+
public static final String DD_HANDLER_SPAN_CONTINUE_SUFFIX = ".continue";
3335

3436
public SpringWebHttpServerDecorator(CharSequence component) {
3537
this.component = component;

‎dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import javax.servlet.http.HttpServletRequest
2929
import javax.servlet.http.HttpServletResponse
3030

3131
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.EXCEPTION
32+
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.FORWARDED
3233
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.LOGIN
3334
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.MATRIX_PARAM
3435
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.NOT_FOUND
@@ -348,7 +349,7 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
348349
}
349350

350351
boolean hasResponseSpan(ServerEndpoint endpoint) {
351-
return endpoint == REDIRECT || endpoint == NOT_FOUND || endpoint == LOGIN
352+
return endpoint == REDIRECT || endpoint == NOT_FOUND || endpoint == LOGIN || endpoint == FORWARDED
352353
}
353354

354355
@Override
@@ -399,6 +400,18 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
399400
defaultTags()
400401
}
401402
}
403+
} else if (endpoint == FORWARDED) {
404+
trace.span {
405+
serviceName expectedServiceName()
406+
operationName 'servlet.dispatch'
407+
resourceName 'servlet.dispatch'
408+
tags {
409+
"$Tags.COMPONENT" 'java-web-servlet-async-dispatcher'
410+
'servlet.context' "/$servletContext"
411+
'servlet.path' '/forwarded'
412+
defaultTags()
413+
}
414+
}
402415
} else {
403416
throw new UnsupportedOperationException("responseSpan not implemented for " + endpoint)
404417
}

‎dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/TestController.groovy

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import org.springframework.web.bind.annotation.ResponseBody
1919
import org.springframework.web.servlet.view.RedirectView
2020

2121
import javax.servlet.http.HttpServletRequest
22+
import java.util.concurrent.CompletableFuture
2223

2324
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_JSON
2425
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_URLENCODED
@@ -49,9 +50,11 @@ class TestController {
4950

5051
@RequestMapping("/forwarded")
5152
@ResponseBody
52-
String forwarded(HttpServletRequest request) {
53-
HttpServerTest.controller(FORWARDED) {
54-
request.getHeader("x-forwarded-for")
53+
CompletableFuture<String> forwarded(HttpServletRequest request) {
54+
CompletableFuture.supplyAsync {
55+
HttpServerTest.controller(FORWARDED) {
56+
request.getHeader("x-forwarded-for")
57+
}
5558
}
5659
}
5760

0 commit comments

Comments
 (0)