Implement instrumentation for async requests
MODIFY
dd-java-agent/instrumentation/google-http-client/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientInstrumentation.java
dd-java-agent/instrumentation/google-http-client/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientInstrumentation.java
@@ -5,13 +5,15 @@ import static java.util.Collections.singletonMap;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
import com.google.api.client.http.HttpRequest;
import com.google.api.client.http.HttpResponse;
import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
import io.opentracing.Scope;
import io.opentracing.Span;
import io.opentracing.log.Fields;
@@ -19,6 +21,8 @@ import io.opentracing.propagation.Format;
import io.opentracing.propagation.TextMap;
import io.opentracing.tag.Tags;
import io.opentracing.util.GlobalTracer;
+import java.util.Collections;
+import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import net.bytebuddy.asm.Advice;
@@ -38,6 +42,12 @@ public class GoogleHttpClientInstrumentation extends Instrumenter.Default {
return named("com.google.api.client.http.HttpRequest");
}
+ @Override
+ public Map<String, String> contextStore() {
+ return Collections.singletonMap(
+ "com.google.api.client.http.HttpRequest", RequestState.class.getName());
+ }
+
@Override
public String[] helperClassNames() {
return new String[] {
@@ -45,45 +55,71 @@ public class GoogleHttpClientInstrumentation extends Instrumenter.Default {
"datadog.trace.agent.decorator.ClientDecorator",
"datadog.trace.agent.decorator.HttpClientDecorator",
packageName + ".GoogleHttpClientDecorator",
+ packageName + ".RequestState",
getClass().getName() + "$GoogleHttpClientAdvice",
+ getClass().getName() + "$GoogleHttpClientAsyncAdvice",
getClass().getName() + "$HeadersInjectAdapter"
};
}
@Override
public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
- return singletonMap(
- isMethod().and(isPublic()).and(named("execute").and(takesArguments(0))),
+ final Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
+ transformers.put(
+ isMethod().and(isPublic()).and(named("execute")).and(takesArguments(0)),
GoogleHttpClientAdvice.class.getName());
+
+ transformers.put(
+ isMethod()
+ .and(isPublic())
+ .and(named("executeAsync"))
+ .and(takesArguments(1))
+ .and(takesArgument(0, (named("java.util.concurrent.Executor")))),
+ GoogleHttpClientAsyncAdvice.class.getName());
+
+ return transformers;
}
public static class GoogleHttpClientAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
- public static Scope methodEnter(@Advice.This final HttpRequest request) {
- final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(HttpRequest.class);
- if (callDepth > 0) {
- return null;
+ public static void methodEnter(@Advice.This final HttpRequest request) {
+
+ final ContextStore<HttpRequest, RequestState> contextStore =
+ InstrumentationContext.get(HttpRequest.class, RequestState.class);
+
+ RequestState state = contextStore.get(request);
+
+ if (state == null) {
+ state = new RequestState(GlobalTracer.get().buildSpan("http.request").start());
+ contextStore.put(request, state);
}
- final Span span = GlobalTracer.get().buildSpan("http.request").start();
- final Scope scope = GlobalTracer.get().scopeManager().activate(span, false);
+ final Span span = state.getSpan();
+
+ try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, false)) {
DECORATE.afterStart(span);
DECORATE.onRequest(span, request);
GlobalTracer.get()
.inject(span.context(), Format.Builtin.HTTP_HEADERS, new HeadersInjectAdapter(request));
- return scope;
+ }
}
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void methodExit(
- @Advice.Enter final Scope scope,
+ @Advice.This final HttpRequest request,
@Advice.Return final HttpResponse response,
@Advice.Thrown final Throwable throwable) {
- if (scope != null) {
- try {
- final Span span = scope.span();
+ final ContextStore<HttpRequest, RequestState> contextStore =
+ InstrumentationContext.get(HttpRequest.class, RequestState.class);
+ final RequestState state = contextStore.get(request);
+
+ if (state != null) {
+ final Span span = state.getSpan();
+
+ try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, false)) {
DECORATE.onResponse(span, response);
+ DECORATE.onError(span, throwable);
// If HttpRequest.setThrowExceptionOnExecuteError is set to false, there are no exceptions
// for a failed request. Thus, check the response code
@@ -91,13 +127,46 @@ public class GoogleHttpClientInstrumentation extends Instrumenter.Default {
Tags.ERROR.set(span, true);
span.log(singletonMap(Fields.MESSAGE, response.getStatusMessage()));
}
+
+ DECORATE.beforeFinish(span);
+ span.finish();
+ }
+ }
+ }
+ }
+
+ public static class GoogleHttpClientAsyncAdvice {
+
+ @Advice.OnMethodEnter(suppress = Throwable.class)
+ public static void methodEnter(@Advice.This final HttpRequest request) {
+ final Span span = GlobalTracer.get().buildSpan("http.request").start();
+
+ final ContextStore<HttpRequest, RequestState> contextStore =
+ InstrumentationContext.get(HttpRequest.class, RequestState.class);
+
+ final RequestState state = new RequestState(span);
+ contextStore.put(request, state);
+ }
+
+ @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+ public static void methodExit(
+ @Advice.This final HttpRequest request, @Advice.Thrown final Throwable throwable) {
+
+ if (throwable != null) {
+
+ final ContextStore<HttpRequest, RequestState> contextStore =
+ InstrumentationContext.get(HttpRequest.class, RequestState.class);
+ final RequestState state = contextStore.get(request);
+
+ if (state != null) {
+ final Span span = state.getSpan();
+
+ try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, false)) {
DECORATE.onError(span, throwable);
DECORATE.beforeFinish(span);
span.finish();
- } finally {
- scope.close();
- CallDepthThreadLocalMap.reset(HttpRequest.class);
+ }
}
}
}

