diff options
Diffstat (limited to 'kamon-core/src/main/scala/kamon/trace/TraceContext.scala')
-rw-r--r-- | kamon-core/src/main/scala/kamon/trace/TraceContext.scala | 49 |
1 files changed, 24 insertions, 25 deletions
diff --git a/kamon-core/src/main/scala/kamon/trace/TraceContext.scala b/kamon-core/src/main/scala/kamon/trace/TraceContext.scala index 48e56153..85b7396f 100644 --- a/kamon-core/src/main/scala/kamon/trace/TraceContext.scala +++ b/kamon-core/src/main/scala/kamon/trace/TraceContext.scala @@ -18,7 +18,9 @@ package kamon.trace import java.io.ObjectStreamException import kamon.trace.TraceContextAware.DefaultTraceContextAware -import kamon.util.RelativeNanoTimestamp +import kamon.util.{ SameThreadExecutionContext, Supplier, Function, RelativeNanoTimestamp } + +import scala.concurrent.Future trait TraceContext { def name: String @@ -35,34 +37,31 @@ trait TraceContext { def addMetadata(key: String, value: String) def startTimestamp: RelativeNanoTimestamp -} - -object TraceContext { - private[kamon] val _traceContextStorage = new ThreadLocal[TraceContext] { - override def initialValue(): TraceContext = EmptyTraceContext - } - - def currentContext: TraceContext = - _traceContextStorage.get() - def setCurrentContext(context: TraceContext): Unit = - _traceContextStorage.set(context) - - def clearCurrentContext: Unit = - _traceContextStorage.remove() + def collect[T](f: TraceContext ⇒ T): Option[T] = + if (nonEmpty) + Some(f(this)) + else None - def withContext[T](context: TraceContext)(code: ⇒ T): T = { - val oldContext = _traceContextStorage.get() - _traceContextStorage.set(context) + def collect[T](f: Function[TraceContext, T]): Option[T] = + if (nonEmpty) + Some(f(this)) + else None - try code finally _traceContextStorage.set(oldContext) + def withNewSegment[T](segmentName: String, category: String, library: String)(code: ⇒ T): T = { + val segment = startSegment(segmentName, category, library) + try code finally segment.finish() } - def map[T](f: TraceContext ⇒ T): Option[T] = { - val current = currentContext - if (current.nonEmpty) - Some(f(current)) - else None + // Java variant. + def withNewSegment[T](segmentName: String, category: String, library: String, code: Supplier[T]): T = + withNewSegment(segmentName, category, library)(code.get) + + def withNewAsyncSegment[T](segmentName: String, category: String, library: String)(code: ⇒ Future[T]): Future[T] = { + val segment = startSegment(segmentName, category, library) + val result = code + code.onComplete(_ ⇒ segment.finish())(SameThreadExecutionContext) + result } } @@ -132,7 +131,7 @@ object TraceContextAware { def default: TraceContextAware = new DefaultTraceContextAware class DefaultTraceContextAware extends TraceContextAware { - @transient val traceContext = TraceContext.currentContext + @transient val traceContext = Tracer.currentContext // // Beware of this hack, it might bite us in the future! |