Rework servlet 2 to not wrap HttpServletResponse objects
MODIFY
dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java
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);
@@ -44,8 +44,7 @@ public class Servlet2Advice {
// For use by HttpServletResponseInstrumentation:
InstrumentationContext.get(HttpServletResponse.class, HttpServletRequest.class)
.put((HttpServletResponse) response, httpServletRequest);
-
- response = new StatusSavingHttpServletResponseWrapper((HttpServletResponse) response);
+ httpServletRequest.setAttribute("dd.http-status", HttpServletResponse.SC_OK);
}
final AgentSpan.Context extractedContext = propagate().extract(httpServletRequest, GETTER);
@@ -73,7 +72,6 @@ public class Servlet2Advice {
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void stopSpan(
@Advice.Argument(0) final ServletRequest request,
- @Advice.Argument(1) final ServletResponse response,
@Advice.Enter final AgentScope scope,
@Advice.Thrown final Throwable throwable) {
// Set user.principal regardless of who created this span.
@@ -89,11 +87,11 @@ public class Servlet2Advice {
return;
}
final AgentSpan span = scope.span();
- DECORATE.onResponse(span, response);
+ if (request instanceof HttpServletRequest) {
+ DECORATE.onResponse(span, (HttpServletRequest) request);
+ }
if (throwable != null) {
- if (response instanceof StatusSavingHttpServletResponseWrapper
- && ((StatusSavingHttpServletResponseWrapper) response).status
- == HttpServletResponse.SC_OK) {
+ if ((Integer) request.getAttribute("dd.http-status") == HttpServletResponse.SC_OK) {
// exception was thrown but status code wasn't set
span.setTag(Tags.HTTP_STATUS, 500);
}

MODIFY
dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Decorator.java
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.api.AgentSpan;
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, HttpServletRequest> {
public static final Servlet2Decorator DECORATE = new Servlet2Decorator();
@Override
@@ -50,11 +49,10 @@ public class Servlet2Decorator
}
@Override
- protected Integer status(final ServletResponse httpServletResponse) {
- if (httpServletResponse instanceof StatusSavingHttpServletResponseWrapper) {
- return ((StatusSavingHttpServletResponseWrapper) httpServletResponse).status;
+ protected Integer status(final HttpServletRequest httpServletRequest) {
+ if (httpServletRequest != null) {
+ return (Integer) httpServletRequest.getAttribute("dd.http-status");
} else {
- // HttpServletResponse doesn't have accessor for status code.
return null;
}
}

MODIFY
dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Instrumentation.java
dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Instrumentation.java
@@ -37,9 +37,7 @@ public final class Servlet2Instrumentation extends Instrumenter.Default {
@Override
public String[] helperClassNames() {
return new String[] {
- packageName + ".Servlet2Decorator",
- packageName + ".HttpServletRequestExtractAdapter",
- packageName + ".StatusSavingHttpServletResponseWrapper",
+ packageName + ".Servlet2Decorator", packageName + ".HttpServletRequestExtractAdapter",
};
}

ADD
dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2ResponseRedirectAdvice.java
+package datadog.trace.instrumentation.servlet2;
+
+import datadog.trace.bootstrap.InstrumentationContext;
+import javax.servlet.http.HttpServletRequest;
+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(HttpServletResponse.class, HttpServletRequest.class)
+ .get(response)
+ .setAttribute("status-code", 302);
+ }
+}

ADD
dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2ResponseStatusAdvice.java
+package datadog.trace.instrumentation.servlet2;
+
+import datadog.trace.bootstrap.InstrumentationContext;
+import javax.servlet.http.HttpServletRequest;
+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(HttpServletResponse.class, HttpServletRequest.class)
+ .get(response)
+ .setAttribute("dd.http-status", status);
+ }
+}

ADD
dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2ResponseStatusInstrumentation.java
+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.http.HttpServletResponse", "javax.servlet.http.HttpServletRequest");
+ }
+
+ /**
+ * 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;
+ }
+}

DELETE
dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/StatusSavingHttpServletResponseWrapper.java
-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 sendRedirect(final String location) throws IOException {
- status = 302;
- super.sendRedirect(location);
- }
-
- @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);
- }
-}

Rework servlet 2 to not wrap HttpServletResponse objects