instrumentation is back in place

2 files changed, 20 insertions, 0 deletions

diff --git a/kamon-trace/src/main/scala/kamon/trace/Segments.scala b/kamon-trace/src/main/scala/kamon/trace/Segments.scala
index 0e1778bf..7cc20181 100644
--- a/kamon-trace/src/main/scala/kamon/trace/Segments.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/Segments.scala
@@ -16,6 +16,8 @@
 
 package kamon.trace
 
+import kamon.trace.Trace.SegmentCompletionHandle
+
 object Segments {
   
   trait Category
@@ -27,4 +29,13 @@ object Segments {
   case class End(attributes: Map[String, String] = Map(), timestamp: Long = System.nanoTime())
   
   case class Segment(start: Start, end: End)
+
+
+
+
+  trait SegmentCompletionHandleAware {
+    var completionHandle: Option[SegmentCompletionHandle]
+  }
+
+  trait ContextAndSegmentCompletionAware extends ContextAware with SegmentCompletionHandleAware
 }
diff --git a/kamon-trace/src/main/scala/kamon/trace/Trace.scala b/kamon-trace/src/main/scala/kamon/trace/Trace.scala
index 02947dd6..6eff69e5 100644
--- a/kamon-trace/src/main/scala/kamon/trace/Trace.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/Trace.scala
@@ -61,9 +61,18 @@ object Trace extends ExtensionId[TraceExtension] with ExtensionIdProvider {
   // TODO: FIX
   def newTraceContext(name: String)(implicit system: ActorSystem): TraceContext = TraceContext(Kamon(Trace), tranid.getAndIncrement, name)
 
+  def startSegment(category: Segments.Category, description: String = "", attributes: Map[String, String] = Map()): SegmentCompletionHandle = {
+    val start = Segments.Start(category, description, attributes)
+    SegmentCompletionHandle(start)
+  }
+
   def startSegment(start: Segments.Start): SegmentCompletionHandle = SegmentCompletionHandle(start)
 
   case class SegmentCompletionHandle(start: Segments.Start) {
+    def complete(): Unit = {
+      val end = Segments.End()
+      println(s"Completing the Segment: $start - $end")
+    }
     def complete(end: Segments.End): Unit = {
       println(s"Completing the Segment: $start - $end")
     }