diff --git a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/build.gradle b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/build.gradle index 0bcee37c285..5c56b57bd45 100644 --- a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/build.gradle +++ b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/build.gradle @@ -14,7 +14,7 @@ addTestSuiteForDir('latestDepTest', 'test') dependencies { compileOnly group: 'org.apache.httpcomponents.client5', name: 'httpclient5', version: '5.0' - testImplementation group: 'org.apache.httpcomponents.client5', name: 'httpclient5', version: '5.0' + testImplementation group: 'org.apache.httpcomponents.client5', name: 'httpclient5', version: '5.0.2' // we need this fix https://github.com/apache/httpcomponents-client/commit/a22d889807eb74a0d4b8dd8d6360802a391a6eb3 latestDepTestImplementation group: 'org.apache.httpcomponents.client5', name: 'httpclient5', version: '+' } diff --git a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpClientInstrumentation.java b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpClientInstrumentation.java index 051b5305771..9df8d93c7d0 100644 --- a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpClientInstrumentation.java +++ b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpClientInstrumentation.java @@ -7,10 +7,12 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; -import datadog.appsec.api.blocking.BlockingException; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import java.util.Collections; +import java.util.Map; import net.bytebuddy.asm.Advice; import net.bytebuddy.description.type.TypeDescription; import net.bytebuddy.implementation.bytecode.assign.Assigner; @@ -68,6 +70,15 @@ public String[] helperClassNames() { }; } + @Override + public Map contextStore() { + // used to mark when a request has been instrumented. + // We don't count depth like we usually do, because sub-requests can be spawned and need to be + // traced + return Collections.singletonMap( + "org.apache.hc.core5.http.ClassicHttpRequest", "java.lang.Boolean"); + } + @Override public void methodAdvice(MethodTransformer transformer) { transformer.applyAdvice( @@ -116,17 +127,13 @@ public void methodAdvice(MethodTransformer transformer) { public static class RequestAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope methodEnter(@Advice.Argument(0) final ClassicHttpRequest request) { - try { - return HelperMethods.doMethodEnter(request); - } catch (BlockingException e) { - HelperMethods.onBlockingRequest(); - // re-throw blocking exceptions - throw e; - } + return HelperMethods.doMethodEnter( + InstrumentationContext.get(ClassicHttpRequest.class, Boolean.class), request); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void methodExit( + @Advice.Argument(0) final ClassicHttpRequest request, @Advice.Enter final AgentScope scope, @Advice.Return final Object result, @Advice.Thrown final Throwable throwable) { @@ -139,17 +146,13 @@ public static class HostRequestAdvice { public static AgentScope methodEnter( @Advice.Argument(0) final HttpHost host, @Advice.Argument(1) final ClassicHttpRequest request) { - try { - return HelperMethods.doMethodEnter(host, request); - } catch (BlockingException e) { - HelperMethods.onBlockingRequest(); - // re-throw blocking exceptions - throw e; - } + return HelperMethods.doMethodEnter( + InstrumentationContext.get(ClassicHttpRequest.class, Boolean.class), host, request); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void methodExit( + @Advice.Argument(1) final ClassicHttpRequest request, @Advice.Enter final AgentScope scope, @Advice.Return final Object result, @Advice.Thrown final Throwable throwable) { @@ -170,24 +173,21 @@ public static AgentScope methodEnter( typing = Assigner.Typing.DYNAMIC, readOnly = false) Object handler) { - try { - final AgentScope scope = HelperMethods.doMethodEnter(host, request); - // Wrap the handler so we capture the status code - if (null != scope && handler instanceof HttpClientResponseHandler) { - handler = - new WrappingStatusSettingResponseHandler( - scope.span(), (HttpClientResponseHandler) handler); - } - return scope; - } catch (BlockingException e) { - HelperMethods.onBlockingRequest(); - // re-throw blocking exceptions - throw e; + final AgentScope scope = + HelperMethods.doMethodEnter( + InstrumentationContext.get(ClassicHttpRequest.class, Boolean.class), host, request); + // Wrap the handler so we capture the status code + if (null != scope && handler instanceof HttpClientResponseHandler) { + handler = + new WrappingStatusSettingResponseHandler( + scope.span(), (HttpClientResponseHandler) handler); } + return scope; } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void methodExit( + @Advice.Argument(1) final ClassicHttpRequest request, @Advice.Enter final AgentScope scope, @Advice.Return final Object result, @Advice.Thrown final Throwable throwable) { diff --git a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java index e6187609433..c962ada1455 100644 --- a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java +++ b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java @@ -7,33 +7,52 @@ import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.HTTP_REQUEST; import static datadog.trace.instrumentation.apachehttpclient5.HttpHeadersInjectAdapter.SETTER; -import datadog.trace.bootstrap.CallDepthThreadLocalMap; +import datadog.trace.bootstrap.ContextStore; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; -import org.apache.hc.client5.http.classic.HttpClient; +import org.apache.hc.core5.http.ClassicHttpRequest; import org.apache.hc.core5.http.HttpHost; import org.apache.hc.core5.http.HttpRequest; import org.apache.hc.core5.http.HttpResponse; public class HelperMethods { - public static AgentScope doMethodEnter(final HttpRequest request) { - final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(HttpClient.class); - if (callDepth > 0) { + + public static AgentScope doMethodEnter( + final ContextStore instrumentationMarker, + final ClassicHttpRequest request) { + if (testAndSet(instrumentationMarker, request)) { return null; } - return activateHttpSpan(request); } - public static AgentScope doMethodEnter(HttpHost host, HttpRequest request) { - final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(HttpClient.class); - if (callDepth > 0) { + public static AgentScope doMethodEnter( + final ContextStore instrumentationMarker, + HttpHost host, + final ClassicHttpRequest request) { + if (testAndSet(instrumentationMarker, request)) { return null; } - return activateHttpSpan(new HostAndRequestAsHttpUriRequest(host, request)); } + // checks current value in context store, + // and ensures it's set to true when this method exists + private static boolean testAndSet( + final ContextStore instrumentationMarker, + final ClassicHttpRequest request) { + if (request == null) { + // we probably don't want to instrument a call with a null request ? + return true; + } + Boolean instrumented = instrumentationMarker.get(request); + if (instrumented == Boolean.TRUE) { + return true; + } + instrumentationMarker.put(request, Boolean.TRUE); + return false; + } + private static AgentScope activateHttpSpan(final HttpRequest request) { final AgentSpan span = startSpan(HTTP_REQUEST); final AgentScope scope = activateSpan(span); @@ -55,8 +74,8 @@ public static void doMethodExit( if (scope == null) { return; } - final AgentSpan span = scope.span(); try { + final AgentSpan span = scope.span(); if (result instanceof HttpResponse) { DECORATE.onResponse(span, (HttpResponse) result); } // else they probably provided a ResponseHandler. @@ -64,13 +83,9 @@ public static void doMethodExit( DECORATE.onError(span, throwable); DECORATE.beforeFinish(span); } finally { + AgentSpan span = scope.span(); scope.close(); span.finish(); - CallDepthThreadLocalMap.reset(HttpClient.class); } } - - public static void onBlockingRequest() { - CallDepthThreadLocalMap.reset(HttpClient.class); - } } diff --git a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/test/groovy/ApacheHttpClientTest.groovy b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/test/groovy/ApacheHttpClientTest.groovy index c1a27127857..13a89e35f59 100644 --- a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/test/groovy/ApacheHttpClientTest.groovy +++ b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/test/groovy/ApacheHttpClientTest.groovy @@ -215,3 +215,72 @@ class ApacheClientResponseHandlerAllV0Test extends ApacheClientResponseHandlerAl class ApacheClientResponseHandlerAllV1ForkedTest extends ApacheClientResponseHandlerAll implements TestingGenericHttpNamingConventions.ClientV1 { } +/** + * Tests that HTTP calls made from within an ExecChainHandler (exec interceptor) are instrumented. + * Reproduces the scenario from https://github.com/DataDog/dd-trace-java/issues/10383: an + * interceptor fetches a token via a separate HttpClient, adds it as a header, then proceeds. + */ +@Timeout(5) +class ApacheClientNestedExecuteTest extends ApacheHttpClientTest { + + @Override + ClassicHttpRequest createRequest(String method, URI uri) { + return new BasicClassicHttpRequest(method, uri) + } + + @Override + CloseableHttpResponse executeRequest(ClassicHttpRequest request, URI uri) { + return client.execute(request) + } + + def "HTTP call from ExecChainHandler (e.g. token fetch) is instrumented"() { + when: + def tokenUri = server.address.resolve("/success") + def mainUri = server.address.resolve("/success") + def requestConfig = RequestConfig.custom() + .setConnectTimeout(CONNECT_TIMEOUT_MS, TimeUnit.MILLISECONDS) + .build() + def tokenClient = HttpClients.custom() + .setConnectionManager(new BasicHttpClientConnectionManager()) + .setDefaultRequestConfig(requestConfig) + .build() + def tokenUriFinal = tokenUri + def clientWithInterceptor = HttpClients.custom() + .setConnectionManager(new BasicHttpClientConnectionManager()) + .setDefaultRequestConfig(requestConfig) + .addExecInterceptorLast("token", { request, scope, chain -> + String token = tokenClient.execute( + new BasicClassicHttpRequest("GET", tokenUriFinal), { resp -> + String.valueOf(resp.code) + } + ) + request.addHeader(new BasicHeader("x-token", token)) + return chain.proceed(request, scope) + }) + .build() + def response = clientWithInterceptor.execute( + new BasicClassicHttpRequest("GET", mainUri), { r -> + r + } + ) + + then: + response != null + assertTraces(3) { + sortSpansByStart() + trace(size(2)) { + sortSpansByStart() + // Token request runs first (inside interceptor), then main request + clientSpan(it, null, "GET", false, false, tokenUri) + clientSpan(it, span(0), "GET", false, false, mainUri) + } + server.distributedRequestTrace(it, trace(0)[0]) + server.distributedRequestTrace(it, trace(0)[1]) + } + + cleanup: + tokenClient?.close() + clientWithInterceptor?.close() + } +} +