diff --git a/Sources/Kafka/Configuration/KafkaConsumerConfig.swift b/Sources/Kafka/Configuration/KafkaConsumerConfig.swift index 5eb64e74..e8c42f84 100644 --- a/Sources/Kafka/Configuration/KafkaConsumerConfig.swift +++ b/Sources/Kafka/Configuration/KafkaConsumerConfig.swift @@ -34,6 +34,24 @@ public struct KafkaConsumerConfig: Sendable { /// Consumer metrics configuration. public var metrics: KafkaConfiguration.ConsumerMetrics = .init() + /// Low watermark for the ``KafkaConsumerEvents`` asynchronous sequence backpressure. + /// When the number of buffered events falls below this threshold, the consumer + /// will resume yielding events to the sequence. + /// Default: `100` + public var eventsBackpressureLowWatermark: Int = 100 + + /// High watermark for the ``KafkaConsumerEvents`` asynchronous sequence backpressure. + /// When the number of buffered events reaches this threshold, the consumer + /// will temporarily stop yielding events to the sequence to prevent out-of-memory errors. + /// Default: `1000` + public var eventsBackpressureHighWatermark: Int = 1000 + + /// Arbitrary configuration properties passed directly to librdkafka. + /// + /// - Warning: Properties set here override typed properties. + /// Intended for testing or advanced configurations not explicitly supported by this library. + internal var additionalConfig: [String: String] = [:] + // MARK: - Properties generated from librdkafka config list /// Client identifier. @@ -1046,14 +1064,6 @@ public struct KafkaConsumerConfig: Sendable { /// librdkafka property name: `"consume.callback.max.messages"` public var consumeCallbackMaxMessages: Int? - /// Additional librdkafka configuration properties not covered by typed properties. - /// Keys and values are passed directly to librdkafka. - /// - /// - Warning: Properties set here override typed properties above. - /// Intended for testing (e.g. `test.mock.num.brokers`) or advanced configurations - /// not explicitly supported by this library. - internal var additionalConfig: [String: String] = [:] - public init() {} internal var config: [String: String] { @@ -1193,9 +1203,7 @@ public struct KafkaConsumerConfig: Sendable { config["statistics.interval.ms"] = String(updateInterval.inMilliseconds) } - for (key, value) in self.additionalConfig { - config[key] = value - } + config.merge(self.additionalConfig) { _, new in new } return config } diff --git a/Sources/Kafka/Configuration/KafkaConsumerConfig.swift.gyb b/Sources/Kafka/Configuration/KafkaConsumerConfig.swift.gyb index 037daf92..9fda76ce 100644 --- a/Sources/Kafka/Configuration/KafkaConsumerConfig.swift.gyb +++ b/Sources/Kafka/Configuration/KafkaConsumerConfig.swift.gyb @@ -32,6 +32,24 @@ public struct KafkaConsumerConfig: Sendable { /// Consumer metrics configuration. public var metrics: KafkaConfiguration.ConsumerMetrics = .init() + /// Low watermark for the ``KafkaConsumerEvents`` asynchronous sequence backpressure. + /// When the number of buffered events falls below this threshold, the consumer + /// will resume yielding events to the sequence. + /// Default: `100` + public var eventsBackpressureLowWatermark: Int = 100 + + /// High watermark for the ``KafkaConsumerEvents`` asynchronous sequence backpressure. + /// When the number of buffered events reaches this threshold, the consumer + /// will temporarily stop yielding events to the sequence to prevent out-of-memory errors. + /// Default: `1000` + public var eventsBackpressureHighWatermark: Int = 1000 + + /// Arbitrary configuration properties passed directly to librdkafka. + /// + /// - Warning: Properties set here override typed properties. + /// Intended for testing or advanced configurations not explicitly supported by this library. + internal var additionalConfig: [String: String] = [:] + // MARK: - Properties generated from librdkafka config list % for rdconf in consumer_librdkafka_configs: @@ -54,7 +72,7 @@ public struct KafkaConsumerConfig: Sendable { /// /// Default: ${rdconf["defaultValue"]} % end - public var ${rdconf["swiftName"]}: ${rdconf["swiftType"].replace("enum", f"KafkaConfig.{rdconf["swiftEnumName"]}")}? + public var ${rdconf["swiftName"]}: ${rdconf["swiftType"].replace("enum", f"KafkaConfig.{rdconf['swiftEnumName']}")}? % end public init() {} @@ -98,6 +116,8 @@ public struct KafkaConsumerConfig: Sendable { config["statistics.interval.ms"] = String(updateInterval.inMilliseconds) } + config.merge(self.additionalConfig) { _, new in new } + return config } @@ -115,9 +135,9 @@ public struct KafkaConsumerConfig: Sendable { % elif rdconf["swiftType"] == "[String]": self.${rdconf["swiftName"]} = configDict["${rdconf["name"]}"]?.components(separatedBy: ",")${" // ignore-unacceptable-language" if rdconf["name"] == "topic.blacklist" else ""} % elif rdconf["swiftType"] == "enum": - self.${rdconf["swiftName"]} = configDict["${rdconf["name"]}"].map(KafkaConfig.${rdconf["swiftEnumName"]}.init) + self.${rdconf["swiftName"]} = configDict["${rdconf["name"]}"].map(KafkaConfig.${rdconf['swiftEnumName']}.init) % elif rdconf["swiftType"] == "[enum]": - self.${rdconf["swiftName"]} = configDict["${rdconf["name"]}"]?.components(separatedBy: ",").map { ${f"KafkaConfig.{rdconf["swiftEnumName"]}"}(description: $0) } + self.${rdconf["swiftName"]} = configDict["${rdconf["name"]}"]?.components(separatedBy: ",").map { ${f"KafkaConfig.{rdconf['swiftEnumName']}"}(description: $0) } % end % end } diff --git a/Sources/Kafka/Configuration/KafkaProducerConfig.swift b/Sources/Kafka/Configuration/KafkaProducerConfig.swift index cf2dfeea..5c30fa36 100644 --- a/Sources/Kafka/Configuration/KafkaProducerConfig.swift +++ b/Sources/Kafka/Configuration/KafkaProducerConfig.swift @@ -39,6 +39,25 @@ public struct KafkaProducerConfig: Sendable { /// Producer metrics configuration. public var metrics: KafkaConfiguration.ProducerMetrics = .init() + /// Low watermark for the ``KafkaProducerEvents`` asynchronous sequence backpressure. + /// When the number of buffered events falls below this threshold, the producer + /// will resume yielding events to the sequence. + /// Default: `1000` + public var eventsBackpressureLowWatermark: Int = 1000 + + /// High watermark for the ``KafkaProducerEvents`` asynchronous sequence backpressure. + /// When the number of buffered events reaches this threshold, the producer + /// will temporarily stop yielding events to the sequence to prevent out-of-memory errors. + /// Polling of librdkafka continues to process delivery reports for `sendAndAwait`. + /// Default: `10000` (corresponds to 1 million messages, matching confluent-kafka-go default) + public var eventsBackpressureHighWatermark: Int = 10000 + + /// Arbitrary configuration properties passed directly to librdkafka. + /// + /// - Warning: Properties set here override typed properties. + /// Intended for testing or advanced configurations not explicitly supported by this library. + internal var additionalConfig: [String: String] = [:] + // MARK: - Properties generated from librdkafka config list /// Client identifier. @@ -1035,14 +1054,6 @@ public struct KafkaProducerConfig: Sendable { /// Default: -1 public var compressionLevel: Int? - /// Additional librdkafka configuration properties not covered by typed properties. - /// Keys and values are passed directly to librdkafka. - /// - /// - Warning: Properties set here override typed properties above. - /// Intended for testing (e.g. `test.mock.num.brokers`) or advanced configurations - /// not explicitly supported by this library. - internal var additionalConfig: [String: String] = [:] - public init() {} internal var config: [String: String] { @@ -1162,9 +1173,7 @@ public struct KafkaProducerConfig: Sendable { config["statistics.interval.ms"] = String(updateInterval.inMilliseconds) } - for (key, value) in self.additionalConfig { - config[key] = value - } + config.merge(self.additionalConfig) { _, new in new } return config } diff --git a/Sources/Kafka/Configuration/KafkaProducerConfig.swift.gyb b/Sources/Kafka/Configuration/KafkaProducerConfig.swift.gyb index bd15bd73..9eda4728 100644 --- a/Sources/Kafka/Configuration/KafkaProducerConfig.swift.gyb +++ b/Sources/Kafka/Configuration/KafkaProducerConfig.swift.gyb @@ -37,6 +37,25 @@ public struct KafkaProducerConfig: Sendable { /// Producer metrics configuration. public var metrics: KafkaConfiguration.ProducerMetrics = .init() + /// Low watermark for the ``KafkaProducerEvents`` asynchronous sequence backpressure. + /// When the number of buffered events falls below this threshold, the producer + /// will resume yielding events to the sequence. + /// Default: `1000` + public var eventsBackpressureLowWatermark: Int = 1000 + + /// High watermark for the ``KafkaProducerEvents`` asynchronous sequence backpressure. + /// When the number of buffered events reaches this threshold, the producer + /// will temporarily stop yielding events to the sequence to prevent out-of-memory errors. + /// Polling of librdkafka continues to process delivery reports for `sendAndAwait`. + /// Default: `10000` (corresponds to 1 million messages, matching confluent-kafka-go default) + public var eventsBackpressureHighWatermark: Int = 10000 + + /// Arbitrary configuration properties passed directly to librdkafka. + /// + /// - Warning: Properties set here override typed properties. + /// Intended for testing or advanced configurations not explicitly supported by this library. + internal var additionalConfig: [String: String] = [:] + // MARK: - Properties generated from librdkafka config list % for rdconf in producer_librdkafka_configs: @@ -59,7 +78,7 @@ public struct KafkaProducerConfig: Sendable { /// /// Default: ${rdconf["defaultValue"]} % end - public var ${rdconf["swiftName"]}: ${rdconf["swiftType"].replace("enum", f"KafkaConfig.{rdconf["swiftEnumName"]}")}? + public var ${rdconf["swiftName"]}: ${rdconf["swiftType"].replace("enum", f"KafkaConfig.{rdconf['swiftEnumName']}")}? % end public init() {} @@ -86,6 +105,8 @@ public struct KafkaProducerConfig: Sendable { config["statistics.interval.ms"] = String(updateInterval.inMilliseconds) } + config.merge(self.additionalConfig) { _, new in new } + return config } @@ -103,9 +124,9 @@ public struct KafkaProducerConfig: Sendable { % elif rdconf["swiftType"] == "[String]": self.${rdconf["swiftName"]} = configDict["${rdconf["name"]}"]?.components(separatedBy: ",")${" // ignore-unacceptable-language" if rdconf["name"] == "topic.blacklist" else ""} % elif rdconf["swiftType"] == "enum": - self.${rdconf["swiftName"]} = configDict["${rdconf["name"]}"].map(KafkaConfig.${rdconf["swiftEnumName"]}.init) + self.${rdconf["swiftName"]} = configDict["${rdconf["name"]}"].map(KafkaConfig.${rdconf['swiftEnumName']}.init) % elif rdconf["swiftType"] == "[enum]": - self.${rdconf["swiftName"]} = configDict["${rdconf["name"]}"]?.components(separatedBy: ",").map { ${f"KafkaConfig.{rdconf["swiftEnumName"]}"}(description: $0) } + self.${rdconf["swiftName"]} = configDict["${rdconf["name"]}"]?.components(separatedBy: ",").map { ${f"KafkaConfig.{rdconf['swiftEnumName']}"}(description: $0) } % end % end } diff --git a/Sources/Kafka/KafkaConsumer.swift b/Sources/Kafka/KafkaConsumer.swift index ccd0e70c..f2ddd746 100644 --- a/Sources/Kafka/KafkaConsumer.swift +++ b/Sources/Kafka/KafkaConsumer.swift @@ -27,7 +27,7 @@ internal struct KafkaConsumerEventsDelegate: Sendable { extension KafkaConsumerEventsDelegate: NIOAsyncSequenceProducerDelegate { func produceMore() { - return // No back pressure + self.stateMachine.withLockedValue { $0.resumeYielding() } } func didTerminate() { @@ -40,7 +40,7 @@ extension KafkaConsumerEventsDelegate: NIOAsyncSequenceProducerDelegate { /// `AsyncSequence` implementation for handling ``KafkaConsumerEvent``s emitted by Kafka. public struct KafkaConsumerEvents: Sendable, AsyncSequence { public typealias Element = KafkaConsumerEvent - typealias BackPressureStrategy = NIOAsyncSequenceProducerBackPressureStrategies.NoBackPressure + typealias BackPressureStrategy = NIOAsyncSequenceProducerBackPressureStrategies.HighLowWatermark typealias WrappedSequence = NIOAsyncSequenceProducer let wrappedSequence: WrappedSequence @@ -144,7 +144,7 @@ public struct KafkaConsumerMessages: Sendable, AsyncSequence { public final class KafkaConsumer: Sendable, Service { typealias ConsumerEventsProducer = NIOAsyncSequenceProducer< KafkaConsumerEvent, - NIOAsyncSequenceProducerBackPressureStrategies.NoBackPressure, + NIOAsyncSequenceProducerBackPressureStrategies.HighLowWatermark, KafkaConsumerEventsDelegate > @@ -294,7 +294,10 @@ public final class KafkaConsumer: Sendable, Service { // it leads to a call to `stateMachine.messageSequenceTerminated()` while it's still in the `.uninitialized` state. let sourceAndSequence = NIOAsyncSequenceProducer.makeSequence( elementType: KafkaConsumerEvent.self, - backPressureStrategy: NIOAsyncSequenceProducerBackPressureStrategies.NoBackPressure(), + backPressureStrategy: NIOAsyncSequenceProducerBackPressureStrategies.HighLowWatermark( + lowWatermark: config.eventsBackpressureLowWatermark, + highWatermark: config.eventsBackpressureHighWatermark + ), finishOnDeinit: true, delegate: KafkaConsumerEventsDelegate(stateMachine: stateMachine) ) @@ -530,12 +533,15 @@ public final class KafkaConsumer: Sendable, Service { self.logger.error("Closing KafkaConsumer failed", metadata: ["error": "\(error)"]) self.stateMachine.withLockedValue { $0.closeFailed() } } - self.pollAndDrainEvents(client: client) + self.pollAndDrainEvents(client: client, yieldToSequence: true) try await Task.sleep(for: self.config.pollInterval) - case .pollForEvents(let client): - self.pollAndDrainEvents(client: client) - try await Task.sleep(for: self.config.pollInterval) - case .suspendPollLoop: + case .pollForEvents(let client, let yieldToSequence): + let hitHighWatermark = self.pollAndDrainEvents(client: client, yieldToSequence: yieldToSequence) + if hitHighWatermark { + self.stateMachine.withLockedValue { $0.pauseYielding() } + // Re-evaluate next action immediately without sleeping + continue + } try await Task.sleep(for: self.config.pollInterval) case .terminatePollLoop(let client): // Final drain: the close process may have completed (isConsumerClosed = true) @@ -550,15 +556,20 @@ public final class KafkaConsumer: Sendable, Service { } /// Poll librdkafka event queue and drain buffered rebalance events. - private func pollAndDrainEvents(client: RDKafkaClient) { + /// - Returns: `true` if hit high watermark backpressure. + @discardableResult + private func pollAndDrainEvents(client: RDKafkaClient, yieldToSequence: Bool) -> Bool { let events = client.consumerEventPoll() + var hitHighWatermark = false for event in events { switch event { case .statistics(let statistics): self.config.metrics.update(with: statistics) case .error(let kafkaError): - if let source = self.eventsSource { - _ = source.yield(.error(kafkaError)) + if yieldToSequence { + if self.eventsSource?.yield(.error(kafkaError)) == .stopProducing { + hitHighWatermark = true + } } self.logger.error( "Kafka client error", @@ -585,8 +596,10 @@ public final class KafkaConsumer: Sendable, Service { KafkaTopicPartition(topic: $0.topic, partition: KafkaPartition(rawValue: $0.partition)) } ) - if let source = self.eventsSource { - _ = source.yield(.rebalance(rebalance)) + if yieldToSequence { + if self.eventsSource?.yield(.rebalance(rebalance)) == .stopProducing { + hitHighWatermark = true + } } self.logger.info( "Consumer rebalance", @@ -596,6 +609,8 @@ public final class KafkaConsumer: Sendable, Service { ] ) } + + return hitHighWatermark } /// Final drain of the event queue and rebalance buffer before shutdown completes. @@ -934,7 +949,8 @@ extension KafkaConsumer { /// /// - Parameter client: Client used for handling the connection to the Kafka cluster. /// - Parameter rebalanceContext: The context for the C rebalance callback. - case running(client: RDKafkaClient, rebalanceContext: RebalanceContext) + /// - Parameter isYieldingPaused: Whether the events sequence buffer is full and yielding is temporarily suspended. + case running(client: RDKafkaClient, rebalanceContext: RebalanceContext, isYieldingPaused: Bool) /// The ``KafkaConsumer`` is being closed. /// /// - Parameter client: Client used for handling the connection to the Kafka cluster. @@ -943,7 +959,8 @@ extension KafkaConsumer { client: RDKafkaClient, rebalanceContext: RebalanceContext, closeInitiated: Bool, - gracefulShutdownRequested: Bool + gracefulShutdownRequested: Bool, + isYieldingPaused: Bool ) /// The ``KafkaConsumer`` has closed its queue and left the group, /// but is waiting for ``KafkaConsumer/triggerGracefulShutdown()`` to be invoked @@ -978,13 +995,12 @@ extension KafkaConsumer { /// Poll for new events. /// /// - Parameter client: Client used for handling the connection to the Kafka cluster. - case pollForEvents(client: RDKafkaClient) + /// - Parameter yieldToSequence: Whether to yield events to the `source`. + case pollForEvents(client: RDKafkaClient, yieldToSequence: Bool) /// Initiate consumer close, then poll for new events. /// /// - Parameter client: Client used for handling the connection to the Kafka cluster. case initiateCloseAndPoll(client: RDKafkaClient) - /// Suspend the poll loop. - case suspendPollLoop /// Terminate the poll loop. /// /// - Parameter client: Client for final drain (may be nil if state was already `.finished`). @@ -1001,16 +1017,23 @@ extension KafkaConsumer { fatalError("\(#function) invoked while still in state \(self.state)") case .initializing: fatalError("Subscribe to consumer group / assign to topic partition pair before reading messages") - case .running(let client, _): - return .pollForEvents(client: client) - case .finishing(let client, let rebalanceContext, let closeInitiated, let gracefulShutdownRequested): + case .running(let client, _, let isYieldingPaused): + return .pollForEvents(client: client, yieldToSequence: !isYieldingPaused) + case .finishing( + let client, + let rebalanceContext, + let closeInitiated, + let gracefulShutdownRequested, + let isYieldingPaused + ): if client.isConsumerClosed { if gracefulShutdownRequested { self.state = .finished return .terminatePollLoop(client: client) } else { self.state = .finishedAwaitingGracefulShutdown(client: client) - return .suspendPollLoop + // During shutdown, we always yield to the sequence to ensure rebalance/close signals are delivered. + return .pollForEvents(client: client, yieldToSequence: true) } } else { if !closeInitiated { @@ -1018,14 +1041,16 @@ extension KafkaConsumer { client: client, rebalanceContext: rebalanceContext, closeInitiated: true, - gracefulShutdownRequested: gracefulShutdownRequested + gracefulShutdownRequested: gracefulShutdownRequested, + isYieldingPaused: isYieldingPaused ) return .initiateCloseAndPoll(client: client) } - return .pollForEvents(client: client) + // During shutdown, we always yield to the sequence to ensure rebalance/close signals are delivered. + return .pollForEvents(client: client, yieldToSequence: true) } - case .finishedAwaitingGracefulShutdown: - return .suspendPollLoop + case .finishedAwaitingGracefulShutdown(let client): + return .pollForEvents(client: client, yieldToSequence: false) case .finished: return .terminatePollLoop(client: nil) } @@ -1053,7 +1078,7 @@ extension KafkaConsumer { fatalError("\(#function) invoked while still in state \(self.state)") case .initializing: return .suspendPollLoop - case .running(let client, _): + case .running(let client, _, _): return .poll(client: client) case .finishing, .finishedAwaitingGracefulShutdown, .finished: return .terminatePollLoop @@ -1077,7 +1102,7 @@ extension KafkaConsumer { case .uninitialized: fatalError("\(#function) invoked while still in state \(self.state)") case .initializing(let client, let rebalanceContext): - self.state = .running(client: client, rebalanceContext: rebalanceContext) + self.state = .running(client: client, rebalanceContext: rebalanceContext, isYieldingPaused: false) return .setUpConnection(client: client) case .running: fatalError("\(#function) should not be invoked more than once") @@ -1107,7 +1132,7 @@ extension KafkaConsumer { fatalError("\(#function) invoked while still in state \(self.state)") case .initializing: return .throwClosedError - case .running(let client, _): + case .running(let client, _, _): return .client(client) case .finishing, .finishedAwaitingGracefulShutdown, .finished: return .throwClosedError @@ -1126,7 +1151,7 @@ extension KafkaConsumer { fatalError("\(#function) invoked while still in state \(self.state)") case .initializing(let client, _): return .client(client) - case .running(let client, _): + case .running(let client, _, _): return .client(client) case .finishing, .finishedAwaitingGracefulShutdown, .finished: return .throwClosedError @@ -1153,21 +1178,23 @@ extension KafkaConsumer { // and transition straight to .finished. self.state = .finished return nil - case .running(let client, let rebalanceContext): + case .running(let client, let rebalanceContext, let isYieldingPaused): self.state = .finishing( client: client, rebalanceContext: rebalanceContext, closeInitiated: true, - gracefulShutdownRequested: true + gracefulShutdownRequested: true, + isYieldingPaused: isYieldingPaused ) return .triggerGracefulShutdown(client: client) - case .finishing(let client, let rebalanceContext, let closeInitiated, _): + case .finishing(let client, let rebalanceContext, let closeInitiated, _, let isYieldingPaused): // Upgrade to graceful shutdown requested self.state = .finishing( client: client, rebalanceContext: rebalanceContext, closeInitiated: closeInitiated, - gracefulShutdownRequested: true + gracefulShutdownRequested: true, + isYieldingPaused: isYieldingPaused ) return nil case .finishedAwaitingGracefulShutdown: @@ -1185,12 +1212,13 @@ extension KafkaConsumer { fatalError("\(#function) invoked while still in state \(self.state)") case .initializing: self.state = .finished - case .running(let client, let rebalanceContext): + case .running(let client, let rebalanceContext, let isYieldingPaused): self.state = .finishing( client: client, rebalanceContext: rebalanceContext, closeInitiated: false, - gracefulShutdownRequested: false + gracefulShutdownRequested: false, + isYieldingPaused: isYieldingPaused ) case .finishing, .finishedAwaitingGracefulShutdown, .finished: break @@ -1200,7 +1228,7 @@ extension KafkaConsumer { /// Called if `consumerClose()` fails to avoid polling forever. mutating func closeFailed() { switch self.state { - case .finishing(let client, _, _, let gracefulShutdownRequested): + case .finishing(let client, _, _, let gracefulShutdownRequested, _): if gracefulShutdownRequested { self.state = .finished } else { @@ -1211,6 +1239,50 @@ extension KafkaConsumer { } } + /// Stop yielding events to the events asynchronous sequence (backpressure). + mutating func pauseYielding() { + switch self.state { + case .running(let client, let rebalanceContext, _): + self.state = .running( + client: client, + rebalanceContext: rebalanceContext, + isYieldingPaused: true + ) + case .finishing(let client, let rebalanceContext, let closeInitiated, let gracefulShutdownRequested, _): + self.state = .finishing( + client: client, + rebalanceContext: rebalanceContext, + closeInitiated: closeInitiated, + gracefulShutdownRequested: gracefulShutdownRequested, + isYieldingPaused: true + ) + default: + break + } + } + + /// Resume yielding events to the events asynchronous sequence (relieve backpressure). + mutating func resumeYielding() { + switch self.state { + case .running(let client, let rebalanceContext, _): + self.state = .running( + client: client, + rebalanceContext: rebalanceContext, + isYieldingPaused: false + ) + case .finishing(let client, let rebalanceContext, let closeInitiated, let gracefulShutdownRequested, _): + self.state = .finishing( + client: client, + rebalanceContext: rebalanceContext, + closeInitiated: closeInitiated, + gracefulShutdownRequested: gracefulShutdownRequested, + isYieldingPaused: false + ) + default: + break + } + } + /// Returns the client if available, for cleanup purposes (e.g., final event drain). /// Unlike ``withClient()``, this does not fatalError in transitional states — /// it returns `nil` when no client is available. @@ -1219,8 +1291,8 @@ extension KafkaConsumer { case .uninitialized, .finished: return nil case .initializing(let client, _), - .running(let client, _), - .finishing(let client, _, _, _), + .running(let client, _, _), + .finishing(let client, _, _, _, _), .finishedAwaitingGracefulShutdown(let client): return client } diff --git a/Sources/Kafka/KafkaProducer.swift b/Sources/Kafka/KafkaProducer.swift index 51a1d5fa..b883f166 100644 --- a/Sources/Kafka/KafkaProducer.swift +++ b/Sources/Kafka/KafkaProducer.swift @@ -27,7 +27,7 @@ internal struct KafkaProducerCloseOnTerminate: Sendable { extension KafkaProducerCloseOnTerminate: NIOAsyncSequenceProducerDelegate { func produceMore() { - return // No back pressure + self.stateMachine.withLockedValue { $0.resumeYielding() } } func didTerminate() { @@ -46,7 +46,7 @@ extension KafkaProducerCloseOnTerminate: NIOAsyncSequenceProducerDelegate { /// `AsyncSequence` implementation for handling ``KafkaProducerEvent``s emitted by Kafka. public struct KafkaProducerEvents: Sendable, AsyncSequence { public typealias Element = KafkaProducerEvent - typealias BackPressureStrategy = NIOAsyncSequenceProducerBackPressureStrategies.NoBackPressure + typealias BackPressureStrategy = NIOAsyncSequenceProducerBackPressureStrategies.HighLowWatermark typealias WrappedSequence = NIOAsyncSequenceProducer let wrappedSequence: WrappedSequence @@ -72,7 +72,7 @@ public struct KafkaProducerEvents: Sendable, AsyncSequence { public final class KafkaProducer: Service, Sendable { typealias Producer = NIOAsyncSequenceProducer< KafkaProducerEvent, - NIOAsyncSequenceProducerBackPressureStrategies.NoBackPressure, + NIOAsyncSequenceProducerBackPressureStrategies.HighLowWatermark, KafkaProducerCloseOnTerminate > @@ -188,7 +188,10 @@ public final class KafkaProducer: Service, Sendable { // it leads to a call to `stateMachine.stopConsuming()` while it's still in the `.uninitialized` state. let sourceAndSequence = NIOAsyncSequenceProducer.makeSequence( elementType: KafkaProducerEvent.self, - backPressureStrategy: NIOAsyncSequenceProducerBackPressureStrategies.NoBackPressure(), + backPressureStrategy: NIOAsyncSequenceProducerBackPressureStrategies.HighLowWatermark( + lowWatermark: config.eventsBackpressureLowWatermark, + highWatermark: config.eventsBackpressureHighWatermark + ), finishOnDeinit: true, delegate: KafkaProducerCloseOnTerminate(stateMachine: stateMachine) ) @@ -257,23 +260,41 @@ public final class KafkaProducer: Service, Sendable { } } try await Task.sleep(for: self.config.pollInterval) - case .pollAndYield(let client, let source): + case .pollAndYield(let client, let source, let yieldToSequence): let events = client.producerEventPoll() + var hitHighWatermark = false for event in events { switch event { case .statistics(let statistics): self.config.metrics.update(with: statistics) case .deliveryReport(let reports): + // Always resume continuations (sendAndAwait) self.resumeContinuations(for: reports) - _ = source?.yield(.deliveryReports(reports)) + + if yieldToSequence { + if source?.yield(.deliveryReports(reports)) == .stopProducing { + hitHighWatermark = true + } + } case .error(let kafkaError): - _ = source?.yield(.error(kafkaError)) + if yieldToSequence { + if source?.yield(.error(kafkaError)) == .stopProducing { + hitHighWatermark = true + } + } self.logger.error( "Kafka client error", metadata: ["error": "\(kafkaError)"] ) } } + + if hitHighWatermark { + self.stateMachine.withLockedValue { $0.pauseYielding() } + // Re-evaluate next action immediately without sleeping (Fixes Double Sleep) + continue + } + try await Task.sleep(for: self.config.pollInterval) case .flushFinishSourceAndTerminatePollLoop(let client, let source): precondition( @@ -464,11 +485,13 @@ extension KafkaProducer { /// - Parameter client: Client used for handling the connection to the Kafka cluster. /// - Parameter source: ``NIOAsyncSequenceProducer/Source`` used for yielding new elements. /// - Parameter topicHandles: Class containing all topic names with their respective `rd_kafka_topic_t` pointer. + /// - Parameter isYieldingPaused: Whether the events sequence buffer is full and yielding is temporarily suspended. case started( client: RDKafkaClient, messageIDCounter: UInt, source: Producer.Source?, - topicHandles: RDKafkaTopicHandles + topicHandles: RDKafkaTopicHandles, + isYieldingPaused: Bool ) /// Producer is still running but the event asynchronous sequence was terminated. /// All incoming events will be dropped. @@ -480,9 +503,11 @@ extension KafkaProducer { /// /// - Parameter client: Client used for handling the connection to the Kafka cluster. /// - Parameter source: ``NIOAsyncSequenceProducer/Source`` used for yielding new elements. + /// - Parameter isYieldingPaused: Whether the events sequence buffer is full and yielding is temporarily suspended. case finishing( client: RDKafkaClient, - source: Producer.Source? + source: Producer.Source?, + isYieldingPaused: Bool ) /// The ``KafkaProducer`` has been shut down and cannot be used anymore. case finished @@ -506,7 +531,8 @@ extension KafkaProducer { client: client, messageIDCounter: 0, source: source, - topicHandles: RDKafkaTopicHandles(client: client) + topicHandles: RDKafkaTopicHandles(client: client), + isYieldingPaused: false ) } @@ -520,7 +546,8 @@ extension KafkaProducer { /// /// - Parameter client: Client used for handling the connection to the Kafka cluster. /// - Parameter source: ``NIOAsyncSequenceProducer/Source`` used for yielding new elements. - case pollAndYield(client: RDKafkaClient, source: Producer.Source?) + /// - Parameter yieldToSequence: Whether to yield events to the `source`. + case pollAndYield(client: RDKafkaClient, source: Producer.Source?, yieldToSequence: Bool) /// Flush any outstanding producer messages. /// Then terminate the poll loop and finish the given `NIOAsyncSequenceProducerSource`. /// @@ -539,11 +566,11 @@ extension KafkaProducer { switch self.state { case .uninitialized: fatalError("\(#function) invoked while still in state \(self.state)") - case .started(let client, _, let source, _): - return .pollAndYield(client: client, source: source) + case .started(let client, _, let source, _, let isYieldingPaused): + return .pollAndYield(client: client, source: source, yieldToSequence: !isYieldingPaused) case .eventConsumptionFinished(let client): return .pollWithoutYield(client: client) - case .finishing(let client, let source): + case .finishing(let client, let source, _): return .flushFinishSourceAndTerminatePollLoop(client: client, source: source) case .finished: return .terminatePollLoop @@ -569,13 +596,14 @@ extension KafkaProducer { switch self.state { case .uninitialized: fatalError("\(#function) invoked while still in state \(self.state)") - case .started(let client, let messageIDCounter, let source, let topicHandles): + case .started(let client, let messageIDCounter, let source, let topicHandles, let isYieldingPaused): let newMessageID = messageIDCounter + 1 self.state = .started( client: client, messageIDCounter: newMessageID, source: source, - topicHandles: topicHandles + topicHandles: topicHandles, + isYieldingPaused: isYieldingPaused ) return .send( client: client, @@ -609,12 +637,12 @@ extension KafkaProducer { fatalError("\(#function) invoked while still in state \(self.state)") case .eventConsumptionFinished: fatalError("messageSequenceTerminated() must not be invoked more than once") - case .started(let client, _, let source, _): + case .started(let client, _, let source, _, _): self.state = .eventConsumptionFinished(client: client) return .finishSource(source: source) - case .finishing(let client, let source): + case .finishing(let client, let source, _): // Setting source to nil to prevent incoming events from buffering in `source` - self.state = .finishing(client: client, source: nil) + self.state = .finishing(client: client, source: nil, isYieldingPaused: false) return .finishSource(source: source) case .finished: break @@ -622,6 +650,50 @@ extension KafkaProducer { return nil } + /// Stop yielding events to the events asynchronous sequence (backpressure). + mutating func pauseYielding() { + switch self.state { + case .started(let client, let messageIDCounter, let source, let topicHandles, _): + self.state = .started( + client: client, + messageIDCounter: messageIDCounter, + source: source, + topicHandles: topicHandles, + isYieldingPaused: true + ) + case .finishing(let client, let source, _): + self.state = .finishing( + client: client, + source: source, + isYieldingPaused: true + ) + default: + break + } + } + + /// Resume yielding events to the events asynchronous sequence (relieve backpressure). + mutating func resumeYielding() { + switch self.state { + case .started(let client, let messageIDCounter, let source, let topicHandles, _): + self.state = .started( + client: client, + messageIDCounter: messageIDCounter, + source: source, + topicHandles: topicHandles, + isYieldingPaused: false + ) + case .finishing(let client, let source, _): + self.state = .finishing( + client: client, + source: source, + isYieldingPaused: false + ) + default: + break + } + } + /// Get action to be taken when wanting to do close the producer. /// /// - Important: This function throws a `fatalError` if called while in the `.initializing` state. @@ -629,10 +701,10 @@ extension KafkaProducer { switch self.state { case .uninitialized: fatalError("\(#function) invoked while still in state \(self.state)") - case .started(let client, _, let source, _): - self.state = .finishing(client: client, source: source) + case .started(let client, _, let source, _, let isYieldingPaused): + self.state = .finishing(client: client, source: source, isYieldingPaused: isYieldingPaused) case .eventConsumptionFinished(let client): - self.state = .finishing(client: client, source: nil) + self.state = .finishing(client: client, source: nil, isYieldingPaused: false) case .finishing, .finished: break } diff --git a/Sources/Kafka/RDKafka/RDKafkaClient.swift b/Sources/Kafka/RDKafka/RDKafkaClient.swift index e03b925f..0172d11e 100644 --- a/Sources/Kafka/RDKafka/RDKafkaClient.swift +++ b/Sources/Kafka/RDKafka/RDKafkaClient.swift @@ -131,12 +131,12 @@ public final class RDKafkaClient: Sendable { ) // Pass message over to librdkafka where it will be queued and sent to the Kafka Cluster. - // Returns 0 on success, error code otherwise. - let error = try topicHandles.withTopicHandlePointer( + // Returns an error object on failure or nil on success. + let errorPointer = try topicHandles.withTopicHandlePointer( topic: message.topic ) { topicHandle in try Self.withMessageKeyAndValueBuffer(for: message) { keyBuffer, valueBuffer in - let errorPointer = try Self.withKafkaCHeaders(for: message.headers) { cHeaders in + try Self.withKafkaCHeaders(for: message.headers) { cHeaders in try self._produceVariadic( topicHandle: topicHandle, partition: Int32(message.partition.rawValue), @@ -147,12 +147,11 @@ public final class RDKafkaClient: Sendable { cHeaders: cHeaders ) } - return rd_kafka_error_code(errorPointer) } } - if error != RD_KAFKA_RESP_ERR_NO_ERROR { - throw KafkaError.rdKafkaError(wrapping: error) + if let errorPointer { + throw KafkaError.rdKafkaError(wrapping: errorPointer) } } @@ -173,7 +172,7 @@ public final class RDKafkaClient: Sendable { cHeaders: [(key: UnsafePointer, value: UnsafeRawBufferPointer?)] ) throws -> OpaquePointer? { let sizeWithoutHeaders = (key != nil) ? 6 : 5 - let size = sizeWithoutHeaders + cHeaders.count + let size = sizeWithoutHeaders + (cHeaders.isEmpty ? 0 : 1) var arguments = Array(repeating: rd_kafka_vu_t(), count: size) var index = 0 @@ -205,23 +204,41 @@ public final class RDKafkaClient: Sendable { arguments[index].u.ptr = opaque index += 1 - for cHeader in cHeaders { - arguments[index].vtype = RD_KAFKA_VTYPE_HEADER - - arguments[index].u.header.name = cHeader.key - arguments[index].u.header.val = cHeader.value?.baseAddress - arguments[index].u.header.size = cHeader.value?.count ?? 0 + var hdrs: OpaquePointer? = nil + if !cHeaders.isEmpty { + hdrs = rd_kafka_headers_new(cHeaders.count) + for cHeader in cHeaders { + let addError = rd_kafka_header_add( + hdrs, + cHeader.key, + -1, + cHeader.value?.baseAddress, + cHeader.value?.count ?? 0 + ) + if addError != RD_KAFKA_RESP_ERR_NO_ERROR { + rd_kafka_headers_destroy(hdrs) + throw KafkaError.rdKafkaError(wrapping: addError) + } + } + arguments[index].vtype = RD_KAFKA_VTYPE_HEADERS + arguments[index].u.headers = hdrs index += 1 } assert(arguments.count == size) - return rd_kafka_produceva( + let result = rd_kafka_produceva( self.kafkaHandle.pointer, arguments, arguments.count ) + + if result != nil, let hdrs { + rd_kafka_headers_destroy(hdrs) + } + + return result } /// Scoped accessor that enables safe access to a ``KafkaProducerMessage``'s key and value raw buffers. diff --git a/Tests/IntegrationTests/KafkaTests.swift b/Tests/IntegrationTests/KafkaTests.swift index 8a15330f..c03dfa5c 100644 --- a/Tests/IntegrationTests/KafkaTests.swift +++ b/Tests/IntegrationTests/KafkaTests.swift @@ -2232,6 +2232,164 @@ func withTestTopic(partitions: Int32 = 1, _ body: (_ testTopic: String) async th } } + @Test func producerBackpressureStopsYielding() async throws { + try await withTestTopic { testTopic in + var backpressureConfig = self.producerConfig + backpressureConfig.eventsBackpressureLowWatermark = 3 + backpressureConfig.eventsBackpressureHighWatermark = 5 + + let (producer, events) = try KafkaProducer.makeProducerWithEvents( + config: backpressureConfig, + logger: .kafkaTest + ) + + let serviceGroupConfiguration = ServiceGroupConfiguration( + services: [producer], + logger: .kafkaTest + ) + let serviceGroup = ServiceGroup(configuration: serviceGroupConfiguration) + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + try await serviceGroup.run() + } + + let message = KafkaProducerMessage( + topic: testTopic, + value: "backpressure-test-payload" + ) + + // Send many messages without consuming events. + // Buffer will fill to high watermark, then yielding pauses. + // Production continues (librdkafka still polled). + for _ in 0..<50 { + try producer.send(message) + try await Task.sleep(for: .milliseconds(20)) + } + + // Consume events — they should be available (buffered up to watermark) + var eventCount = 0 + for await _ in events { + eventCount += 1 + if eventCount >= 3 { + break + } + } + #expect(eventCount >= 3) + + await serviceGroup.triggerGracefulShutdown() + } + } + } + + @Test func sendAndAwaitWorksDuringBackpressure() async throws { + try await withTestTopic { testTopic in + var backpressureConfig = self.producerConfig + backpressureConfig.queueBufferingMaxMessages = 100 + backpressureConfig.eventsBackpressureLowWatermark = 3 + backpressureConfig.eventsBackpressureHighWatermark = 5 + + let (producer, events) = try KafkaProducer.makeProducerWithEvents( + config: backpressureConfig, + logger: .kafkaTest + ) + + let serviceGroupConfiguration = ServiceGroupConfiguration( + services: [producer], + logger: .kafkaTest + ) + let serviceGroup = ServiceGroup(configuration: serviceGroupConfiguration) + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + try await serviceGroup.run() + } + + // Never consume events — force watermark hit + _ = events + + let message = KafkaProducerMessage( + topic: testTopic, + value: "sendAndAwait-backpressure" + ) + + // Fill events buffer past high watermark + for _ in 0..<20 { + try producer.send(message) + try await Task.sleep(for: .milliseconds(50)) + } + + // sendAndAwait must still work — event loop polls and services continuations + // even when yielding to the events sequence is paused + let report = try await producer.sendAndAwait(message) + switch report.status { + case .acknowledged(let ack): + #expect(ack.topic == testTopic) + case .failure(let error): + Issue.record("sendAndAwait failed during backpressure: \(error)") + } + + await serviceGroup.triggerGracefulShutdown() + } + } + } + + @Test func producerBackpressureCausesQueueFull() async throws { + try await withTestTopic { testTopic in + var backpressureConfig = self.producerConfig + // Tiny librdkafka queue + backpressureConfig.queueBufferingMaxMessages = 10 + // Hold messages for 5 seconds before sending — simulates slow broker ack + backpressureConfig.lingerMs = 5000 + // Small watermarks so Swift-side buffer fills quickly + backpressureConfig.eventsBackpressureLowWatermark = 3 + backpressureConfig.eventsBackpressureHighWatermark = 5 + + let (producer, events) = try KafkaProducer.makeProducerWithEvents( + config: backpressureConfig, + logger: .kafkaTest + ) + + let serviceGroupConfiguration = ServiceGroupConfiguration( + services: [producer], + logger: .kafkaTest + ) + let serviceGroup = ServiceGroup(configuration: serviceGroupConfiguration) + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + try await serviceGroup.run() + } + + // Never consume events + _ = events + + let message = KafkaProducerMessage( + topic: testTopic, + value: "queue-full-test" + ) + + // Flood the producer. With linger.ms=5000 and queue size=10, + // messages won't be acked for 5 seconds. The librdkafka internal + // queue will fill to 10, and send() will return QueueFull. + var lastError: Error? + for _ in 0..<100 { + do { + try producer.send(message) + } catch { + lastError = error + break + } + } + + let kafkaError = try #require(lastError as? KafkaError) + #expect(kafkaError.rdKafkaCode == .queueFull) + + await serviceGroup.triggerGracefulShutdown() + } + } + } + // MARK: - Helpers func produceMessages(topic: String, count: UInt) async throws -> [KafkaProducerMessage] { diff --git a/Tests/KafkaTests/KafkaProducerTests.swift b/Tests/KafkaTests/KafkaProducerTests.swift index b41da644..3a4d21ea 100644 --- a/Tests/KafkaTests/KafkaProducerTests.swift +++ b/Tests/KafkaTests/KafkaProducerTests.swift @@ -378,6 +378,199 @@ import Foundation } } + @Test func producerBackpressureStopsYielding() async throws { + var backpressureConfig = self.config + // Set small watermarks so we can observe the pause quickly + backpressureConfig.eventsBackpressureLowWatermark = 3 + backpressureConfig.eventsBackpressureHighWatermark = 5 + + let (producer, events) = try KafkaProducer.makeProducerWithEvents( + config: backpressureConfig, + logger: .kafkaTest + ) + + let serviceGroupConfiguration = ServiceGroupConfiguration(services: [producer], logger: .kafkaTest) + let serviceGroup = ServiceGroup(configuration: serviceGroupConfiguration) + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + try await serviceGroup.run() + } + + let message = KafkaProducerMessage( + topic: "backpressure-topic", + value: "hello" + ) + + // Send many messages without consuming events. + // The event loop will hit the high watermark and stop yielding, + // but continue polling librdkafka (so sendAndAwait still works). + for _ in 0..<50 { + try producer.send(message) + try await Task.sleep(for: .milliseconds(20)) + } + + // Now consume events — should get at most highWatermark buffered events + // (the sequence stops growing once watermark is hit) + var eventCount = 0 + for await _ in events { + eventCount += 1 + if eventCount >= 3 { + break + } + } + #expect(eventCount >= 3) + + await serviceGroup.triggerGracefulShutdown() + } + } + + @Test func sendAndAwaitWorksWhileBackpressured() async throws { + var backpressureConfig = self.config + backpressureConfig.queueBufferingMaxMessages = 100 + backpressureConfig.eventsBackpressureLowWatermark = 2 + backpressureConfig.eventsBackpressureHighWatermark = 5 + + let (producer, events) = try KafkaProducer.makeProducerWithEvents( + config: backpressureConfig, + logger: .kafkaTest + ) + + let serviceGroupConfiguration = ServiceGroupConfiguration(services: [producer], logger: .kafkaTest) + let serviceGroup = ServiceGroup(configuration: serviceGroupConfiguration) + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + try await serviceGroup.run() + } + + // Never consume events — force the watermark to be hit + _ = events + + let message = KafkaProducerMessage( + topic: "backpressure-topic", + value: "sendAndAwait-test" + ) + + // Fill the events buffer past the high watermark by sending fire-and-forget messages. + // The event loop will transition to yieldToSequence: false. + for _ in 0..<20 { + try producer.send(message) + try await Task.sleep(for: .milliseconds(50)) + } + + // Now call sendAndAwait — it should NOT block indefinitely because the event loop + // still polls librdkafka and resumes continuations even when yielding is paused. + let report = try await producer.sendAndAwait(message) + switch report.status { + case .acknowledged: + break // success — continuations are serviced during backpressure + case .failure(let error): + // Acceptable: mock broker may not ack. The key assertion is that this + // does NOT block indefinitely — it completes one way or another. + _ = error + } + + await serviceGroup.triggerGracefulShutdown() + } + } + + @Test func producerResumesAfterEventsDrained() async throws { + var backpressureConfig = self.config + backpressureConfig.eventsBackpressureLowWatermark = 3 + backpressureConfig.eventsBackpressureHighWatermark = 5 + + let (producer, events) = try KafkaProducer.makeProducerWithEvents( + config: backpressureConfig, + logger: .kafkaTest + ) + + let serviceGroupConfiguration = ServiceGroupConfiguration(services: [producer], logger: .kafkaTest) + let serviceGroup = ServiceGroup(configuration: serviceGroupConfiguration) + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + try await serviceGroup.run() + } + + let message = KafkaProducerMessage( + topic: "backpressure-topic", + value: "resume-test" + ) + + // Phase 1: Send messages without consuming to fill buffer past high watermark + for _ in 0..<20 { + try producer.send(message) + try await Task.sleep(for: .milliseconds(20)) + } + + // Phase 2: Drain events and then send more — verify new events arrive + // (yielding resumes after low watermark crossed via produceMore callback) + var totalEvents = 0 + for await _ in events { + totalEvents += 1 + + // After draining initial batch, send more to verify yielding resumed + if totalEvents == 5 { + for _ in 0..<5 { + try producer.send(message) + } + } + + // If we receive events beyond the initial batch, yielding resumed successfully + if totalEvents >= 8 { + break + } + } + #expect(totalEvents >= 8) + + await serviceGroup.triggerGracefulShutdown() + } + } + + @Test func gracefulShutdownWhileBackpressured() async throws { + var backpressureConfig = self.config + backpressureConfig.queueBufferingMaxMessages = 10 + backpressureConfig.eventsBackpressureLowWatermark = 3 + backpressureConfig.eventsBackpressureHighWatermark = 5 + + let (producer, events) = try KafkaProducer.makeProducerWithEvents( + config: backpressureConfig, + logger: .kafkaTest + ) + + let serviceGroupConfiguration = ServiceGroupConfiguration(services: [producer], logger: .kafkaTest) + let serviceGroup = ServiceGroup(configuration: serviceGroupConfiguration) + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + try await serviceGroup.run() + } + + // Never consume events + _ = events + + let message = KafkaProducerMessage( + topic: "backpressure-topic", + value: "shutdown-test" + ) + + // Fill buffer to trigger backpressure + for _ in 0..<30 { + do { + try producer.send(message) + } catch { + break + } + try await Task.sleep(for: .milliseconds(10)) + } + + // Trigger shutdown while backpressured — must complete cleanly + await serviceGroup.triggerGracefulShutdown() + } + // If we reach here, shutdown completed successfully despite backpressure + } + // MARK: - KafkaProducerEvent.error Tests @Test func producerEventErrorPatternMatch() {