Implement instrumentation for async requests

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);
-      DECORATE.afterStart(span);
-      DECORATE.onRequest(span, request);
-      GlobalTracer.get()
-          .inject(span.context(), Format.Builtin.HTTP_HEADERS, new HeadersInjectAdapter(request));
-      return scope;
+      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));
+      }
     }
 
     @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.onError(span, throwable);
 
           DECORATE.beforeFinish(span);
           span.finish();
-        } finally {
-          scope.close();
-          CallDepthThreadLocalMap.reset(HttpRequest.class);
+        }
+      }
+    }
+  }
+
+  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();
+          }
         }
       }
     }

dd-java-agent/instrumentation/google-http-client/src/main/java/datadog/trace/instrumentation/googlehttpclient/RequestState.java

@@ -0,0 +1,10 @@
+package datadog.trace.instrumentation.googlehttpclient;
+
+import io.opentracing.Span;
+import lombok.Data;
+import lombok.NonNull;
+
+@Data
+public class RequestState {
+  @NonNull public Span span;
+}

dd-java-agent/instrumentation/google-http-client/src/test/groovy/AbstractGoogleHttpClientTest.groovy

@@ -0,0 +1,69 @@
+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"
+  }
+}

dd-java-agent/instrumentation/google-http-client/src/test/groovy/GoogleHttpClientAsyncTest.groovy

@@ -0,0 +1,9 @@
+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()
+  }
+}

dd-java-agent/instrumentation/google-http-client/src/test/groovy/GoogleHttpClientTest.groovy

@@ -1,67 +1,9 @@
-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();
   }
 }