wip

4 files changed, 11 insertions, 23 deletions

diff --git a/kamon-trace/src/main/scala/kamon/trace/Trace.scala b/kamon-trace/src/main/scala/kamon/trace/Trace.scala
index 6e01ad26..1956509c 100644
--- a/kamon-trace/src/main/scala/kamon/trace/Trace.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/Trace.scala
@@ -60,7 +60,6 @@ class TraceManager extends Actor with ActorLogging {
       context.child(tracerName).getOrElse(newTracer(tracerName)) ! segment
 
     case trace: UowTrace =>
-      println("Delivering a trace to: " + listeners)
       listeners foreach(_ ! trace)
   }
 
diff --git a/kamon-trace/src/main/scala/kamon/trace/UowTracing.scala b/kamon-trace/src/main/scala/kamon/trace/UowTracing.scala
index 009a6da2..273e0ca1 100644
--- a/kamon-trace/src/main/scala/kamon/trace/UowTracing.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/UowTracing.scala
@@ -41,6 +41,7 @@ class UowTraceAggregator(reporting: ActorRef, aggregationTimeout: Duration) exte
   def receive = {
     case start: Start =>
       this.start = start.timestamp
+      segments = segments :+ start;
       name = start.name
     case finish: Finish       =>
       end = finish.timestamp
@@ -58,7 +59,6 @@ class UowTraceAggregator(reporting: ActorRef, aggregationTimeout: Duration) exte
 
   def finishTracing(): Unit = {
     reporting ! UowTrace(name, "", start, end, segments)
-    println("Recorded Segments: " + segments)
     context.stop(self)
   }
 }
diff --git a/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorRefTellInstrumentation.scala b/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorRefTellInstrumentation.scala
index 3caba77c..057e339d 100644
--- a/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorRefTellInstrumentation.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorRefTellInstrumentation.scala
@@ -34,11 +34,9 @@ class ActorCellInvokeInstrumentation {
 class EnvelopeTracingContext {
 
   @DeclareMixin("akka.dispatch.Envelope")
-  def mixin: ContextAware =  new ContextAware {
-    val traceContext: Option[TraceContext] = Trace.context()
-  }
+  def mixin: ContextAware = ContextAware.default
 
-  @Pointcut("execution(akka.dispatch.ContextAware.new(..)) && this(ctx)")
+  @Pointcut("execution(akka.dispatch.Envelope.new(..)) && this(ctx)")
   def requestRecordInit(ctx: ContextAware): Unit = {}
 
   @After("requestRecordInit(ctx)")
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 3e5a7cce..5c16a648 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, Trace}
+import kamon.trace.{ContextAware, TraceContext, Trace}
 
 @Aspect
 class RunnableTracing {
@@ -12,32 +12,30 @@ class RunnableTracing {
    *  while their run method is executed.
    */
   @DeclareMixin("scala.concurrent.impl.CallbackRunnable || scala.concurrent.impl.Future.PromiseCompletingRunnable")
-  def onCompleteCallbacksRunnable: TraceContextAwareRunnable = new TraceContextAwareRunnable {
-    val traceContext: Option[TraceContext] = Trace.traceContext.value
-  }
+  def onCompleteCallbacksRunnable: ContextAware = ContextAware.default
 
 
   /**
    *  Pointcuts
    */
 
-  @Pointcut("execution(kamon.instrumentation.TraceContextAwareRunnable+.new(..)) && this(runnable)")
-  def instrumentedRunnableCreation(runnable: TraceContextAwareRunnable): Unit = {}
+  @Pointcut("execution(kamon.trace.ContextAware+.new(..)) && this(runnable)")
+  def instrumentedRunnableCreation(runnable: ContextAware): Unit = {}
 
-  @Pointcut("execution(* kamon.instrumentation.TraceContextAwareRunnable+.run()) && this(runnable)")
-  def runnableExecution(runnable: TraceContextAwareRunnable) = {}
+  @Pointcut("execution(* kamon.trace.ContextAware+.run()) && this(runnable)")
+  def runnableExecution(runnable: ContextAware) = {}
 
 
 
   @After("instrumentedRunnableCreation(runnable)")
-  def beforeCreation(runnable: TraceContextAwareRunnable): Unit = {
+  def beforeCreation(runnable: ContextAware): Unit = {
     // Force traceContext initialization.
     runnable.traceContext
   }
 
 
   @Around("runnableExecution(runnable)")
-  def around(pjp: ProceedingJoinPoint, runnable: TraceContextAwareRunnable): Any = {
+  def around(pjp: ProceedingJoinPoint, runnable: ContextAware): Any = {
     import pjp._
 
     Trace.traceContext.withValue(runnable.traceContext) {
@@ -45,11 +43,4 @@ class RunnableTracing {
     }
   }
 
-}
-
-/**
- *  Marker interface, just to make sure we don't instrument all the Runnables in the classpath.
- */
-trait TraceContextAwareRunnable {
-  def traceContext: Option[TraceContext]
 }
\ No newline at end of file