ruststream_rdkafka/
retry.rs1use std::collections::HashMap;
11use std::sync::{Arc, Mutex};
12use std::time::Duration;
13
14use rdkafka::Offset;
15use rdkafka::consumer::{Consumer as _, StreamConsumer};
16use rdkafka::producer::FutureRecord;
17use rdkafka::util::Timeout;
18use ruststream::Headers;
19
20use crate::broker::SharedConn;
21use crate::convert;
22use crate::error::KafkaError;
23use crate::tracker::TrackingContext;
24
25pub const RETRY_COUNT_HEADER: &str = "kafka-retry-count";
31
32pub const DLQ_SOURCE_TOPIC_HEADER: &str = "kafka-dlq-source-topic";
38pub const DLQ_SOURCE_PARTITION_HEADER: &str = "kafka-dlq-source-partition";
40pub const DLQ_SOURCE_OFFSET_HEADER: &str = "kafka-dlq-source-offset";
42
43#[derive(Debug, Clone, PartialEq, Eq)]
46#[non_exhaustive]
47pub enum Retry {
48 Topic(String),
52 SeekBack,
56 Drop,
58}
59
60pub(crate) struct RetryContext {
62 policy: Option<Retry>,
63 max_deliveries: Option<u32>,
64 dead_letter: Option<String>,
65 conn: SharedConn,
66 consumer: Arc<StreamConsumer<TrackingContext>>,
67 seeks: Mutex<HashMap<(String, i32, i64), u32>>,
71}
72
73impl std::fmt::Debug for RetryContext {
74 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75 f.debug_struct("RetryContext")
76 .field("policy", &self.policy)
77 .field("max_deliveries", &self.max_deliveries)
78 .field("dead_letter", &self.dead_letter)
79 .finish_non_exhaustive()
80 }
81}
82
83impl RetryContext {
84 pub(crate) fn new(
85 policy: Option<Retry>,
86 max_deliveries: Option<u32>,
87 dead_letter: Option<String>,
88 conn: SharedConn,
89 consumer: Arc<StreamConsumer<TrackingContext>>,
90 ) -> Self {
91 Self {
92 policy,
93 max_deliveries,
94 dead_letter,
95 conn,
96 consumer,
97 seeks: Mutex::new(HashMap::new()),
98 }
99 }
100
101 pub(crate) fn policy(&self) -> Option<&Retry> {
102 self.policy.as_ref()
103 }
104
105 pub(crate) fn dead_letter(&self) -> Option<&str> {
106 self.dead_letter.as_deref()
107 }
108
109 pub(crate) fn over_cap(&self, delivery: u32) -> bool {
111 self.max_deliveries.is_some_and(|cap| delivery > cap)
112 }
113
114 pub(crate) fn next_seek_delivery(&self, topic: &str, partition: i32, offset: i64) -> u32 {
117 let mut seeks = self.seeks.lock().expect("seek counter mutex poisoned");
118 let seeks_done = *seeks
119 .entry((topic.to_owned(), partition, offset))
120 .or_insert(0);
121 drop(seeks);
122 seeks_done + 2
123 }
124
125 pub(crate) fn record_seek(&self, topic: &str, partition: i32, offset: i64) {
126 let mut seeks = self.seeks.lock().expect("seek counter mutex poisoned");
127 *seeks
128 .entry((topic.to_owned(), partition, offset))
129 .or_insert(0) += 1;
130 }
131
132 pub(crate) fn forget_seeks(&self, topic: &str, partition: i32, offset: i64) {
133 let mut seeks = self.seeks.lock().expect("seek counter mutex poisoned");
134 seeks.remove(&(topic.to_owned(), partition, offset));
135 }
136
137 pub(crate) fn seek_back(
139 &self,
140 topic: &str,
141 partition: i32,
142 offset: i64,
143 ) -> Result<(), KafkaError> {
144 self.consumer
145 .seek(
146 topic,
147 partition,
148 Offset::Offset(offset),
149 Duration::from_secs(5),
150 )
151 .map_err(KafkaError::consume)
152 }
153
154 pub(crate) async fn republish(
157 &self,
158 topic: &str,
159 payload: &[u8],
160 headers: &Headers,
161 ) -> Result<(), KafkaError> {
162 let state = self.conn.get().ok_or(KafkaError::NotConnected)?;
163 let parts = convert::headers_for_publish(headers)?;
164 let mut record = FutureRecord::<[u8], [u8]>::to(topic).payload(payload);
165 if let Some(key) = &parts.key {
166 record = record.key(key.as_ref());
167 }
168 if let Some(partition) = parts.partition {
169 record = record.partition(partition);
170 }
171 if let Some(native) = parts.headers {
172 record = record.headers(native);
173 }
174 state
175 .producer()
176 .send(record, Timeout::Never)
177 .await
178 .map(|_delivery| ())
179 .map_err(|(err, _record)| KafkaError::publish(err))
180 }
181}