Support unregistering all metrics in a sensor (#21)

* Support unregistering all metrics in a sensor

* Address a comment

gradle.properties

@@ -1,6 +1,6 @@
 group=io.tehuti
 archivesBaseName=tehuti
-version=0.8.8
+version=0.8.9
 
 signing.enabled=false
 signing.keyId=

src/main/java/io/tehuti/metrics/JmxReporter.java

@@ -98,8 +98,9 @@ public void removeMetric(TehutiMetric metric) {
         synchronized (lock) {
             String[] names = Utils.splitMetricName(prefix + metric.name());
             String qualifiedName = names[0] + "." + names[1];
-            if (mbeans.remove(qualifiedName) != null) {
-                unregister(mbeans.get(qualifiedName));
+            TehutiMbean mbean = mbeans.remove(qualifiedName);
+            if (mbean != null) {
+                unregister(mbean);
             }
         }
     }

src/main/java/io/tehuti/metrics/MetricsRepository.java

@@ -137,6 +137,17 @@ public synchronized Sensor sensor(String name, MetricConfig config, Sensor... pa
         return s;
     }
 
+    /**
+     * Remove a sensor with the given unique name. Unregister all metrics with this sensor too.
+     * @param name The name of the sensor
+     */
+    public synchronized void removeSensor(String name) {
+        Sensor s = this.sensors.remove(name);
+        if (s != null) {
+            s.removeAll();
+        }
+    }
+
     /**
      * Add a metric to monitor an object that implements measurable. This metric won't be associated with any sensor.
      * This is a way to expose existing values as metrics.
@@ -213,6 +224,14 @@ synchronized void registerMetric(TehutiMetric metric) {
             reporter.addMetric(metric);
     }
 
+    synchronized void unregisterMetric(TehutiMetric metric) {
+        if (this.metrics.remove(metric.name()) != null) {
+            for (MetricsReporter reporter : reporters) {
+                reporter.removeMetric(metric);
+            }
+        }
+    }
+
     /**
      * Get all the metrics currently maintained indexed by metric name
      * @return A map of all metrics in this metricsRepository

src/main/java/io/tehuti/metrics/Sensor.java

@@ -216,6 +216,18 @@ public synchronized Metric add(String name, String description, MeasurableStat s
         return metric;
     }
 
+    /**
+     * Unregister all metrics with this sensor
+     */
+    synchronized void removeAll() {
+        for (TehutiMetric metric : this.metrics) {
+            this.registry.unregisterMetric(metric);
+        }
+        this.metrics.clear();
+        this.stats.clear();
+        this.statConfigs.clear();
+    }
+
     synchronized List<? extends Metric> metrics() {
         return Collections.unmodifiableList(this.metrics);
     }

src/test/java/io/tehuti/metrics/MetricsTest.java

@@ -3,16 +3,18 @@
  * file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file
  * to You under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the
  * License. You may obtain a copy of the License at
- * 
+ *
  * http://www.apache.org/licenses/LICENSE-2.0
- * 
+ *
  * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
  * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
  * specific language governing permissions and limitations under the License.
  */
 package io.tehuti.metrics;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.fail;
 
 import io.tehuti.utils.Time;
@@ -265,6 +267,25 @@ public void testPercentiles() {
         assertEquals(0.0, p75.value(), 1.0);
     }
 
+    @Test
+    public void testRemoveMetrics() {
+        Sensor sensor = metricsRepository.sensor("test");
+        sensor.add("test.avg", new Avg());
+        sensor.add("test.count", new Count());
+        for (int i = 0; i < 10; i++)
+            sensor.record(i);
+
+        assertNotNull(metricsRepository.getSensor(sensor.name()));
+        assertEquals("Avg(0...9) = 4.5", 4.5, metricsRepository.getMetric("test.avg").value(), EPS);
+        assertEquals("Count(0...9) = 10", 10.0, metricsRepository.getMetric("test.count").value(), EPS);
+
+        metricsRepository.removeSensor(sensor.name());
+
+        assertNull(metricsRepository.getMetric("test.avg"));
+        assertNull(metricsRepository.getMetric("test.count"));
+        assertNull(metricsRepository.getSensor(sensor.name()));
+    }
+
     public static class ConstantMeasurable implements Measurable {
         public double value = 0.0;