From 5c8eb362ea6a957449419cec6c48b54159d4c68e Mon Sep 17 00:00:00 2001 From: Ivan Topolnak Date: Mon, 11 Nov 2013 18:11:33 -0300 Subject: Fix UowDirectives --- kamon-spray/src/main/scala/kamon/spray/UowDirectives.scala | 3 +-- kamon-trace/src/main/scala/kamon/trace/Trace.scala | 4 ++++ 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) -- cgit v1.2.3