KafkaProducer: Drop acks on sequence termination (#84)

felixschlegel · web-flow · commit 38ece7b97e19 · 2023-07-11T16:00:53.000+01:00
* `KafkaProducer`: Drop acks on sequence termination

Motivation:

We do not to buffer any incoming messages when the asynchronous
sequences we expose are terminated.

This PR proposes adding a new state to `KafkaProducer` that drops any incoming acknowledgements after the asynchronous
sequence has been terminated.

&gt; **Note**: we are not shutting down entirely as terminating the `run()`
loop early is seen as an error is `swift-service-lifecycle`.

Modifications:

* add new `NIOAsyncSequenceProducerDelegate` `KafkaProducerCloseOnTerminate`
* `KafkaProducer`: add new stateMachine logic
  `stopConsuming` that does not stop the poll loop but drops all
  incoming messages / acknowledgements
* add test
  `KafkaProducerTests.testSendFailsAfterTerminatingAcknowledgementSequence`

* Review Franz

Modifications:

* rename `acks` -&gt; `acknowledgments`
* rename `iter` -&gt; `iterator`
* `KafkaProducerTests`: remove unnecessary `TaskGroup` tasks

* Fix rebase issues
diff --git a/Sources/SwiftKafka/KafkaConsumer.swift b/Sources/SwiftKafka/KafkaConsumer.swift
@@ -18,6 +18,14 @@ import NIOConcurrencyHelpers
 import NIOCore
 import ServiceLifecycle
 
+// MARK: - NoDelegate
+
+// `NIOAsyncSequenceProducerDelegate` that does nothing.
+internal struct NoDelegate: NIOAsyncSequenceProducerDelegate {
+    func produceMore() {}
+    func didTerminate() {}
+}
+
 // MARK: - KafkaConsumerMessages
 
 /// `AsyncSequence` implementation for handling messages received from the Kafka cluster (``KafkaConsumerMessage``).
diff --git a/Sources/SwiftKafka/KafkaProducer.swift b/Sources/SwiftKafka/KafkaProducer.swift
@@ -18,13 +18,37 @@ import NIOConcurrencyHelpers
 import NIOCore
 import ServiceLifecycle
 
+// MARK: - KafkaProducerCloseOnTerminate
+
+/// `NIOAsyncSequenceProducerDelegate` that terminates the closes the producer when
+/// `didTerminate()` is invoked.
+internal struct KafkaProducerCloseOnTerminate: Sendable {
+    let stateMachine: NIOLockedValueBox<KafkaProducer.StateMachine>
+}
+
+extension KafkaProducerCloseOnTerminate: NIOAsyncSequenceProducerDelegate {
+    func produceMore() {
+        return // No back pressure
+    }
+
+    func didTerminate() {
+        let action = self.stateMachine.withLockedValue { $0.stopConsuming() }
+        switch action {
+        case .finishSource(let source):
+            source?.finish()
+        case .none:
+            break
+        }
+    }
+}
+
 // MARK: - KafkaMessageAcknowledgements
 
 /// `AsyncSequence` implementation for handling messages acknowledged by the Kafka cluster (``KafkaAcknowledgedMessage``).
 public struct KafkaMessageAcknowledgements: AsyncSequence {
     public typealias Element = Result<KafkaAcknowledgedMessage, KafkaAcknowledgedMessageError>
     typealias BackPressureStrategy = NIOAsyncSequenceProducerBackPressureStrategies.NoBackPressure
-    typealias WrappedSequence = NIOAsyncSequenceProducer<Element, BackPressureStrategy, NoDelegate>
+    typealias WrappedSequence = NIOAsyncSequenceProducer<Element, BackPressureStrategy, KafkaProducerCloseOnTerminate>
     let wrappedSequence: WrappedSequence
 
     /// `AsynceIteratorProtocol` implementation for handling messages acknowledged by the Kafka cluster (``KafkaAcknowledgedMessage``).
@@ -51,7 +75,7 @@ public final class KafkaProducer: Service, Sendable {
     typealias Producer = NIOAsyncSequenceProducer<
         Result<KafkaAcknowledgedMessage, KafkaAcknowledgedMessageError>,
         NIOAsyncSequenceProducerBackPressureStrategies.NoBackPressure,
-        NoDelegate
+        KafkaProducerCloseOnTerminate
     >
 
     /// State of the ``KafkaProducer``.
@@ -122,7 +146,8 @@ public final class KafkaProducer: Service, Sendable {
     ///
     /// Use the asynchronous sequence to consume message acknowledgements.
     ///
-    /// - Important: When the asynchronous sequence is deinited the producer will be shutdown.
+    /// - Important: When the asynchronous sequence is deinited the producer will be shutdown and disallow sending more messages.
+    /// Additionally, make sure to consume the asynchronous sequence otherwise the acknowledgements will be buffered in memory indefinitely.
     ///
     /// - Parameter config: The ``KafkaProducerConfiguration`` for configuring the ``KafkaProducer``.
     /// - Parameter topicConfig: The ``KafkaTopicConfiguration`` used for newly created topics.
@@ -140,7 +165,7 @@ public final class KafkaProducer: Service, Sendable {
         let sourceAndSequence = NIOAsyncSequenceProducer.makeSequence(
             elementType: Result<KafkaAcknowledgedMessage, KafkaAcknowledgedMessageError>.self,
             backPressureStrategy: NIOAsyncSequenceProducerBackPressureStrategies.NoBackPressure(),
-            delegate: NoDelegate()
+            delegate: KafkaProducerCloseOnTerminate(stateMachine: stateMachine)
         )
         let source = sourceAndSequence.source
 
@@ -183,7 +208,10 @@ public final class KafkaProducer: Service, Sendable {
         while !Task.isCancelled {
             let nextAction = self.stateMachine.withLockedValue { $0.nextPollLoopAction() }
             switch nextAction {
-            case .poll(let client, let source):
+            case .pollWithoutYield(let client):
+                // Drop any incoming acknowledgments
+                let _ = client.eventPoll()
+            case .pollAndYield(let client, let source):
                 let events = client.eventPoll()
                 for event in events {
                     switch event {
@@ -258,6 +286,11 @@ extension KafkaProducer {
                 source: Producer.Source?,
                 topicHandles: RDKafkaTopicHandles
             )
+            /// Producer is still running but the acknowledgement asynchronous sequence was terminated.
+            /// All incoming acknowledgements will be dropped.
+            ///
+            /// - Parameter client: Client used for handling the connection to the Kafka cluster.
+            case consumptionStopped(client: KafkaClient)
             /// ``KafkaProducer/triggerGracefulShutdown()`` was invoked so we are flushing
             /// any messages that wait to be sent and serve any remaining queued callbacks.
             ///
@@ -293,11 +326,15 @@ extension KafkaProducer {
 
         /// Action to be taken when wanting to poll.
         enum PollLoopAction {
-            /// Poll client for new consumer messages.
+            /// Poll client.
+            ///
+            /// - Parameter client: Client used for handling the connection to the Kafka cluster.
+            case pollWithoutYield(client: KafkaClient)
+            /// Poll client and yield acknowledgments if any received.
             ///
             /// - Parameter client: Client used for handling the connection to the Kafka cluster.
             /// - Parameter source: ``NIOAsyncSequenceProducer/Source`` used for yielding new elements.
-            case poll(client: KafkaClient, source: Producer.Source?)
+            case pollAndYield(client: KafkaClient, source: Producer.Source?)
             /// Terminate the poll loop and finish the given `NIOAsyncSequenceProducerSource`.
             ///
             /// - Parameter source: ``NIOAsyncSequenceProducer/Source`` used for yielding new elements.
@@ -315,10 +352,12 @@ extension KafkaProducer {
             case .uninitialized:
                 fatalError("\(#function) invoked while still in state \(self.state)")
             case .started(let client, _, let source, _):
-                return .poll(client: client, source: source)
+                return .pollAndYield(client: client, source: source)
+            case .consumptionStopped(let client):
+                return .pollWithoutYield(client: client)
             case .flushing(let client, let source):
                 if client.outgoingQueueSize > 0 {
-                    return .poll(client: client, source: source)
+                    return .pollAndYield(client: client, source: source)
                 } else {
                     self.state = .finished
                     return .terminatePollLoopAndFinishSource(source: source)
@@ -360,13 +399,44 @@ extension KafkaProducer {
                     newMessageID: newMessageID,
                     topicHandles: topicHandles
                 )
+            case .consumptionStopped:
+                throw KafkaError.connectionClosed(reason: "Sequence consuming acknowledgements was abruptly terminated, producer closed")
             case .flushing:
                 throw KafkaError.connectionClosed(reason: "Producer in the process of flushing and shutting down")
             case .finished:
                 throw KafkaError.connectionClosed(reason: "Tried to produce a message with a closed producer")
             }
         }
 
+        /// Action to take after invoking ``KafkaProducer/StateMachine/stopConsuming()``.
+        enum StopConsumingAction {
+            /// Finish the given `NIOAsyncSequenceProducerSource`.
+            ///
+            /// - Parameter source: ``NIOAsyncSequenceProducer/Source`` used for yielding new elements.
+            case finishSource(source: Producer.Source?)
+        }
+
+        /// The acknowledgements asynchronous sequence was terminated.
+        /// All incoming acknowledgements will be dropped.
+        mutating func stopConsuming() -> StopConsumingAction? {
+            switch self.state {
+            case .uninitialized:
+                fatalError("\(#function) invoked while still in state \(self.state)")
+            case .consumptionStopped:
+                fatalError("stopConsuming() must not be invoked more than once")
+            case .started(let client, _, let source, _):
+                self.state = .consumptionStopped(client: client)
+                return .finishSource(source: source)
+            case .flushing(let client, let source):
+                // Setting source to nil to prevent incoming acknowledgements from buffering in `source`
+                self.state = .flushing(client: client, source: nil)
+                return .finishSource(source: source)
+            case .finished:
+                break
+            }
+            return nil
+        }
+
         /// 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.
@@ -376,6 +446,8 @@ extension KafkaProducer {
                 fatalError("\(#function) invoked while still in state \(self.state)")
             case .started(let client, _, let source, _):
                 self.state = .flushing(client: client, source: source)
+            case .consumptionStopped(let client):
+                self.state = .flushing(client: client, source: nil)
             case .flushing, .finished:
                 break
             }
diff --git a/Sources/SwiftKafka/Utilities/NIOAsyncSequenceBackPressureStrategies+NoBackPressure.swift b/Sources/SwiftKafka/Utilities/NIOAsyncSequenceBackPressureStrategies+NoBackPressure.swift
@@ -21,9 +21,3 @@ extension NIOAsyncSequenceProducerBackPressureStrategies {
         func didConsume(bufferDepth: Int) -> Bool { true }
     }
 }
-
-/// `NIOAsyncSequenceProducerDelegate` that does nothing.
-internal struct NoDelegate: NIOAsyncSequenceProducerDelegate {
-    func produceMore() {}
-    func didTerminate() {}
-}
diff --git a/Tests/IntegrationTests/SwiftKafkaTests.swift b/Tests/IntegrationTests/SwiftKafkaTests.swift
@@ -85,7 +85,7 @@ final class SwiftKafkaTests: XCTestCase {
 
     func testProduceAndConsumeWithConsumerGroup() async throws {
         let testMessages = Self.createTestMessages(topic: self.uniqueTestTopic, count: 10)
-        let (producer, acks) = try KafkaProducer.makeProducerWithAcknowledgements(config: self.producerConfig, logger: .kafkaTest)
+        let (producer, acknowledgments) = try KafkaProducer.makeProducerWithAcknowledgements(config: self.producerConfig, logger: .kafkaTest)
 
         let consumerConfig = KafkaConsumerConfiguration(
             consumptionStrategy: .group(groupID: "subscription-test-group-id", topics: [self.uniqueTestTopic]),
@@ -115,7 +115,7 @@ final class SwiftKafkaTests: XCTestCase {
             group.addTask {
                 try await Self.sendAndAcknowledgeMessages(
                     producer: producer,
-                    acknowledgements: acks,
+                    acknowledgements: acknowledgments,
                     messages: testMessages
                 )
             }
@@ -153,7 +153,7 @@ final class SwiftKafkaTests: XCTestCase {
 
     func testProduceAndConsumeWithAssignedTopicPartition() async throws {
         let testMessages = Self.createTestMessages(topic: self.uniqueTestTopic, count: 10)
-        let (producer, acks) = try KafkaProducer.makeProducerWithAcknowledgements(config: self.producerConfig, logger: .kafkaTest)
+        let (producer, acknowledgments) = try KafkaProducer.makeProducerWithAcknowledgements(config: self.producerConfig, logger: .kafkaTest)
 
         let consumerConfig = KafkaConsumerConfiguration(
             consumptionStrategy: .partition(
@@ -187,7 +187,7 @@ final class SwiftKafkaTests: XCTestCase {
             group.addTask {
                 try await Self.sendAndAcknowledgeMessages(
                     producer: producer,
-                    acknowledgements: acks,
+                    acknowledgements: acknowledgments,
                     messages: testMessages
                 )
             }
@@ -225,7 +225,7 @@ final class SwiftKafkaTests: XCTestCase {
 
     func testProduceAndConsumeWithCommitSync() async throws {
         let testMessages = Self.createTestMessages(topic: self.uniqueTestTopic, count: 10)
-        let (producer, acks) = try KafkaProducer.makeProducerWithAcknowledgements(config: self.producerConfig, logger: .kafkaTest)
+        let (producer, acknowledgments) = try KafkaProducer.makeProducerWithAcknowledgements(config: self.producerConfig, logger: .kafkaTest)
 
         let consumerConfig = KafkaConsumerConfiguration(
             consumptionStrategy: .group(groupID: "commit-sync-test-group-id", topics: [self.uniqueTestTopic]),
@@ -256,7 +256,7 @@ final class SwiftKafkaTests: XCTestCase {
             group.addTask {
                 try await Self.sendAndAcknowledgeMessages(
                     producer: producer,
-                    acknowledgements: acks,
+                    acknowledgements: acknowledgments,
                     messages: testMessages
                 )
             }
diff --git a/Tests/SwiftKafkaTests/KafkaProducerTests.swift b/Tests/SwiftKafkaTests/KafkaProducerTests.swift

Original file line number	Diff line number	Diff line change
`@@ -21,9 +21,3 @@ extension NIOAsyncSequenceProducerBackPressureStrategies {`
`21`	`21`	`func didConsume(bufferDepth: Int) -> Bool { true }`
`22`	`22`	`}`
`23`	`23`	`}`
`24`		`-`
`25`		-/// `NIOAsyncSequenceProducerDelegate` that does nothing.
`26`		`-internal struct NoDelegate: NIOAsyncSequenceProducerDelegate {`
`27`		`- func produceMore() {}`
`28`		`- func didTerminate() {}`
`29`		`-}`