Fix Play-ws stream() bug

dd-java-agent/instrumentation/play-ws-1/play-ws-1.gradle

@@ -38,6 +38,9 @@ def scalaVersion = '2.12'
 dependencies {
   compileOnly group: 'com.typesafe.play', name: "play-ahc-ws-standalone_$scalaVersion", version: '1.0.2'
 
+  compile project(':dd-java-agent:instrumentation:play-ws')
+  testCompile project(path: ':dd-java-agent:instrumentation:play-ws', configuration: 'testArtifacts')
+
   testCompile project(':dd-java-agent:instrumentation:java-concurrent')
 
   // These are to ensure cross compatibility

dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws1/AsyncHandlerWrapper.java

@@ -1,7 +1,7 @@
 package datadog.trace.instrumentation.playws1;
 
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
-import static datadog.trace.instrumentation.playws1.PlayWSClientDecorator.DECORATE;
+import static datadog.trace.instrumentation.playws.PlayWSClientDecorator.DECORATE;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.context.TraceScope;

dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws1/PlayWSClientInstrumentation.java

@@ -1,76 +1,21 @@
 package datadog.trace.instrumentation.playws1;
 
-import static datadog.trace.agent.tooling.ClassLoaderMatcher.classLoaderHasNoResources;
-import static datadog.trace.agent.tooling.bytebuddy.matcher.DDElementMatchers.hasInterface;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
-import static datadog.trace.instrumentation.playws1.HeadersInjectAdapter.SETTER;
-import static datadog.trace.instrumentation.playws1.PlayWSClientDecorator.DECORATE;
-import static java.util.Collections.singletonMap;
-import static net.bytebuddy.matcher.ElementMatchers.isMethod;
-import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
-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 static datadog.trace.instrumentation.playws.HeadersInjectAdapter.SETTER;
+import static datadog.trace.instrumentation.playws.PlayWSClientDecorator.DECORATE;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-import java.util.Map;
+import datadog.trace.instrumentation.playws.BasePlayWSClientInstrumentation;
 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;
+import play.shaded.ahc.org.asynchttpclient.handler.StreamedAsyncHandler;
 
 @AutoService(Instrumenter.class)
-public class PlayWSClientInstrumentation extends Instrumenter.Default {
-  public PlayWSClientInstrumentation() {
-    super("play-ws");
-  }
-
-  @Override
-  public ElementMatcher<ClassLoader> classLoaderMatcher() {
-    // Optimization for expensive typeMatcher.
-    return not(
-        classLoaderHasNoResources("play/shaded/ahc/org/asynchttpclient/AsyncHttpClient.class"));
-  }
-
-  @Override
-  public ElementMatcher<? super TypeDescription> typeMatcher() {
-    // CachingAsyncHttpClient rejects overrides to AsyncHandler
-    // It also delegates to another AsyncHttpClient
-    return nameStartsWith("play.")
-        .<TypeDescription>and(
-            hasInterface(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"))),
-        PlayWSClientInstrumentation.class.getName() + "$ClientAdvice");
-  }
-
-  @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 class PlayWSClientInstrumentation extends BasePlayWSClientInstrumentation {
   public static class ClientAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static AgentSpan methodEnter(
@@ -83,7 +28,11 @@ public static AgentSpan methodEnter(
       DECORATE.onRequest(span, request);
       propagate().inject(span, request, SETTER);
 
-      asyncHandler = new AsyncHandlerWrapper(asyncHandler, span);
+      if (asyncHandler instanceof StreamedAsyncHandler) {
+        asyncHandler = new StreamedAsyncHandlerWrapper((StreamedAsyncHandler) asyncHandler, span);
+      } else {
+        asyncHandler = new AsyncHandlerWrapper(asyncHandler, span);
+      }
 
       return span;
     }

dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws1/StreamedAsyncHandlerWrapper.java

@@ -0,0 +1,20 @@
+package datadog.trace.instrumentation.playws1;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import org.reactivestreams.Publisher;
+import play.shaded.ahc.org.asynchttpclient.handler.StreamedAsyncHandler;
+
+public class StreamedAsyncHandlerWrapper extends AsyncHandlerWrapper
+    implements StreamedAsyncHandler {
+  private final StreamedAsyncHandler streamedDelegate;
+
+  public StreamedAsyncHandlerWrapper(final StreamedAsyncHandler delegate, final AgentSpan span) {
+    super(delegate, span);
+    streamedDelegate = delegate;
+  }
+
+  @Override
+  public State onStream(final Publisher publisher) {
+    return streamedDelegate.onStream(publisher);
+  }
+}

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

@@ -1,74 +1,131 @@
-import akka.actor.ActorSystem
-import akka.stream.ActorMaterializer
-import akka.stream.ActorMaterializerSettings
-import datadog.trace.agent.test.base.HttpClientTest
-import datadog.trace.instrumentation.playws1.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 scala.collection.JavaConverters
+import scala.concurrent.Await
+import scala.concurrent.ExecutionContext
+import scala.concurrent.Future
+import scala.concurrent.duration.Duration
 import spock.lang.Shared
 
 import java.util.concurrent.TimeUnit
 
-class PlayWSClientTest extends HttpClientTest<PlayWSClientDecorator> {
-  @Shared
-  ActorSystem system
-
+class PlayJavaWSClientTest extends PlayWSClientTestBase {
   @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)
+  @Override
+  int doRequest(String method, URI uri, Map<String, String> headers, Closure callback) {
+    StandaloneWSRequest wsRequest = wsClient.url(uri.toURL().toString()).setFollowRedirects(true)
 
-    AsyncHttpClientConfig asyncHttpClientConfig =
-      new DefaultAsyncHttpClientConfig.Builder()
-        .setMaxRequestRetry(0)
-        .setShutdownQuietPeriod(0)
-        .setShutdownTimeout(0)
-        .build()
+    headers.entrySet().each { entry -> wsRequest.addHeader(entry.getKey(), entry.getValue()) }
+    StandaloneWSResponse wsResponse = wsRequest.setMethod(method).execute()
+      .whenComplete({ response, throwable ->
+        callback?.call()
+      }).toCompletableFuture().get(5, TimeUnit.SECONDS)
 
-    AsyncHttpClient asyncHttpClient = new DefaultAsyncHttpClient(asyncHttpClientConfig)
+    return wsResponse.getStatus()
+  }
 
+  def setupSpec() {
     wsClient = new StandaloneAhcWSClient(asyncHttpClient, materializer)
   }
 
   def cleanupSpec() {
     wsClient?.close()
-    system?.terminate()
   }
+}
+
+class PlayJavaStreamedWSClientTest extends PlayWSClientTestBase {
+  @Shared
+  StandaloneWSClient wsClient
 
   @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)
+    StandaloneWSResponse wsResponse = wsRequest.setMethod(method).stream()
       .whenComplete({ response, throwable ->
         callback?.call()
       }).toCompletableFuture().get(5, TimeUnit.SECONDS)
 
+    // The status can be ready before the body so explicity call wait for body to be ready
+    wsResponse.getBodyAsSource().runFold("", { acc, out -> "" }, materializer)
+      .toCompletableFuture().get(5, TimeUnit.SECONDS)
     return wsResponse.getStatus()
   }
 
+  def setupSpec() {
+    wsClient = new StandaloneAhcWSClient(asyncHttpClient, materializer)
+  }
+
+  def cleanupSpec() {
+    wsClient?.close()
+  }
+}
+
+class PlayScalaWSClientTest extends PlayWSClientTestBase {
+  @Shared
+  play.api.libs.ws.StandaloneWSClient wsClient
+
   @Override
-  PlayWSClientDecorator decorator() {
-    return PlayWSClientDecorator.DECORATE
+  int doRequest(String method, URI uri, Map<String, String> headers, Closure callback) {
+    Future<play.api.libs.ws.StandaloneWSResponse> futureResponse = wsClient.url(uri.toURL().toString())
+      .withMethod(method)
+      .withFollowRedirects(true)
+      .withHttpHeaders(JavaConverters.mapAsScalaMap(headers).toSeq())
+      .execute()
+      .transform({ theTry ->
+        callback?.call()
+        theTry
+      }, ExecutionContext.global())
+
+    play.api.libs.ws.StandaloneWSResponse wsResponse = Await.result(futureResponse, Duration.apply(5, TimeUnit.SECONDS))
+
+    return wsResponse.status()
   }
 
-  String expectedOperationName() {
-    return "play-ws.request"
+  def setupSpec() {
+    wsClient = new play.api.libs.ws.ahc.StandaloneAhcWSClient(asyncHttpClient, materializer)
   }
 
+  def cleanupSpec() {
+    wsClient?.close()
+  }
+}
+
+class PlayScalaStreamedWSClientTest extends PlayWSClientTestBase {
+  @Shared
+  play.api.libs.ws.StandaloneWSClient wsClient
+
   @Override
-  boolean testCircularRedirects() {
-    return false
+  int doRequest(String method, URI uri, Map<String, String> headers, Closure callback) {
+    Future<play.api.libs.ws.StandaloneWSResponse> futureResponse = wsClient.url(uri.toURL().toString())
+      .withMethod(method)
+      .withFollowRedirects(true)
+      .withHttpHeaders(JavaConverters.mapAsScalaMap(headers).toSeq())
+      .stream()
+      .transform({ theTry ->
+        callback?.call()
+        theTry
+      }, ExecutionContext.global())
+
+    play.api.libs.ws.StandaloneWSResponse wsResponse = Await.result(futureResponse, Duration.apply(5, TimeUnit.SECONDS))
+
+    // The status can be ready before the body so explicity call wait for body to be ready
+    Await.result(
+      wsResponse.bodyAsSource().runFold("", { acc, out -> "" }, materializer),
+      Duration.apply(5, TimeUnit.SECONDS))
+    return wsResponse.status()
+  }
+
+  def setupSpec() {
+    wsClient = new play.api.libs.ws.ahc.StandaloneAhcWSClient(asyncHttpClient, materializer)
+  }
+
+  def cleanupSpec() {
+    wsClient?.close()
   }
 }

dd-java-agent/instrumentation/play-ws-2.1/play-ws-2.1.gradle

@@ -7,12 +7,11 @@ apply from: "${rootDir}/gradle/java.gradle"
 
 apply plugin: 'org.unbroken-dome.test-sets'
 
-// TODO: Uncomment when there are more releases, right now 2.1.0 is the latest release
-//testSets {
-//  latestDepTest {
-//    dirName = 'test'
-//  }
-//}
+testSets {
+  latestDepTest {
+    dirName = 'test'
+  }
+}
 
 muzzle {
   // 2.0.5 was a bad release
@@ -62,6 +61,9 @@ def scalaVersion = '2.12'
 dependencies {
   compileOnly group: 'com.typesafe.play', name: "play-ahc-ws-standalone_$scalaVersion", version: '2.1.0'
 
+  compile project(':dd-java-agent:instrumentation:play-ws')
+  testCompile project(path: ':dd-java-agent:instrumentation:play-ws', configuration: 'testArtifacts')
+
   testCompile project(':dd-java-agent:instrumentation:java-concurrent')
 
   // These are to ensure cross compatibility
@@ -71,7 +73,5 @@ dependencies {
 
   testCompile group: 'com.typesafe.play', name: "play-ahc-ws-standalone_$scalaVersion", version: '2.1.0'
 
-  // TODO: Uncomment when there are more releases, right now 2.1.0 is the latest release
-
-  // latestDepTestCompile group: 'com.typesafe.play', name: "play-ahc-ws-standalone_$scalaVersion", version: '+'
+  latestDepTestCompile group: 'com.typesafe.play', name: "play-ahc-ws-standalone_$scalaVersion", version: '+'
 }

dd-java-agent/instrumentation/play-ws-2.1/src/main/java/datadog/trace/instrumentation/playws21/AsyncHandlerWrapper.java

@@ -1,7 +1,7 @@
 package datadog.trace.instrumentation.playws21;
 
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
-import static datadog.trace.instrumentation.playws21.PlayWSClientDecorator.DECORATE;
+import static datadog.trace.instrumentation.playws.PlayWSClientDecorator.DECORATE;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.context.TraceScope;