Tag WSClient requests with information about requests

play-zipkin-tracing/build.sbt

@@ -83,7 +83,9 @@ lazy val play = (project in file("play")).
     libraryDependencies ++= Seq(
       "com.typesafe.play" %% "play" % playVersion % Provided,
       "com.typesafe.play" %% "play-ws" % playVersion % Provided,
-      "com.typesafe.play" %% "play-guice" % playVersion % Test
+      "com.typesafe.play" %% "play-guice" % playVersion % Test,
+      "com.typesafe.play" %% "play-ahc-ws" % playVersion % Test,
+      "de.leanovate.play-mockws" %% "play-mockws" % "2.7.0" % Test
     )
   ).dependsOn(
     core % "test->test;compile->compile",

play-zipkin-tracing/play/src/main/scala/brave/play/TraceWSClient.scala

@@ -2,8 +2,8 @@ package brave.play
 
 import java.io.{File, IOException}
 import java.net.URI
-import javax.inject.Inject
 
+import javax.inject.Inject
 import akka.stream.scaladsl.Source
 import akka.util.ByteString
 import play.api.libs.ws._
@@ -65,12 +65,22 @@ private class TraceWSRequest(spanName: String, request: WSRequest, tracer: Zipki
   override def withUrl(url: String): TraceWSRequest = new TraceWSRequest(spanName, request.withUrl(url), tracer, traceData)
   override def withMethod(method: String): TraceWSRequest = new TraceWSRequest(spanName, request.withMethod(method), tracer, traceData)
 
-  override def execute(): Future[Response] = tracer.traceFuture(spanName){ data =>
-    request.addHttpHeaders(tracer.toMap(data.span).toSeq: _*).execute()
-  }(traceData)
-  override def stream(): Future[Response] = tracer.traceFuture(spanName){ data =>
-    request.addHttpHeaders(tracer.toMap(data.span).toSeq: _*).stream()
-  }(traceData)
+  private def executeTraced(executeRequestFn: WSRequest => Future[Response]): Future[Response] = {
+    val tracingTags = List(
+      "http.method" -> this.method,
+      "http.path" -> this.uri.getPath,
+      "http.host" -> this.uri.getHost
+    )
+
+    tracer.traceFuture(spanName, tracingTags: _*) { data =>
+      val tracedRequest = request.addHttpHeaders(tracer.toMap(data.span).toSeq: _*)
+      executeRequestFn.apply(tracedRequest)
+    }(traceData)
+  }
+
+  override def execute(): Future[Response] = executeTraced { _.execute() }
+
+  override def stream(): Future[Response] = executeTraced { _.stream() }
 
   override def uri: URI = request.uri
   override def contentType: Option[String] = request.contentType

play-zipkin-tracing/play/src/test/scala/brave/play/module/TraceWSClientSpec.scala

@@ -0,0 +1,36 @@
+package brave.play.module
+
+import brave.play.{TestZipkinTraceService, TraceWSClient, ZipkinTraceServiceLike}
+import mockws.MockWS
+import org.scalatest.{AsyncFlatSpec, Matchers}
+import play.api.inject.bind
+import play.api.inject.guice.GuiceApplicationBuilder
+import play.api.libs.ws.WSClient
+import brave.play
+
+class TraceWSClientSpec extends AsyncFlatSpec with Matchers {
+  val zipkinService = new TestZipkinTraceService
+
+  val injector = new GuiceApplicationBuilder()
+    .bindings(bind[ZipkinTraceServiceLike].to(zipkinService))
+    .overrides(bind[WSClient].to(MockWS(PartialFunction.empty)))
+    .injector
+
+  val client = injector.instanceOf[TraceWSClient]
+
+  it should "set tags in span" in {
+    implicit val parentTraceData = play.TraceData(zipkinService.toSpan(Map.empty)((_, _) => None))
+
+    val result = client.url("test-span", "http://localhost/testpath").execute()
+
+    result.map { _ =>
+      parentTraceData.span.finish()
+      val span = zipkinService.reporter.spans.find(_.name() == "test-span")
+
+      span shouldNot be(empty)
+      span.map(_.tags().get("http.method")) should contain("GET")
+      span.map(_.tags().get("http.host")) should contain("localhost")
+      span.map(_.tags().get("http.path")) should contain("/testpath")
+    }
+  }
+}