1 files changed, 26 insertions, 0 deletions
diff --git a/kamon-core/src/main/scala/kamon/TraceContextSwap.scala b/kamon-core/src/main/scala/kamon/TraceContextSwap.scala
new file mode 100644
index 00000000..68ee808b
--- /dev/null
+++ b/kamon-core/src/main/scala/kamon/TraceContextSwap.scala
@@ -0,0 +1,26 @@
+package kamon
+
+/**
+ *  Provides support for making a TraceContext available as ThreadLocal and cleanning up afterwards.
+ */
+trait TraceContextSwap {
+
+  def withContext[A](ctx: Option[TraceContext], body: => A): A = withContext(ctx, body, body)
+
+  def withContext[A](ctx: Option[TraceContext], primary: => A, fallback: => A): A = {
+    ctx match {
+      case Some(context) => {
+        Kamon.set(context)
+        val bodyResult = primary
+        Kamon.clear
+
+        bodyResult
+      }
+      case None => fallback
+    }
+
+  }
+
+}
+
+object TraceContextSwap extends TraceContextSwap