switched back to scalatest and created simple instrumentation for futures body

3 files changed, 42 insertions, 1 deletions

diff --git a/src/main/resources/META-INF/aop.xml b/src/main/resources/META-INF/aop.xml
index 99eee806..ae3914cd 100644
--- a/src/main/resources/META-INF/aop.xml
+++ b/src/main/resources/META-INF/aop.xml
@@ -11,6 +11,7 @@
         <aspect name="akka.instrumentation.ActorRefTellInstrumentation"/>
         <aspect name="akka.instrumentation.ActorCellInvokeInstrumentation"/>
         <aspect name="kamon.instrumentation.PromiseInstrumentation"/>
+        <aspect name="kamon.instrumentation.PromiseCompletingRunnableInstrumentation" />
 
         <include within="*"/>
         <exclude within="javax.*"/>
diff --git a/src/main/scala/kamon/instrumentation/PromiseCompletingRunnableInstrumentation.scala b/src/main/scala/kamon/instrumentation/PromiseCompletingRunnableInstrumentation.scala
new file mode 100644
index 00000000..dde0d857
--- /dev/null
+++ b/src/main/scala/kamon/instrumentation/PromiseCompletingRunnableInstrumentation.scala
@@ -0,0 +1,41 @@
+package kamon.instrumentation
+
+import org.aspectj.lang.annotation.{Around, Before, Pointcut, Aspect}
+import kamon.TraceContext
+import org.aspectj.lang.ProceedingJoinPoint
+
+@Aspect("perthis(promiseCompletingRunnableCreation())")
+class PromiseCompletingRunnableInstrumentation {
+
+  private var traceContext: Option[TraceContext] = None
+
+  @Pointcut("execution(scala.concurrent.impl.Future.PromiseCompletingRunnable.new(..))")
+  def promiseCompletingRunnableCreation(): Unit = {}
+
+  @Before("promiseCompletingRunnableCreation()")
+  def catchTheTraceContext = {
+    TraceContext.current match {
+      case Some(ctx) => traceContext = Some(ctx.fork)
+      case None      => traceContext = None
+    }
+  }
+
+
+  @Pointcut("execution(* scala.concurrent.impl.Future.PromiseCompletingRunnable.run())")
+  def runnableExecution() = {}
+
+  @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
index 45e9c414..b87abe44 100644
--- a/src/main/scala/kamon/instrumentation/PromiseInstrumentation.scala
+++ b/src/main/scala/kamon/instrumentation/PromiseInstrumentation.scala
@@ -16,7 +16,6 @@ class PromiseInstrumentation {
 
   @Before("promiseCreation()")
   def catchTheTraceContext = {
-    println(s"During promise creation the context is: ${TraceContext.current}")
     TraceContext.current match {
       case Some(ctx) => traceContext = Some(ctx.fork)
       case None      => traceContext = None