Settled with a @Aspectj style for instrumentation on runnables

5 files changed, 63 insertions, 85 deletions

diff --git a/src/main/aspectj/kamon/instrumentation/PromiseCompletingRunnableInstrumentation.aj b/src/main/aspectj/kamon/instrumentation/PromiseCompletingRunnableInstrumentation.aj
deleted file mode 100644
index 5ae26d75..00000000
--- a/src/main/aspectj/kamon/instrumentation/PromiseCompletingRunnableInstrumentation.aj
+++ /dev/null
@@ -1,28 +0,0 @@
-package kamon.instrumentation;
-
-import kamon.TraceContext;
-import scala.concurrent.impl.Future;
-import scala.Option;
-
-privileged public aspect PromiseCompletingRunnableInstrumentation {
-
-    declare parents : Future.PromiseCompletingRunnable extends TraceContextHolder;
-
-    pointcut run(scala.concurrent.impl.Future.PromiseCompletingRunnable runnable)
-            : execution(* scala.concurrent.impl.Future.PromiseCompletingRunnable.run()) && this(runnable);
-
-    void around(Object runnable)
-        : run(runnable) {
-
-        TraceContextHolder contextHolder = (TraceContextHolder) runnable;
-
-        if(contextHolder.context().isDefined()) {
-            TraceContext.set(contextHolder.context().get());
-            proceed(contextHolder);
-            TraceContext.clear();
-
-        } else {
-            proceed(contextHolder);
-        }
-    }
-}
diff --git a/src/main/resources/META-INF/aop.xml b/src/main/resources/META-INF/aop.xml
index cafba5ee..d97a00ea 100644
--- a/src/main/resources/META-INF/aop.xml
+++ b/src/main/resources/META-INF/aop.xml
@@ -1,9 +1,9 @@
 <!DOCTYPE aspectj PUBLIC "-//AspectJ//DTD//EN" "http://www.eclipse.org/aspectj/dtd/aspectj.dtd">
 
 <aspectj>
-    <weaver options="-verbose -showWeaveInfo -nowarn">
+    <!--<weaver options="-verbose -showWeaveInfo">
        <dump within="*"/>
-    </weaver>
+    </weaver>-->
 
     <aspects>
 
@@ -12,7 +12,6 @@
         <aspect name="akka.PoolMonitorAspect"/>-->
         <aspect name="akka.instrumentation.ActorRefTellInstrumentation"/>
         <aspect name="akka.instrumentation.ActorCellInvokeInstrumentation"/>
-        <aspect name="kamon.instrumentation.PromiseInstrumentation"/>
         <aspect name="kamon.instrumentation.PromiseCompletingRunnableInstrumentation" />
 
         <include within="*"/>
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