+ core: initial support for TraceLocal storage

2 files changed, 27 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])
+}