فهرست منبع

Merge pull request #1075 from DataDog/landerson/wsclient

Play WS Client instrumentation
Laplie Anderson 5 سال پیش
والد
کامیت
a3cd46708b
13فایلهای تغییر یافته به همراه817 افزوده شده و 0 حذف شده
  1. 51 0
      dd-java-agent/instrumentation/play-ws-1/play-ws-1.gradle
  2. 80 0
      dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws/AsyncHandlerWrapper.java
  3. 14 0
      dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws/HeadersInjectAdapter.java
  4. 46 0
      dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws/PlayWSClientDecorator.java
  5. 91 0
      dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws/PlayWSClientInstrumentation.java
  6. 74 0
      dd-java-agent/instrumentation/play-ws-1/src/test/groovy/PlayWSClientTest.groovy
  7. 75 0
      dd-java-agent/instrumentation/play-ws-2/play-ws-2.gradle
  8. 159 0
      dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws/AsyncHandlerWrapper.java
  9. 14 0
      dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws/HeadersInjectAdapter.java
  10. 46 0
      dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws/PlayWSClientDecorator.java
  11. 91 0
      dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws/PlayWSClientInstrumentation.java
  12. 74 0
      dd-java-agent/instrumentation/play-ws-2/src/test/groovy/PlayWSClientTest.groovy
  13. 2 0
      settings.gradle

+ 51 - 0
dd-java-agent/instrumentation/play-ws-1/play-ws-1.gradle

@@ -0,0 +1,51 @@
+// Set properties before any plugins get loaded
+ext {
+  minJavaVersionForTests = JavaVersion.VERSION_1_8
+}
+
+apply from: "${rootDir}/gradle/java.gradle"
+
+apply plugin: 'org.unbroken-dome.test-sets'
+
+testSets {
+  latestDepTest {
+    dirName = 'test'
+  }
+}
+
+muzzle {
+  pass {
+    group = 'com.typesafe.play'
+    module = 'play-ahc-ws-standalone_2.11'
+    versions = '[1.0.0,2.0.0)'
+    assertInverse = true
+  }
+  pass {
+    group = 'com.typesafe.play'
+    module = 'play-ahc-ws-standalone_2.12'
+    versions = '[1.0.0,2.0.0)'
+    assertInverse = true
+  }
+  fail {
+    group = 'com.typesafe.play'
+    module = 'play-ahc-ws-standalone_2.13'
+    versions = '[,]'
+  }
+}
+
+def scalaVersion = '2.12'
+
+dependencies {
+  compileOnly group: 'com.typesafe.play', name: "play-ahc-ws-standalone_$scalaVersion", version: '1.0.2'
+
+  testCompile project(':dd-java-agent:instrumentation:java-concurrent')
+
+  // These are to ensure cross compatibility
+  testCompile project(':dd-java-agent:instrumentation:netty-4.0')
+  testCompile project(':dd-java-agent:instrumentation:netty-4.1')
+  testCompile project(':dd-java-agent:instrumentation:akka-http-10.0')
+
+  testCompile group: 'com.typesafe.play', name: "play-ahc-ws-standalone_$scalaVersion", version: '1.0.2'
+
+  latestDepTestCompile group: 'com.typesafe.play', name: "play-ahc-ws-standalone_$scalaVersion", version: '1.+'
+}

+ 80 - 0
dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws/AsyncHandlerWrapper.java

