refactor: uniquely identify Counters by name and label sets

Previously, counters could only be identified by their name. This change allows multiple counters to share the same name, with each unique set of labels defining a distinct time series.

Additionally, the internal data structure for a stored metric has been refactored, providing a more robust and programmatic representation.

Signed-off-by: Melissa Kilby <[email protected]>

Author: incertum

Sources/Prometheus/PrometheusCollectorRegistry.swift

@@ -47,9 +47,14 @@ public final class PrometheusCollectorRegistry: Sendable {
         }
     }
 
-    private enum Metric {
-        case counter(Counter, help: String)
-        case counterWithLabels([String], [LabelsKey: Counter], help: String)
+    private typealias Metric<T> = (metric: T, help: String)
+
+    private struct MetricContainer<T> {
+        var metrics: [LabelsKey: Metric<T>]
+    }
+
+    private enum MetricType {
+        case counterMetrics(String, MetricContainer<Counter>)
         case gauge(Gauge, help: String)
         case gaugeWithLabels([String], [LabelsKey: Gauge], help: String)
         case durationHistogram(DurationHistogram, help: String)
@@ -58,7 +63,7 @@ public final class PrometheusCollectorRegistry: Sendable {
         case valueHistogramWithLabels([String], [LabelsKey: ValueHistogram], [Double], help: String)
     }
 
-    private let box = NIOLockedValueBox([String: Metric]())
+    private let box = NIOLockedValueBox([String: MetricType]())
 
     /// Create a new collector registry
     public init() {}
@@ -75,25 +80,7 @@ public final class PrometheusCollectorRegistry: Sendable {
     ///                   If the parameter is omitted or an empty string is passed, the `# HELP` line will not be generated for this metric.
     /// - Returns: A ``Counter`` that is registered with this ``PrometheusCollectorRegistry``
     public func makeCounter(name: String, help: String) -> Counter {
-        let name = name.ensureValidMetricName()
-        let help = help.ensureValidHelpText()
-        return self.box.withLockedValue { store -> Counter in
-            guard let value = store[name] else {
-                let counter = Counter(name: name, labels: [])
-                store[name] = .counter(counter, help: help)
-                return counter
-            }
-            guard case .counter(let counter, _) = value else {
-                fatalError(
-                    """
-                    Could not make Counter with name: \(name), since another metric type
-                    already exists for the same name.
-                    """
-                )
-            }
-
-            return counter
-        }
+        return self.makeCounter(name: name, labels: [], help: help)
     }
 
     /// Creates a new ``Counter`` collector or returns the already existing one with the same name.
@@ -104,7 +91,7 @@ public final class PrometheusCollectorRegistry: Sendable {
     /// - Parameter name: A name to identify ``Counter``'s value.
     /// - Returns: A ``Counter`` that is registered with this ``PrometheusCollectorRegistry``
     public func makeCounter(name: String) -> Counter {
-        return self.makeCounter(name: name, help: "")
+        return self.makeCounter(name: name, labels: [], help: "")
     }
 
     /// Creates a new ``Counter`` collector or returns the already existing one with the same name,
@@ -116,7 +103,7 @@ public final class PrometheusCollectorRegistry: Sendable {
     /// - Parameter descriptor: An ``MetricNameDescriptor`` that provides the fully qualified name for the metric.
     /// - Returns: A ``Counter`` that is registered with this ``PrometheusCollectorRegistry``
     public func makeCounter(descriptor: MetricNameDescriptor) -> Counter {
-        return self.makeCounter(name: descriptor.name, help: descriptor.helpText ?? "")
+        return self.makeCounter(name: descriptor.name, labels: [], help: descriptor.helpText ?? "")
     }
 
     /// Creates a new ``Counter`` collector or returns the already existing one with the same name.
@@ -131,51 +118,49 @@ public final class PrometheusCollectorRegistry: Sendable {
     ///                   If the parameter is omitted or an empty string is passed, the `# HELP` line will not be generated for this metric.
     /// - Returns: A ``Counter`` that is registered with this ``PrometheusCollectorRegistry``
     public func makeCounter(name: String, labels: [(String, String)], help: String) -> Counter {
-        guard !labels.isEmpty else {
-            return self.makeCounter(name: name, help: help)
-        }
-
         let name = name.ensureValidMetricName()
         let labels = labels.ensureValidLabelNames()
         let help = help.ensureValidHelpText()
+        let key = LabelsKey(labels)
 
         return self.box.withLockedValue { store -> Counter in
-            guard let value = store[name] else {
-                let labelNames = labels.allLabelNames
+            guard let entry = store[name] else {
+                // First time a Counter is registered with this name.
                 let counter = Counter(name: name, labels: labels)
-
-                store[name] = .counterWithLabels(labelNames, [LabelsKey(labels): counter], help: help)
-                return counter
-            }
-            guard case .counterWithLabels(let labelNames, var dimensionLookup, let help) = value else {
-                fatalError(
-                    """
-                    Could not make Counter with name: \(name) and labels: \(labels), since another
-                    metric type already exists for the same name.
-                    """
+                let newMetric: Metric<Counter> = (metric: counter, help: help)
+                let newContainer = MetricContainer<Counter>(
+                    metrics: [key: newMetric]
                 )
-            }
-
-            let key = LabelsKey(labels)
-            if let counter = dimensionLookup[key] {
+                store[name] = .counterMetrics(name, newContainer)
                 return counter
             }
+            switch entry {
+            case .counterMetrics(let existingName, var MetricContainer):
+                if let existingMetric = MetricContainer.metrics[key] {
+                    return existingMetric.metric
+                }
 
-            // check if all labels match the already existing ones.
-            if labelNames != labels.allLabelNames {
+                // Even if the metric name is identical, each label set defines a unique time series
+                let newCounter = Counter(name: name, labels: labels)
+                let netNewMetric: Metric<Counter> = (metric: newCounter, help: help)
+                MetricContainer.metrics[key] = netNewMetric
+
+                // Write the modified entry back to the store.
+                store[name] = .counterMetrics(existingName, MetricContainer)
+
+                return newCounter
+
+            default:
+                // A metric with this name exists, but it's not a Counter. This is a programming error.
+                // While Prometheus wouldn't stop you, it may result in unpredictable behavior with tools like Grafana or PromQL.
                 fatalError(
                     """
-                    Could not make Counter with name: \(name) and labels: \(labels), since the
-                    label names don't match the label names of previously registered Counters with
-                    the same name.
+                    Metric type mismatch:
+                    Could not register a Counter with name '\(name)',
+                    since a different metric type (\(entry.self)) was already registered with this name.
                     """
                 )
             }
-
-            let counter = Counter(name: name, labels: labels)
-            dimensionLookup[key] = counter
-            store[name] = .counterWithLabels(labelNames, dimensionLookup, help: help)
-            return counter
         }
     }
 
@@ -703,17 +688,19 @@ public final class PrometheusCollectorRegistry: Sendable {
     public func unregisterCounter(_ counter: Counter) {
         self.box.withLockedValue { store in
             switch store[counter.name] {
-            case .counter(let storedCounter, _):
-                guard storedCounter === counter else { return }
-                store.removeValue(forKey: counter.name)
-            case .counterWithLabels(let labelNames, var dimensions, let help):
-                let labelsKey = LabelsKey(counter.labels)
-                guard dimensions[labelsKey] === counter else { return }
-                dimensions.removeValue(forKey: labelsKey)
-                if dimensions.isEmpty {
-                    store.removeValue(forKey: counter.name)
+            case .counterMetrics(let name, var MetricContainer):
+                let key = LabelsKey(counter.labels)
+                guard let existingMetric = MetricContainer.metrics[key],
+                    existingMetric.metric === counter
+                else {
+                    return
+                }
+                MetricContainer.metrics.removeValue(forKey: key)
+
+                if MetricContainer.metrics.isEmpty {
+                    store.removeValue(forKey: name)
                 } else {
-                    store[counter.name] = .counterWithLabels(labelNames, dimensions, help: help)
+                    store[name] = .counterMetrics(name, MetricContainer)
                 }
             default:
                 return
@@ -806,52 +793,52 @@ public final class PrometheusCollectorRegistry: Sendable {
         let prefixHelp = "HELP"
         let prefixType = "TYPE"
 
-        for (label, metric) in metrics {
+        for (name, metric) in metrics {
             switch metric {
-            case .counter(let counter, let help):
-                help.isEmpty ? () : buffer.addLine(prefix: prefixHelp, name: label, value: help)
-                buffer.addLine(prefix: prefixType, name: label, value: "counter")
-                counter.emit(into: &buffer)
-
-            case .counterWithLabels(_, let counters, let help):
-                help.isEmpty ? () : buffer.addLine(prefix: prefixHelp, name: label, value: help)
-                buffer.addLine(prefix: prefixType, name: label, value: "counter")
-                for counter in counters.values {
-                    counter.emit(into: &buffer)
+            case .counterMetrics(_, let MetricContainer):
+                // An entry should not be empty, as a safeguard skip if it is.
+                guard let _ = MetricContainer.metrics.first?.value else {
+                    continue
+                }
+                for Metric in MetricContainer.metrics.values {
+                    let help = Metric.help
+                    help.isEmpty ? () : buffer.addLine(prefix: prefixHelp, name: name, value: help)
+                    buffer.addLine(prefix: prefixType, name: name, value: "counter")
+                    Metric.metric.emit(into: &buffer)
                 }
 
             case .gauge(let gauge, let help):
-                help.isEmpty ? () : buffer.addLine(prefix: prefixHelp, name: label, value: help)
-                buffer.addLine(prefix: prefixType, name: label, value: "gauge")
+                help.isEmpty ? () : buffer.addLine(prefix: prefixHelp, name: name, value: help)
+                buffer.addLine(prefix: prefixType, name: name, value: "gauge")
                 gauge.emit(into: &buffer)
 
             case .gaugeWithLabels(_, let gauges, let help):
-                help.isEmpty ? () : buffer.addLine(prefix: prefixHelp, name: label, value: help)
-                buffer.addLine(prefix: prefixType, name: label, value: "gauge")
+                help.isEmpty ? () : buffer.addLine(prefix: prefixHelp, name: name, value: help)
+                buffer.addLine(prefix: prefixType, name: name, value: "gauge")
                 for gauge in gauges.values {
                     gauge.emit(into: &buffer)
                 }
 
             case .durationHistogram(let histogram, let help):
-                help.isEmpty ? () : buffer.addLine(prefix: prefixHelp, name: label, value: help)
-                buffer.addLine(prefix: prefixType, name: label, value: "histogram")
+                help.isEmpty ? () : buffer.addLine(prefix: prefixHelp, name: name, value: help)
+                buffer.addLine(prefix: prefixType, name: name, value: "histogram")
                 histogram.emit(into: &buffer)
 
             case .durationHistogramWithLabels(_, let histograms, _, let help):
-                help.isEmpty ? () : buffer.addLine(prefix: prefixHelp, name: label, value: help)
-                buffer.addLine(prefix: prefixType, name: label, value: "histogram")
+                help.isEmpty ? () : buffer.addLine(prefix: prefixHelp, name: name, value: help)
+                buffer.addLine(prefix: prefixType, name: name, value: "histogram")
                 for histogram in histograms.values {
                     histogram.emit(into: &buffer)
                 }
 
             case .valueHistogram(let histogram, let help):
-                help.isEmpty ? () : buffer.addLine(prefix: prefixHelp, name: label, value: help)
-                buffer.addLine(prefix: prefixType, name: label, value: "histogram")
+                help.isEmpty ? () : buffer.addLine(prefix: prefixHelp, name: name, value: help)
+                buffer.addLine(prefix: prefixType, name: name, value: "histogram")
                 histogram.emit(into: &buffer)
 
             case .valueHistogramWithLabels(_, let histograms, _, let help):
-                help.isEmpty ? () : buffer.addLine(prefix: prefixHelp, name: label, value: help)
-                buffer.addLine(prefix: prefixType, name: label, value: "histogram")
+                help.isEmpty ? () : buffer.addLine(prefix: prefixHelp, name: name, value: help)
+                buffer.addLine(prefix: prefixType, name: name, value: "histogram")
                 for histogram in histograms.values {
                     histogram.emit(into: &buffer)
                 }

Tests/PrometheusTests/CounterTests.swift

@@ -165,6 +165,107 @@ final class CounterTests: XCTestCase {
         )
     }
 
+    func testCounterWithSharedMetricNamDistinctLabelSets() {
+        let client = PrometheusCollectorRegistry()
+
+        let counter0 = client.makeCounter(
+            name: "foo",
+            labels: [],
+            help: "Base metric name with no labels"
+        )
+
+        let counter1 = client.makeCounter(
+            name: "foo",
+            labels: [("bar", "baz")],
+            help: "Base metric name with one label set variant"
+        )
+
+        let counter2 = client.makeCounter(
+            name: "foo",
+            labels: [("bar", "newBaz"), ("newKey1", "newValue1")],
+            help: "Base metric name with a different label set variant"
+        )
+
+        var buffer = [UInt8]()
+        counter0.increment()
+        counter1.increment(by: Int64(9))
+        counter2.increment(by: Int64(4))
+        counter1.increment(by: Int64(3))
+        counter0.increment()
+        counter2.increment(by: Int64(20))
+        client.emit(into: &buffer)
+        var outputString = String(decoding: buffer, as: Unicode.UTF8.self)
+        var actualLines = Set(outputString.components(separatedBy: .newlines).filter { !$0.isEmpty })
+        var expectedLines = Set([
+            "# HELP foo Base metric name with no labels",
+            "# TYPE foo counter",
+            "foo 2",
+
+            "# HELP foo Base metric name with one label set variant",
+            "# TYPE foo counter",
+            #"foo{bar="baz"} 12"#,
+
+            "# HELP foo Base metric name with a different label set variant",
+            "# TYPE foo counter",
+            #"foo{bar="newBaz",newKey1="newValue1"} 24"#,
+        ])
+        XCTAssertEqual(actualLines, expectedLines)
+
+        // Counters are unregistered in a cascade.
+        client.unregisterCounter(counter0)
+        buffer.removeAll(keepingCapacity: true)
+        client.emit(into: &buffer)
+        outputString = String(decoding: buffer, as: Unicode.UTF8.self)
+        actualLines = Set(outputString.components(separatedBy: .newlines).filter { !$0.isEmpty })
+        expectedLines = Set([
+            "# HELP foo Base metric name with one label set variant",
+            "# TYPE foo counter",
+            #"foo{bar="baz"} 12"#,
+
+            "# HELP foo Base metric name with a different label set variant",
+            "# TYPE foo counter",
+            #"foo{bar="newBaz",newKey1="newValue1"} 24"#,
+        ])
+        XCTAssertEqual(actualLines, expectedLines)
+
+        client.unregisterCounter(counter1)
+        buffer.removeAll(keepingCapacity: true)
+        client.emit(into: &buffer)
+        outputString = String(decoding: buffer, as: Unicode.UTF8.self)
+        actualLines = Set(outputString.components(separatedBy: .newlines).filter { !$0.isEmpty })
+        expectedLines = Set([
+            "# HELP foo Base metric name with a different label set variant",
+            "# TYPE foo counter",
+            #"foo{bar="newBaz",newKey1="newValue1"} 24"#,
+        ])
+        XCTAssertEqual(actualLines, expectedLines)
+
+        client.unregisterCounter(counter2)
+        buffer.removeAll(keepingCapacity: true)
+        client.emit(into: &buffer)
+        outputString = String(decoding: buffer, as: Unicode.UTF8.self)
+        actualLines = Set(outputString.components(separatedBy: .newlines).filter { !$0.isEmpty })
+        expectedLines = Set([])
+        XCTAssertEqual(actualLines, expectedLines)
+
+        let counterGaugeSameName = client.makeGauge(
+            name: "foo",
+            labels: [],
+            help: "Base metric name used for new metric of type gauge"
+        )
+        buffer.removeAll(keepingCapacity: true)
+        client.emit(into: &buffer)
+        XCTAssertEqual(
+            String(decoding: buffer, as: Unicode.UTF8.self),
+            """
+            # HELP foo Base metric name used for new metric of type gauge
+            # TYPE foo gauge
+            foo 0.0
+
+            """
+        )
+    }
+
     func testWithMetricNameDescriptorWithFullComponentMatrix() {
         // --- Test Constants ---
         let helpTextValue = "https://help.url/sub"