Skip to main content

objectstore_inventory_tracker/
kafka.rs

1//! A [`Producer`] backed by [`sentry_arroyo`]'s Kafka producer.
2
3use std::collections::HashMap;
4use std::time::Duration;
5
6use rdkafka::client::ClientContext;
7use rdkafka::error::KafkaError;
8use rdkafka::producer::{BaseRecord, DeliveryResult, ProducerContext};
9use sentry_arroyo::backends::ProducerError;
10use sentry_arroyo::backends::kafka::config::KafkaConfig as ArroyoKafkaConfig;
11use sentry_arroyo::backends::kafka::producer::KafkaProducer as ArroyoKafkaProducer;
12
13use crate::producer::Producer;
14
15/// Connection settings for [`KafkaProducer`].
16#[derive(Clone, Debug, serde::Deserialize, serde::Serialize)]
17#[serde(default)]
18pub struct KafkaConfig {
19    /// Topic to produce to.
20    pub topic: String,
21
22    /// Broker addresses.
23    pub bootstrap_servers: Vec<String>,
24
25    /// Additional librdkafka properties, passed through verbatim.
26    ///
27    /// SASL credentials, compression, and buffering limits go here.
28    pub override_params: HashMap<String, String>,
29}
30
31impl Default for KafkaConfig {
32    fn default() -> Self {
33        Self {
34            topic: "shared-resources-inventory".to_owned(),
35            bootstrap_servers: Vec::new(),
36            override_params: HashMap::new(),
37        }
38    }
39}
40
41/// Something that went wrong producing to Kafka.
42#[derive(Debug, thiserror::Error)]
43pub enum KafkaProducerError {
44    /// The producer could not be created from the given configuration.
45    #[error("failed to create kafka producer")]
46    InvalidConfig(#[source] KafkaError),
47
48    /// The record could not be enqueued.
49    ///
50    /// This may mean the local queue is full, which may be caused by the broker being
51    /// unreachable or backed up.
52    #[error("failed to enqueue inventory record")]
53    SendFailed(#[source] ProducerError),
54
55    /// The local queue was not emptied before the flush timeout elapsed.
56    #[error("failed to flush inventory records")]
57    FlushFailed(#[source] KafkaError),
58}
59
60/// Called for each record that fails to deliver.
61///
62/// Invoked from a librdkafka background thread, so it must not block.
63pub type OnDeliveryFailure = Box<dyn Fn(&KafkaError) + Send + Sync>;
64
65/// Reports delivery outcomes.
66///
67/// Delivery is asynchronous, so a successful [`Producer::send`] only means the record was
68/// enqueued locally. If the broker later rejects messages for some reason, this type
69/// provides a callback.
70struct DeliveryReporter {
71    on_failure: Option<OnDeliveryFailure>,
72}
73
74impl ClientContext for DeliveryReporter {}
75
76impl ProducerContext for DeliveryReporter {
77    type DeliveryOpaque = ();
78
79    fn delivery(&self, result: &DeliveryResult<'_>, _opaque: Self::DeliveryOpaque) {
80        if let Err((error, _)) = result {
81            // `&dyn Error` rather than `%error` so the source chain is captured.
82            tracing::warn!(
83                error = error as &dyn std::error::Error,
84                "failed to deliver inventory record"
85            );
86            if let Some(on_failure) = &self.on_failure {
87                on_failure(error);
88            }
89        }
90    }
91}
92
93/// Produces inventory records onto a Kafka topic.
94///
95/// Sends are non-blocking: records go onto librdkafka's internal queue and a background
96/// thread delivers them. When that queue is full, [`Producer::send`] returns an error
97/// rather than waiting.
98pub struct KafkaProducer {
99    topic: String,
100    producer: ArroyoKafkaProducer<DeliveryReporter>,
101}
102
103impl std::fmt::Debug for KafkaProducer {
104    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105        f.debug_struct("KafkaProducer")
106            .field("topic", &self.topic)
107            .field("in_flight", &self.producer.in_flight_count())
108            .finish_non_exhaustive()
109    }
110}
111
112impl KafkaProducer {
113    /// Creates a producer from `config`, optionally reporting delivery failures to
114    /// `on_delivery_failure`.
115    ///
116    /// Failures are always logged. The callback is how a caller additionally counts them
117    /// without this crate having to depend on a metrics backend.
118    pub fn try_new(
119        config: KafkaConfig,
120        on_delivery_failure: Option<OnDeliveryFailure>,
121    ) -> Result<Self, KafkaProducerError> {
122        // Overrides are applied after the broker list, so a caller can override it too.
123        // Not `new_producer_config`: it enables per-second stats, which rdkafka logs at info.
124        let client_config =
125            ArroyoKafkaConfig::new_config(config.bootstrap_servers, Some(config.override_params));
126
127        let producer = ArroyoKafkaProducer::new_with_context(
128            client_config,
129            DeliveryReporter {
130                on_failure: on_delivery_failure,
131            },
132        )
133        .map_err(KafkaProducerError::InvalidConfig)?;
134
135        Ok(Self {
136            topic: config.topic,
137            producer,
138        })
139    }
140}
141
142impl Producer for KafkaProducer {
143    type Error = KafkaProducerError;
144
145    fn send(&self, key: &[u8], payload: Vec<u8>) -> Result<(), Self::Error> {
146        let record: BaseRecord<'_, [u8], [u8]> =
147            BaseRecord::to(&self.topic).key(key).payload(&payload);
148
149        self.producer
150            .produce_record(record)
151            .map_err(KafkaProducerError::SendFailed)
152    }
153
154    fn join_blocking(&self, timeout: Duration) -> Result<(), Self::Error> {
155        self.producer
156            .flush_for(timeout)
157            .map_err(KafkaProducerError::FlushFailed)
158    }
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164
165    fn config() -> KafkaConfig {
166        KafkaConfig {
167            topic: "shared-resources-inventory".into(),
168            bootstrap_servers: vec!["127.0.0.1:9092".into()],
169            override_params: HashMap::from([("compression.type".into(), "lz4".into())]),
170        }
171    }
172
173    #[test]
174    fn producer_targets_the_configured_topic() {
175        let producer = KafkaProducer::try_new(config(), None).unwrap();
176        assert_eq!(producer.topic, "shared-resources-inventory");
177    }
178
179    #[test]
180    fn every_bootstrap_server_is_passed_through() {
181        let mut config = config();
182        config.bootstrap_servers = vec!["a:9092".into(), "b:9092".into(), "c:9092".into()];
183
184        // librdkafka parses the broker list at creation, so building successfully is the
185        // assertion that the joined value was well formed.
186        assert!(KafkaProducer::try_new(config, None).is_ok());
187    }
188
189    #[test]
190    fn bad_config_is_an_error() {
191        let mut config = config();
192        config
193            .override_params
194            .insert("not.a.real.property".into(), "1".into());
195
196        assert!(matches!(
197            KafkaProducer::try_new(config, None),
198            Err(KafkaProducerError::InvalidConfig(_))
199        ));
200    }
201
202    #[test]
203    fn joining_an_empty_queue_returns_immediately() {
204        let producer = KafkaProducer::try_new(config(), None).unwrap();
205
206        let start = std::time::Instant::now();
207        producer.join_blocking(Duration::from_secs(5)).unwrap();
208        assert!(start.elapsed() < Duration::from_secs(1));
209    }
210
211    #[test]
212    fn delivery_failures_reach_the_callback() {
213        use std::sync::Arc;
214        use std::sync::atomic::{AtomicUsize, Ordering};
215
216        let failures = Arc::new(AtomicUsize::new(0));
217        let counter = Arc::clone(&failures);
218
219        let mut config = config();
220        // Nothing is listening on this port, and a short timeout means librdkafka gives
221        // up and reports the record as undeliverable while the test is still running.
222        config.bootstrap_servers = vec!["127.0.0.1:1".into()];
223        config
224            .override_params
225            .insert("message.timeout.ms".into(), "300".into());
226
227        let producer = KafkaProducer::try_new(
228            config,
229            Some(Box::new(move |_| {
230                counter.fetch_add(1, Ordering::SeqCst);
231            })),
232        )
233        .unwrap();
234
235        producer.send(b"key", b"payload".to_vec()).unwrap();
236        // Draining drives the delivery callbacks; the record cannot succeed.
237        let _ = producer.join_blocking(Duration::from_secs(5));
238
239        assert_eq!(failures.load(Ordering::SeqCst), 1);
240    }
241}