@@ -0,0 +1,80 @@
+package datadog.trace.instrumentation.playws;
+
+import static datadog.trace.instrumentation.api.AgentTracer.propagate;
+import static datadog.trace.instrumentation.playws.PlayWSClientDecorator.DECORATE;
+
+import datadog.trace.context.TraceScope;
+import datadog.trace.instrumentation.api.AgentSpan;
+import play.shaded.ahc.org.asynchttpclient.AsyncHandler;
+import play.shaded.ahc.org.asynchttpclient.HttpResponseBodyPart;
+import play.shaded.ahc.org.asynchttpclient.HttpResponseHeaders;
+import play.shaded.ahc.org.asynchttpclient.HttpResponseStatus;
+import play.shaded.ahc.org.asynchttpclient.Response;
+
+public class AsyncHandlerWrapper implements AsyncHandler {
+  private final AsyncHandler delegate;
+  private final AgentSpan span;
+  private final TraceScope.Continuation continuation;
+
+  private final Response.ResponseBuilder builder = new Response.ResponseBuilder();
+
+  public AsyncHandlerWrapper(final AsyncHandler delegate, final AgentSpan span) {
+    this.delegate = delegate;
+    this.span = span;
+    continuation = propagate().capture();
+  }
+
+  @Override
+  public State onBodyPartReceived(final HttpResponseBodyPart content) throws Exception {
+    builder.accumulate(content);
+    return delegate.onBodyPartReceived(content);
+  }
+
+  @Override
+  public State onStatusReceived(final HttpResponseStatus status) throws Exception {
+    builder.reset();
+    builder.accumulate(status);
+    return delegate.onStatusReceived(status);
+  }
+
+  @Override
+  public State onHeadersReceived(final HttpResponseHeaders httpHeaders) throws Exception {
+    builder.accumulate(httpHeaders);
+    return delegate.onHeadersReceived(httpHeaders);
+  }
+
+  @Override
+  public Object onCompleted() throws Exception {
+    final Response response = builder.build();
+    if (response != null) {
+      DECORATE.onResponse(span, response);
+    }
+    DECORATE.beforeFinish(span);
+    span.finish();
+
+    if (continuation != null) {
+      try (final TraceScope scope = continuation.activate()) {
+        scope.setAsyncPropagation(true);
+        return delegate.onCompleted();
+      }
+    } else {
+      return delegate.onCompleted();
+    }
+  }
+
+  @Override
+  public void onThrowable(final Throwable throwable) {
+    DECORATE.onError(span, throwable);
+    DECORATE.beforeFinish(span);
+    span.finish();
+
+    if (continuation != null) {
+      try (final TraceScope scope = continuation.activate()) {
+        scope.setAsyncPropagation(true);
+        delegate.onThrowable(throwable);
+      }
+    } else {
+      delegate.onThrowable(throwable);
+    }
+  }
+}

+ 14 - 0
dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws/HeadersInjectAdapter.java

@@ -0,0 +1,14 @@
+package datadog.trace.instrumentation.playws;
+
+import datadog.trace.instrumentation.api.AgentPropagation;
+import play.shaded.ahc.org.asynchttpclient.Request;
+
+public class HeadersInjectAdapter implements AgentPropagation.Setter<Request> {
+
+  public static final HeadersInjectAdapter SETTER = new HeadersInjectAdapter();
+
+  @Override
+  public void set(final Request carrier, final String key, final String value) {
+    carrier.getHeaders().add(key, value);
+  }
+}

+ 46 - 0
dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws/PlayWSClientDecorator.java

@@ -0,0 +1,46 @@
+package datadog.trace.instrumentation.playws;
+
+import datadog.trace.agent.decorator.HttpClientDecorator;
+import java.net.URI;
+import java.net.URISyntaxException;
+import play.shaded.ahc.org.asynchttpclient.Request;
+import play.shaded.ahc.org.asynchttpclient.Response;
+
+public class PlayWSClientDecorator extends HttpClientDecorator<Request, Response> {
+  public static final PlayWSClientDecorator DECORATE = new PlayWSClientDecorator();
+
+  @Override
+  protected String method(final Request request) {
+    return request.getMethod();
+  }
+
+  @Override
+  protected URI url(final Request request) throws URISyntaxException {
+    return request.getUri().toJavaNetURI();
+  }
+
+  @Override
+  protected String hostname(final Request request) {
+    return request.getUri().getHost();
+  }
+
+  @Override
+  protected Integer port(final Request request) {
+    return request.getUri().getPort();
+  }
+
+  @Override
+  protected Integer status(final Response response) {
+    return response.getStatusCode();
+  }
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"play-ws"};
+  }
+
+  @Override
+  protected String component() {
+    return "play-ws";
+  }
+}

+ 91 - 0
dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws/PlayWSClientInstrumentation.java

