Add instrumentation for Apache HttpAsyncClient

Extract http client tests to shared class.

dd-java-agent/instrumentation/apache-httpasyncclient-4/apache-httpasyncclient-4.gradle

@@ -0,0 +1,38 @@
+ext {
+  minJavaVersionForTests = JavaVersion.VERSION_1_8
+}
+
+muzzle {
+  pass {
+    group = "org.apache.httpcomponents"
+    module = "httpasyncclient"
+    versions = "[4.0,)"
+    assertInverse = true
+  }
+}
+
+apply from: "${rootDir}/gradle/java.gradle"
+apply plugin: 'org.unbroken-dome.test-sets'
+
+testSets {
+  latestDepTest {
+    dirName = 'test'
+  }
+}
+
+dependencies {
+  compileOnly group: 'org.apache.httpcomponents', name: 'httpasyncclient', version: '4.0'
+
+
+  compile project(':dd-java-agent:agent-tooling')
+
+  compile deps.bytebuddy
+  compile deps.opentracing
+  annotationProcessor deps.autoservice
+  implementation deps.autoservice
+
+  testCompile project(':dd-java-agent:testing')
+  testCompile group: 'org.apache.httpcomponents', name: 'httpasyncclient', version: '4.0'
+
+  latestDepTestCompile group: 'org.apache.httpcomponents', name: 'httpasyncclient', version: '+'
+}

dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/ApacheHttpAsyncClientDecorator.java

@@ -0,0 +1,73 @@
+package datadog.trace.instrumentation.apachehttpasyncclient;
+
+import datadog.trace.agent.decorator.HttpClientDecorator;
+import java.net.URI;
+import org.apache.http.HttpRequest;
+import org.apache.http.HttpResponse;
+import org.apache.http.RequestLine;
+import org.apache.http.StatusLine;
+import org.apache.http.protocol.HttpContext;
+import org.apache.http.protocol.HttpCoreContext;
+
+public class ApacheHttpAsyncClientDecorator extends HttpClientDecorator<HttpRequest, HttpContext> {
+  public static final ApacheHttpAsyncClientDecorator DECORATE =
+      new ApacheHttpAsyncClientDecorator();
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"httpasyncclient", "apache-httpasyncclient"};
+  }
+
+  @Override
+  protected String component() {
+    return "apache-httpasyncclient";
+  }
+
+  @Override
+  protected String method(final HttpRequest request) {
+    final RequestLine requestLine = request.getRequestLine();
+    return requestLine == null ? null : requestLine.getMethod();
+  }
+
+  @Override
+  protected String url(final HttpRequest request) {
+    final RequestLine requestLine = request.getRequestLine();
+    return requestLine == null ? null : requestLine.getUri();
+  }
+
+  @Override
+  protected String hostname(final HttpRequest request) {
+    final RequestLine requestLine = request.getRequestLine();
+    if (requestLine != null) {
+      final URI uri = URI.create(requestLine.getUri());
+      if (uri != null) {
+        return uri.getHost();
+      }
+    }
+    return null;
+  }
+
+  @Override
+  protected Integer port(final HttpRequest request) {
+    final RequestLine requestLine = request.getRequestLine();
+    if (requestLine != null) {
+      final URI uri = URI.create(requestLine.getUri());
+      if (uri != null) {
+        return uri.getPort();
+      }
+    }
+    return null;
+  }
+
+  @Override
+  protected Integer status(final HttpContext context) {
+    final Object responseObject = context.getAttribute(HttpCoreContext.HTTP_RESPONSE);
+    if (responseObject instanceof HttpResponse) {
+      final StatusLine statusLine = ((HttpResponse) responseObject).getStatusLine();
+      if (statusLine != null) {
+        return statusLine.getStatusCode();
+      }
+    }
+    return null;
+  }
+}

dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/ApacheHttpAsyncClientInstrumentation.java

