2 files changed, 22 insertions, 11 deletions
diff --git a/kamon-akka/src/main/scala/kamon/akka/instrumentation/ActorCellInstrumentation.scala b/kamon-akka/src/main/scala/kamon/akka/instrumentation/ActorCellInstrumentation.scala
index 4783484f..d1d420fe 100644
--- a/kamon-akka/src/main/scala/kamon/akka/instrumentation/ActorCellInstrumentation.scala
+++ b/kamon-akka/src/main/scala/kamon/akka/instrumentation/ActorCellInstrumentation.scala
@@ -34,11 +34,14 @@ class ActorCellInstrumentation {
 
   @After("actorCellCreation(cell, system, ref, props, dispatcher, parent)")
   def afterCreation(cell: ActorCell, system: ActorSystem, ref: ActorRef, props: Props, dispatcher: MessageDispatcher, parent: ActorRef): Unit = {
-    Kamon.metrics.register(ActorMetrics, ref.path.elements.mkString("/")).map { registration ⇒
+    val actorEntity = Entity(ref.path.elements.mkString("/"), ActorMetrics.category)
+
+    if (Kamon.metrics.shouldTrack(actorEntity)) {
+      val actorMetricsRecorder = Kamon.metrics.entity(ActorMetrics, actorEntity)
       val cellMetrics = cell.asInstanceOf[ActorCellMetrics]
 
-      cellMetrics.entity = registration.entity
-      cellMetrics.recorder = Some(registration.recorder)
+      cellMetrics.entity = actorEntity
+      cellMetrics.recorder = Some(actorMetricsRecorder)
     }
 
   }
@@ -90,14 +93,14 @@ class ActorCellInstrumentation {
   def afterStop(cell: ActorCell): Unit = {
     val cellMetrics = cell.asInstanceOf[ActorCellMetrics]
     cellMetrics.recorder.map { _ ⇒
-      Kamon.metrics.unregister(cellMetrics.entity)
+      Kamon.metrics.removeEntity(cellMetrics.entity)
     }
 
     // The Stop can't be captured from the RoutedActorCell so we need to put this piece of cleanup here.
     if (cell.isInstanceOf[RoutedActorCell]) {
       val routedCellMetrics = cell.asInstanceOf[RoutedActorCellMetrics]
       routedCellMetrics.routerRecorder.map { _ ⇒
-        Kamon.metrics.unregister(routedCellMetrics.routerEntity)
+        Kamon.metrics.removeEntity(routedCellMetrics.routerEntity)
       }
     }
   }
@@ -124,11 +127,13 @@ class RoutedActorCellInstrumentation {
 
   @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 = {
-    Kamon.metrics.register(RouterMetrics, ref.path.elements.mkString("/")).map { registration ⇒
+    val routerEntity = Entity(ref.path.elements.mkString("/"), RouterMetrics.category)
+
+    if (Kamon.metrics.shouldTrack(routerEntity)) {
       val cellMetrics = cell.asInstanceOf[RoutedActorCellMetrics]
 
-      cellMetrics.routerEntity = registration.entity
-      cellMetrics.routerRecorder = Some(registration.recorder)
+      cellMetrics.routerEntity = routerEntity
+      cellMetrics.routerRecorder = Some(Kamon.metrics.entity(RouterMetrics, routerEntity))
     }
   }
 
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 7b15c443..eb6b5293 100644
--- a/kamon-akka/src/main/scala/kamon/akka/instrumentation/DispatcherInstrumentation.scala
+++ b/kamon-akka/src/main/scala/kamon/akka/instrumentation/DispatcherInstrumentation.scala
@@ -59,10 +59,16 @@ class DispatcherInstrumentation {
   private def registerDispatcher(dispatcherName: String, executorService: ExecutorService, system: ActorSystem): Unit =
     executorService match {
       case fjp: AkkaForkJoinPool ⇒
-        Kamon.metrics.register(ForkJoinPoolDispatcherMetrics.factory(fjp), dispatcherName)
+        val dispatcherEntity = Entity(dispatcherName, AkkaDispatcherMetrics.Category)
+
+        if (Kamon.metrics.shouldTrack(dispatcherEntity))
+          Kamon.metrics.entity(ForkJoinPoolDispatcherMetrics.factory(fjp), dispatcherName)
 
       case tpe: ThreadPoolExecutor ⇒
-        Kamon.metrics.register(ThreadPoolExecutorDispatcherMetrics.factory(tpe), dispatcherName)
+        val dispatcherEntity = Entity(dispatcherName, AkkaDispatcherMetrics.Category)
+
+        if (Kamon.metrics.shouldTrack(dispatcherEntity))
+          Kamon.metrics.entity(ThreadPoolExecutorDispatcherMetrics.factory(tpe), dispatcherName)
 
       case others ⇒ // Currently not interested in other kinds of dispatchers.
     }
@@ -120,7 +126,7 @@ class DispatcherInstrumentation {
     import lazyExecutor.lookupData
 
     if (lookupData.actorSystem != null)
-      Kamon.metrics.unregister(Entity(lookupData.dispatcherName, AkkaDispatcherMetrics.Category))
+      Kamon.metrics.removeEntity(Entity(lookupData.dispatcherName, AkkaDispatcherMetrics.Category))
   }
 
 }