diff options
Diffstat (limited to 'kamon-core/src/main/scala/kamon/ContextStorage.scala')
-rw-r--r-- | kamon-core/src/main/scala/kamon/ContextStorage.scala | 49 |
1 files changed, 49 insertions, 0 deletions
diff --git a/kamon-core/src/main/scala/kamon/ContextStorage.scala b/kamon-core/src/main/scala/kamon/ContextStorage.scala new file mode 100644 index 00000000..47d0b567 --- /dev/null +++ b/kamon-core/src/main/scala/kamon/ContextStorage.scala @@ -0,0 +1,49 @@ +package kamon + +import kamon.context.{Context, Storage} +import kamon.trace.Span + +import scala.util.control.NonFatal + +trait ContextStorage { + private val _contextStorage = Storage.ThreadLocal() + + def currentContext(): Context = + _contextStorage.current() + + def currentSpan(): Span = + _contextStorage.current().get(Span.ContextKey) + + def storeContext(context: Context): Storage.Scope = + _contextStorage.store(context) + + def withContext[T](context: Context)(f: => T): T = { + val scope = _contextStorage.store(context) + try { + f + } finally { + scope.close() + } + } + + def withContextKey[T, K](key: Context.Key[K], value: K)(f: => T): T = + withContext(currentContext().withKey(key, value))(f) + + def withSpan[T](span: Span)(f: => T): T = + withSpan(span, true)(f) + + def withSpan[T](span: Span, finishSpan: Boolean)(f: => T): T = { + try { + withContextKey(Span.ContextKey, span)(f) + } catch { + case NonFatal(t) => + span.addError(t.getMessage, t) + throw t + + } finally { + if(finishSpan) + span.finish() + } + } + +} |