ADD
dd-java-agent/instrumentation/google-http-client/src/main/java/datadog/trace/instrumentation/googlehttpclient/RequestState.java
+package datadog.trace.instrumentation.googlehttpclient;
+
+import io.opentracing.Span;
+import lombok.Data;
+import lombok.NonNull;
+
+@Data
+public class RequestState {
+ @NonNull public Span span;
+}

ADD
dd-java-agent/instrumentation/google-http-client/src/test/groovy/AbstractGoogleHttpClientTest.groovy
+import com.google.api.client.http.GenericUrl
+import com.google.api.client.http.HttpRequest
+import com.google.api.client.http.HttpResponse
+import com.google.api.client.http.javanet.NetHttpTransport
+import datadog.trace.agent.test.base.HttpClientTest
+import datadog.trace.api.DDSpanTypes
+import datadog.trace.instrumentation.googlehttpclient.GoogleHttpClientDecorator
+import spock.lang.Shared
+
+abstract class AbstractGoogleHttpClientTest extends HttpClientTest<GoogleHttpClientDecorator> {
+
+ @Shared
+ def requestFactory = new NetHttpTransport().createRequestFactory();
+
+ @Override
+ int doRequest(String method, URI uri, Map<String, String> headers, Closure callback) {
+ doRequest(method, uri, headers, callback, false);
+ }
+
+ int doRequest(String method, URI uri, Map<String, String> headers, Closure callback, boolean throwExceptionOnError) {
+ GenericUrl genericUrl = new GenericUrl(uri)
+
+ HttpRequest request = requestFactory.buildRequest(method, genericUrl, null)
+ request.getHeaders().putAll(headers)
+ request.setThrowExceptionOnExecuteError(throwExceptionOnError)
+
+ HttpResponse response = executeRequest(request);
+ callback?.call()
+
+ return response.getStatusCode()
+ }
+
+ abstract HttpResponse executeRequest(HttpRequest request);
+
+ @Override
+ GoogleHttpClientDecorator decorator() {
+ return GoogleHttpClientDecorator.DECORATE
+ }
+
+ @Override
+ boolean testRedirects() {
+ // Circular redirects don't throw an exception with Google Http Client
+ return false
+ }
+
+ def "error traces when exception is not thrown"() {
+ given:
+ def uri = server.address.resolve("/error")
+
+ when:
+ def status = doRequest(method, uri)
+
+ then:
+ status == 500
+ assertTraces(2) {
+ server.distributedRequestTrace(it, 0, trace(1).last())
+ trace(1, size(1)) {
+ span(0) {
+ resourceName "$method $uri.path"
+ spanType DDSpanTypes.HTTP_CLIENT
+ errored true
+ }
+ }
+ }
+
+ where:
+ method = "GET"
+ }
+}

ADD
dd-java-agent/instrumentation/google-http-client/src/test/groovy/GoogleHttpClientAsyncTest.groovy
+import com.google.api.client.http.HttpRequest
+import com.google.api.client.http.HttpResponse
+
+class GoogleHttpClientAsyncTest extends AbstractGoogleHttpClientTest {
+ @Override
+ HttpResponse executeRequest(HttpRequest request) {
+ return request.executeAsync().get()
+ }
+}

MODIFY
dd-java-agent/instrumentation/google-http-client/src/test/groovy/GoogleHttpClientTest.groovy
dd-java-agent/instrumentation/google-http-client/src/test/groovy/GoogleHttpClientTest.groovy
-import com.google.api.client.http.GenericUrl
import com.google.api.client.http.HttpRequest
import com.google.api.client.http.HttpResponse
-import com.google.api.client.http.javanet.NetHttpTransport
-import datadog.trace.agent.test.base.HttpClientTest
-import datadog.trace.api.DDSpanTypes
-import datadog.trace.instrumentation.googlehttpclient.GoogleHttpClientDecorator
-import spock.lang.Shared
-
-class GoogleHttpClientTest extends HttpClientTest<GoogleHttpClientDecorator> {
-
- @Shared
- def requestFactory = new NetHttpTransport().createRequestFactory();
+class GoogleHttpClientTest extends AbstractGoogleHttpClientTest {
@Override
- int doRequest(String method, URI uri, Map<String, String> headers, Closure callback) {
- doRequest(method, uri, headers, callback, false);
- }
-
- int doRequest(String method, URI uri, Map<String, String> headers, Closure callback, boolean throwExceptionOnError) {
- GenericUrl genericUrl = new GenericUrl(uri)
-
- HttpRequest request = requestFactory.buildRequest(method, genericUrl, null)
- request.getHeaders().putAll(headers)
- request.setThrowExceptionOnExecuteError(throwExceptionOnError)
-
- HttpResponse response = request.execute()
- callback?.call()
-
- return response.getStatusCode()
- }
-
- @Override
- GoogleHttpClientDecorator decorator() {
- return GoogleHttpClientDecorator.DECORATE
- }
-
- @Override
- boolean testRedirects() {
- // Circular redirects don't throw an exception with Google Http Client
- return false
- }
-
- def "error traces when exception is not thrown"() {
- given:
- def uri = server.address.resolve("/error")
-
- when:
- def status = doRequest(method, uri)
-
- then:
- status == 500
- assertTraces(2) {
- server.distributedRequestTrace(it, 0, trace(1).last())
- trace(1, size(1)) {
- span(0) {
- resourceName "$method $uri.path"
- spanType DDSpanTypes.HTTP_CLIENT
- errored true
- }
- }
- }
-
- where:
- method = "GET"
+ HttpResponse executeRequest(HttpRequest request) {
+ return request.execute();
}
}

Implement instrumentation for async requests