= akka: correctly set LookupData for routers used by BalancingPool, fixes #199.

1 files changed, 13 insertions, 1 deletions

diff --git a/kamon-akka/src/main/scala/kamon/akka/instrumentation/DispatcherInstrumentation.scala b/kamon-akka/src/main/scala/kamon/akka/instrumentation/DispatcherInstrumentation.scala
index 0b963c16..e90838ad 100644
--- a/kamon-akka/src/main/scala/kamon/akka/instrumentation/DispatcherInstrumentation.scala
+++ b/kamon-akka/src/main/scala/kamon/akka/instrumentation/DispatcherInstrumentation.scala
@@ -18,7 +18,7 @@ package akka.kamon.instrumentation
 
 import java.util.concurrent.{ ExecutorService, ThreadPoolExecutor }
 
-import akka.actor.{ ActorSystem, ActorSystemImpl }
+import akka.actor.{ ActorContext, Props, ActorSystem, ActorSystemImpl }
 import akka.dispatch.ForkJoinExecutorConfigurator.AkkaForkJoinPool
 import akka.dispatch._
 import akka.kamon.instrumentation.LookupDataAware.LookupData
@@ -139,6 +139,18 @@ class DispatcherInstrumentation {
       }
   }
 
+  @Pointcut("execution(* akka.routing.BalancingPool.newRoutee(..)) && args(props, context)")
+  def createNewRouteeOnBalancingPool(props: Props, context: ActorContext): Unit = {}
+
+  @Around("createNewRouteeOnBalancingPool(props, context)")
+  def aroundCreateNewRouteeOnBalancingPool(pjp: ProceedingJoinPoint, props: Props, context: ActorContext): Any = {
+    val deployPath = context.self.path.elements.drop(1).mkString("/", "/", "")
+    val dispatcherId = s"BalancingPool-$deployPath"
+
+    LookupDataAware.withLookupData(LookupData(dispatcherId, context.system)) {
+      pjp.proceed()
+    }
+  }
 }
 
 @Aspect