Initial ask pattern tracing

10 files changed, 111 insertions, 19 deletions

diff --git a/kamon-core/src/main/resources/META-INF/aop.xml b/kamon-core/src/main/resources/META-INF/aop.xml
index 3532b19b..af1dad7c 100644
--- a/kamon-core/src/main/resources/META-INF/aop.xml
+++ b/kamon-core/src/main/resources/META-INF/aop.xml
@@ -2,7 +2,7 @@
 
 <aspectj>
     <weaver options="-verbose -showWeaveInfo">
-       <!--<dump within="*" beforeandafter="true"/>-->
+       <!--<dump within="*"/>-->
     </weaver>
 
     <aspects>
@@ -12,7 +12,7 @@
         <aspect name="kamon.trace.instrumentation.RunnableTracing" />
         <aspect name="kamon.instrumentation.ActorSystemInstrumentation"/>
         <aspect name="kamon.trace.instrumentation.ActorLoggingInstrumentation"/>
-
+        <aspect name="akka.pattern.tracing.AskPatternTracing"/>
         <!--<aspect name="kamon.instrumentation.MessageQueueInstrumentation" />-->
         <!--<aspect name="kamon.instrumentation.InceptionAspect"/>-->
         <!-- ExecutorService Instrumentation for Akka. -->
@@ -21,14 +21,11 @@
         <!--<aspect name ="kamon.instrumentation.ForkJoinPoolInstrumentation"/>-->
 
 
-        <include within="*"/>
-        <exclude within="javax..*"/>
-        <exclude within="com.newrelic..*"/>
-        <exclude within="org.aspectj..*"/>
-        <exclude within="scala..*"/>
-        <exclude within="scalaz..*"/>
-        <exclude within="scalad..*"/>
-        <exclude within="play..*"/>
+        <!--<exclude within="*"/>-->
+        <include within="scala.concurrent..*"/>
+        <include within="akka..*"/>
+        <include within="spray..*"/>
+        <include within="kamon..*"/>
     </aspects>
 
 </aspectj>
