1 files changed, 5 insertions, 38 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()
   }
 }