diff --git a/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java b/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java index 87735d6a17..2264d9945b 100644 --- a/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java +++ b/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 = diff --git a/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Instrumentation.java b/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Instrumentation.java index 5f309738d4..8490d7a237 100644 --- a/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Instrumentation.java +++ b/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 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 diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletRequestExtractAdapter.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletRequestExtractAdapter.java index e8b1fff393..3b904d12ff 100644 --- a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletRequestExtractAdapter.java +++ b/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 keys(final HttpServletRequest carrier) { - final ArrayList keys = Collections.list(carrier.getHeaderNames()); + final List keys = Collections.list(carrier.getHeaderNames()); keys.addAll(Collections.list(carrier.getAttributeNames())); return keys; } diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java index dbc1a4419d..407520288a 100644 --- a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java +++ b/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 = diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Instrumentation.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Instrumentation.java index 0fb9455ac6..42c982b169 100644 --- a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Instrumentation.java +++ b/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 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 diff --git a/dd-java-agent/instrumentation/servlet/servlet.gradle b/dd-java-agent/instrumentation/servlet/servlet.gradle index 1f4583dd52..50f9d1e4b1 100644 --- a/dd-java-agent/instrumentation/servlet/servlet.gradle +++ b/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' + } } diff --git a/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/ServletContextInstrumentation.java b/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/ServletContextInstrumentation.java index 3426135c6b..9c3df78076 100644 --- a/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/ServletContextInstrumentation.java +++ b/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"); - } - - @Override - public boolean defaultEnabled() { - return false; + super("servlet", "servlet-dispatcher"); } @Override diff --git a/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletInstrumentation.java b/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletInstrumentation.java index ad4c7a3190..dbcfcf488e 100644 --- a/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletInstrumentation.java +++ b/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 contextStore() { - return singletonMap( - "javax.servlet.http.HttpServletResponse", "javax.servlet.http.HttpServletRequest"); - } - public static class HttpServletAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope start( - @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); + public static AgentScope start(@Advice.Origin final Method method) { if (activeSpan() == null) { // Don't want to generate a new top-level span diff --git a/dd-java-agent/instrumentation/servlet/src/test/groovy/HttpServletResponseTest.groovy b/dd-java-agent/instrumentation/servlet/src/test/groovy/HttpServletResponseTest.groovy index 1bc6f2612e..f82e5904ac 100644 --- a/dd-java-agent/instrumentation/servlet/src/test/groovy/HttpServletResponseTest.groovy +++ b/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") {