datum_mq/
offset.rs

1use std::{collections::BTreeMap, fmt, sync::Arc};
2
3use crate::{MqError, MqResult};
4
5/// Kafka topic/partition identifier.
6#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
7pub struct TopicPartition {
8    pub topic: String,
9    pub partition: i32,
10}
11
12impl TopicPartition {
13    #[must_use]
14    pub fn new(topic: impl Into<String>, partition: i32) -> Self {
15        Self {
16            topic: topic.into(),
17            partition,
18        }
19    }
20}
21
22/// A Kafka commit position. `offset` is the next offset to consume.
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub struct OffsetCommit {
25    pub topic: String,
26    pub partition: i32,
27    pub offset: i64,
28}
29
30impl OffsetCommit {
31    #[must_use]
32    pub fn topic_partition(&self) -> TopicPartition {
33        TopicPartition::new(self.topic.clone(), self.partition)
34    }
35}
36
37#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
38pub(crate) struct OffsetPosition {
39    pub(crate) topic: String,
40    pub(crate) partition: i32,
41    pub(crate) offset: i64,
42}
43
44impl OffsetPosition {
45    pub(crate) fn new(topic: impl Into<String>, partition: i32, offset: i64) -> Self {
46        Self {
47            topic: topic.into(),
48            partition,
49            offset,
50        }
51    }
52
53    pub(crate) fn topic_partition(&self) -> TopicPartition {
54        TopicPartition::new(self.topic.clone(), self.partition)
55    }
56}
57
58pub(crate) trait OffsetCommitter: Send + Sync {
59    fn stage_offset(&self, position: OffsetPosition) -> MqResult<()>;
60    fn commit_offset(&self, position: OffsetPosition) -> MqResult<()>;
61}
62
63/// Context value emitted by `KafkaSource::committable`.
64#[derive(Clone)]
65pub struct KafkaOffset {
66    position: OffsetPosition,
67    committer: Option<Arc<dyn OffsetCommitter>>,
68}
69
70impl KafkaOffset {
71    pub(crate) fn new(position: OffsetPosition, committer: Arc<dyn OffsetCommitter>) -> Self {
72        Self {
73            position,
74            committer: Some(committer),
75        }
76    }
77
78    #[must_use]
79    pub fn detached(topic: impl Into<String>, partition: i32, offset: i64) -> Self {
80        Self {
81            position: OffsetPosition::new(topic, partition, offset),
82            committer: None,
83        }
84    }
85
86    #[must_use]
87    pub fn topic(&self) -> &str {
88        &self.position.topic
89    }
90
91    #[must_use]
92    pub fn partition(&self) -> i32 {
93        self.position.partition
94    }
95
96    #[must_use]
97    pub fn offset(&self) -> i64 {
98        self.position.offset
99    }
100
101    #[must_use]
102    pub fn next_offset(&self) -> i64 {
103        self.position.offset + 1
104    }
105
106    #[must_use]
107    pub fn topic_partition(&self) -> TopicPartition {
108        self.position.topic_partition()
109    }
110
111    /// Marks this offset as completed downstream and commits the highest contiguous
112    /// completed position for the partition.
113    pub fn commit(&self) -> MqResult<()> {
114        let Some(committer) = &self.committer else {
115            return Err(MqError::ControlUnavailable);
116        };
117        committer.commit_offset(self.position.clone())
118    }
119
120    /// Marks this offset as completed downstream without sending a Kafka commit.
121    ///
122    /// Use this for batched commits: stage earlier offsets in a contiguous batch,
123    /// then call [`KafkaOffset::commit`] on the last offset after the downstream
124    /// checkpoint succeeds.
125    pub fn mark_processed(&self) -> MqResult<()> {
126        let Some(committer) = &self.committer else {
127            return Err(MqError::ControlUnavailable);
128        };
129        committer.stage_offset(self.position.clone())
130    }
131}
132
133impl fmt::Debug for KafkaOffset {
134    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
135        f.debug_struct("KafkaOffset")
136            .field("topic", &self.position.topic)
137            .field("partition", &self.position.partition)
138            .field("offset", &self.position.offset)
139            .finish_non_exhaustive()
140    }
141}
142
143/// Batch of offsets for external checkpointing or batched downstream commit sinks.
144#[derive(Debug, Clone, Default)]
145pub struct KafkaOffsetBatch {
146    offsets: Vec<KafkaOffset>,
147}
148
149impl KafkaOffsetBatch {
150    #[must_use]
151    pub fn new() -> Self {
152        Self::default()
153    }
154
155    #[must_use]
156    pub fn from_offsets<I>(offsets: I) -> Self
157    where
158        I: IntoIterator<Item = KafkaOffset>,
159    {
160        Self {
161            offsets: offsets.into_iter().collect(),
162        }
163    }
164
165    pub fn push(&mut self, offset: KafkaOffset) {
166        self.offsets.push(offset);
167    }
168
169    #[must_use]
170    pub fn len(&self) -> usize {
171        self.offsets.len()
172    }
173
174    #[must_use]
175    pub fn is_empty(&self) -> bool {
176        self.offsets.is_empty()
177    }
178
179    pub fn commit(&self) -> MqResult<()> {
180        let mut by_partition = BTreeMap::<TopicPartition, &KafkaOffset>::new();
181        for offset in &self.offsets {
182            by_partition
183                .entry(offset.topic_partition())
184                .and_modify(|current| {
185                    if offset.offset() > current.offset() {
186                        *current = offset;
187                    }
188                })
189                .or_insert(offset);
190        }
191        for offset in by_partition.values() {
192            offset.commit()?;
193        }
194        Ok(())
195    }
196}
197
198/// Committable watermark for a batch of Kafka records.
199///
200/// `KafkaBatchOffset` is emitted by batch-oriented sources. It advances one
201/// offset watermark per topic/partition after the downstream batch checkpoint
202/// succeeds, avoiding one offset handle per record while preserving
203/// at-least-once commit ordering.
204#[derive(Clone)]
205pub struct KafkaBatchOffset {
206    positions: Vec<OffsetPosition>,
207    committer: KafkaBatchCommitter,
208}
209
210#[derive(Clone)]
211enum KafkaBatchCommitter {
212    Offset(Option<Arc<dyn OffsetCommitter>>),
213    Custom(Arc<dyn Fn() -> MqResult<()> + Send + Sync>),
214}
215
216impl KafkaBatchOffset {
217    pub(crate) fn custom<F>(positions: Vec<OffsetPosition>, commit: F) -> Self
218    where
219        F: Fn() -> MqResult<()> + Send + Sync + 'static,
220    {
221        Self {
222            positions,
223            committer: KafkaBatchCommitter::Custom(Arc::new(commit)),
224        }
225    }
226
227    #[must_use]
228    pub fn detached<I>(positions: I) -> Self
229    where
230        I: IntoIterator<Item = OffsetCommit>,
231    {
232        Self {
233            positions: positions
234                .into_iter()
235                .map(|position| {
236                    OffsetPosition::new(position.topic, position.partition, position.offset - 1)
237                })
238                .collect(),
239            committer: KafkaBatchCommitter::Offset(None),
240        }
241    }
242
243    #[must_use]
244    pub fn len(&self) -> usize {
245        self.positions.len()
246    }
247
248    #[must_use]
249    pub fn is_empty(&self) -> bool {
250        self.positions.is_empty()
251    }
252
253    /// Marks the batch as completed downstream and commits eligible Kafka
254    /// watermarks according to the source committer settings.
255    pub fn commit(&self) -> MqResult<()> {
256        match &self.committer {
257            KafkaBatchCommitter::Offset(Some(committer)) => {
258                for position in &self.positions {
259                    committer.commit_offset(position.clone())?;
260                }
261                Ok(())
262            }
263            KafkaBatchCommitter::Offset(None) => Err(MqError::ControlUnavailable),
264            KafkaBatchCommitter::Custom(commit) => commit(),
265        }
266    }
267}
268
269impl fmt::Debug for KafkaBatchOffset {
270    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
271        f.debug_struct("KafkaBatchOffset")
272            .field("positions", &self.positions)
273            .finish_non_exhaustive()
274    }
275}
276
277#[cfg(test)]
278mod tests {
279    use super::*;
280    use std::sync::Mutex;
281
282    fn pos(offset: i64) -> OffsetPosition {
283        OffsetPosition::new("topic", 0, offset)
284    }
285
286    #[derive(Debug, Default)]
287    struct FakeCommitter {
288        staged: Mutex<Vec<OffsetPosition>>,
289        committed: Mutex<Vec<OffsetPosition>>,
290    }
291
292    impl OffsetCommitter for FakeCommitter {
293        fn stage_offset(&self, position: OffsetPosition) -> MqResult<()> {
294            self.staged.lock().expect("staged lock").push(position);
295            Ok(())
296        }
297
298        fn commit_offset(&self, position: OffsetPosition) -> MqResult<()> {
299            self.committed
300                .lock()
301                .expect("committed lock")
302                .push(position);
303            Ok(())
304        }
305    }
306
307    #[test]
308    fn offset_batch_commits_last_offset_per_partition() {
309        let committer = Arc::new(FakeCommitter::default());
310        let offsets = [
311            KafkaOffset::new(pos(0), committer.clone()),
312            KafkaOffset::new(OffsetPosition::new("topic", 1, 0), committer.clone()),
313            KafkaOffset::new(pos(1), committer.clone()),
314            KafkaOffset::new(OffsetPosition::new("topic", 1, 1), committer.clone()),
315        ];
316
317        KafkaOffsetBatch::from_offsets(offsets)
318            .commit()
319            .expect("batch commits");
320
321        let staged = committer.staged.lock().expect("staged lock");
322        let committed = committer.committed.lock().expect("committed lock");
323        assert_eq!(staged.len(), 0);
324        assert_eq!(committed.len(), 2);
325        assert!(committed.iter().all(|position| position.offset == 1));
326    }
327}
datum_mq/offset.rs

datum_mq/
offset.rs
