Move request-response linking to main servlet advice
servlet-service is off by default, so we can't rely on it being called there.
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
@@ -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 =

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
@@ -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

MODIFY
dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletRequestExtractAdapter.java
dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServletRequestExtractAdapter.java
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;
}

MODIFY
dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java
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 =

MODIFY
dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Instrumentation.java
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

MODIFY
dd-java-agent/instrumentation/servlet/servlet.gradle
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'
+ }
}

MODIFY
dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/ServletContextInstrumentation.java
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

MODIFY
dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletInstrumentation.java
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(
- @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

MODIFY
dd-java-agent/instrumentation/servlet/src/test/groovy/HttpServletResponseTest.groovy
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") {

Move request-response linking to main servlet advice

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