objectstore_inventory_tracker/
kafka.rs1use 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#[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] ProducerError),
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: 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 pub fn try_new(
119 config: KafkaConfig,
120 on_delivery_failure: Option<OnDeliveryFailure>,
121 ) -> Result<Self, KafkaProducerError> {
122 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 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 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 let _ = producer.join_blocking(Duration::from_secs(5));
238
239 assert_eq!(failures.load(Ordering::SeqCst), 1);
240 }
241}