Set ignoreActiveSpan for Servlet 2 and capture status code

This change wraps the servlet response object in order to collect the status code.  This can be risky if the code is expecting a raw type object.  We are extending a standard “HttpServletResponseWrapper” though, so it should be easy to work around.

This change also moves to the standard HttpServerTest for Servlet 2.

dd-java-agent/instrumentation/servlet-2/servlet-2.gradle

@@ -11,6 +11,14 @@ muzzle {
 
 apply from: "${rootDir}/gradle/java.gradle"
 
+apply plugin: 'org.unbroken-dome.test-sets'
+
+testSets {
+  latestDepTest {
+    dirName = 'test'
+  }
+}
+
 dependencies {
   compileOnly group: 'javax.servlet', name: 'servlet-api', version: '2.3'
 
@@ -26,4 +34,7 @@ dependencies {
   }
   testCompile group: 'org.eclipse.jetty', name: 'jetty-server', version: '7.0.0.v20091005'
   testCompile group: 'org.eclipse.jetty', name: 'jetty-servlet', version: '7.0.0.v20091005'
+
+  latestDepTestCompile group: 'org.eclipse.jetty', name: 'jetty-server', version: '7.+'
+  latestDepTestCompile group: 'org.eclipse.jetty', name: 'jetty-servlet', version: '7.+'
 }

dd-java-agent/instrumentation/servlet-2/src/main/java/datadog/trace/instrumentation/servlet2/AbstractServlet2Instrumentation.java

@@ -24,9 +24,10 @@ public abstract class AbstractServlet2Instrumentation extends Instrumenter.Defau
       "datadog.trace.agent.decorator.BaseDecorator",
       "datadog.trace.agent.decorator.ServerDecorator",
       "datadog.trace.agent.decorator.HttpServerDecorator",
+      packageName + ".Servlet2Decorator",
       packageName + ".HttpServletRequestExtractAdapter",
       packageName + ".HttpServletRequestExtractAdapter$MultivaluedMapFlatIterator",
-      packageName + ".Servlet2Decorator",
+      packageName + ".StatusSavingHttpServletResponseWrapper",
     };
   }
 }

dd-java-agent/instrumentation/servlet-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java

@@ -8,23 +8,35 @@ import io.opentracing.Scope;
 import io.opentracing.Span;
 import io.opentracing.SpanContext;
 import io.opentracing.propagation.Format;
