+ core: cleanup the simple trace implementation

1 files changed, 38 insertions, 32 deletions

diff --git a/kamon-core/src/main/scala/kamon/trace/MetricsOnlyContext.scala b/kamon-core/src/main/scala/kamon/trace/MetricsOnlyContext.scala
index 04e61407..f478d971 100644
--- a/kamon-core/src/main/scala/kamon/trace/MetricsOnlyContext.scala
+++ b/kamon-core/src/main/scala/kamon/trace/MetricsOnlyContext.scala
@@ -4,50 +4,42 @@ import java.util.concurrent.ConcurrentLinkedQueue
 
 import akka.actor.ActorSystem
 import akka.event.LoggingAdapter
-import kamon.Kamon
+import kamon.{ RelativeNanoTimestamp, NanoInterval }
 import kamon.metric.TraceMetrics.TraceMetricRecorder
-import kamon.metric.{ MetricsExtension, TraceMetrics, Metrics }
+import kamon.metric.{ MetricsExtension, TraceMetrics }
 
 import scala.annotation.tailrec
 
-class MetricsOnlyContext(
-  traceName: String,
-  val token: String,
-  izOpen: Boolean,
-  val levelOfDetail: LevelOfDetail,
-  val origin: TraceContextOrigin,
-  nanoTimeztamp: Long,
-  log: LoggingAdapter,
-  metricsExtension: MetricsExtension,
-  val system: ActorSystem)
+private[trace] class MetricsOnlyContext(traceName: String, val token: String, izOpen: Boolean, val levelOfDetail: LevelOfDetail, val origin: TraceContextOrigin,
+  val startRelativeTimestamp: RelativeNanoTimestamp, log: LoggingAdapter, metricsExtension: MetricsExtension, val system: ActorSystem)
     extends TraceContext {
 
   @volatile private var _name = traceName
   @volatile private var _isOpen = izOpen
-  @volatile protected var _elapsedNanoTime = 0L
+  @volatile protected var _elapsedTime = NanoInterval.default
 
-  private val _nanoTimestamp = nanoTimeztamp
   private val _finishedSegments = new ConcurrentLinkedQueue[SegmentLatencyData]()
-  private[kamon] val traceLocalStorage: TraceLocalStorage = new TraceLocalStorage
+  private val _traceLocalStorage = new TraceLocalStorage
 
   def rename(newName: String): Unit =
-    if (isOpen) _name = newName
-    else if (log.isWarningEnabled) log.warning("Can't rename trace from [{}] to [{}] because the trace is already closed.", name, newName)
+    if (isOpen)
+      _name = newName
+    else if (log.isWarningEnabled)
+      log.warning("Can't rename trace from [{}] to [{}] because the trace is already closed.", name, newName)
 
   def name: String = _name
   def isEmpty: Boolean = false
   def isOpen: Boolean = _isOpen
-  def nanoTimestamp: Long = _nanoTimestamp
-  def elapsedNanoTime: Long = _elapsedNanoTime
   def addMetadata(key: String, value: String): Unit = {}
 
   def finish(): Unit = {
     _isOpen = false
-    _elapsedNanoTime = System.nanoTime() - _nanoTimestamp
-    val metricRecorder = metricsExtension.register(TraceMetrics(name), TraceMetrics.Factory)
+    val traceElapsedTime = NanoInterval.since(startRelativeTimestamp)
+    _elapsedTime = traceElapsedTime
 
+    val metricRecorder = metricsExtension.register(TraceMetrics(name), TraceMetrics.Factory)
     metricRecorder.map { traceMetrics ⇒
-      traceMetrics.elapsedTime.record(elapsedNanoTime)
+      traceMetrics.elapsedTime.record(traceElapsedTime.nanos)
       drainFinishedSegments(traceMetrics)
     }
   }
@@ -58,12 +50,12 @@ class MetricsOnlyContext(
   @tailrec private def drainFinishedSegments(metricRecorder: TraceMetricRecorder): Unit = {
     val segment = _finishedSegments.poll()
     if (segment != null) {
-      metricRecorder.segmentRecorder(segment.identity).record(segment.duration)
+      metricRecorder.segmentRecorder(segment.identity).record(segment.duration.nanos)
       drainFinishedSegments(metricRecorder)
     }
   }
 
-  protected def finishSegment(segmentName: String, category: String, library: String, duration: Long): Unit = {
+  protected def finishSegment(segmentName: String, category: String, library: String, duration: NanoInterval): Unit = {
     _finishedSegments.add(SegmentLatencyData(SegmentMetricIdentity(segmentName, category, library), duration))
 
     if (isClosed) {
@@ -73,27 +65,41 @@ class MetricsOnlyContext(
     }
   }
 
+  // Should only be used by the TraceLocal utilities.
+  def traceLocalStorage: TraceLocalStorage = _traceLocalStorage
+
+  // Handle with care and make sure that the trace is closed before calling this method, otherwise NanoInterval.default
+  // will be returned.
+  def elapsedTime: NanoInterval = _elapsedTime
+
   class MetricsOnlySegment(segmentName: String, val category: String, val library: String) extends Segment {
-    protected val segmentStartNanoTime = System.nanoTime()
+    private val _startTimestamp = RelativeNanoTimestamp.now
     @volatile private var _segmentName = segmentName
-    @volatile private var _elapsedNanoTime = 0L
-    @volatile protected var _isOpen = true
+    @volatile private var _elapsedTime = NanoInterval.default
+    @volatile private var _isOpen = true
 
     def name: String = _segmentName
     def isEmpty: Boolean = false
     def addMetadata(key: String, value: String): Unit = {}
     def isOpen: Boolean = _isOpen
-    def elapsedNanoTime: Long = _elapsedNanoTime
 
     def rename(newName: String): Unit =
-      if (isOpen) _segmentName = newName
-      else if (log.isWarningEnabled) log.warning("Can't rename segment from [{}] to [{}] because the segment is already closed.", name, newName)
+      if (isOpen)
+        _segmentName = newName
+      else if (log.isWarningEnabled)
+        log.warning("Can't rename segment from [{}] to [{}] because the segment is already closed.", name, newName)
 
     def finish: Unit = {
       _isOpen = false
-      _elapsedNanoTime = System.nanoTime() - segmentStartNanoTime
+      val segmentElapsedTime = NanoInterval.since(_startTimestamp)
+      _elapsedTime = segmentElapsedTime
 
-      finishSegment(name, category, library, elapsedNanoTime)
+      finishSegment(name, category, library, segmentElapsedTime)
     }
+
+    // Handle with care and make sure that the segment is closed before calling this method, otherwise
+    // NanoInterval.default will be returned.
+    def elapsedTime: NanoInterval = _elapsedTime
+    def startTimestamp: RelativeNanoTimestamp = _startTimestamp
   }
 }
\ No newline at end of file