! core: move the TraceContext manipulation API to the Tracer companion object.

5 files changed, 39 insertions, 40 deletions

diff --git a/kamon-core/src/main/scala/kamon/trace/TraceContext.scala b/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
index 48e56153..613e97e6 100644
--- a/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
+++ b/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
@@ -17,6 +17,7 @@
 package kamon.trace
 
 import java.io.ObjectStreamException
+import java.util.function.Function
 import kamon.trace.TraceContextAware.DefaultTraceContextAware
 import kamon.util.RelativeNanoTimestamp
 
@@ -35,36 +36,16 @@ trait TraceContext {
   def addMetadata(key: String, value: String)
 
   def startTimestamp: RelativeNanoTimestamp
-}
-
-object TraceContext {
-  private[kamon] val _traceContextStorage = new ThreadLocal[TraceContext] {
-    override def initialValue(): TraceContext = EmptyTraceContext
-  }
-
-  def currentContext: TraceContext =
-    _traceContextStorage.get()
-
-  def setCurrentContext(context: TraceContext): Unit =
-    _traceContextStorage.set(context)
-
-  def clearCurrentContext: Unit =
-    _traceContextStorage.remove()
 
-  def withContext[T](context: TraceContext)(code: ⇒ T): T = {
-    val oldContext = _traceContextStorage.get()
-    _traceContextStorage.set(context)
-
-    try code finally _traceContextStorage.set(oldContext)
-  }
-
-  def map[T](f: TraceContext ⇒ T): Option[T] = {
-    val current = currentContext
-    if (current.nonEmpty)
-      Some(f(current))
+  def collect[T](f: TraceContext ⇒ T): Option[T] =
+    if (nonEmpty)
+      Some(f(this))
     else None
-  }
 
+  def collect[T](f: Function[TraceContext, T]): Option[T] =
+    if (nonEmpty)
+      Some(f(this))
+    else None
 }
 
 trait Segment {
@@ -132,7 +113,7 @@ object TraceContextAware {
   def default: TraceContextAware = new DefaultTraceContextAware
 
   class DefaultTraceContextAware extends TraceContextAware {
-    @transient val traceContext = TraceContext.currentContext
+    @transient val traceContext = Tracer.currentContext
 
     //
     // Beware of this hack, it might bite us in the future!
diff --git a/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala b/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala
index 057f564e..0a0a120a 100644
--- a/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala
+++ b/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala
@@ -42,12 +42,12 @@ object TraceLocal {
 
   object HttpContextKey extends TraceLocal.TraceLocalKey { type ValueType = HttpContext }
 
-  def store(key: TraceLocalKey)(value: key.ValueType): Unit = TraceContext.currentContext match {
+  def store(key: TraceLocalKey)(value: key.ValueType): Unit = Tracer.currentContext match {
     case ctx: MetricsOnlyContext ⇒ ctx.traceLocalStorage.store(key)(value)
     case EmptyTraceContext       ⇒ // Can't store in the empty context.
   }
 
-  def retrieve(key: TraceLocalKey): Option[key.ValueType] = TraceContext.currentContext match {
+  def retrieve(key: TraceLocalKey): Option[key.ValueType] = Tracer.currentContext match {
     case ctx: MetricsOnlyContext ⇒ ctx.traceLocalStorage.retrieve(key)
     case EmptyTraceContext       ⇒ None // Can't retrieve anything from the empty context.
   }
diff --git a/kamon-core/src/main/scala/kamon/trace/Tracer.scala b/kamon-core/src/main/scala/kamon/trace/Tracer.scala
index 49ab518f..cf83b36b 100644
--- a/kamon-core/src/main/scala/kamon/trace/Tracer.scala
+++ b/kamon-core/src/main/scala/kamon/trace/Tracer.scala
@@ -35,6 +35,28 @@ trait Tracer {
   def unsubscribe(subscriber: ActorRef): Unit
 }
 
+object Tracer {
+  private[kamon] val _traceContextStorage = new ThreadLocal[TraceContext] {
+    override def initialValue(): TraceContext = EmptyTraceContext
+  }
+
+  def currentContext: TraceContext =
+    _traceContextStorage.get()
+
+  def setCurrentContext(context: TraceContext): Unit =
+    _traceContextStorage.set(context)
+
+  def clearCurrentContext: Unit =
+    _traceContextStorage.remove()
+
+  def withContext[T](context: TraceContext)(code: ⇒ T): T = {
+    val oldContext = _traceContextStorage.get()
+    _traceContextStorage.set(context)
+
+    try code finally _traceContextStorage.set(oldContext)
+  }
+}
+
 private[kamon] class TracerImpl(metricsExtension: Metrics, config: Config) extends Tracer {
   private val _settings = TraceSettings(config)
   private val _hostnamePrefix = Try(InetAddress.getLocalHost.getHostName).getOrElse("unknown-localhost")
diff --git a/kamon-core/src/main/scala/kamon/trace/logging/LogbackTraceTokenConverter.scala b/kamon-core/src/main/scala/kamon/trace/logging/LogbackTraceTokenConverter.scala
index 961c3099..8177ed14 100644
--- a/kamon-core/src/main/scala/kamon/trace/logging/LogbackTraceTokenConverter.scala
+++ b/kamon-core/src/main/scala/kamon/trace/logging/LogbackTraceTokenConverter.scala
@@ -17,14 +17,10 @@ package kamon.trace.logging
 
 import ch.qos.logback.classic.pattern.ClassicConverter
 import ch.qos.logback.classic.spi.ILoggingEvent
-import kamon.trace.TraceContext
+import kamon.trace.Tracer
 
 class LogbackTraceTokenConverter extends ClassicConverter {
-  def convert(event: ILoggingEvent): String = {
-    val ctx = TraceContext.currentContext
-    if (ctx.isEmpty)
-      "undefined"
-    else
-      ctx.token
-  }
+
+  def convert(event: ILoggingEvent): String =
+    Tracer.currentContext.collect(_.token).getOrElse("undefined")
 }
diff --git a/kamon-core/src/main/scala/kamon/trace/logging/MdcKeysSupport.scala b/kamon-core/src/main/scala/kamon/trace/logging/MdcKeysSupport.scala
index 4970d97e..5097fe6b 100644
--- a/kamon-core/src/main/scala/kamon/trace/logging/MdcKeysSupport.scala
+++ b/kamon-core/src/main/scala/kamon/trace/logging/MdcKeysSupport.scala
@@ -17,14 +17,14 @@
 package kamon.trace.logging
 
 import kamon.trace.TraceLocal.AvailableToMdc
-import kamon.trace.{ EmptyTraceContext, MetricsOnlyContext, TraceContext }
+import kamon.trace.{ Tracer, EmptyTraceContext, MetricsOnlyContext, TraceContext }
 
 import org.slf4j.MDC
 
 trait MdcKeysSupport {
 
   def withMdc[A](thunk: ⇒ A): A = {
-    val keys = copyToMdc(TraceContext.currentContext)
+    val keys = copyToMdc(Tracer.currentContext)
     try thunk finally keys.foreach(key ⇒ MDC.remove(key))
   }