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

1 files changed, 26 insertions, 27 deletions

diff --git a/kamon-core/src/main/scala/kamon/trace/TraceRecorder.scala b/kamon-core/src/main/scala/kamon/trace/TraceRecorder.scala
index 778edc42..9b0ba038 100644
--- a/kamon-core/src/main/scala/kamon/trace/TraceRecorder.scala
+++ b/kamon-core/src/main/scala/kamon/trace/TraceRecorder.scala
@@ -23,65 +23,64 @@ import kamon.macros.InlineTraceContextMacro
 import scala.util.Try
 import java.net.InetAddress
 import akka.actor.ActorSystem
-import kamon.trace.TraceContext.SegmentIdentity
 
 object TraceRecorder {
-  private val traceContextStorage = new ThreadLocal[Option[TraceContext]] {
-    override def initialValue(): Option[TraceContext] = None
+  private val traceContextStorage = new ThreadLocal[TraceContext] {
+    override def initialValue(): TraceContext = EmptyTraceContext
   }
 
   private val tokenCounter = new AtomicLong
   private val hostnamePrefix = Try(InetAddress.getLocalHost.getHostName).getOrElse("unknown-localhost")
 
-  def newToken = "%s-%s".format(hostnamePrefix, tokenCounter.incrementAndGet())
+  def newToken: String = hostnamePrefix + "-" + String.valueOf(tokenCounter.incrementAndGet())
 
-  private def newTraceContext(name: String, token: Option[String], metadata: Map[String, String],
-    system: ActorSystem): TraceContext = {
-
-    // In the future this should select between implementations.
-    val finalToken = token.getOrElse(newToken)
-    new SimpleMetricCollectionContext(name, finalToken, metadata, TraceContextOrigin.Local, system)
+  private def newTraceContext(name: String, token: Option[String], system: ActorSystem): TraceContext = {
+    new DefaultTraceContext(
+      name, token.getOrElse(newToken),
+      izOpen = true,
+      LevelOfDetail.OnlyMetrics,
+      TraceContextOrigin.Local,
+      startNanoTime = System.nanoTime)(system)
   }
 
   def joinRemoteTraceContext(traceName: String, traceToken: String, startMilliTime: Long, isOpen: Boolean, system: ActorSystem): TraceContext = {
-    new SimpleMetricCollectionContext(
+    /*new SimpleMetricCollectionContext(
       traceName,
       traceToken,
       Map.empty,
       TraceContextOrigin.Remote,
       system,
       startMilliTime,
-      isOpen)
+      isOpen)*/
+    ???
   }
 
-  def setContext(context: Option[TraceContext]): Unit = traceContextStorage.set(context)
+  def setContext(context: TraceContext): Unit = traceContextStorage.set(context)
 
-  def clearContext: Unit = traceContextStorage.set(None)
+  def clearContext: Unit = traceContextStorage.set(EmptyTraceContext)
 
-  def currentContext: Option[TraceContext] = traceContextStorage.get()
+  def currentContext: TraceContext = traceContextStorage.get()
 
-  def start(name: String, token: Option[String] = None, metadata: Map[String, String] = Map.empty)(implicit system: ActorSystem) = {
-    val ctx = newTraceContext(name, token, metadata, system)
-    traceContextStorage.set(Some(ctx))
+  // TODO: Remove this method.
+  def start(name: String, token: Option[String] = None)(implicit system: ActorSystem) = {
+    //val ctx = newTraceContext(name, token, metadata, system)
+    //traceContextStorage.set(Some(ctx))
   }
 
-  def startSegment(identity: SegmentIdentity, metadata: Map[String, String] = Map.empty): Option[SegmentCompletionHandle] =
-    currentContext.map(_.startSegment(identity, metadata))
-
-  def rename(name: String): Unit = currentContext.map(_.rename(name))
+  def rename(name: String): Unit = currentContext.rename(name)
 
-  def withNewTraceContext[T](name: String, token: Option[String] = None, metadata: Map[String, String] = Map.empty)(thunk: ⇒ T)(implicit system: ActorSystem): T =
-    withTraceContext(Some(newTraceContext(name, token, metadata, system)))(thunk)
+  def withNewTraceContext[T](name: String, token: Option[String] = None)(thunk: ⇒ T)(implicit system: ActorSystem): T =
+    withTraceContext(newTraceContext(name, token, system))(thunk)
 
-  def withTraceContext[T](context: Option[TraceContext])(thunk: ⇒ T): T = {
+  def withTraceContext[T](context: TraceContext)(thunk: ⇒ T): T = {
     val oldContext = currentContext
     setContext(context)
 
     try thunk finally setContext(oldContext)
   }
 
-  def withInlineTraceContextReplacement[T](traceCtx: Option[TraceContext])(thunk: ⇒ T): T = macro InlineTraceContextMacro.withInlineTraceContextImpl[T, Option[TraceContext]]
+  def withInlineTraceContextReplacement[T](traceCtx: TraceContext)(thunk: ⇒ T): T = macro InlineTraceContextMacro.withInlineTraceContextImpl[T, TraceContext]
 
-  def finish(metadata: Map[String, String] = Map.empty): Unit = currentContext.map(_.finish(metadata))
+  def finish(): Unit = currentContext.finish()
 
 }