initial implementation of ThreadLocalStorage for Context

2 files changed, 188 insertions, 0 deletions

diff --git a/kamon-core/src/main/scala/kamon/context/Context.scala b/kamon-core/src/main/scala/kamon/context/Context.scala
new file mode 100644
index 00000000..f7c78388
--- /dev/null
+++ b/kamon-core/src/main/scala/kamon/context/Context.scala
@@ -0,0 +1,147 @@
+package kamon.context
+
+class Context private (private val keys: Map[Key[_], Any]) {
+  def get[T](key: Key[T]): T =
+    keys.get(key).getOrElse(key.emptyValue).asInstanceOf[T]
+
+  def withKey[T](key: Key[T], value: T): Context =
+    new Context(keys.updated(key, value))
+}
+
+object Context {
+  val Empty = new Context(Map.empty)
+
+  def apply(): Context = Empty
+  def create(): Context = Empty
+}
+
+
+trait Key[T] {
+  def name: String
+  def emptyValue: T
+  def broadcast: Boolean
+}
+
+object Key {
+
+  def local[T](name: String, emptyValue: T): Key[T] =
+    new Default[T](name, emptyValue, false)
+
+  def broadcast[T](name: String, emptyValue: T): Key[T] =
+    new Default[T](name, emptyValue, true)
+
+
+  private class Default[T](val name: String, val emptyValue: T, val broadcast: Boolean) extends Key[T] {
+    override def hashCode(): Int =
+      name.hashCode
+
+    override def equals(that: Any): Boolean =
+      that.isInstanceOf[Default[_]] && that.asInstanceOf[Default[_]].name == this.name
+  }
+}
+
+trait Storage {
+  def current(): Context
+  def store(context: Context): Scope
+
+  trait Scope {
+    def context: Context
+    def close(): Unit
+  }
+}
+
+object Storage {
+
+  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)
+      }
+    }
+  }
+}
+
+trait KeyCodec[T] {
+  def encode(context: Context): T
+  def decode(carrier: T, context: Context): Context
+}
+
+/*
+object Example {
+  // this is defined somewhere statically, only once.
+  val User = Key.local[Option[User]]("user", None)
+  val Client = Key.local[Option[User]]("client", null)
+  val Span = Key.broadcast[Span]("span", EmptySpan)
+  val storage = Kamon.contextStorage // or something similar.
+
+  storage.get(Span) // returns a Span instance or EmptySpan.
+  storage.get(User) // Returns Option[User] or None if not set.
+  storage.get(Client) // Returns Option[Client] or null if not set.
+
+  // Context Propagation works the very same way as before.
+
+  val scope = storage.store(context)
+  // do something here
+  scope.close()
+
+  // Configuration for codecs would be handled sort of like this:
+
+  //  kamon.context.propagation {
+  //    http-header-codecs {
+  //      "span" = kamon.trace.propagation.B3
+  //    }
+  //
+  //    binary-codecs {
+  //      "span" = kamon.trace.propagation.Binary
+  //    }
+  //  }
+
+
+
+
+
+}*/
+
+
+
+/*
+
+
+
+
+
+class Context(private val keys: Map[Key[_], Any]) {
+  
+
+
+
+
+}
+
+object Context {
+
+
+}
+
+sealed trait Key[T] {
+  def name: String
+}
+
+object Key {
+
+  def local[T](name: String): Key[T] = Local(name)
+
+  case class Local[T](name: String) extends Key[T]
+}*/
diff --git a/kamon-core/src/test/scala/kamon/context/ThreadLocalStorageSpec.scala b/kamon-core/src/test/scala/kamon/context/ThreadLocalStorageSpec.scala
new file mode 100644
index 00000000..39f316ba
--- /dev/null
+++ b/kamon-core/src/test/scala/kamon/context/ThreadLocalStorageSpec.scala
@@ -0,0 +1,41 @@
+package kamon.context
+
+
+import org.scalatest.{Matchers, WordSpec}
+
+class ThreadLocalStorageSpec extends WordSpec with Matchers {
+
+  "the Storage.ThreadLocal implementation of Context storage" should {
+    "return a empty context when no context has been set" in {
+      TLS.current() shouldBe Context.Empty
+    }
+
+    "return the empty value for keys that have not been set in the context" in {
+      TLS.current().get(TestKey) shouldBe 42
+      TLS.current().get(AnotherKey) shouldBe 99
+      TLS.current().get(BroadcastKey) shouldBe "i travel around"
+
+      ScopeWithKey.get(TestKey) shouldBe 43
+      ScopeWithKey.get(AnotherKey) shouldBe 99
+      ScopeWithKey.get(BroadcastKey) shouldBe "i travel around"
+    }
+
+    "allow setting a context as current and remove it when closing the Scope" in {
+      TLS.current() shouldBe Context.Empty
+
+      val scope = TLS.store(ScopeWithKey)
+        TLS.current() shouldBe theSameInstanceAs(ScopeWithKey)
+      scope.close()
+
+      TLS.current() shouldBe Context.Empty
+    }
+
+
+  }
+
+  val TLS: Storage = new Storage.ThreadLocal
+  val TestKey = Key.local("test-key", 42)
+  val AnotherKey = Key.local("another-key", 99)
+  val BroadcastKey = Key.broadcast("broadcast", "i travel around")
+  val ScopeWithKey = Context.create().withKey(TestKey, 43)
+}