make sure the TraceContext is passed around when processing system messages

5 files changed, 246 insertions, 5 deletions

diff --git a/kamon-trace/src/main/resources/META-INF/aop.xml b/kamon-trace/src/main/resources/META-INF/aop.xml
index 40bea5fc..d9916724 100644
--- a/kamon-trace/src/main/resources/META-INF/aop.xml
+++ b/kamon-trace/src/main/resources/META-INF/aop.xml
@@ -3,6 +3,9 @@
 <aspectj>
   <aspects>
     <!-- Actors -->
+    <aspect name="akka.actor.RepointableActorRefTraceContextMixin"/>
+    <aspect name="akka.actor.SystemMessageTraceContextMixin"/>
+    <aspect name="akka.actor.ActorSystemMessagePassingTracing"/>
     <aspect name="kamon.trace.instrumentation.EnvelopeTraceContextMixin"/>
     <aspect name="kamon.trace.instrumentation.BehaviourInvokeTracing"/>
     <aspect name="kamon.trace.instrumentation.ActorLoggingTracing"/>
diff --git a/kamon-trace/src/main/scala/akka/actor/ActorSystemMessagePassingTracing.scala b/kamon-trace/src/main/scala/akka/actor/ActorSystemMessagePassingTracing.scala
new file mode 100644
index 00000000..b6b7b048
--- /dev/null
+++ b/kamon-trace/src/main/scala/akka/actor/ActorSystemMessagePassingTracing.scala
@@ -0,0 +1,65 @@
+package akka.actor
+
+import org.aspectj.lang.annotation._
+import kamon.trace.{Trace, ContextAware}
+import akka.dispatch.sysmsg.EarliestFirstSystemMessageList
+import org.aspectj.lang.ProceedingJoinPoint
+
+@Aspect
+class SystemMessageTraceContextMixin {
+
+  @DeclareMixin("akka.dispatch.sysmsg.SystemMessage+")
+  def mixin: ContextAware = ContextAware.default
+
+  @Pointcut("execution(akka.dispatch.sysmsg.SystemMessage+.new(..)) && this(ctx)")
+  def envelopeCreation(ctx: ContextAware): Unit = {}
+
+  @After("envelopeCreation(ctx)")
+  def afterEnvelopeCreation(ctx: ContextAware): Unit = {
+    // Necessary to force the initialization of ContextAware at the moment of creation.
+    ctx.traceContext
+  }
+}
+
+@Aspect
+class RepointableActorRefTraceContextMixin {
+
+  @DeclareMixin("akka.actor.RepointableActorRef")
+  def mixin: ContextAware = ContextAware.default
+
+  @Pointcut("execution(akka.actor.RepointableActorRef.new(..)) && this(ctx)")
+  def envelopeCreation(ctx: ContextAware): Unit = {}
+
+  @After("envelopeCreation(ctx)")
+  def afterEnvelopeCreation(ctx: ContextAware): Unit = {
+    // Necessary to force the initialization of ContextAware at the moment of creation.
+    ctx.traceContext
+  }
+
+  @Pointcut("execution(* akka.actor.RepointableActorRef.point(..)) && this(repointableActorRef)")
+  def repointableActorRefCreation(repointableActorRef: ContextAware): Unit = {}
+
+  @Around("repointableActorRefCreation(repointableActorRef)")
+  def afterRepointableActorRefCreation(pjp: ProceedingJoinPoint, repointableActorRef: ContextAware): Any = {
+    Trace.withContext(repointableActorRef.traceContext) {
+      pjp.proceed()
+    }
+  }
+
+}
+
+@Aspect
+class ActorSystemMessagePassingTracing {
+
+  @Pointcut("execution(* akka.actor.ActorCell.invokeAll$1(..)) && args(messages, *)")
+  def systemMessageProcessing(messages: EarliestFirstSystemMessageList): Unit = {}
+
+  @Around("systemMessageProcessing(messages)")
+  def aroundSystemMessageInvoke(pjp: ProceedingJoinPoint, messages: EarliestFirstSystemMessageList): Any = {
+    if(messages.nonEmpty) {
+      val ctx = messages.head.asInstanceOf[ContextAware].traceContext
+      Trace.withContext(ctx)(pjp.proceed())
+
+    } else pjp.proceed()
+  }
+}
diff --git a/kamon-trace/src/test/scala/kamon/trace/instrumentation/ActorSystemMessagePassingInstrumentationSpec.scala b/kamon-trace/src/test/scala/kamon/trace/instrumentation/ActorSystemMessagePassingInstrumentationSpec.scala
new file mode 100644
index 00000000..a845ad0c
--- /dev/null
+++ b/kamon-trace/src/test/scala/kamon/trace/instrumentation/ActorSystemMessagePassingInstrumentationSpec.scala
@@ -0,0 +1,168 @@
+package kamon.trace.instrumentation
+
+import akka.testkit.{ImplicitSender, TestKit}
+import akka.actor._
+import org.scalatest.WordSpecLike
+import kamon.trace.Trace
+import scala.util.control.NonFatal
+import akka.actor.SupervisorStrategy.{Escalate, Stop, Restart, Resume}
+import scala.concurrent.duration._
+
+class ActorSystemMessagePassingInstrumentationSpec extends TestKit(ActorSystem("actor-message-passing-tracing-spec")) with WordSpecLike with ImplicitSender {
+  implicit val executionContext = system.dispatcher
+
+
+  "the system message passing instrumentation" should {
+    "keep the TraceContext while processing the Create message in top level actors" in new TraceContextFixture {
+      Trace.withContext(testTraceContext) {
+        system.actorOf(Props(new Actor {
+
+          testActor ! Trace.context()
+
+          def receive: Actor.Receive = { case any => }
+        }))
+      }
+
+      expectMsg(testTraceContext)
+    }
+
+
+    "keep the TraceContext while processing the Create message in non top level actors" in new TraceContextFixture {
+      Trace.withContext(testTraceContext) {
+        system.actorOf(Props(new Actor {
+          def receive: Actor.Receive = {
+            case any =>
+              context.actorOf(Props(new Actor {
+
+                testActor ! Trace.context()
+
+                def receive: Actor.Receive = { case any => }
+              }))
+          }
+        })) ! "any"
+      }
+
+      expectMsg(testTraceContext)
+    }
+
+
+    "keep the TraceContext in the supervision cycle" when {
+      "the actor is resumed" in new TraceContextFixture {
+        val supervisor = supervisorWithDirective(Resume)
+
+        Trace.withContext(testTraceContext) {
+          supervisor ! "fail"
+        }
+
+        expectMsg(testTraceContext) // From the parent executing the supervision strategy
+
+        // Ensure we didn't tie the actor with the context
+        supervisor ! "context"
+        expectMsg(None)
+      }
+
+      "the actor is restarted" in new TraceContextFixture {
+        val supervisor = supervisorWithDirective(Restart, sendPreRestart = true, sendPostRestart = true)
+
+        Trace.withContext(testTraceContext) {
+          supervisor ! "fail"
+        }
+
+        expectMsg(testTraceContext) // From the parent executing the supervision strategy
+        expectMsg(testTraceContext) // From the preRestart hook
+        expectMsg(testTraceContext) // From the postRestart hook
+
+        // Ensure we didn't tie the actor with the context
+        supervisor ! "context"
+        expectMsg(None)
+      }
+
+      "the actor is stopped" in new TraceContextFixture {
+        val supervisor = supervisorWithDirective(Stop, sendPostStop = true)
+
+        Trace.withContext(testTraceContext) {
+          supervisor ! "fail"
+        }
+
+        expectMsg(testTraceContext) // From the parent executing the supervision strategy
+        expectMsg(testTraceContext) // From the postStop hook
+        expectNoMsg(1 second)
+      }
+
+      "the failure is escalated" in new TraceContextFixture {
+        val supervisor = supervisorWithDirective(Escalate, sendPostStop = true)
+
+        Trace.withContext(testTraceContext) {
+          supervisor ! "fail"
+        }
+
+        expectMsg(testTraceContext) // From the parent executing the supervision strategy
+        expectMsg(testTraceContext) // From the grandparent executing the supervision strategy
+        expectMsg(testTraceContext) // From the postStop hook in the child
+        expectMsg(testTraceContext) // From the postStop hook in the parent
+        expectNoMsg(1 second)
+      }
+    }
+  }
+
+  def supervisorWithDirective(directive: SupervisorStrategy.Directive, sendPreRestart: Boolean = false, sendPostRestart: Boolean = false,
+                               sendPostStop: Boolean = false, sendPreStart: Boolean = false): ActorRef = {
+    class GrandParent extends Actor {
+      val child = context.actorOf(Props(new Parent))
+
+      override def supervisorStrategy: SupervisorStrategy = OneForOneStrategy() {
+        case NonFatal(throwable) => testActor ! Trace.context(); Stop
+      }
+
+      def receive = {
+        case any => child forward any
+      }
+    }
+
+    class Parent extends Actor {
+      val child = context.actorOf(Props(new Child))
+
+      override def supervisorStrategy: SupervisorStrategy = OneForOneStrategy() {
+        case NonFatal(throwable) => testActor ! Trace.context(); directive
+      }
+
+      def receive: Actor.Receive = {
+        case any => child forward any
+      }
+
+      override def postStop(): Unit = {
+        if(sendPostStop) testActor ! Trace.context()
+        super.postStop()
+      }
+    }
+
+    class Child extends Actor {
+      def receive = {
+        case "fail"     => 1/0
+        case "context"  => sender ! Trace.context()
+      }
+
+      override def preRestart(reason: Throwable, message: Option[Any]): Unit = {
+        if(sendPreRestart) testActor ! Trace.context()
+        super.preRestart(reason, message)
+      }
+
+      override def postRestart(reason: Throwable): Unit = {
+        if(sendPostRestart) testActor ! Trace.context()
+        super.postRestart(reason)
+      }
+
+      override def postStop(): Unit = {
+        if(sendPostStop) testActor ! Trace.context()
+        super.postStop()
+      }
+
+      override def preStart(): Unit = {
+        if(sendPreStart) testActor ! Trace.context()
+        super.preStart()
+      }
+    }
+
+    system.actorOf(Props(new GrandParent))
+  }
+}
diff --git a/kamon-trace/src/test/scala/kamon/trace/instrumentation/FutureTracingSpec.scala b/kamon-trace/src/test/scala/kamon/trace/instrumentation/FutureTracingSpec.scala
index 339fe6be..a5554836 100644
--- a/kamon-trace/src/test/scala/kamon/trace/instrumentation/FutureTracingSpec.scala
+++ b/kamon-trace/src/test/scala/kamon/trace/instrumentation/FutureTracingSpec.scala
@@ -58,10 +58,5 @@ class FutureTracingSpec extends WordSpec with Matchers with ScalaFutures with Pa
       }
     }
   }
-
-  trait TraceContextFixture {
-    val random = new Random(System.nanoTime)
-    val testTraceContext = Some(TraceContext(Actor.noSender, random.nextInt))
-  }
 }
 
diff --git a/kamon-trace/src/test/scala/kamon/trace/instrumentation/TraceContextFixture.scala b/kamon-trace/src/test/scala/kamon/trace/instrumentation/TraceContextFixture.scala
new file mode 100644
index 00000000..62f7ec84
--- /dev/null
+++ b/kamon-trace/src/test/scala/kamon/trace/instrumentation/TraceContextFixture.scala
@@ -0,0 +1,10 @@
+package kamon.trace.instrumentation
+
+import scala.util.Random
+import kamon.trace.TraceContext
+import akka.actor.Actor
+
+trait TraceContextFixture {
+  val random = new Random(System.nanoTime)
+  val testTraceContext = Some(TraceContext(Actor.noSender, random.nextInt))
+}
\ No newline at end of file