Feature: expose librdkafka statistics as swift metrics

Package.swift

@@ -48,9 +48,11 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-nio.git", from: "2.55.0"),
         .package(url: "https://github.com/swift-server/swift-service-lifecycle.git", from: "2.0.0-alpha.1"),
         .package(url: "https://github.com/apple/swift-log.git", from: "1.0.0"),
+        .package(url: "https://github.com/apple/swift-metrics", from: "2.4.1"),
         // The zstd Swift package produces warnings that we cannot resolve:
         // https://github.com/facebook/zstd/issues/3328
         .package(url: "https://github.com/facebook/zstd.git", from: "1.5.0"),
+        .package(url: "https://github.com/swift-extras/swift-extras-json.git", .upToNextMajor(from: "0.6.0")),
     ],
     targets: [
         .target(
@@ -80,6 +82,8 @@ let package = Package(
                 .product(name: "NIOCore", package: "swift-nio"),
                 .product(name: "ServiceLifecycle", package: "swift-service-lifecycle"),
                 .product(name: "Logging", package: "swift-log"),
+                .product(name: "Metrics", package: "swift-metrics"),
+                .product(name: "ExtrasJSON", package: "swift-extras-json"),
             ]
         ),
         .target(

Sources/Kafka/Configuration/KafkaConfiguration+Metrics.swift

@@ -0,0 +1,65 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the swift-kafka-client open source project
+//
+// Copyright (c) 2023 Apple Inc. and the swift-kafka-client project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of swift-kafka-client project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import Metrics
+
+extension KafkaConfiguration {
+    // MARK: - Metrics
+
+    /// Use to configure metrics.
+    public struct MetricsOptions: Sendable {
+        /// librdkafka's internal monotonic clock (microseconds)
+        public var ts: Gauge?
+        /// Wall clock time in seconds since the epoch
+        public var time: Timer?
+        /// Time since this client instance was created
+        public var age: Timer?
+        /// Number of ops (callbacks, events, etc) waiting in queue for application to serve
+        public var replyQueue: Gauge?
+        /// Current number of messages in producer queues
+        public var msgCount: Gauge?
+        /// Current total size of messages in producer queues
+        public var msgSize: Gauge?
+        /// Threshold: maximum number of messages allowed allowed on the producer queues
+        public var msgMax: Gauge?
+        /// Threshold: maximum total size of messages allowed on the producer queues
+        public var msgSizeMax: Gauge?
+
+        /// Total number of requests sent to Kafka brokers
+        public var tx: Gauge?
+        /// Total number of bytes transmitted to Kafka brokers
+        public var txBytex: Gauge?
+        /// Total number of responses received from Kafka brokers
+        public var rx: Gauge?
+        /// Total number of bytes received from Kafka brokers
+        public var rxBytex: Gauge?
+
+        /// Total number of messages transmitted (produced) to Kafka brokers
+        public var txMessages: Gauge?
+        /// Total number of message bytes (including framing, such as per-Message framing and MessageSet/batch framing) transmitted to Kafka brokers
+        public var txMessagesBytex: Gauge?
+        /// Total number of messages consumed, not including ignored messages (due to offset, etc), from Kafka brokers.
+        public var rxMessages: Gauge?
+        /// Total number of message bytes (including framing) received from Kafka brokers
+        public var rxMessagesBytex: Gauge?
+
+        /// Number of topics in the metadata cache.
+        public var metadataCacheCount: Gauge?
+    }
+
+    public enum Metrics: Sendable {
+        case disable
+        case enable(updateInterval: KafkaConfiguration.KeyRefreshAttempts, options: MetricsOptions)
+    }
+}

Sources/Kafka/Configuration/KafkaConfiguration.swift

@@ -283,4 +283,22 @@ public enum KafkaConfiguration {
         /// Use the IPv6 address family.
         public static let v6 = IPAddressFamily(description: "v6")
     }
+
+    /// Minimum time between key refresh attempts.
+    public struct KeyRefreshAttempts: Sendable, Hashable {
+        internal let rawValue: UInt
+
+        private init(rawValue: UInt) {
+            self.rawValue = rawValue
+        }
+
+        /// (Lowest granularity is milliseconds)
+        public static func value(_ value: Duration) -> KeyRefreshAttempts {
+            precondition(
+                value.canBeRepresentedAsMilliseconds,
+                "Lowest granularity is milliseconds"
+            )
+            return .init(rawValue: UInt(value.inMilliseconds))
+        }
+    }
 }

Sources/Kafka/Configuration/KafkaConsumerConfiguration.swift

@@ -207,6 +207,9 @@ public struct KafkaConsumerConfiguration {
     /// Reconnect options.
     public var reconnect: KafkaConfiguration.ReconnectOptions = .init()
 
+    /// Options for librdkafka metrics updates
+    public var metrics: KafkaConfiguration.Metrics = .disable
+
     /// Security protocol to use (plaintext, ssl, sasl_plaintext, sasl_ssl).
     /// Default: `.plaintext`
     public var securityProtocol: KafkaConfiguration.SecurityProtocol = .plaintext
@@ -270,6 +273,10 @@ extension KafkaConsumerConfiguration {
         resultDict["broker.address.family"] = broker.addressFamily.description
         resultDict["reconnect.backoff.ms"] = String(reconnect.backoff.rawValue)
         resultDict["reconnect.backoff.max.ms"] = String(reconnect.maximumBackoff.inMilliseconds)
+
+        if case .enable(let interval, _) = metrics {
+            resultDict["statistics.interval.ms"] = String(interval.rawValue)
+        }
 
         // Merge with SecurityProtocol configuration dictionary
         resultDict.merge(securityProtocol.dictionary) { _, _ in

Sources/Kafka/Configuration/KafkaProducerConfiguration.swift

@@ -161,6 +161,9 @@ public struct KafkaProducerConfiguration {
     /// Reconnect options.
     public var reconnect: KafkaConfiguration.ReconnectOptions = .init()
 
+    /// Options for librdkafka metrics updates
+    public var metrics: KafkaConfiguration.Metrics = .disable
+
     /// Security protocol to use (plaintext, ssl, sasl_plaintext, sasl_ssl).
     /// Default: `.plaintext`
     public var securityProtocol: KafkaConfiguration.SecurityProtocol = .plaintext
@@ -211,6 +214,10 @@ extension KafkaProducerConfiguration {
         resultDict["broker.address.family"] = self.broker.addressFamily.description
         resultDict["reconnect.backoff.ms"] = String(self.reconnect.backoff.rawValue)
         resultDict["reconnect.backoff.max.ms"] = String(self.reconnect.maximumBackoff.inMilliseconds)
+
+        if case .enable(let interval, _) = metrics {
+            resultDict["statistics.interval.ms"] = String(interval.rawValue)
+        }
 
         // Merge with SecurityProtocol configuration dictionary
         resultDict.merge(self.securityProtocol.dictionary) { _, _ in

Sources/Kafka/KafkaConsumer.swift

@@ -121,6 +121,12 @@ public final class KafkaConsumer: Sendable, Service {
         NIOAsyncSequenceProducerBackPressureStrategies.NoBackPressure,
         KafkaConsumerCloseOnTerminate
     >
+    typealias ProducerEvents = NIOAsyncSequenceProducer<
+        KafkaConsumerEvent,
+        NIOAsyncSequenceProducerBackPressureStrategies.NoBackPressure,
+        KafkaConsumerCloseOnTerminate
+    >
+
     /// The configuration object of the consumer client.
     private let configuration: KafkaConsumerConfiguration
     /// A logger.
@@ -241,6 +247,9 @@ public final class KafkaConsumer: Sendable, Service {
         if configuration.isAutoCommitEnabled == false {
             subscribedEvents.append(.offsetCommit)
         }
+        if case .enable = configuration.metrics {
+            subscribedEvents.append(.statistics)
+        }
 
         let client = try RDKafkaClient.makeClient(
             type: .consumer,
@@ -249,20 +258,21 @@ public final class KafkaConsumer: Sendable, Service {
             logger: logger
         )
 
-        let consumer = try KafkaConsumer(
-            client: client,
-            stateMachine: stateMachine,
-            configuration: configuration,
-            logger: logger
-        )
-
         let sourceAndSequence = NIOAsyncSequenceProducer.makeSequence(
             elementType: KafkaConsumerEvent.self,
             backPressureStrategy: NIOAsyncSequenceProducerBackPressureStrategies.NoBackPressure(),
             delegate: KafkaConsumerCloseOnTerminate(stateMachine: stateMachine)
         )
 
         let eventsSequence = KafkaConsumerEvents(wrappedSequence: sourceAndSequence.sequence)
+
+        let consumer = try KafkaConsumer(
+            client: client,
+            stateMachine: stateMachine,
+            configuration: configuration,
+            logger: logger
+        )
+
         return (consumer, eventsSequence)
     }
 
@@ -333,6 +343,10 @@ public final class KafkaConsumer: Sendable, Service {
                             source.finish()
                             throw error
                         }
+                    case .statistics(let statistics):
+                        if case let .enable(_, options) = self.configuration.metrics {
+                            statistics.fill(options)
+                        }
                     default:
                         break // Ignore
                     }

Sources/Kafka/KafkaConsumerEvent.swift

@@ -14,11 +14,16 @@
 
 /// An enumeration representing events that can be received through the ``KafkaConsumerEvents`` asynchronous sequence.
 public enum KafkaConsumerEvent: Sendable, Hashable {
+//    /// Statistics from librdkafka
+//    case statistics(KafkaStatistics)
     /// - Important: Always provide a `default` case when switiching over this `enum`.
     case DO_NOT_SWITCH_OVER_THIS_EXHAUSITVELY
 
     internal init(_ event: RDKafkaClient.KafkaEvent) {
         switch event {
+        case .statistics(let stat):
+            fatalError("Cannot cast \(event) to KafkaConsumerEvent")
+//            self = .statistics(stat)
         case .deliveryReport:
             fatalError("Cannot cast \(event) to KafkaConsumerEvent")
         case .consumerMessages:

Sources/Kafka/KafkaProducer.swift

@@ -162,11 +162,17 @@ public final class KafkaProducer: Service, Sendable {
             delegate: KafkaProducerCloseOnTerminate(stateMachine: stateMachine)
         )
         let source = sourceAndSequence.source
+
+        var subscribedEvents: [RDKafkaEvent] = [.log, .deliveryReport]
+        // Listen to statistics events when statistics enabled
+        if case .enable = configuration.metrics {
+            subscribedEvents.append(.statistics)
+        }
 
         let client = try RDKafkaClient.makeClient(
             type: .producer,
             configDictionary: configuration.dictionary,
-            events: [.log, .deliveryReport],
+            events: subscribedEvents,
             logger: logger
         )
 
@@ -208,6 +214,12 @@ public final class KafkaProducer: Service, Sendable {
             case .pollAndYield(let client, let source):
                 let events = client.eventPoll()
                 for event in events {
+                    if case let .statistics(kafkaStatistics) = event {
+                        if case let .enable(_, options) = self.configuration.metrics {
+                            kafkaStatistics.fill(options)
+                        }
+                        continue
+                    }
                     let producerEvent = KafkaProducerEvent(event)
                     // Ignore YieldResult as we don't support back pressure in KafkaProducer
                     _ = source?.yield(producerEvent)

Sources/Kafka/KafkaProducerEvent.swift

@@ -16,13 +16,18 @@
 public enum KafkaProducerEvent: Sendable, Hashable {
     /// A collection of delivery reports received from the Kafka cluster indicating the status of produced messages.
     case deliveryReports([KafkaDeliveryReport])
+//    /// Statistics from librdkafka
+//    case statistics(KafkaStatistics)
     /// - Important: Always provide a `default` case when switching over this `enum`.
     case DO_NOT_SWITCH_OVER_THIS_EXHAUSITVELY
 
     internal init(_ event: RDKafkaClient.KafkaEvent) {
         switch event {
         case .deliveryReport(results: let results):
             self = .deliveryReports(results)
+        case .statistics(let stat):
+//            self = .statistics(stat)
+            fatalError("Cannot cast \(event) to KafkaProducerEvent")
         case .consumerMessages:
             fatalError("Cannot cast \(event) to KafkaProducerEvent")
         }

Sources/Kafka/RDKafka/RDKafkaClient.swift

@@ -169,6 +169,7 @@ final class RDKafkaClient: Sendable {
     enum KafkaEvent {
         case deliveryReport(results: [KafkaDeliveryReport])
         case consumerMessages(result: Result<KafkaConsumerMessage, Error>)
+        case statistics(KafkaStatistics)
     }
 
     /// Poll the event `rd_kafka_queue_t` for new events.
@@ -199,6 +200,8 @@ final class RDKafkaClient: Sendable {
                 self.handleLogEvent(event)
             case .offsetCommit:
                 self.handleOffsetCommitEvent(event)
+            case .statistics:
+                events.append(self.handleStatistics(event))
             case .none:
                 // Finished reading events, return early
                 return events
@@ -250,6 +253,11 @@ final class RDKafkaClient: Sendable {
         // The returned message(s) MUST NOT be freed with rd_kafka_message_destroy().
     }
 
+    private func handleStatistics(_ event: OpaquePointer?) -> KafkaEvent {
+        let jsonStr = String(cString: rd_kafka_event_stats(event))
+        return .statistics(KafkaStatistics(jsonString: jsonStr))
+    }
+
     /// Handle event of type `RDKafkaEvent.log`.
     ///
     /// - Parameter event: Pointer to underlying `rd_kafka_event_t`.

Sources/Kafka/Utilities/KafkaStatistics.swift

@@ -0,0 +1,33 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the swift-kafka-gsoc open source project
+//
+// Copyright (c) 2023 Apple Inc. and the swift-kafka-gsoc project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of swift-kafka-gsoc project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import ExtrasJSON
+
+struct KafkaStatistics: Sendable, Hashable {
+    let jsonString: String
+
+    func fill(_ options: KafkaConfiguration.MetricsOptions) {
+        do {
+            let json = try XJSONDecoder().decode(KafkaStatisticsJson.self, from: self.jsonString.utf8)
+            if let age = options.age,
+               let jsonAge = json.age {
+                age.recordMicroseconds(jsonAge)
+            }
+
+            // TODO: other metrics
+        } catch {
+            fatalError("Statistics json decode error \(error)")
+        }
+    }
+}