Added new respondWithUow directive to add the X-UOW header back in the response.

Changed ActorLogging to just propagate the trace context (does not add to the MDC anymore).
Created UowConverter to let logback know where the extract the UOW from.

2 files changed, 12 insertions, 10 deletions

diff --git a/kamon-trace/src/main/scala/kamon/trace/UowTracing.scala b/kamon-trace/src/main/scala/kamon/trace/UowTracing.scala
index 3d1a7778..fcddcf90 100644
--- a/kamon-trace/src/main/scala/kamon/trace/UowTracing.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/UowTracing.scala
@@ -79,4 +79,12 @@ class UowTraceAggregator(reporting: ActorRef, aggregationTimeout: Duration) exte
 
 object UowTraceAggregator {
   def props(reporting: ActorRef, aggregationTimeout: Duration) = Props(classOf[UowTraceAggregator], reporting, aggregationTimeout)
-}
\ No newline at end of file
+}
+
+package logback {
+  import ch.qos.logback.classic.pattern.ClassicConverter
+  import ch.qos.logback.classic.spi.ILoggingEvent
+  class UowConverter extends ClassicConverter {
+    def convert(event: ILoggingEvent): String = Trace.context().map(_.uow).getOrElse("undefined")
+  }
+}
diff --git a/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorLoggingTracing.scala b/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorLoggingTracing.scala
index 783fd6e1..0d06eb2c 100644
--- a/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorLoggingTracing.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorLoggingTracing.scala
@@ -17,8 +17,7 @@ package kamon.trace.instrumentation
 
 import org.aspectj.lang.annotation.{ Around, Pointcut, DeclareMixin, Aspect }
 import org.aspectj.lang.ProceedingJoinPoint
-import org.slf4j.MDC
-import kamon.trace.{ TraceContext, ContextAware, Trace }
+import kamon.trace.{ ContextAware, Trace }
 
 @Aspect
 class ActorLoggingTracing {
@@ -31,13 +30,8 @@ class ActorLoggingTracing {
 
   @Around("withMdcInvocation(logSource, logEvent, logStatement)")
   def aroundWithMdcInvocation(pjp: ProceedingJoinPoint, logSource: String, logEvent: ContextAware, logStatement: () ⇒ _): Unit = {
-    logEvent.traceContext match {
-      case Some(ctx) ⇒
-        MDC.put("uow", ctx.uow)
-        pjp.proceed()
-        MDC.remove("uow")
-
-      case None ⇒ pjp.proceed()
+    Trace.withContext(logEvent.traceContext) {
+      pjp.proceed()
     }
   }
 }