make sure the TraceContext is passed around when processing system messages

2 files changed, 68 insertions, 0 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()
+  }
+}