Skip to main content

objectstore_inventory_tracker/
kafka.rs

1//! A [`Producer`] backed by [`rdkafka`].
2
3use std::collections::HashMap;
4use std::time::Duration;
5
6use rdkafka::ClientConfig;
7use rdkafka::client::ClientContext;
8use rdkafka::error::KafkaError;
9use rdkafka::producer::{
10    BaseRecord, DeliveryResult, Producer as RdKafkaProducer, ProducerContext, ThreadedProducer,
11};
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] KafkaError),
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: ThreadedProducer<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        let mut client_config = ClientConfig::new();
123        client_config.set("bootstrap.servers", config.bootstrap_servers.join(","));
124
125        // Applied after the broker list so that a caller can override it if they need to.
126        for (key, value) in &config.override_params {
127            client_config.set(key, value);
128        }
129
130        let producer = client_config
131            .create_with_context(DeliveryReporter {
132                on_failure: on_delivery_failure,
133            })
134            .map_err(KafkaProducerError::InvalidConfig)?;
135
136        Ok(Self {
137            topic: config.topic,
138            producer,
139        })
140    }
141}
142
143impl Producer for KafkaProducer {
144    type Error = KafkaProducerError;
145
146    fn send(&self, key: &[u8], payload: Vec<u8>) -> Result<(), Self::Error> {
147        let record: BaseRecord<'_, [u8], [u8]> =
148            BaseRecord::to(&self.topic).key(key).payload(&payload);
149
150        self.producer
151            .send(record)
152            .map_err(|(error, _)| KafkaProducerError::SendFailed(error))
153    }
154
155    fn join_blocking(&self, timeout: Duration) -> Result<(), Self::Error> {
156        self.producer
157            .flush(timeout)
158            .map_err(KafkaProducerError::FlushFailed)
159    }
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165
166    fn config() -> KafkaConfig {
167        KafkaConfig {
168            topic: "shared-resources-inventory".into(),
169            bootstrap_servers: vec!["127.0.0.1:9092".into()],
170            override_params: HashMap::from([("compression.type".into(), "lz4".into())]),
171        }
172    }
173
174    #[test]
175    fn producer_targets_the_configured_topic() {
176        let producer = KafkaProducer::try_new(config(), None).unwrap();
177        assert_eq!(producer.topic, "shared-resources-inventory");
178    }
179
180    #[test]
181    fn every_bootstrap_server_is_passed_through() {
182        let mut config = config();
183        config.bootstrap_servers = vec!["a:9092".into(), "b:9092".into(), "c:9092".into()];
184
185        // librdkafka parses the broker list at creation, so building successfully is the
186        // assertion that the joined value was well formed.
187        assert!(KafkaProducer::try_new(config, None).is_ok());
188    }
189
190    #[test]
191    fn bad_config_is_an_error() {
192        let mut config = config();
193        config
194            .override_params
195            .insert("not.a.real.property".into(), "1".into());
196
197        assert!(matches!(
198            KafkaProducer::try_new(config, None),
199            Err(KafkaProducerError::InvalidConfig(_))
200        ));
201    }
202
203    #[test]
204    fn joining_an_empty_queue_returns_immediately() {
205        let producer = KafkaProducer::try_new(config(), None).unwrap();
206
207        let start = std::time::Instant::now();
208        producer.join_blocking(Duration::from_secs(5)).unwrap();
209        assert!(start.elapsed() < Duration::from_secs(1));
210    }
211
212    #[test]
213    fn delivery_failures_reach_the_callback() {
214        use std::sync::Arc;
215        use std::sync::atomic::{AtomicUsize, Ordering};
216
217        let failures = Arc::new(AtomicUsize::new(0));
218        let counter = Arc::clone(&failures);
219
220        let mut config = config();
221        // Nothing is listening on this port, and a short timeout means librdkafka gives
222        // up and reports the record as undeliverable while the test is still running.
223        config.bootstrap_servers = vec!["127.0.0.1:1".into()];
224        config
225            .override_params
226            .insert("message.timeout.ms".into(), "300".into());
227
228        let producer = KafkaProducer::try_new(
229            config,
230            Some(Box::new(move |_| {
231                counter.fetch_add(1, Ordering::SeqCst);
232            })),
233        )
234        .unwrap();
235
236        producer.send(b"key", b"payload".to_vec()).unwrap();
237        // Draining drives the delivery callbacks; the record cannot succeed.
238        let _ = producer.join_blocking(Duration::from_secs(5));
239
240        assert_eq!(failures.load(Ordering::SeqCst), 1);
241    }
242}