From daae198b0863c1176c96231d19dd714f0147c61b Mon Sep 17 00:00:00 2001 From: Tyler Benson Date: Thu, 30 Jan 2020 21:53:28 -0500 Subject: [PATCH] Set dispatcher span on request instead of clear Clearing the span caused traces to be broken up and reported independently when calling forward/include. --- .../RequestDispatcherInstrumentation.java | 13 +++-- .../test/groovy/RequestDispatcherTest.groovy | 54 +++++++++++++++++-- .../test/groovy/RequestDispatcherUtils.java | 23 +++++++- 3 files changed, 79 insertions(+), 11 deletions(-) diff --git a/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java b/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java index 86ddbd3c69..3048b700d8 100644 --- a/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java @@ -71,6 +71,7 @@ public final class RequestDispatcherInstrumentation extends Instrumenter.Default public static AgentScope start( @Advice.Origin("#m") final String method, @Advice.This final RequestDispatcher dispatcher, + @Advice.Local("_requestSpan") Object requestSpan, @Advice.Argument(0) final ServletRequest request) { if (activeSpan() == null) { // Don't want to generate a new top-level span @@ -87,8 +88,9 @@ public final class RequestDispatcherInstrumentation extends Instrumenter.Default // In case we lose context, inject trace into to the request. propagate().inject(span, request, SETTER); - // temporarily remove from request to avoid spring resource name bubbling up: - request.removeAttribute(DD_SPAN_ATTRIBUTE); + // temporarily replace from request to avoid spring resource name bubbling up: + requestSpan = request.getAttribute(DD_SPAN_ATTRIBUTE); + request.setAttribute(DD_SPAN_ATTRIBUTE, span); return activateSpan(span, true).setAsyncPropagation(true); } @@ -96,14 +98,17 @@ public final class RequestDispatcherInstrumentation extends Instrumenter.Default @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stop( @Advice.Enter final AgentScope scope, + @Advice.Local("_requestSpan") final Object requestSpan, @Advice.Argument(0) final ServletRequest request, @Advice.Thrown final Throwable throwable) { if (scope == null) { return; } - // now add it back... - request.setAttribute(DD_SPAN_ATTRIBUTE, scope.span()); + if (requestSpan != null) { + // now add it back... + request.setAttribute(DD_SPAN_ATTRIBUTE, requestSpan); + } DECORATE.onError(scope, throwable); DECORATE.beforeFinish(scope); diff --git a/dd-java-agent/instrumentation/servlet/src/test/groovy/RequestDispatcherTest.groovy b/dd-java-agent/instrumentation/servlet/src/test/groovy/RequestDispatcherTest.groovy index f352051753..9f272cb46f 100644 --- a/dd-java-agent/instrumentation/servlet/src/test/groovy/RequestDispatcherTest.groovy +++ b/dd-java-agent/instrumentation/servlet/src/test/groovy/RequestDispatcherTest.groovy @@ -1,15 +1,23 @@ +import datadog.opentracing.DDSpan import datadog.trace.agent.test.AgentTestRunner import javax.servlet.ServletException import javax.servlet.http.HttpServletRequest import javax.servlet.http.HttpServletResponse +import static datadog.opentracing.propagation.DatadogHttpCodec.SAMPLING_PRIORITY_KEY +import static datadog.opentracing.propagation.DatadogHttpCodec.SPAN_ID_KEY +import static datadog.opentracing.propagation.DatadogHttpCodec.TRACE_ID_KEY +import static datadog.trace.agent.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE import static datadog.trace.agent.test.utils.TraceUtils.basicSpan import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace class RequestDispatcherTest extends AgentTestRunner { - def dispatcher = new RequestDispatcherUtils(Mock(HttpServletRequest), Mock(HttpServletResponse)) + def request = Mock(HttpServletRequest) + def response = Mock(HttpServletResponse) + def mockSpan = Mock(DDSpan) + def dispatcher = new RequestDispatcherUtils(request, response) def "test dispatch no-parent"() { when: @@ -17,7 +25,17 @@ class RequestDispatcherTest extends AgentTestRunner { dispatcher.include("") then: - assertTraces(0) {} + assertTraces(2) { + trace(0, 1) { + basicSpan(it, 0, "forward-child") + } + trace(1, 1) { + basicSpan(it, 0, "include-child") + } + } + + and: + 0 * _ } def "test dispatcher #method with parent"() { @@ -28,7 +46,7 @@ class RequestDispatcherTest extends AgentTestRunner { then: assertTraces(1) { - trace(0, 2) { + trace(0, 3) { basicSpan(it, 0, "parent") span(1) { operationName "servlet.$operation" @@ -39,9 +57,22 @@ class RequestDispatcherTest extends AgentTestRunner { defaultTags() } } + basicSpan(it, 2, "$operation-child", span(1)) } } + then: + 1 * request.setAttribute(TRACE_ID_KEY, _) + 1 * request.setAttribute(SPAN_ID_KEY, _) + 1 * request.setAttribute(SAMPLING_PRIORITY_KEY, _) + then: + 1 * request.getAttribute(DD_SPAN_ATTRIBUTE) >> mockSpan + then: + 1 * request.setAttribute(DD_SPAN_ATTRIBUTE, { it.spanName == "servlet.$operation" }) + then: + 1 * request.setAttribute(DD_SPAN_ATTRIBUTE, mockSpan) + 0 * _ + where: operation | method "forward" | "forward" @@ -55,7 +86,7 @@ class RequestDispatcherTest extends AgentTestRunner { def "test dispatcher #method exception"() { setup: def ex = new ServletException("some error") - def dispatcher = new RequestDispatcherUtils(Mock(HttpServletRequest), Mock(HttpServletResponse), ex) + def dispatcher = new RequestDispatcherUtils(request, response, ex) when: runUnderTrace("parent") { @@ -67,7 +98,7 @@ class RequestDispatcherTest extends AgentTestRunner { th == ex assertTraces(1) { - trace(0, 2) { + trace(0, 3) { basicSpan(it, 0, "parent", null, ex) span(1) { operationName "servlet.$operation" @@ -80,9 +111,22 @@ class RequestDispatcherTest extends AgentTestRunner { errorTags(ex.class, ex.message) } } + basicSpan(it, 2, "$operation-child", span(1)) } } + then: + 1 * request.setAttribute(TRACE_ID_KEY, _) + 1 * request.setAttribute(SPAN_ID_KEY, _) + 1 * request.setAttribute(SAMPLING_PRIORITY_KEY, _) + then: + 1 * request.getAttribute(DD_SPAN_ATTRIBUTE) >> mockSpan + then: + 1 * request.setAttribute(DD_SPAN_ATTRIBUTE, { it.spanName == "servlet.$operation" }) + then: + 1 * request.setAttribute(DD_SPAN_ATTRIBUTE, mockSpan) + 0 * _ + where: operation | method "forward" | "forward" diff --git a/dd-java-agent/instrumentation/servlet/src/test/groovy/RequestDispatcherUtils.java b/dd-java-agent/instrumentation/servlet/src/test/groovy/RequestDispatcherUtils.java index 2f08ad74ed..2f102b4708 100644 --- a/dd-java-agent/instrumentation/servlet/src/test/groovy/RequestDispatcherUtils.java +++ b/dd-java-agent/instrumentation/servlet/src/test/groovy/RequestDispatcherUtils.java @@ -1,9 +1,12 @@ +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace; + import java.io.IOException; import java.io.InputStream; import java.net.MalformedURLException; import java.net.URL; import java.util.Enumeration; import java.util.Set; +import java.util.concurrent.Callable; import javax.servlet.RequestDispatcher; import javax.servlet.Servlet; import javax.servlet.ServletContext; @@ -164,7 +167,15 @@ public class RequestDispatcherUtils { class TestDispatcher implements RequestDispatcher { @Override public void forward(final ServletRequest servletRequest, final ServletResponse servletResponse) - throws ServletException, IOException { + throws ServletException { + runUnderTrace( + "forward-child", + new Callable() { + @Override + public Object call() throws Exception { + return null; + } + }); if (toThrow != null) { throw toThrow; } @@ -172,7 +183,15 @@ public class RequestDispatcherUtils { @Override public void include(final ServletRequest servletRequest, final ServletResponse servletResponse) - throws ServletException, IOException { + throws ServletException { + runUnderTrace( + "include-child", + new Callable() { + @Override + public Object call() throws Exception { + return null; + } + }); if (toThrow != null) { throw toThrow; }