Set dispatcher span on request instead of clear

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

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);

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"

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<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;
       }