wip

2 files changed, 10 insertions, 21 deletions

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