1 files changed, 15 insertions, 27 deletions
diff --git a/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala b/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala
index b9f09111..c58e9f0c 100644
--- a/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala
+++ b/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala
@@ -16,15 +16,15 @@
 
 package kamon.play.instrumentation
 
-import org.aspectj.lang.annotation.{ Around, Pointcut, Aspect }
+import kamon.Kamon
+import kamon.play.Play
+import kamon.trace.SegmentMetricIdentityLabel
 import org.aspectj.lang.ProceedingJoinPoint
+import org.aspectj.lang.annotation.{ Around, Aspect, Pointcut }
 import kamon.trace.TraceRecorder
-import kamon.metric.TraceMetrics.HttpClientRequest
 import play.api.libs.ws.WS.WSRequest
 import scala.concurrent.Future
 import play.api.libs.ws.Response
-import scala.util.{ Failure, Success }
-import scala.concurrent.ExecutionContext.Implicits.global
 
 @Aspect
 class WSInstrumentation {
@@ -34,27 +34,15 @@ class WSInstrumentation {
 
   @Around("onExecuteRequest(request)")
   def aroundExecuteRequest(pjp: ProceedingJoinPoint, request: WSRequest): Any = {
-    import WSInstrumentation._
-
-    val completionHandle = TraceRecorder.startSegment(HttpClientRequest(request.url), basicRequestAttributes(request))
-
-    val response = pjp.proceed().asInstanceOf[Future[Response]]
-
-    response.onComplete {
-      case Failure(t) ⇒ completionHandle.map(_.finish(Map("completed-with-error" -> t.getMessage)))
-      case Success(_) ⇒ completionHandle.map(_.finish(Map.empty))
-    }
-
-    response
+    TraceRecorder.withTraceContextAndSystem { (ctx, system) ⇒
+      val playExtension = Kamon(Play)(system)
+      val executor = playExtension.defaultDispatcher
+      val segmentName = playExtension.generateHttpClientSegmentName(request)
+      val segment = ctx.startSegment(segmentName, SegmentMetricIdentityLabel.HttpClient)
+      val response = pjp.proceed().asInstanceOf[Future[Response]]
+
+      response.map(result ⇒ segment.finish())(executor)
+      response
+    } getOrElse (pjp.proceed())
   }
-}
-
-object WSInstrumentation {
-
-  def basicRequestAttributes(request: WSRequest): Map[String, String] = {
-    Map[String, String](
-      "host" -> request.header("host").getOrElse("Unknown"),
-      "path" -> request.method)
-  }
-}
-
+}
+\ No newline at end of file