Settled with a @Aspectj style for instrumentation on runnables

3 files changed, 61 insertions, 54 deletions

diff --git a/src/main/scala/kamon/instrumentation/PromiseCompletingRunnableInstrumentation.scala b/src/main/scala/kamon/instrumentation/PromiseCompletingRunnableInstrumentation.scala
new file mode 100644
index 00000000..ce19a7e6
--- /dev/null
+++ b/src/main/scala/kamon/instrumentation/PromiseCompletingRunnableInstrumentation.scala
@@ -0,0 +1,61 @@
+package kamon.instrumentation
+
+import org.aspectj.lang.annotation._
+import kamon.TraceContext
+import org.aspectj.lang.ProceedingJoinPoint
+import scala.Some
+
+/**
+ *  Marker interface, just to make sure we don't instrument all the Runnables in the classpath.
+ */
+trait TraceContextAwareRunnable extends Runnable {}
+
+
+@Aspect("perthis(instrumentedRunnableCreation())")
+class PromiseCompletingRunnableInstrumentation {
+
+  /**
+   *  These are the Runnables that need to be instrumented and make the TraceContext available
+   *  while their run method is executed.
+   */
+  @DeclareMixin("scala.concurrent.impl.CallbackRunnable || scala.concurrent.impl.Future.PromiseCompletingRunnable")
+  def onCompleteCallbacksRunnable: TraceContextAwareRunnable = null
+
+
+  /**
+   *  Pointcuts
+   */
+
+  @Pointcut("execution(kamon.instrumentation.TraceContextAwareRunnable+.new(..))")
+  def instrumentedRunnableCreation(): Unit = {}
+
+  @Pointcut("execution(* kamon.instrumentation.TraceContextAwareRunnable.run())")
+  def runnableExecution() = {}
+
+
+  /**
+   *  Aspect members
+   */
+
+  private val traceContext = TraceContext.current
+
+
+  /**
+   *  Advices
+   */
+
+  @Around("runnableExecution()")
+  def around(pjp: ProceedingJoinPoint) = {
+    import pjp._
+
+    traceContext match {
+      case Some(ctx) => {
+        TraceContext.set(ctx)
+        proceed()
+        TraceContext.clear
+      }
+      case None => proceed()
+    }
+  }
+
+}
diff --git a/src/main/scala/kamon/instrumentation/PromiseInstrumentation.scala b/src/main/scala/kamon/instrumentation/PromiseInstrumentation.scala
deleted file mode 100644
index a39bc4ad..00000000
--- a/src/main/scala/kamon/instrumentation/PromiseInstrumentation.scala
+++ /dev/null
@@ -1,47 +0,0 @@
-package kamon.instrumentation
-
-import org.aspectj.lang.annotation._
-import kamon.TraceContext
-import scala.util.Try
-import scala.concurrent.ExecutionContext
-import org.aspectj.lang.ProceedingJoinPoint
-import scala.Some
-
-@Aspect("perthis(promiseCreation())")
-class PromiseInstrumentation {
-
-  private var traceContext: Option[TraceContext] = None
-
-  @Pointcut("execution(scala.concurrent.impl.Promise.DefaultPromise.new(..))")
-  def promiseCreation(): Unit = {}
-
-  @Before("promiseCreation()")
-  def catchTheTraceContext = {
-    TraceContext.current match {
-      case Some(ctx) => traceContext = Some(ctx.fork)
-      case None      => traceContext = None
-    }
-  }
-
-  @Pointcut("execution(* scala.concurrent.Future.onComplete(..)) && args(func, executor)")
-  def registeringOnCompleteCallback(func: Try[Any] => Any, executor: ExecutionContext) = {}
-
-  @Around("registeringOnCompleteCallback(func, executor)")
-  def around(pjp: ProceedingJoinPoint, func: Try[Any] => Any, executor: ExecutionContext) = {
-    import pjp._
-
-    val wrappedFunction = traceContext match {
-      case Some(ctx) => (tryResult: Try[Any]) => {
-        TraceContext.set(ctx)
-        val result = func(tryResult)
-        TraceContext.clear
-
-        result
-      }
-      case None => func
-    }
-
-    proceed(getArgs.updated(0, wrappedFunction))
-  }
-
-}
diff --git a/src/main/scala/kamon/instrumentation/TraceContextHolder.scala b/src/main/scala/kamon/instrumentation/TraceContextHolder.scala
deleted file mode 100644
index a9a66d63..00000000
--- a/src/main/scala/kamon/instrumentation/TraceContextHolder.scala
+++ /dev/null
@@ -1,7 +0,0 @@
-package kamon.instrumentation;
-import kamon.TraceContext
-import scala.Option
-
-class TraceContextHolder {
-    val context:Option[TraceContext] = TraceContext.current
-}
\ No newline at end of file