4 files changed, 17 insertions, 13 deletions

diff --git a/kamon-spray/src/main/scala/kamon/spray/UowDirectives.scala b/kamon-spray/src/main/scala/kamon/spray/UowDirectives.scala
index b4ff1877..44db55e7 100644
--- a/kamon-spray/src/main/scala/kamon/spray/UowDirectives.scala
+++ b/kamon-spray/src/main/scala/kamon/spray/UowDirectives.scala
@@ -21,6 +21,7 @@ import java.util.concurrent.atomic.AtomicLong
 import scala.util.Try
 import java.net.InetAddress
 import kamon.trace.Trace
+import spray.http.HttpHeaders.RawHeader
 
 trait UowDirectives extends BasicDirectives {
   def uow: Directive0 = mapRequest { request ⇒
@@ -30,6 +31,7 @@ trait UowDirectives extends BasicDirectives {
     Trace.transformContext(_.copy(uow = generatedUow))
     request
   }
+  def respondWithUow = mapHttpResponseHeaders(headers ⇒ Trace.context().map(ctx ⇒ RawHeader("X-UOW", ctx.uow) :: headers).getOrElse(headers))
 }
 
 object UowDirectives {
@@ -37,4 +39,4 @@ object UowDirectives {
   val hostnamePrefix = Try(InetAddress.getLocalHost.getHostName).getOrElse("unknown-localhost")
   def newUow = "%s-%s".format(hostnamePrefix, uowCounter.incrementAndGet())
 
-}
\ No newline at end of file
+}
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()
     }
   }
 }
diff --git a/project/Build.scala b/project/Build.scala
index 3a6d4f13..3faf7adc 100644
--- a/project/Build.scala
+++ b/project/Build.scala
@@ -31,7 +31,7 @@ object Build extends Build {
     .settings(aspectJSettings: _*)
     .settings(
       libraryDependencies ++=
-        compile(akkaActor, aspectJ) ++
+        compile(akkaActor, aspectJ, logback) ++
         test(scalatest, akkaTestKit, sprayTestkit, akkaSlf4j, logback))
     .dependsOn(kamonCore)
 
@@ -94,4 +94,4 @@ object Build extends Build {
 
 
   val noPublishing = Seq(publish := (), publishLocal := ())
-}
\ No newline at end of file
+}