Skip to content

Commit 959c2cc

Browse files
vandonrvandonr
committed
instrument inner requests in apache http client
1 parent 61376da commit 959c2cc

3 files changed

Lines changed: 140 additions & 29 deletions

File tree

‎dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpClientInstrumentation.java‎

Lines changed: 43 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,17 @@
1010
import datadog.appsec.api.blocking.BlockingException;
1111
import datadog.trace.agent.tooling.Instrumenter;
1212
import datadog.trace.agent.tooling.InstrumenterModule;
13+
import datadog.trace.bootstrap.InstrumentationContext;
1314
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
15+
import java.util.Collections;
16+
import java.util.Map;
1417
import net.bytebuddy.asm.Advice;
1518
import net.bytebuddy.description.type.TypeDescription;
1619
import net.bytebuddy.implementation.bytecode.assign.Assigner;
1720
import net.bytebuddy.matcher.ElementMatcher;
1821
import org.apache.hc.core5.http.ClassicHttpRequest;
1922
import org.apache.hc.core5.http.HttpHost;
23+
import org.apache.hc.core5.http.HttpRequest;
2024
import org.apache.hc.core5.http.io.HttpClientResponseHandler;
2125
import org.apache.hc.core5.http.protocol.HttpContext;
2226

@@ -68,6 +72,11 @@ public String[] helperClassNames() {
6872
};
6973
}
7074

75+
@Override
76+
public Map<String, String> contextStore() {
77+
return Collections.singletonMap("org.apache.hc.core5.http.HttpRequest", "java.lang.Integer");
78+
}
79+
7180
@Override
7281
public void methodAdvice(MethodTransformer transformer) {
7382
transformer.applyAdvice(
@@ -117,20 +126,28 @@ public static class RequestAdvice {
117126
@Advice.OnMethodEnter(suppress = Throwable.class)
118127
public static AgentScope methodEnter(@Advice.Argument(0) final ClassicHttpRequest request) {
119128
try {
120-
return HelperMethods.doMethodEnter(request);
129+
return HelperMethods.doMethodEnter(
130+
InstrumentationContext.get(HttpRequest.class, Integer.class), request);
121131
} catch (BlockingException e) {
122-
HelperMethods.onBlockingRequest();
132+
HelperMethods.onBlockingRequest(
133+
InstrumentationContext.get(HttpRequest.class, Integer.class), request);
123134
// re-throw blocking exceptions
124135
throw e;
125136
}
126137
}
127138

128139
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
129140
public static void methodExit(
141+
@Advice.Argument(0) final ClassicHttpRequest request,
130142
@Advice.Enter final AgentScope scope,
131143
@Advice.Return final Object result,
132144
@Advice.Thrown final Throwable throwable) {
133-
HelperMethods.doMethodExit(scope, result, throwable);
145+
HelperMethods.doMethodExit(
146+
InstrumentationContext.get(HttpRequest.class, Integer.class),
147+
request,
148+
scope,
149+
result,
150+
throwable);
134151
}
135152
}
136153

@@ -140,20 +157,28 @@ public static AgentScope methodEnter(
140157
@Advice.Argument(0) final HttpHost host,
141158
@Advice.Argument(1) final ClassicHttpRequest request) {
142159
try {
143-
return HelperMethods.doMethodEnter(host, request);
160+
return HelperMethods.doMethodEnter(
161+
InstrumentationContext.get(HttpRequest.class, Integer.class), host, request);
144162
} catch (BlockingException e) {
145-
HelperMethods.onBlockingRequest();
163+
HelperMethods.onBlockingRequest(
164+
InstrumentationContext.get(HttpRequest.class, Integer.class), request);
146165
// re-throw blocking exceptions
147166
throw e;
148167
}
149168
}
150169

151170
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
152171
public static void methodExit(
172+
@Advice.Argument(1) final ClassicHttpRequest request,
153173
@Advice.Enter final AgentScope scope,
154174
@Advice.Return final Object result,
155175
@Advice.Thrown final Throwable throwable) {
156-
HelperMethods.doMethodExit(scope, result, throwable);
176+
HelperMethods.doMethodExit(
177+
InstrumentationContext.get(HttpRequest.class, Integer.class),
178+
request,
179+
scope,
180+
result,
181+
throwable);
157182
}
158183
}
159184

@@ -171,7 +196,9 @@ public static AgentScope methodEnter(
171196
readOnly = false)
172197
Object handler) {
173198
try {
174-
final AgentScope scope = HelperMethods.doMethodEnter(host, request);
199+
final AgentScope scope =
200+
HelperMethods.doMethodEnter(
201+
InstrumentationContext.get(HttpRequest.class, Integer.class), host, request);
175202
// Wrap the handler so we capture the status code
176203
if (null != scope && handler instanceof HttpClientResponseHandler) {
177204
handler =
@@ -180,18 +207,25 @@ public static AgentScope methodEnter(
180207
}
181208
return scope;
182209
} catch (BlockingException e) {
183-
HelperMethods.onBlockingRequest();
210+
HelperMethods.onBlockingRequest(
211+
InstrumentationContext.get(HttpRequest.class, Integer.class), request);
184212
// re-throw blocking exceptions
185213
throw e;
186214
}
187215
}
188216

189217
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
190218
public static void methodExit(
219+
@Advice.Argument(1) final ClassicHttpRequest request,
191220
@Advice.Enter final AgentScope scope,
192221
@Advice.Return final Object result,
193222
@Advice.Thrown final Throwable throwable) {
194-
HelperMethods.doMethodExit(scope, result, throwable);
223+
HelperMethods.doMethodExit(
224+
InstrumentationContext.get(HttpRequest.class, Integer.class),
225+
request,
226+
scope,
227+
result,
228+
throwable);
195229
}
196230
}
197231
}

