= play: minor refactor in RequestInstrumentation

1 files changed, 9 insertions, 16 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 975510e9..ebac38d9 100644
--- a/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
+++ b/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
@@ -16,7 +16,7 @@
 package kamon.play.instrumentation
 
 import kamon.Kamon
-import kamon.play.{ PlayExtension, Play }
+import kamon.play.{ Play, PlayExtension }
 import kamon.trace.{ TraceContextAware, TraceRecorder }
 import org.aspectj.lang.ProceedingJoinPoint
 import org.aspectj.lang.annotation._
@@ -52,31 +52,24 @@ class RequestInstrumentation {
     val essentialAction = (requestHeader: RequestHeader) ⇒ {
 
       val incomingContext = TraceRecorder.currentContext
-      val playExtension = Kamon(Play)(Akka.system())
-      val executor = playExtension.defaultDispatcher
+      val executor = Kamon(Play)(Akka.system()).defaultDispatcher
 
       next(requestHeader).map {
         result ⇒
-          TraceRecorder.currentContext.map { ctx ⇒
-            recordHttpServerMetrics(result, ctx.name, playExtension)
-          }
-
           TraceRecorder.finish()
 
-          incomingContext match {
-            case None ⇒ result
-            case Some(traceContext) ⇒
-              val playExtension = Kamon(Play)(traceContext.system)
-              if (playExtension.includeTraceToken) {
-                result.withHeaders(playExtension.traceTokenHeaderName -> traceContext.token)
-              } else result
-          }
+          incomingContext.map { ctx ⇒
+            val playExtension = Kamon(Play)(ctx.system)
+            recordHttpServerMetrics(result, ctx.name, playExtension)
+            if (playExtension.includeTraceToken) result.withHeaders(playExtension.traceTokenHeaderName -> ctx.token)
+            else result
+          }.getOrElse(result)
       }(executor)
     }
     pjp.proceed(Array(EssentialAction(essentialAction)))
   }
 
-  def recordHttpServerMetrics(result: Result, traceName: String, playExtension: PlayExtension): Unit =
+  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)")