Simplify logging instrumentation and integrate it with kamon-core, kamon-ouw is gone.

5 files changed, 72 insertions, 14 deletions

diff --git a/kamon-core/src/main/scala/kamon/TraceContext.scala b/kamon-core/src/main/scala/kamon/TraceContext.scala
index 0dfc1630..155b7760 100644
--- a/kamon-core/src/main/scala/kamon/TraceContext.scala
+++ b/kamon-core/src/main/scala/kamon/TraceContext.scala
@@ -9,13 +9,7 @@ import kamon.newrelic.NewRelicReporting
 import kamon.trace.UowTracing.Start
 
 // TODO: Decide if we need or not an ID, generating it takes time and it doesn't seem necessary.
-case class TraceContext(id: Long, entries: ActorRef, userContext: Option[Any] = None) {
-  //implicit val timeout = Timeout(30, TimeUnit.SECONDS)
-  implicit val as = Kamon.actorSystem.dispatcher
-
-  def append(entry: TraceEntry) = entries ! entry
-  def close = entries ! "Close" // TODO type this thing!.
-}
+case class TraceContext(id: Long, tracer: ActorRef, uow: String = "", userContext: Option[Any] = None)
 
 object TraceContext {
   val reporter = Kamon.actorSystem.actorOf(Props[NewRelicReporting])
diff --git a/kamon-core/src/main/scala/kamon/instrumentation/ActorLoggingInstrumentation.scala b/kamon-core/src/main/scala/kamon/instrumentation/ActorLoggingInstrumentation.scala
new file mode 100644
index 00000000..3cf0c6fc
--- /dev/null
+++ b/kamon-core/src/main/scala/kamon/instrumentation/ActorLoggingInstrumentation.scala
@@ -0,0 +1,32 @@
+package kamon.instrumentation
+
+import org.aspectj.lang.annotation.{Around, Pointcut, DeclareMixin, Aspect}
+import kamon.{Tracer, TraceContext}
+import org.aspectj.lang.ProceedingJoinPoint
+import org.slf4j.MDC
+
+
+@Aspect
+class ActorLoggingInstrumentation {
+
+
+  @DeclareMixin("akka.event.Logging.LogEvent+")
+  def traceContextMixin: ContextAware = new ContextAware {
+    def traceContext: Option[TraceContext] = Tracer.context()
+  }
+
+  @Pointcut("execution(* akka.event.slf4j.Slf4jLogger.withMdc(..)) && args(logSource, logEvent, logStatement)")
+  def withMdcInvocation(logSource: String, logEvent: ContextAware, logStatement: () => _): Unit = {}
+
+  @Around("withMdcInvocation(logSource, logEvent, logStatement)")
+  def putTraceContextInMDC(pjp: ProceedingJoinPoint, logSource: String, logEvent: ContextAware, logStatement: () => _): Unit = {
+    logEvent.traceContext match {
+      case Some(ctx) =>
+        MDC.put("uow", ctx.uow)
+        pjp.proceed()
+        MDC.remove("uow")
+
+      case None => pjp.proceed()
+    }
+  }
+}
diff --git a/kamon-core/src/main/scala/kamon/instrumentation/SprayServerInstrumentation.scala b/kamon-core/src/main/scala/kamon/instrumentation/SprayServerInstrumentation.scala
index 32eabe71..5117e7e7 100644
--- a/kamon-core/src/main/scala/kamon/instrumentation/SprayServerInstrumentation.scala
+++ b/kamon-core/src/main/scala/kamon/instrumentation/SprayServerInstrumentation.scala
@@ -32,7 +32,7 @@ class SprayServerInstrumentation {
     Tracer.start
     val discard = openRequest.asInstanceOf[ContextAware].traceContext
 
-    Tracer.context().map(_.entries ! Rename(request.uri.path.toString()))
+    Tracer.context().map(_.tracer ! Rename(request.uri.path.toString()))
   }
 
   @Pointcut("execution(* spray.can.server.OpenRequestComponent$DefaultOpenRequest.handleResponseEndAndReturnNextOpenRequest(..)) && target(openRequest)")
@@ -42,7 +42,7 @@ class SprayServerInstrumentation {
   def afterFinishingRequest(openRequest: OpenRequest): Unit = {
     val original = openRequest.asInstanceOf[ContextAware].traceContext
 
-    Tracer.context().map(_.entries ! Finish())
+    Tracer.context().map(_.tracer ! Finish())
 
     if(Tracer.context() != original) {
       println(s"OMG DIFFERENT Original: [${original}] - Came in: [${Tracer.context}]")
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/test/SimpleRequestProcessor.scala b/kamon-core/src/main/scala/test/SimpleRequestProcessor.scala
index 5b216b39..7d4cec52 100644
--- a/kamon-core/src/main/scala/test/SimpleRequestProcessor.scala
+++ b/kamon-core/src/main/scala/test/SimpleRequestProcessor.scala
@@ -6,8 +6,9 @@ import spray.routing.SimpleRoutingApp
 import akka.util.Timeout
 import spray.httpx.RequestBuilding
 import scala.concurrent.{Await, Future}
+import kamon.trace.UowDirectives
 
-object SimpleRequestProcessor extends App with SimpleRoutingApp with RequestBuilding {
+object SimpleRequestProcessor extends App with SimpleRoutingApp with RequestBuilding with UowDirectives {
   import scala.concurrent.duration._
   import spray.client.pipelining._
   import akka.pattern.ask
@@ -28,11 +29,13 @@ object SimpleRequestProcessor extends App with SimpleRoutingApp with RequestBuil
         }
       } ~
       path("reply" / Segment) { reqID =>
-        complete {
-          if (Tracer.context().isEmpty)
-            println("ROUTE NO CONTEXT")
+        uow {
+          complete {
+            if (Tracer.context().isEmpty)
+              println("ROUTE NO CONTEXT")
 
-          (replier ? reqID).mapTo[String]
+            (replier ? reqID).mapTo[String]
+          }
         }
       } ~
       path("ok") {
@@ -78,6 +81,7 @@ class Replier extends Actor with ActorLogging {
       if(Tracer.context.isEmpty)
         log.warning("PROCESSING A MESSAGE WITHOUT CONTEXT")
 
+      log.info("Processing at the Replier")
       sender ! anything
   }
 }