Added sendig header for kafka message for Producer.

userver-framework · Feb 27, 2025 · 6e7567e · 6e7567e
1 parent 903fcf3
commit 6e7567e
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 8 deletions.
diff --git a/kafka/include/userver/kafka/producer.hpp b/kafka/include/userver/kafka/producer.hpp
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <cstdint>
+#include <unordered_map>
 
 #include <userver/engine/task/task_processor_fwd.hpp>
 #include <userver/engine/task/task_with_result.hpp>
@@ -44,6 +45,9 @@ struct Secret;
 ///
 /// @see https://docs.confluent.io/platform/current/clients/producer.html
 class Producer final {
+public:
+    using Headers = std::unordered_map<std::string, std::string>;
+
 public:
     /// @brief Creates the Kafka Producer.
     ///
@@ -95,6 +99,7 @@ class Producer final {
         const std::string& topic_name,
         std::string_view key,
         std::string_view message,
+        Headers headers = {},
         std::optional<std::uint32_t> partition = std::nullopt
     ) const;
 
@@ -111,6 +116,7 @@ class Producer final {
         std::string topic_name,
         std::string key,
         std::string message,
+        Headers headers = {},
         std::optional<std::uint32_t> partition = std::nullopt
     ) const;
 
@@ -124,6 +130,7 @@ class Producer final {
         const std::string& topic_name,
         std::string_view key,
         std::string_view message,
+        Headers headers = {},
         std::optional<std::uint32_t> partition
     ) const;
 

diff --git a/kafka/src/kafka/impl/producer_impl.cpp b/kafka/src/kafka/impl/producer_impl.cpp
@@ -126,10 +126,11 @@ DeliveryResult ProducerImpl::Send(
     const std::string& topic_name,
     std::string_view key,
     std::string_view message,
+    const Headers& headers,
     std::optional<std::uint32_t> partition
 ) const {
     LOG_INFO() << fmt::format("Message to topic '{}' is requested to send", topic_name);
-    auto delivery_result_future = ScheduleMessageDelivery(topic_name, key, message, partition);
+    auto delivery_result_future = ScheduleMessageDelivery(topic_name, key, message, headers, partition);
 
     WaitUntilDeliveryReported(delivery_result_future);
 
@@ -140,6 +141,7 @@ engine::Future<DeliveryResult> ProducerImpl::ScheduleMessageDelivery(
     const std::string& topic_name,
     std::string_view key,
     std::string_view message,
+    const Headers& headers,
     std::optional<std::uint32_t> partition
 ) const {
     auto waiter = std::make_unique<DeliveryWaiter>();
@@ -173,13 +175,26 @@ engine::Future<DeliveryResult> ProducerImpl::ScheduleMessageDelivery(
 #pragma clang diagnostic ignored "-Wgnu-statement-expression"
 #endif
 
+    // We assume that message takes ownership of headers, if rd_kafka_producev()
+    // succeeded. See rd_kafka RD_KAFKA_V_HEADERS reference for details
+    struct KafkaHeaders final {
+        using Ptr = std::unique_ptr<rd_kafka_headers_t, KafkaHeaders>;
+        void operator()(rd_kafka_headers_t* headers) const { rd_kafka_headers_destroy(headers); }
+    };
+
+    KafkaHeaders::Ptr kafka_headers{rd_kafka_headers_new(headers.size())};
+    for (const auto& [header, value] : headers) {
+        rd_kafka_header_add(kafka_headers.get(), header.c_str(), header.size(), value.c_str(), value.size());
+    }
+
     // NOLINTBEGIN(clang-analyzer-cplusplus.NewDeleteLeaks,cppcoreguidelines-pro-type-const-cast)
     const rd_kafka_resp_err_t enqueue_error = rd_kafka_producev(
         producer_.GetHandle(),
         RD_KAFKA_V_TOPIC(topic_name.c_str()),
         RD_KAFKA_V_KEY(key.data(), key.size()),
         RD_KAFKA_V_VALUE(const_cast<char*>(message.data()), message.size()),
         RD_KAFKA_V_MSGFLAGS(0),
+        RD_KAFKA_V_HEADERS(kafka_headers.get()),
         RD_KAFKA_V_PARTITION(partition.value_or(RD_KAFKA_PARTITION_UA)),
         RD_KAFKA_V_OPAQUE(waiter.get()),
         RD_KAFKA_V_END
@@ -191,6 +206,7 @@ engine::Future<DeliveryResult> ProducerImpl::ScheduleMessageDelivery(
 #endif
 
     if (enqueue_error == RD_KAFKA_RESP_ERR_NO_ERROR) {
+        [[maybe_unused]] const auto released = kafka_headers.release();
         [[maybe_unused]] auto _ = waiter.release();
     } else {
         LOG_WARNING(

diff --git a/kafka/src/kafka/impl/producer_impl.hpp b/kafka/src/kafka/impl/producer_impl.hpp
@@ -3,6 +3,7 @@
 #include <chrono>
 #include <cstdint>
 #include <optional>
+#include <unordered_map>
 
 #include <librdkafka/rdkafka.h>
 
@@ -20,6 +21,9 @@ namespace kafka::impl {
 class Configuration;
 
 class ProducerImpl final {
+public:
+    using Headers = std::unordered_map<std::string, std::string>;
+
 public:
     explicit ProducerImpl(Configuration&& configuration);
 
@@ -31,6 +35,7 @@ class ProducerImpl final {
         const std::string& topic_name,
         std::string_view key,
         std::string_view message,
+        const Headers& headers,
         std::optional<std::uint32_t> partition
     ) const;
 
@@ -51,6 +56,7 @@ class ProducerImpl final {
         const std::string& topic_name,
         std::string_view key,
         std::string_view message,
+        const Headers& headers,
         std::optional<std::uint32_t> partition
     ) const;
 

diff --git a/kafka/src/kafka/producer.cpp b/kafka/src/kafka/producer.cpp
@@ -86,11 +86,16 @@ void Producer::Send(
     const std::string& topic_name,
     std::string_view key,
     std::string_view message,
+    Headers headers,
     std::optional<std::uint32_t> partition
 ) const {
-    utils::Async(producer_task_processor_, "producer_send", [this, &topic_name, key, message, partition] {
-        SendImpl(topic_name, key, message, partition);
-    }).Get();
+    utils::Async(
+        producer_task_processor_,
+        "producer_send",
+        [this, &topic_name, key, message, headers = std::move(headers), partition] {
+            SendImpl(topic_name, key, headers, message, headers, partition);
+        }
+    ).Get();
 }
 
 engine::TaskWithResult<void> Producer::SendAsync(
@@ -102,9 +107,12 @@ engine::TaskWithResult<void> Producer::SendAsync(
     return utils::Async(
         producer_task_processor_,
         "producer_send_async",
-        [this, topic_name = std::move(topic_name), key = std::move(key), message = std::move(message), partition] {
-            SendImpl(topic_name, key, message, partition);
-        }
+        [this,
+         topic_name = std::move(topic_name),
+         key = std::move(key),
+         message = std::move(message),
+         headers = std::move(headers),
+         partition] { SendImpl(topic_name, key, message, headers, partition); }
     );
 }
 
@@ -114,11 +122,12 @@ void Producer::SendImpl(
     const std::string& topic_name,
     std::string_view key,
     std::string_view message,
+    const Headers& headers,
     std::optional<std::uint32_t> partition
 ) const {
     tracing::Span::CurrentSpan().AddTag("kafka_producer", name_);
 
-    const impl::DeliveryResult delivery_result = producer_->Send(topic_name, key, message, partition);
+    const impl::DeliveryResult delivery_result = producer_->Send(topic_name, key, message, headers, partition);
     if (!delivery_result.IsSuccess()) {
         ThrowSendError(delivery_result);
     }