@@ -0,0 +1,91 @@
+package datadog.trace.instrumentation.playws;
+
+import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
+import static datadog.trace.instrumentation.api.AgentTracer.propagate;
+import static datadog.trace.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.playws.HeadersInjectAdapter.SETTER;
+import static datadog.trace.instrumentation.playws.PlayWSClientDecorator.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.not;
+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.instrumentation.api.AgentSpan;
+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 play.shaded.ahc.org.asynchttpclient.AsyncHandler;
+import play.shaded.ahc.org.asynchttpclient.Request;
+
+@AutoService(Instrumenter.class)
+public class PlayWSClientInstrumentation extends Instrumenter.Default {
+  public PlayWSClientInstrumentation() {
+    super("play-ws");
+  }
+
+  @Override
+  public ElementMatcher<? super TypeDescription> typeMatcher() {
+    // CachingAsyncHttpClient rejects overrides to AsyncHandler
+    // It also delegates to another AsyncHttpClient
+    return safeHasSuperType(named("play.shaded.ahc.org.asynchttpclient.AsyncHttpClient"))
+        .and(not(named("play.api.libs.ws.ahc.cache.CachingAsyncHttpClient")));
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isMethod()
+            .and(named("execute"))
+            .and(takesArguments(2))
+            .and(takesArgument(0, named("play.shaded.ahc.org.asynchttpclient.Request")))
+            .and(takesArgument(1, named("play.shaded.ahc.org.asynchttpclient.AsyncHandler"))),
+        ClientAdvice.class.getName());
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      "datadog.trace.agent.decorator.BaseDecorator",
+      "datadog.trace.agent.decorator.ClientDecorator",
+      "datadog.trace.agent.decorator.HttpClientDecorator",
+      packageName + ".PlayWSClientDecorator",
+      packageName + ".HeadersInjectAdapter",
+      packageName + ".AsyncHandlerWrapper"
+    };
+  }
+
+  public static class ClientAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentSpan methodEnter(
+        @Advice.Argument(0) final Request request,
+        @Advice.Argument(value = 1, readOnly = false) AsyncHandler asyncHandler) {
+
+      final AgentSpan span = startSpan("play-ws.request");
+
+      DECORATE.afterStart(span);
+      DECORATE.onRequest(span, request);
+      propagate().inject(span, request, SETTER);
+
+      asyncHandler = new AsyncHandlerWrapper(asyncHandler, span);
+
+      return span;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void methodExit(
+        @Advice.Enter final AgentSpan clientSpan, @Advice.Thrown final Throwable throwable) {
+
+      if (throwable != null) {
+        DECORATE.onError(clientSpan, throwable);
+        DECORATE.beforeFinish(clientSpan);
+        clientSpan.finish();
+      }
+    }
+  }
+}

+ 74 - 0
dd-java-agent/instrumentation/play-ws-1/src/test/groovy/PlayWSClientTest.groovy

@@ -0,0 +1,74 @@
+import akka.actor.ActorSystem
+import akka.stream.ActorMaterializer
+import akka.stream.ActorMaterializerSettings
+import datadog.trace.agent.test.base.HttpClientTest
+import datadog.trace.instrumentation.playws.PlayWSClientDecorator
+import play.libs.ws.StandaloneWSClient
+import play.libs.ws.StandaloneWSRequest
+import play.libs.ws.StandaloneWSResponse
+import play.libs.ws.ahc.StandaloneAhcWSClient
+import play.shaded.ahc.org.asynchttpclient.AsyncHttpClient
+import play.shaded.ahc.org.asynchttpclient.AsyncHttpClientConfig
+import play.shaded.ahc.org.asynchttpclient.DefaultAsyncHttpClient
+import play.shaded.ahc.org.asynchttpclient.DefaultAsyncHttpClientConfig
+import spock.lang.Shared
+
+import java.util.concurrent.TimeUnit
+
+class PlayWSClientTest extends HttpClientTest<PlayWSClientDecorator> {
+  @Shared
+  ActorSystem system
+
+  @Shared
+  StandaloneWSClient wsClient
+
+  def setupSpec() {
+    String name = "play-ws"
+    system = ActorSystem.create(name)
+    ActorMaterializerSettings settings = ActorMaterializerSettings.create(system)
+    ActorMaterializer materializer = ActorMaterializer.create(settings, system, name)
+
+    AsyncHttpClientConfig asyncHttpClientConfig =
+      new DefaultAsyncHttpClientConfig.Builder()
+        .setMaxRequestRetry(0)
+        .setShutdownQuietPeriod(0)
+        .setShutdownTimeout(0)
+        .build()
+
+    AsyncHttpClient asyncHttpClient = new DefaultAsyncHttpClient(asyncHttpClientConfig)
+
+    wsClient = new StandaloneAhcWSClient(asyncHttpClient, materializer)
+  }
+
+  def cleanupSpec() {
+    wsClient?.close()
+    system?.terminate()
+  }
+
+  @Override
+  int doRequest(String method, URI uri, Map<String, String> headers, Closure callback) {
+    StandaloneWSRequest wsRequest = wsClient.url(uri.toURL().toString()).setFollowRedirects(true)
+
+    headers.entrySet().each { entry -> wsRequest.addHeader(entry.getKey(), entry.getValue()) }
+    StandaloneWSResponse wsResponse = wsRequest.execute(method)
+      .whenComplete({ response, throwable ->
+        callback?.call()
+      }).toCompletableFuture().get(5, TimeUnit.SECONDS)
+
+    return wsResponse.getStatus()
+  }
+
+  @Override
+  PlayWSClientDecorator decorator() {
+    return PlayWSClientDecorator.DECORATE
+  }
+
+  String expectedOperationName() {
+    return "play-ws.request"
+  }
+
+  @Override
+  boolean testCircularRedirects() {
+    return false
+  }
+}

