! core: replace Option[TraceContext] by empty object pattern and implement basic segments with renaming.

1 files changed, 8 insertions, 6 deletions

diff --git a/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala b/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala
index 3ff074b6..0766af74 100644
--- a/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala
+++ b/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala
@@ -24,18 +24,20 @@ object TraceLocal {
     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))
+  def store(key: TraceLocalKey)(value: key.ValueType): Unit = TraceRecorder.currentContext match {
+    case ctx: DefaultTraceContext ⇒ ctx.traceLocalStorage.store(key)(value)
+    case EmptyTraceContext        ⇒ // Can't store in the empty context.
+  }
 
+  def retrieve(key: TraceLocalKey): Option[key.ValueType] = TraceRecorder.currentContext match {
+    case ctx: DefaultTraceContext ⇒ ctx.traceLocalStorage.retrieve(key)
+    case EmptyTraceContext        ⇒ None // Can't retrieve anything from the empty context.
+  }
 }
 
 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])
 }