Skip to main content

ruststream_rdkafka/
retry.rs

1//! Per-subscription retry and dead-letter policies for negative acknowledgement.
2//!
3//! Kafka has no per-message negative acknowledgement, so `nack(true)` natively means "leave
4//! the offset unsettled and re-consume later". These policies give it an immediate meaning:
5//! republish to a retry topic (with an attempt counter riding in a header), seek the partition
6//! back to re-consume in place, or drop - with an optional dead-letter topic on the drop path
7//! and a poison cap on the number of deliveries. Everything here is off by default; without a
8//! policy, settlement behaves exactly as the commit mode describes.
9
10use 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
25/// Header carrying the number of retry republishes a message has been through.
26///
27/// Set by [`Retry::Topic`] on every republish (an ASCII decimal), read back to enforce
28/// [`KafkaTopic::max_deliveries`](crate::KafkaTopic::max_deliveries) across hops. A message on
29/// its retry topic with the value `2` is on its third delivery.
30pub const RETRY_COUNT_HEADER: &str = "kafka-retry-count";
31
32/// Headers stamped onto a dead-lettered message with the origin of the failed delivery.
33///
34/// `kafka-dlq-source-topic`, `kafka-dlq-source-partition`, and `kafka-dlq-source-offset` name
35/// the topic, partition, and offset the message failed on, so a dead-letter consumer can trace
36/// it back without parsing payloads.
37pub const DLQ_SOURCE_TOPIC_HEADER: &str = "kafka-dlq-source-topic";
38/// See [`DLQ_SOURCE_TOPIC_HEADER`].
39pub const DLQ_SOURCE_PARTITION_HEADER: &str = "kafka-dlq-source-partition";
40/// See [`DLQ_SOURCE_TOPIC_HEADER`].
41pub const DLQ_SOURCE_OFFSET_HEADER: &str = "kafka-dlq-source-offset";
42
43/// What `nack(true)` does on this subscription (see
44/// [`KafkaTopic::retry`](crate::KafkaTopic::retry)).
45#[derive(Debug, Clone, PartialEq, Eq)]
46#[non_exhaustive]
47pub enum Retry {
48    /// Republish the message to this topic (stamping [`RETRY_COUNT_HEADER`]), then settle the
49    /// original offset. Republish-first ordering makes a crash between the two steps a
50    /// duplicate, never a loss.
51    Topic(String),
52    /// Seek the partition back to the message's offset and re-consume it in place. Redelivery
53    /// is immediate, but everything after the offset on that partition replays too
54    /// (at-least-once duplicates), and the delivery count only survives within the session.
55    SeekBack,
56    /// Treat `nack(true)` like the drop path (dead-letter when configured, settle otherwise).
57    Drop,
58}
59
60/// The resolved retry/dead-letter wiring one subscription hands to its deliveries.
61pub(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    /// In-session delivery counts for `SeekBack`, keyed by the seeked offset. Entries are
68    /// removed when the offset resolves to the drop path; a poison offset therefore holds at
69    /// most one entry at a time.
70    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    /// Whether delivery number `delivery` (1-based) exceeds the poison cap.
110    pub(crate) fn over_cap(&self, delivery: u32) -> bool {
111        self.max_deliveries.is_some_and(|cap| delivery > cap)
112    }
113
114    /// Counts a seek-back redelivery of `offset` and returns the delivery number it will be
115    /// (the original delivery is number one).
116    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    /// Seeks the partition back so `offset` (and everything after it) re-consumes.
138    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    /// Republishes `payload` with `headers` to `topic` and awaits the delivery report, so the
155    /// caller settles the original only after the copy is durably accepted.
156    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}