Fix UowDirectives

2 files changed, 5 insertions, 2 deletions

diff --git a/kamon-spray/src/main/scala/kamon/spray/UowDirectives.scala b/kamon-spray/src/main/scala/kamon/spray/UowDirectives.scala
index 685388df..0c0ea214 100644
--- a/kamon-spray/src/main/scala/kamon/spray/UowDirectives.scala
+++ b/kamon-spray/src/main/scala/kamon/spray/UowDirectives.scala
@@ -12,8 +12,7 @@ trait UowDirectives extends BasicDirectives {
     val uowHeader = request.headers.find(_.name == "X-UOW")
 
     val generatedUow = uowHeader.map(_.value).getOrElse(UowDirectives.newUow)
-    // TODO: Tracer will always have a context at this point, just rename the uow.
-    Trace.set(Trace.context().getOrElse(Trace.newTraceContext()).copy(uow = generatedUow))
+    Trace.transformContext(_.copy(uow = generatedUow))
     request
   }
 }
diff --git a/kamon-trace/src/main/scala/kamon/trace/Trace.scala b/kamon-trace/src/main/scala/kamon/trace/Trace.scala
index 1956509c..d5bb1ddf 100644
--- a/kamon-trace/src/main/scala/kamon/trace/Trace.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/Trace.scala
@@ -33,6 +33,10 @@ object Trace extends ExtensionId[TraceExtension] with ExtensionIdProvider {
     set(ctx)
   }
 
+  def transformContext(f: TraceContext => TraceContext): Unit = {
+    context.map(f).foreach(set(_))
+  }
+
   def finish(): Option[TraceContext] = {
     val ctx = context()
     ctx.map(_.finish)