Put FutureTracing instrumentation in place, simpler and more explicit test

2 files changed, 6 insertions, 10 deletions

diff --git a/kamon-trace/src/main/scala/kamon/trace/Trace.scala b/kamon-trace/src/main/scala/kamon/trace/Trace.scala
index 0ef438d4..fbd4ea8b 100644
--- a/kamon-trace/src/main/scala/kamon/trace/Trace.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/Trace.scala
@@ -35,6 +35,8 @@ object Trace extends ExtensionId[TraceExtension] with ExtensionIdProvider {
     ctx
   }
 
+  def withValue[T](ctx: Option[TraceContext])(thunk: => T): T = traceContext.withValue(ctx)(thunk)
+
   def transformContext(f: TraceContext => TraceContext): Unit = {
     context.map(f).foreach(set(_))
   }
diff --git a/kamon-trace/src/main/scala/kamon/trace/instrumentation/RunnableTracing.scala b/kamon-trace/src/main/scala/kamon/trace/instrumentation/FutureTracing.scala
index 5c16a648..10daa2da 100644
--- a/kamon-trace/src/main/scala/kamon/trace/instrumentation/RunnableTracing.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/instrumentation/FutureTracing.scala
@@ -5,7 +5,7 @@ import org.aspectj.lang.ProceedingJoinPoint
 import kamon.trace.{ContextAware, TraceContext, Trace}
 
 @Aspect
-class RunnableTracing {
+class FutureTracing {
 
   /**
    *  These are the Runnables that need to be instrumented and make the TraceContext available
@@ -15,16 +15,11 @@ class RunnableTracing {
   def onCompleteCallbacksRunnable: ContextAware = ContextAware.default
 
 
-  /**
-   *  Pointcuts
-   */
-
   @Pointcut("execution(kamon.trace.ContextAware+.new(..)) && this(runnable)")
   def instrumentedRunnableCreation(runnable: ContextAware): Unit = {}
 
   @Pointcut("execution(* kamon.trace.ContextAware+.run()) && this(runnable)")
-  def runnableExecution(runnable: ContextAware) = {}
-
+  def futureRunnableExecution(runnable: ContextAware) = {}
 
 
   @After("instrumentedRunnableCreation(runnable)")
@@ -33,12 +28,11 @@ class RunnableTracing {
     runnable.traceContext
   }
 
-
-  @Around("runnableExecution(runnable)")
+  @Around("futureRunnableExecution(runnable)")
   def around(pjp: ProceedingJoinPoint, runnable: ContextAware): Any = {
     import pjp._
 
-    Trace.traceContext.withValue(runnable.traceContext) {
+    Trace.withValue(runnable.traceContext) {
       proceed()
     }
   }