etl_core/checkpoint/
ack.rs

1//! Batch acknowledgement handles (Vector-finalizer style).
2
3use crate::checkpoint::sync::{Arc, AtomicU8, Ordering};
4use crate::record::PartitionId;
5use std::fmt;
6
7/// Identity of one source poll batch within a partition and assignment
8/// epoch. Epochs are bumped on every rebalance so acknowledgements from
9/// revoked assignments are recognized as stale and discarded.
10#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
11pub struct BatchId {
12    /// Partition the batch was polled from.
13    pub partition: PartitionId,
14    /// Assignment epoch the batch belongs to.
15    pub epoch: u32,
16    /// Per-partition, per-epoch monotonically increasing sequence number.
17    pub seq: u64,
18}
19
20/// Resolution status of a batch. Forms a "worst-of" lattice: once any
21/// record of the batch fails, the batch is failed.
22#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
23#[repr(u8)]
24pub enum AckStatus {
25    /// Every record was durably delivered, filtered, or intentionally
26    /// skipped by policy.
27    Delivered = 0,
28    /// At least one record could not be delivered; the partition watermark
29    /// must not advance past this batch.
30    Failed = 1,
31}
32
33/// Message sent to the checkpointer when the last [`AckRef`] clone of a
34/// batch drops.
35#[derive(Clone, Copy, Debug)]
36pub struct AckMsg {
37    /// Which batch resolved.
38    pub id: BatchId,
39    /// Highest source offset contained in the batch (inclusive). The
40    /// committable position after this batch is `last_offset + 1`.
41    pub last_offset: i64,
42    /// Worst status observed across the batch's records.
43    pub status: AckStatus,
44}
45
46/// Where a batch's resolution message is delivered: normally the
47/// checkpointer's unbounded channel; under `--cfg loom` an additional
48/// recorder variant lets the drop/fail races be model-checked without an
49/// unmodelled channel implementation.
50#[derive(Clone)]
51pub(crate) enum AckTx {
52    /// The checkpointer's unbounded channel.
53    Channel(crossbeam_channel::Sender<AckMsg>),
54    /// Loom-test recorder.
55    #[cfg(loom)]
56    Recorder(Arc<loom::sync::Mutex<Vec<AckMsg>>>),
57}
58
59impl AckTx {
60    fn send(&self, msg: AckMsg) {
61        match self {
62            // A send error means the checkpointer is gone (shutdown
63            // teardown); there is nothing left to acknowledge to.
64            AckTx::Channel(tx) => {
65                let _ = tx.send(msg);
66            }
67            #[cfg(loom)]
68            AckTx::Recorder(rec) => rec.lock().unwrap().push(msg),
69        }
70    }
71}
72
73impl fmt::Debug for AckTx {
74    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
75        f.write_str("AckTx")
76    }
77}
78
79/// Shared acknowledgement state of one source poll batch.
80///
81/// Dropping the final handle sends an [`AckMsg`] on the checkpointer's
82/// unbounded channel. `Arc`'s release/acquire on the refcount orders all
83/// `fail()` calls before the drop-time status read, so no additional
84/// synchronization is needed.
85#[derive(Debug)]
86struct AckState {
87    id: BatchId,
88    last_offset: i64,
89    status: AtomicU8,
90    tx: AckTx,
91}
92
93impl Drop for AckState {
94    fn drop(&mut self) {
95        let status = if self.status.load(Ordering::Relaxed) == AckStatus::Delivered as u8 {
96            AckStatus::Delivered
97        } else {
98            AckStatus::Failed
99        };
100        self.tx.send(AckMsg {
101            id: self.id,
102            last_offset: self.last_offset,
103            status,
104        });
105    }
106}
107
108/// Refcounted handle to a batch's acknowledgement state. Clone = one atomic
109/// increment, no allocation; records hold one each.
110#[derive(Clone, Debug)]
111pub struct AckRef(Arc<AckState>);
112
113impl AckRef {
114    /// Create the handle for a new batch. Called by the source driver once
115    /// per poll batch; everything downstream only clones.
116    pub(crate) fn new(id: BatchId, last_offset: i64, tx: AckTx) -> Self {
117        AckRef(Arc::new(AckState {
118            id,
119            last_offset,
120            status: AtomicU8::new(AckStatus::Delivered as u8),
121            tx,
122        }))
123    }
124
125    /// Mark the batch failed. Idempotent; any single failure poisons the
126    /// batch (worst-of lattice), stalling the partition watermark.
127    pub fn fail(&self) {
128        self.0
129            .status
130            .store(AckStatus::Failed as u8, Ordering::Relaxed);
131    }
132
133    /// Identity of the batch this handle belongs to.
134    #[must_use]
135    pub fn batch_id(&self) -> BatchId {
136        self.0.id
137    }
138
139    /// Build a handle plus the receiver its resolution message arrives on.
140    /// Test-support constructor used across the workspace's test suites.
141    #[doc(hidden)]
142    #[must_use]
143    pub fn test_pair() -> (Self, crossbeam_channel::Receiver<AckMsg>) {
144        let (tx, rx) = crossbeam_channel::unbounded();
145        (
146            Self::new(
147                BatchId {
148                    partition: PartitionId(0),
149                    epoch: 0,
150                    seq: 0,
151                },
152                0,
153                AckTx::Channel(tx),
154            ),
155            rx,
156        )
157    }
158}
159
160#[cfg(all(test, not(loom)))]
161mod tests {
162    use super::*;
163
164    fn make(tx: crossbeam_channel::Sender<AckMsg>, seq: u64, last_offset: i64) -> AckRef {
165        AckRef::new(
166            BatchId {
167                partition: PartitionId(1),
168                epoch: 7,
169                seq,
170            },
171            last_offset,
172            AckTx::Channel(tx),
173        )
174    }
175
176    #[test]
177    fn resolves_once_on_last_drop() {
178        let (tx, rx) = crossbeam_channel::unbounded();
179        let ack = make(tx, 3, 99);
180        let clones: Vec<_> = (0..8).map(|_| ack.clone()).collect();
181        drop(ack);
182        assert!(rx.try_recv().is_err(), "must not resolve while clones live");
183        drop(clones);
184        let msg = rx.try_recv().expect("resolved on last drop");
185        assert_eq!(msg.id.seq, 3);
186        assert_eq!(msg.last_offset, 99);
187        assert_eq!(msg.status, AckStatus::Delivered);
188        assert!(rx.try_recv().is_err(), "resolves exactly once");
189    }
190
191    #[test]
192    fn any_failure_poisons_the_batch() {
193        let (tx, rx) = crossbeam_channel::unbounded();
194        let ack = make(tx, 0, 10);
195        let sibling = ack.clone();
196        sibling.fail();
197        drop(sibling);
198        drop(ack);
199        assert_eq!(rx.try_recv().unwrap().status, AckStatus::Failed);
200    }
201
202    #[test]
203    fn dropped_checkpointer_does_not_panic() {
204        let (tx, rx) = crossbeam_channel::unbounded();
205        let ack = make(tx, 0, 0);
206        drop(rx);
207        drop(ack); // send fails silently — shutdown teardown ordering
208    }
209}
210
211/// A fail-safe collection of acknowledgement handles for data that has not
212/// been durably written yet (encoded chunks, sink batches).
213///
214/// A bare [`AckRef`] resolves *Delivered* on drop — the right default for
215/// records, where drop means "consumed" (filtered out, skipped by policy,
216/// or written). Collections travelling the sink path are different: there,
217/// drop without an explicit outcome means the data was torn down (task
218/// aborted, queue dropped, worker cancelled), and resolving Delivered would
219/// commit offsets for rows that never reached the sink. `AckSet` therefore
220/// **fails** every handle it still holds when dropped; the happy path calls
221/// [`AckSet::deliver`] after the durable write. Over-failing is always safe
222/// under at-least-once — it costs replay, never loss.
223#[derive(Debug, Default)]
224pub struct AckSet(Vec<AckRef>);
225
226impl AckSet {
227    /// An empty set.
228    #[must_use]
229    pub fn new() -> Self {
230        AckSet(Vec::new())
231    }
232
233    /// Add a handle to the set.
234    pub fn push(&mut self, ack: AckRef) {
235        self.0.push(ack);
236    }
237
238    /// Move every handle out of `other` into `self`.
239    pub fn absorb(&mut self, mut other: AckSet) {
240        self.0.append(&mut other.0);
241    }
242
243    /// Whether the set holds no handles.
244    #[must_use]
245    pub fn is_empty(&self) -> bool {
246        self.0.is_empty()
247    }
248
249    /// Number of handles held.
250    #[must_use]
251    pub fn len(&self) -> usize {
252        self.0.len()
253    }
254
255    /// Resolve the set as delivered: the data it covers is durably
256    /// written. Consumes the set; the handles drop with their default
257    /// (Delivered) resolution.
258    pub fn deliver(mut self) {
259        self.0.clear();
260    }
261}
262
263impl Drop for AckSet {
264    fn drop(&mut self) {
265        for ack in &self.0 {
266            ack.fail();
267        }
268    }
269}
270
271impl From<Vec<AckRef>> for AckSet {
272    fn from(acks: Vec<AckRef>) -> Self {
273        AckSet(acks)
274    }
275}
276
277#[cfg(all(test, not(loom)))]
278mod ack_set_tests {
279    use super::*;
280
281    #[test]
282    fn dropping_a_set_fails_its_batches() {
283        let (ack, rx) = AckRef::test_pair();
284        let mut set = AckSet::new();
285        set.push(ack.clone());
286        drop(ack);
287        drop(set);
288        assert_eq!(rx.try_recv().unwrap().status, AckStatus::Failed);
289    }
290
291    #[test]
292    fn delivering_a_set_resolves_delivered() {
293        let (ack, rx) = AckRef::test_pair();
294        let mut set = AckSet::new();
295        set.push(ack.clone());
296        drop(ack);
297        set.deliver();
298        assert_eq!(rx.try_recv().unwrap().status, AckStatus::Delivered);
299    }
300
301    #[test]
302    fn absorb_moves_handles_without_resolving() {
303        let (ack, rx) = AckRef::test_pair();
304        let mut a = AckSet::new();
305        a.push(ack.clone());
306        drop(ack);
307        let mut b = AckSet::new();
308        b.absorb(a);
309        assert!(rx.try_recv().is_err(), "absorb must not resolve");
310        assert_eq!(b.len(), 1);
311        b.deliver();
312        assert_eq!(rx.try_recv().unwrap().status, AckStatus::Delivered);
313    }
314}
etl_core/checkpoint/ack.rs

etl_core/checkpoint/
ack.rs
