KafkaConsumer: drop messages on sequence termination (#85)

* `KafkaConsumer`: drop msgs on sequence termination

Motivation:

We want to avoid running out of memory when the user stops reading the
`KafkaConsumer.messages` `AsyncSequence`. Therefore we drop any incoming
messages that we receive after the `AsyncSequence` was terminated.

Modifications:

* `KafkaConsumer`
    * enter idle polling state when `KafkaConsumer.messages`
      `AsyncSequence` was terminated to ensure that the `func run()`
      async does not return early and result in an error for its `ServiceGroup`
* `KafkaConsumerConfiguration`:
    * proxy `enable.auto.commit` and make it `false` by default
      Reason: our library should be in charge of deciding when a message
      offset is to be commited because only we know if our library user received the message through the
      `KafkaConsumer.messages` `AsyncSequence`
* `KafkaSharedConfiguration.partition`: document default value of
  parameter `offset`
* add new test `SwiftKafkaTests.testCommittedOffsetsAreCorrect()`

* KafkaConsumerCloseOnTerminate: make 'checked' `Sendable`

* Review Franz

Motivation:

Instead of storing a messages offset after every yield to the
`messages` `AsyncSequence`, we want to do that after every invocation to
the `message` `AsyncSequenceIterator`'s `next()` method to ensure that
the user has actually received the message. Yielding only enqueued a
message but did not ensure that the user actually consumed it through
the `AsyncSequence`.

Modifications:

* `KafkaConsumerConfig`
    * revert to plain `librdkafka` `enable.auto.commit` implementation
* `KafkaConsumer`
    * use `rd_kafka_offsets_store()` instead of `rd_kafka_commit` to store a
      message's offset after it has been consumed through the `messages`
      `AsyncSequence` -> more lightweight than `rd_kafka_commit` + we
      can still leverage `librdkafka`'s auto-commit logic
    * commit message offsets on `ConsumerMessagesAsyncIterator.next()`
      instead of after yielding (explaination in Motivation section)

* Review Franz & George

Modifications:

* make `KafkaConsumerMessages` `AsyncSequence` throwing
* remove associated `Producer.Source` from
  `KafkaConsumer.StateMachine.State.consumptionStopped`
* rename `KafkaConsumer.StateMachine.stopConsuming()` ->
  `messageSequenceTerminated()`

* Review Franz

Modifications:

* `KafkaConsumerMessages` `AsyncSequnce`: deallocate `wrappedIterator`
  when sequence terminates (abruptly)

* Remove unused try

Author: felixschlegel

README.md

@@ -99,7 +99,7 @@ await withThrowingTaskGroup(of: Void.self) { group in
 
     // Task receiving messages
     group.addTask {
-        for await message in consumer.messages {
+        for try await message in consumer.messages {
             // Do something with message
         }
     }
@@ -135,7 +135,7 @@ await withThrowingTaskGroup(of: Void.self) { group in
 
     // Task receiving messages
     group.addTask {
-        for await message in consumer.messages {
+        for try await message in consumer.messages {
             // Do something with message
         }
     }
@@ -172,7 +172,7 @@ await withThrowingTaskGroup(of: Void.self) { group in
 
     // Task receiving messages
     group.addTask {
-        for await message in consumer.messages {
+        for try await message in consumer.messages {
             // Do something with message
             // ...
             try await consumer.commitSync(message)

Sources/SwiftKafka/Configuration/KafkaConsumerConfiguration.swift

@@ -498,6 +498,7 @@ extension KafkaSharedConfiguration {
         /// - Parameter topic: The name of the Kafka topic.
         /// - Parameter partition: The partition of the topic to consume from.
         /// - Parameter offset: The offset to start consuming from.
+        /// Defaults to the end of the Kafka partition queue (meaning wait for next produced message).
         public static func partition(
             topic: String,
             partition: KafkaPartition,

Sources/SwiftKafka/KafkaClient.swift

@@ -277,7 +277,33 @@ final class KafkaClient: Sendable {
         }
     }
 
-    /// Non-blocking commit of a the `message`'s offset to Kafka.
+    /// Store `message`'s offset for next auto-commit.
+    ///
+    /// - Important: `enable.auto.offset.store` must be set to `false` when using this API.
+    func storeMessageOffset(_ message: KafkaConsumerMessage) throws {
+        // The offset committed is always the offset of the next requested message.
+        // Thus, we increase the offset of the current message by one before committing it.
+        // See: https://github.com/edenhill/librdkafka/issues/2745#issuecomment-598067945
+        let changesList = RDKafkaTopicPartitionList()
+        changesList.setOffset(
+            topic: message.topic,
+            partition: message.partition,
+            offset: Int64(message.offset + 1)
+        )
+
+        let error = changesList.withListPointer { listPointer in
+            rd_kafka_offsets_store(
+                self.kafkaHandle,
+                listPointer
+            )
+        }
+
+        if error != RD_KAFKA_RESP_ERR_NO_ERROR {
+            throw KafkaError.rdKafkaError(wrapping: error)
+        }
+    }
+
+    /// Non-blocking **awaitable** commit of a the `message`'s offset to Kafka.
     ///
     /// - Parameter message: Last received message that shall be marked as read.
     func commitSync(_ message: KafkaConsumerMessage) async throws {

Sources/SwiftKafka/KafkaConsumer.swift

@@ -18,45 +18,86 @@ import NIOConcurrencyHelpers
 import NIOCore
 import ServiceLifecycle
 
-// MARK: - NoDelegate
+// MARK: - KafkaConsumerCloseOnTerminate
 
-// `NIOAsyncSequenceProducerDelegate` that does nothing.
-internal struct NoDelegate: NIOAsyncSequenceProducerDelegate {
-    func produceMore() {}
-    func didTerminate() {}
+/// `NIOAsyncSequenceProducerDelegate` that terminates the closes the producer when
+/// `didTerminate()` is invoked.
+internal struct KafkaConsumerCloseOnTerminate: Sendable {
+    let stateMachine: NIOLockedValueBox<KafkaConsumer.StateMachine>
+}
+
+extension KafkaConsumerCloseOnTerminate: NIOAsyncSequenceProducerDelegate {
+    func produceMore() {
+        return // No back pressure
+    }
+
+    func didTerminate() {
+        self.stateMachine.withLockedValue { $0.messageSequenceTerminated() }
+    }
 }
 
 // MARK: - KafkaConsumerMessages
 
 /// `AsyncSequence` implementation for handling messages received from the Kafka cluster (``KafkaConsumerMessage``).
 public struct KafkaConsumerMessages: Sendable, AsyncSequence {
+    let stateMachine: NIOLockedValueBox<KafkaConsumer.StateMachine>
+
     public typealias Element = KafkaConsumerMessage
     typealias BackPressureStrategy = NIOAsyncSequenceProducerBackPressureStrategies.NoBackPressure
-    typealias WrappedSequence = NIOAsyncSequenceProducer<Element, BackPressureStrategy, NoDelegate>
+    typealias WrappedSequence = NIOThrowingAsyncSequenceProducer<
+        Element,
+        Error,
+        BackPressureStrategy,
+        KafkaConsumerCloseOnTerminate
+    >
     let wrappedSequence: WrappedSequence
 
     /// `AsynceIteratorProtocol` implementation for handling messages received from the Kafka cluster (``KafkaConsumerMessage``).
     public struct ConsumerMessagesAsyncIterator: AsyncIteratorProtocol {
-        var wrappedIterator: WrappedSequence.AsyncIterator
+        let stateMachine: NIOLockedValueBox<KafkaConsumer.StateMachine>
+        var wrappedIterator: WrappedSequence.AsyncIterator?
 
-        public mutating func next() async -> Element? {
-            await self.wrappedIterator.next()
+        public mutating func next() async throws -> Element? {
+            guard let element = try await self.wrappedIterator?.next() else {
+                self.deallocateIterator()
+                return nil
+            }
+
+            let action = self.stateMachine.withLockedValue { $0.storeOffset() }
+            switch action {
+            case .storeOffset(let client):
+                do {
+                    try client.storeMessageOffset(element)
+                } catch {
+                    self.deallocateIterator()
+                    throw error
+                }
+            }
+            return element
+        }
+
+        private mutating func deallocateIterator() {
+            self.wrappedIterator = nil
         }
     }
 
     public func makeAsyncIterator() -> ConsumerMessagesAsyncIterator {
-        return ConsumerMessagesAsyncIterator(wrappedIterator: self.wrappedSequence.makeAsyncIterator())
+        return ConsumerMessagesAsyncIterator(
+            stateMachine: self.stateMachine,
+            wrappedIterator: self.wrappedSequence.makeAsyncIterator()
+        )
     }
 }
 
 // MARK: - KafkaConsumer
 
 /// Receive messages from the Kafka cluster.
 public final class KafkaConsumer: Sendable, Service {
-    typealias Producer = NIOAsyncSequenceProducer<
+    typealias Producer = NIOThrowingAsyncSequenceProducer<
         KafkaConsumerMessage,
+        Error,
         NIOAsyncSequenceProducerBackPressureStrategies.NoBackPressure,
-        NoDelegate
+        KafkaConsumerCloseOnTerminate
     >
     /// The configuration object of the consumer client.
     private let config: KafkaConsumerConfiguration
@@ -90,13 +131,14 @@ public final class KafkaConsumer: Sendable, Service {
 
         self.stateMachine = NIOLockedValueBox(StateMachine(logger: self.logger))
 
-        let sourceAndSequence = NIOAsyncSequenceProducer.makeSequence(
+        let sourceAndSequence = NIOThrowingAsyncSequenceProducer.makeSequence(
             elementType: KafkaConsumerMessage.self,
             backPressureStrategy: NIOAsyncSequenceProducerBackPressureStrategies.NoBackPressure(),
-            delegate: NoDelegate()
+            delegate: KafkaConsumerCloseOnTerminate(stateMachine: self.stateMachine)
         )
 
         self.messages = KafkaConsumerMessages(
+            stateMachine: self.stateMachine,
             wrappedSequence: sourceAndSequence.sequence
         )
 
@@ -141,7 +183,8 @@ public final class KafkaConsumer: Sendable, Service {
     /// Assign the``KafkaConsumer`` to a specific `partition` of a `topic`.
     /// - Parameter topic: Name of the topic that this ``KafkaConsumer`` will read from.
     /// - Parameter partition: Partition that this ``KafkaConsumer`` will read from.
-    /// - Parameter offset: The topic offset where reading begins. Defaults to the offset of the last read message.
+    /// - Parameter offset: The offset to start consuming from.
+    /// Defaults to the end of the Kafka partition queue (meaning wait for next produced message).
     /// - Throws: A ``KafkaError`` if the consumer could not be assigned to the topic + partition pair.
     private func assign(
         topic: String,
@@ -190,10 +233,11 @@ public final class KafkaConsumer: Sendable, Service {
                     }
                 }
                 try await Task.sleep(for: self.config.pollInterval)
-            case .pollUntilClosed(let client):
-                // Ignore poll result, we are closing down and just polling to commit
-                // outstanding consumer state
-                let _ = client.eventPoll()
+            case .pollWithoutYield(let client):
+                // Ignore poll result.
+                // We are just polling to serve any remaining events queued inside of `librdkafka`.
+                // All remaining queued consumer messages will get dropped and not be committed (marked as read).
+                _ = client.eventPoll()
                 try await Task.sleep(for: self.config.pollInterval)
             case .terminatePollLoop:
                 return
@@ -228,24 +272,26 @@ public final class KafkaConsumer: Sendable, Service {
     private func triggerGracefulShutdown() {
         let action = self.stateMachine.withLockedValue { $0.finish() }
         switch action {
+        case .triggerGracefulShutdown(let client):
+            self._triggerGracefulShutdown(
+                client: client,
+                logger: self.logger
+            )
         case .triggerGracefulShutdownAndFinishSource(let client, let source):
-            self._triggerGracefulShutdownAndFinishSource(
+            source.finish()
+            self._triggerGracefulShutdown(
                 client: client,
-                source: source,
                 logger: self.logger
             )
         case .none:
             return
         }
     }
 
-    private func _triggerGracefulShutdownAndFinishSource(
+    private func _triggerGracefulShutdown(
         client: KafkaClient,
-        source: Producer.Source,
         logger: Logger
     ) {
-        source.finish()
-
         do {
             try client.consumerClose()
         } catch {
@@ -288,6 +334,11 @@ extension KafkaConsumer {
                 client: KafkaClient,
                 source: Producer.Source
             )
+            /// Consumer is still running but the messages asynchronous sequence was terminated.
+            /// All incoming messages will be dropped.
+            ///
+            /// - Parameter client: Client used for handling the connection to the Kafka cluster.
+            case consumptionStopped(client: KafkaClient)
             /// The ``KafkaConsumer/triggerGracefulShutdown()`` has been invoked.
             /// We are now in the process of commiting our last state to the broker.
             ///
@@ -325,11 +376,12 @@ extension KafkaConsumer {
                 client: KafkaClient,
                 source: Producer.Source
             )
-            /// The ``KafkaConsumer`` is in the process of closing down, but still needs to poll
-            /// to commit its state to the broker.
+            /// The ``KafkaConsumer`` stopped consuming messages or
+            /// is in the process of shutting down.
+            /// Poll to serve any queued events and commit outstanding state to the broker.
             ///
             /// - Parameter client: Client used for handling the connection to the Kafka cluster.
-            case pollUntilClosed(client: KafkaClient)
+            case pollWithoutYield(client: KafkaClient)
             /// Terminate the poll loop.
             case terminatePollLoop
         }
@@ -346,12 +398,14 @@ extension KafkaConsumer {
                 fatalError("Subscribe to consumer group / assign to topic partition pair before reading messages")
             case .consuming(let client, let source):
                 return .pollForAndYieldMessage(client: client, source: source)
+            case .consumptionStopped(let client):
+                return .pollWithoutYield(client: client)
             case .finishing(let client):
                 if client.isConsumerClosed {
                     self.state = .finished
                     return .terminatePollLoop
                 } else {
-                    return .pollUntilClosed(client: client)
+                    return .pollWithoutYield(client: client)
                 }
             case .finished:
                 return .terminatePollLoop
@@ -378,11 +432,51 @@ extension KafkaConsumer {
                     source: source
                 )
                 return .setUpConnection(client: client)
-            case .consuming, .finishing, .finished:
+            case .consuming, .consumptionStopped, .finishing, .finished:
                 fatalError("\(#function) should only be invoked upon initialization of KafkaConsumer")
             }
         }
 
+        /// The messages asynchronous sequence was terminated.
+        /// All incoming messages will be dropped.
+        mutating func messageSequenceTerminated() {
+            switch self.state {
+            case .uninitialized:
+                fatalError("\(#function) invoked while still in state \(self.state)")
+            case .initializing:
+                fatalError("Call to \(#function) before setUpConnection() was invoked")
+            case .consumptionStopped:
+                fatalError("messageSequenceTerminated() must not be invoked more than once")
+            case .consuming(let client, _):
+                self.state = .consumptionStopped(client: client)
+            case .finishing, .finished:
+                break
+            }
+        }
+
+        /// Action to take when wanting to store a message offset (to be auto-committed by `librdkafka`).
+        enum StoreOffsetAction {
+            /// Store the message offset with the given `client`.
+            /// - Parameter client: Client used for handling the connection to the Kafka cluster.
+            case storeOffset(client: KafkaClient)
+        }
+
+        /// Get action to take when wanting to store a message offset (to be auto-committed by `librdkafka`).
+        func storeOffset() -> StoreOffsetAction {
+            switch self.state {
+            case .uninitialized:
+                fatalError("\(#function) invoked while still in state \(self.state)")
+            case .initializing:
+                fatalError("Subscribe to consumer group / assign to topic partition pair before committing offsets")
+            case .consumptionStopped:
+                fatalError("Cannot store offset when consumption has been stopped")
+            case .consuming(let client, _):
+                return .storeOffset(client: client)
+            case .finishing, .finished:
+                fatalError("\(#function) invoked while still in state \(self.state)")
+            }
+        }
+
         /// Action to be taken when wanting to do a synchronous commit.
         enum CommitSyncAction {
             /// Do a synchronous commit.
@@ -405,6 +499,8 @@ extension KafkaConsumer {
                 fatalError("\(#function) invoked while still in state \(self.state)")
             case .initializing:
                 fatalError("Subscribe to consumer group / assign to topic partition pair before committing offsets")
+            case .consumptionStopped:
+                fatalError("Cannot commit when consumption has been stopped")
             case .consuming(let client, _):
                 return .commitSync(client: client)
             case .finishing, .finished:
@@ -414,6 +510,12 @@ extension KafkaConsumer {
 
         /// Action to be taken when wanting to do close the consumer.
         enum FinishAction {
+            /// Shut down the ``KafkaConsumer``.
+            ///
+            /// - Parameter client: Client used for handling the connection to the Kafka cluster.
+            case triggerGracefulShutdown(
+                client: KafkaClient
+            )
             /// Shut down the ``KafkaConsumer`` and finish the given `source` object.
             ///
             /// - Parameter client: Client used for handling the connection to the Kafka cluster.
@@ -440,6 +542,9 @@ extension KafkaConsumer {
                     client: client,
                     source: source
                 )
+            case .consumptionStopped(let client):
+                self.state = .finishing(client: client)
+                return .triggerGracefulShutdown(client: client)
             case .finishing, .finished:
                 return nil
             }

Sources/SwiftKafka/KafkaProducer.swift

@@ -425,7 +425,7 @@ extension KafkaProducer {
             case .uninitialized:
                 fatalError("\(#function) invoked while still in state \(self.state)")
             case .consumptionStopped:
-                fatalError("stopConsuming() must not be invoked more than once")
+                fatalError("messageSequenceTerminated() must not be invoked more than once")
             case .started(let client, _, let source, _):
                 self.state = .consumptionStopped(client: client)
                 return .finishSource(source: source)

Sources/SwiftKafka/RDKafka/RDKafka.swift

@@ -33,7 +33,13 @@ struct RDKafka {
         let clientType = type == .producer ? RD_KAFKA_PRODUCER : RD_KAFKA_CONSUMER
 
         let rdConfig = try RDKafkaConfig.createFrom(configDictionary: configDictionary)
+        // Manually override some of the configuration options
+        // Handle logs in event queue
         try RDKafkaConfig.set(configPointer: rdConfig, key: "log.queue", value: "true")
+        // KafkaConsumer is manually storing read offsets
+        if type == .consumer {
+            try RDKafkaConfig.set(configPointer: rdConfig, key: "enable.auto.offset.store", value: "false")
+        }
         RDKafkaConfig.setEvents(configPointer: rdConfig, events: events)
 
         let errorChars = UnsafeMutablePointer<CChar>.allocate(capacity: KafkaClient.stringSize)