diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/AsyncResultExtensions.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/AsyncResultExtensions.java index 4feec759ca6..1c33522a530 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/AsyncResultExtensions.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/AsyncResultExtensions.java @@ -61,17 +61,17 @@ public Object apply(Object result, AgentSpan span) { } return null; } + } - private BiConsumer finishSpan(AgentSpan span) { - return (o, throwable) -> { - if (throwable != null) { - span.addThrowable( - throwable instanceof ExecutionException || throwable instanceof CompletionException - ? throwable.getCause() - : throwable); - } - span.finish(); - }; - } + public static BiConsumer finishSpan(AgentSpan span) { + return (o, throwable) -> { + if (throwable != null) { + span.addThrowable( + throwable instanceof ExecutionException || throwable instanceof CompletionException + ? throwable.getCause() + : throwable); + } + span.finish(); + }; } } diff --git a/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie b/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie index 8fad583b62f..95b54a47715 100644 --- a/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie +++ b/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie @@ -357,6 +357,7 @@ 0 org.springframework.util.StreamUtils # Included for IAST Spring mvc unvalidated redirect and xss vulnerability 0 org.springframework.web.method.support.HandlerMethodReturnValueHandlerComposite +0 org.springframework.web.method.support.InvocableHandlerMethod 2 org.xml.* 2 org.yaml.snakeyaml.* # Need for IAST sink diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/build.gradle b/dd-java-agent/instrumentation/spring-webmvc-3.1/build.gradle index 0fb7ea9927d..dfb49f6f882 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/build.gradle +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/build.gradle @@ -3,27 +3,8 @@ muzzle { group = 'org.springframework' module = 'spring-webmvc' versions = "[3.1.0.RELEASE,6)" - skipVersions += [ - '1.2.1', - '1.2.2', - '1.2.3', - '1.2.4'] // broken releases... missing dependencies skipVersions += '3.2.1.RELEASE' // missing a required class. (bad release?) extraDependency "javax.servlet:javax.servlet-api:3.0.1" - // assertInverse = true - } - - // FIXME: webmvc depends on web, so we need a separate integration for spring-web specifically. - fail { - group = 'org.springframework' - module = 'spring-web' - versions = "[,]" - skipVersions += [ - '1.2.1', - '1.2.2', - '1.2.3', - '1.2.4'] // broken releases... missing dependencies - extraDependency "javax.servlet:javax.servlet-api:3.0.1" } pass { diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/test/boot/AppConfig.groovy b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/test/boot/AppConfig.groovy index 88b756f210f..61fac0edede 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/test/boot/AppConfig.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/test/boot/AppConfig.groovy @@ -8,6 +8,7 @@ import org.springframework.boot.web.server.WebServerFactoryCustomizer import org.springframework.boot.web.servlet.filter.OrderedRequestContextFilter import org.springframework.context.annotation.Bean import org.springframework.core.Ordered +import org.springframework.scheduling.annotation.EnableAsync import org.springframework.web.filter.RequestContextFilter import org.springframework.web.servlet.config.annotation.PathMatchConfigurer import org.springframework.web.servlet.config.annotation.WebMvcConfigurerAdapter @@ -15,6 +16,7 @@ import org.springframework.web.util.UrlPathHelper // Component scan defeats the purpose of configuring with specific classes @SpringBootApplication(scanBasePackages = "doesnotexist") +@EnableAsync class AppConfig extends WebMvcConfigurerAdapter { @Bean diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/test/boot/SpringBootBasedTest.groovy b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/test/boot/SpringBootBasedTest.groovy index 0751a2ff471..25f579997e0 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/test/boot/SpringBootBasedTest.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/test/boot/SpringBootBasedTest.groovy @@ -27,6 +27,7 @@ import javax.servlet.http.HttpServletRequest import javax.servlet.http.HttpServletResponse import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.EXCEPTION +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.FORWARDED import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.LOGIN import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.MATRIX_PARAM import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.NOT_FOUND @@ -290,7 +291,7 @@ class SpringBootBasedTest extends HttpServerTest } boolean hasResponseSpan(ServerEndpoint endpoint) { - return endpoint == REDIRECT || endpoint == NOT_FOUND || endpoint == LOGIN + return endpoint == REDIRECT || endpoint == NOT_FOUND || endpoint == LOGIN || endpoint == FORWARDED } @Override @@ -341,6 +342,18 @@ class SpringBootBasedTest extends HttpServerTest defaultTags() } } + } else if (endpoint == FORWARDED) { + trace.span { + serviceName expectedServiceName() + operationName 'servlet.dispatch' + resourceName 'servlet.dispatch' + tags { + "$Tags.COMPONENT" 'java-web-servlet-async-dispatcher' + 'servlet.context' "/$servletContext" + 'servlet.path' '/forwarded' + defaultTags() + } + } } else { throw new UnsupportedOperationException("responseSpan not implemented for " + endpoint) } diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/test/boot/TestController.groovy b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/test/boot/TestController.groovy index 62b0ae0b1e1..7c4d9a93005 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/test/boot/TestController.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/test/boot/TestController.groovy @@ -19,6 +19,7 @@ import org.springframework.web.bind.annotation.ResponseBody import org.springframework.web.servlet.view.RedirectView import javax.servlet.http.HttpServletRequest +import java.util.concurrent.CompletableFuture import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_JSON import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_URLENCODED @@ -48,9 +49,11 @@ class TestController { @RequestMapping("/forwarded") @ResponseBody - String forwarded(HttpServletRequest request) { - HttpServerTest.controller(FORWARDED) { - request.getHeader("x-forwarded-for") + CompletableFuture forwarded(HttpServletRequest request) { + CompletableFuture.supplyAsync { + HttpServerTest.controller(FORWARDED) { + request.getHeader("x-forwarded-for") + } } } diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java index 3722ce989d9..60e8e4944ef 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java @@ -7,6 +7,8 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_CONTINUE_SUFFIX; +import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_PREFIX_KEY; import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.isPublic; @@ -22,6 +24,7 @@ import net.bytebuddy.asm.Advice; import net.bytebuddy.description.type.TypeDescription; import net.bytebuddy.matcher.ElementMatcher; +import org.springframework.web.method.HandlerMethod; @AutoService(InstrumenterModule.class) public final class HandlerAdapterInstrumentation extends InstrumenterModule.Tracing @@ -64,7 +67,10 @@ public static class ControllerAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope nameResourceAndStartSpan( @Advice.Argument(0) final HttpServletRequest request, - @Advice.Argument(2) final Object handler) { + @Advice.Argument(2) final Object handler, + @Advice.Local("handlerSpanKey") String handlerSpanKey) { + handlerSpanKey = ""; + // Name the parent span based on the matching pattern Object parentSpan = request.getAttribute(DD_SPAN_ATTRIBUTE); if (parentSpan instanceof AgentSpan) { @@ -77,24 +83,47 @@ public static AgentScope nameResourceAndStartSpan( // Now create a span for handler/controller execution. + final String handlerKey; + if (handler instanceof HandlerMethod) { + handlerKey = ((HandlerMethod) handler).getBean().getClass().getName(); + } else { + handlerKey = handler.getClass().getName(); + } + handlerSpanKey = DD_HANDLER_SPAN_PREFIX_KEY + handlerKey; + final Object existingSpan = request.getAttribute(handlerSpanKey); + if (existingSpan instanceof AgentSpan) { + return activateSpan((AgentSpan) existingSpan); + } + final AgentSpan span = startSpan(DECORATE.spanName()).setMeasured(true); DECORATE.afterStart(span); DECORATE.onHandle(span, handler); - + request.setAttribute(handlerSpanKey, span); return activateSpan(span); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( - @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) { + @Advice.Argument(0) final HttpServletRequest request, + @Advice.Enter final AgentScope scope, + @Advice.Thrown final Throwable throwable, + @Advice.Local("handlerSpanKey") String handlerSpanKey) { if (scope == null) { return; } - - DECORATE.onError(scope, throwable); - DECORATE.beforeFinish(scope); + boolean finish = + !Boolean.TRUE.equals( + request.getAttribute(handlerSpanKey + DD_HANDLER_SPAN_CONTINUE_SUFFIX)); + final AgentSpan span = scope.span(); scope.close(); - scope.span().finish(); + if (throwable != null) { + DECORATE.onError(span, throwable); + finish = true; + } + if (finish) { + DECORATE.beforeFinish(span); + span.finish(); + } } } } diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/InvocableHandlerMethodInstrumentation.java b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/InvocableHandlerMethodInstrumentation.java new file mode 100644 index 00000000000..a4b26041f02 --- /dev/null +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/InvocableHandlerMethodInstrumentation.java @@ -0,0 +1,76 @@ +package datadog.trace.instrumentation.springweb; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_CONTINUE_SUFFIX; +import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_PREFIX_KEY; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtensions; +import java.util.concurrent.CompletionStage; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.implementation.bytecode.assign.Assigner; +import org.springframework.web.context.request.NativeWebRequest; +import org.springframework.web.context.request.ServletWebRequest; +import org.springframework.web.method.support.InvocableHandlerMethod; + +@AutoService(InstrumenterModule.class) +public class InvocableHandlerMethodInstrumentation extends InstrumenterModule.Tracing + implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { + public InvocableHandlerMethodInstrumentation() { + super("spring-web"); + } + + @Override + public String instrumentedType() { + return "org.springframework.web.method.support.InvocableHandlerMethod"; + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + named("invokeForRequest"), getClass().getName() + "$WrapContinuableResultAdvice"); + } + + @Override + public String[] helperClassNames() { + return new String[] { + packageName + ".SpringWebHttpServerDecorator", packageName + ".ServletRequestURIAdapter", + }; + } + + public static class WrapContinuableResultAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + public static void after( + @Advice.Argument(value = 0, typing = Assigner.Typing.DYNAMIC) + final NativeWebRequest nativeWebRequest, + @Advice.Return(readOnly = false) Object result, + @Advice.This final InvocableHandlerMethod self) { + if (!(nativeWebRequest instanceof ServletWebRequest) + || !(result instanceof CompletionStage)) { + return; + } + ServletWebRequest servletWebRequest = (ServletWebRequest) nativeWebRequest; + final String handlerSpanKey = + DD_HANDLER_SPAN_PREFIX_KEY + self.getBean().getClass().getName(); + + if (Boolean.TRUE.equals( + servletWebRequest.getAttribute( + handlerSpanKey + DD_HANDLER_SPAN_CONTINUE_SUFFIX, ServletWebRequest.SCOPE_REQUEST))) { + return; + } + + Object span = servletWebRequest.getAttribute(handlerSpanKey, ServletWebRequest.SCOPE_REQUEST); + if (!(span instanceof AgentSpan)) { + return; + } + servletWebRequest.setAttribute( + handlerSpanKey + DD_HANDLER_SPAN_CONTINUE_SUFFIX, true, ServletWebRequest.SCOPE_REQUEST); + result = + ((CompletionStage) result) + .whenComplete(AsyncResultExtensions.finishSpan((AgentSpan) span)); + } + } +} diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/SpringWebHttpServerDecorator.java b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/SpringWebHttpServerDecorator.java index da34fd37ce8..2b0bf4dd18e 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/SpringWebHttpServerDecorator.java +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/SpringWebHttpServerDecorator.java @@ -30,6 +30,8 @@ public class SpringWebHttpServerDecorator new SpringWebHttpServerDecorator(UTF8BytesString.create("spring-web-controller")); public static final SpringWebHttpServerDecorator DECORATE_RENDER = new SpringWebHttpServerDecorator(UTF8BytesString.create("spring-webmvc")); + public static final String DD_HANDLER_SPAN_PREFIX_KEY = "dd.handler.span."; + public static final String DD_HANDLER_SPAN_CONTINUE_SUFFIX = ".continue"; public SpringWebHttpServerDecorator(CharSequence component) { this.component = component; diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy index fc9ee36389c..a80247071ee 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy @@ -29,6 +29,7 @@ import javax.servlet.http.HttpServletRequest import javax.servlet.http.HttpServletResponse import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.EXCEPTION +import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.FORWARDED import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.LOGIN import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.MATRIX_PARAM import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.NOT_FOUND @@ -348,7 +349,7 @@ class SpringBootBasedTest extends HttpServerTest } boolean hasResponseSpan(ServerEndpoint endpoint) { - return endpoint == REDIRECT || endpoint == NOT_FOUND || endpoint == LOGIN + return endpoint == REDIRECT || endpoint == NOT_FOUND || endpoint == LOGIN || endpoint == FORWARDED } @Override @@ -399,6 +400,18 @@ class SpringBootBasedTest extends HttpServerTest defaultTags() } } + } else if (endpoint == FORWARDED) { + trace.span { + serviceName expectedServiceName() + operationName 'servlet.dispatch' + resourceName 'servlet.dispatch' + tags { + "$Tags.COMPONENT" 'java-web-servlet-async-dispatcher' + 'servlet.context' "/$servletContext" + 'servlet.path' '/forwarded' + defaultTags() + } + } } else { throw new UnsupportedOperationException("responseSpan not implemented for " + endpoint) } diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/TestController.groovy b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/TestController.groovy index 7a505d8be63..329e18ed6d6 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/TestController.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/TestController.groovy @@ -19,6 +19,7 @@ import org.springframework.web.bind.annotation.ResponseBody import org.springframework.web.servlet.view.RedirectView import javax.servlet.http.HttpServletRequest +import java.util.concurrent.CompletableFuture import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_JSON import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_URLENCODED @@ -49,9 +50,11 @@ class TestController { @RequestMapping("/forwarded") @ResponseBody - String forwarded(HttpServletRequest request) { - HttpServerTest.controller(FORWARDED) { - request.getHeader("x-forwarded-for") + CompletableFuture forwarded(HttpServletRequest request) { + CompletableFuture.supplyAsync { + HttpServerTest.controller(FORWARDED) { + request.getHeader("x-forwarded-for") + } } } diff --git a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java/datadog/trace/instrumentation/springweb6/InvocableHandlerMethodInstrumentation.java b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java/datadog/trace/instrumentation/springweb6/InvocableHandlerMethodInstrumentation.java new file mode 100644 index 00000000000..b74163d170b --- /dev/null +++ b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java/datadog/trace/instrumentation/springweb6/InvocableHandlerMethodInstrumentation.java @@ -0,0 +1,33 @@ +package datadog.trace.instrumentation.springweb6; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; + +@AutoService(InstrumenterModule.class) +public class InvocableHandlerMethodInstrumentation extends InstrumenterModule.Tracing + implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { + public InvocableHandlerMethodInstrumentation() { + super("spring-web"); + } + + @Override + public String instrumentedType() { + return "org.springframework.web.method.support.InvocableHandlerMethod"; + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + named("invokeForRequest"), packageName + ".WrapContinuableResultAdvice"); + } + + @Override + public String[] helperClassNames() { + return new String[] { + packageName + ".SpringWebHttpServerDecorator", packageName + ".ServletRequestURIAdapter", + }; + } +} diff --git a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java index f4b7e8937fe..4753681acd7 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java +++ b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java @@ -4,23 +4,28 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.instrumentation.springweb6.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_CONTINUE_SUFFIX; +import static datadog.trace.instrumentation.springweb6.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_PREFIX_KEY; +import static datadog.trace.instrumentation.springweb6.SpringWebHttpServerDecorator.DECORATE; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import jakarta.servlet.http.HttpServletRequest; import net.bytebuddy.asm.Advice; +import org.springframework.web.method.HandlerMethod; public class ControllerAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope nameResourceAndStartSpan( @Advice.Argument(0) final HttpServletRequest request, - @Advice.Argument(2) final Object handler) { + @Advice.Argument(2) final Object handler, + @Advice.Local("handlerSpanKey") String handlerSpanKey) { + handlerSpanKey = ""; // Name the parent span based on the matching pattern Object parentSpan = request.getAttribute(DD_SPAN_ATTRIBUTE); if (parentSpan instanceof AgentSpan) { - SpringWebHttpServerDecorator.DECORATE.onRequest( - (AgentSpan) parentSpan, request, request, null); + DECORATE.onRequest((AgentSpan) parentSpan, request, request, null); } if (activeSpan() == null) { @@ -29,24 +34,46 @@ public static AgentScope nameResourceAndStartSpan( // Now create a span for handler/controller execution. - final AgentSpan span = - startSpan(SpringWebHttpServerDecorator.DECORATE.spanName()).setMeasured(true); - SpringWebHttpServerDecorator.DECORATE.afterStart(span); - SpringWebHttpServerDecorator.DECORATE.onHandle(span, handler); + final String handlerKey; + if (handler instanceof HandlerMethod) { + handlerKey = ((HandlerMethod) handler).getBean().getClass().getName(); + } else { + handlerKey = handler.getClass().getName(); + } + handlerSpanKey = DD_HANDLER_SPAN_PREFIX_KEY + handlerKey; + final Object existingSpan = request.getAttribute(handlerSpanKey); + if (existingSpan instanceof AgentSpan) { + return activateSpan((AgentSpan) existingSpan); + } + final AgentSpan span = startSpan(DECORATE.spanName()).setMeasured(true); + DECORATE.afterStart(span); + DECORATE.onHandle(span, handler); + request.setAttribute(handlerSpanKey, span); return activateSpan(span); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( - @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) { + @Advice.Enter final AgentScope scope, + @Advice.Argument(0) final HttpServletRequest request, + @Advice.Thrown final Throwable throwable, + @Advice.Local("handlerSpanKey") String handlerSpanKey) { if (scope == null) { return; } - - SpringWebHttpServerDecorator.DECORATE.onError(scope, throwable); - SpringWebHttpServerDecorator.DECORATE.beforeFinish(scope); + boolean finish = + !Boolean.TRUE.equals( + request.getAttribute(handlerSpanKey + DD_HANDLER_SPAN_CONTINUE_SUFFIX)); + final AgentSpan span = scope.span(); scope.close(); - scope.span().finish(); + if (throwable != null) { + DECORATE.onError(span, throwable); + finish = true; + } + if (finish) { + DECORATE.beforeFinish(span); + span.finish(); + } } } diff --git a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/SpringWebHttpServerDecorator.java b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/SpringWebHttpServerDecorator.java index 91c25b90f1a..3d05ae9e1a0 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/SpringWebHttpServerDecorator.java +++ b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/SpringWebHttpServerDecorator.java @@ -34,6 +34,9 @@ public class SpringWebHttpServerDecorator public static final SpringWebHttpServerDecorator DECORATE_RENDER = new SpringWebHttpServerDecorator(UTF8BytesString.create("spring-webmvc")); + public static final String DD_HANDLER_SPAN_PREFIX_KEY = "dd.handler.span."; + public static final String DD_HANDLER_SPAN_CONTINUE_SUFFIX = ".continue"; + public SpringWebHttpServerDecorator(CharSequence component) { this.component = component; } diff --git a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/WrapContinuableResultAdvice.java b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/WrapContinuableResultAdvice.java new file mode 100644 index 00000000000..4939751da1b --- /dev/null +++ b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/WrapContinuableResultAdvice.java @@ -0,0 +1,44 @@ +package datadog.trace.instrumentation.springweb6; + +import static datadog.trace.instrumentation.springweb6.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_CONTINUE_SUFFIX; +import static datadog.trace.instrumentation.springweb6.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_PREFIX_KEY; + +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtensions; +import java.util.concurrent.CompletionStage; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.implementation.bytecode.assign.Assigner; +import org.springframework.web.context.request.NativeWebRequest; +import org.springframework.web.context.request.ServletWebRequest; +import org.springframework.web.method.support.InvocableHandlerMethod; + +public class WrapContinuableResultAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + public static void after( + @Advice.Argument(value = 0, typing = Assigner.Typing.DYNAMIC) + final NativeWebRequest nativeWebRequest, + @Advice.Return(readOnly = false) Object result, + @Advice.This final InvocableHandlerMethod self) { + if (!(nativeWebRequest instanceof ServletWebRequest) || !(result instanceof CompletionStage)) { + return; + } + + ServletWebRequest servletWebRequest = (ServletWebRequest) nativeWebRequest; + final String handlerSpanKey = DD_HANDLER_SPAN_PREFIX_KEY + self.getBean().getClass().getName(); + + if (Boolean.TRUE.equals( + servletWebRequest.getAttribute( + handlerSpanKey + DD_HANDLER_SPAN_CONTINUE_SUFFIX, ServletWebRequest.SCOPE_REQUEST))) { + return; + } + Object span = servletWebRequest.getAttribute(handlerSpanKey, ServletWebRequest.SCOPE_REQUEST); + if (!(span instanceof AgentSpan)) { + return; + } + servletWebRequest.setAttribute( + handlerSpanKey + DD_HANDLER_SPAN_CONTINUE_SUFFIX, true, ServletWebRequest.SCOPE_REQUEST); + result = + ((CompletionStage) result) + .whenComplete(AsyncResultExtensions.finishSpan((AgentSpan) span)); + } +} diff --git a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/TestController.groovy b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/TestController.groovy index 06ddc978215..29f41d2c101 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/TestController.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/TestController.groovy @@ -12,6 +12,8 @@ import org.springframework.web.servlet.view.RedirectView import jakarta.servlet.http.HttpServletRequest +import java.util.concurrent.CompletableFuture + import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.* @Controller @@ -27,9 +29,11 @@ class TestController { @RequestMapping("/forwarded") @ResponseBody - String forwarded(HttpServletRequest request) { - HttpServerTest.controller(FORWARDED) { - request.getHeader("x-forwarded-for") + CompletableFuture forwarded(HttpServletRequest request) { + CompletableFuture.supplyAsync { + HttpServerTest.controller(FORWARDED) { + request.getHeader("x-forwarded-for") + } } }