implement Span propagation on top of Kamon.Context

1 files changed, 39 insertions, 0 deletions

diff --git a/kamon-core/src/main/scala/kamon/context/Storage.scala b/kamon-core/src/main/scala/kamon/context/Storage.scala
new file mode 100644
index 00000000..6b92ff85
--- /dev/null
+++ b/kamon-core/src/main/scala/kamon/context/Storage.scala
@@ -0,0 +1,39 @@
+package kamon.context
+
+trait Storage {
+  def current(): Context
+  def store(context: Context): Storage.Scope
+}
+
+object Storage {
+
+  trait Scope {
+    def context: Context
+    def close(): Unit
+  }
+
+
+  class ThreadLocal extends Storage {
+    private val tls = new java.lang.ThreadLocal[Context]() {
+      override def initialValue(): Context = Context.Empty
+    }
+
+    override def current(): Context =
+      tls.get()
+
+    override def store(context: Context): Scope = {
+      val newContext = context
+      val previousContext = tls.get()
+      tls.set(newContext)
+
+      new Scope {
+        override def context: Context = newContext
+        override def close(): Unit = tls.set(previousContext)
+      }
+    }
+  }
+
+  object ThreadLocal {
+    def apply(): ThreadLocal = new ThreadLocal()
+  }
+}
\ No newline at end of file