add approx. transactional api

Author: blindspotbounty

Package.swift

@@ -44,6 +44,7 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-nio.git", from: "2.55.0"),
         .package(url: "https://github.com/swift-server/swift-service-lifecycle.git", from: "2.0.0-alpha.1"),
         .package(url: "https://github.com/apple/swift-log.git", from: "1.0.0"),
+        .package(url: "https://github.com/ordo-one/package-concurrency-helpers", .upToNextMajor(from: "1.0.0")),
         // The zstd Swift package produces warnings that we cannot resolve:
         // https://github.com/facebook/zstd/issues/3328
         .package(url: "https://github.com/facebook/zstd.git", from: "1.5.0"),
@@ -73,6 +74,7 @@ let package = Package(
             name: "SwiftKafka",
             dependencies: [
                 "Crdkafka",
+                .product(name: "ConcurrencyHelpers", package: "package-concurrency-helpers", moduleAliases: ["ConcurrencyHelpers" : "BlockingCallWrapper"]),
                 .product(name: "NIOCore", package: "swift-nio"),
                 .product(name: "ServiceLifecycle", package: "swift-service-lifecycle"),
                 .product(name: "Logging", package: "swift-log"),

Sources/SwiftKafka/Configuration/KafkaProducerConfiguration.swift

@@ -36,6 +36,8 @@ public struct KafkaProducerConfiguration {
     /// 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.
     /// Default: `false`
     public var enableIdempotence: Bool = false
+    
+    public var transactionalId: String?
 
     /// Producer queue options.
     public var queue: KafkaConfiguration.QueueOptions = .init()
@@ -108,6 +110,9 @@ extension KafkaProducerConfiguration {
         var resultDict: [String: String] = [:]
 
         resultDict["enable.idempotence"] = String(self.enableIdempotence)
+        if let transactionalId {
+            resultDict["transactional.id"] = transactionalId
+        }
         resultDict["queue.buffering.max.messages"] = String(self.queue.bufferingMaxMessages)
         resultDict["queue.buffering.max.kbytes"] = String(self.queue.bufferingMaxKBytes)
         resultDict["queue.buffering.max.ms"] = String(self.queue.bufferingMaxMilliseconds)

Sources/SwiftKafka/KafkaConsumer.swift

@@ -301,6 +301,10 @@ public final class KafkaConsumer: Sendable, Service {
             }
         }
     }
+    
+    func client() throws -> RDKafkaClient {
+        return try stateMachine.withLockedValue { try $0.client() }
+    }
 }
 
 // MARK: - KafkaConsumer + StateMachine
@@ -548,5 +552,22 @@ extension KafkaConsumer {
                 return nil
             }
         }
+        
+        func client() throws -> RDKafkaClient {
+            switch self.state {
+            case .uninitialized:
+                fatalError("\(#function) invoked while still in state \(self.state)")
+            case .initializing(let client, _):
+                return client
+            case .consuming(let client, _):
+                return client
+            case .consumptionStopped(let client):
+                return client
+            case .finishing(let client):
+                return client
+            case .finished:
+                throw KafkaError.client(reason: "Client is stopped")
+            }
+        }
     }
 }

Sources/SwiftKafka/KafkaError.swift

@@ -54,7 +54,7 @@ public struct KafkaError: Error, CustomStringConvertible {
     }
 
     static func rdKafkaError(
-        wrapping error: rd_kafka_resp_err_t, file: String = #fileID, line: UInt = #line
+        wrapping error: rd_kafka_resp_err_t, isFatal: Bool = false, file: String = #fileID, line: UInt = #line
     ) -> KafkaError {
         let errorMessage = String(cString: rd_kafka_err2str(error))
         return KafkaError(
@@ -143,6 +143,36 @@ public struct KafkaError: Error, CustomStringConvertible {
             )
         )
     }
