WIP: Instrumentation on Future.apply

2 files changed, 27 insertions, 26 deletions

diff --git a/src/main/scala/kamon/instrumentation/PromiseInstrumentation.scala b/src/main/scala/kamon/instrumentation/PromiseInstrumentation.scala
index 45e9c414..087e43bf 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 {
@@ -43,4 +44,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
+    }
+  }
+
 }
diff --git a/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala b/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala
deleted file mode 100644
index 2eb8d07a..00000000
--- a/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala
+++ /dev/null
@@ -1,25 +0,0 @@
-package kamon.instrumentation
-
-import scala.concurrent.{Await, Future}
-import org.specs2.mutable.Specification
-import scala.concurrent.ExecutionContext.Implicits.global
-import scala.concurrent.duration.{FiniteDuration, DurationLong}
-import org.specs2.time.{ Duration => SpecsDuration }
-
-
-class FutureInstrumentationSpec extends Specification {
-  import Await.result
-  implicit def specsDuration2Akka(duration: SpecsDuration): FiniteDuration = new DurationLong(duration.inMillis).millis
-
-  "a instrumented Future" should {
-    "preserve the transaction context available during the future creation" in {
-
-    }
-
-    "use the same context available at creation when executing the onComplete callback" in {
-      val future = Future { "hello" }
-
-      result(future, 100 millis) === "hello"
-    }
-  }
-}