@@ -0,0 +1,255 @@
+package datadog.trace.instrumentation.apachehttpasyncclient;
+
+import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
+import static datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator.DECORATE;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.context.TraceScope;
+import io.opentracing.Scope;
+import io.opentracing.Span;
+import io.opentracing.Tracer;
+import io.opentracing.propagation.Format;
+import io.opentracing.propagation.TextMap;
+import io.opentracing.util.GlobalTracer;
+import java.io.IOException;
+import java.util.Iterator;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.apache.http.HttpException;
+import org.apache.http.HttpHost;
+import org.apache.http.HttpRequest;
+import org.apache.http.concurrent.FutureCallback;
+import org.apache.http.nio.ContentEncoder;
+import org.apache.http.nio.IOControl;
+import org.apache.http.nio.protocol.HttpAsyncRequestProducer;
+import org.apache.http.protocol.HttpContext;
+
+@AutoService(Instrumenter.class)
+public class ApacheHttpAsyncClientInstrumentation extends Instrumenter.Default {
+
+  public ApacheHttpAsyncClientInstrumentation() {
+    super("httpclient", "apache-httpclient", "apache-http-client");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return safeHasSuperType(named("org.apache.http.nio.client.HttpAsyncClient"));
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      getClass().getName() + "$HttpHeadersInjectAdapter",
+      getClass().getName() + "$DelegatingRequestProducer",
+      getClass().getName() + "$TraceContinuedFutureCallback",
+      "datadog.trace.agent.decorator.BaseDecorator",
+      "datadog.trace.agent.decorator.ClientDecorator",
+      "datadog.trace.agent.decorator.HttpClientDecorator",
+      packageName + ".ApacheHttpAsyncClientDecorator",
+    };
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isMethod()
+            .and(named("execute"))
+            .and(takesArguments(4))
+            .and(takesArgument(0, named("org.apache.http.nio.protocol.HttpAsyncRequestProducer")))
+            .and(takesArgument(1, named("org.apache.http.nio.protocol.HttpAsyncResponseConsumer")))
+            .and(takesArgument(2, named("org.apache.http.protocol.HttpContext")))
+            .and(takesArgument(3, named("org.apache.http.concurrent.FutureCallback"))),
+        ClientAdvice.class.getName());
+  }
+
+  public static class ClientAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static Span methodEnter(
+        @Advice.Argument(value = 0, readOnly = false) HttpAsyncRequestProducer requestProducer,
+        @Advice.Argument(2) final HttpContext context,
+        @Advice.Argument(value = 3, readOnly = false) FutureCallback<?> futureCallback) {
+
+      final Tracer tracer = GlobalTracer.get();
+      final Scope parentScope = tracer.scopeManager().active();
+      final Span clientSpan = tracer.buildSpan("http.request").start();
+
+      DECORATE.afterStart(clientSpan);
+      requestProducer = new DelegatingRequestProducer(clientSpan, requestProducer);
+      futureCallback =
+          new TraceContinuedFutureCallback(parentScope, clientSpan, context, futureCallback);
+
+      return clientSpan;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void methodExit(
+        @Advice.Enter final Span span,
+        @Advice.Return final Object result,
+        @Advice.Thrown final Throwable throwable) {
+      if (throwable != null) {
+        DECORATE.onError(span, throwable);
+        DECORATE.beforeFinish(span);
+        span.finish();
+      }
+    }
+  }
+
+  public static class DelegatingRequestProducer implements HttpAsyncRequestProducer {
+    final Span span;
+    final HttpAsyncRequestProducer delegate;
+
+    public DelegatingRequestProducer(final Span span, final HttpAsyncRequestProducer delegate) {
+      this.span = span;
+      this.delegate = delegate;
+    }
+
+    @Override
+    public HttpHost getTarget() {
+      return delegate.getTarget();
+    }
+
+    @Override
+    public HttpRequest generateRequest() throws IOException, HttpException {
+      final HttpRequest request = delegate.generateRequest();
+      DECORATE.onRequest(span, request);
+
+      GlobalTracer.get()
+          .inject(
+              span.context(), Format.Builtin.HTTP_HEADERS, new HttpHeadersInjectAdapter(request));
+
+      return request;
+    }
+
+    @Override
+    public void produceContent(final ContentEncoder encoder, final IOControl ioctrl)
+        throws IOException {
+      delegate.produceContent(encoder, ioctrl);
+    }
+
+    @Override
+    public void requestCompleted(final HttpContext context) {
+      delegate.requestCompleted(context);
+    }
+
+    @Override
+    public void failed(final Exception ex) {
+      delegate.failed(ex);
+    }
+
+    @Override
+    public boolean isRepeatable() {
+      return delegate.isRepeatable();
+    }
+
+    @Override
+    public void resetRequest() throws IOException {
+      delegate.resetRequest();
+    }
+
+    @Override
+    public void close() throws IOException {
+      delegate.close();
+    }
+  }
+
+  public static class TraceContinuedFutureCallback<T> implements FutureCallback<T> {
+    private final TraceScope.Continuation parentContinuation;
+    private final Span clientSpan;
+    private final HttpContext context;
+    private final FutureCallback<T> delegate;
+
+    public TraceContinuedFutureCallback(
+        final Scope parentScope,
+        final Span clientSpan,
+        final HttpContext context,
+        final FutureCallback<T> delegate) {
+      if (parentScope instanceof TraceScope) {
+        parentContinuation = ((TraceScope) parentScope).capture();
+      } else {
+        parentContinuation = null;
+      }
+      this.clientSpan = clientSpan;
+      this.context = context;
+      this.delegate = delegate;
+    }
+
+    @Override
+    public void completed(final T result) {
+      DECORATE.onResponse(clientSpan, context);
+      DECORATE.beforeFinish(clientSpan);
+      clientSpan.finish(); // Finish span before calling delegate
+
+      if (parentContinuation == null) {
+        delegate.completed(result);
+      } else {
+        try (final TraceScope scope = parentContinuation.activate()) {
+          scope.setAsyncPropagation(true);
+          delegate.completed(result);
+        }
+      }
+    }
+
+    @Override
+    public void failed(final Exception ex) {
+      DECORATE.onResponse(clientSpan, context);
+      DECORATE.onError(clientSpan, ex);
+      DECORATE.beforeFinish(clientSpan);
+      clientSpan.finish(); // Finish span before calling delegate
+
+      if (parentContinuation == null) {
+        delegate.failed(ex);
+      } else {
+        try (final TraceScope scope = parentContinuation.activate()) {
+          scope.setAsyncPropagation(true);
+          delegate.failed(ex);
+        }
+      }
+    }
+
+    @Override
+    public void cancelled() {
+      DECORATE.onResponse(clientSpan, context);
+      DECORATE.beforeFinish(clientSpan);
+      clientSpan.finish(); // Finish span before calling delegate
+
+      if (parentContinuation == null) {
+        delegate.cancelled();
+      } else {
+        try (final TraceScope scope = parentContinuation.activate()) {
+          scope.setAsyncPropagation(true);
+          delegate.cancelled();
+        }
+      }
+    }
+  }
+
+  public static class HttpHeadersInjectAdapter implements TextMap {
+
+    private final HttpRequest httpRequest;
+
+    public HttpHeadersInjectAdapter(final HttpRequest httpRequest) {
+      this.httpRequest = httpRequest;
+    }
+
+    @Override
+    public void put(final String key, final String value) {
+      httpRequest.addHeader(key, value);
+    }
+
+    @Override
+    public Iterator<Map.Entry<String, String>> iterator() {
+      throw new UnsupportedOperationException(
+          "This class should be used only with tracer#inject()");
+    }
+  }
+}

dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/ApacheHttpClientRedirectInstrumentation.java

@@ -0,0 +1,65 @@
+package datadog.trace.instrumentation.apachehttpasyncclient;
+
+import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.implementation.bytecode.assign.Assigner;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.apache.http.Header;
+import org.apache.http.HttpRequest;
+
+/**
+ * Early versions don't copy headers over on redirect. This instrumentation copies our headers over
+ * manually. Inspired by
+ * https://github.com/elastic/apm-agent-java/blob/master/apm-agent-plugins/apm-apache-httpclient-plugin/src/main/java/co/elastic/apm/agent/httpclient/ApacheHttpAsyncClientRedirectInstrumentation.java
+ */
+@AutoService(Instrumenter.class)
+public class ApacheHttpClientRedirectInstrumentation extends Instrumenter.Default {
+
+  public ApacheHttpClientRedirectInstrumentation() {
+    super("httpasyncclient", "apache-httpasyncclient");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return safeHasSuperType(named("org.apache.http.client.RedirectStrategy"));
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isMethod()
+            .and(named("getRedirect"))
+            .and(takesArgument(0, named("org.apache.http.HttpRequest"))),
+        ClientRedirectAdvice.class.getName());
+  }
+
+  public static class ClientRedirectAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    private static void onAfterExecute(
+        @Advice.Argument(value = 0) final HttpRequest original,
+        @Advice.Return(typing = Assigner.Typing.DYNAMIC) final HttpRequest redirect) {
+      if (redirect == null) {
+        return;
+      }
+
+      for (final Header header : original.getAllHeaders()) {
+        final String name = header.getName().toLowerCase();
+        if (name.startsWith("x-datadog-") || name.startsWith("x-b3-")) {
+          if (!redirect.containsHeader(header.getName())) {
+            redirect.setHeader(header.getName(), header.getValue());
+          }
+        }
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/apache-httpasyncclient-4/src/test/groovy/ApacheHttpAsyncClientCallbackTest.groovy

@@ -0,0 +1,76 @@
+import datadog.trace.agent.test.base.HttpClientTest
+import datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator
+import io.opentracing.util.GlobalTracer
+import org.apache.http.HttpResponse
+import org.apache.http.client.methods.HttpGet
+import org.apache.http.concurrent.FutureCallback
+import org.apache.http.impl.nio.client.HttpAsyncClients
+import org.apache.http.message.BasicHeader
+import spock.lang.AutoCleanup
+import spock.lang.Shared
+
+import java.util.concurrent.CompletableFuture
+
+class ApacheHttpAsyncClientCallbackTest extends HttpClientTest<ApacheHttpAsyncClientDecorator> {
+
+  @AutoCleanup
+  @Shared
+  def client = HttpAsyncClients.createDefault()
+
+  def setupSpec() {
+    client.start()
+  }
+
+  @Override
+  int doRequest(String method, URI uri, Map<String, String> headers, Closure callback) {
+    assert method == "GET"
+
+    def hasParent = GlobalTracer.get().activeSpan() != null
+
+    HttpGet request = new HttpGet(uri)
+    headers.entrySet().each {
+      request.addHeader(new BasicHeader(it.key, it.value))
+    }
+
+    def responseFuture = new CompletableFuture<>()
+
+    client.execute(request, new FutureCallback<HttpResponse>() {
+
+      @Override
+      void completed(HttpResponse result) {
+        if (hasParent && GlobalTracer.get().activeSpan() == null) {
+          responseFuture.completeExceptionally(new Exception("Missing span in scope"))
+        } else {
+          responseFuture.complete(result.statusLine.statusCode)
+        }
+        callback?.call()
+      }
+
+      @Override
+      void failed(Exception ex) {
+        if (hasParent && GlobalTracer.get().activeSpan() == null) {
+          responseFuture.completeExceptionally(new Exception("Missing span in scope"))
+        } else {
+          responseFuture.completeExceptionally(ex)
+        }
+      }
+
+      @Override
+      void cancelled() {
+        responseFuture.cancel(true)
+      }
+    })
+
+    return responseFuture.get()
+  }
+
+  @Override
+  ApacheHttpAsyncClientDecorator decorator() {
+    return ApacheHttpAsyncClientDecorator.DECORATE
+  }
+
+  @Override
+  Integer statusOnRedirectError() {
+    return 302
+  }
+}

dd-java-agent/instrumentation/apache-httpasyncclient-4/src/test/groovy/ApacheHttpAsyncClientTest.groovy

@@ -0,0 +1,59 @@
+import datadog.trace.agent.test.base.HttpClientTest
+import datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator
+import org.apache.http.HttpResponse
+import org.apache.http.client.methods.HttpGet
+import org.apache.http.concurrent.FutureCallback
+import org.apache.http.impl.nio.client.HttpAsyncClients
+import org.apache.http.message.BasicHeader
+import spock.lang.AutoCleanup
+import spock.lang.Shared
+
+class ApacheHttpAsyncClientTest extends HttpClientTest<ApacheHttpAsyncClientDecorator> {
+
+  @AutoCleanup
+  @Shared
+  def client = HttpAsyncClients.createDefault()
+
+  def setupSpec() {
+    client.start()
+  }
+
+  @Override
+  int doRequest(String method, URI uri, Map<String, String> headers, Closure callback) {
+    assert method == "GET"
+    HttpGet request = new HttpGet(uri)
+    headers.entrySet().each {
+      request.addHeader(new BasicHeader(it.key, it.value))
+    }
+
+    def handler = callback == null ? null : new FutureCallback<HttpResponse>() {
+
+      @Override
+      void completed(HttpResponse result) {
+        callback()
+      }
+
+      @Override
+      void failed(Exception ex) {
+      }
+
+      @Override
+      void cancelled() {
+      }
+    }
+
+    def response = client.execute(request, handler).get()
+    response.entity.getContent().close() // Make sure the connection is closed.
+    response.statusLine.statusCode
+  }
+
+  @Override
+  ApacheHttpAsyncClientDecorator decorator() {
+    return ApacheHttpAsyncClientDecorator.DECORATE
+  }
+
+  @Override
+  Integer statusOnRedirectError() {
+    return 302
+  }
+}

dd-java-agent/instrumentation/apache-httpclient-4/apache-httpclient-4.gradle

@@ -1,5 +1,3 @@
-apply from: "${rootDir}/gradle/java.gradle"
-
 muzzle {
   fail {
     group = "commons-httpclient"
@@ -21,10 +19,13 @@ muzzle {
   }
 }
 
+apply from: "${rootDir}/gradle/java.gradle"
 apply plugin: 'org.unbroken-dome.test-sets'
 
 testSets {
-  latestDepTest
+  latestDepTest {
+    dirName = 'test'
+  }
 }
 
 dependencies {

dd-java-agent/instrumentation/apache-httpclient-4/src/latestDepTest/groovy/ApacheHttpClientTest.groovy

@@ -1,243 +0,0 @@
-import datadog.opentracing.DDSpan
-import datadog.trace.agent.test.AgentTestRunner
-import datadog.trace.agent.test.asserts.TraceAssert
-import datadog.trace.api.Config
-import datadog.trace.api.DDSpanTypes
-import io.opentracing.tag.Tags
-import org.apache.http.HttpResponse
-import org.apache.http.client.ClientProtocolException
-import org.apache.http.client.HttpClient
-import org.apache.http.client.config.RequestConfig
-import org.apache.http.client.methods.HttpGet
-import org.apache.http.impl.client.BasicResponseHandler
-import org.apache.http.impl.client.HttpClientBuilder
-import org.apache.http.message.BasicHeader
-import spock.lang.AutoCleanup
-import spock.lang.Shared
-
-import static datadog.trace.agent.test.server.http.TestHttpServer.httpServer
-import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
-import static datadog.trace.agent.test.utils.TraceUtils.withConfigOverride
-
-class ApacheHttpClientTest extends AgentTestRunner {
-
-  @AutoCleanup
-  @Shared
-  def server = httpServer {
-    handlers {
-      prefix("success") {
-        handleDistributedRequest()
-        String msg = "Hello."
-        response.status(200).send(msg)
-      }
-      prefix("redirect") {
-        handleDistributedRequest()
-        redirect(server.address.resolve("/success").toURL().toString())
-      }
-      prefix("another-redirect") {
-        handleDistributedRequest()
-        redirect(server.address.resolve("/redirect").toURL().toString())
-      }
-    }
-  }
-  @Shared
-  int port = server.address.port
-  @Shared
-  def successUrl = server.address.resolve("/success")
-  @Shared
-  def redirectUrl = server.address.resolve("/redirect")
-  @Shared
-  def twoRedirectsUrl = server.address.resolve("/another-redirect")
-  @Shared
-  def handler = new BasicResponseHandler()
-
-  final HttpClientBuilder builder = HttpClientBuilder.create()
-  final HttpClient client = builder.build()
-
-  def "trace request with propagation"() {
-    when:
-
-    String response = withConfigOverride(Config.HTTP_CLIENT_HOST_SPLIT_BY_DOMAIN, "$renameService") {
-      runUnderTrace("parent") {
-        if (responseHandler) {
-          client.execute(new HttpGet(successUrl), responseHandler)
-        } else {
-          client.execute(new HttpGet(successUrl)).entity.content.text
-        }
-      }
-    }
-
-    then:
-    response == "Hello."
-    // one trace on the server, one trace on the client
-    assertTraces(2) {
-      server.distributedRequestTrace(it, 0, TEST_WRITER[1][1])
-      trace(1, 2) {
-        parentSpan(it, 0)
-        successClientSpan(it, 1, span(0), renameService)
-      }
-    }
-
-    where:
-    responseHandler << [null, handler]
-    renameService << [false, true]
-  }
-
-  def "trace redirected request with propagation many redirects allowed"() {
-    setup:
-    final RequestConfig.Builder requestConfigBuilder = new RequestConfig.Builder()
-    requestConfigBuilder.setMaxRedirects(10)
-
-    HttpGet request = new HttpGet(redirectUrl)
-    request.setConfig(requestConfigBuilder.build())
-
-    when:
-    HttpResponse response = withConfigOverride(Config.HTTP_CLIENT_HOST_SPLIT_BY_DOMAIN, "$renameService") {
-      runUnderTrace("parent") {
-        client.execute(request)
-      }
-    }
-
-    then:
-    response.getStatusLine().getStatusCode() == 200
-    // two traces on the server, one trace on the client
-    assertTraces(3) {
-      server.distributedRequestTrace(it, 0, TEST_WRITER[2][1])
-      server.distributedRequestTrace(it, 1, TEST_WRITER[2][1])
-      trace(2, 2) {
-        parentSpan(it, 0)
-        successClientSpan(it, 1, span(0), renameService, 200, "redirect")
-      }
-    }
-
-    where:
-    renameService << [false, true]
-  }
-
-  def "trace redirected request with propagation 1 redirect allowed"() {
-    setup:
-    final RequestConfig.Builder requestConfigBuilder = new RequestConfig.Builder()
-    requestConfigBuilder.setMaxRedirects(1)
-    HttpGet request = new HttpGet(redirectUrl)
-    request.setConfig(requestConfigBuilder.build())
-
-    when:
-    HttpResponse response = withConfigOverride(Config.HTTP_CLIENT_HOST_SPLIT_BY_DOMAIN, "$renameService") {
-      runUnderTrace("parent") {
-        client.execute(request)
-      }
-    }
-
-    then:
-    response.getStatusLine().getStatusCode() == 200
-    // two traces on the server, one trace on the client
-    assertTraces(3) {
-      server.distributedRequestTrace(it, 0, TEST_WRITER[2][1])
-      server.distributedRequestTrace(it, 1, TEST_WRITER[2][1])
-      trace(2, 2) {
-        parentSpan(it, 0)
-        successClientSpan(it, 1, span(0), renameService, 200, "redirect")
-      }
-    }
-
-    where:
-    renameService << [false, true]
-  }
-
-  def "trace redirected request with propagation too many redirects"() {
-    setup:
-    final RequestConfig.Builder requestConfigBuilder = new RequestConfig.Builder()
-    requestConfigBuilder.setMaxRedirects(1)
-
-    HttpGet request = new HttpGet(twoRedirectsUrl)
-    request.setConfig(requestConfigBuilder.build())
-
-    when:
-    withConfigOverride(Config.HTTP_CLIENT_HOST_SPLIT_BY_DOMAIN, "$renameService") {
-      runUnderTrace("parent") {
-        client.execute(request)
-      }
-    }
-
-    then:
-    def exception = thrown(ClientProtocolException)
-    // two traces on the server, one trace on the client
-    assertTraces(3) {
-      server.distributedRequestTrace(it, 0, TEST_WRITER[2][1])
-      server.distributedRequestTrace(it, 1, TEST_WRITER[2][1])
-      trace(2, 2) {
-        parentSpan(it, 0, exception)
-        successClientSpan(it, 1, span(0), renameService, null, "another-redirect", exception)
-      }
-    }
-
-    where:
-    renameService << [false, true]
-  }
-
-  def "trace request without propagation"() {
-    setup:
-    HttpGet request = new HttpGet(successUrl)
-    request.addHeader(new BasicHeader("is-dd-server", "false"))
-
-    when:
-    HttpResponse response = withConfigOverride(Config.HTTP_CLIENT_HOST_SPLIT_BY_DOMAIN, "$renameService") {
-      runUnderTrace("parent") {
-        client.execute(request)
-      }
-    }
-
-    then:
-    response.getStatusLine().getStatusCode() == 200
-    // only one trace (client).
-    assertTraces(1) {
-      trace(0, 2) {
-        parentSpan(it, 0)
-        successClientSpan(it, 1, span(0), renameService)
-      }
-    }
-
-    where:
-    renameService << [false, true]
-  }
-
-  def parentSpan(TraceAssert trace, int index, Throwable exception = null) {
-    trace.span(index) {
-      parent()
-      serviceName "unnamed-java-app"
-      operationName "parent"
-      resourceName "parent"
-      errored exception != null
-      tags {
-        defaultTags()
-        if (exception) {
-          errorTags(exception.class)
-        }
-      }
-    }
-  }
-
-  def successClientSpan(TraceAssert trace, int index, DDSpan parent, boolean renameService, status = 200, route = "success", Throwable exception = null) {
-    trace.span(index) {
-      childOf parent
-      serviceName renameService ? "localhost" : "unnamed-java-app"
-      operationName "http.request"
-      resourceName "GET /$route"
-      spanType DDSpanTypes.HTTP_CLIENT
-      errored exception != null
-      tags {
-        defaultTags()
-        if (exception) {
-          errorTags(exception.class)
-        }
-        "$Tags.COMPONENT.key" "apache-httpclient"
-        "$Tags.HTTP_STATUS.key" status
-        "$Tags.HTTP_URL.key" "http://localhost:$port/$route"
-        "$Tags.PEER_HOSTNAME.key" "localhost"
-        "$Tags.PEER_PORT.key" server.address.port
-        "$Tags.HTTP_METHOD.key" "GET"
-        "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
-      }
-    }
-  }
-}

dd-java-agent/instrumentation/apache-httpclient-4/src/test/groovy/ApacheHttpClientResponseHandlerTest.groovy

@@ -0,0 +1,43 @@
+import datadog.trace.agent.test.base.HttpClientTest
+import datadog.trace.instrumentation.apachehttpclient.ApacheHttpClientDecorator
+import org.apache.http.HttpResponse
+import org.apache.http.client.ResponseHandler
+import org.apache.http.client.methods.HttpGet
+import org.apache.http.impl.client.DefaultHttpClient
+import org.apache.http.message.BasicHeader
+import spock.lang.Shared
+
+class ApacheHttpClientResponseHandlerTest extends HttpClientTest<ApacheHttpClientDecorator> {
+
+  @Shared
+  def client = new DefaultHttpClient()
+
+  @Shared
+  def handler = new ResponseHandler<Integer>() {
+    @Override
+    Integer handleResponse(HttpResponse response) {
+      return response.statusLine.statusCode
+    }
+  }
+
+  @Override
+  int doRequest(String method, URI uri, Map<String, String> headers, Closure callback) {
+    assert method == "GET"
+    HttpGet request = new HttpGet(uri)
+    headers.entrySet().each {
+      request.addHeader(new BasicHeader(it.key, it.value))
+    }
+
+    def status = client.execute(request, handler)
+
+    // handler execution is included within the client span, so we can't call the callback there.
+    callback?.call()
+
+    return status
+  }
+
+  @Override
+  ApacheHttpClientDecorator decorator() {
+    return ApacheHttpClientDecorator.DECORATE
+  }
+}

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

@@ -1,148 +1,31 @@
-import datadog.opentracing.DDSpan
-import datadog.trace.agent.test.AgentTestRunner
-import datadog.trace.agent.test.asserts.TraceAssert
-import datadog.trace.api.Config
-import datadog.trace.api.DDSpanTypes
-import io.opentracing.tag.Tags
-import org.apache.http.HttpResponse
-import org.apache.http.client.HttpClient
+import datadog.trace.agent.test.base.HttpClientTest
+import datadog.trace.instrumentation.apachehttpclient.ApacheHttpClientDecorator
 import org.apache.http.client.methods.HttpGet
-import org.apache.http.impl.client.BasicResponseHandler
 import org.apache.http.impl.client.DefaultHttpClient
 import org.apache.http.message.BasicHeader
-import spock.lang.AutoCleanup
 import spock.lang.Shared
 
-import static datadog.trace.agent.test.server.http.TestHttpServer.httpServer
-import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
-import static datadog.trace.agent.test.utils.TraceUtils.withConfigOverride
+class ApacheHttpClientTest extends HttpClientTest<ApacheHttpClientDecorator> {
 
-class ApacheHttpClientTest extends AgentTestRunner {
-
-  @AutoCleanup
   @Shared
-  def server = httpServer {
-    handlers {
-      prefix("success") {
-        handleDistributedRequest()
-        String msg = "Hello."
-        response.status(200).send(msg)
-      }
-      prefix("redirect") {
-        handleDistributedRequest()
-        redirect(server.address.resolve("/success").toURL().toString())
-      }
-      prefix("another-redirect") {
-        handleDistributedRequest()
-        redirect(server.address.resolve("/redirect").toURL().toString())
-      }
-    }
-  }
-  @Shared
-  int port = server.address.port
-  @Shared
-  def successUrl = server.address.resolve("/success")
-  @Shared
-  def redirectUrl = server.address.resolve("/redirect")
-  @Shared
-  def twoRedirectsUrl = server.address.resolve("/another-redirect")
-  @Shared
-  def handler = new BasicResponseHandler()
+  def client = new DefaultHttpClient()
 
-  final HttpClient client = new DefaultHttpClient()
-
-  def "trace request with propagation"() {
-    when:
-
-    String response = withConfigOverride(Config.HTTP_CLIENT_HOST_SPLIT_BY_DOMAIN, "$renameService") {
-      runUnderTrace("parent") {
-        if (responseHandler) {
-          client.execute(new HttpGet(successUrl), responseHandler)
-        } else {
-          client.execute(new HttpGet(successUrl)).entity.content.text
-        }
-      }
+  @Override
+  int doRequest(String method, URI uri, Map<String, String> headers, Closure callback) {
+    assert method == "GET"
+    HttpGet request = new HttpGet(uri)
+    headers.entrySet().each {
+      request.addHeader(new BasicHeader(it.key, it.value))
     }
 
-    then:
-    response == "Hello."
-    // one trace on the server, one trace on the client
-    assertTraces(2) {
-      server.distributedRequestTrace(it, 0, TEST_WRITER[1][1])
-      trace(1, 2) {
-        parentSpan(it, 0)
-        successClientSpan(it, 1, span(0), renameService)
-      }
+    def response = client.execute(request)
+    callback?.call()
+    response.entity.getContent().close() // Make sure the connection is closed.
+    response.statusLine.statusCode
   }
 
-    where:
-    responseHandler << [null, handler]
-    renameService << [false, true]
-  }
-
-  def "trace request without propagation"() {
-    setup:
-    HttpGet request = new HttpGet(successUrl)
-    request.addHeader(new BasicHeader("is-dd-server", "false"))
-
-    when:
-    HttpResponse response = withConfigOverride(Config.HTTP_CLIENT_HOST_SPLIT_BY_DOMAIN, "$renameService") {
-      runUnderTrace("parent") {
-        client.execute(request)
-      }
-    }
-
-    then:
-    response.getStatusLine().getStatusCode() == 200
-    // only one trace (client).
-    assertTraces(1) {
-      trace(0, 2) {
-        parentSpan(it, 0)
-        successClientSpan(it, 1, span(0), renameService)
-      }
-    }
-
-    where:
-    renameService << [false, true]
-  }
-
-  def parentSpan(TraceAssert trace, int index, Throwable exception = null) {
-    trace.span(index) {
-      parent()
-      serviceName "unnamed-java-app"
-      operationName "parent"
-      resourceName "parent"
-      errored exception != null
-      tags {
-        defaultTags()
-        if (exception) {
-          errorTags(exception.class)
-        }
-      }
-    }
-  }
-
-  def successClientSpan(TraceAssert trace, int index, DDSpan parent, boolean renameService, status = 200, route = "success", Throwable exception = null) {
-    trace.span(index) {
-      childOf parent
-      serviceName renameService ? "localhost" : "unnamed-java-app"
-      operationName "http.request"
-      resourceName "GET /$route"
-      spanType DDSpanTypes.HTTP_CLIENT
-      errored exception != null
-      tags {
-        defaultTags()
-        if (exception) {
-          errorTags(exception.class)
-        }
-        "$Tags.COMPONENT.key" "apache-httpclient"
-        "$Tags.HTTP_STATUS.key" status
-        "$Tags.HTTP_URL.key" "http://localhost:$port/$route"
-        "$Tags.PEER_HOSTNAME.key" "localhost"
-        "$Tags.PEER_PORT.key" server.address.port
-        "$Tags.HTTP_METHOD.key" "GET"
-        "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
-      }
-    }
+  @Override
+  ApacheHttpClientDecorator decorator() {
+    return ApacheHttpClientDecorator.DECORATE
   }
 }

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpClientTest.groovy

@@ -0,0 +1,312 @@
+package datadog.trace.agent.test.base
+
+import datadog.opentracing.DDSpan
+import datadog.trace.agent.decorator.HttpClientDecorator
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.agent.test.asserts.TraceAssert
+import datadog.trace.api.Config
+import datadog.trace.api.DDSpanTypes
+import io.opentracing.tag.Tags
+import spock.lang.AutoCleanup
+import spock.lang.Shared
+
+import java.util.concurrent.ExecutionException
+
+import static datadog.trace.agent.test.server.http.TestHttpServer.httpServer
+import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
+import static datadog.trace.agent.test.utils.TraceUtils.withConfigOverride
+
+abstract class HttpClientTest<T extends HttpClientDecorator> extends AgentTestRunner {
+
+  @AutoCleanup
+  @Shared
+  def server = httpServer {
+    handlers {
+      prefix("success") {
+        handleDistributedRequest()
+        String msg = "Hello."
+        response.status(200).send(msg)
+      }
+      prefix("error") {
+        handleDistributedRequest()
+        String msg = "Sorry."
+        response.status(500).send(msg)
+      }
+      prefix("redirect") {
+        handleDistributedRequest()
+        redirect(server.address.resolve("/success").toURL().toString())
+      }
+      prefix("another-redirect") {
+        handleDistributedRequest()
+        redirect(server.address.resolve("/redirect").toURL().toString())
+      }
+      prefix("circular-redirect") {
+        handleDistributedRequest()
+        redirect(server.address.resolve("/circular-redirect").toURL().toString())
+      }
+    }
+  }
+
+  @Shared
+  T decorator = decorator()
+
+  /**
+   * Make the request and return the status code response
+   * @param method
+   * @return
+   */
+  abstract int doRequest(String method, URI uri, Map<String, String> headers = [:], Closure callback = null)
+
+  abstract T decorator()
+
+  Integer statusOnRedirectError() {
+    return null
+  }
+
+  def "basic #method request"() {
+    when:
+    def status = doRequest(method, server.address.resolve("/success"))
+
+    then:
+    status == 200
+    assertTraces(2) {
+      server.distributedRequestTrace(it, 0, trace(1).get(0))
+      trace(1, 1) {
+        clientSpan(it, 0, null, false)
+      }
+    }
+
+    where:
+    method = "GET"
+  }
+
+  def "basic #method request with parent"() {
+    when:
+    def status = runUnderTrace("parent") {
+      doRequest(method, server.address.resolve("/success"))
+    }
+
+    then:
+    status == 200
+    assertTraces(2) {
+      server.distributedRequestTrace(it, 0, trace(1).get(1))
+      trace(1, 2) {
+        parentSpan(it, 0)
+        clientSpan(it, 1, span(0), false)
+      }
+    }
+
+    where:
+    method = "GET"
+  }
+
+  def "basic #method request with split-by-domain"() {
+    when:
+    def status = withConfigOverride(Config.HTTP_CLIENT_HOST_SPLIT_BY_DOMAIN, "true") {
+      doRequest(method, server.address.resolve("/success"))
+    }
+
+    then:
+    status == 200
+    assertTraces(2) {
+      server.distributedRequestTrace(it, 0, trace(1).get(0))
+      trace(1, 1) {
+        clientSpan(it, 0, null, true)
+      }
+    }
+
+    where:
+    method = "GET"
+  }
+
+  def "trace request without propagation"() {
+    when:
+    def status = withConfigOverride(Config.HTTP_CLIENT_HOST_SPLIT_BY_DOMAIN, "$renameService") {
+      runUnderTrace("parent") {
+        doRequest(method, server.address.resolve("/success"), ["is-dd-server": "false"])
+      }
+    }
+
+    then:
+    status == 200
+    // only one trace (client).
+    assertTraces(1) {
+      trace(0, 2) {
+        parentSpan(it, 0)
+        clientSpan(it, 1, span(0), renameService)
+      }
+    }
+
+    where:
+    method = "GET"
+    renameService << [false, true]
+  }
+
+  def "trace request with callback and parent"() {
+    when:
+    def status = runUnderTrace("parent") {
+      doRequest(method, server.address.resolve("/success"), ["is-dd-server": "false"]) {
+        runUnderTrace("child") {}
+      }
+    }
+
+    then:
+    status == 200
+    // only one trace (client).
+    assertTraces(1) {
+      trace(0, 3) {
+        parentSpan(it, 0)
+        span(1) {
+          operationName "child"
+          childOf span(0)
+        }
+        clientSpan(it, 2, span(0), false)
+      }
+    }
+
+    where:
+    method = "GET"
+  }
+
+  def "trace request with callback and no parent"() {
+    when:
+    def status = doRequest(method, server.address.resolve("/success"), ["is-dd-server": "false"]) {
+      runUnderTrace("child") {
+        // Ensure consistent ordering of traces for assertion.
+        TEST_WRITER.waitForTraces(1)
+      }
+    }
+
+    then:
+    status == 200
+    // only one trace (client).
+    assertTraces(2) {
+      trace(0, 1) {
+        clientSpan(it, 0, null, false)
+      }
+      trace(1, 1) {
+        span(0) {
+          operationName "child"
+          parent()
+        }
+      }
+    }
+
+    where:
+    method = "GET"
+  }
+
+  def "basic #method request with 1 redirect"() {
+    setup:
+    def uri = server.address.resolve("/redirect")
+
+    when:
+    def status = doRequest(method, uri)
+
+    then:
+    status == 200
+    assertTraces(3) {
+      server.distributedRequestTrace(it, 0, trace(2).get(0))
+      server.distributedRequestTrace(it, 1, trace(2).get(0))
+      trace(2, 1) {
+        clientSpan(it, 0, null, false, uri)
+      }
+    }
+
+    where:
+    method = "GET"
+  }
+
+  def "basic #method request with 2 redirects"() {
+    setup:
+    def uri = server.address.resolve("/another-redirect")
+
+    when:
+    def status = doRequest(method, uri)
+
+    then:
+    status == 200
+    assertTraces(4) {
+      server.distributedRequestTrace(it, 0, trace(3).get(0))
+      server.distributedRequestTrace(it, 1, trace(3).get(0))
+      server.distributedRequestTrace(it, 2, trace(3).get(0))
+      trace(3, 1) {
+        clientSpan(it, 0, null, false, uri)
+      }
+    }
+
+    where:
+    method = "GET"
+  }
+
+  def "basic #method request with circular redirects"() {
+    setup:
+    def uri = server.address.resolve("/circular-redirect")
+
+    when:
+    doRequest(method, uri)//, ["is-dd-server": "false"])
+
+    then:
+    def ex = thrown(Exception)
+    def thrownException = ex instanceof ExecutionException ? ex.cause : ex
+
+    and:
+    assertTraces(3) {
+      server.distributedRequestTrace(it, 0, trace(2).get(0))
+      server.distributedRequestTrace(it, 1, trace(2).get(0))
+      trace(2, 1) {
+        clientSpan(it, 0, null, false, uri, statusOnRedirectError(), thrownException)
+      }
+    }
+
+    where:
+    method = "GET"
+  }
+
+  void parentSpan(TraceAssert trace, int index, Throwable exception = null) {
+    trace.span(index) {
+      parent()
+      serviceName "unnamed-java-app"
+      operationName "parent"
+      resourceName "parent"
+      errored exception != null
+      tags {
+        defaultTags()
+        if (exception) {
+          errorTags(exception.class)
+        }
+      }
+    }
+  }
+
+  // parent span must be cast otherwise it breaks debugging classloading (junit loads it early)
+  void clientSpan(TraceAssert trace, int index, Object parentSpan, boolean renameService, URI uri = server.address.resolve("/success"), Integer status = 200, Throwable exception = null) {
+    trace.span(index) {
+      if (parentSpan == null) {
+        parent()
+      } else {
+        childOf((DDSpan) parentSpan)
+      }
+      serviceName renameService ? "localhost" : "unnamed-java-app"
+      operationName "http.request"
+      resourceName "GET $uri.path"
+      spanType DDSpanTypes.HTTP_CLIENT
+      errored exception != null
+      tags {
+        defaultTags()
+        if (exception) {
+          errorTags(exception.class, exception.message)
+        }
+        "$Tags.COMPONENT.key" decorator.component()
+        if (status) {
+          "$Tags.HTTP_STATUS.key" status
+        }
+        "$Tags.HTTP_URL.key" "$uri"
+        "$Tags.PEER_HOSTNAME.key" "localhost"
+        "$Tags.PEER_PORT.key" server.address.port
+        "$Tags.HTTP_METHOD.key" "GET"
+        "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
+      }
+    }
+  }
+}

dd-java-agent/testing/testing.gradle

@@ -4,6 +4,7 @@ minimumBranchCoverage = 0.5
 minimumInstructionCoverage = 0.5
 excludedClassesConverage += [
   'datadog.trace.agent.test.asserts.*Assert',
+  'datadog.trace.agent.test.base.*',
   'datadog.trace.agent.test.AgentTestRunner.ErrorCountingListener',
   'datadog.trace.agent.test.utils.*',
   // Avoid applying jacoco instrumentation to classes instrumented by tested agent

settings.gradle

@@ -25,6 +25,7 @@ include ':dd-smoke-tests:wildfly'
 
 // instrumentation:
 include ':dd-java-agent:instrumentation:akka-http-10.0'
+include ':dd-java-agent:instrumentation:apache-httpasyncclient-4'
 include ':dd-java-agent:instrumentation:apache-httpclient-4'
 include ':dd-java-agent:instrumentation:aws-java-sdk-1.11.0'
 include ':dd-java-agent:instrumentation:aws-java-sdk-2.2'