+ 75 - 0
dd-java-agent/instrumentation/play-ws-2/play-ws-2.gradle

@@ -0,0 +1,75 @@
+// Set properties before any plugins get loaded
+ext {
+  minJavaVersionForTests = JavaVersion.VERSION_1_8
+}
+
+apply from: "${rootDir}/gradle/java.gradle"
+
+apply plugin: 'org.unbroken-dome.test-sets'
+
+testSets {
+  latestDepTest {
+    dirName = 'test'
+  }
+}
+
+muzzle {
+  // 2.0.5 was a bad release
+
+  fail {
+    group = 'com.typesafe.play'
+    module = 'play-ahc-ws-standalone_2.11'
+    versions = '[,2.0.0)'
+  }
+  pass {
+    group = 'com.typesafe.play'
+    module = 'play-ahc-ws-standalone_2.11'
+    versions = '[2.0.0,2.0.4]'
+  }
+  pass {
+    group = 'com.typesafe.play'
+    module = 'play-ahc-ws-standalone_2.11'
+    versions = '[2.0.6,]'
+  }
+
+  fail {
+    group = 'com.typesafe.play'
+    module = 'play-ahc-ws-standalone_2.12'
+    versions = '[,2.0.0)'
+  }
+  pass {
+    group = 'com.typesafe.play'
+    module = 'play-ahc-ws-standalone_2.12'
+    versions = '[2.0.0,2.0.4]'
+  }
+  pass {
+    group = 'com.typesafe.play'
+    module = 'play-ahc-ws-standalone_2.12'
+    versions = '[2.0.6,]'
+  }
+
+  // No Scala 2.13 versions below 2.0.6 exist
+  pass {
+    group = 'com.typesafe.play'
+    module = 'play-ahc-ws-standalone_2.13'
+    versions = '[2.0.6,]'
+  }
+}
+
+def scalaVersion = '2.12'
+
+dependencies {
+  compileOnly group: 'com.typesafe.play', name: "play-ahc-ws-standalone_$scalaVersion", version: '2.0.0'
+
+  testCompile project(':dd-java-agent:instrumentation:java-concurrent')
+
+  // These are to ensure cross compatibility
+  testCompile project(':dd-java-agent:instrumentation:netty-4.0')
+  testCompile project(':dd-java-agent:instrumentation:netty-4.1')
+  testCompile project(':dd-java-agent:instrumentation:akka-http-10.0')
+
+  testCompile group: 'com.typesafe.play', name: "play-ahc-ws-standalone_$scalaVersion", version: '2.0.0'
+
+  // TODO: Revisit when 2.1.x are out of the release candidate stage
+  latestDepTestCompile group: 'com.typesafe.play', name: "play-ahc-ws-standalone_$scalaVersion", version: '2.0.+'
+}

+ 159 - 0
dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws/AsyncHandlerWrapper.java

