ruststream_rdkafka/
retry.rs1use std::collections::HashMap;
11use std::sync::{Arc, Mutex};
12use std::time::Duration;
13
14use rdkafka::consumer::{Consumer as _, StreamConsumer};
15use rdkafka::producer::FutureRecord;
16use rdkafka::util::Timeout;
17use rdkafka::{Offset, TopicPartitionList};
18use ruststream::Headers;
19
20use crate::broker::ConnState;
21use crate::convert;
22use crate::error::KafkaError;
23use crate::tracker::{CommitTracker, 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 state: Arc<ConnState>,
66 consumer: Arc<StreamConsumer<TrackingContext>>,
67 tracker: Arc<CommitTracker>,
68 seeks: Mutex<HashMap<(String, i32, i64), u32>>,
72}
73
74impl std::fmt::Debug for RetryContext {
75 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
76 f.debug_struct("RetryContext")
77 .field("policy", &self.policy)
78 .field("max_deliveries", &self.max_deliveries)
79 .field("dead_letter", &self.dead_letter)
80 .finish_non_exhaustive()
81 }
82}
83
84impl RetryContext {
85 pub(crate) fn new(
86 policy: Option<Retry>,
87 max_deliveries: Option<u32>,
88 dead_letter: Option<String>,
89 state: Arc<ConnState>,
90 consumer: Arc<StreamConsumer<TrackingContext>>,
91 tracker: Arc<CommitTracker>,
92 ) -> Self {
93 Self {
94 policy,
95 max_deliveries,
96 dead_letter,
97 state,
98 consumer,
99 tracker,
100 seeks: Mutex::new(HashMap::new()),
101 }
102 }
103
104 pub(crate) fn policy(&self) -> Option<&Retry> {
105 self.policy.as_ref()
106 }
107
108 pub(crate) fn dead_letter(&self) -> Option<&str> {
109 self.dead_letter.as_deref()
110 }
111
112 pub(crate) fn over_cap(&self, delivery: u32) -> bool {
114 self.max_deliveries.is_some_and(|cap| delivery > cap)
115 }
116
117 pub(crate) fn next_seek_delivery(&self, topic: &str, partition: i32, offset: i64) -> u32 {
120 let mut seeks = self.seeks.lock().expect("seek counter mutex poisoned");
121 let seeks_done = *seeks
122 .entry((topic.to_owned(), partition, offset))
123 .or_insert(0);
124 drop(seeks);
125 seeks_done + 2
126 }
127
128 pub(crate) fn record_seek(&self, topic: &str, partition: i32, offset: i64) {
129 let mut seeks = self.seeks.lock().expect("seek counter mutex poisoned");
130 *seeks
131 .entry((topic.to_owned(), partition, offset))
132 .or_insert(0) += 1;
133 }
134
135 pub(crate) fn forget_seeks(&self, topic: &str, partition: i32, offset: i64) {
136 let mut seeks = self.seeks.lock().expect("seek counter mutex poisoned");
137 seeks.remove(&(topic.to_owned(), partition, offset));
138 }
139
140 pub(crate) fn seek_back(
142 &self,
143 topic: &str,
144 partition: i32,
145 offset: i64,
146 ) -> Result<(), KafkaError> {
147 self.tracker.reposition(topic, partition);
152 let mut rewound = TopicPartitionList::new();
153 rewound
154 .add_partition_offset(topic, partition, Offset::Offset(offset))
155 .map_err(KafkaError::consume)?;
156 crate::seek::clear_stored_offsets(&self.consumer, &rewound)?;
157 self.consumer
158 .seek(
159 topic,
160 partition,
161 Offset::Offset(offset),
162 Duration::from_secs(5),
163 )
164 .map_err(KafkaError::consume)
165 }
166
167 pub(crate) async fn republish(
170 &self,
171 topic: &str,
172 payload: &[u8],
173 headers: &Headers,
174 ) -> Result<(), KafkaError> {
175 self.state.ensure_open(topic)?;
176 let parts = convert::headers_for_publish(headers)?;
177 let mut record = FutureRecord::<[u8], [u8]>::to(topic).payload(payload);
178 if let Some(key) = &parts.key {
179 record = record.key(key.as_ref());
180 }
181 if let Some(partition) = parts.partition {
182 record = record.partition(partition);
183 }
184 if let Some(native) = parts.headers {
185 record = record.headers(native);
186 }
187 self.state
188 .producer()
189 .send(record, Timeout::Never)
190 .await
191 .map(|_delivery| ())
192 .map_err(|(err, _record)| KafkaError::publish(err))
193 }
194}