= play: fixes #84

5 files changed, 20 insertions, 23 deletions

diff --git a/kamon-play/src/main/scala/kamon/play/instrumentation/LoggerLikeInstrumentation.scala b/kamon-play/src/main/scala/kamon/play/instrumentation/LoggerLikeInstrumentation.scala
index b7afeb76..92686ff0 100644
--- a/kamon-play/src/main/scala/kamon/play/instrumentation/LoggerLikeInstrumentation.scala
+++ b/kamon-play/src/main/scala/kamon/play/instrumentation/LoggerLikeInstrumentation.scala
@@ -15,15 +15,16 @@
 
 package kamon.play.instrumentation
 
-import kamon.trace.{ TraceContext, TraceContextAware }
+import kamon.trace.{ TraceContext, TraceContextAware, TraceRecorder }
 import org.aspectj.lang.ProceedingJoinPoint
 import org.aspectj.lang.annotation._
 import org.slf4j.MDC
+import play.api.LoggerLike
 
 @Aspect
 class LoggerLikeInstrumentation {
 
-  import LoggerLikeInstrumentation._
+  import kamon.play.instrumentation.LoggerLikeInstrumentation._
 
   @DeclareMixin("play.api.LoggerLike+")
   def mixinContextAwareToLoggerLike: TraceContextAware = TraceContextAware.default
@@ -41,16 +42,17 @@ class LoggerLikeInstrumentation {
   def tracePointcut(): Unit = {}
 
   @Around("(infoPointcut() || warnPointcut() || errorPointcut() || tracePointcut()) && this(logger)")
-  def aroundLog(pjp: ProceedingJoinPoint, logger: TraceContextAware): Any = {
-    withMDC(logger.traceContext) {
+  def aroundLog(pjp: ProceedingJoinPoint, logger: LoggerLike): Any = {
+    withMDC {
       pjp.proceed()
     }
   }
 }
 
 object LoggerLikeInstrumentation {
-  def withMDC[A](currentContext: Option[TraceContext])(block: ⇒ A): A = {
-    val keys = currentContext.map(extractProperties).map(putAndExtractKeys)
+
+  @inline final def withMDC[A](block: ⇒ A): A = {
+    val keys = TraceRecorder.currentContext.map(extractProperties).map(putAndExtractKeys)
 
     try block finally keys.map(k ⇒ k.foreach(MDC.remove(_)))
   }
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 2308e326..c761e72f 100644
--- a/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
+++ b/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
@@ -15,8 +15,6 @@
 
 package kamon.play.instrumentation
 
-import java.util.Locale
-
 import kamon.Kamon
 import kamon.play.{ Play, PlayExtension }
 import kamon.trace.{ TraceContextAware, TraceRecorder }
@@ -27,8 +25,6 @@ import play.api.mvc.Results._
 import play.api.mvc._
 import play.libs.Akka
 
-import scala.collection.concurrent.TrieMap
-
 @Aspect
 class RequestInstrumentation {
 
@@ -89,6 +85,10 @@ class RequestInstrumentation {
 }
 
 object RequestInstrumentation {
+
+  import java.util.Locale
+  import scala.collection.concurrent.TrieMap
+
   private val cache = TrieMap.empty[String, String]
 
   def normaliseTraceName(requestHeader: RequestHeader): Option[String] = requestHeader.tags.get(Routes.ROUTE_VERB).map({ verb ⇒
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 14bcd8c9..87467050 100644
--- a/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala
+++ b/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala
@@ -24,7 +24,6 @@ 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.Future
 
@@ -39,18 +38,14 @@ class WSInstrumentation {
 
     import kamon.play.instrumentation.WSInstrumentation._
 
-    TraceRecorder.currentContext match {
-      case ctx @ Some(_) ⇒
-        TraceRecorder.withTraceContext(ctx) {
-          val executor = Kamon(Play)(Akka.system()).defaultDispatcher
-          val segmentHandle = TraceRecorder.startSegment(HttpClientRequest(request.url), basicRequestAttributes(request))
-          val response = pjp.proceed().asInstanceOf[Future[WSResponse]]
+    TraceRecorder.currentContext.map { ctx ⇒
+      val executor = Kamon(Play)(ctx.system).defaultDispatcher
+      val segmentHandle = TraceRecorder.startSegment(HttpClientRequest(request.url), basicRequestAttributes(request))
+      val response = pjp.proceed().asInstanceOf[Future[WSResponse]]
 
-          response.map(result ⇒ segmentHandle.map(_.finish()))(executor)
-          response
-        }
-      case None ⇒ pjp.proceed()
-    }
+      response.map(result ⇒ segmentHandle.map(_.finish()))(executor)
+      response
+    }.getOrElse(pjp.proceed())
   }
 }
 
diff --git a/kamon-play/src/test/scala/kamon/play/WSInstrumentationSpec.scala b/kamon-play/src/test/scala/kamon/play/WSInstrumentationSpec.scala
index 8b0e0135..3378d477 100644
--- a/kamon-play/src/test/scala/kamon/play/WSInstrumentationSpec.scala
+++ b/kamon-play/src/test/scala/kamon/play/WSInstrumentationSpec.scala
@@ -59,7 +59,6 @@ class WSInstrumentationSpec extends WordSpecLike with Matchers with OneServerPer
       }(Akka.system())
 
       val snapshot = takeSnapshotOf("trace-outside-action")
-      println(snapshot) // force load
       snapshot.elapsedTime.numberOfMeasurements should be(1)
       snapshot.segments.size should be(1)
       snapshot.segments(HttpClientRequest("http://localhost:19001/outside")).numberOfMeasurements should be(1)
diff --git a/project/Settings.scala b/project/Settings.scala
index 98891bc8..0f48b808 100644
--- a/project/Settings.scala
+++ b/project/Settings.scala
@@ -23,6 +23,7 @@ object Settings {
       "-g:vars",
       "-feature",
       "-unchecked",
+      "-optimise",
       "-deprecation",
       "-target:jvm-1.6",
       "-language:postfixOps",