@@ -0,0 +1,159 @@
+package datadog.trace.instrumentation.playws;
+
+import static datadog.trace.instrumentation.api.AgentTracer.propagate;
+import static datadog.trace.instrumentation.playws.PlayWSClientDecorator.DECORATE;
+
+import datadog.trace.context.TraceScope;
+import datadog.trace.instrumentation.api.AgentSpan;
+import java.net.InetSocketAddress;
+import java.util.List;
+import play.shaded.ahc.io.netty.channel.Channel;
+import play.shaded.ahc.io.netty.handler.codec.http.HttpHeaders;
+import play.shaded.ahc.org.asynchttpclient.AsyncHandler;
+import play.shaded.ahc.org.asynchttpclient.HttpResponseBodyPart;
+import play.shaded.ahc.org.asynchttpclient.HttpResponseStatus;
+import play.shaded.ahc.org.asynchttpclient.Response;
+import play.shaded.ahc.org.asynchttpclient.netty.request.NettyRequest;
+
+public class AsyncHandlerWrapper implements AsyncHandler {
+  private final AsyncHandler delegate;
+  private final AgentSpan span;
+  private final TraceScope.Continuation continuation;
+
+  private final Response.ResponseBuilder builder = new Response.ResponseBuilder();
+
+  public AsyncHandlerWrapper(final AsyncHandler delegate, final AgentSpan span) {
+    this.delegate = delegate;
+    this.span = span;
+    continuation = propagate().capture();
+  }
+
+  @Override
+  public State onBodyPartReceived(final HttpResponseBodyPart content) throws Exception {
+    builder.accumulate(content);
+    return delegate.onBodyPartReceived(content);
+  }
+
+  @Override
+  public State onStatusReceived(final HttpResponseStatus status) throws Exception {
+    builder.reset();
+    builder.accumulate(status);
+    return delegate.onStatusReceived(status);
+  }
+
+  @Override
+  public State onHeadersReceived(final HttpHeaders httpHeaders) throws Exception {
+    builder.accumulate(httpHeaders);
+    return delegate.onHeadersReceived(httpHeaders);
+  }
+
+  @Override
+  public Object onCompleted() throws Exception {
+    final Response response = builder.build();
+    if (response != null) {
+      DECORATE.onResponse(span, response);
+    }
+    DECORATE.beforeFinish(span);
+    span.finish();
+
+    if (continuation != null) {
+      try (final TraceScope scope = continuation.activate()) {
+        scope.setAsyncPropagation(true);
+        return delegate.onCompleted();
+      }
+    } else {
+      return delegate.onCompleted();
+    }
+  }
+
+  @Override
+  public void onThrowable(final Throwable throwable) {
+    DECORATE.onError(span, throwable);
+    DECORATE.beforeFinish(span);
+    span.finish();
+
+    if (continuation != null) {
+      try (final TraceScope scope = continuation.activate()) {
+        scope.setAsyncPropagation(true);
+        delegate.onThrowable(throwable);
+      }
+    } else {
+      delegate.onThrowable(throwable);
+    }
+  }
+
+  @Override
+  public State onTrailingHeadersReceived(final HttpHeaders headers) throws Exception {
+    return delegate.onTrailingHeadersReceived(headers);
+  }
+
+  @Override
+  public void onHostnameResolutionAttempt(final String name) {
+    delegate.onHostnameResolutionAttempt(name);
+  }
+
+  @Override
+  public void onHostnameResolutionSuccess(final String name, final List list) {
+    delegate.onHostnameResolutionSuccess(name, list);
+  }
+
+  @Override
+  public void onHostnameResolutionFailure(final String name, final Throwable cause) {
+    delegate.onHostnameResolutionFailure(name, cause);
+  }
+
+  @Override
+  public void onTcpConnectAttempt(final InetSocketAddress remoteAddress) {
+    delegate.onTcpConnectAttempt(remoteAddress);
+  }
+
+  @Override
+  public void onTcpConnectSuccess(final InetSocketAddress remoteAddress, final Channel connection) {
+    delegate.onTcpConnectSuccess(remoteAddress, connection);
+  }
+
+  @Override
+  public void onTcpConnectFailure(final InetSocketAddress remoteAddress, final Throwable cause) {
+    delegate.onTcpConnectFailure(remoteAddress, cause);
+  }
+
+  @Override
+  public void onTlsHandshakeAttempt() {
+    delegate.onTlsHandshakeAttempt();
+  }
+
+  @Override
+  public void onTlsHandshakeSuccess() {
+    delegate.onTlsHandshakeSuccess();
+  }
+
+  @Override
+  public void onTlsHandshakeFailure(final Throwable cause) {
+    delegate.onTlsHandshakeFailure(cause);
+  }
+
+  @Override
+  public void onConnectionPoolAttempt() {
+    delegate.onConnectionPoolAttempt();
+  }
+
+  @Override
+  public void onConnectionPooled(final Channel connection) {
+    delegate.onConnectionPooled(connection);
+  }
+
+  @Override
+  public void onConnectionOffer(final Channel connection) {
+    delegate.onConnectionOffer(connection);
+  }
+
+  @Override
+  public void onRequestSend(final NettyRequest request) {
+    delegate.onRequestSend(request);
+  }
+
+  @Override
+  public void onRetry() {
+    delegate.onRetry();
+  }
+}

