MBS-13416: rewrite statsd client

GitOrigin-RevId: 4d90c515a7d10df397da64464b3d1cac4782b247

Author: Vladislav Alekseev

Package.resolved

@@ -10,6 +10,15 @@
           "version": "1.0.21"
         }
       },
+      {
+        "package": "Socket",
+        "repositoryURL": "https://github.com/Kitura/BlueSocket.git",
+        "state": {
+          "branch": null,
+          "revision": "c46a3d41f5b2401d18bcb46d0101cdc5cd13e307",
+          "version": "1.0.52"
+        }
+      },
       {
         "package": "Glob",
         "repositoryURL": "https://github.com/Bouke/Glob",

Package.swift

@@ -49,6 +49,7 @@ let package = Package(
     dependencies: [
         .package(name: "Glob", url: "https://github.com/Bouke/Glob", .exact("1.0.5")),
         .package(name: "Signals", url: "https://github.com/IBM-Swift/BlueSignals.git", .exact("1.0.21")),
+        .package(name: "Socket", url: "https://github.com/Kitura/BlueSocket.git", .exact("1.0.52")),
         .package(name: "swift-argument-parser", url: "https://github.com/apple/swift-argument-parser", from: "1.0.3"),
     ],
     targets: [
@@ -374,16 +375,19 @@ let package = Package(
                 "AtomicModels",
                 "IO",
                 "MetricsUtils",
+                .product(name: "Socket", package: "Socket"),
                 "SocketModels",
-                "Waitable",
             ],
             path: "Sources/Statsd"
         ),
         .testTarget(
             name: "StatsdTests",
             dependencies: [
                 "MetricsRecording",
+                .product(name: "Socket", package: "Socket"),
+                "SocketModels",
                 "Statsd",
+                "TestHelpers",
             ],
             path: "Tests/StatsdTests"
         ),
@@ -427,6 +431,7 @@ let package = Package(
         .target(
             name: "TestHelpers",
             dependencies: [
+                "AtomicModels",
             ],
             path: "Tests/TestHelpers"
         ),

Sources/Statsd/AppleStatsdClient.swift

@@ -1,11 +1,16 @@
+import Dispatch
 import Foundation
 
 public protocol StatsdClient: AnyObject {
-    var stateUpdateHandler: ((StatsdClientState) -> ())? { get set }
-    var state: StatsdClientState { get }
+    func send(
+        content: Data,
+        queue: DispatchQueue,
+        completion: @escaping (Error?) -> ()
+    )
 
-    func start(queue: DispatchQueue)
-    func cancel()
-    
-    func send(content: Data, completion: @escaping (Error?) -> ())
+    func tearDown(
+        queue: DispatchQueue,
+        timeout: TimeInterval,
+        completion: @escaping () -> ()
+    )
 }

Sources/Statsd/StatsdClient.swift

@@ -0,0 +1,96 @@
+import AtomicModels
+import Dispatch
+import Foundation
+import Socket
+import SocketModels
+
+public final class StatsdClientImpl: StatsdClient {
+    private let statsdSocketAddress: SocketAddress
+    private let group = DispatchGroup()
+    private let tornDown = AtomicValue(false)
+    
+    public struct InvalidSocketAddressError: Error, CustomStringConvertible {
+        public let address: SocketAddress
+        
+        public var description: String {
+            "Invalid socket address: \(address)"
+        }
+    }
+    
+    public init(
+        statsdSocketAddress: SocketAddress
+    ) throws {
+        self.statsdSocketAddress = statsdSocketAddress
+    }
+    
+    public func send(
+        content: Data,
+        queue: DispatchQueue,
+        completion: @escaping (Error?) -> ()
+    ) {
+        guard tornDown.currentValue() == false else {
+            return
+        }
+        
+        guard let address = Socket.createAddress(
+            for: statsdSocketAddress.host,
+            on: Int32(statsdSocketAddress.port.value)
+        ) else {
+            queue.async { [statsdSocketAddress] in
+                completion(InvalidSocketAddressError(address: statsdSocketAddress))
+            }
+            return
+        }
+        
+        queue.async { [group] in
+            group.enter()
+            defer {
+                group.leave()
+            }
+            
+            do {
+                let socket = try Socket.create(
+                    type: .datagram,
+                    proto: .udp
+                )
+                defer {
+                    socket.close()
+                }
+                
+                try socket.write(
+                    from: content,
+                    to: address
+                )
+                completion(nil)
+            } catch {
+                completion(error)
+            }
+        }
+    }
+    
+    public func tearDown(
+        queue: DispatchQueue,
+        timeout: TimeInterval,
+        completion: @escaping () -> ()
+    ) {
+        self.tornDown.set(true)
+        
+        let tearDownCompleted = AtomicValue(false)
+        
+        let complete = {
+            tearDownCompleted.withExclusiveAccess {
+                if $0 == false {
+                    $0 = true
+                    completion()
+                }
+            }
+        }
+        
+        group.notify(queue: queue) {
+            complete()
+        }
+        queue.asyncAfter(deadline: .now() + timeout) {
+            complete()
+        }
+    }
+}

Sources/Statsd/StatsdClientImpl.swift

@@ -7,9 +7,7 @@ public final class StatsdMetricHandlerImpl: StatsdMetricHandler {
     private let statsdDomain: [String]
     private let statsdClient: StatsdClient
     private let serialQueue: DispatchQueue
-    
-    private let metricsBuffer = AtomicCollection([StatsdMetric]())
-    private let metricsBeingSent = AtomicCollection([UUID]())
+    private let group = DispatchGroup()
 
     public init(
         statsdDomain: [String],
@@ -19,63 +17,32 @@ public final class StatsdMetricHandlerImpl: StatsdMetricHandler {
         self.statsdDomain = statsdDomain
         self.statsdClient = statsdClient
         self.serialQueue = serialQueue
-        
-        self.statsdClient.stateUpdateHandler = { [weak self] state in
-            guard let strongSelf = self else { return }
-            
-            switch state {
-            case .notReady:
-                break
-            case .ready:
-                strongSelf.metricsBuffer.withExclusiveAccess {
-                    $0.forEach(strongSelf.send)
-                    $0.removeAll()
-                }
-            case .failed:
-                strongSelf.statsdClient.cancel()
-            }
-        }
-        statsdClient.start(queue: serialQueue)
     }
 
     public func handle(metric: StatsdMetric) {
-        // swiftlint:disable:next async
-        serialQueue.async { [weak self] in
-            guard let strongSelf = self else { return }
-            let state = strongSelf.statsdClient.state
-            
-            switch state {
-            case .notReady:
-                strongSelf.metricsBuffer.withExclusiveAccess { $0.append(metric) }
-            case .ready:
-                strongSelf.send(metric: metric)
-            case .failed:
-                break
-            }
+        let metricData = Data(metric.build(domain: self.statsdDomain).utf8)
+        
+        group.enter()
+        
+        statsdClient.send(
+            content: metricData,
+            queue: serialQueue
+        ) { [group] _ in
+            group.leave()
         }
     }
 
     public func tearDown(timeout: TimeInterval) {
-        let tearDownRequestTimestamp = Date()
-        _ = metricsBuffer.waitWhen(count: 0, before: tearDownRequestTimestamp.addingTimeInterval(timeout))
+        group.enter()
 
-        let timeoutRemainder = timeout - Date().timeIntervalSince(tearDownRequestTimestamp)
-        if timeoutRemainder > 0 {
-            _ = metricsBeingSent.waitWhen(count: 0, before: Date(timeIntervalSinceNow: timeoutRemainder))
-        }
-        
-        statsdClient.cancel()
-    }
-    
-    private func send(metric: StatsdMetric) {
-        let seed = UUID()
-        metricsBeingSent.withExclusiveAccess { $0.append(seed) }
-        self.statsdClient.send(
-            content: Data(metric.build(domain: self.statsdDomain).utf8)
-        ) { _ in
-            self.metricsBeingSent.withExclusiveAccess {
-                $0.removeAll(where: { $0 == seed })
+        statsdClient.tearDown(
+            queue: serialQueue,
+            timeout: timeout,
+            completion: { [group] in
+                group.leave()
             }
-        }
+        )
+        
+        _ = group.wait(timeout: .now() + timeout)
     }
 }

Sources/Statsd/StatsdClientState.swift

@@ -87,7 +87,7 @@ final class AppleEasyOutputStreamTests: XCTestCase {
 
         closeReadingEndAfterReadingSomeData(outputStreamProvider)
 
-        wait(for: [streamEndHandlerInvoked], timeout: 5)
+        wait(for: streamEndHandlerInvoked, timeout: 5)
         stream.close()
     }
 
@@ -110,7 +110,7 @@ final class AppleEasyOutputStreamTests: XCTestCase {
             try stream.enqueueWrite(data: data)
         )
 
-        wait(for: [streamErrorHandlerInvoked], timeout: 5)
+        wait(for: streamErrorHandlerInvoked, timeout: 5)
         stream.close()
     }