Add instrumentation to detect the route at the beginning of the spring request

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java

@@ -19,6 +19,8 @@ public interface AgentSpan {
 
   AgentSpan getLocalRootSpan();
 
+  boolean isSameTrace(AgentSpan otherSpan);
+
   Context context();
 
   void finish();

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java

@@ -182,6 +182,12 @@ public AgentSpan getLocalRootSpan() {
       return this;
     }
 
+    @Override
+    public boolean isSameTrace(final AgentSpan otherSpan) {
+      // Not sure if this is the best idea...
+      return otherSpan instanceof NoopAgentSpan;
+    }
+
     @Override
     public Context context() {
       return NoopContext.INSTANCE;

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/OpenTracing32.java

@@ -198,6 +198,20 @@ public AgentSpan getLocalRootSpan() {
       return this;
     }
 
+    @Override
+    public boolean isSameTrace(final AgentSpan otherAgentSpan) {
+      if (otherAgentSpan instanceof OT32Span) {
+        final Span otherSpan = ((OT32Span) otherAgentSpan).span;
+        if (span instanceof DDSpan && otherSpan instanceof DDSpan) {
+          // minor optimization to avoid BigInteger.toString()
+          return ((DDSpan) span).getTraceId().equals(((DDSpan) otherSpan).getTraceId());
+        } else {
+          return span.context().toTraceId().equals(otherSpan.context().toTraceId());
+        }
+      }
+      return false;
+    }
+
     @Override
     public OT32Context context() {
       final SpanContext context = span.context();

dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Decorator.java

@@ -2,12 +2,17 @@
 
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator;
+import java.io.IOException;
 import java.net.URI;
 import java.net.URISyntaxException;
+import javax.servlet.Filter;
+import javax.servlet.ServletContext;
 import javax.servlet.ServletException;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
+import lombok.extern.slf4j.Slf4j;
 
+@Slf4j
 public class Servlet3Decorator
     extends HttpServerDecorator<HttpServletRequest, HttpServletRequest, HttpServletResponse> {
   public static final Servlet3Decorator DECORATE = new Servlet3Decorator();
@@ -58,12 +63,32 @@ protected Integer status(final HttpServletResponse httpServletResponse) {
   public AgentSpan onRequest(final AgentSpan span, final HttpServletRequest request) {
     assert span != null;
     if (request != null) {
-      span.setTag("servlet.context", request.getContextPath());
       span.setTag("servlet.path", request.getServletPath());
+      span.setTag("servlet.context", request.getContextPath());
+      onContext(span, request, request.getServletContext());
     }
     return super.onRequest(span, request);
   }
 
+  /**
+   * This method executes the filter created by
+   * datadog.trace.instrumentation.springweb.DispatcherServletInstrumentation$HandlerMappingAdvice.
+   * This was easier and less "hacky" than other ways to add the filter to the front of the filter
+   * chain.
+   */
+  private void onContext(
+      final AgentSpan span, final HttpServletRequest request, final ServletContext context) {
+    final Object attribute = context.getAttribute("dd.dispatcher-filter");
+    if (attribute instanceof Filter) {
+      request.setAttribute(DD_SPAN_ATTRIBUTE, span);
+      try {
+        ((Filter) attribute).doFilter(request, null, null);
+      } catch (final IOException | ServletException e) {
+        log.debug("Exception unexpectedly thrown by filter", e);
+      }
+    }
+  }
+
   @Override
   public AgentSpan onError(final AgentSpan span, final Throwable throwable) {
     if (throwable instanceof ServletException && throwable.getCause() != null) {

dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java

@@ -79,12 +79,27 @@ public static AgentScope start(
         @Advice.This final RequestDispatcher dispatcher,
         @Advice.Local("_requestSpan") Object requestSpan,
         @Advice.Argument(0) final ServletRequest request) {
-      if (activeSpan() == null) {
+      final AgentSpan parentSpan = activeSpan();
+
+      final Object servletSpanObject = request.getAttribute(DD_SPAN_ATTRIBUTE);
+      final AgentSpan servletSpan =
+          servletSpanObject instanceof AgentSpan ? (AgentSpan) servletSpanObject : null;
+
+      if (parentSpan == null && servletSpan == null) {
         // Don't want to generate a new top-level span
         return null;
       }
+      final AgentSpan.Context parent;
+      if (servletSpan == null || (parentSpan != null && servletSpan.isSameTrace(parentSpan))) {
+        // Use the parentSpan if the servletSpan is null or part of the same trace.
+        parent = parentSpan.context();
+      } else {
+        // parentSpan is part of a different trace, so lets ignore it.
+        // This can happen with the way Tomcat does error handling.
+        parent = servletSpan.context();
+      }
 
-      final AgentSpan span = startSpan("servlet." + method);
+      final AgentSpan span = startSpan("servlet." + method, parent);
       DECORATE.afterStart(span);
 
       final String target =

dd-java-agent/instrumentation/servlet/src/test/groovy/RequestDispatcherTest.groovy

@@ -25,6 +25,7 @@ class RequestDispatcherTest extends AgentTestRunner {
     dispatcher.include("")
 
     then:
+    2 * request.getAttribute(DD_SPAN_ATTRIBUTE)
     assertTraces(2) {
       trace(0, 1) {
         basicSpan(it, 0, "forward-child")
@@ -45,6 +46,7 @@ class RequestDispatcherTest extends AgentTestRunner {
     }
 
     then:
+    1 * request.getAttribute(DD_SPAN_ATTRIBUTE)
     assertTraces(1) {
       trace(0, 3) {
         basicSpan(it, 0, "parent")
@@ -97,6 +99,7 @@ class RequestDispatcherTest extends AgentTestRunner {
     def th = thrown(ServletException)
     th == ex
 
+    1 * request.getAttribute(DD_SPAN_ATTRIBUTE)
     assertTraces(1) {
       trace(0, 3) {
         basicSpan(it, 0, "parent", null, ex)

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

@@ -48,6 +48,7 @@ dependencies {
   }
 
   // Include servlet instrumentation for verifying the tomcat requests
+  testCompile project(':dd-java-agent:instrumentation:servlet')
   testCompile project(':dd-java-agent:instrumentation:servlet:request-3')
 
   testCompile group: 'javax.validation', name: 'validation-api', version: '1.1.0.Final'

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

@@ -5,23 +5,31 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DECORATE;
 import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DECORATE_RENDER;
+import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isProtected;
 import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
+import javax.servlet.ServletContext;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
-import org.springframework.web.method.HandlerMethod;
+import org.springframework.context.ApplicationContext;
+import org.springframework.web.servlet.DispatcherServlet;
+import org.springframework.web.servlet.HandlerMapping;
 import org.springframework.web.servlet.ModelAndView;
 
 @AutoService(Instrumenter.class)
@@ -36,23 +44,38 @@ public ElementMatcher<TypeDescription> typeMatcher() {
     return named("org.springframework.web.servlet.DispatcherServlet");
   }
 
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap(
+        "org.springframework.web.servlet.DispatcherServlet",
+        packageName + ".HandlerMappingResourceNameFilter");
+  }
+
   @Override
   public String[] helperClassNames() {
     return new String[] {
       packageName + ".SpringWebHttpServerDecorator",
       packageName + ".SpringWebHttpServerDecorator$1",
+      packageName + ".HandlerMappingResourceNameFilter",
     };
   }
 
   @Override
   public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
     final Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
+    transformers.put(
+        isMethod()
+            .and(isProtected())
+            .and(named("onRefresh"))
+            .and(takesArgument(0, named("org.springframework.context.ApplicationContext")))
+            .and(takesArguments(1)),
+        DispatcherServletInstrumentation.class.getName() + "$HandlerMappingAdvice");
     transformers.put(
         isMethod()
             .and(isProtected())
             .and(named("render"))
             .and(takesArgument(0, named("org.springframework.web.servlet.ModelAndView"))),
-        DispatcherServletInstrumentation.class.getName() + "$DispatcherAdvice");
+        DispatcherServletInstrumentation.class.getName() + "$RenderAdvice");
     transformers.put(
         isMethod()
             .and(isProtected())
@@ -62,7 +85,37 @@ public Map<? extends ElementMatcher<? super MethodDescription>, String> transfor
     return transformers;
   }
 
-  public static class DispatcherAdvice {
+  /**
+   * This advice creates a filter that has reference to the handlerMappings from DispatcherServlet
+   * which allows the mappings to be evaluated at the beginning of the filter chain. This evaluation
+   * is done inside the Servlet3Decorator.onContext method.
+   */
+  public static class HandlerMappingAdvice {
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void afterRefresh(
+        @Advice.This final DispatcherServlet dispatcher,
+        @Advice.Argument(0) final ApplicationContext springCtx,
+        @Advice.FieldValue("handlerMappings") final List<HandlerMapping> handlerMappings,
+        @Advice.Thrown final Throwable throwable) {
+      final ServletContext servletContext = springCtx.getBean(ServletContext.class);
+      if (handlerMappings != null && servletContext != null) {
+        final ContextStore<DispatcherServlet, HandlerMappingResourceNameFilter> contextStore =
+            InstrumentationContext.get(
+                DispatcherServlet.class, HandlerMappingResourceNameFilter.class);
+        HandlerMappingResourceNameFilter filter = contextStore.get(dispatcher);
+        if (filter == null) {
+          filter = new HandlerMappingResourceNameFilter();
+          contextStore.put(dispatcher, filter);
+        }
+        filter.setHandlerMappings(handlerMappings);
+        servletContext.setAttribute(
+            "dd.dispatcher-filter", filter); // used by Servlet3Decorator.onContext
+      }
+    }
+  }
+
+  public static class RenderAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static AgentScope onEnter(@Advice.Argument(0) final ModelAndView mv) {
@@ -79,11 +132,6 @@ public static void stopSpan(
       DECORATE_RENDER.beforeFinish(scope);
       scope.close();
     }
-
-    // Make this advice match consistently with HandlerAdapterInstrumentation
-    private void muzzleCheck(final HandlerMethod method) {
-      method.getMethod();
-    }
   }
 
   public static class ErrorHandlerAdvice {

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

@@ -0,0 +1,65 @@
+package datadog.trace.instrumentation.springweb;
+
+import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE;
+import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DECORATE;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import java.util.List;
+import javax.servlet.Filter;
+import javax.servlet.FilterChain;
+import javax.servlet.FilterConfig;
+import javax.servlet.ServletRequest;
+import javax.servlet.ServletResponse;
+import javax.servlet.http.HttpServletRequest;
+import org.springframework.web.servlet.HandlerExecutionChain;
+import org.springframework.web.servlet.HandlerMapping;
+
+public class HandlerMappingResourceNameFilter implements Filter {
+  private volatile List<HandlerMapping> handlerMappings;
+
+  @Override
+  public void init(final FilterConfig filterConfig) {}
+
+  @Override
+  public void doFilter(
+      final ServletRequest servletRequest,
+      final ServletResponse servletResponse,
+      final FilterChain filterChain) {
+    if (servletRequest instanceof HttpServletRequest && handlerMappings != null) {
+      final HttpServletRequest request = (HttpServletRequest) servletRequest;
+      try {
+        if (findMapping(request)) {
+          // Name the parent span based on the matching pattern
+          final Object parentSpan = request.getAttribute(DD_SPAN_ATTRIBUTE);
+          if (parentSpan instanceof AgentSpan) {
+            // Let the parent span resource name be set with the attribute set in findMapping.
+            DECORATE.onRequest((AgentSpan) parentSpan, request);
+          }
+        }
+      } catch (final Exception e) {
+      }
+    }
+  }
+
+  @Override
+  public void destroy() {}
+
+  /**
+   * When a HandlerMapping matches a request, it sets HandlerMapping.BEST_MATCHING_PATTERN_ATTRIBUTE
+   * as an attribute on the request. This attribute is read by
+   * SpringWebHttpServerDecorator.onRequest and set as the resource name.
+   */
+  private boolean findMapping(final HttpServletRequest request) throws Exception {
+    for (final HandlerMapping mapping : handlerMappings) {
+      final HandlerExecutionChain handler = mapping.getHandler(request);
+      if (handler != null) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  public void setHandlerMappings(final List<HandlerMapping> handlerMappings) {
+    this.handlerMappings = handlerMappings;
+  }
+}

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

@@ -1,4 +1,4 @@
-package test
+package test.boot
 
 import org.apache.catalina.connector.Connector
 import org.springframework.boot.autoconfigure.SpringBootApplication
@@ -25,7 +25,7 @@ class AppConfig extends WebMvcConfigurerAdapter {
       .defaultContentTypeStrategy(new ContentNegotiationStrategy() {
         @Override
         List<MediaType> resolveMediaTypes(NativeWebRequest webRequest) throws HttpMediaTypeNotAcceptableException {
-          return [MediaType.TEXT_PLAIN, MediaType.APPLICATION_JSON]
+          return [MediaType.TEXT_PLAIN]
         }
       })
   }