relay_kafka/producer/
utils.rs

1use std::error::Error;
2
3use rdkafka::message::{Header, OwnedHeaders, ToBytes};
4use rdkafka::producer::{DeliveryResult, ProducerContext};
5use rdkafka::{ClientContext, Message};
6use relay_statsd::metric;
7use sentry_arroyo::backends::kafka::producer::ProducerContext as ArroyoProducerContext;
8
9use crate::statsd::KafkaCounters;
10
11/// A thin wrapper around [`OwnedHeaders`].
12///
13/// Unlike [`OwnedHeaders`], this will not allocate on creation.
14/// Allocations are tuned for the use-case in a [`super::Producer`].
15pub struct KafkaHeaders(Option<OwnedHeaders>);
16
17impl KafkaHeaders {
18    pub fn new() -> Self {
19        Self(None)
20    }
21
22    pub fn insert<V>(&mut self, header: Header<'_, &V>)
23    where
24        V: ToBytes + ?Sized,
25    {
26        self.extend(Some(header));
27    }
28
29    pub fn into_inner(self) -> Option<OwnedHeaders> {
30        self.0
31    }
32}
33
34impl<'a, 'b, V> Extend<Header<'a, &'b V>> for KafkaHeaders
35where
36    V: ToBytes + ?Sized,
37{
38    fn extend<T: IntoIterator<Item = Header<'a, &'b V>>>(&mut self, iter: T) {
39        let mut iter = iter.into_iter();
40
41        // Probe if the iterator is empty, if it is empty, no need to do anything.
42        let Some(first) = iter.next() else {
43            return;
44        };
45
46        let mut headers = self.0.take().unwrap_or_else(|| {
47            // Get a size hint from the iterator, +2 for the already removed
48            // first element and reserving space for 1 extra header which is conditionally
49            // added by the `Producer` in this crate.
50            //
51            // This means we might allocate a little bit too much, but we never have to resize
52            // and allocate a second time, a good trade-off.
53            let size = iter.size_hint().0 + 2;
54            OwnedHeaders::new_with_capacity(size)
55        });
56        headers = headers.insert(first);
57        for remaining in iter {
58            headers = headers.insert(remaining);
59        }
60
61        self.0 = Some(headers);
62    }
63}
64
65impl<'a, 'b, V> FromIterator<Header<'a, &'b V>> for KafkaHeaders
66where
67    V: ToBytes + ?Sized,
68{
69    fn from_iter<I: IntoIterator<Item = Header<'a, &'b V>>>(iter: I) -> Self {
70        let mut c = Self::new();
71        c.extend(iter);
72        c
73    }
74}
75
76/// Kafka client and producer context that logs statistics and producer errors.
77///
78/// Cloned once per client a producer builds, so that the
79/// callbacks keep reporting across a config reload.
80#[derive(Clone)]
81pub struct Context {
82    /// Producer name for deployment identification
83    producer_name: String,
84    /// Delegate statistics to Arroyo.
85    arroyo_statistics: ArroyoProducerContext,
86}
87
88impl Context {
89    pub fn new(producer_name: String) -> Self {
90        Self {
91            arroyo_statistics: ArroyoProducerContext::new(producer_name.clone()),
92            producer_name,
93        }
94    }
95
96    pub fn producer_name(&self) -> &str {
97        &self.producer_name
98    }
99}
100
101impl ClientContext for Context {
102    /// Report client statistics as statsd metrics.
103    ///
104    /// This method is only called if `statistics.interval.ms` is configured.
105    fn stats(&self, statistics: rdkafka::Statistics) {
106        self.arroyo_statistics.stats(statistics);
107    }
108}
109
110impl ProducerContext for Context {
111    type DeliveryOpaque = ();
112
113    /// This method is called after attempting to send a message to Kafka.
114    /// It's called asynchronously for every message, so we want to handle errors explicitly here.
115    fn delivery(&self, result: &DeliveryResult, _delivery_opaque: Self::DeliveryOpaque) {
116        // TODO: any `Accepted` outcomes (e.g. spans) should be logged here instead of on the caller side,
117        // such that we do not over-report in the error case.
118
119        match result {
120            Ok(message) => {
121                metric!(
122                    counter(KafkaCounters::ProduceStatusSuccess) += 1,
123                    topic = message.topic(),
124                    producer_name = self.producer_name.as_str(),
125                );
126                metric!(
127                    counter(KafkaCounters::ProcessingMessageProduced) += 1,
128                    topic = message.topic(),
129                    producer_name = self.producer_name.as_str(),
130                );
131            }
132            Err((error, message)) => {
133                relay_log::error!(
134                    error = error as &dyn Error,
135                    payload_len = message.payload_len(),
136                    tags.topic = message.topic(),
137                    "failed to produce message to Kafka (delivery callback)",
138                );
139
140                metric!(
141                    counter(KafkaCounters::ProduceStatusError) += 1,
142                    topic = message.topic(),
143                    producer_name = self.producer_name.as_str(),
144                );
145            }
146        }
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use rdkafka::ClientContext;
153    use rdkafka::statistics::{Broker, Statistics};
154    use relay_statsd::with_capturing_test_client;
155
156    use super::Context;
157
158    #[test]
159    fn test_statistics_backend() {
160        let context = Context::new("test-producer".to_owned());
161        let statistics = Statistics {
162            msg_cnt: 42,
163            brokers: [(
164                "broker-id".to_owned(),
165                Broker {
166                    name: "broker-name".to_owned(),
167                    state: "UP".to_owned(),
168                    outbuf_cnt: 7,
169                    ..Default::default()
170                },
171            )]
172            .into(),
173            ..Default::default()
174        };
175
176        let metrics = with_capturing_test_client(|| context.stats(statistics));
177        let prefix = "arroyo.producer.librdkafka.";
178        let broker_metric = "broker_outbuf_requests";
179
180        assert!(metrics.contains(&format!(
181            "{prefix}message_count:42|g|#producer_name:test-producer"
182        )));
183        assert!(metrics.contains(&format!("{prefix}{broker_metric}:7|g|#")));
184        assert!(metrics.iter().all(|metric| metric.starts_with(prefix)));
185    }
186}
relay_kafka/producer/utils.rs

relay_kafka/producer/
utils.rs