+import io.opentracing.tag.Tags;
 import io.opentracing.util.GlobalTracer;
 import java.security.Principal;
 import javax.servlet.ServletRequest;
 import javax.servlet.ServletResponse;
 import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.implementation.bytecode.assign.Assigner;
 
 public class Servlet2Advice {
+  public static final String SERVLET_SPAN = "datadog.servlet.span";
 
   @Advice.OnMethodEnter(suppress = Throwable.class)
   public static Scope startSpan(
-      @Advice.This final Object servlet, @Advice.Argument(0) final ServletRequest req) {
+      @Advice.This final Object servlet,
-    if (GlobalTracer.get().activeSpan() != null || !(req instanceof HttpServletRequest)) {
+      @Advice.Argument(0) final ServletRequest req,
+      @Advice.Argument(value = 1, readOnly = false, typing = Assigner.Typing.DYNAMIC)
+          ServletResponse resp) {
+    final Object spanAttr = req.getAttribute(SERVLET_SPAN);
+    if (!(req instanceof HttpServletRequest) || spanAttr != null) {
       // Tracing might already be applied by the FilterChain.  If so ignore this.
       return null;
     }
 
+    if (resp instanceof HttpServletResponse) {
+      resp = new StatusSavingHttpServletResponseWrapper((HttpServletResponse) resp);
+    }
+
     final HttpServletRequest httpServletRequest = (HttpServletRequest) req;
     final SpanContext extractedContext =
         GlobalTracer.get()
@@ -35,6 +47,7 @@ public class Servlet2Advice {
     final Scope scope =
         GlobalTracer.get()
             .buildSpan("servlet.request")
+            .ignoreActiveSpan()
             .asChildOf(extractedContext)
             .withTag("span.origin.type", servlet.getClass().getName())
             .startActive(true);
@@ -47,6 +60,8 @@ public class Servlet2Advice {
     if (scope instanceof TraceScope) {
       ((TraceScope) scope).setAsyncPropagation(true);
     }
+
+    req.setAttribute(SERVLET_SPAN, span);
     return scope;
   }
 
@@ -68,9 +83,18 @@ public class Servlet2Advice {
     }
 
     if (scope != null) {
-      DECORATE.onResponse(scope.span(), response);
+      final Span span = scope.span();
-      DECORATE.onError(scope.span(), throwable);
+      DECORATE.onResponse(span, response);
-      DECORATE.beforeFinish(scope.span());
+      if (throwable != null) {
+        if (response instanceof StatusSavingHttpServletResponseWrapper
+            && ((StatusSavingHttpServletResponseWrapper) response).status
+                == HttpServletResponse.SC_OK) {
+          // exception was thrown but status code wasn't set
+          Tags.HTTP_STATUS.set(span, 500);
+        }
+        DECORATE.onError(span, throwable);
+      }
+      DECORATE.beforeFinish(span);
 
       if (scope instanceof TraceScope) {
         ((TraceScope) scope).setAsyncPropagation(false);

dd-java-agent/instrumentation/servlet-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Decorator.java

@@ -43,14 +43,19 @@ public class Servlet2Decorator
 
   @Override
   protected Integer peerPort(final HttpServletRequest httpServletRequest) {
+    // HttpServletResponse doesn't have accessor for remote port.
     return null;
   }
 
   @Override
   protected Integer status(final ServletResponse httpServletResponse) {
+    if (httpServletResponse instanceof StatusSavingHttpServletResponseWrapper) {
+      return ((StatusSavingHttpServletResponseWrapper) httpServletResponse).status;
+    } else {
       // HttpServletResponse doesn't have accessor for status code.
       return null;
     }
+  }
 
   @Override
   public Span onRequest(final Span span, final HttpServletRequest request) {

dd-java-agent/instrumentation/servlet-2/src/main/java/datadog/trace/instrumentation/servlet2/StatusSavingHttpServletResponseWrapper.java

@@ -0,0 +1,37 @@
+package datadog.trace.instrumentation.servlet2;
+
+import java.io.IOException;
+import javax.servlet.http.HttpServletResponse;
+import javax.servlet.http.HttpServletResponseWrapper;
+
+public class StatusSavingHttpServletResponseWrapper extends HttpServletResponseWrapper {
+  public int status = 200;
+
+  public StatusSavingHttpServletResponseWrapper(final HttpServletResponse response) {
+    super(response);
+  }
+
+  @Override
+  public void sendError(final int status) throws IOException {
+    this.status = status;
+    super.sendError(status);
+  }
+
+  @Override
+  public void sendError(final int status, final String message) throws IOException {
+    this.status = status;
+    super.sendError(status, message);
+  }
+
+  @Override
+  public void setStatus(final int status) {
+    this.status = status;
+    super.setStatus(status);
+  }
+
+  @Override
+  public void setStatus(final int status, final String message) {
+    this.status = status;
+    super.setStatus(status, message);
+  }
+}

dd-java-agent/instrumentation/servlet-2/src/test/groovy/JettyServlet2Test.groovy

@@ -1,198 +1,117 @@
-import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.agent.test.asserts.TraceAssert
-import datadog.trace.agent.test.utils.OkHttpUtils
+import datadog.trace.agent.test.base.HttpServerTest
-import datadog.trace.agent.test.utils.PortUtils
 import datadog.trace.api.DDSpanTypes
-import datadog.trace.api.DDTags
+import datadog.trace.instrumentation.servlet2.Servlet2Decorator
-import okhttp3.Credentials
+import io.opentracing.tag.Tags
-import okhttp3.Interceptor
+import javax.servlet.http.HttpServletRequest
-import okhttp3.OkHttpClient
-import okhttp3.Request
-import okhttp3.Response
-import org.eclipse.jetty.http.HttpHeaders
-import org.eclipse.jetty.http.security.Constraint
-import org.eclipse.jetty.security.ConstraintMapping
-import org.eclipse.jetty.security.ConstraintSecurityHandler
-import org.eclipse.jetty.security.HashLoginService
-import org.eclipse.jetty.security.LoginService
-import org.eclipse.jetty.security.authentication.BasicAuthenticator
 import org.eclipse.jetty.server.Server
+import org.eclipse.jetty.server.handler.ErrorHandler
 import org.eclipse.jetty.servlet.ServletContextHandler
 
-class JettyServlet2Test extends AgentTestRunner {
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.AUTH_REQUIRED
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.ERROR
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.EXCEPTION
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.REDIRECT
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS
 
-  OkHttpClient client = OkHttpUtils.clientBuilder().addNetworkInterceptor(new Interceptor() {
+class JettyServlet2Test extends HttpServerTest<Servlet2Decorator> {
-    @Override
-    Response intercept(Interceptor.Chain chain) throws IOException {
-      def response = chain.proceed(chain.request())
-      TEST_WRITER.waitForTraces(1)
-      return response
-    }
-  })
-    .build()
 
-  int port
+  private static final CONTEXT = "ctx"
   private Server jettyServer
-  private ServletContextHandler servletContext
 
-  def setup() {
+  @Override
-    port = PortUtils.randomOpenPort()
+  void startServer(int port) {
     jettyServer = new Server(port)
-    servletContext = new ServletContextHandler()
+    jettyServer.connectors.each { it.resolveNames = true } // get localhost instead of 127.0.0.1
-    servletContext.contextPath = "/ctx"
+    ServletContextHandler servletContext = new ServletContextHandler(null, "/$CONTEXT")
+    servletContext.errorHandler = new ErrorHandler() {
+      protected void handleErrorPage(HttpServletRequest request, Writer writer, int code, String message) throws IOException {
+        Throwable th = (Throwable) request.getAttribute("javax.servlet.error.exception")
+        writer.write(th ? th.message : message)
+      }
+    }
 
-    ConstraintSecurityHandler security = setupAuthentication(jettyServer)
+    // FIXME: Add tests for security/authentication.
+//    ConstraintSecurityHandler security = setupAuthentication(jettyServer)
+//    servletContext.setSecurityHandler(security)
 
-    servletContext.setSecurityHandler(security)
+    servletContext.addServlet(TestServlet2.Sync, SUCCESS.path)
-    servletContext.addServlet(TestServlet2.Sync, "/sync")
+    servletContext.addServlet(TestServlet2.Sync, ERROR.path)
-    servletContext.addServlet(TestServlet2.Sync, "/auth/sync")
+    servletContext.addServlet(TestServlet2.Sync, EXCEPTION.path)
+    servletContext.addServlet(TestServlet2.Sync, REDIRECT.path)
+    servletContext.addServlet(TestServlet2.Sync, AUTH_REQUIRED.path)
 
     jettyServer.setHandler(servletContext)
     jettyServer.start()
   }
 
-  def cleanup() {
+  @Override
+  void stopServer() {
     jettyServer.stop()
     jettyServer.destroy()
   }
 
-  def "test #path servlet call (auth: #auth, distributed tracing: #distributedTracing)"() {
+  @Override
-    setup:
+  URI buildAddress() {
-    def requestBuilder = new Request.Builder()
+    return new URI("http://localhost:$port/$CONTEXT/")
-      .url("http://localhost:$port/ctx/$path")
-      .get()
-    if (distributedTracing) {
-      requestBuilder.header("x-datadog-trace-id", "123")
-      requestBuilder.header("x-datadog-parent-id", "456")
   }
-    if (auth) {
+
-      requestBuilder.header(HttpHeaders.AUTHORIZATION, Credentials.basic("user", "password"))
+  @Override
+  Servlet2Decorator decorator() {
+    return Servlet2Decorator.DECORATE
   }
-    def response = client.newCall(requestBuilder.build()).execute()
 
-    expect:
+  @Override
-    response.body().string().trim() == expectedResponse
+  String expectedServiceName() {
+    CONTEXT
+  }
 
-    assertTraces(1) {
+  @Override
-      trace(0, 1) {
+  String expectedOperationName() {
-        span(0) {
+    return "servlet.request"
-          if (distributedTracing) {
+  }
-            traceId "123"
+
-            parentId "456"
+  @Override
+  boolean testNotFound() {
+    false
+  }
+
+  // parent span must be cast otherwise it breaks debugging classloading (junit loads it early)
+  void serverSpan(TraceAssert trace, int index, String traceID = null, String parentID = null, String method = "GET", ServerEndpoint endpoint = SUCCESS) {
+    trace.span(index) {
+      serviceName expectedServiceName()
+      operationName expectedOperationName()
+      resourceName endpoint.status == 404 ? "404" : "$method ${endpoint.resolve(address).path}"
+      spanType DDSpanTypes.HTTP_SERVER
+      errored endpoint.errored
+      if (parentID != null) {
+        traceId traceID
+        parentId parentID
       } else {
         parent()
       }
-          serviceName "ctx"
-          operationName "servlet.request"
-          resourceName "GET /ctx/$path"
-          spanType DDSpanTypes.HTTP_SERVER
-          errored false
       tags {
-            "http.url" "http://localhost:$port/ctx/$path"
+        "servlet.context" "/$CONTEXT"
-            "http.method" "GET"
+        "span.origin.type" TestServlet2.Sync.name
-            "span.kind" "server"
-            "component" "java-web-servlet"
-            "peer.hostname" "127.0.0.1"
-            "peer.ipv4" "127.0.0.1"
-            "span.origin.type" "TestServlet2\$Sync"
-            "servlet.context" "/ctx"
-            if (auth) {
-              "$DDTags.USER_NAME" "user"
-            }
-            defaultTags(distributedTracing)
-          }
-        }
-      }
-    }
 
-    where:
+        defaultTags(true)
-    path        | expectedResponse | auth  | distributedTracing
+        "$Tags.COMPONENT.key" serverDecorator.component()
-    "sync"      | "Hello Sync"     | false | false
+        if (endpoint.errored) {
-    "auth/sync" | "Hello Sync"     | true  | false
+          "$Tags.ERROR.key" endpoint.errored
-    "sync"      | "Hello Sync"     | false | true
+          "error.msg" { it == null || it == EXCEPTION.body }
-    "auth/sync" | "Hello Sync"     | true  | true
+          "error.type" { it == null || it == Exception.name }
+          "error.stack" { it == null || it instanceof String }
         }
-
+        "$Tags.HTTP_STATUS.key" endpoint.status
-  def "test #path error servlet call"() {
+        "$Tags.HTTP_URL.key" "${endpoint.resolve(address)}"
-    setup:
+        "$Tags.PEER_HOSTNAME.key" "localhost"
-    def request = new Request.Builder()
+        // No peer port
-      .url("http://localhost:$port/ctx/$path?error=true")
+        "$Tags.PEER_HOST_IPV4.key" "127.0.0.1"
-      .get()
+        "$Tags.HTTP_METHOD.key" method
-      .build()
+        "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_SERVER
-    def response = client.newCall(request).execute()
-
-    expect:
-    response.body().string().trim() != expectedResponse
-
-    assertTraces(1) {
-      trace(0, 1) {
-        span(0) {
-          serviceName "ctx"
-          operationName "servlet.request"
-          resourceName "GET /ctx/$path"
-          spanType DDSpanTypes.HTTP_SERVER
-          errored true
-          parent()
-          tags {
-            "http.url" "http://localhost:$port/ctx/$path"
-            "http.method" "GET"
-            "span.kind" "server"
-            "component" "java-web-servlet"
-            "peer.hostname" "127.0.0.1"
-            "peer.ipv4" "127.0.0.1"
-            "span.origin.type" "TestServlet2\$Sync"
-            "servlet.context" "/ctx"
-            errorTags(RuntimeException, "some $path error")
-            defaultTags()
       }
     }
   }
-    }
-
-    where:
-    path   | expectedResponse
-    "sync" | "Hello Sync"
-  }
-
-  def "test #path non-throwing-error servlet call"() {
-    // This doesn't actually detect the error because we can't get the status code via the old servlet API.
-    setup:
-    def request = new Request.Builder()
-      .url("http://localhost:$port/ctx/$path?non-throwing-error=true")
-      .get()
-      .build()
-    def response = client.newCall(request).execute()
-
-    expect:
-    response.body().string().trim() != expectedResponse
-
-    assertTraces(1) {
-      trace(0, 1) {
-        span(0) {
-          serviceName "ctx"
-          operationName "servlet.request"
-          resourceName "GET /ctx/$path"
-          spanType DDSpanTypes.HTTP_SERVER
-          errored false
-          parent()
-          tags {
-            "http.url" "http://localhost:$port/ctx/$path"
-            "http.method" "GET"
-            "span.kind" "server"
-            "component" "java-web-servlet"
-            "peer.hostname" "127.0.0.1"
-            "peer.ipv4" "127.0.0.1"
-            "span.origin.type" "TestServlet2\$Sync"
-            "servlet.context" "/ctx"
-            defaultTags()
-          }
-        }
-      }
-    }
-
-    where:
-    path   | expectedResponse
-    "sync" | "Hello Sync"
-  }
 
   /**
    * Setup simple authentication for tests
@@ -204,26 +123,26 @@ class JettyServlet2Test extends AgentTestRunner {
    * @param jettyServer server to attach login service
    * @return SecurityHandler that can be assigned to servlet
    */
-  private ConstraintSecurityHandler setupAuthentication(Server jettyServer) {
+//  private ConstraintSecurityHandler setupAuthentication(Server jettyServer) {
-    ConstraintSecurityHandler security = new ConstraintSecurityHandler()
+//    ConstraintSecurityHandler security = new ConstraintSecurityHandler()
-
+//
-    Constraint constraint = new Constraint()
+//    Constraint constraint = new Constraint()
-    constraint.setName("auth")
+//    constraint.setName("auth")
-    constraint.setAuthenticate(true)
+//    constraint.setAuthenticate(true)
-    constraint.setRoles("role")
+//    constraint.setRoles("role")
-
+//
-    ConstraintMapping mapping = new ConstraintMapping()
+//    ConstraintMapping mapping = new ConstraintMapping()
-    mapping.setPathSpec("/auth/*")
+//    mapping.setPathSpec("/auth/*")
-    mapping.setConstraint(constraint)
+//    mapping.setConstraint(constraint)
-
+//
-    security.setConstraintMappings(mapping)
+//    security.setConstraintMappings(mapping)
-    security.setAuthenticator(new BasicAuthenticator())
+//    security.setAuthenticator(new BasicAuthenticator())
-
+//
-    LoginService loginService = new HashLoginService("TestRealm",
+//    LoginService loginService = new HashLoginService("TestRealm",
-      "src/test/resources/realm.properties")
+//      "src/test/resources/realm.properties")
-    security.setLoginService(loginService)
+//    security.setLoginService(loginService)
-    jettyServer.addBean(loginService)
+//    jettyServer.addBean(loginService)
-
+//
-    security
+//    security
-  }
+//  }
 }

dd-java-agent/instrumentation/servlet-2/src/test/groovy/ServletTestInstrumentation.java

@@ -0,0 +1,28 @@
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.test.base.HttpServerTestAdvice;
+import datadog.trace.agent.tooling.Instrumenter;
+import net.bytebuddy.agent.builder.AgentBuilder;
+
+@AutoService(Instrumenter.class)
+public class ServletTestInstrumentation implements Instrumenter {
+
+  @Override
+  public AgentBuilder instrument(final AgentBuilder agentBuilder) {
+    return agentBuilder
+        // Jetty 7.0
+        .type(named("org.eclipse.jetty.server.HttpConnection"))
+        .transform(
+            new AgentBuilder.Transformer.ForAdvice()
+                .advice(
+                    named("handleRequest"), HttpServerTestAdvice.ServerEntryAdvice.class.getName()))
+        // Jetty 7.latest
+        .type(named("org.eclipse.jetty.server.AbstractHttpConnection"))
+        .transform(
+            new AgentBuilder.Transformer.ForAdvice()
+                .advice(
+                    named("headerComplete"),
+                    HttpServerTestAdvice.ServerEntryAdvice.class.getName()));
+  }
+}

dd-java-agent/instrumentation/servlet-2/src/test/groovy/TestServlet2.groovy

@@ -1,21 +1,37 @@
+import datadog.trace.agent.test.base.HttpServerTest
 import groovy.servlet.AbstractHttpServlet
-
 import javax.servlet.http.HttpServletRequest
 import javax.servlet.http.HttpServletResponse
 
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.ERROR
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.EXCEPTION
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.REDIRECT
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS
+
 class TestServlet2 {
 
   static class Sync extends AbstractHttpServlet {
     @Override
-    void doGet(HttpServletRequest req, HttpServletResponse resp) {
+    protected void service(HttpServletRequest req, HttpServletResponse resp) {
-      if (req.getParameter("error") != null) {
+      req.getRequestDispatcher()
-        throw new RuntimeException("some sync error")
+      HttpServerTest.ServerEndpoint endpoint = HttpServerTest.ServerEndpoint.forPath(req.servletPath)
+      HttpServerTest.controller(endpoint) {
+        resp.contentType = "text/plain"
+        switch (endpoint) {
+          case SUCCESS:
+            resp.status = endpoint.status
+            resp.writer.print(endpoint.body)
+            break
+          case REDIRECT:
+            resp.sendRedirect(endpoint.body)
+            break
+          case ERROR:
+            resp.sendError(endpoint.status, endpoint.body)
+            break
+          case EXCEPTION:
+            throw new Exception(endpoint.body)
         }
-      if (req.getParameter("non-throwing-error") != null) {
-        resp.sendError(500, "some sync error")
-        return
       }
-      resp.writer.print("Hello Sync")
     }
   }
 }

dd-java-agent/instrumentation/servlet-3/src/test/groovy/JettyServlet3Test.groovy

@@ -43,7 +43,7 @@ abstract class JettyServlet3Test extends AbstractServlet3Test<ServletContextHand
     servletContext.errorHandler = new ErrorHandler() {
       protected void handleErrorPage(HttpServletRequest request, Writer writer, int code, String message) throws IOException {
         Throwable th = (Throwable) request.getAttribute("javax.servlet.error.exception")
-        writer.write(th.message)
+        writer.write(th ? th.message : message)
       }
     }
 //    setupAuthentication(jettyServer, servletContext)

dd-java-agent/instrumentation/servlet-3/src/test/groovy/TestServlet3.groovy

@@ -22,13 +22,15 @@ class TestServlet3 {
         resp.contentType = "text/plain"
         switch (endpoint) {
           case SUCCESS:
-          case ERROR:
             resp.status = endpoint.status
             resp.writer.print(endpoint.body)
             break
           case REDIRECT:
             resp.sendRedirect(endpoint.body)
             break
+          case ERROR:
+            resp.sendError(endpoint.status, endpoint.body)
+            break
           case EXCEPTION:
             throw new Exception(endpoint.body)
         }

dd-java-agent/instrumentation/servlet-3/src/test/groovy/TomcatServlet3Test.groovy

@@ -125,7 +125,7 @@ class ErrorHandlerValve extends ErrorReportValve {
       return
     }
     try {
-      response.writer.print(t.cause.message)
+      response.writer.print(t ? t.cause.message : response.message)
     } catch (IOException e) {
       e.printStackTrace()
     }