Perf improvements (#50)

glbrntt · web-flow · commit 3caacd32263a · 2021-05-26T12:50:06.000+02:00
Motivation:

Metrics collecting systems should be efficient enough so that their
impact on the system being instrumented is negligible.

Modifications:

- Store `PromMetric`s in a dictionary keyed by label in
  `PrometheusClient`; this gives fast lookup when checking if a metric
  already exists
- Remove the `metricTypeMap`, we can recover the same information from
  `PromMetric`
- Remove calls to `getMetricInstrumet` in `PrometheusMetricsFactory`
  these were redundant as the same check is done in `createCounter` etc.
- In each `createCounter` (etc.) call we now hold the lock for the
  duration of the call to avoid races between checking the cache and
  creating and storing a new metric.
- The `PrometheusLabelSanitizer` now checks whether input needs
  santizing instead of santizing all input
- Sanitizing is done in one step, mapping each utf8 code point to a
  sanitized code point rather than lowercasing and then checking the
  validity of each character.

Result:

- Sanitizing pre-sanitized labels is ~100x faster
- Sanitizing non-sanitized label is ~20x faster
- Incrementing 10 counters is ~20x faster
- Incrementing 100 counters is ~40x faster
- Incrementing 1000 counters is ~250x faster
- (Similar results for other metrics.)
diff --git a/Sources/Prometheus/Prometheus.swift b/Sources/Prometheus/Prometheus.swift
@@ -7,19 +7,14 @@ import NIO
 public class PrometheusClient {
 
     /// Metrics tracked by this Prometheus instance
-    private var metrics: [PromMetric]
-    
-    /// To keep track of the type of a metric since  it can not change
-    /// througout the lifetime of the program
-    private var metricTypeMap: [String: PromMetricType]
+    private var metrics: [String: PromMetric]
     
     /// Lock used for thread safety
     private let lock: Lock
     
     /// Create a PrometheusClient instance
     public init() {
-        self.metrics = []
-        self.metricTypeMap = [:]
+        self.metrics = [:]
         self.lock = Lock()
     }
     
@@ -31,7 +26,7 @@ public class PrometheusClient {
     ///     - succeed: Closure that will be called with a newline separated string with metrics for all Metrics this PrometheusClient handles
     public func collect(_ succeed: (String) -> ()) {
         self.lock.withLock {
-            succeed(self.metrics.isEmpty ? "": "\(self.metrics.map { $0.collect() }.joined(separator: "\n"))\n")
+            succeed(self.metrics.isEmpty ? "": "\(self.metrics.values.map { $0.collect() }.joined(separator: "\n"))\n")
         }
     }
     
@@ -50,7 +45,7 @@ public class PrometheusClient {
     public func collect(_ succeed: (ByteBuffer) -> ()) {
         self.lock.withLock {
             var buffer = ByteBufferAllocator().buffer(capacity: 0)
-            self.metrics.forEach {
+            self.metrics.values.forEach {
                 $0.collect(into: &buffer)
                 buffer.writeString("\n")
             }
@@ -72,13 +67,21 @@ public class PrometheusClient {
         // `metricTypeMap` is left untouched as those must be consistent
         // throughout the lifetime of a program.
         return lock.withLock {
-            self.metrics.removeAll { $0._type == metric._type && $0.name == metric.name }
+            self.metrics.removeValue(forKey: metric.name)
         }
     }
     
-    public func getMetricInstance<T>(with name: String, andType type: PromMetricType) -> T? where T: PromMetric {
+    public func getMetricInstance<Metric>(with name: String, andType type: PromMetricType) -> Metric? where Metric: PromMetric {
         return lock.withLock {
-            self.metrics.compactMap { $0 as? T }.filter { $0.name == name && $0._type == type }.first
+            self._getMetricInstance(with: name, andType: type)
+        }
+    }
+    
+    private func _getMetricInstance<Metric>(with name: String, andType type: PromMetricType) -> Metric? where Metric: PromMetric {
+        if let metric = self.metrics[name], metric._type == type {
+            return metric as? Metric
+        } else {
+            return nil
         }
     }
     
@@ -101,17 +104,14 @@ public class PrometheusClient {
         initialValue: T = 0,
         withLabelType labelType: U.Type) -> PromCounter<T, U>
     {
-        if let counter: PromCounter<T, U> = getMetricInstance(with: name, andType: .counter) {
-            return counter
-        }
-
         return self.lock.withLock {
-            if let type = metricTypeMap[name] {
-                precondition(type == .counter, "Label \(name) was associated with \(type) before. Can not be used for a counter now.")
+            if let cachedCounter: PromCounter<T, U> = self._getMetricInstance(with: name, andType: .counter) {
+                return cachedCounter
             }
+
             let counter = PromCounter<T, U>(name, helpText, initialValue, self)
-            self.metricTypeMap[name] = .counter
-            self.metrics.append(counter)
+            let oldInstrument = self.metrics.updateValue(counter, forKey: name)
+            precondition(oldInstrument == nil, "Label \(oldInstrument!.name) is already associated with a \(oldInstrument!._type).")
             return counter
         }
     }
@@ -153,17 +153,14 @@ public class PrometheusClient {
         initialValue: T = 0,
         withLabelType labelType: U.Type) -> PromGauge<T, U>
     {
-        if let gauge: PromGauge<T, U> = getMetricInstance(with: name, andType: .gauge) {
-            return gauge
-        }
-
         return self.lock.withLock {
-            if let type = metricTypeMap[name] {
-                precondition(type == .gauge, "Label \(name) was associated with \(type) before. Can not be used for a gauge now.")
+            if let cachedGauge: PromGauge<T, U> = self._getMetricInstance(with: name, andType: .gauge) {
+                return cachedGauge
             }
+
             let gauge = PromGauge<T, U>(name, helpText, initialValue, self)
-            self.metricTypeMap[name] = .gauge
-            self.metrics.append(gauge)
+            let oldInstrument = self.metrics.updateValue(gauge, forKey: name)
+            precondition(oldInstrument == nil, "Label \(oldInstrument!.name) is already associated with a \(oldInstrument!._type).")
             return gauge
         }
     }
@@ -205,17 +202,14 @@ public class PrometheusClient {
         buckets: Buckets = .defaultBuckets,
         labels: U.Type) -> PromHistogram<T, U>
     {
-        if let histogram: PromHistogram<T, U> = getMetricInstance(with: name, andType: .histogram) {
-            return histogram
-        }
-
         return self.lock.withLock {
-            if let type = metricTypeMap[name] {
-                precondition(type == .histogram, "Label \(name) was associated with \(type) before. Can not be used for a histogram now.")
+            if let cachedHistogram: PromHistogram<T, U> = self._getMetricInstance(with: name, andType: .histogram) {
+                return cachedHistogram
             }
+
             let histogram = PromHistogram<T, U>(name, helpText, U(), buckets, self)
-            self.metricTypeMap[name] = .histogram
-            self.metrics.append(histogram)
+            let oldInstrument = self.metrics.updateValue(histogram, forKey: name)
+            precondition(oldInstrument == nil, "Label \(oldInstrument!.name) is already associated with a \(oldInstrument!._type).")
             return histogram
         }
     }
@@ -257,17 +251,14 @@ public class PrometheusClient {
         quantiles: [Double] = Prometheus.defaultQuantiles,
         labels: U.Type) -> PromSummary<T, U>
     {
-        if let summary: PromSummary<T, U> = getMetricInstance(with: name, andType: .summary) {
-            return summary
-        }
-
         return self.lock.withLock {
-            if let type = metricTypeMap[name] {
-                precondition(type == .summary, "Label \(name) was associated with \(type) before. Can not be used for a summary now.")
+            if let cachedSummary: PromSummary<T, U> = self._getMetricInstance(with: name, andType: .summary) {
+                return cachedSummary
             }
+            
             let summary = PromSummary<T, U>(name, helpText, U(), quantiles, self)
-            self.metricTypeMap[name] = .summary
-            self.metrics.append(summary)
+            let oldInstrument = self.metrics.updateValue(summary, forKey: name)
+            precondition(oldInstrument == nil, "Label \(oldInstrument!.name) is already associated with a \(oldInstrument!._type).")
             return summary
         }
     }
diff --git a/Sources/Prometheus/PrometheusMetrics.swift b/Sources/Prometheus/PrometheusMetrics.swift
@@ -128,14 +128,57 @@ public protocol LabelSanitizer {
 ///
 /// See `https://prometheus.io/docs/concepts/data_model/#metric-names-and-labels` for more info.
 public struct PrometheusLabelSanitizer: LabelSanitizer {
-    let allowedCharacters = "abcdefghijklmnopqrstuvwxyz0123456789_:"
-    
+    private static let uppercaseAThroughZ = UInt8(ascii: "A") ... UInt8(ascii: "Z")
+    private static let lowercaseAThroughZ = UInt8(ascii: "a") ... UInt8(ascii: "z")
+    private static let zeroThroughNine = UInt8(ascii: "0") ... UInt8(ascii: "9")
+
     public init() { }
 
     public func sanitize(_ label: String) -> String {
-        return String(label
-            .lowercased()
-            .map { (c: Character) -> Character in if allowedCharacters.contains(c) { return c }; return "_" })
+        if PrometheusLabelSanitizer.isSanitized(label) {
+            return label
+        } else {
+            return PrometheusLabelSanitizer.sanitizeLabel(label)
+        }
+    }
+
+    /// Returns a boolean indicating whether the label is already sanitized.
+    private static func isSanitized(_ label: String) -> Bool {
+        return label.utf8.allSatisfy(PrometheusLabelSanitizer.isValidCharacter(_:))
+    }
+    
+    /// Returns a boolean indicating whether the character may be used in a label.
+    private static func isValidCharacter(_ codePoint: String.UTF8View.Element) -> Bool {
+        switch codePoint {
+        case PrometheusLabelSanitizer.lowercaseAThroughZ,
+             PrometheusLabelSanitizer.zeroThroughNine,
+             UInt8(ascii: ":"),
+             UInt8(ascii: "_"):
+            return true
+        default:
+            return false
+        }
+    }
+
+    private static func sanitizeLabel(_ label: String) -> String {
+        let sanitized: [UInt8] = label.utf8.map { character in
+            if PrometheusLabelSanitizer.isValidCharacter(character) {
+                return character
+            } else {
+                return PrometheusLabelSanitizer.sanitizeCharacter(character)
+            }
+        }
+        
+        return String(decoding: sanitized, as: UTF8.self)
+    }
+    
+    private static func sanitizeCharacter(_ character: UInt8) -> UInt8 {
+        if PrometheusLabelSanitizer.uppercaseAThroughZ.contains(character) {
+            // Uppercase, so shift to lower case.
+            return character + (UInt8(ascii: "a") - UInt8(ascii: "A"))
+        } else {
+            return UInt8(ascii: "_")
+        }
     }
 }
 
@@ -191,13 +234,8 @@ public struct PrometheusMetricsFactory: PrometheusWrappedMetricsFactory {
     
     public func makeCounter(label: String, dimensions: [(String, String)]) -> CounterHandler {
         let label = configuration.labelSanitizer.sanitize(label)
-        let createHandler = { (counter: PromCounter) -> CounterHandler in
-            return MetricsCounter(counter: counter, dimensions: dimensions)
-        }
-        if let counter: PromCounter<Int64, DimensionLabels> = client.getMetricInstance(with: label, andType: .counter) {
-            return createHandler(counter)
-        }
-        return createHandler(client.createCounter(forType: Int64.self, named: label, withLabelType: DimensionLabels.self))
+        let counter = client.createCounter(forType: Int64.self, named: label, withLabelType: DimensionLabels.self)
+        return MetricsCounter(counter: counter, dimensions: dimensions)
     }
     
     public func makeRecorder(label: String, dimensions: [(String, String)], aggregate: Bool) -> RecorderHandler {
@@ -207,24 +245,14 @@ public struct PrometheusMetricsFactory: PrometheusWrappedMetricsFactory {
     
     private func makeGauge(label: String, dimensions: [(String, String)]) -> RecorderHandler {
         let label = configuration.labelSanitizer.sanitize(label)
-        let createHandler = { (gauge: PromGauge) -> RecorderHandler in
-            return MetricsGauge(gauge: gauge, dimensions: dimensions)
-        }
-        if let gauge: PromGauge<Double, DimensionLabels> = client.getMetricInstance(with: label, andType: .gauge) {
-            return createHandler(gauge)
-        }
-        return createHandler(client.createGauge(forType: Double.self, named: label, withLabelType: DimensionLabels.self))
+        let gauge = client.createGauge(forType: Double.self, named: label, withLabelType: DimensionLabels.self)
+        return MetricsGauge(gauge: gauge, dimensions: dimensions)
     }
     
     private func makeHistogram(label: String, dimensions: [(String, String)]) -> RecorderHandler {
         let label = configuration.labelSanitizer.sanitize(label)
-        let createHandler = { (histogram: PromHistogram) -> RecorderHandler in
-            return MetricsHistogram(histogram: histogram, dimensions: dimensions)
-        }
-        if let histogram: PromHistogram<Double, DimensionHistogramLabels> = client.getMetricInstance(with: label, andType: .histogram) {
-            return createHandler(histogram)
-        }
-        return createHandler(client.createHistogram(forType: Double.self, named: label, labels: DimensionHistogramLabels.self))
+        let histogram = client.createHistogram(forType: Double.self, named: label, labels: DimensionHistogramLabels.self)
+        return MetricsHistogram(histogram: histogram, dimensions: dimensions)
     }
     
     public func makeTimer(label: String, dimensions: [(String, String)]) -> TimerHandler {
@@ -240,26 +268,16 @@ public struct PrometheusMetricsFactory: PrometheusWrappedMetricsFactory {
     /// This method creates `Summary` backed timer implementation
     private func makeSummaryTimer(label: String, dimensions: [(String, String)], quantiles: [Double]) -> TimerHandler {
         let label = configuration.labelSanitizer.sanitize(label)
-        let createHandler = { (summary: PromSummary) -> TimerHandler in
-            return MetricsSummary(summary: summary, dimensions: dimensions)
-        }
-        if let summary: PromSummary<Int64, DimensionSummaryLabels> = client.getMetricInstance(with: label, andType: .summary) {
-            return createHandler(summary)
-        }
-        return createHandler(client.createSummary(forType: Int64.self, named: label, quantiles: quantiles, labels: DimensionSummaryLabels.self))
+        let summary = client.createSummary(forType: Int64.self, named: label, quantiles: quantiles, labels: DimensionSummaryLabels.self)
+        return MetricsSummary(summary: summary, dimensions: dimensions)
     }
 
     /// There's two different ways to back swift-api `Timer` with Prometheus classes.
     /// This method creates `Histogram` backed timer implementation
     private func makeHistogramTimer(label: String, dimensions: [(String, String)], buckets: Buckets) -> TimerHandler {
-        let createHandler = { (histogram: PromHistogram) -> TimerHandler in
-            MetricsHistogramTimer(histogram: histogram, dimensions: dimensions)
-        }
-        // PromHistogram should be reused when created for the same label, so we try to look it up
-        if let histogram: PromHistogram<Int64, DimensionHistogramLabels> = client.getMetricInstance(with: label, andType: .histogram) {
-            return createHandler(histogram)
-        }
-        return createHandler(client.createHistogram(forType: Int64.self, named: label, buckets: buckets, labels: DimensionHistogramLabels.self))
+        let label = configuration.labelSanitizer.sanitize(label)
+        let histogram = client.createHistogram(forType: Int64.self, named: label, buckets: buckets, labels: DimensionHistogramLabels.self)
+        return MetricsHistogramTimer(histogram: histogram, dimensions: dimensions)
     }
 }
 
diff --git a/Tests/SwiftPrometheusTests/PrometheusMetricsTests.swift b/Tests/SwiftPrometheusTests/PrometheusMetricsTests.swift
@@ -121,14 +121,19 @@ final class PrometheusMetricsTests: XCTestCase {
         prom.collect(promise.succeed)
         var buffer = try promise.futureResult.wait()
 
-        XCTAssertEqual(buffer.readString(length: buffer.readableBytes),
-            """
-            # TYPE my_counter counter
-            my_counter 10
-            my_counter{x="x", a="aaa"} 4
-            # TYPE my_gauge gauge
-            my_gauge 100.0\n
-            """)
+        let collected = buffer.readString(length: buffer.readableBytes)!
+
+        // We can't guarantee order so check the output contains the expected metrics.
+        XCTAssertTrue(collected.contains("""
+                                         # TYPE my_counter counter
+                                         my_counter 10
+                                         my_counter{x="x", a="aaa"} 4
+                                         """))
+
+        XCTAssertTrue(collected.contains("""
+                                         # TYPE my_gauge gauge
+                                         my_gauge 100.0
+                                         """))
     }
 
     func testCollectAFewMetricsIntoString() {
@@ -141,16 +146,19 @@ final class PrometheusMetricsTests: XCTestCase {
 
         let promise = self.eventLoop.makePromise(of: String.self)
         prom.collect(promise.succeed)
-        let string = try! promise.futureResult.wait()
-
-        XCTAssertEqual(string,
-            """
-            # TYPE my_counter counter
-            my_counter 10
-            my_counter{x="x", a="aaa"} 4
-            # TYPE my_gauge gauge
-            my_gauge 100.0\n
-            """)
+        let collected = try! promise.futureResult.wait()
+
+        // We can't guarantee order so check the output contains the expected metrics.
+        XCTAssertTrue(collected.contains("""
+                                         # TYPE my_counter counter
+                                         my_counter 10
+                                         my_counter{x="x", a="aaa"} 4
+                                         """))
+
+        XCTAssertTrue(collected.contains("""
+                                         # TYPE my_gauge gauge
+                                         my_gauge 100.0
+                                         """))
     }
 
     func testHistogramBackedTimer() {
