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::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
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    state: Arc<ConnState>,
66    consumer: Arc<StreamConsumer<TrackingContext>>,
67    tracker: Arc<CommitTracker>,
68    /// In-session delivery counts for `SeekBack`, keyed by the seeked offset. Entries are
69    /// removed when the offset resolves to the drop path; a poison offset therefore holds at
70    /// most one entry at a time.
71    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    /// Whether delivery number `delivery` (1-based) exceeds the poison cap.
113    pub(crate) fn over_cap(&self, delivery: u32) -> bool {
114        self.max_deliveries.is_some_and(|cap| delivery > cap)
115    }
116
117    /// Counts a seek-back redelivery of `offset` and returns the delivery number it will be
118    /// (the original delivery is number one).
119    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    /// Seeks the partition back so `offset` (and everything after it) re-consumes.
141    pub(crate) fn seek_back(
142        &self,
143        topic: &str,
144        partition: i32,
145        offset: i64,
146    ) -> Result<(), KafkaError> {
147        // The same resets the explicit seeker performs, for the same reason: deliveries already
148        // in flight from beyond this offset belong to the read position being replaced, so
149        // neither their settles nor librdkafka's own stored position may commit past the
150        // records this rewind replays.
151        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    /// Republishes `payload` with `headers` to `topic` and awaits the delivery report, so the
168    /// caller settles the original only after the copy is durably accepted.
169    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}