basic separation of concerns between sub-projects

5 files changed, 121 insertions, 55 deletions

diff --git a/kamon-trace/src/main/scala/kamon/trace/Tracer.scala b/kamon-trace/src/main/scala/kamon/trace/Trace.scala
index 4ea89850..232b7420 100644
--- a/kamon-trace/src/main/scala/kamon/trace/Tracer.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/Trace.scala
@@ -6,6 +6,7 @@ import akka.actor._
 import scala.Some
 import kamon.trace.Trace.Register
 import scala.concurrent.duration._
+import java.util.concurrent.atomic.AtomicLong
 
 object Trace extends ExtensionId[TraceExtension] with ExtensionIdProvider {
   def lookup(): ExtensionId[_ <: Extension] = Trace
@@ -14,10 +15,31 @@ object Trace extends ExtensionId[TraceExtension] with ExtensionIdProvider {
 
   /*** Protocol */
   case object Register
+
+
+
+  /** User API */
+  private[trace] val traceContext = new DynamicVariable[Option[TraceContext]](None)
+  private[trace] val tranid = new AtomicLong()
+
+
+  def context() = traceContext.value
+  def set(ctx: TraceContext) = traceContext.value = Some(ctx)
+
+  def start(name: String)(implicit system: ActorSystem) = set(newTraceContext)
+
+  def finish(): Option[TraceContext] = {
+    val ctx = context()
+    ctx.map(_.finish)
+    ctx
+  }
+
+  // TODO: FIX
+  def newTraceContext()(implicit system: ActorSystem): TraceContext = TraceContext(Kamon(Trace), tranid.getAndIncrement)
 }
 
 class TraceExtension(system: ExtendedActorSystem) extends Kamon.Extension {
-  def manager: ActorRef = ???
+  def manager: ActorRef = system.actorOf(Props[TraceManager])
 }
 
 class TraceManager extends Actor {
@@ -35,15 +57,3 @@ class TraceManager extends Actor {
       listeners foreach(_ ! trace)
   }
 }
-
-
-object Tracer {
-  val traceContext = new DynamicVariable[Option[TraceContext]](None)
-
-
-  def context() = traceContext.value
-  def set(ctx: TraceContext) = traceContext.value = Some(ctx)
-
-  def start = set(newTraceContext)
-  def newTraceContext(): TraceContext = TraceContext()(Kamon.actorSystem)
-}
diff --git a/kamon-trace/src/main/scala/kamon/trace/TraceContext.scala b/kamon-trace/src/main/scala/kamon/trace/TraceContext.scala
index c3f1f2c2..f8491c12 100644
--- a/kamon-trace/src/main/scala/kamon/trace/TraceContext.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/TraceContext.scala
@@ -4,56 +4,32 @@ import java.util.UUID
 import akka.actor._
 import java.util.concurrent.atomic.AtomicLong
 import scala.concurrent.duration._
+import kamon.Kamon
+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.
-case class TraceContext(id: Long, tracer: ActorRef, uow: String = "", userContext: Option[Any] = None)
+protected[kamon] case class TraceContext(private val collector: ActorRef, id: Long, uow: String = "", userContext: Option[Any] = None) {
+  collector ! Start(id)
 
-object TraceContext {
-
-  def apply()(implicit system: ActorSystem) =  {
-    val n = traceIdCounter.incrementAndGet()
-    val actor = system.actorOf(UowTraceAggregator.props(reporter, 30 seconds), s"tracer-${n}")
-    actor ! Start()
-
-    new TraceContext(n, actor) // TODO: Move to a kamon specific supervisor, like /user/kamon/tracer
+  def finish: Unit = {
+    collector ! Finish(id)
   }
-}
-
 
 
-class TraceAccumulator extends Actor {
-  def receive = {
-    case a  => println("Trace Accumulated: "+a)
-  }
 }
 
 
-trait TraceEntry
-case class CodeBlockExecutionTime(name: String, begin: Long, end: Long) extends TraceEntry
-case class TransactionTrace(id: UUID, start: Long, end: Long, entries: Seq[TraceEntry])
-
-object Collector
-
-trait TraceEntryStorage {
-  def store(entry: TraceEntry): Boolean
+trait ContextAware {
+  def traceContext: Option[TraceContext]
 }
 
-class TransactionContext(val id: UUID, private val storage: TraceEntryStorage) {
-  def store(entry: TraceEntry) = storage.store(entry)
-}
-
-object ThreadLocalTraceEntryStorage extends TraceEntryStorage {
-
-  private val storage = new ThreadLocal[List[TraceEntry]] {
-    override def initialValue(): List[TraceEntry] = Nil
-  }
-
-  def update(f: List[TraceEntry] => List[TraceEntry]) = storage set f(storage.get)
-
-  def store(entry: TraceEntry): Boolean = {
-    update(entry :: _)
-    true
+object ContextAware {
+  def default: ContextAware = new ContextAware {
+    val traceContext: Option[TraceContext] = Trace.context()
   }
 }
 
-
+trait TimedContextAware {
+  def timestamp: Long
+  def traceContext: Option[TraceContext]
+}
diff --git a/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorLoggingInstrumentation.scala b/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorLoggingInstrumentation.scala
new file mode 100644
index 00000000..77993cdd
--- /dev/null
+++ b/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorLoggingInstrumentation.scala
@@ -0,0 +1,31 @@
+package kamon.trace.instrumentation
+
+import org.aspectj.lang.annotation.{Around, Pointcut, DeclareMixin, Aspect}
+import org.aspectj.lang.ProceedingJoinPoint
+import org.slf4j.MDC
+import kamon.trace.{TraceContext, ContextAware, Trace}
+
+@Aspect
+class ActorLoggingInstrumentation {
+
+
+  @DeclareMixin("akka.event.Logging.LogEvent+")
+  def traceContextMixin: ContextAware = new ContextAware {
+    def traceContext: Option[TraceContext] = Trace.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-trace/src/main/scala/kamon/trace/instrumentation/ActorRefTellInstrumentation.scala b/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorRefTellInstrumentation.scala
new file mode 100644
index 00000000..3caba77c
--- /dev/null
+++ b/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorRefTellInstrumentation.scala
@@ -0,0 +1,49 @@
+package kamon.trace.instrumentation
+
+import org.aspectj.lang.annotation._
+import org.aspectj.lang.ProceedingJoinPoint
+import akka.actor.{Props, ActorSystem, ActorRef}
+import akka.dispatch.{Envelope, MessageDispatcher}
+import com.codahale.metrics.Timer
+import kamon.trace.{ContextAware, TraceContext, Trace}
+
+case class TraceableMessage(traceContext: Option[TraceContext], message: Any, timer: Timer.Context)
+case class DefaultTracingAwareEnvelopeContext(traceContext: Option[TraceContext] = Trace.traceContext.value, timestamp: Long = System.nanoTime) extends ContextAware
+
+@Aspect
+class ActorCellInvokeInstrumentation {
+
+  @Pointcut("execution(akka.actor.ActorCell.new(..)) && args(system, ref, props, dispatcher, parent)")
+  def actorCellCreation(system: ActorSystem, ref: ActorRef, props: Props, dispatcher: MessageDispatcher, parent: ActorRef): Unit = {}
+
+  @Pointcut("(execution(* akka.actor.ActorCell.invoke(*)) || execution(* akka.routing.RoutedActorCell.sendMessage(*))) && args(envelope)")
+  def invokingActorBehaviourAtActorCell(envelope: Envelope) = {}
+
+  @Around("invokingActorBehaviourAtActorCell(envelope)")
+  def around(pjp: ProceedingJoinPoint, envelope: Envelope): Unit = {
+    //safe cast
+    val msgContext = envelope.asInstanceOf[ContextAware].traceContext
+
+    Trace.traceContext.withValue(msgContext) {
+      pjp.proceed()
+    }
+  }
+}
+
+@Aspect
+class EnvelopeTracingContext {
+
+  @DeclareMixin("akka.dispatch.Envelope")
+  def mixin: ContextAware =  new ContextAware {
+    val traceContext: Option[TraceContext] = Trace.context()
+  }
+
+  @Pointcut("execution(akka.dispatch.ContextAware.new(..)) && this(ctx)")
+  def requestRecordInit(ctx: ContextAware): Unit = {}
+
+  @After("requestRecordInit(ctx)")
+  def whenCreatedRequestRecord(ctx: ContextAware): Unit = {
+    // Necessary to force the initialization of TracingAwareRequestContext at the moment of creation.
+    ctx.traceContext
+  }
+}
diff --git a/kamon-trace/src/main/scala/kamon/trace/instrumentation/RunnableTracing.scala b/kamon-trace/src/main/scala/kamon/trace/instrumentation/RunnableTracing.scala
index 236fd4fc..3e5a7cce 100644
--- a/kamon-trace/src/main/scala/kamon/trace/instrumentation/RunnableTracing.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/instrumentation/RunnableTracing.scala
@@ -2,7 +2,7 @@ package kamon.trace.instrumentation
 
 import org.aspectj.lang.annotation._
 import org.aspectj.lang.ProceedingJoinPoint
-import kamon.trace.TraceContext
+import kamon.trace.{TraceContext, Trace}
 
 @Aspect
 class RunnableTracing {
@@ -13,7 +13,7 @@ class RunnableTracing {
    */
   @DeclareMixin("scala.concurrent.impl.CallbackRunnable || scala.concurrent.impl.Future.PromiseCompletingRunnable")
   def onCompleteCallbacksRunnable: TraceContextAwareRunnable = new TraceContextAwareRunnable {
-    val traceContext: Option[TraceContext] = Tracer.traceContext.value
+    val traceContext: Option[TraceContext] = Trace.traceContext.value
   }
 
 
@@ -40,7 +40,7 @@ class RunnableTracing {
   def around(pjp: ProceedingJoinPoint, runnable: TraceContextAwareRunnable): Any = {
     import pjp._
 
-    Tracer.traceContext.withValue(runnable.traceContext) {
+    Trace.traceContext.withValue(runnable.traceContext) {
       proceed()
     }
   }