diff options
Diffstat (limited to 'src/main/scala/kamon/instrumentation/PromiseInstrumentation.scala')
-rw-r--r-- | src/main/scala/kamon/instrumentation/PromiseInstrumentation.scala | 28 |
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 + } + } + } |