+ core: initial support for TraceLocal storage

3 files changed, 77 insertions, 0 deletions

diff --git a/kamon-core/src/main/scala/kamon/trace/TraceContext.scala b/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
index 9980a022..307cf17a 100644
--- a/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
+++ b/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
@@ -32,6 +32,8 @@ trait TraceContext {
   def levelOfDetail: TracingLevelOfDetail
   def startSegment(identity: SegmentIdentity, metadata: Map[String, String]): SegmentCompletionHandle
   def finish(metadata: Map[String, String])
+
+  private[kamon] val traceLocalStorage: TraceLocalStorage = new TraceLocalStorage
 }
 
 object TraceContext {
diff --git a/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala b/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala
new file mode 100644
index 00000000..c79fa632
--- /dev/null
+++ b/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala
@@ -0,0 +1,25 @@
+package kamon.trace
+
+import scala.collection.concurrent.TrieMap
+import kamon.trace.TraceLocal.TraceLocalKey
+
+object TraceLocal {
+  trait TraceLocalKey {
+    type ValueType
+  }
+
+  def store(key: TraceLocalKey)(value: key.ValueType): Unit =
+    TraceRecorder.currentContext.map(_.traceLocalStorage.store(key)(value))
+
+  def retrieve(key: TraceLocalKey): Option[key.ValueType] =
+    TraceRecorder.currentContext.flatMap(_.traceLocalStorage.retrieve(key))
+
+}
+
+class TraceLocalStorage {
+  val underlyingStorage = TrieMap[TraceLocal.TraceLocalKey, Any]()
+
+  def store(key: TraceLocalKey)(value: key.ValueType): Unit = underlyingStorage.put(key, value)
+
+  def retrieve(key: TraceLocalKey): Option[key.ValueType] = underlyingStorage.get(key).map(_.asInstanceOf[key.ValueType])
+}
diff --git a/kamon-core/src/test/scala/kamon/trace/TraceLocalSpec.scala b/kamon-core/src/test/scala/kamon/trace/TraceLocalSpec.scala
new file mode 100644
index 00000000..53ff8a24
--- /dev/null
+++ b/kamon-core/src/test/scala/kamon/trace/TraceLocalSpec.scala
@@ -0,0 +1,50 @@
+package kamon.trace
+
+import akka.testkit.TestKit
+import akka.actor.ActorSystem
+import org.scalatest.{ OptionValues, Matchers, WordSpecLike }
+import org.scalatest.concurrent.PatienceConfiguration
+
+class TraceLocalSpec extends TestKit(ActorSystem("trace-local-spec")) with WordSpecLike with Matchers
+    with PatienceConfiguration with OptionValues {
+
+  object SampleTraceLocalKey extends TraceLocal.TraceLocalKey { type ValueType = String }
+
+  "the TraceLocal storage" should {
+    "allow storing and retrieving values" in {
+      TraceRecorder.withNewTraceContext("store-and-retrieve-trace-local") {
+        val testString = "Hello World"
+
+        TraceLocal.store(SampleTraceLocalKey)(testString)
+        TraceLocal.retrieve(SampleTraceLocalKey).value should equal(testString)
+      }
+    }
+
+    "return None when retrieving a non existent key" in {
+      TraceRecorder.withNewTraceContext("non-existent-key") {
+        TraceLocal.retrieve(SampleTraceLocalKey) should equal(None)
+      }
+    }
+
+    "return None when retrieving a key without a current TraceContext" in {
+      TraceLocal.retrieve(SampleTraceLocalKey) should equal(None)
+    }
+
+    "be attached to the TraceContext when it is propagated" in {
+      val testString = "Hello World"
+      val testContext = TraceRecorder.withNewTraceContext("manually-propagated-trace-local") {
+        TraceLocal.store(SampleTraceLocalKey)(testString)
+        TraceLocal.retrieve(SampleTraceLocalKey).value should equal(testString)
+        TraceRecorder.currentContext
+      }
+
+      /** No TraceLocal should be available here */
+      TraceLocal.retrieve(SampleTraceLocalKey) should equal(None)
+
+      TraceRecorder.withTraceContext(testContext) {
+        TraceLocal.retrieve(SampleTraceLocalKey).value should equal(testString)
+      }
+    }
+  }
+
+}