Move request-response linking to main servlet advice

servlet-service is off by default, so we can't rely on it being called there.

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

@@ -9,6 +9,7 @@ import static datadog.trace.instrumentation.servlet2.HttpServletRequestExtractAd
 import static datadog.trace.instrumentation.servlet2.Servlet2Decorator.DECORATE;
 
 import datadog.trace.api.DDTags;
+import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.instrumentation.api.AgentScope;
 import datadog.trace.instrumentation.api.AgentSpan;
 import datadog.trace.instrumentation.api.Tags;
@@ -36,12 +37,16 @@ public class Servlet2Advice {
       return null;
     }
 
+    final HttpServletRequest httpServletRequest = (HttpServletRequest) request;
+
     if (response instanceof HttpServletResponse) {
+      // For use by HttpServletResponseInstrumentation:
+      InstrumentationContext.get(HttpServletResponse.class, HttpServletRequest.class)
+          .put((HttpServletResponse) response, httpServletRequest);
+
       response = new StatusSavingHttpServletResponseWrapper((HttpServletResponse) response);
     }
 
-    final HttpServletRequest httpServletRequest = (HttpServletRequest) request;
-
     final AgentSpan.Context extractedContext = propagate().extract(httpServletRequest, GETTER);
 
     final AgentSpan span =

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

