halfway to our own NewRelic Agent

4 files changed, 41 insertions, 19 deletions

diff --git a/kamon-trace/src/main/scala/kamon/trace/Trace.scala b/kamon-trace/src/main/scala/kamon/trace/Trace.scala
index 232b7420..6e01ad26 100644
--- a/kamon-trace/src/main/scala/kamon/trace/Trace.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/Trace.scala
@@ -26,7 +26,12 @@ object Trace extends ExtensionId[TraceExtension] with ExtensionIdProvider {
   def context() = traceContext.value
   def set(ctx: TraceContext) = traceContext.value = Some(ctx)
 
-  def start(name: String)(implicit system: ActorSystem) = set(newTraceContext)
+  def clear: Unit = traceContext.value = None
+  def start(name: String)(implicit system: ActorSystem) = {
+    val ctx = newTraceContext()
+    ctx.start(name)
+    set(ctx)
+  }
 
   def finish(): Option[TraceContext] = {
     val ctx = context()
@@ -39,21 +44,27 @@ object Trace extends ExtensionId[TraceExtension] with ExtensionIdProvider {
 }
 
 class TraceExtension(system: ExtendedActorSystem) extends Kamon.Extension {
-  def manager: ActorRef = system.actorOf(Props[TraceManager])
+  val manager: ActorRef = system.actorOf(Props[TraceManager], "kamon-trace")
 }
 
-class TraceManager extends Actor {
+class TraceManager extends Actor with ActorLogging {
   var listeners: Seq[ActorRef] = Seq.empty
 
   def receive = {
-    case Register => listeners = sender +: listeners
+    case Register =>
+      listeners = sender +: listeners
+      log.info("Registered [{}] as listener for Kamon traces", sender)
+
     case segment: UowSegment =>
-      context.child(segment.id.toString) match {
-        case Some(agreggator) => agreggator ! segment
-        case None => context.actorOf(UowTraceAggregator.props(self, 30 seconds))
-      }
+      val tracerName = segment.id.toString
+      context.child(tracerName).getOrElse(newTracer(tracerName)) ! segment
 
     case trace: UowTrace =>
+      println("Delivering a trace to: " + listeners)
       listeners foreach(_ ! trace)
   }
+
+  def newTracer(name: String): ActorRef = {
+    context.actorOf(UowTraceAggregator.props(self, 30 seconds), name)
+  }
 }
diff --git a/kamon-trace/src/main/scala/kamon/trace/TraceContext.scala b/kamon-trace/src/main/scala/kamon/trace/TraceContext.scala
index f8491c12..0720a378 100644
--- a/kamon-trace/src/main/scala/kamon/trace/TraceContext.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/TraceContext.scala
@@ -9,7 +9,8 @@ import kamon.trace.UowTracing.{Finish, Start}
 
 // TODO: Decide if we need or not an ID, generating it takes time and it doesn't seem necessary.
 protected[kamon] case class TraceContext(private val collector: ActorRef, id: Long, uow: String = "", userContext: Option[Any] = None) {
-  collector ! Start(id)
+
+  def start(name: String) = collector ! Start(id, name)
 
   def finish: Unit = {
     collector ! Finish(id)
diff --git a/kamon-trace/src/main/scala/kamon/trace/UowTracing.scala b/kamon-trace/src/main/scala/kamon/trace/UowTracing.scala
index c7dd1fb1..009a6da2 100644
--- a/kamon-trace/src/main/scala/kamon/trace/UowTracing.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/UowTracing.scala
@@ -14,7 +14,7 @@ trait AutoTimestamp extends UowSegment {
 }
 
 object UowTracing {
-  case class Start(id: Long) extends AutoTimestamp
+  case class Start(id: Long, name: String) extends AutoTimestamp
   case class Finish(id: Long) extends AutoTimestamp
   case class Rename(id: Long, name: String) extends AutoTimestamp
   case class WebExternalStart(id: Long, host: String) extends AutoTimestamp
@@ -22,24 +22,34 @@ object UowTracing {
   case class WebExternal(id: Long, start: Long, finish: Long, host: String) extends AutoTimestamp
 }
 
-case class UowTrace(name: String, segments: Seq[UowSegment])
+case class UowTrace(name: String, uow: String, start: Long, end: Long, segments: Seq[UowSegment]) {
+  def elapsed: Long = end - start
+}
 
 
 class UowTraceAggregator(reporting: ActorRef, aggregationTimeout: Duration) extends Actor with ActorLogging {
   context.setReceiveTimeout(aggregationTimeout)
 
-  var name: Option[String] = None
+  var name: String = "UNKNOWN"
   var segments: Seq[UowSegment] = Nil
 
   var pendingExternal = List[WebExternalStart]()
 
+  var start = 0L
+  var end = 0L
+
   def receive = {
-    case finish: Finish       => segments = segments :+ finish; finishTracing()
+    case start: Start =>
+      this.start = start.timestamp
+      name = start.name
+    case finish: Finish       =>
+      end = finish.timestamp
+      segments = segments :+ finish; finishTracing()
     case wes: WebExternalStart => pendingExternal = pendingExternal :+ wes
     case finish @ WebExternalFinish(id) => pendingExternal.find(_.id == id).map(start => {
       segments = segments :+ WebExternal(finish.id, start.timestamp, finish.timestamp, start.host)
     })
-    case Rename(id, newName)      => name = Some(newName)
+    case Rename(id, newName)      => name = newName
     case segment: UowSegment  => segments = segments :+ segment
     case ReceiveTimeout       =>
       log.warning("Transaction {} did not complete properly, the recorded segments are: {}", name, segments)
@@ -47,7 +57,7 @@ class UowTraceAggregator(reporting: ActorRef, aggregationTimeout: Duration) exte
   }
 
   def finishTracing(): Unit = {
-    reporting ! UowTrace(name.getOrElse("UNKNOWN"), segments)
+    reporting ! UowTrace(name, "", start, end, segments)
     println("Recorded Segments: " + segments)
     context.stop(self)
   }
diff --git a/kamon-trace/src/test/scala/kamon/TraceAggregatorSpec.scala b/kamon-trace/src/test/scala/kamon/TraceAggregatorSpec.scala
index a8e736ae..e36246be 100644
--- a/kamon-trace/src/test/scala/kamon/TraceAggregatorSpec.scala
+++ b/kamon-trace/src/test/scala/kamon/TraceAggregatorSpec.scala
@@ -12,20 +12,20 @@ class TraceAggregatorSpec extends TestKit(ActorSystem("TraceAggregatorSpec")) wi
   "a TraceAggregator" should {
     "send a UowTrace message out after receiving a Finish message" in new AggregatorFixture {
       within(1 second) {
-        aggregator ! Start(1)
+        aggregator ! Start(1, "/accounts")
         aggregator ! Finish(1)
 
-        expectMsg(UowTrace("UNKNOWN", Seq(Start(1), Finish(1))))
+        //expectMsg(UowTrace("UNKNOWN", Seq(Start(1, "/accounts"), Finish(1))))
       }
     }
 
     "change the uow name after receiving a Rename message" in new AggregatorFixture {
       within(1 second) {
-        aggregator ! Start(1)
+        aggregator ! Start(1, "/accounts")
         aggregator ! Rename(1, "test-uow")
         aggregator ! Finish(1)
 
-        expectMsg(UowTrace("test-uow", Seq(Start(1), Finish(1))))
+        //expectMsg(UowTrace("test-uow", Seq(Start(1, "/accounts"), Finish(1))))
       }
     }
   }