diff --git a/kamon-trace/src/main/resources/reference.conf b/kamon-trace/src/main/resources/reference.conf
new file mode 100644
index 00000000..03f13f01
--- /dev/null
+++ b/kamon-trace/src/main/resources/reference.conf
@@ -0,0 +1,3 @@
+kamon.trace {
+  ask-pattern-tracing = off
+}
\ No newline at end of file
diff --git a/kamon-trace/src/main/scala/kamon/trace/Trace.scala b/kamon-trace/src/main/scala/kamon/trace/Trace.scala
index d5bb1ddf..0ef438d4 100644
--- a/kamon-trace/src/main/scala/kamon/trace/Trace.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/Trace.scala
@@ -27,10 +27,12 @@ object Trace extends ExtensionId[TraceExtension] with ExtensionIdProvider {
   def set(ctx: TraceContext) = traceContext.value = Some(ctx)
 
   def clear: Unit = traceContext.value = None
-  def start(name: String)(implicit system: ActorSystem) = {
-    val ctx = newTraceContext()
+  def start(name: String)(implicit system: ActorSystem): TraceContext = {
+    val ctx = newTraceContext(name)
     ctx.start(name)
     set(ctx)
+
+    ctx
   }
 
   def transformContext(f: TraceContext => TraceContext): Unit = {
@@ -44,7 +46,7 @@ object Trace extends ExtensionId[TraceExtension] with ExtensionIdProvider {
   }
 
   // TODO: FIX
-  def newTraceContext()(implicit system: ActorSystem): TraceContext = TraceContext(Kamon(Trace), tranid.getAndIncrement)
+  def newTraceContext(name: String)(implicit system: ActorSystem): TraceContext = TraceContext(Kamon(Trace), tranid.getAndIncrement, name)
 }
 
 class TraceExtension(system: ExtendedActorSystem) extends Kamon.Extension {
diff --git a/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorLoggingInstrumentation.scala b/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorLoggingInstrumentation.scala
index 77993cdd..fecdc7e3 100644
--- a/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorLoggingInstrumentation.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorLoggingInstrumentation.scala
@@ -10,9 +10,7 @@ class ActorLoggingInstrumentation {
 
 
   @DeclareMixin("akka.event.Logging.LogEvent+")
-  def traceContextMixin: ContextAware = new ContextAware {
-    def traceContext: Option[TraceContext] = Trace.context()
-  }
+  def traceContextMixin: ContextAware = ContextAware.default
 
   @Pointcut("execution(* akka.event.slf4j.Slf4jLogger.withMdc(..)) && args(logSource, logEvent, logStatement)")
   def withMdcInvocation(logSource: String, logEvent: ContextAware, logStatement: () => _): Unit = {}
diff --git a/kamon-trace/src/main/scala/kamon/trace/instrumentation/AskPatternTracing.scala b/kamon-trace/src/main/scala/kamon/trace/instrumentation/AskPatternTracing.scala
new file mode 100644
index 00000000..625ed1ae
--- /dev/null
+++ b/kamon-trace/src/main/scala/kamon/trace/instrumentation/AskPatternTracing.scala
@@ -0,0 +1,34 @@
+package akka.pattern.tracing
+
+import org.aspectj.lang.annotation.{After, AfterReturning, Pointcut, Aspect}
+import akka.pattern.{AskTimeoutException, PromiseActorRef}
+import akka.event.Logging.Warning
+import scala.compat.Platform.EOL
+import akka.actor.ActorRefProvider
+
+@Aspect
+class AskPatternTracing {
+
+  class StackTraceCaptureException extends Throwable
+
+  @Pointcut(value = "execution(* akka.pattern.PromiseActorRef$.apply(..)) && args(provider, *)", argNames = "provider")
+  def promiseActorRefApply(provider: ActorRefProvider): Unit = {
+    provider.settings.config.getBoolean("kamon.trace.ask-pattern-tracing")
+  }
+
+  @AfterReturning(pointcut = "promiseActorRefApply(provider)", returning = "promiseActor")
+  def hookAskTimeoutWarning(provider: ActorRefProvider, promiseActor: PromiseActorRef): Unit = {
+    val future = promiseActor.result.future
+    val system = promiseActor.provider.guardian.underlying.system
+    implicit val ec = system.dispatcher
+    val stack = new StackTraceCaptureException
+
+    future onFailure {
+      case timeout: AskTimeoutException =>
+        val stackString = stack.getStackTrace.drop(3).mkString("", EOL, EOL)
+
+        system.eventStream.publish(Warning("AskPatternTracing", classOf[AskPatternTracing],
+          "Timeout triggered for ask pattern registered at: " + stackString))
+    }
+  }
+}
diff --git a/kamon-trace/src/test/resources/application.conf b/kamon-trace/src/test/resources/application.conf
new file mode 100644
index 00000000..e8217fc2
--- /dev/null
+++ b/kamon-trace/src/test/resources/application.conf
@@ -0,0 +1,3 @@
+akka {
+  loggers = ["akka.event.slf4j.Slf4jLogger"]
+}
\ No newline at end of file
diff --git a/kamon-trace/src/test/resources/logback.xml b/kamon-trace/src/test/resources/logback.xml
new file mode 100644
index 00000000..2ae1e3bd
--- /dev/null
+++ b/kamon-trace/src/test/resources/logback.xml
@@ -0,0 +1,12 @@
+<configuration scan="true">
+    <appender name="STDOUT" class="ch.qos.logback.core.ConsoleAppender">
+        <encoder>
+            <pattern>%date{HH:mm:ss.SSS} %-5level [%X{uow}][%X{requestId}] [%thread] %logger{55} - %msg%n</pattern>
+        </encoder>
+    </appender>
+
+    <root level="debug">
+        <appender-ref ref="STDOUT" />
+    </root>
+
+</configuration>
diff --git a/kamon-trace/src/test/scala/kamon/ActorInstrumentationSpec.scala b/kamon-trace/src/test/scala/kamon/ActorInstrumentationSpec.scala
index f5d88f06..d675c4f4 100644
--- a/kamon-trace/src/test/scala/kamon/ActorInstrumentationSpec.scala
+++ b/kamon-trace/src/test/scala/kamon/ActorInstrumentationSpec.scala
@@ -65,7 +65,7 @@ class ActorInstrumentationSpec extends TestKit(ActorSystem("ActorInstrumentation
   }
 
   trait TraceContextEchoFixture {
-    val testTraceContext = Trace.newTraceContext()
+    val testTraceContext = Trace.newTraceContext("")
     val echo = system.actorOf(Props[TraceContextEcho])
 
     Trace.set(testTraceContext)
@@ -75,7 +75,7 @@ class ActorInstrumentationSpec extends TestKit(ActorSystem("ActorInstrumentation
     override val echo = system.actorOf(Props[TraceContextEcho].withRouter(RoundRobinRouter(nrOfInstances = 10)))
 
     def tellWithNewContext(target: ActorRef, message: Any): TraceContext = {
-      val context = Trace.newTraceContext()
+      val context = Trace.newTraceContext("")
       Trace.set(context)
 
       target ! message
diff --git a/kamon-trace/src/test/scala/kamon/AskPatternTracingSpec.scala b/kamon-trace/src/test/scala/kamon/AskPatternTracingSpec.scala
new file mode 100644
index 00000000..c2566725
--- /dev/null
+++ b/kamon-trace/src/test/scala/kamon/AskPatternTracingSpec.scala
@@ -0,0 +1,43 @@
+package kamon
+
+import akka.testkit.TestKit
+import akka.actor.{Props, Actor, ActorSystem}
+import org.scalatest.{Matchers, WordSpecLike}
+import akka.event.Logging.Warning
+import scala.concurrent.duration._
+import akka.pattern.ask
+import akka.util.Timeout
+import kamon.trace.{Trace, ContextAware}
+import org.scalatest.OptionValues._
+
+
+class AskPatternTracingSpec extends TestKit(ActorSystem("ask-pattern-tracing-spec")) with WordSpecLike with Matchers {
+
+  "the AskPatternTracing" should {
+    "log a warning with a stack trace and TraceContext taken from the moment the ask was triggered" in {
+      implicit val ec = system.dispatcher
+      implicit val timeout = Timeout(10 milliseconds)
+      val noReply = system.actorOf(Props[NoReply])
+      system.eventStream.subscribe(testActor, classOf[Warning])
+
+      within(500 milliseconds) {
+        val initialCtx = Trace.start("ask-test")
+        noReply ? "hello"
+
+        val warn = expectMsgPF() {
+          case warn: Warning if warn.message.toString.contains("Timeout triggered for ask pattern") => warn
+        }
+        val capturedCtx = warn.asInstanceOf[ContextAware].traceContext
+
+        capturedCtx should be('defined)
+        capturedCtx.value should equal (initialCtx)
+      }
+    }
+  }
+}
+
+class NoReply extends Actor {
+  def receive = {
+    case any =>
+  }
+}
diff --git a/project/Build.scala b/project/Build.scala
index 323f7f02..67d37422 100644
--- a/project/Build.scala
+++ b/project/Build.scala
@@ -29,7 +29,7 @@ object Build extends Build {
     .settings(
       libraryDependencies ++=
         compile(akkaActor, aspectJ) ++
-        test(scalatest, akkaTestKit, sprayTestkit))
+        test(scalatest, akkaTestKit, sprayTestkit, akkaSlf4j, logback))
     .dependsOn(kamonCore)
 
   lazy val kamonMetrics = Project("kamon-metrics", file("kamon-metrics"))