Poll for messages using TaskExecutor

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 #165

Author: mimischi

Sources/Kafka/KafkaConsumer.swift

@@ -12,6 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+import Dispatch
 import Logging
 import NIOConcurrencyHelpers
 import NIOCore
@@ -72,6 +73,7 @@ public struct KafkaConsumerMessages: Sendable, AsyncSequence {
     public struct AsyncIterator: AsyncIteratorProtocol {
         private let stateMachineHolder: MachineHolder
         let pollInterval: Duration
+        let queue: NaiveQueueExecutor
 
         private final class MachineHolder: Sendable {  // only for deinit
             let stateMachine: LockedMachine
@@ -88,21 +90,30 @@ public struct KafkaConsumerMessages: Sendable, AsyncSequence {
         init(stateMachine: LockedMachine, pollInterval: Duration) {
             self.stateMachineHolder = .init(stateMachine: stateMachine)
             self.pollInterval = pollInterval
+            self.queue = NaiveQueueExecutor(DispatchQueue(label: "com.swift-server.swift-kafka.message-consumer"))
         }
 
         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
                     }
-                    try await Task.sleep(for: self.pollInterval)
+
+                    #if swift(>=6.0)
+                    // Wait on a separate thread for the next message.
+                    return try await withTaskExecutorPreference(queue) {
+                        try client.consumerPoll(for: Int32(self.pollInterval.inMilliseconds))
+                    }
+                    #else
+                    // No messages. Sleep a little.
+                    return 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

@@ -446,8 +446,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/NaiveQueueExecutor.swift

@@ -0,0 +1,36 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+import Dispatch
+
+final class NaiveQueueExecutor: 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)
+  }
+}