= core: avoid error when TraceContext is empty

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 = {