ref(kafka): Optimize allocations when creating Kafka headers (#4739)

`OwnedHeaders` actually allocates, with a recent refactor this means we
allocate for every message now, even if there are no headers to be
added.

The wrapper type now handles allocations properly and also optimizes for
the case when we add the rate limited header.

---------

Co-authored-by: Sebastian Zivota <[email protected]>

Author: Dav1dde

relay-kafka/src/producer/mod.rs

@@ -1,4 +1,4 @@
-//! This module contains the kafka producer related code.
+//! This module contains the Kafka producer related code.
 
 use std::borrow::Cow;
 use std::collections::{BTreeMap, HashMap};
@@ -7,7 +7,7 @@ use std::sync::Arc;
 use std::time::{Duration, Instant};
 
 use rdkafka::ClientConfig;
-use rdkafka::message::{Header, OwnedHeaders};
+use rdkafka::message::Header;
 use rdkafka::producer::{BaseRecord, Producer as _};
 use relay_statsd::metric;
 use thiserror::Error;
@@ -16,6 +16,7 @@ use uuid::Uuid;
 use crate::config::{KafkaParams, KafkaTopic};
 use crate::debounced::Debounced;
 use crate::limits::KafkaRateLimits;
+use crate::producer::utils::KafkaHeaders;
 use crate::statsd::{KafkaCounters, KafkaGauges, KafkaHistograms};
 
 mod utils;
@@ -318,15 +319,14 @@ impl Producer {
             topic = topic_name,
         );
 
-        let mut kafka_headers = OwnedHeaders::new();
-        if let Some(headers) = headers {
-            for (key, value) in headers {
-                kafka_headers = kafka_headers.insert(Header {
-                    key,
-                    value: Some(value),
-                });
-            }
-        }
+        let mut headers = headers
+            .unwrap_or(&BTreeMap::new())
+            .iter()
+            .map(|(key, value)| Header {
+                key,
+                value: Some(value),
+            })
+            .collect::<KafkaHeaders>();
 
         let mut key = key;
         if let Some(ref limiter) = self.rate_limiter {
@@ -338,17 +338,18 @@ impl Producer {
                 );
 
                 key = Uuid::new_v4().into_bytes();
-                kafka_headers = kafka_headers.insert(Header {
+                headers.insert(Header {
                     key: "sentry-reshuffled",
                     value: Some("1"),
                 });
             }
         }
 
-        let record = BaseRecord::to(topic_name)
-            .key(&key)
-            .payload(payload)
-            .headers(kafka_headers);
+        let mut record = BaseRecord::to(topic_name).key(&key).payload(payload);
+
+        if let Some(headers) = headers.into_inner() {
+            record = record.headers(headers);
+        }
 
         self.metrics.debounce(now, || {
             metric!(

relay-kafka/src/producer/utils.rs

@@ -1,11 +1,77 @@
 use std::error::Error;
 
+use rdkafka::message::{Header, OwnedHeaders, ToBytes};
 use rdkafka::producer::{DeliveryResult, ProducerContext};
 use rdkafka::{ClientContext, Message};
 use relay_statsd::metric;
 
 use crate::statsd::{KafkaCounters, KafkaGauges};
 
+/// A thin wrapper around [`OwnedHeaders`].
+///
+/// Unlike [`OwnedHeaders`], this will not allocate on creation.
+/// Allocations are tuned for the use-case in a [`super::Producer`].
+pub struct KafkaHeaders(Option<OwnedHeaders>);
+
+impl KafkaHeaders {
+    pub fn new() -> Self {
+        Self(None)
+    }
+
+    pub fn insert<V>(&mut self, header: Header<'_, &V>)
+    where
+        V: ToBytes + ?Sized,
+    {
+        self.extend(Some(header));
+    }
+
+    pub fn into_inner(self) -> Option<OwnedHeaders> {
+        self.0
+    }
+}
+
+impl<'a, 'b, V> Extend<Header<'a, &'b V>> for KafkaHeaders
+where
+    V: ToBytes + ?Sized,
+{
+    fn extend<T: IntoIterator<Item = Header<'a, &'b V>>>(&mut self, iter: T) {
+        let mut iter = iter.into_iter();
+
+        // Probe if the iterator is empty, if it is empty, no need to do anything.
+        let Some(first) = iter.next() else {
+            return;
+        };
+
+        let mut headers = self.0.take().unwrap_or_else(|| {
+            // Get a size hint from the iterator, +2 for the already removed
+            // first element and reserving space for 1 extra header which is conditionally
+            // added by the `Producer` in this crate.
+            //
+            // This means we might allocate a little bit too much, but we never have to resize
+            // and allocate a second time, a good trade-off.
+            let size = iter.size_hint().0 + 2;
+            OwnedHeaders::new_with_capacity(size)
+        });
+        headers = headers.insert(first);
+        for remaining in iter {
+            headers = headers.insert(remaining);
+        }
+
+        self.0 = Some(headers);
+    }
+}
+
+impl<'a, 'b, V> FromIterator<Header<'a, &'b V>> for KafkaHeaders
+where
+    V: ToBytes + ?Sized,
+{
+    fn from_iter<I: IntoIterator<Item = Header<'a, &'b V>>>(iter: I) -> Self {
+        let mut c = Self::new();
+        c.extend(iter);
+        c
+    }
+}
+
 /// Kafka client and producer context that logs statistics and producer errors.
 #[derive(Debug)]
 pub struct Context;