Use the parent interface
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
@@ -44,7 +44,9 @@ public class Servlet2Advice {
// For use by HttpServletResponseInstrumentation:
InstrumentationContext.get(HttpServletResponse.class, HttpServletRequest.class)
.put((HttpServletResponse) response, httpServletRequest);
- httpServletRequest.setAttribute("dd.http-status", HttpServletResponse.SC_OK);
+
+ // 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);
@@ -72,6 +74,7 @@ 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.
@@ -87,11 +90,18 @@ public class Servlet2Advice {
return;
}
final AgentSpan span = scope.span();
- if (request instanceof HttpServletRequest) {
- DECORATE.onResponse(span, (HttpServletRequest) request);
+
+ if (response instanceof HttpServletResponse) {
+ DECORATE.onResponse(
+ span, InstrumentationContext.get(ServletResponse.class, Integer.class).get(response));
+ } else {
+ DECORATE.onResponse(span, null);
}
+
if (throwable != null) {
- if ((Integer) request.getAttribute("dd.http-status") == HttpServletResponse.SC_OK) {
+ 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);
}

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
@@ -7,7 +7,7 @@ import java.net.URISyntaxException;
import javax.servlet.http.HttpServletRequest;
public class Servlet2Decorator
- extends HttpServerDecorator<HttpServletRequest, HttpServletRequest, HttpServletRequest> {
+ extends HttpServerDecorator<HttpServletRequest, HttpServletRequest, Integer> {
public static final Servlet2Decorator DECORATE = new Servlet2Decorator();
@Override
@@ -49,12 +49,8 @@ public class Servlet2Decorator
}
@Override
- protected Integer status(final HttpServletRequest httpServletRequest) {
- if (httpServletRequest != null) {
- return (Integer) httpServletRequest.getAttribute("dd.http-status");
- } else {
- return null;
- }
+ protected Integer status(final Integer status) {
+ return status;
}
@Override

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
@@ -10,6 +10,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
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;
@@ -43,8 +44,11 @@ public final class Servlet2Instrumentation extends Instrumenter.Default {
@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;
}
/**

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

MODIFY
dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2ResponseStatusAdvice.java
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.ServletResponse;
import javax.servlet.http.HttpServletResponse;
import net.bytebuddy.asm.Advice;
@@ -9,8 +9,6 @@ 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);
+ InstrumentationContext.get(ServletResponse.class, Integer.class).put(response, status);
}
}

MODIFY
dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2ResponseStatusInstrumentation.java
dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2ResponseStatusInstrumentation.java
@@ -33,8 +33,7 @@ public final class Servlet2ResponseStatusInstrumentation extends Instrumenter.De
@Override
public Map<String, String> contextStore() {
- return singletonMap(
- "javax.servlet.http.HttpServletResponse", "javax.servlet.http.HttpServletRequest");
+ return singletonMap("javax.servlet.ServletResponse", Integer.class.getName());
}
/**

Use the parent interface