Move request-response linking to main servlet advice

servlet-service is off by default, so we can't rely on it being called there.
This commit is contained in:
Tyler Benson 2020-01-13 10:48:35 -08:00
parent e440eba9a1
commit 49425e9963
9 changed files with 50 additions and 29 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@ -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") {