‎dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java‎

Lines changed: 45 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -7,33 +7,43 @@
77
import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.HTTP_REQUEST;
88
import static datadog.trace.instrumentation.apachehttpclient5.HttpHeadersInjectAdapter.SETTER;
99

10-
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
10+
import datadog.trace.bootstrap.ContextStore;
1111
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1212
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
13-
import org.apache.hc.client5.http.classic.HttpClient;
1413
import org.apache.hc.core5.http.HttpHost;
1514
import org.apache.hc.core5.http.HttpRequest;
1615
import org.apache.hc.core5.http.HttpResponse;
1716

1817
public class HelperMethods {
19-
public static AgentScope doMethodEnter(final HttpRequest request) {
20-
final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(HttpClient.class);
21-
if (callDepth > 0) {
18+
19+
public static AgentScope doMethodEnter(
20+
final ContextStore<HttpRequest, Integer> depthStore, final HttpRequest request) {
21+
int depth = incrementDepth(depthStore, request);
22+
if (depth > 1) {
2223
return null;
2324
}
24-
2525
return activateHttpSpan(request);
2626
}
2727

28-
public static AgentScope doMethodEnter(HttpHost host, HttpRequest request) {
29-
final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(HttpClient.class);
30-
if (callDepth > 0) {
28+
public static AgentScope doMethodEnter(
29+
final ContextStore<HttpRequest, Integer> depthStore,
30+
HttpHost host,
31+
final HttpRequest request) {
32+
int depth = incrementDepth(depthStore, request);
33+
if (depth > 1) {
3134
return null;
3235
}
33-
3436
return activateHttpSpan(new HostAndRequestAsHttpUriRequest(host, request));
3537
}
3638

39+
private static int incrementDepth(
40+
final ContextStore<HttpRequest, Integer> depthStore, final HttpRequest request) {
41+
Integer depth = depthStore.get(request);
42+
depth = depth == null ? 1 : (depth + 1);
43+
depthStore.put(request, depth);
44+
return depth;
45+
}
46+
3747
private static AgentScope activateHttpSpan(final HttpRequest request) {
3848
final AgentSpan span = startSpan(HTTP_REQUEST);
3949
final AgentScope scope = activateSpan(span);
@@ -51,26 +61,41 @@ private static AgentScope activateHttpSpan(final HttpRequest request) {
5161
}
5262

5363
public static void doMethodExit(
54-
final AgentScope scope, final Object result, final Throwable throwable) {
55-
if (scope == null) {
56-
return;
57-
}
58-
final AgentSpan span = scope.span();
64+
final ContextStore<HttpRequest, Integer> depthStore,
65+
final HttpRequest request,
66+
final AgentScope scope,
67+
final Object result,
68+
final Throwable throwable) {
5969
try {
70+
if (scope == null) {
71+
return;
72+
}
73+
final AgentSpan span = scope.span();
6074
if (result instanceof HttpResponse) {
6175
DECORATE.onResponse(span, (HttpResponse) result);
6276
} // else they probably provided a ResponseHandler.
6377

6478
DECORATE.onError(span, throwable);
6579
DECORATE.beforeFinish(span);
6680
} finally {
67-
scope.close();
68-
span.finish();
69-
CallDepthThreadLocalMap.reset(HttpClient.class);
81+
if (scope != null) {
82+
AgentSpan span = scope.span();
83+
scope.close();
84+
span.finish();
85+
}
86+
Integer depth = depthStore.get(request);
87+
if (depth != null && depth > 0) {
88+
depthStore.put(request, depth - 1);
89+
}
7090
}
7191
}
7292

73-
public static void onBlockingRequest() {
74-
CallDepthThreadLocalMap.reset(HttpClient.class);
93+
/**
94+
* Cleans up state when a BlockingException is thrown from methodEnter. Since the exception
95+
* unwinds the whole stack, we this request's depth to 0.
96+
*/
97+
public static void onBlockingRequest(
98+
final ContextStore<HttpRequest, Integer> depthStore, final HttpRequest request) {
99+
depthStore.put(request, 0);
75100
}
76101
}

‎dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/test/groovy/ApacheHttpClientTest.groovy‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -215,3 +215,55 @@ class ApacheClientResponseHandlerAllV0Test extends ApacheClientResponseHandlerAl
215215
class ApacheClientResponseHandlerAllV1ForkedTest extends ApacheClientResponseHandlerAll implements TestingGenericHttpNamingConventions.ClientV1 {
216216
}
217217

218+
/**
219+
* Tests that HTTP calls made from within an exec interceptor (or other nested execute() context)
220+
* using a different HttpClient instance are instrumented. See
221+
* https://github.com/DataDog/dd-trace-java/issues/10383
222+
*/
223+
@Timeout(5)
224+
class ApacheClientNestedExecuteTest extends ApacheHttpClientTest<ClassicHttpRequest> {
225+
226+
@Override
227+
ClassicHttpRequest createRequest(String method, URI uri) {
228+
return new BasicClassicHttpRequest(method, uri)
229+
}
230+
231+
@Override
232+
CloseableHttpResponse executeRequest(ClassicHttpRequest request, URI uri) {
233+
return client.execute(request)
234+
}
235+
236+
def "nested execute from different client (e.g. interceptor token fetch) creates spans for both requests"() {
237+
when:
238+
def tokenUri = server.address.resolve("/success")
239+
def mainUri = server.address.resolve("/success")
240+
def innerCode = new int[1]
241+
// Use a separate client for the inner request (as in issue: token client in interceptor)
242+
def tokenClient = HttpClients.custom()
243+
.setConnectionManager(new BasicHttpClientConnectionManager())
244+
.setDefaultRequestConfig(RequestConfig.custom()
245+
.setConnectTimeout(CONNECT_TIMEOUT_MS, TimeUnit.MILLISECONDS)
246+
.build()).build()
247+
def response = client.execute(new BasicClassicHttpRequest("GET", mainUri), { r ->
248+
innerCode[0] = tokenClient.execute(new BasicClassicHttpRequest("GET", tokenUri), { inner -> inner.code })
249+
return r
250+
})
251+
tokenClient.close()
252+
253+
then:
254+
response != null
255+
innerCode[0] == 200
256+
assertTraces(3) {
257+
sortSpansByStart()
258+
trace(size(2)) {
259+
sortSpansByStart()
260+
// Main request starts first, then token request (nested) - order by start time
261+
clientSpan(it, null, "GET", false, false, mainUri)
262+
clientSpan(it, span(0), "GET", false, false, tokenUri)
263+
}
264+
server.distributedRequestTrace(it, trace(0)[0])
265+
server.distributedRequestTrace(it, trace(0)[1])
266+
}
267+
}
268+
}
269+

0 commit comments

Comments
 (0)