Set dispatcher span on request instead of clear
Clearing the span caused traces to be broken up and reported independently when calling forward/include.
MODIFY
dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java
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;
}
+ if (requestSpan != null) {
// now add it back...
- request.setAttribute(DD_SPAN_ATTRIBUTE, scope.span());
+ request.setAttribute(DD_SPAN_ATTRIBUTE, requestSpan);
+ }
DECORATE.onError(scope, throwable);
DECORATE.beforeFinish(scope);

MODIFY
dd-java-agent/instrumentation/servlet/src/test/groovy/RequestDispatcherTest.groovy
dd-java-agent/instrumentation/servlet/src/test/groovy/RequestDispatcherTest.groovy
+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"

MODIFY
dd-java-agent/instrumentation/servlet/src/test/groovy/RequestDispatcherUtils.java
dd-java-agent/instrumentation/servlet/src/test/groovy/RequestDispatcherUtils.java
+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<Object>() {
+ @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<Object>() {
+ @Override
+ public Object call() throws Exception {
+ return null;
+ }
+ });
if (toThrow != null) {
throw toThrow;
}

Set dispatcher span on request instead of clear

Clearing the span caused traces to be broken up and reported independently when calling forward/include.