wip in aspects for actor tracing

3 files changed, 55 insertions, 24 deletions

diff --git a/src/main/scala/akka/ActorAspect.scala b/src/main/scala/akka/ActorAspect.scala
index 9d64f205..05a7bc0a 100644
--- a/src/main/scala/akka/ActorAspect.scala
+++ b/src/main/scala/akka/ActorAspect.scala
@@ -11,7 +11,7 @@ class ActorAspect extends Metrics {
    @Pointcut("execution(* akka.actor.ActorCell+.receiveMessage(..))")
    protected def actorReceive:Unit = {}
 
-   @Around("actorReceive() && this(actor)")
+   @Around("sendingMessageToActorRef() && this(actor)")
    def around(pjp: ProceedingJoinPoint, actor: akka.actor.ActorCell): AnyRef = {
 
      //println("The path is: "+actor.self.path.)
diff --git a/src/main/scala/akka/instrumentation/ActorInstrumentation.scala b/src/main/scala/akka/instrumentation/ActorInstrumentation.scala
deleted file mode 100644
index ea599891..00000000
--- a/src/main/scala/akka/instrumentation/ActorInstrumentation.scala
+++ /dev/null
@@ -1,23 +0,0 @@
-package akka.instrumentation
-
-import org.aspectj.lang.annotation.{Before, Around, Pointcut, Aspect}
-import org.aspectj.lang.ProceedingJoinPoint
-import kamon.metric.Metrics
-import akka.actor.ActorCell
-
-@Aspect
-class ActorInstrumentation {
-  println("Created ActorAspect")
-
-  @Pointcut("execution(* kamon.executor.PingActor.receive(..))")
-  protected def actorReceive:Unit = {}
-
-  @Before("actorReceive() && args(message)")
-  def around(message: Any) = {
-    println("Around the actor cell receive")
-    //pjp.proceed(Array(Wrapper(message)))
-    //pjp.proceed
-  }
-}
-
-case class Wrapper(content: Any)
\ No newline at end of file
diff --git a/src/main/scala/akka/instrumentation/ActorRefTellInstrumentation.scala b/src/main/scala/akka/instrumentation/ActorRefTellInstrumentation.scala
new file mode 100644
index 00000000..783a6c45
--- /dev/null
+++ b/src/main/scala/akka/instrumentation/ActorRefTellInstrumentation.scala
@@ -0,0 +1,54 @@
+package akka.instrumentation
+
+import org.aspectj.lang.annotation.{Around, Pointcut, Aspect}
+import org.aspectj.lang.ProceedingJoinPoint
+import akka.actor.{ActorRef, ActorCell}
+import kamon.TraceContext
+import kamon.actor.TraceableMessage
+import akka.dispatch.Envelope
+
+@Aspect
+class ActorRefTellInstrumentation {
+  println("Created ActorAspect")
+
+  @Pointcut("execution(* akka.actor.ScalaActorRef+.$bang(..)) && args(message, sender)")
+  def sendingMessageToActorRef(message: Any, sender: ActorRef) = {}
+
+  @Around("sendingMessageToActorRef(message, sender)")
+  def around(pjp: ProceedingJoinPoint, message: Any, sender: ActorRef): Unit  = {
+    import pjp._
+
+    TraceContext.current match {
+      case Some(ctx) => {
+        val traceableMessage = TraceableMessage(ctx.fork, message)
+        proceed(getArgs.updated(0, traceableMessage))
+      }
+      case None => proceed
+    }
+  }
+}
+
+@Aspect
+class ActorCellInvokeInstrumentation {
+
+  @Pointcut("execution(* akka.actor.ActorCell.invoke(*)) && args(envelope)")
+  def invokingActorBehaviourAtActorCell(envelope: Envelope) = {}
+
+
+  @Around("invokingActorBehaviourAtActorCell(envelope)")
+  def around(pjp: ProceedingJoinPoint, envelope: Envelope) = {
+    import pjp._
+
+    envelope match {
+      case Envelope(TraceableMessage(ctx, msg), sender) => {
+        TraceContext.set(ctx)
+
+        val originalEnvelope = envelope.copy(message = msg)
+        proceed(getArgs.updated(0, originalEnvelope))
+
+        TraceContext.clear
+      }
+      case _ => proceed
+    }
+  }
+}
\ No newline at end of file