+ 14 - 0
dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws/HeadersInjectAdapter.java

@@ -0,0 +1,14 @@
+package datadog.trace.instrumentation.playws;
+
+import datadog.trace.instrumentation.api.AgentPropagation;
+import play.shaded.ahc.org.asynchttpclient.Request;
+
+public class HeadersInjectAdapter implements AgentPropagation.Setter<Request> {
+
+  public static final HeadersInjectAdapter SETTER = new HeadersInjectAdapter();
+
+  @Override
+  public void set(final Request carrier, final String key, final String value) {
+    carrier.getHeaders().add(key, value);
+  }
+}

+ 46 - 0
dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws/PlayWSClientDecorator.java

@@ -0,0 +1,46 @@
+package datadog.trace.instrumentation.playws;
+
+import datadog.trace.agent.decorator.HttpClientDecorator;
+import java.net.URI;
+import java.net.URISyntaxException;
+import play.shaded.ahc.org.asynchttpclient.Request;
+import play.shaded.ahc.org.asynchttpclient.Response;
+
+public class PlayWSClientDecorator extends HttpClientDecorator<Request, Response> {
+  public static final PlayWSClientDecorator DECORATE = new PlayWSClientDecorator();
+
+  @Override
+  protected String method(final Request request) {
+    return request.getMethod();
+  }
+
+  @Override
+  protected URI url(final Request request) throws URISyntaxException {
+    return request.getUri().toJavaNetURI();
+  }
+
+  @Override
+  protected String hostname(final Request request) {
+    return request.getUri().getHost();
+  }
+
+  @Override
+  protected Integer port(final Request request) {
+    return request.getUri().getPort();
+  }
+
+  @Override
+  protected Integer status(final Response response) {
+    return response.getStatusCode();
+  }
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"play-ws"};
+  }
+
+  @Override
+  protected String component() {
+    return "play-ws";
+  }
+}

+ 91 - 0
dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws/PlayWSClientInstrumentation.java

@@ -0,0 +1,91 @@
+package datadog.trace.instrumentation.playws;
+
+import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
+import static datadog.trace.instrumentation.api.AgentTracer.propagate;
+import static datadog.trace.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.playws.HeadersInjectAdapter.SETTER;
+import static datadog.trace.instrumentation.playws.PlayWSClientDecorator.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.not;
+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.instrumentation.api.AgentSpan;
+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 play.shaded.ahc.org.asynchttpclient.AsyncHandler;
+import play.shaded.ahc.org.asynchttpclient.Request;
+
+@AutoService(Instrumenter.class)
+public class PlayWSClientInstrumentation extends Instrumenter.Default {
+  public PlayWSClientInstrumentation() {
+    super("play-ws");
+  }
+
+  @Override
+  public ElementMatcher<? super TypeDescription> typeMatcher() {
+    // CachingAsyncHttpClient rejects overrides to AsyncHandler
+    // It also delegates to another AsyncHttpClient
+    return safeHasSuperType(named("play.shaded.ahc.org.asynchttpclient.AsyncHttpClient"))
+        .and(not(named("play.api.libs.ws.ahc.cache.CachingAsyncHttpClient")));
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isMethod()
+            .and(named("execute"))
+            .and(takesArguments(2))
+            .and(takesArgument(0, named("play.shaded.ahc.org.asynchttpclient.Request")))
+            .and(takesArgument(1, named("play.shaded.ahc.org.asynchttpclient.AsyncHandler"))),
+        ClientAdvice.class.getName());
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      "datadog.trace.agent.decorator.BaseDecorator",
+      "datadog.trace.agent.decorator.ClientDecorator",
+      "datadog.trace.agent.decorator.HttpClientDecorator",
+      packageName + ".PlayWSClientDecorator",
+      packageName + ".HeadersInjectAdapter",
+      packageName + ".AsyncHandlerWrapper"
+    };
+  }
+
+  public static class ClientAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentSpan methodEnter(
+        @Advice.Argument(0) final Request request,
+        @Advice.Argument(value = 1, readOnly = false) AsyncHandler asyncHandler) {
+
+      final AgentSpan span = startSpan("play-ws.request");
+
+      DECORATE.afterStart(span);
+      DECORATE.onRequest(span, request);
+      propagate().inject(span, request, SETTER);
+
+      asyncHandler = new AsyncHandlerWrapper(asyncHandler, span);
+
+      return span;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void methodExit(
+        @Advice.Enter final AgentSpan clientSpan, @Advice.Thrown final Throwable throwable) {
+
+      if (throwable != null) {
+        DECORATE.onError(clientSpan, throwable);
+        DECORATE.beforeFinish(clientSpan);
+        clientSpan.finish();
+      }
+    }
+  }
+}

