objectstore_inventory_tracker/
kafka.rs1use 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#[derive(Clone, Debug, serde::Deserialize, serde::Serialize)]
17#[serde(default)]
18pub struct KafkaConfig {
19 pub topic: String,
21
22 pub bootstrap_servers: Vec<String>,
24
25 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#[derive(Debug, thiserror::Error)]
43pub enum KafkaProducerError {
44 #[error("failed to create kafka producer")]
46 InvalidConfig(#[source] KafkaError),
47
48 #[error("failed to enqueue inventory record")]
53 SendFailed(#[source] KafkaError),
54
55 #[error("failed to flush inventory records")]
57 FlushFailed(#[source] KafkaError),
58}
59
60pub type OnDeliveryFailure = Box<dyn Fn(&KafkaError) + Send + Sync>;
64
65struct 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 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
93pub 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 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 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 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 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 let _ = producer.join_blocking(Duration::from_secs(5));
239
240 assert_eq!(failures.load(Ordering::SeqCst), 1);
241 }
242}