3 files changed, 93 insertions, 0 deletions
diff --git a/kamon-core/src/main/scala/kamon/trace/UowDirectives.scala b/kamon-core/src/main/scala/kamon/trace/UowDirectives.scala
new file mode 100644
index 00000000..392f53b8
--- /dev/null
+++ b/kamon-core/src/main/scala/kamon/trace/UowDirectives.scala
@@ -0,0 +1,28 @@
+package kamon.trace
+
+import spray.routing.directives.BasicDirectives
+import spray.routing._
+import kamon.Tracer
+import java.util.concurrent.atomic.AtomicLong
+import scala.util.Try
+import java.net.InetAddress
+
+trait UowDirectives extends BasicDirectives {
+  def uow: Directive0 = mapRequest { request =>
+    val uowHeader = request.headers.find(_.name == "X-UOW")
+
+    val generatedUow = uowHeader.map(_.value).getOrElse(UowDirectives.newUow)
+    Tracer.set(Tracer.context().getOrElse(Tracer.newTraceContext()).copy(uow = generatedUow))
+
+    request
+  }
+}
+
+object UowDirectives {
+  val uowCounter = new AtomicLong
+
+  val hostnamePrefix = Try(InetAddress.getLocalHost.getHostName).getOrElse("unknown-localhost")
+
+  def newUow = "%s-%s".format(hostnamePrefix, uowCounter.incrementAndGet())
+
+}
+\ No newline at end of file
diff --git a/kamon-core/src/main/scala/kamon/trace/UowTracing.scala b/kamon-core/src/main/scala/kamon/trace/UowTracing.scala
new file mode 100644
index 00000000..b09478cc
--- /dev/null
+++ b/kamon-core/src/main/scala/kamon/trace/UowTracing.scala
@@ -0,0 +1,57 @@
+package kamon.trace
+
+import akka.actor._
+import scala.concurrent.duration.Duration
+import kamon.trace.UowTracing._
+
+sealed trait UowSegment {
+  def timestamp: Long
+}
+
+trait AutoTimestamp extends UowSegment {
+  val timestamp = System.nanoTime
+}
+
+object UowTracing {
+  case class Start() extends AutoTimestamp
+  case class Finish() extends AutoTimestamp
+  case class Rename(name: String) extends AutoTimestamp
+  case class WebExternalStart(id: Long, host: String) extends AutoTimestamp
+  case class WebExternalFinish(id: Long) extends AutoTimestamp
+  case class WebExternal(start: Long, finish: Long, host: String) extends AutoTimestamp
+}
+
+case class UowTrace(name: String, segments: Seq[UowSegment])
+
+
+class UowTraceAggregator(reporting: ActorRef, aggregationTimeout: Duration) extends Actor with ActorLogging {
+  context.setReceiveTimeout(aggregationTimeout)
+
+  var name: Option[String] = None
+  var segments: Seq[UowSegment] = Nil
+
+  var pendingExternal = List[WebExternalStart]()
+
+  def receive = {
+    case finish: Finish       => segments = segments :+ finish; finishTracing()
+    case wes: WebExternalStart => pendingExternal = pendingExternal :+ wes
+    case finish @ WebExternalFinish(id) => pendingExternal.find(_.id == id).map(start => {
+      segments = segments :+ WebExternal(start.timestamp, finish.timestamp, start.host)
+    })
+    case Rename(newName)      => name = Some(newName)
+    case segment: UowSegment  => segments = segments :+ segment
+    case ReceiveTimeout       =>
+      log.warning("Transaction {} did not complete properly, the recorded segments are: {}", name, segments)
+      context.stop(self)
+  }
+
+  def finishTracing(): Unit = {
+    reporting ! UowTrace(name.getOrElse("UNKNOWN"), segments)
+    println("Recorded Segments: " + segments)
+    context.stop(self)
+  }
+}
+
+object UowTraceAggregator {
+  def props(reporting: ActorRef, aggregationTimeout: Duration) = Props(classOf[UowTraceAggregator], reporting, aggregationTimeout)
+}
+\ No newline at end of file
diff --git a/kamon-core/src/main/scala/kamon/trace/context/TracingAwareContext.scala b/kamon-core/src/main/scala/kamon/trace/context/TracingAwareContext.scala
new file mode 100644
index 00000000..3766dd22
--- /dev/null
+++ b/kamon-core/src/main/scala/kamon/trace/context/TracingAwareContext.scala
@@ -0,0 +1,8 @@
+package kamon.trace.context
+
+import kamon.TraceContext
+
+trait TracingAwareContext {
+  def traceContext: Option[TraceContext]
+  def timestamp: Long
+}
+\ No newline at end of file