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

9 files changed, 170 insertions, 151 deletions

diff --git a/kamon-core/src/main/scala/kamon/metric/TraceMetrics.scala b/kamon-core/src/main/scala/kamon/metric/TraceMetrics.scala
index 10dbcc01..54626b6c 100644
--- a/kamon-core/src/main/scala/kamon/metric/TraceMetrics.scala
+++ b/kamon-core/src/main/scala/kamon/metric/TraceMetrics.scala
@@ -35,7 +35,7 @@ object TraceMetrics extends MetricGroupCategory {
   case class TraceMetricRecorder(elapsedTime: Histogram, private val segmentRecorderFactory: () ⇒ Histogram)
       extends MetricGroupRecorder {
 
-    private val segments = TrieMap[MetricIdentity, Histogram]()
+    val segments = TrieMap[MetricIdentity, Histogram]()
 
     def segmentRecorder(segmentIdentity: MetricIdentity): Histogram =
       segments.getOrElseUpdate(segmentIdentity, segmentRecorderFactory.apply())
diff --git a/kamon-core/src/main/scala/kamon/trace/TraceContext.scala b/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
index 64ee70be..a5855308 100644
--- a/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
+++ b/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
@@ -23,113 +23,67 @@ import kamon.Kamon
 import kamon.metric._
 import java.util.concurrent.ConcurrentLinkedQueue
 import kamon.trace.TraceContextAware.DefaultTraceContextAware
-import kamon.trace.TraceContext.SegmentIdentity
 import kamon.metric.TraceMetrics.TraceMetricRecorder
 
 import scala.annotation.tailrec
 
-trait TraceContext {
+sealed trait TraceContext {
   def name: String
   def token: String
-  def system: ActorSystem
   def rename(name: String): Unit
-  def levelOfDetail: TracingLevelOfDetail
-  def startSegment(identity: SegmentIdentity, metadata: Map[String, String]): SegmentCompletionHandle
-  def finish(metadata: Map[String, String])
+  def finish(): Unit
   def origin: TraceContextOrigin
-  def startMilliTime: Long
   def isOpen: Boolean
-
-  private[kamon] val traceLocalStorage: TraceLocalStorage = new TraceLocalStorage
-}
-
-object TraceContext {
-  type SegmentIdentity = MetricIdentity
-}
-
-trait SegmentCompletionHandle {
-  def finish(metadata: Map[String, String] = Map.empty)
+  def isEmpty: Boolean
+  def startSegment(segmentName: String, label: String): Segment
 }
 
-case class SegmentData(identity: MetricIdentity, duration: Long, metadata: Map[String, String])
-
-sealed trait TracingLevelOfDetail
-case object OnlyMetrics extends TracingLevelOfDetail
-case object SimpleTrace extends TracingLevelOfDetail
-case object FullTrace extends TracingLevelOfDetail
-
-sealed trait TraceContextOrigin
-object TraceContextOrigin {
-  case object Local extends TraceContextOrigin
-  case object Remote extends TraceContextOrigin
-}
-
-trait TraceContextAware extends Serializable {
-  def captureNanoTime: Long
-  def traceContext: Option[TraceContext]
+sealed trait Segment {
+  def name: String
+  def rename(newName: String): Unit
+  def label: String
+  def finish(): Unit
 }
 
-object TraceContextAware {
-  def default: TraceContextAware = new DefaultTraceContextAware
-
-  class DefaultTraceContextAware extends TraceContextAware {
-    @transient val captureNanoTime = System.nanoTime()
-    @transient val traceContext = TraceRecorder.currentContext
-
-    //
-    // Beware of this hack, it might bite us in the future!
-    //
-    // When using remoting/cluster all messages carry the TraceContext in the envelope in which they
-    // are sent but that doesn't apply to System Messages. We are certain that the TraceContext is
-    // available (if any) when the system messages are read and this will make sure that it is correctly
-    // captured and propagated.
-    @throws[ObjectStreamException]
-    private def readResolve: AnyRef = {
-      new DefaultTraceContextAware
-    }
+case object EmptyTraceContext extends TraceContext {
+  def name: String = "empty-trace"
+  def token: String = ""
+  def rename(name: String): Unit = {}
+  def finish(): Unit = {}
+  def origin: TraceContextOrigin = TraceContextOrigin.Local
+  def isOpen: Boolean = false
+  def isEmpty: Boolean = true
+  def startSegment(segmentName: String, label: String): Segment = EmptySegment
+
+  case object EmptySegment extends Segment {
+    val name: String = "empty-segment"
+    val label: String = "empty-label"
+    def rename(newName: String): Unit = {}
+    def finish: Unit = {}
   }
 }
 
-trait SegmentCompletionHandleAware extends TraceContextAware {
-  @volatile var segmentCompletionHandle: Option[SegmentCompletionHandle] = None
-}
-
-object SegmentCompletionHandleAware {
-  def default: SegmentCompletionHandleAware = new DefaultSegmentCompletionHandleAware
-
-  class DefaultSegmentCompletionHandleAware extends DefaultTraceContextAware with SegmentCompletionHandleAware {}
-}
-
-class SimpleMetricCollectionContext(traceName: String, val token: String, metadata: Map[String, String],
-    val origin: TraceContextOrigin, val system: ActorSystem, val startMilliTime: Long = System.currentTimeMillis,
-    izOpen: Boolean = true) extends TraceContext {
+class DefaultTraceContext(traceName: String, val token: String, izOpen: Boolean, val levelOfDetail: LevelOfDetail,
+    val origin: TraceContextOrigin, startNanoTime: Long)(implicit system: ActorSystem) extends TraceContext {
 
+  val isEmpty: Boolean = false
   @volatile private var _name = traceName
   @volatile private var _isOpen = izOpen
 
-  val levelOfDetail = OnlyMetrics
-  val startNanoTime = System.nanoTime()
-  val finishedSegments = new ConcurrentLinkedQueue[SegmentData]()
-  val metricsExtension = Kamon(Metrics)(system)
+  private val _startNanoTime = startNanoTime
+  private val finishedSegments = new ConcurrentLinkedQueue[SegmentData]()
+  private val metricsExtension = Kamon(Metrics)(system)
+  private[kamon] val traceLocalStorage: TraceLocalStorage = new TraceLocalStorage
 
   def name: String = _name
+  def rename(newName: String): Unit =
+    if (isOpen) _name = newName // TODO: log a warning about renaming a closed trace.
 
-  def rename(newName: String): Unit = _name = newName
-
-  def isOpen(): Boolean = _isOpen
+  def isOpen: Boolean = _isOpen
 
-  def finish(metadata: Map[String, String]): Unit = {
+  def finish(): Unit = {
     _isOpen = false
-
-    val elapsedNanoTime =
-      if (origin == TraceContextOrigin.Local)
-        // Everything is local, nanoTime is still the best resolution we can use.
-        System.nanoTime() - startNanoTime
-      else
-        // For a remote TraceContext we can only rely on the startMilliTime and we need to scale it to nanoseconds
-        // to be consistent with unit used for all latency measurements.
-        (System.currentTimeMillis() - startMilliTime) * 1000000L
-
+    val elapsedNanoTime = System.nanoTime() - _startNanoTime
     val metricRecorder = metricsExtension.register(TraceMetrics(name), TraceMetrics.Factory)
 
     metricRecorder.map { traceMetrics ⇒
@@ -138,6 +92,8 @@ class SimpleMetricCollectionContext(traceName: String, val token: String, metada
     }
   }
 
+  def startSegment(segmentName: String, segmentLabel: String): Segment = new DefaultSegment(segmentName, segmentLabel)
+
   @tailrec private def drainFinishedSegments(metricRecorder: TraceMetricRecorder): Unit = {
     val segment = finishedSegments.poll()
     if (segment != null) {
@@ -146,8 +102,8 @@ class SimpleMetricCollectionContext(traceName: String, val token: String, metada
     }
   }
 
-  private def finishSegment(identity: MetricIdentity, duration: Long, metadata: Map[String, String]): Unit = {
-    finishedSegments.add(SegmentData(identity, duration, metadata))
+  private def finishSegment(segmentName: String, label: String, duration: Long): Unit = {
+    finishedSegments.add(SegmentData(SegmentMetricIdentity(segmentName, label), duration))
 
     if (!_isOpen) {
       metricsExtension.register(TraceMetrics(name), TraceMetrics.Factory).map { traceMetrics ⇒
@@ -156,16 +112,72 @@ class SimpleMetricCollectionContext(traceName: String, val token: String, metada
     }
   }
 
-  def startSegment(identity: SegmentIdentity, metadata: Map[String, String]): SegmentCompletionHandle =
-    new SimpleMetricCollectionCompletionHandle(identity, metadata)
+  class DefaultSegment(segmentName: String, val label: String) extends Segment {
+    private val _segmentStartNanoTime = System.nanoTime()
+    @volatile private var _segmentName = segmentName
+    @volatile private var _isOpen = true
 
-  class SimpleMetricCollectionCompletionHandle(identity: MetricIdentity, startMetadata: Map[String, String]) extends SegmentCompletionHandle {
-    val segmentStartNanoTime = System.nanoTime()
+    def name: String = _segmentName
+    def rename(newName: String): Unit = _segmentName = newName
 
-    def finish(metadata: Map[String, String] = Map.empty): Unit = {
+    def finish: Unit = {
       val segmentFinishNanoTime = System.nanoTime()
-      finishSegment(identity, (segmentFinishNanoTime - segmentStartNanoTime), startMetadata ++ metadata)
+      finishSegment(segmentName, label, (segmentFinishNanoTime - _segmentStartNanoTime))
     }
   }
 }
 
+case class SegmentMetricIdentity(name: String, label: String) extends MetricIdentity
+case class SegmentData(identity: SegmentMetricIdentity, duration: Long)
+
+object SegmentMetricIdentityLabel {
+  val HttpClient = "http-client"
+}
+
+sealed trait LevelOfDetail
+object LevelOfDetail {
+  case object OnlyMetrics extends LevelOfDetail
+  case object SimpleTrace extends LevelOfDetail
+  case object FullTrace extends LevelOfDetail
+}
+
+sealed trait TraceContextOrigin
+object TraceContextOrigin {
+  case object Local extends TraceContextOrigin
+  case object Remote extends TraceContextOrigin
+}
+
+trait TraceContextAware extends Serializable {
+  def captureNanoTime: Long
+  def traceContext: TraceContext
+}
+
+object TraceContextAware {
+  def default: TraceContextAware = new DefaultTraceContextAware
+
+  class DefaultTraceContextAware extends TraceContextAware {
+    @transient val captureNanoTime = System.nanoTime()
+    @transient val traceContext = TraceRecorder.currentContext
+
+    //
+    // Beware of this hack, it might bite us in the future!
+    //
+    // When using remoting/cluster all messages carry the TraceContext in the envelope in which they
+    // are sent but that doesn't apply to System Messages. We are certain that the TraceContext is
+    // available (if any) when the system messages are read and this will make sure that it is correctly
+    // captured and propagated.
+    @throws[ObjectStreamException]
+    private def readResolve: AnyRef = {
+      new DefaultTraceContextAware
+    }
+  }
+}
+
+trait SegmentAware extends TraceContextAware {
+  @volatile var segment: Segment = EmptyTraceContext.EmptySegment
+}
+
+object SegmentAware {
+  def default: SegmentAware = new DefaultSegmentAware
+  class DefaultSegmentAware extends DefaultTraceContextAware with SegmentAware {}
+}
\ No newline at end of file
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])
 }
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()
 
 }
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 4b7dbb28..f052f009 100644
--- a/kamon-core/src/main/scala/kamon/trace/logging/LogbackTraceTokenConverter.scala
+++ b/kamon-core/src/main/scala/kamon/trace/logging/LogbackTraceTokenConverter.scala
@@ -20,5 +20,11 @@ import ch.qos.logback.classic.spi.ILoggingEvent
 import kamon.trace.TraceRecorder
 
 class LogbackTraceTokenConverter extends ClassicConverter {
-  def convert(event: ILoggingEvent): String = TraceRecorder.currentContext.map(_.token).getOrElse("undefined")
+  def convert(event: ILoggingEvent): String = {
+    val ctx = TraceRecorder.currentContext
+    if (ctx.isEmpty)
+      "undefined"
+    else
+      ctx.token
+  }
 }
diff --git a/kamon-core/src/test/scala/kamon/instrumentation/akka/ActorSystemMessageInstrumentationSpec.scala b/kamon-core/src/test/scala/kamon/instrumentation/akka/ActorSystemMessageInstrumentationSpec.scala
index cb39f5e6..d79ccbe0 100644
--- a/kamon-core/src/test/scala/kamon/instrumentation/akka/ActorSystemMessageInstrumentationSpec.scala
+++ b/kamon-core/src/test/scala/kamon/instrumentation/akka/ActorSystemMessageInstrumentationSpec.scala
@@ -3,7 +3,7 @@ package kamon.instrumentation.akka
 import akka.actor.SupervisorStrategy.{ Escalate, Restart, Resume, Stop }
 import akka.actor._
 import akka.testkit.{ ImplicitSender, TestKit }
-import kamon.trace.TraceRecorder
+import kamon.trace.{ EmptyTraceContext, TraceRecorder }
 import org.scalatest.WordSpecLike
 
 import scala.concurrent.duration._
@@ -59,7 +59,7 @@ class ActorSystemMessageInstrumentationSpec extends TestKit(ActorSystem("actor-s
 
         // Ensure we didn't tie the actor with the context
         supervisor ! "context"
-        expectMsg(None)
+        expectMsg(EmptyTraceContext)
       }
 
       "the actor is restarted" in {
@@ -76,7 +76,7 @@ class ActorSystemMessageInstrumentationSpec extends TestKit(ActorSystem("actor-s
 
         // Ensure we didn't tie the actor with the context
         supervisor ! "context"
-        expectMsg(None)
+        expectMsg(EmptyTraceContext)
       }
 
       "the actor is stopped" in {
diff --git a/kamon-core/src/test/scala/kamon/instrumentation/akka/AskPatternInstrumentationSpec.scala b/kamon-core/src/test/scala/kamon/instrumentation/akka/AskPatternInstrumentationSpec.scala
index d914ffe8..17312ba3 100644
--- a/kamon-core/src/test/scala/kamon/instrumentation/akka/AskPatternInstrumentationSpec.scala
+++ b/kamon-core/src/test/scala/kamon/instrumentation/akka/AskPatternInstrumentationSpec.scala
@@ -54,7 +54,6 @@ class AskPatternInstrumentationSpec extends TestKitBase with WordSpecLike with M
       }
       val capturedCtx = warn.asInstanceOf[TraceContextAware].traceContext
 
-      capturedCtx should be('defined)
       capturedCtx should equal(testTraceContext)
     }
   }
diff --git a/kamon-core/src/test/scala/kamon/metric/TraceMetricsSpec.scala b/kamon-core/src/test/scala/kamon/metric/TraceMetricsSpec.scala
index 8a87408d..6453dd77 100644
--- a/kamon-core/src/test/scala/kamon/metric/TraceMetricsSpec.scala
+++ b/kamon-core/src/test/scala/kamon/metric/TraceMetricsSpec.scala
@@ -5,8 +5,7 @@ import akka.testkit.{ ImplicitSender, TestKitBase }
 import com.typesafe.config.ConfigFactory
 import kamon.Kamon
 import kamon.metric.TraceMetrics.TraceMetricsSnapshot
-import kamon.trace.TraceContext.SegmentIdentity
-import kamon.trace.TraceRecorder
+import kamon.trace.{ SegmentMetricIdentity, TraceRecorder }
 import org.scalatest.{ Matchers, WordSpecLike }
 
 class TraceMetricsSpec extends TestKitBase with WordSpecLike with Matchers with ImplicitSender {
@@ -54,39 +53,37 @@ class TraceMetricsSpec extends TestKitBase with WordSpecLike with Matchers with
 
     "record the elapsed time for segments that occur inside a given trace" in {
       TraceRecorder.withNewTraceContext("trace-with-segments") {
-        val segmentHandle = TraceRecorder.startSegment(TraceMetricsTestSegment("test-segment"))
-        segmentHandle.get.finish()
+        val segment = TraceRecorder.currentContext.startSegment("test-segment", "test-label")
+        segment.finish()
         TraceRecorder.finish()
       }
 
       val snapshot = takeSnapshotOf("trace-with-segments")
       snapshot.elapsedTime.numberOfMeasurements should be(1)
       snapshot.segments.size should be(1)
-      snapshot.segments(TraceMetricsTestSegment("test-segment")).numberOfMeasurements should be(1)
+      snapshot.segments(SegmentMetricIdentity("test-segment", "test-label")).numberOfMeasurements should be(1)
     }
 
     "record the elapsed time for segments that finish after their correspondent trace has finished" in {
-      val segmentHandle = TraceRecorder.withNewTraceContext("closing-segment-after-trace") {
-        val sh = TraceRecorder.startSegment(TraceMetricsTestSegment("test-segment"))
+      val segment = TraceRecorder.withNewTraceContext("closing-segment-after-trace") {
+        val s = TraceRecorder.currentContext.startSegment("test-segment", "test-label")
         TraceRecorder.finish()
-        sh
+        s
       }
 
       val beforeFinishSegmentSnapshot = takeSnapshotOf("closing-segment-after-trace")
       beforeFinishSegmentSnapshot.elapsedTime.numberOfMeasurements should be(1)
       beforeFinishSegmentSnapshot.segments.size should be(0)
 
-      segmentHandle.get.finish()
+      segment.finish()
 
       val afterFinishSegmentSnapshot = takeSnapshotOf("closing-segment-after-trace")
       afterFinishSegmentSnapshot.elapsedTime.numberOfMeasurements should be(0)
       afterFinishSegmentSnapshot.segments.size should be(1)
-      afterFinishSegmentSnapshot.segments(TraceMetricsTestSegment("test-segment")).numberOfMeasurements should be(1)
+      afterFinishSegmentSnapshot.segments(SegmentMetricIdentity("test-segment", "test-label")).numberOfMeasurements should be(1)
     }
   }
 
-  case class TraceMetricsTestSegment(name: String) extends SegmentIdentity
-
   def takeSnapshotOf(traceName: String): TraceMetricsSnapshot = {
     val recorder = Kamon(Metrics).register(TraceMetrics(traceName), TraceMetrics.Factory)
     val collectionContext = Kamon(Metrics).buildDefaultCollectionContext
diff --git a/kamon-core/src/test/scala/kamon/trace/TraceContextManipulationSpec.scala b/kamon-core/src/test/scala/kamon/trace/TraceContextManipulationSpec.scala
index 4d0049f1..e2031a72 100644
--- a/kamon-core/src/test/scala/kamon/trace/TraceContextManipulationSpec.scala
+++ b/kamon-core/src/test/scala/kamon/trace/TraceContextManipulationSpec.scala
@@ -3,7 +3,6 @@ package kamon.trace
 import akka.actor.ActorSystem
 import akka.testkit.{ ImplicitSender, TestKitBase }
 import com.typesafe.config.ConfigFactory
-import kamon.trace.TraceContext.SegmentIdentity
 import org.scalatest.{ Matchers, WordSpecLike }
 
 class TraceContextManipulationSpec extends TestKitBase with WordSpecLike with Matchers with ImplicitSender {
@@ -38,7 +37,7 @@ class TraceContextManipulationSpec extends TestKitBase with WordSpecLike with Ma
     "allow starting a trace within a specified block of code, and only within that block of code" in {
       val createdContext = TraceRecorder.withNewTraceContext("start-context") {
         TraceRecorder.currentContext should not be empty
-        TraceRecorder.currentContext.get
+        TraceRecorder.currentContext
       }
 
       TraceRecorder.currentContext shouldBe empty
@@ -48,7 +47,7 @@ class TraceContextManipulationSpec extends TestKitBase with WordSpecLike with Ma
     "allow starting a trace within a specified block of code, providing a trace-token and only within that block of code" in {
       val createdContext = TraceRecorder.withNewTraceContext("start-context-with-token", Some("token-1")) {
         TraceRecorder.currentContext should not be empty
-        TraceRecorder.currentContext.get
+        TraceRecorder.currentContext
       }
 
       TraceRecorder.currentContext shouldBe empty
@@ -70,7 +69,7 @@ class TraceContextManipulationSpec extends TestKitBase with WordSpecLike with Ma
     "allow renaming a trace" in {
       val createdContext = TraceRecorder.withNewTraceContext("trace-before-rename") {
         TraceRecorder.rename("renamed-trace")
-        TraceRecorder.currentContext.get
+        TraceRecorder.currentContext
       }
 
       TraceRecorder.currentContext shouldBe empty
@@ -79,17 +78,22 @@ class TraceContextManipulationSpec extends TestKitBase with WordSpecLike with Ma
 
     "allow creating a segment within a trace" in {
       val createdContext = TraceRecorder.withNewTraceContext("trace-with-segments") {
-        val segmentHandle = TraceRecorder.startSegment(TraceManipulationTestSegment("segment-1"))
-
-        TraceRecorder.currentContext.get
+        val segment = TraceRecorder.currentContext.startSegment("segment-1", "segment-1-label")
+        TraceRecorder.currentContext
       }
 
       TraceRecorder.currentContext shouldBe empty
       createdContext.name shouldBe ("trace-with-segments")
-
     }
-  }
 
-  case class TraceManipulationTestSegment(name: String) extends SegmentIdentity
+    "allow renaming a segment" in {
+      TraceRecorder.withNewTraceContext("trace-with-renamed-segment") {
+        val segment = TraceRecorder.currentContext.startSegment("original-segment-name", "segment-label")
+        segment.name should be("original-segment-name")
 
+        segment.rename("new-segment-name")
+        segment.name should be("new-segment-name")
+      }
+    }
+  }
 }