Unregister the metric fully when last labelled metric is unregistered (#110)

Author: simonjbeaumont

Sources/Prometheus/PrometheusCollectorRegistry.swift

@@ -409,7 +409,11 @@ public final class PrometheusCollectorRegistry: Sendable {
                 let labelsKey = LabelsKey(counter.labels)
                 guard dimensions[labelsKey] === counter else { return }
                 dimensions.removeValue(forKey: labelsKey)
-                store[counter.name] = .counterWithLabels(labelNames, dimensions)
+                if dimensions.isEmpty {
+                    store.removeValue(forKey: counter.name)
+                } else {
+                    store[counter.name] = .counterWithLabels(labelNames, dimensions)
+                }
             default:
                 return
             }
@@ -431,7 +435,11 @@ public final class PrometheusCollectorRegistry: Sendable {
                 let dimensionsKey = LabelsKey(gauge.labels)
                 guard dimensions[dimensionsKey] === gauge else { return }
                 dimensions.removeValue(forKey: dimensionsKey)
-                store[gauge.name] = .gaugeWithLabels(labelNames, dimensions)
+                if dimensions.isEmpty {
+                    store.removeValue(forKey: gauge.name)
+                } else {
+                    store[gauge.name] = .gaugeWithLabels(labelNames, dimensions)
+                }
             default:
                 return
             }
@@ -453,7 +461,11 @@ public final class PrometheusCollectorRegistry: Sendable {
                 let dimensionsKey = LabelsKey(histogram.labels)
                 guard dimensions[dimensionsKey] === histogram else { return }
                 dimensions.removeValue(forKey: dimensionsKey)
-                store[histogram.name] = .durationHistogramWithLabels(labelNames, dimensions, buckets)
+                if dimensions.isEmpty {
+                    store.removeValue(forKey: histogram.name)
+                } else {
+                    store[histogram.name] = .durationHistogramWithLabels(labelNames, dimensions, buckets)
+                }
             default:
                 return
             }
@@ -475,7 +487,11 @@ public final class PrometheusCollectorRegistry: Sendable {
                 let dimensionsKey = LabelsKey(histogram.labels)
                 guard dimensions[dimensionsKey] === histogram else { return }
                 dimensions.removeValue(forKey: dimensionsKey)
-                store[histogram.name] = .valueHistogramWithLabels(labelNames, dimensions, buckets)
+                if dimensions.isEmpty {
+                    store.removeValue(forKey: histogram.name)
+                } else {
+                    store[histogram.name] = .valueHistogramWithLabels(labelNames, dimensions, buckets)
+                }
             default:
                 return
             }

Tests/PrometheusTests/PrometheusCollectorRegistryTests.swift

@@ -231,4 +231,34 @@ final class PrometheusCollectorRegistryTests: XCTestCase {
         )
     }
 
+    func testUnregisterReregisterWithoutLabels() {
+        let registry = PrometheusCollectorRegistry()
+        registry.unregisterCounter(registry.makeCounter(name: "name"))
+        registry.unregisterGauge(registry.makeGauge(name: "name"))
+        registry.unregisterDurationHistogram(registry.makeDurationHistogram(name: "name", buckets: []))
+        registry.unregisterValueHistogram(registry.makeValueHistogram(name: "name", buckets: []))
+        _ = registry.makeCounter(name: "name")
+    }
+
+    func testUnregisterReregisterWithLabels() {
+        let registry = PrometheusCollectorRegistry()
+
+        registry.unregisterCounter(registry.makeCounter(name: "name", labels: [("a", "1")]))
+        registry.unregisterCounter(registry.makeCounter(name: "name", labels: [("b", "1")]))
+
+        registry.unregisterGauge(registry.makeGauge(name: "name", labels: [("a", "1")]))
+        registry.unregisterGauge(registry.makeGauge(name: "name", labels: [("b", "1")]))
+
+        registry.unregisterDurationHistogram(
+            registry.makeDurationHistogram(name: "name", labels: [("a", "1")], buckets: [])
+        )
+        registry.unregisterDurationHistogram(
+            registry.makeDurationHistogram(name: "name", labels: [("b", "1")], buckets: [])
+        )
+
+        registry.unregisterValueHistogram(registry.makeValueHistogram(name: "name", labels: [("a", "1")], buckets: []))
+        registry.unregisterValueHistogram(registry.makeValueHistogram(name: "name", labels: [("b", "1")], buckets: []))
+
+        _ = registry.makeCounter(name: "name", labels: [("a", "1")])
+    }
 }

Tests/PrometheusTests/PrometheusMetricsFactoryTests.swift

@@ -67,8 +67,6 @@ final class PrometheusMetricsFactoryTests: XCTestCase {
         XCTAssertEqual(
             String(decoding: buffer, as: Unicode.UTF8.self),
             """
-            # TYPE foo counter
-
             """
         )
     }
@@ -126,8 +124,6 @@ final class PrometheusMetricsFactoryTests: XCTestCase {
         XCTAssertEqual(
             String(decoding: buffer, as: Unicode.UTF8.self),
             """
-            # TYPE foo gauge
-
             """
         )
     }