@@ -49,6 +49,12 @@ public final class Servlet2Instrumentation extends Instrumenter.Default {
                 named("javax.servlet.FilterChain").or(named("javax.servlet.http.HttpServlet"))));
   }
 
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap(
+        "javax.servlet.http.HttpServletResponse", "javax.servlet.http.HttpServletRequest");
+  }
+
   /**
    * Here we are instrumenting the public method for HttpServlet. This should ensure that this
    * advice is always called before HttpServletInstrumentation which is instrumenting the protected

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

@@ -1,7 +1,6 @@
 package datadog.trace.instrumentation.servlet3;
 
 import datadog.trace.instrumentation.api.AgentPropagation;
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import javax.servlet.http.HttpServletRequest;
@@ -14,7 +13,7 @@ public class HttpServletRequestExtractAdapter
 
   @Override
   public List<String> keys(final HttpServletRequest carrier) {
-    final ArrayList<String> keys = Collections.list(carrier.getHeaderNames());
+    final List<String> keys = Collections.list(carrier.getHeaderNames());
     keys.addAll(Collections.list(carrier.getAttributeNames()));
     return keys;
   }

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

@@ -9,6 +9,7 @@ import static datadog.trace.instrumentation.servlet3.HttpServletRequestExtractAd
 import static datadog.trace.instrumentation.servlet3.Servlet3Decorator.DECORATE;
 
 import datadog.trace.api.DDTags;
+import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.instrumentation.api.AgentScope;
 import datadog.trace.instrumentation.api.AgentSpan;
 import datadog.trace.instrumentation.api.Tags;
@@ -24,7 +25,10 @@ public class Servlet3Advice {
 
   @Advice.OnMethodEnter(suppress = Throwable.class)
   public static AgentScope onEnter(
-      @Advice.This final Object servlet, @Advice.Argument(0) final ServletRequest request) {
+      @Advice.This final Object servlet,
+      @Advice.Argument(0) final ServletRequest request,
+      @Advice.Argument(1) final ServletResponse response) {
+
     final boolean hasActiveTrace = activeSpan() != null;
     final boolean hasServletTrace = request.getAttribute(DD_SPAN_ATTRIBUTE) instanceof AgentSpan;
     final boolean invalidRequest = !(request instanceof HttpServletRequest);
@@ -35,6 +39,10 @@ public class Servlet3Advice {
 
     final HttpServletRequest httpServletRequest = (HttpServletRequest) request;
 
+    // For use by HttpServletResponseInstrumentation:
+    InstrumentationContext.get(HttpServletResponse.class, HttpServletRequest.class)
+        .put((HttpServletResponse) response, httpServletRequest);
+
     final AgentSpan.Context extractedContext = propagate().extract(httpServletRequest, GETTER);
 
     final AgentSpan span =

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

@@ -41,6 +41,12 @@ public final class Servlet3Instrumentation extends Instrumenter.Default {
                 named("javax.servlet.FilterChain").or(named("javax.servlet.http.HttpServlet"))));
   }
 
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap(
+        "javax.servlet.http.HttpServletResponse", "javax.servlet.http.HttpServletRequest");
+  }
+
   /**
    * Here we are instrumenting the public method for HttpServlet. This should ensure that this
    * advice is always called before HttpServletInstrumentation which is instrumenting the protected

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

@@ -18,4 +18,12 @@ dependencies {
   compileOnly group: 'javax.servlet', name: 'servlet-api', version: '2.3'
 
   testCompile group: 'javax.servlet', name: 'servlet-api', version: '2.3'
+
+  // servlet request instrumentation required for linking request to response.
+  testCompile project(':dd-java-agent:instrumentation:servlet:request-2')
+
+  // Don't want to conflict with jetty from the test server.
+  testCompile(project(':dd-java-agent:testing')) {
+    exclude group: 'org.eclipse.jetty', module: 'jetty-server'
+  }
 }

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

@@ -22,12 +22,7 @@ import net.bytebuddy.matcher.ElementMatcher;
 @AutoService(Instrumenter.class)
 public final class ServletContextInstrumentation extends Instrumenter.Default {
   public ServletContextInstrumentation() {
-    super("servlet-beta", "servlet-dispatcher");
+    super("servlet", "servlet-dispatcher");
-  }
-
-  @Override
-  public boolean defaultEnabled() {
-    return false;
   }
 
   @Override

dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletInstrumentation.java

@@ -17,13 +17,10 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.api.DDTags;
-import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.instrumentation.api.AgentScope;
 import datadog.trace.instrumentation.api.AgentSpan;
 import java.lang.reflect.Method;
 import java.util.Map;
-import javax.servlet.http.HttpServletRequest;
-import javax.servlet.http.HttpServletResponse;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
@@ -67,22 +64,10 @@ public final class HttpServletInstrumentation extends Instrumenter.Default {
         HttpServletAdvice.class.getName());
   }
 
-  @Override
-  public Map<String, String> contextStore() {
-    return singletonMap(
-        "javax.servlet.http.HttpServletResponse", "javax.servlet.http.HttpServletRequest");
-  }
-
   public static class HttpServletAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static AgentScope start(
+    public static AgentScope start(@Advice.Origin final Method method) {
-        @Advice.Origin final Method method,
-        @Advice.Argument(0) final HttpServletRequest request,
-        @Advice.Argument(1) final HttpServletResponse response) {
-      // For use by HttpServletResponseInstrumentation:
-      InstrumentationContext.get(HttpServletResponse.class, HttpServletRequest.class)
-          .put(response, request);
 
       if (activeSpan() == null) {
         // Don't want to generate a new top-level span

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

@@ -3,12 +3,15 @@ import groovy.servlet.AbstractHttpServlet
 import spock.lang.Subject
 
 import javax.servlet.ServletOutputStream
+import javax.servlet.ServletRequest
+import javax.servlet.ServletResponse
 import javax.servlet.http.Cookie
 import javax.servlet.http.HttpServletRequest
 import javax.servlet.http.HttpServletResponse
 
 import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
 import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
+import static java.util.Collections.emptyEnumeration
 
 class HttpServletResponseTest extends AgentTestRunner {
 
@@ -17,12 +20,16 @@ class HttpServletResponseTest extends AgentTestRunner {
   def request = Mock(HttpServletRequest) {
     getMethod() >> "GET"
     getProtocol() >> "TEST"
+    getHeaderNames() >> emptyEnumeration()
+    getAttributeNames() >> emptyEnumeration()
   }
 
   def setup() {
     def servlet = new AbstractHttpServlet() {}
     // We need to call service so HttpServletAdvice can link the request to the response.
-    servlet.service(request, response)
+    servlet.service((ServletRequest) request, (ServletResponse) response)
+    assert response.__datadogContext$javax$servlet$http$HttpServletResponse != null
+    TEST_WRITER.clear()
   }
 
   def "test send no-parent"() {
@@ -89,7 +96,9 @@ class HttpServletResponseTest extends AgentTestRunner {
     }
     def servlet = new AbstractHttpServlet() {}
     // We need to call service so HttpServletAdvice can link the request to the response.
-    servlet.service(request, response)
+    servlet.service((ServletRequest) request, (ServletResponse) response)
+    assert response.__datadogContext$javax$servlet$http$HttpServletResponse != null
+    TEST_WRITER.clear()
 
     when:
     runUnderTrace("parent") {