Poll for messages using TaskExecutor (swift-server#178)

We currently sleep for `pollInterval` when no new messages have been
polled from the cluster. This leads to unnecessary slowness of the
client. Instead of doing that, we now break up the polling of messages
into two distinct approaches:

1. Attempt to poll synchronously: if there a message is polled, we
return it. If there is no message, we immediately go to step 2.
2. We create a `DispatchQueue` and run the `consumerPoll` on it using
`withTaskExecutorPreference`. We make the `consumerPoll` call wait for
up to `pollInterval` before bailing.

This prevents us from sleeping on the running thread, and frees up
cycles to do other work if required.

Resolves swift-server#165

---------

Co-authored-by: Konrad `ktoso` Malawski <[email protected]>

Author: mimischi

Sources/Kafka/KafkaConsumer.swift

@@ -72,6 +72,9 @@ public struct KafkaConsumerMessages: Sendable, AsyncSequence {
     public struct AsyncIterator: AsyncIteratorProtocol {
         private let stateMachineHolder: MachineHolder
         let pollInterval: Duration
+        #if swift(>=6.0)
+        private let queue: DispatchQueueTaskExecutor
+        #endif
 
         private final class MachineHolder: Sendable {  // only for deinit
             let stateMachine: LockedMachine
@@ -88,21 +91,35 @@ public struct KafkaConsumerMessages: Sendable, AsyncSequence {
         init(stateMachine: LockedMachine, pollInterval: Duration) {
             self.stateMachineHolder = .init(stateMachine: stateMachine)
             self.pollInterval = pollInterval
+            #if swift(>=6.0)
+            self.queue = DispatchQueueTaskExecutor(
+                DispatchQueue(label: "com.swift-server.swift-kafka.message-consumer")
+            )
+            #endif
         }
 
         public func next() async throws -> Element? {
-            // swift-kafka-client issue: https://github.com/swift-server/swift-kafka-client/issues/165
-            // Currently use Task.sleep() if no new messages, should use task executor preference when implemented:
-            // https://github.com/apple/swift-evolution/blob/main/proposals/0417-task-executor-preference.md
             while !Task.isCancelled {
                 let action = self.stateMachineHolder.stateMachine.withLockedValue { $0.nextConsumerPollLoopAction() }
 
                 switch action {
                 case .poll(let client):
-                    if let message = try client.consumerPoll() {  // non-blocking call
+                    // Attempt to fetch a message synchronously. Bail
+                    // immediately if no message is waiting for us.
+                    if let message = try client.consumerPoll() {
                         return message
                     }
+
+                    #if swift(>=6.0)
+                    // Wait on a separate thread for the next message.
+                    // The call below will block for `pollInterval`.
+                    return try await withTaskExecutorPreference(queue) {
+                        try client.consumerPoll(for: Int32(self.pollInterval.inMilliseconds))
+                    }
+                    #else
+                    // No messages. Sleep a little.
                     try await Task.sleep(for: self.pollInterval)
+                    #endif
                 case .suspendPollLoop:
                     try await Task.sleep(for: self.pollInterval)  // not started yet
                 case .terminatePollLoop:

Sources/Kafka/RDKafka/RDKafkaClient.swift

@@ -447,8 +447,8 @@ public final class RDKafkaClient: Sendable {
     ///
     /// - Returns: A ``KafkaConsumerMessage`` or `nil` if there are no new messages.
     /// - Throws: A ``KafkaError`` if the received message is an error message or malformed.
-    func consumerPoll() throws -> KafkaConsumerMessage? {
-        guard let messagePointer = rd_kafka_consumer_poll(self.kafkaHandle.pointer, 0) else {
+    func consumerPoll(for pollTimeoutMs: Int32 = 0) throws -> KafkaConsumerMessage? {
+        guard let messagePointer = rd_kafka_consumer_poll(self.kafkaHandle.pointer, pollTimeoutMs) else {
             // No error, there might be no more messages
             return nil
         }

Sources/Kafka/Utilities/DispatchQueueTaskExecutor.swift

@@ -0,0 +1,39 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the swift-kafka-client open source project
+//
+// Copyright (c) 2024 Apple Inc. and the swift-kafka-client project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of swift-kafka-client project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+#if swift(>=6.0)
+import Dispatch
+
+final class DispatchQueueTaskExecutor: TaskExecutor {
+    let queue: DispatchQueue
+
+    init(_ queue: DispatchQueue) {
+        self.queue = queue
+    }
+
+    public func enqueue(_ _job: consuming ExecutorJob) {
+        let job = UnownedJob(_job)
+        queue.async {
+            job.runSynchronously(
+                on: self.asUnownedTaskExecutor()
+            )
+        }
+    }
+
+    @inlinable
+    public func asUnownedTaskExecutor() -> UnownedTaskExecutor {
+        UnownedTaskExecutor(ordinary: self)
+    }
+}
+#endif