Merge branch 'master' of github.com:dpsoft/Kamon

1 files changed, 27 insertions, 1 deletions

diff --git a/src/main/scala/kamon/instrumentation/PromiseInstrumentation.scala b/src/main/scala/kamon/instrumentation/PromiseInstrumentation.scala
index b87abe44..f1aec404 100644
--- a/src/main/scala/kamon/instrumentation/PromiseInstrumentation.scala
+++ b/src/main/scala/kamon/instrumentation/PromiseInstrumentation.scala
@@ -1,10 +1,11 @@
 package kamon.instrumentation
 
-import org.aspectj.lang.annotation.{Around, Before, Pointcut, Aspect}
+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 {
@@ -42,4 +43,29 @@ class PromiseInstrumentation {
 
     proceed(getArgs.updated(0, wrappedFunction))
   }
+
+  @Pointcut("execution(* scala.concurrent.impl.Future$.apply(..)) && args(body, executor)")
+  def registeringApplyOnFuture(body: () => Any, executor: ExecutionContext) = {}
+
+  @Around("registeringApplyOnFuture(body, executor)")
+  def aroundApplyOnFuture(pjp:ProceedingJoinPoint, body: () => Any, executor: ExecutionContext) = {
+    import pjp._
+
+    val wrappedBody = wrapFutureBody(traceContext)(body)
+    proceed(getArgs.updated(0, wrappedBody))
+  }
+
+  def wrapFutureBody[A](ctx:Option[TraceContext])(block: => A) : A = {
+    ctx match {
+        case Some(ctx) => {
+          println("Wrapping body")
+          TraceContext.set(ctx)
+          val result = block
+          TraceContext.clear
+          result
+        }
+        case None => block
+    }
+  }
+
 }