+    
+    static func transactionAborted(
+        reason: String, file: String = #fileID, line: UInt = #line
+    ) -> KafkaError {
+        return KafkaError(
+            backing: .init(
+                code: .transactionAborted, reason: reason, file: file, line: line
+            )
+        )
+    }
+    
+    static func transactionIncomplete(
+        reason: String, file: String = #fileID, line: UInt = #line
+    ) -> KafkaError {
+        return KafkaError(
+            backing: .init(
+                code: .transactionIncomplete, reason: reason, file: file, line: line
+            )
+        )
+    }
+    
+    static func transactionOutOfAttempts(
+        numOfAttempts: UInt64, file: String = #fileID, line: UInt = #line
+    ) -> KafkaError {
+        return KafkaError(
+            backing: .init(
+                code: .transactionOutOfAttempts, reason: "Out of \(numOfAttempts) attempts", file: file, line: line
+            )
+        )
+    }
 }
 
 extension KafkaError {
@@ -162,6 +192,10 @@ extension KafkaError {
             case messageConsumption
             case topicCreation
             case topicDeletion
+            case transactionAborted
+            case transactionIncomplete
+            case notInTransaction // FIXME: maybe add subcode ?
+            case transactionOutOfAttempts
         }
 
         fileprivate var backingCode: BackingCode
@@ -188,6 +222,12 @@ extension KafkaError {
         public static let topicCreation = ErrorCode(.topicCreation)
         /// Deleting a topic failed.
         public static let topicDeletion = ErrorCode(.topicDeletion)
+        /// Transaction was aborted (can be re-tried from scratch).
+        public static let transactionAborted = ErrorCode(.transactionAborted)
+        /// Transaction could not be completed
+        public static let transactionIncomplete = ErrorCode(.transactionIncomplete)
+        /// Out of provided number of attempts
+        public static let transactionOutOfAttempts = ErrorCode(.transactionOutOfAttempts)
 
         public var description: String {
             return String(describing: self.backingCode)
@@ -206,17 +246,21 @@ extension KafkaError {
         let file: String
 
         let line: UInt
+        
+        let isFatal: Bool
 
         fileprivate init(
             code: KafkaError.ErrorCode,
             reason: String,
             file: String,
-            line: UInt
+            line: UInt,
+            isFatal: Bool = false
         ) {
             self.code = code
             self.reason = reason
             self.file = file
             self.line = line
+            self.isFatal = isFatal
         }
 
         // Only the error code matters for equality.

Sources/SwiftKafka/KafkaProducer.swift

@@ -265,6 +265,43 @@ public final class KafkaProducer: Service, Sendable {
             return KafkaProducerMessageID(rawValue: newMessageID)
         }
     }
+    
+    func initTransactions(timeout: Duration = .seconds(5)) async throws {
+        guard config.transactionalId != nil else {
+            throw KafkaError.config(
+                reason: "Could not initialize transactions because transactionalId is not set in config")
+        }
+        // FIXME: maybe add state 'startedWithTransactions'?
+        let client = try self.stateMachine.withLockedValue { try $0.transactionsClient() }
+        try await client.initTransactions(timeout: timeout)
+    }
+    
+    func beginTransaction() throws {
+        let client = try self.stateMachine.withLockedValue { try $0.transactionsClient() }
+        try client.beginTransaction()
+    }
+    
+    func send(
+        offsets: RDKafkaTopicPartitionList,
+        forConsumer consumer: KafkaConsumer,
+        timeout: Duration = .kafkaUntilEndOfTransactionTimeout,
+        attempts: UInt64 = .max
+    ) async throws {
+        let client = try self.stateMachine.withLockedValue { try $0.transactionsClient() }
+        let consumerClient = try consumer.client()
+        try await consumerClient.withKafkaHandlePointer {
+            try await client.send(attempts: attempts, offsets: offsets, forConsumerKafkaHandle: $0, timeout: timeout)
+        }
+    }
+    
+    func abortTransaction(
+        timeout: Duration = .kafkaUntilEndOfTransactionTimeout,
+        attempts: UInt64) async throws {
+        let client = try self.stateMachine.withLockedValue { try $0.transactionsClient() }
+            try await client.abortTransaction(attempts: attempts, timeout: timeout)
+    }
+    
+    
 }
 
 // MARK: - KafkaProducer + StateMachine
@@ -455,5 +492,22 @@ extension KafkaProducer {
                 break
             }
         }
+
+        // TODO:
+        // 1. add client()
+        // 2. initTransactions() -> change state to startedWithTransactions
+        // 3. transactionsClient() -> return client only for startedWithTransactions
+        
+        
+        func transactionsClient() throws -> RDKafkaClient {
+            switch self.state {
+            case .uninitialized:
+                fatalError("\(#function) invoked while still in state \(self.state)")
+            case .started(let client, _, _, _):
+                return client
+            default:
+                throw KafkaError.connectionClosed(reason: "Producer is stopping or finished")
+            }
+        }
     }
 }

