From 9cfb5936afef80ea57f135762432d29427406d9d Mon Sep 17 00:00:00 2001
From: Bartosz Spyrko-Smietanko <bspyrkos@redhat.com>
Date: Tue, 23 Mar 2021 13:51:28 +0000
Subject: [PATCH] #420 Clean up metrics member to metric mappings on shutdown

---
 .../metrics/MemberToMetricMappings.java       | 13 +++++++
 .../setup/MetricCdiInjectionExtension.java    | 38 ++++++++++++++-----
 2 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/implementation/src/main/java/io/smallrye/metrics/MemberToMetricMappings.java b/implementation/src/main/java/io/smallrye/metrics/MemberToMetricMappings.java
index d489a9d8..5d8b5054 100644
--- a/implementation/src/main/java/io/smallrye/metrics/MemberToMetricMappings.java
+++ b/implementation/src/main/java/io/smallrye/metrics/MemberToMetricMappings.java
@@ -86,4 +86,17 @@ public void addMetric(MemberInfo member, MetricID metricID, MetricType metricTyp
         SmallRyeMetricsLogging.log.matchingMemberToMetric(member, metricID, metricType);
     }
 
+    public void removeMappingsFor(MemberInfo member, MetricID metricID) {
+        removeMapping(counters, member, metricID);
+        removeMapping(concurrentGauges, member, metricID);
+        removeMapping(meters, member, metricID);
+        removeMapping(timers, member, metricID);
+        removeMapping(simpleTimers, member, metricID);
+    }
+
+    private void removeMapping(Map<MemberInfo, Set<MetricID>> map, MemberInfo member, MetricID metricID) {
+        if (map.containsKey(member)) {
+            map.get(member).remove(metricID);
+        }
+    }
 }
diff --git a/implementation/src/main/java/io/smallrye/metrics/setup/MetricCdiInjectionExtension.java b/implementation/src/main/java/io/smallrye/metrics/setup/MetricCdiInjectionExtension.java
index e5493ce2..f6367320 100644
--- a/implementation/src/main/java/io/smallrye/metrics/setup/MetricCdiInjectionExtension.java
+++ b/implementation/src/main/java/io/smallrye/metrics/setup/MetricCdiInjectionExtension.java
@@ -53,10 +53,13 @@
 import org.eclipse.microprofile.metrics.annotation.SimplyTimed;
 import org.eclipse.microprofile.metrics.annotation.Timed;
 
+import io.smallrye.metrics.MemberToMetricMappings;
 import io.smallrye.metrics.MetricProducer;
 import io.smallrye.metrics.MetricRegistries;
+import io.smallrye.metrics.MetricsRegistryImpl;
 import io.smallrye.metrics.MetricsRequestHandler;
 import io.smallrye.metrics.SmallRyeMetricsLogging;
+import io.smallrye.metrics.elementdesc.MemberInfo;
 import io.smallrye.metrics.elementdesc.adapter.BeanInfoAdapter;
 import io.smallrye.metrics.elementdesc.adapter.cdi.CDIBeanInfoAdapter;
 import io.smallrye.metrics.elementdesc.adapter.cdi.CDIMemberInfoAdapter;
@@ -82,7 +85,7 @@ public class MetricCdiInjectionExtension implements Extension {
 
     private final List<Class<?>> metricsInterfaces;
 
-    private final List<MetricID> metricIDs = new ArrayList<>();
+    private final Map<MemberInfo, List<MetricID>> registeredMetrics = new HashMap<>();
 
     public MetricCdiInjectionExtension() {
         metricsInterfaces = new ArrayList<>();
@@ -167,10 +170,10 @@ void registerMetrics(@Observes AfterDeploymentValidation adv, BeanManager manage
         for (Map.Entry<Bean<?>, List<AnnotatedMember<?>>> entry : metricsFromAnnotatedMethods.entrySet()) {
             Bean<?> bean = entry.getKey();
             for (AnnotatedMember<?> method : entry.getValue()) {
-                metricIDs.addAll(MetricsMetadata.registerMetrics(registry,
-                        resolver,
-                        beanInfoAdapter.convert(bean.getBeanClass()),
-                        memberInfoAdapter.convert(method.getJavaMember())));
+                final MemberInfo info = memberInfoAdapter.convert(method.getJavaMember());
+                final List<MetricID> metricIDs = MetricsMetadata.registerMetrics(registry, resolver,
+                        beanInfoAdapter.convert(bean.getBeanClass()), info);
+                registeredMetrics.put(info, metricIDs);
             }
         }
 
@@ -179,10 +182,10 @@ void registerMetrics(@Observes AfterDeploymentValidation adv, BeanManager manage
             for (Class<?> metricsInterface : metricsInterfaces) {
                 for (Method method : metricsInterface.getDeclaredMethods()) {
                     if (!method.isDefault() && !Modifier.isStatic(method.getModifiers())) {
-                        metricIDs.addAll(MetricsMetadata.registerMetrics(registry,
-                                resolver,
-                                beanInfoAdapter.convert(metricsInterface),
-                                memberInfoAdapter.convert(method)));
+                        final MemberInfo info = memberInfoAdapter.convert(method);
+                        final List<MetricID> metricIDs = MetricsMetadata.registerMetrics(registry, resolver,
+                                beanInfoAdapter.convert(metricsInterface), info);
+                        registeredMetrics.put(info, metricIDs);
                     }
                 }
             }
@@ -196,7 +199,22 @@ void registerMetrics(@Observes AfterDeploymentValidation adv, BeanManager manage
 
     void unregisterMetrics(@Observes BeforeShutdown shutdown) {
         MetricRegistry registry = MetricRegistries.get(MetricRegistry.Type.APPLICATION);
-        metricIDs.forEach(metricId -> registry.remove(metricId));
+
+        final MemberToMetricMappings memberToMetricMappings;
+        if (registry instanceof MetricsRegistryImpl) {
+            memberToMetricMappings = ((MetricsRegistryImpl) registry).getMemberToMetricMappings();
+        } else {
+            memberToMetricMappings = null;
+        }
+
+        registeredMetrics.forEach(((memberInfo, metricIDs) -> {
+            metricIDs.forEach(metricID -> {
+                registry.remove(metricID);
+                if (memberToMetricMappings != null) {
+                    memberToMetricMappings.removeMappingsFor(memberInfo, metricID);
+                }
+            });
+        }));
     }
 
     private Optional<String> getImplementationVersion() {