diff options
Diffstat (limited to 'kamon-core/src/main/scala/kamon/metric')
-rw-r--r-- | kamon-core/src/main/scala/kamon/metric/EntityRecorder.scala | 2 | ||||
-rw-r--r-- | kamon-core/src/main/scala/kamon/metric/MetricsExtension.scala | 10 |
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 = |