Sources/SwiftKafka/RDKafka/RDKafkaClient.swift

@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 import Crdkafka
+import BlockingCallWrapper
 import Dispatch
 import Logging
 
@@ -436,6 +437,14 @@ final class RDKafkaClient: Sendable {
     func withKafkaHandlePointer<T>(_ body: (OpaquePointer) throws -> T) rethrows -> T {
         return try body(self.kafkaHandle)
     }
+    
+    /// Scoped accessor that enables safe access to the pointer of the client's Kafka handle with async closure.
+    /// - Warning: Do not escape the pointer from the closure for later use.
+    /// - Parameter body: The closure will use the Kafka handle pointer.
+    @discardableResult
+    func withKafkaHandlePointer<T>(_ body: (OpaquePointer) async throws -> T) async rethrows -> T {
+        return try await body(self.kafkaHandle)
+    }
 
     /// Convert an unsafe`rd_kafka_message_t` object to a safe ``KafkaAcknowledgementResult``.
     /// - Parameter messagePointer: An `UnsafePointer` pointing to the `rd_kafka_message_t` object in memory.
@@ -462,4 +471,138 @@ final class RDKafkaClient: Sendable {
 
         return messageResult
     }
+    
+    func initTransactions(timeout: Duration) async throws {
+        let result = await forBlockingFunc {
+            rd_kafka_init_transactions(self.kafkaHandle, Int32(timeout.totalMilliseconds))
+        }
+        
+        if result != nil {
+            let code = rd_kafka_error_code(result)
+            rd_kafka_error_destroy(result)
+            throw KafkaError.rdKafkaError(wrapping: code)
+        }
+    }
+    
+    func beginTransaction() throws {
+        let result = rd_kafka_begin_transaction(kafkaHandle)
+        if result != nil {
+            let code = rd_kafka_error_code(result)
+            rd_kafka_error_destroy(result)
+            throw KafkaError.rdKafkaError(wrapping: code)
+        }
+    }
+    
+    func send(
+        attempts: UInt64,
+        offsets: RDKafkaTopicPartitionList,
+        forConsumerKafkaHandle consumer: OpaquePointer,
+        timeout: Duration) async throws {
+            try await offsets.withListPointer { topicPartitionList in
+                
+                let consumerMetadata = rd_kafka_consumer_group_metadata(consumer)
+                defer { rd_kafka_consumer_group_metadata_destroy(consumerMetadata) }
+                
+                // TODO: actually it should be withing some timeout (like transaction timeout or session timeout)
+                for idx in 0..<attempts {
+                    let error = await forBlockingFunc {
+                        rd_kafka_send_offsets_to_transaction(self.kafkaHandle, topicPartitionList,
+                                                             consumerMetadata, timeout.totalMillisecondsOrMinusOne)
+                    }
+                    
+                    /* check if offset commit is completed successfully  */
+                    if error == nil {
+                        return
+                    }
+                    defer { rd_kafka_error_destroy(error) }
+                    
+                    /* check if offset commit is retriable */
+                    if rd_kafka_error_is_retriable(error) == 1 {
+                        continue
+                    }
+                    
+                    /* check if transaction need to be aborted */
+                    if rd_kafka_error_txn_requires_abort(error) == 1 {
+                        do {
+                            try await abortTransaction(attempts: attempts - idx, timeout: timeout)
+                            throw KafkaError.transactionAborted(reason: "Transaction aborted and can be started from scratch")
+                        } catch {
+                            throw KafkaError.transactionIncomplete(
+                                reason: "Could not complete or abort transaction with error \(error)")
+                        }
+                    }
+                    let isFatal = (rd_kafka_error_is_fatal(error) == 1) // fatal when Producer/Consumer must be restarted
+                    throw KafkaError.rdKafkaError(wrapping: rd_kafka_error_code(error), isFatal: isFatal)
+                }
+            }
+            throw KafkaError.transactionOutOfAttempts(numOfAttempts: attempts)
+        }
+    
+    func abortTransaction(attempts: UInt64, timeout: Duration) async throws {
+        for _ in 0..<attempts {
+            let error = await forBlockingFunc {
+                rd_kafka_abort_transaction(self.kafkaHandle, timeout.totalMillisecondsOrMinusOne)
+            }
+            /* check if transaction abort is completed successfully  */
+            if error == nil {
+                return
+            }
+            defer { rd_kafka_error_destroy(error) }
+
+            /* check if transaction abort is retriable */
+            if rd_kafka_error_is_retriable(error) == 1 {
+                continue
+            }
+            let isFatal = (rd_kafka_error_is_fatal(error) == 1) // fatal when Producer/Consumer must be restarted
+            throw KafkaError.rdKafkaError(wrapping: rd_kafka_error_code(error), isFatal: isFatal)
+        }
+        throw KafkaError.transactionOutOfAttempts(numOfAttempts: attempts)
+    }
+    
+    func commitTransaction(attempts: UInt64, timeout: Duration) async throws {
+        for idx in 0..<attempts {
+            let error = await forBlockingFunc {
+                rd_kafka_commit_transaction(self.kafkaHandle, timeout.totalMillisecondsOrMinusOne)
+            }
+            /* check if transaction is completed successfully  */
+            if error == nil {
+                return
+            }
+            /* check if transaction is retriable */
+            if rd_kafka_error_is_retriable(error) == 1 {
+                continue
+            }
+            defer { rd_kafka_error_destroy(error) }
+            
+            /* check if transaction need to be aborted */
+            if rd_kafka_error_txn_requires_abort(error) == 1 {
+                do {
+                    try await abortTransaction(attempts: attempts - idx, timeout: timeout)
+                    throw KafkaError.transactionAborted(reason: "Transaction aborted and can be started from scratch")
+                } catch {
+                    throw KafkaError.transactionIncomplete(
+                        reason: "Could not complete or abort transaction with error \(error)")
+                }
+            }
+            /* check if error is fatal */
+            let isFatal = (rd_kafka_error_is_fatal(error) == 1) // fatal when Producer/Consumer must be restarted
+            throw KafkaError.rdKafkaError(wrapping: rd_kafka_error_code(error), isFatal: isFatal)
+        }
+        throw KafkaError.transactionOutOfAttempts(numOfAttempts: attempts)
+    }
+}
+
+// TODO: tmp, should be in other PRs
+extension Duration {
+    // Internal usage only: librdkafka accepts Int32 as timeouts
+    var totalMilliseconds: Int32 {
+        return Int32(self.components.seconds * 1000 + self.components.attoseconds / 1_000_000_000_000_000)
+    }
+    
+    var totalMillisecondsOrMinusOne: Int32 {
+        return max(totalMilliseconds, -1)
+    }
+    
+    public static var kafkaUntilEndOfTransactionTimeout: Duration = .milliseconds(-1)
+    public static var kafkaNoWaitTransaction: Duration = .zero
 }