+ akka: avoid runtime exceptions logged on dispatchers shutdown

1 files changed, 8 insertions, 6 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 931d8a45..cc177aac 100644
--- a/kamon-akka/src/main/scala/kamon/akka/instrumentation/DispatcherInstrumentation.scala
+++ b/kamon-akka/src/main/scala/kamon/akka/instrumentation/DispatcherInstrumentation.scala
@@ -24,7 +24,7 @@ import akka.dispatch._
 import akka.kamon.instrumentation.LookupDataAware.LookupData
 import kamon.Kamon
 import kamon.akka.{ AkkaDispatcherMetrics, ThreadPoolExecutorDispatcherMetrics, ForkJoinPoolDispatcherMetrics }
-import kamon.metric.Entity
+import kamon.metric.{MetricsModule, Entity}
 import org.aspectj.lang.ProceedingJoinPoint
 import org.aspectj.lang.annotation._
 
@@ -121,6 +121,7 @@ class DispatcherInstrumentation {
   @Pointcut("execution(* akka.dispatch.Dispatcher.LazyExecutorServiceDelegate.shutdown()) && this(lazyExecutor)")
   def lazyExecutorShutdown(lazyExecutor: LookupDataAware): Unit = {}
 
+
   @After("lazyExecutorShutdown(lazyExecutor)")
   def afterLazyExecutorShutdown(lazyExecutor: LookupDataAware): Unit = {
     import lazyExecutor.lookupData
@@ -128,17 +129,18 @@ class DispatcherInstrumentation {
     if (lookupData.actorSystem != null)
       lazyExecutor.asInstanceOf[ExecutorServiceDelegate].executor match {
         case fjp: AkkaForkJoinPool ⇒
-          Kamon.metrics.removeEntity(Entity(lookupData.actorSystem.name + "/" + lookupData.dispatcherName,
-            AkkaDispatcherMetrics.Category, tags = Map("dispatcher-type" -> "fork-join-pool")))
+          lookupData.metrics.removeEntity(Entity(lookupData.actorSystem.name + "/" + lookupData.dispatcherName,
+          AkkaDispatcherMetrics.Category, tags = Map("dispatcher-type" -> "fork-join-pool")))
 
         case tpe: ThreadPoolExecutor ⇒
-          Kamon.metrics.removeEntity(Entity(lookupData.actorSystem.name + "/" + lookupData.dispatcherName,
-            AkkaDispatcherMetrics.Category, tags = Map("dispatcher-type" -> "thread-pool-executor")))
+          lookupData.metrics.removeEntity(Entity(lookupData.actorSystem.name + "/" + lookupData.dispatcherName,
+          AkkaDispatcherMetrics.Category, tags = Map("dispatcher-type" -> "thread-pool-executor")))
 
         case other ⇒ // nothing to remove.
       }
   }
 
+
 }
 
 @Aspect
@@ -167,7 +169,7 @@ trait LookupDataAware {
 }
 
 object LookupDataAware {
-  case class LookupData(dispatcherName: String, actorSystem: ActorSystem)
+  case class LookupData(dispatcherName: String, actorSystem: ActorSystem, metrics:MetricsModule = Kamon.metrics)
 
   private val _currentDispatcherLookupData = new ThreadLocal[LookupData]