Merge branch 'issue#271/fix-balancing-pool-metrics'

1 files changed, 53 insertions, 0 deletions

diff --git a/kamon-akka/src/main/scala/kamon/akka/instrumentation/RouterInstrumentation.scala b/kamon-akka/src/main/scala/kamon/akka/instrumentation/RouterInstrumentation.scala
new file mode 100644
index 00000000..c11abc34
--- /dev/null
+++ b/kamon-akka/src/main/scala/kamon/akka/instrumentation/RouterInstrumentation.scala
@@ -0,0 +1,53 @@
+package akka.kamon.instrumentation
+
+import akka.actor.{ Props, ActorRef, ActorSystem, Cell }
+import akka.dispatch.{ Envelope, MessageDispatcher }
+import akka.routing.RoutedActorCell
+import org.aspectj.lang.ProceedingJoinPoint
+import org.aspectj.lang.annotation._
+
+@Aspect
+class RoutedActorCellInstrumentation {
+
+  def routerInstrumentation(cell: Cell): RouterMonitor =
+    cell.asInstanceOf[RouterInstrumentationAware].routerInstrumentation
+
+  @Pointcut("execution(akka.routing.RoutedActorCell.new(..)) && this(cell) && args(system, ref, props, dispatcher, routeeProps, supervisor)")
+  def routedActorCellCreation(cell: RoutedActorCell, system: ActorSystem, ref: ActorRef, props: Props, dispatcher: MessageDispatcher, routeeProps: Props, supervisor: ActorRef): Unit = {}
+
+  @After("routedActorCellCreation(cell, system, ref, props, dispatcher, routeeProps, supervisor)")
+  def afterRoutedActorCellCreation(cell: RoutedActorCell, system: ActorSystem, ref: ActorRef, props: Props, dispatcher: MessageDispatcher, routeeProps: Props, supervisor: ActorRef): Unit = {
+    cell.asInstanceOf[RouterInstrumentationAware].setRouterInstrumentation(
+      RouterMonitor.createRouterInstrumentation(cell))
+  }
+
+  @Pointcut("execution(* akka.routing.RoutedActorCell.sendMessage(*)) && this(cell) && args(envelope)")
+  def sendMessageInRouterActorCell(cell: RoutedActorCell, envelope: Envelope) = {}
+
+  @Around("sendMessageInRouterActorCell(cell, envelope)")
+  def aroundSendMessageInRouterActorCell(pjp: ProceedingJoinPoint, cell: RoutedActorCell, envelope: Envelope): Any = {
+    routerInstrumentation(cell).processMessage(pjp)
+  }
+}
+
+trait RouterInstrumentationAware {
+  def routerInstrumentation: RouterMonitor
+  def setRouterInstrumentation(ai: RouterMonitor): Unit
+}
+
+object RouterInstrumentationAware {
+  def apply(): RouterInstrumentationAware = new RouterInstrumentationAware {
+    private var _ri: RouterMonitor = _
+
+    def setRouterInstrumentation(ai: RouterMonitor): Unit = _ri = ai
+    def routerInstrumentation: RouterMonitor = _ri
+  }
+}
+
+@Aspect
+class MetricsIntoRouterCellsMixin {
+
+  @DeclareMixin("akka.routing.RoutedActorCell")
+  def mixinActorCellMetricsToRoutedActorCell: RouterInstrumentationAware = RouterInstrumentationAware()
+
+}
\ No newline at end of file