Some sort of basic logging with UOW

2 files changed, 104 insertions, 11 deletions

diff --git a/kamon-uow/src/main/scala/kamon/logging/UowActorLogging.scala b/kamon-uow/src/main/scala/kamon/logging/UowActorLogging.scala
index 64c5241a..784bd674 100644
--- a/kamon-uow/src/main/scala/kamon/logging/UowActorLogging.scala
+++ b/kamon-uow/src/main/scala/kamon/logging/UowActorLogging.scala
@@ -1,24 +1,115 @@
 package kamon.logging
 
-import akka.actor.Actor
-import kamon.{Tracer, Kamon}
-import akka.event.{BusLogging, LogSource, LoggingAdapter}
+import akka.actor.{ActorSystem, Actor}
+import kamon.{Tracer}
+import akka.event.{LoggingBus, LogSource, LoggingAdapter}
+import akka.event.Logging._
+import akka.event.slf4j.{Logger, SLF4JLogging}
+import akka.event.Logging.Info
+import akka.event.Logging.Warning
+import akka.event.Logging.Error
+import akka.event.Logging.Debug
+import org.slf4j.MDC
+import akka.util.Helpers
 
 trait UowActorLogging {
   this: Actor =>
 
   val log = {
     val (str, clazz) = LogSource(this, context.system)
-    new BusLogging(context.system.eventStream, str, clazz) with UowLoggingAdapter
+    new ExtendedBusLogging(context.system.eventStream, str, clazz)
   }
+}
+
+trait UowLogging {
+  self: Any =>
+  def system: ActorSystem
 
+  val log = {
+    val (str, clazz) = LogSource(self.getClass, system)
+    new ExtendedBusLogging(system.eventStream, str, clazz)
+  }
 }
 
-trait UowLoggingAdapter extends LoggingAdapter {
+class ExtendedBusLogging(val bus: LoggingBus, val logSource: String, val logClass: Class[_]) extends LoggingAdapter {
+
+  import akka.event.Logging._
+
+  def isErrorEnabled = bus.logLevel >= ErrorLevel
+  def isWarningEnabled = bus.logLevel >= WarningLevel
+  def isInfoEnabled = bus.logLevel >= InfoLevel
+  def isDebugEnabled = bus.logLevel >= DebugLevel
+
   def currentUow: String = Tracer.context().flatMap(_.userContext).map(_.toString).getOrElse("")
-  protected abstract override def notifyWarning(message: String)  = super.notifyWarning(s"[${currentUow}] - $message")
-  protected abstract override def notifyInfo(message: String)     = super.notifyInfo(   s"[${currentUow}] - $message")
-  protected abstract override def notifyDebug(message: String)    = super.notifyDebug(  s"[${currentUow}] - $message")
-  protected abstract override def notifyError(message: String)    = super.notifyError(  s"[${currentUow}] - $message")
-  protected abstract override def notifyError(cause: Throwable, message: String) = super.notifyError(cause, s"[${currentUow}] - $message")
+  def extras = Map("uow" -> currentUow)
+
+  protected def notifyError(message: String): Unit = bus.publish(Error(logSource, logClass, RichLogEvent(message, extras)))
+  protected def notifyError(cause: Throwable, message: String): Unit = bus.publish(Error(cause, logSource, logClass, RichLogEvent(message, extras)))
+  protected def notifyWarning(message: String): Unit = bus.publish(Warning(logSource, logClass, RichLogEvent(message, extras)))
+  protected def notifyInfo(message: String): Unit = bus.publish(Info(logSource, logClass, RichLogEvent(message, extras)))
+  protected def notifyDebug(message: String): Unit = bus.publish(Debug(logSource, logClass, RichLogEvent(message, extras)))
+}
+
+case class RichLogEvent(message: String, extras: Map[String, Any])
+
+
+
+class ExtendedSlf4jLogger extends Actor with SLF4JLogging {
+
+  val mdcThreadAttributeName = "sourceThread"
+  val mdcAkkaSourceAttributeName = "akkaSource"
+  val mdcAkkaTimestamp = "akkaTimestamp"
+
+  def receive = {
+
+    case event @ Error(cause, logSource, logClass, message) ⇒
+      withMdc(logSource, event) {
+        cause match {
+          case Error.NoCause | null ⇒ withRichEventProcessing(message) { Logger(logClass, logSource).error(if (message != null) message.toString else null) }
+          case cause                ⇒ withRichEventProcessing(message) { Logger(logClass, logSource).error(if (message != null) message.toString else cause.getLocalizedMessage, cause) }
+        }
+      }
+
+    case event @ Warning(logSource, logClass, message) ⇒
+      withMdc(logSource, event) { withRichEventProcessing(message) { Logger(logClass, logSource).warn("{}", message.asInstanceOf[AnyRef]) } }
+
+    case event @ Info(logSource, logClass, message) ⇒
+      withMdc(logSource, event) { withRichEventProcessing(message) { Logger(logClass, logSource).info("{}", message.asInstanceOf[AnyRef]) } }
+
+    case event @ Debug(logSource, logClass, message) ⇒
+      withMdc(logSource, event) { withRichEventProcessing(message) { Logger(logClass, logSource).debug("{}", message.asInstanceOf[AnyRef]) } }
+
+    case InitializeLogger(_) ⇒
+      log.info("Slf4jLogger started")
+      sender ! LoggerInitialized
+  }
+
+  def withRichEventProcessing(message: Any)(delegate: => Unit): Unit = message match {
+    case RichLogEvent(event, extras) => {
+      extras.foreach { case (k, v) => MDC.put(k, v.toString) }
+      delegate
+      MDC.clear()
+    }
+    case _ => delegate
+  }
+
+  @inline
+  final def withMdc(logSource: String, logEvent: LogEvent)(logStatement: ⇒ Unit) {
+    MDC.put(mdcAkkaSourceAttributeName, logSource)
+    MDC.put(mdcThreadAttributeName, logEvent.thread.getName)
+    MDC.put(mdcAkkaTimestamp, formatTimestamp(logEvent.timestamp))
+    try logStatement finally {
+      MDC.remove(mdcAkkaSourceAttributeName)
+      MDC.remove(mdcThreadAttributeName)
+      MDC.remove(mdcAkkaTimestamp)
+    }
+  }
+
+  /**
+   * Override this method to provide a differently formatted timestamp
+   * @param timestamp a "currentTimeMillis"-obtained timestamp
+   * @return the given timestamp as a UTC String
+   */
+  protected def formatTimestamp(timestamp: Long): String =
+    Helpers.currentTimeMillisToUTCString(timestamp)
 }
diff --git a/kamon-uow/src/main/scala/kamon/logging/UowDirectives.scala b/kamon-uow/src/main/scala/kamon/logging/UowDirectives.scala
index 58473f13..6d38f3d7 100644
--- a/kamon-uow/src/main/scala/kamon/logging/UowDirectives.scala
+++ b/kamon-uow/src/main/scala/kamon/logging/UowDirectives.scala
@@ -9,7 +9,9 @@ import kamon.{Tracer, Kamon}
 
 trait UowDirectives extends BasicDirectives {
   def uow: Directive0 = mapRequest { request =>
-    val generatedUow = Some(UowDirectives.newUow)
+    val uowHeader = request.headers.find(_.name == "X-UOW")
+
+    val generatedUow = uowHeader.map(_.value).orElse(Some(UowDirectives.newUow))
     println("Generated UOW: "+generatedUow)
     Tracer.set(Tracer.newTraceContext().copy(userContext = generatedUow))