diff options
| -rw-r--r-- | kamon-spray/src/main/scala/kamon/spray/UowDirectives.scala | 3 | ||||
| -rw-r--r-- | 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) |