= play:

* remove from publishErrorMessage method
* refactor onError method in RequestInstrumentation
* refactor WSInstrumentation in order to propagate the TraceContext when a WS call is executed outside an Action
* improve tests
* closes #33

2 files changed, 37 insertions, 32 deletions

diff --git a/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala b/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
index ebac38d9..8d8de230 100644
--- a/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
+++ b/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
@@ -20,6 +20,7 @@ import kamon.play.{ Play, PlayExtension }
 import kamon.trace.{ TraceContextAware, TraceRecorder }
 import org.aspectj.lang.ProceedingJoinPoint
 import org.aspectj.lang.annotation._
+import play.api.mvc.Results._
 import play.api.mvc._
 import play.libs.Akka
 
@@ -34,7 +35,7 @@ class RequestInstrumentation {
     Kamon(Play)(Akka.system())
   }
 
-  @Before("execution(* play.api.GlobalSettings+.onRouteRequest(..)) && args(requestHeader)")
+  @Before("call(* play.api.GlobalSettings+.onRouteRequest(..)) && args(requestHeader)")
   def onRouteRequest(requestHeader: RequestHeader): Unit = {
     val system = Akka.system()
     val playExtension = Kamon(Play)(system)
@@ -57,10 +58,9 @@ class RequestInstrumentation {
       next(requestHeader).map {
         result ⇒
           TraceRecorder.finish()
-
           incomingContext.map { ctx ⇒
             val playExtension = Kamon(Play)(ctx.system)
-            recordHttpServerMetrics(result, ctx.name, playExtension)
+            recordHttpServerMetrics(result.header, ctx.name, playExtension)
             if (playExtension.includeTraceToken) result.withHeaders(playExtension.traceTokenHeaderName -> ctx.token)
             else result
           }.getOrElse(result)
@@ -69,16 +69,11 @@ class RequestInstrumentation {
     pjp.proceed(Array(EssentialAction(essentialAction)))
   }
 
-  private def recordHttpServerMetrics(result: Result, traceName: String, playExtension: PlayExtension): Unit =
-    playExtension.httpServerMetrics.recordResponse(traceName, result.header.status.toString, 1L)
-
-  @Around("execution(* play.api.GlobalSettings+.onError(..)) && args(request, ex)")
-  def aroundOnError(pjp: ProceedingJoinPoint, request: TraceContextAware, ex: Throwable): Any = request.traceContext match {
-    case None ⇒ pjp.proceed()
-    case Some(ctx) ⇒ {
-      val actorSystem = ctx.system
-      Kamon(Play)(actorSystem).publishErrorMessage(actorSystem, ex.getMessage, ex)
-      pjp.proceed()
-    }
+  @Before("execution(* play.api.GlobalSettings+.onError(..)) && args(request, ex)")
+  def beforeOnError(request: TraceContextAware, ex: Throwable): Unit = request.traceContext.map {
+    ctx ⇒ recordHttpServerMetrics(InternalServerError.header, ctx.name, Kamon(Play)(ctx.system))
   }
+
+  private def recordHttpServerMetrics(header: ResponseHeader, traceName: String, playExtension: PlayExtension): Unit =
+    playExtension.httpServerMetrics.recordResponse(traceName, header.status.toString)
 }
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 2862ba19..3e4d6110 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,16 @@
 
 package kamon.play.instrumentation
 
-import org.aspectj.lang.annotation.{ Around, Pointcut, Aspect }
-import org.aspectj.lang.ProceedingJoinPoint
-import kamon.trace.TraceRecorder
 import kamon.metric.TraceMetrics.HttpClientRequest
-import play.api.libs.ws.WSRequest
-import scala.concurrent.Future
-import play.api.libs.ws.WSResponse
-import scala.util.{ Failure, Success }
+import kamon.trace.{ TraceContext, TraceRecorder }
+import org.aspectj.lang.ProceedingJoinPoint
+import org.aspectj.lang.annotation.{ Around, Aspect, Pointcut }
+import play.api.libs.ws.ning.NingWSRequest
+import play.api.libs.ws.{ WSRequest, WSResponse }
+import play.libs.Akka
+
 import scala.concurrent.ExecutionContext.Implicits.global
+import scala.concurrent.Future
 
 @Aspect
 class WSInstrumentation {
@@ -34,27 +35,36 @@ class WSInstrumentation {
 
   @Around("onExecuteRequest(request)")
   def aroundExecuteRequest(pjp: ProceedingJoinPoint, request: WSRequest): Any = {
-    import WSInstrumentation._
 
-    val completionHandle = TraceRecorder.startSegment(HttpClientRequest(request.url), basicRequestAttributes(request))
+    import kamon.play.instrumentation.WSInstrumentation._
 
-    val response = pjp.proceed().asInstanceOf[Future[WSResponse]]
+    withOrNewTraceContext(TraceRecorder.currentContext)(request) {
+      val response = pjp.proceed().asInstanceOf[Future[WSResponse]]
+      val segmentHandle = TraceRecorder.startSegment(HttpClientRequest(request.url), basicRequestAttributes(request))
 
-    response.onComplete {
-      case Failure(t) ⇒ completionHandle.map(_.finish(Map("completed-with-error" -> t.getMessage)))
-      case Success(_) ⇒ completionHandle.map(_.finish(Map.empty))
+      response.map {
+        r ⇒
+          segmentHandle.map(_.finish())
+          TraceRecorder.finish()
+      }
+      response
     }
-
-    response
   }
 }
 
 object WSInstrumentation {
 
+  def uri(request: WSRequest): java.net.URI = request.asInstanceOf[NingWSRequest].builder.build().getURI
+
   def basicRequestAttributes(request: WSRequest): Map[String, String] = {
     Map[String, String](
-      "host" -> request.header("host").getOrElse("Unknown"),
-      "path" -> request.method)
+      "host" -> uri(request).getHost,
+      "path" -> uri(request).getPath,
+      "method" -> request.method)
   }
-}
 
+  def withOrNewTraceContext[T](context: Option[TraceContext])(request: WSRequest)(thunk: ⇒ T): T = {
+    if (context.isDefined) TraceRecorder.withTraceContext(context) { thunk }
+    else TraceRecorder.withNewTraceContext(request.url, metadata = basicRequestAttributes(request)) { thunk }(Akka.system())
+  }
+}
\ No newline at end of file