+ #21 Add LabelSanitizer in order to survive : in label names (e.g. from Vapor) (#22)

* + #21 Add LabelSanitizer in order to survive : in label names (e.g. from Vapor)

* linux tests generated

* fix docs comment, we're a statsd impl, not prom

* Simplify types; metric name sanitizer is just a function now

Author: ktoso

Sources/StatsdClient/StatsdClient.swift

@@ -32,9 +32,14 @@ public final class StatsdClient: MetricsFactory {
     ///     - eventLoopGroupProvider: The `EventLoopGroupProvider` to use, uses`createNew` strategy by default.
     ///     - host: The `statsd` server host.
     ///     - port: The `statsd` server port.
-    public init(eventLoopGroupProvider: EventLoopGroupProvider = .createNew, host: String, port: Int) throws {
+    public init(
+        eventLoopGroupProvider: EventLoopGroupProvider = .createNew,
+        host: String,
+        port: Int,
+        metricNameSanitizer: @escaping StatsdClient.MetricNameSanitizer = StatsdClient.defaultMetricNameSanitizer
+    ) throws {
         let address = try SocketAddress.makeAddressResolvingHost(host, port: port)
-        self.client = Client(eventLoopGroupProvider: eventLoopGroupProvider, address: address)
+        self.client = Client(eventLoopGroupProvider: eventLoopGroupProvider, address: address, metricNameSanitizer: metricNameSanitizer)
     }
 
     /// Shutdown the client. This is a noop when using a `shared` `EventLoopGroupProvider` strategy.
@@ -77,7 +82,7 @@ public final class StatsdClient: MetricsFactory {
     }
 
     private func make<Item>(label: String, dimensions: [(String, String)], registry: inout [String: Item], maker: (String, [(String, String)]) -> Item) -> Item {
-        let id = StatsdUtils.id(label: label, dimensions: dimensions)
+        let id = StatsdUtils.id(label: label, dimensions: dimensions, sanitizer: self.client.metricNameSanitizer)
         if let item = registry[id] {
             return item
         }
@@ -128,7 +133,7 @@ private final class StatsdCounter: CounterHandler, Equatable {
     var value = NIOAtomic<Int64>.makeAtomic(value: 0)
 
     init(label: String, dimensions: [(String, String)], client: Client) {
-        self.id = StatsdUtils.id(label: label, dimensions: dimensions)
+        self.id = StatsdUtils.id(label: label, dimensions: dimensions, sanitizer: client.metricNameSanitizer)
         self.client = client
     }
 
@@ -174,7 +179,7 @@ private final class StatsdRecorder: RecorderHandler, Equatable {
     let client: Client
 
     init(label: String, dimensions: [(String, String)], aggregate: Bool, client: Client) {
-        self.id = StatsdUtils.id(label: label, dimensions: dimensions)
+        self.id = StatsdUtils.id(label: label, dimensions: dimensions, sanitizer: client.metricNameSanitizer)
         self.aggregate = aggregate
         self.client = client
     }
@@ -219,7 +224,7 @@ private final class StatsdTimer: TimerHandler, Equatable {
     let client: Client
 
     init(label: String, dimensions: [(String, String)], client: Client) {
-        self.id = StatsdUtils.id(label: label, dimensions: dimensions)
+        self.id = StatsdUtils.id(label: label, dimensions: dimensions, sanitizer: client.metricNameSanitizer)
         self.client = client
     }
 
@@ -242,6 +247,8 @@ private final class Client {
     private let eventLoopGroupProvider: StatsdClient.EventLoopGroupProvider
     private let eventLoopGroup: EventLoopGroup
 
+    internal let metricNameSanitizer: StatsdClient.MetricNameSanitizer
+
     private let address: SocketAddress
 
     private let isShutdown = NIOAtomic<Bool>.makeAtomic(value: false)
@@ -255,14 +262,19 @@ private final class Client {
         case connected(Channel)
     }
 
-    init(eventLoopGroupProvider: StatsdClient.EventLoopGroupProvider, address: SocketAddress) {
+    init(
+        eventLoopGroupProvider: StatsdClient.EventLoopGroupProvider,
+        address: SocketAddress,
+        metricNameSanitizer: @escaping StatsdClient.MetricNameSanitizer
+    ) {
         self.eventLoopGroupProvider = eventLoopGroupProvider
         switch self.eventLoopGroupProvider {
         case .shared(let group):
             self.eventLoopGroup = group
         case .createNew:
             self.eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
         }
+        self.metricNameSanitizer = metricNameSanitizer
         self.address = address
     }
 
@@ -349,11 +361,43 @@ private final class Client {
     }
 }
 
+// MARK: - Metric Name Sanitizer
+
+extension StatsdClient {
+    /// Used to sanitize labels (and dimensions) into a format compatible with statsd's wire format.
+    ///
+    /// By default `StatsdClient` uses the `StatsdClient.defaultMetricNameSanitizer`.
+    public typealias MetricNameSanitizer = (String) -> String
+
+    /// Default implementation of `LabelSanitizer` that sanitizes any ":" occurrences by replacing them with a replacement character.
+    /// Defaults to replacing the illegal characters with "_", e.g. "offending:example" becomes "offending_example".
+    ///
+    /// See `https://github.com/b/statsd_spec` for more info.
+    public static let defaultMetricNameSanitizer: StatsdClient.MetricNameSanitizer = { label in
+        let illegalCharacter: Character = ":"
+        let replacementCharacter: Character = "_"
+
+        guard label.contains(illegalCharacter) else {
+            return label
+        }
+
+        // replacingOccurrences would be used, but is in Foundation which we try to not depend on here
+        return String(label.compactMap { (c: Character) -> Character? in
+            c != illegalCharacter ? c : replacementCharacter
+        })
+    }
+}
+
 // MARK: - Utility
 
 private enum StatsdUtils {
-    static func id(label: String, dimensions: [(String, String)]) -> String {
-        return dimensions.isEmpty ? label : dimensions.reduce(label) { a, b in "\(a).\(b.0).\(b.1)" }
+    static func id(label: String, dimensions: [(String, String)], sanitizer sanitize: StatsdClient.MetricNameSanitizer) -> String {
+        if dimensions.isEmpty {
+            return sanitize(label)
+        } else {
+            let labelWithDimensions = dimensions.reduce(label) { a, b in "\(a).\(b.0).\(b.1)" }
+            return sanitize(labelWithDimensions)
+        }
     }
 }
 

Tests/StatsdClientTests/StatsdClientTests+XCTest.swift

@@ -35,6 +35,7 @@ extension StatsdClientTests {
             ("testGaugeDouble", testGaugeDouble),
             ("testRecorderInteger", testRecorderInteger),
             ("testRecorderDouble", testRecorderDouble),
+            ("testLabelSanitizer", testLabelSanitizer),
             ("testCouncurrency", testCouncurrency),
             ("testNumberOfConnections", testNumberOfConnections),
         ]

Tests/StatsdClientTests/StatsdClientTests.swift

@@ -246,6 +246,27 @@ class StatsdClientTests: XCTestCase {
         assertTimeoutResult(semaphore.wait(timeout: .now() + .seconds(1)))
     }
 
+    func testLabelSanitizer() {
+        let server = TestServer(host: host, port: port)
+        XCTAssertNoThrow(try server.connect().wait())
+        defer { XCTAssertNoThrow(try server.shutdown()) }
+
+        let illegalID = "hello/:who"
+        let sanitizedID = "hello/_who"
+        let value = Double.random(in: 0 ... 1000)
+
+        let semaphore = DispatchSemaphore(value: 0)
+        server.onData { _, data in
+            defer { semaphore.signal() }
+            XCTAssertEqual(data, "\(sanitizedID):\(value)|h", "expected entries to match")
+        }
+
+        let recorder = Recorder(label: illegalID)
+        recorder.record(value)
+
+        assertTimeoutResult(semaphore.wait(timeout: .now() + .seconds(1)))
+    }
+
     func testCouncurrency() {
         let server = TestServer(host: host, port: port)
         XCTAssertNoThrow(try server.connect().wait())