+ core: refactor MDC facilities and closes #100

2 files changed, 18 insertions, 50 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 e2ffd3f9..3c79fae4 100644
--- a/kamon-play/src/main/scala/kamon/play/instrumentation/LoggerLikeInstrumentation.scala
+++ b/kamon-play/src/main/scala/kamon/play/instrumentation/LoggerLikeInstrumentation.scala
@@ -15,19 +15,12 @@
 
 package kamon.play.instrumentation
 
-import kamon.trace._
+import kamon.trace.logging.MdcKeysSupport
 import org.aspectj.lang.ProceedingJoinPoint
 import org.aspectj.lang.annotation._
-import org.slf4j.MDC
-import play.api.LoggerLike
 
 @Aspect
-class LoggerLikeInstrumentation {
-
-  import kamon.play.instrumentation.LoggerLikeInstrumentation._
-
-  @DeclareMixin("play.api.LoggerLike+")
-  def mixinContextAwareToLoggerLike: TraceContextAware = TraceContextAware.default
+class LoggerLikeInstrumentation extends MdcKeysSupport {
 
   @Pointcut("execution(* play.api.LoggerLike+.info(..))")
   def infoPointcut(): Unit = {}
@@ -41,35 +34,9 @@ class LoggerLikeInstrumentation {
   @Pointcut("execution(* play.api.LoggerLike+.trace(..))")
   def tracePointcut(): Unit = {}
 
-  @Around("(infoPointcut() || warnPointcut() || errorPointcut() || tracePointcut()) && this(logger)")
-  def aroundLog(pjp: ProceedingJoinPoint, logger: LoggerLike): Any = {
-    withMDC {
-      pjp.proceed()
-    }
-  }
-}
-
-object LoggerLikeInstrumentation {
-
-  @inline final def withMDC[A](block: ⇒ A): A = {
-    val keys = putAndExtractKeys(extractProperties(TraceRecorder.currentContext))
-
-    try block finally keys.foreach(k ⇒ MDC.remove(k))
-  }
-
-  def putAndExtractKeys(values: Iterable[Map[String, Any]]): Iterable[String] = values.map {
-    value ⇒ value.map { case (key, value) ⇒ MDC.put(key, value.toString); key }
-  }.flatten
-
-  def extractProperties(traceContext: TraceContext): Iterable[Map[String, Any]] = traceContext match {
-    case ctx: DefaultTraceContext ⇒
-      ctx.traceLocalStorage.underlyingStorage.values.collect {
-        case traceLocalValue @ (p: Product) ⇒ {
-          val properties = p.productIterator
-          traceLocalValue.getClass.getDeclaredFields.filter(field ⇒ field.getName != "$outer").map(_.getName -> properties.next).toMap
-        }
-      }
-    case EmptyTraceContext ⇒ Iterable.empty[Map[String, Any]]
+  @Around("(infoPointcut() || warnPointcut() || errorPointcut() || tracePointcut())")
+  def aroundLog(pjp: ProceedingJoinPoint): Any = withMdc {
+    pjp.proceed()
   }
 }
 
diff --git a/kamon-play/src/test/scala/kamon/play/LoggerLikeInstrumentationSpec.scala b/kamon-play/src/test/scala/kamon/play/LoggerLikeInstrumentationSpec.scala
index c41f7004..04ae9729 100644
--- a/kamon-play/src/test/scala/kamon/play/LoggerLikeInstrumentationSpec.scala
+++ b/kamon-play/src/test/scala/kamon/play/LoggerLikeInstrumentationSpec.scala
@@ -16,10 +16,11 @@
 package kamon.play
 
 import ch.qos.logback.classic.spi.ILoggingEvent
-import ch.qos.logback.classic.{ AsyncAppender, LoggerContext }
+import ch.qos.logback.classic.{AsyncAppender, LoggerContext}
 import ch.qos.logback.core.read.ListAppender
 import ch.qos.logback.core.status.NopStatusListener
 import kamon.trace.TraceLocal
+import kamon.trace.TraceLocal.AvailableToMdc
 import org.scalatest.BeforeAndAfter
 import org.scalatestplus.play._
 import org.slf4j
@@ -28,8 +29,9 @@ import play.api.mvc.Results.Ok
 import play.api.mvc._
 import play.api.test.Helpers._
 import play.api.test._
+import scala.concurrent.duration._
 
-import scala.concurrent.Future
+import scala.concurrent.{Await, Future}
 
 class LoggerLikeInstrumentationSpec extends PlaySpec with OneServerPerSuite with BeforeAndAfter {
 
@@ -41,11 +43,8 @@ class LoggerLikeInstrumentationSpec extends PlaySpec with OneServerPerSuite with
   val headerValue = "My header value"
   val otherValue = "My other value"
 
-  case class LocalStorageValue(header: String, other: String)
-
-  object TraceLocalKey extends TraceLocal.TraceLocalKey {
-    type ValueType = LocalStorageValue
-  }
+  val TraceLocalHeaderKey = AvailableToMdc("header")
+  val TraceLocalOtherKey = AvailableToMdc("other")
 
   before {
     LoggingHandler.startLogging()
@@ -60,7 +59,8 @@ class LoggerLikeInstrumentationSpec extends PlaySpec with OneServerPerSuite with
     case ("GET", "/logging") ⇒
       Action.async {
         Future {
-          TraceLocal.store(TraceLocalKey)(LocalStorageValue(headerValue, otherValue))
+          TraceLocal.store(TraceLocalHeaderKey)(headerValue)
+          TraceLocal.store(TraceLocalOtherKey)(otherValue)
           LoggingHandler.info(infoMessage)
           Ok("OK")
         }(executor)
@@ -68,12 +68,13 @@ class LoggerLikeInstrumentationSpec extends PlaySpec with OneServerPerSuite with
   })
 
   "the LoggerLike instrumentation" should {
-    "be put the properties of TraceLocal into the MDC as key -> value in a request" in {
+    "allow retrieve a value from the MDC when was created a key of type AvailableToMdc in the current request" in {
       LoggingHandler.appenderStart()
 
-      val Some(result) = route(FakeRequest(GET, "/logging"))
-      Thread.sleep(500) // wait to complete the future
-      TraceLocal.retrieve(TraceLocalKey) must be(Some(LocalStorageValue(headerValue, otherValue)))
+      Await.result(route(FakeRequest(GET, "/logging")).get, 500 millis)
+
+      TraceLocal.retrieve(TraceLocalHeaderKey) must be(Some(headerValue))
+      TraceLocal.retrieve(TraceLocalOtherKey) must be(Some(otherValue))
 
       LoggingHandler.appenderStop()