! core: revised semantics of .register on Metrics module to match our docs.

2 files changed, 8 insertions, 4 deletions

diff --git a/kamon-core/src/main/scala/kamon/metric/EntityRecorder.scala b/kamon-core/src/main/scala/kamon/metric/EntityRecorder.scala
index a7bba2e1..6e0a4248 100644
--- a/kamon-core/src/main/scala/kamon/metric/EntityRecorder.scala
+++ b/kamon-core/src/main/scala/kamon/metric/EntityRecorder.scala
@@ -38,7 +38,7 @@ abstract class GenericEntityRecorder(instrumentFactory: InstrumentFactory) exten
 
   private val _instruments = TrieMap.empty[MetricKey, Instrument]
   private def register[T <: Instrument](key: MetricKey, instrument: ⇒ T): T =
-    _instruments.atomicGetOrElseUpdate(key, instrument).asInstanceOf[T]
+    _instruments.atomicGetOrElseUpdate(key, instrument, _.cleanup).asInstanceOf[T]
 
   protected def histogram(name: String): Histogram =
     register(HistogramKey(name), instrumentFactory.createHistogram(name))
diff --git a/kamon-core/src/main/scala/kamon/metric/MetricsExtension.scala b/kamon-core/src/main/scala/kamon/metric/MetricsExtension.scala
index c14d4e16..88352e21 100644
--- a/kamon-core/src/main/scala/kamon/metric/MetricsExtension.scala
+++ b/kamon-core/src/main/scala/kamon/metric/MetricsExtension.scala
@@ -86,14 +86,18 @@ class MetricsExtensionImpl(system: ExtendedActorSystem) extends MetricsExtension
 
     if (shouldTrack(entity)) {
       val instrumentFactory = settings.instrumentFactories.get(recorderFactory.category).getOrElse(settings.defaultInstrumentFactory)
-      val recorder = _trackedEntities.atomicGetOrElseUpdate(entity, recorderFactory.createRecorder(instrumentFactory)).asInstanceOf[T]
+      val recorder = _trackedEntities.atomicGetOrElseUpdate(entity, recorderFactory.createRecorder(instrumentFactory), _.cleanup).asInstanceOf[T]
+
       Some(EntityRegistration(entity, recorder))
     } else None
   }
 
   def register[T <: EntityRecorder](entity: Entity, recorder: T): EntityRegistration[T] = {
-    import TriemapAtomicGetOrElseUpdate.Syntax
-    EntityRegistration(entity, _trackedEntities.atomicGetOrElseUpdate(entity, recorder).asInstanceOf[T])
+    _trackedEntities.put(entity, recorder).map { oldRecorder ⇒
+      oldRecorder.cleanup
+    }
+
+    EntityRegistration(entity, recorder)
   }
 
   def unregister(entity: Entity): Unit =