diff options
| -rw-r--r-- | kamon-core/src/main/scala/kamon/instrumentation/akka/AskPatternInstrumentation.scala | 31 |
1 files changed, 16 insertions, 15 deletions
diff --git a/kamon-core/src/main/scala/kamon/instrumentation/akka/AskPatternInstrumentation.scala b/kamon-core/src/main/scala/kamon/instrumentation/akka/AskPatternInstrumentation.scala index 93b53764..5b4fbbc8 100644 --- a/kamon-core/src/main/scala/kamon/instrumentation/akka/AskPatternInstrumentation.scala +++ b/kamon-core/src/main/scala/kamon/instrumentation/akka/AskPatternInstrumentation.scala @@ -18,7 +18,7 @@ package akka.kamon.instrumentation import kamon.Kamon import kamon.extension.akka.Akka -import kamon.trace.TraceContextAware +import kamon.trace.{ TraceContext, EmptyTraceContext, TraceContextAware } import akka.actor.{ ActorSystem, ActorRef } import akka.event.Logging.Warning import akka.pattern.AskTimeoutException @@ -40,21 +40,22 @@ class AskPatternInstrumentation { def askableActorRefAsk(ctx: TraceContextAware, actor: ActorRef): Unit = {} @Around("askableActorRefAsk(ctx, actor)") - def hookAskTimeoutWarning(pjp: ProceedingJoinPoint, ctx: TraceContextAware, actor: ActorRef): AnyRef = { - implicit val system = ctx.traceContext.system - val akkaExtension = Kamon(Akka)(system) - - val future = pjp.proceed().asInstanceOf[Future[AnyRef]] - - val handler = akkaExtension.askPatternTimeoutWarning match { - case "off" ⇒ None - case "lightweight" ⇒ Some(errorHandler(callInfo = Some(CallInfo(s"${actor.path.name} ?", pjp.getSourceLocation)))) - case "heavyweight" ⇒ Some(errorHandler(stack = Some(new StackTraceCaptureException))) - } - - handler.map(h ⇒ future.onFailure(h)(akkaExtension.dispatcher)) + def hookAskTimeoutWarning(pjp: ProceedingJoinPoint, ctx: TraceContextAware, actor: ActorRef): AnyRef = ctx.traceContext match { + case EmptyTraceContext ⇒ pjp.proceed() + case ctx: TraceContext ⇒ + implicit val system = ctx.system + val akkaExtension = Kamon(Akka)(system) + + val future = pjp.proceed().asInstanceOf[Future[AnyRef]] + + val handler = akkaExtension.askPatternTimeoutWarning match { + case "off" ⇒ None + case "lightweight" ⇒ Some(errorHandler(callInfo = Some(CallInfo(s"${actor.path.name} ?", pjp.getSourceLocation)))) + case "heavyweight" ⇒ Some(errorHandler(stack = Some(new StackTraceCaptureException))) + } - future + handler.map(future.onFailure(_)(akkaExtension.dispatcher)) + future } def errorHandler(callInfo: Option[CallInfo] = None, stack: Option[StackTraceCaptureException] = None)(implicit system: ActorSystem): ErrorHandler = { |