Rework servlet 2 to not wrap HttpServletResponse objects

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

@@ -2,10 +2,10 @@ muzzle {
   pass {
     group = "javax.servlet"
     module = "servlet-api"
-    versions = "[2.3, 3.0)"
+    versions = "[2.2, 3.0)"
     assertInverse = true
   }
-  
+
   fail {
     group = "javax.servlet"
     module = 'javax.servlet-api'
@@ -24,7 +24,7 @@ testSets {
 }
 
 dependencies {
-  compileOnly group: 'javax.servlet', name: 'servlet-api', version: '2.3'
+  compileOnly group: 'javax.servlet', name: 'servlet-api', version: '2.2'
 
   testCompile(project(':dd-java-agent:testing')) {
     exclude group: 'org.eclipse.jetty', module: 'jetty-server'

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

@@ -28,8 +28,8 @@ public class Servlet2Advice {
   public static AgentScope onEnter(
       @Advice.This final Object servlet,
       @Advice.Argument(0) final ServletRequest request,
-      @Advice.Argument(value = 1, readOnly = false, typing = Assigner.Typing.DYNAMIC)
-          ServletResponse response) {
+      @Advice.Argument(value = 1, typing = Assigner.Typing.DYNAMIC)
+          final ServletResponse response) {
 
     final boolean hasServletTrace = request.getAttribute(DD_SPAN_ATTRIBUTE) instanceof AgentSpan;
     final boolean invalidRequest = !(request instanceof HttpServletRequest);
@@ -45,7 +45,8 @@ public static AgentScope onEnter(
       InstrumentationContext.get(HttpServletResponse.class, HttpServletRequest.class)
           .put((HttpServletResponse) response, httpServletRequest);
 
-      response = new StatusSavingHttpServletResponseWrapper((HttpServletResponse) response);
+      // Default value for checking for uncaught error later
+      InstrumentationContext.get(ServletResponse.class, Integer.class).put(response, 200);
     }
 
     final AgentSpan.Context extractedContext = propagate().extract(httpServletRequest, GETTER);
@@ -89,10 +90,17 @@ public static void stopSpan(
       return;
     }
     final AgentSpan span = scope.span();
-    DECORATE.onResponse(span, response);
+
+    if (response instanceof HttpServletResponse) {
+      DECORATE.onResponse(
+          span, InstrumentationContext.get(ServletResponse.class, Integer.class).get(response));
+    } else {
+      DECORATE.onResponse(span, null);
+    }
+
     if (throwable != null) {
-      if (response instanceof StatusSavingHttpServletResponseWrapper
-          && ((StatusSavingHttpServletResponseWrapper) response).status
+      if (response instanceof HttpServletResponse
+          && InstrumentationContext.get(ServletResponse.class, Integer.class).get(response)
               == HttpServletResponse.SC_OK) {
         // exception was thrown but status code wasn't set
         span.setTag(Tags.HTTP_STATUS, 500);

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

@@ -4,11 +4,10 @@
 import datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator;
 import java.net.URI;
 import java.net.URISyntaxException;
-import javax.servlet.ServletResponse;
 import javax.servlet.http.HttpServletRequest;
 
 public class Servlet2Decorator
-    extends HttpServerDecorator<HttpServletRequest, HttpServletRequest, ServletResponse> {
+    extends HttpServerDecorator<HttpServletRequest, HttpServletRequest, Integer> {
   public static final Servlet2Decorator DECORATE = new Servlet2Decorator();
 
   @Override
@@ -50,13 +49,8 @@ protected Integer peerPort(final HttpServletRequest httpServletRequest) {
   }
 
   @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;
-    }
+  protected Integer status(final Integer status) {
+    return status;
   }
 
   @Override

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

@@ -10,6 +10,7 @@
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
+import java.util.HashMap;
 import java.util.Map;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
@@ -37,16 +38,17 @@ public ElementMatcher<TypeDescription> typeMatcher() {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      packageName + ".Servlet2Decorator",
-      packageName + ".HttpServletRequestExtractAdapter",
-      packageName + ".StatusSavingHttpServletResponseWrapper",
+      packageName + ".Servlet2Decorator", packageName + ".HttpServletRequestExtractAdapter",
     };
   }
 
   @Override
   public Map<String, String> contextStore() {
-    return singletonMap(
+    final Map<String, String> contextStores = new HashMap<>();
+    contextStores.put(
         "javax.servlet.http.HttpServletResponse", "javax.servlet.http.HttpServletRequest");
+    contextStores.put("javax.servlet.ServletResponse", Integer.class.getName());
+    return contextStores;
   }
 
   /**

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

@@ -0,0 +1,13 @@
+package datadog.trace.instrumentation.servlet2;
+
+import datadog.trace.bootstrap.InstrumentationContext;
+import javax.servlet.ServletResponse;
+import javax.servlet.http.HttpServletResponse;
+import net.bytebuddy.asm.Advice;
+
+public class Servlet2ResponseRedirectAdvice {
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static void onEnter(@Advice.This final HttpServletResponse response) {
+    InstrumentationContext.get(ServletResponse.class, Integer.class).put(response, 302);
+  }
+}

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

@@ -0,0 +1,14 @@
+package datadog.trace.instrumentation.servlet2;
+
+import datadog.trace.bootstrap.InstrumentationContext;
+import javax.servlet.ServletResponse;
+import javax.servlet.http.HttpServletResponse;
+import net.bytebuddy.asm.Advice;
+
+public class Servlet2ResponseStatusAdvice {
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static void onEnter(
+      @Advice.This final HttpServletResponse response, @Advice.Argument(0) final Integer status) {
+    InstrumentationContext.get(ServletResponse.class, Integer.class).put(response, status);
+  }
+}

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

@@ -0,0 +1,51 @@
+package datadog.trace.instrumentation.servlet2;
+
+import static datadog.trace.agent.tooling.ClassLoaderMatcher.hasClassesNamed;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.DDElementMatchers.safeHasSuperType;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import java.util.HashMap;
+import java.util.Map;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public final class Servlet2ResponseStatusInstrumentation extends Instrumenter.Default {
+  public Servlet2ResponseStatusInstrumentation() {
+    super("servlet", "servlet-2");
+  }
+
+  // this is required to make sure servlet 2 instrumentation won't apply to servlet 3
+  @Override
+  public ElementMatcher<ClassLoader> classLoaderMatcher() {
+    return not(hasClassesNamed("javax.servlet.AsyncEvent", "javax.servlet.AsyncListener"));
+  }
+
+  @Override
+  public ElementMatcher<? super TypeDescription> typeMatcher() {
+    return safeHasSuperType(named("javax.servlet.http.HttpServletResponse"));
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap("javax.servlet.ServletResponse", Integer.class.getName());
+  }
+
+  /**
+   * Unlike Servlet2Instrumentation it doesn't matter if the HttpServletResponseInstrumentation
+   * applies first
+   */
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    final Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
+    transformers.put(
+        named("sendError").or(named("setStatus")), packageName + ".Servlet2ResponseStatusAdvice");
+    transformers.put(named("sendRedirect"), packageName + ".Servlet2ResponseRedirectAdvice");
+    return transformers;
+  }
+}