+ 74 - 0
dd-java-agent/instrumentation/play-ws-2/src/test/groovy/PlayWSClientTest.groovy

@@ -0,0 +1,74 @@
+import akka.actor.ActorSystem
+import akka.stream.ActorMaterializer
+import akka.stream.ActorMaterializerSettings
+import datadog.trace.agent.test.base.HttpClientTest
+import datadog.trace.instrumentation.playws.PlayWSClientDecorator
+import play.libs.ws.StandaloneWSClient
+import play.libs.ws.StandaloneWSRequest
+import play.libs.ws.StandaloneWSResponse
+import play.libs.ws.ahc.StandaloneAhcWSClient
+import play.shaded.ahc.org.asynchttpclient.AsyncHttpClient
+import play.shaded.ahc.org.asynchttpclient.AsyncHttpClientConfig
+import play.shaded.ahc.org.asynchttpclient.DefaultAsyncHttpClient
+import play.shaded.ahc.org.asynchttpclient.DefaultAsyncHttpClientConfig
+import spock.lang.Shared
+
+import java.util.concurrent.TimeUnit
+
+class PlayWSClientTest extends HttpClientTest<PlayWSClientDecorator> {
+  @Shared
+  ActorSystem system
+
+  @Shared
+  StandaloneWSClient wsClient
+
+  def setupSpec() {
+    String name = "play-ws"
+    system = ActorSystem.create(name)
+    ActorMaterializerSettings settings = ActorMaterializerSettings.create(system)
+    ActorMaterializer materializer = ActorMaterializer.create(settings, system, name)
+
+    AsyncHttpClientConfig asyncHttpClientConfig =
+      new DefaultAsyncHttpClientConfig.Builder()
+        .setMaxRequestRetry(0)
+        .setShutdownQuietPeriod(0)
+        .setShutdownTimeout(0)
+        .build()
+
+    AsyncHttpClient asyncHttpClient = new DefaultAsyncHttpClient(asyncHttpClientConfig)
+
+    wsClient = new StandaloneAhcWSClient(asyncHttpClient, materializer)
+  }
+
+  def cleanupSpec() {
+    wsClient?.close()
+    system?.terminate()
+  }
+
+  @Override
+  int doRequest(String method, URI uri, Map<String, String> headers, Closure callback) {
+    StandaloneWSRequest wsRequest = wsClient.url(uri.toURL().toString()).setFollowRedirects(true)
+
+    headers.entrySet().each { entry -> wsRequest.addHeader(entry.getKey(), entry.getValue()) }
+    StandaloneWSResponse wsResponse = wsRequest.execute(method)
+      .whenComplete({ response, throwable ->
+        callback?.call()
+      }).toCompletableFuture().get(5, TimeUnit.SECONDS)
+
+    return wsResponse.getStatus()
+  }
+
+  @Override
+  PlayWSClientDecorator decorator() {
+    return PlayWSClientDecorator.DECORATE
+  }
+
+  String expectedOperationName() {
+    return "play-ws.request"
+  }
+
+  @Override
+  boolean testCircularRedirects() {
+    return false
+  }
+}

+ 2 - 0
settings.gradle

@@ -82,6 +82,8 @@ include ':dd-java-agent:instrumentation:okhttp-3'
 include ':dd-java-agent:instrumentation:osgi-classloading'
 include ':dd-java-agent:instrumentation:play-2.4'
 include ':dd-java-agent:instrumentation:play-2.6'
+include ':dd-java-agent:instrumentation:play-ws-1'
+include ':dd-java-agent:instrumentation:play-ws-2'
 include ':dd-java-agent:instrumentation:rabbitmq-amqp-2.7'
 include ':dd-java-agent:instrumentation:ratpack-1.4'
 include ':dd-java-agent:instrumentation:rxjava-1'