Support CompletableFuture on spring webmvc controllers

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 <T> BiConsumer<T, Throwable> finishSpan(AgentSpan span) {
-      return (o, throwable) -> {
-        if (throwable != null) {
-          span.addThrowable(
-              throwable instanceof ExecutionException || throwable instanceof CompletionException
-                  ? throwable.getCause()
-                  : throwable);
-        }
-        span.finish();
-      };
-    }
+  public static <T> BiConsumer<T, Throwable> finishSpan(AgentSpan span) {
+    return (o, throwable) -> {
+      if (throwable != null) {
+        span.addThrowable(
+            throwable instanceof ExecutionException || throwable instanceof CompletionException
+                ? throwable.getCause()
+                : throwable);
+      }
+      span.finish();
+    };
   }
 }

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

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 {

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

@@ -8,13 +8,15 @@ 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
 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

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<ConfigurableApplicationContext>
   }
 
   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<ConfigurableApplicationContext>
           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)
     }

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<String> forwarded(HttpServletRequest request) {
+    CompletableFuture.supplyAsync {
+      HttpServerTest.controller(FORWARDED) {
+        request.getHeader("x-forwarded-for")
+      }
     }
   }
 

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();
+      }
     }
   }
 }

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

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;

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<ConfigurableApplicationContext>
   }
 
   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<ConfigurableApplicationContext>
           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)
     }

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<String> forwarded(HttpServletRequest request) {
+    CompletableFuture.supplyAsync {
+      HttpServerTest.controller(FORWARDED) {
+        request.getHeader("x-forwarded-for")
+      }
     }
   }