Add headers to KafkaAcknowledgedMessage

When a message is being sent to Kafka the acknowledgement should also contains the headers.
swift-server · Nov 15, 2024 · 39949eb · 39949eb
1 parent 9986c5b
commit 39949eb
Show file tree

Hide file tree

Showing 4 changed files with 83 additions and 81 deletions.
diff --git a/Sources/Kafka/KafkaAcknowledgedMessage.swift b/Sources/Kafka/KafkaAcknowledgedMessage.swift
@@ -27,6 +27,8 @@ public struct KafkaAcknowledgedMessage {
     public var value: ByteBuffer
     /// The offset of the message in its partition.
     public var offset: KafkaOffset
+    /// The headers of the message.
+    public var headers: [KafkaHeader]
 
     /// Initialize ``KafkaAcknowledgedMessage`` from `rd_kafka_message_t` pointer.
     /// - Throws: A ``KafkaAcknowledgedMessageError`` for failed acknowledgements or malformed messages.
@@ -53,7 +55,7 @@ public struct KafkaAcknowledgedMessage {
         self.topic = topic
 
         self.partition = KafkaPartition(rawValue: Int(rdKafkaMessage.partition))
-
+        self.headers = try RDKafkaClient.getHeaders(for: messagePointer)
         if let keyPointer = rdKafkaMessage.key {
             let keyBufferPointer = UnsafeRawBufferPointer(
                 start: keyPointer,

diff --git a/Sources/Kafka/KafkaConsumerMessage.swift b/Sources/Kafka/KafkaConsumerMessage.swift
@@ -66,7 +66,7 @@ public struct KafkaConsumerMessage {
 
         self.partition = KafkaPartition(rawValue: Int(rdKafkaMessage.partition))
 
-        self.headers = try Self.getHeaders(for: messagePointer)
+        self.headers = try RDKafkaClient.getHeaders(for: messagePointer)
 
         if let keyPointer = rdKafkaMessage.key {
             let keyBufferPointer = UnsafeRawBufferPointer(
@@ -91,82 +91,3 @@ extension KafkaConsumerMessage: Hashable {}
 // MARK: - KafkaConsumerMessage + Sendable
 
 extension KafkaConsumerMessage: Sendable {}
-
-// MARK: - Helpers
-
-extension KafkaConsumerMessage {
-    /// Extract ``KafkaHeader``s from a `rd_kafka_message_t` pointer.
-    ///
-    /// - Parameters:
-    ///    - for: Pointer to the `rd_kafka_message_t` object to extract the headers from.
-    private static func getHeaders(
-        for messagePointer: UnsafePointer<rd_kafka_message_t>
-    ) throws -> [KafkaHeader] {
-        var result: [KafkaHeader] = []
-        var headers: OpaquePointer?
-
-        var readStatus = rd_kafka_message_headers(messagePointer, &headers)
-
-        if readStatus == RD_KAFKA_RESP_ERR__NOENT {
-            // No Header Entries
-            return result
-        }
-
-        guard readStatus == RD_KAFKA_RESP_ERR_NO_ERROR else {
-            throw KafkaError.rdKafkaError(wrapping: readStatus)
-        }
-
-        guard let headers else {
-            return result
-        }
-
-        let headerCount = rd_kafka_header_cnt(headers)
-        result.reserveCapacity(headerCount)
-
-        var headerIndex = 0
-
-        while readStatus != RD_KAFKA_RESP_ERR__NOENT && headerIndex < headerCount {
-            var headerKeyPointer: UnsafePointer<CChar>?
-            var headerValuePointer: UnsafeRawPointer?
-            var headerValueSize = 0
-
-            readStatus = rd_kafka_header_get_all(
-                headers,
-                headerIndex,
-                &headerKeyPointer,
-                &headerValuePointer,
-                &headerValueSize
-            )
-
-            if readStatus == RD_KAFKA_RESP_ERR__NOENT {
-                // No Header Entries
-                return result
-            }
-
-            guard readStatus == RD_KAFKA_RESP_ERR_NO_ERROR else {
-                throw KafkaError.rdKafkaError(wrapping: readStatus)
-            }
-
-            guard let headerKeyPointer else {
-                fatalError("Found null pointer when reading KafkaConsumerMessage header key")
-            }
-            let headerKey = String(cString: headerKeyPointer)
-
-            var headerValue: ByteBuffer?
-            if let headerValuePointer, headerValueSize > 0 {
-                let headerValueBufferPointer = UnsafeRawBufferPointer(
-                    start: headerValuePointer,
-                    count: headerValueSize
-                )
-                headerValue = ByteBuffer(bytes: headerValueBufferPointer)
-            }
-
-            let newHeader = KafkaHeader(key: headerKey, value: headerValue)
-            result.append(newHeader)
-
-            headerIndex += 1
-        }
-
-        return result
-    }
-}
diff --git a/Sources/Kafka/RDKafka/RDKafkaClient.swift b/Sources/Kafka/RDKafka/RDKafkaClient.swift
@@ -15,6 +15,7 @@
 import Crdkafka
 import Dispatch
 import Logging
+import NIOCore
 
 import class Foundation.JSONDecoder
 
@@ -616,4 +617,79 @@ public final class RDKafkaClient: Sendable {
     func withKafkaHandlePointer<T>(_ body: (OpaquePointer) throws -> T) rethrows -> T {
         try body(self.kafkaHandle.pointer)
     }
+
+    /// Extract ``KafkaHeader``s from a `rd_kafka_message_t` pointer.
+    ///
+    /// - Parameters:
+    ///    - for: Pointer to the `rd_kafka_message_t` object to extract the headers from.
+    internal static func getHeaders(
+        for messagePointer: UnsafePointer<rd_kafka_message_t>
+    ) throws -> [KafkaHeader] {
+        var result: [KafkaHeader] = []
+        var headers: OpaquePointer?
+
+        var readStatus = rd_kafka_message_headers(messagePointer, &headers)
+
+        if readStatus == RD_KAFKA_RESP_ERR__NOENT {
+            // No Header Entries
+            return result
+        }
+
+        guard readStatus == RD_KAFKA_RESP_ERR_NO_ERROR else {
+            throw KafkaError.rdKafkaError(wrapping: readStatus)
+        }
+
+        guard let headers else {
+            return result
+        }
+
+        let headerCount = rd_kafka_header_cnt(headers)
+        result.reserveCapacity(headerCount)
+
+        var headerIndex = 0
+
+        while readStatus != RD_KAFKA_RESP_ERR__NOENT && headerIndex < headerCount {
+            var headerKeyPointer: UnsafePointer<CChar>?
+            var headerValuePointer: UnsafeRawPointer?
+            var headerValueSize = 0
+
+            readStatus = rd_kafka_header_get_all(
+                headers,
+                headerIndex,
+                &headerKeyPointer,
+                &headerValuePointer,
+                &headerValueSize
+            )
+
+            if readStatus == RD_KAFKA_RESP_ERR__NOENT {
+                // No Header Entries
+                return result
+            }
+
+            guard readStatus == RD_KAFKA_RESP_ERR_NO_ERROR else {
+                throw KafkaError.rdKafkaError(wrapping: readStatus)
+            }
+
+            guard let headerKeyPointer else {
+                fatalError("Found null pointer when reading KafkaConsumerMessage header key")
+            }
+            let headerKey = String(cString: headerKeyPointer)
+
+            var headerValue: ByteBuffer?
+            if let headerValuePointer, headerValueSize > 0 {
+                let headerValueBufferPointer = UnsafeRawBufferPointer(
+                    start: headerValuePointer,
+                    count: headerValueSize
+                )
+                headerValue = ByteBuffer(bytes: headerValueBufferPointer)
+            }
+
+            let newHeader = KafkaHeader(key: headerKey, value: headerValue)
+            result.append(newHeader)
+
+            headerIndex += 1
+        }
+
+        return result
+    }
 }
diff --git a/Tests/KafkaTests/KafkaProducerTests.swift b/Tests/KafkaTests/KafkaProducerTests.swift
@@ -82,8 +82,10 @@ final class KafkaProducerTests: XCTestCase {
             }
 
             let expectedTopic = "test-topic"
+            let headers = [KafkaHeader(key: "some", value: ByteBuffer.init(string: "test"))]
             let message = KafkaProducerMessage(
                 topic: expectedTopic,
+                headers: headers,
                 key: "key",
                 value: "Hello, World!"
             )
@@ -118,6 +120,7 @@ final class KafkaProducerTests: XCTestCase {
             XCTAssertEqual(expectedTopic, receivedMessage.topic)
             XCTAssertEqual(ByteBuffer(string: message.key!), receivedMessage.key)
             XCTAssertEqual(ByteBuffer(string: message.value), receivedMessage.value)
+            XCTAssertEqual(headers, receivedMessage.headers)
 
             // Shutdown the serviceGroup
             await serviceGroup.triggerGracefulShutdown()