KafkaProducer: refactor flushing (#94)

* `KafkaProducer`: refactor flushing

Motivation:

With our own flushing implementation we were flushing until `rd_kakfa_outq_len`
reached ``0`. However `rd_kakfa_outq_len` also takes other events such
as statistics into account which led to a race where we were flushing
for a very long time because more and more other events were produced to
the queue while we were flushing.

Modifications:

* `KafkaClient`:
    * remove `var outgoingQueueSize`
    * add new method `flush(timeoutMilliseconds:)` that executes the
      blocking `rd_kafka_flush` call on a `DispatchQueue` but vends this
      as an `async func`
* `KafkaProducer`:
    * rename `KafkaProducer.StateMachine.State.flushing` to
      `KafkaProducer.StateMachine.State.finishing`
    * invoke `KafkaClient.flush` before terminating poll loop

* `KafkaProducerConfiguration`: add flush timeout

* Make `KafkaProducer.flushTimeoutMilliseconds` `Int`

* Review Franz

Modifications:

* add messages to `preconditions`
* remove extra line in `RDKafkaClient.flush`'s parameter documentation

Author: felixschlegel

Sources/SwiftKafka/Configuration/KafkaProducerConfiguration.swift

@@ -20,6 +20,17 @@ public struct KafkaProducerConfiguration {
     /// Default: `.milliseconds(100)`
     public var pollInterval: Duration = .milliseconds(100)
 
+    /// Maximum timeout for flushing outstanding produce requests when the ``KakfaProducer`` is shutting down.
+    /// Default: `10000`
+    public var flushTimeoutMilliseconds: Int = 10000 {
+        didSet {
+            precondition(
+                0...Int(Int32.max) ~= self.flushTimeoutMilliseconds,
+                "Flush timeout outside of valid range \(0...Int32.max)"
+            )
+        }
+    }
+
     // MARK: - Producer-specific Config Properties
 
     /// When set to true, the producer will ensure that messages are successfully produced exactly once and in the original produce order. The following configuration properties are adjusted automatically (if not modified by the user) when idempotence is enabled: max.in.flight.requests.per.connection=5 (must be less than or equal to 5), retries=INT32_MAX (must be greater than 0), acks=all, queuing.strategy=fifo. Producer instantation will fail if user-supplied configuration is incompatible.

Sources/SwiftKafka/KafkaProducer.swift

@@ -222,7 +222,12 @@ public final class KafkaProducer: Service, Sendable {
                     }
                 }
                 try await Task.sleep(for: self.config.pollInterval)
-            case .terminatePollLoopAndFinishSource(let source):
+            case .flushFinishSourceAndTerminatePollLoop(let client, let source):
+                precondition(
+                    0...Int(Int32.max) ~= self.config.flushTimeoutMilliseconds,
+                    "Flush timeout outside of valid range \(0...Int32.max)"
+                )
+                try await client.flush(timeoutMilliseconds: Int32(self.config.flushTimeoutMilliseconds))
                 source?.finish()
                 return
             case .terminatePollLoop:
@@ -297,7 +302,7 @@ extension KafkaProducer {
             ///
             /// - Parameter client: Client used for handling the connection to the Kafka cluster.
             /// - Parameter source: ``NIOAsyncSequenceProducer/Source`` used for yielding new elements.
-            case flushing(
+            case finishing(
                 client: RDKafkaClient,
                 source: Producer.Source?
             )
@@ -336,10 +341,12 @@ 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?)
-            /// Terminate the poll loop and finish the given `NIOAsyncSequenceProducerSource`.
+            /// Flush any outstanding producer messages.
+            /// Then terminate the poll loop and finish the given `NIOAsyncSequenceProducerSource`.
             ///
+            /// - Parameter client: Client used for handling the connection to the Kafka cluster.
             /// - Parameter source: ``NIOAsyncSequenceProducer/Source`` used for yielding new elements.
-            case terminatePollLoopAndFinishSource(source: Producer.Source?)
+            case flushFinishSourceAndTerminatePollLoop(client: RDKafkaClient, source: Producer.Source?)
             /// Terminate the poll loop.
             case terminatePollLoop
         }
@@ -356,13 +363,8 @@ extension KafkaProducer {
                 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 .pollAndYield(client: client, source: source)
-                } else {
-                    self.state = .finished
-                    return .terminatePollLoopAndFinishSource(source: source)
-                }
+            case .finishing(let client, let source):
+                return .flushFinishSourceAndTerminatePollLoop(client: client, source: source)
             case .finished:
                 return .terminatePollLoop
             }
@@ -402,8 +404,8 @@ extension KafkaProducer {
                 )
             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 .finishing:
+                throw KafkaError.connectionClosed(reason: "Producer in the process of finishing")
             case .finished:
                 throw KafkaError.connectionClosed(reason: "Tried to produce a message with a closed producer")
             }
@@ -428,9 +430,9 @@ extension KafkaProducer {
             case .started(let client, _, let source, _):
                 self.state = .consumptionStopped(client: client)
                 return .finishSource(source: source)
-            case .flushing(let client, let source):
+            case .finishing(let client, let source):
                 // Setting source to nil to prevent incoming acknowledgements from buffering in `source`
-                self.state = .flushing(client: client, source: nil)
+                self.state = .finishing(client: client, source: nil)
                 return .finishSource(source: source)
             case .finished:
                 break
@@ -446,10 +448,10 @@ extension KafkaProducer {
             case .uninitialized:
                 fatalError("\(#function) invoked while still in state \(self.state)")
             case .started(let client, _, let source, _):
-                self.state = .flushing(client: client, source: source)
+                self.state = .finishing(client: client, source: source)
             case .consumptionStopped(let client):
-                self.state = .flushing(client: client, source: nil)
-            case .flushing, .finished:
+                self.state = .finishing(client: client, source: nil)
+            case .finishing, .finished:
                 break
             }
         }

Sources/SwiftKafka/RDKafka/RDKafkaClient.swift

@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 import Crdkafka
+import Dispatch
 import Logging
 
 /// Base class for ``KafkaProducer`` and ``KafkaConsumer``,
@@ -389,6 +390,27 @@ final class RDKafkaClient: Sendable {
         }
     }
 
+    /// Flush any outstanding produce requests.
+    ///
+    /// Parameters:
+    ///     - timeoutMilliseconds: Maximum time to wait for outstanding messages to be flushed.
+    func flush(timeoutMilliseconds: Int32) async throws {
+        // rd_kafka_flush is blocking and there is no convenient way to make it non-blocking.
+        // We therefore execute rd_kafka_flush on a DispatchQueue to ensure it gets executed
+        // on a separate thread that is not part of Swift Concurrency's cooperative thread pool.
+        let queue = DispatchQueue(label: "com.swift-server.swift-kafka.flush")
+        try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
+            queue.async {
+                let error = rd_kafka_flush(self.kafkaHandle, timeoutMilliseconds)
+                if error != RD_KAFKA_RESP_ERR_NO_ERROR {
+                    continuation.resume(throwing: KafkaError.rdKafkaError(wrapping: error))
+                } else {
+                    continuation.resume()
+                }
+            }
+        }
+    }
+
     /// Close the consumer asynchronously. This means revoking its assignemnt, committing offsets to broker and
     /// leaving the consumer group (if applicable).
     ///
@@ -407,14 +429,6 @@ final class RDKafkaClient: Sendable {
         rd_kafka_consumer_closed(self.kafkaHandle) == 1
     }
 
-    /// Returns the current out queue length.
-    ///
-    /// This means the number of producer messages that wait to be sent + the number of any
-    /// callbacks that are waiting to be executed by invoking `rd_kafka_poll`.
-    var outgoingQueueSize: Int32 {
-        return rd_kafka_outq_len(self.kafkaHandle)
-    }
-
     /// Scoped accessor that enables safe access to the pointer of the client's Kafka handle.
     /// - Warning: Do not escape the pointer from the closure for later use.
     /// - Parameter body: The closure will use the Kafka handle pointer.