laminar_core/checkpoint/
barrier.rs

1//! Checkpoint barrier protocol.
2//!
3//! The coordinator injects barriers into sources via [`CheckpointBarrierInjector`].
4//! Sources deliver barriers alongside events. The fast path (no pending barrier)
5//! is a single `AtomicU64` load (~10ns).
6
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::sync::Arc;
9
10/// Barrier flags — packed into the `flags` field.
11pub mod flags {
12    /// No special behavior.
13    pub const NONE: u64 = 0;
14    /// This barrier requires a full snapshot (not incremental).
15    pub const FULL_SNAPSHOT: u64 = 1 << 0;
16    /// This barrier is the final barrier before shutdown.
17    pub const DRAIN: u64 = 1 << 1;
18    /// Cancel any in-progress checkpoint with this ID.
19    pub const CANCEL: u64 = 1 << 2;
20}
21
22/// A checkpoint barrier that flows through the dataflow graph.
23///
24/// This is a 24-byte `#[repr(C)]` value type that can be cheaply copied
25/// and embedded in channel messages. It carries the checkpoint identity
26/// and behavior flags.
27///
28/// ## Layout (24 bytes)
29///
30/// | Field          | Offset | Size | Description                |
31/// |----------------|--------|------|----------------------------|
32/// | checkpoint_id  | 0      | 8    | Unique checkpoint ID       |
33/// | epoch          | 8      | 8    | Monotonic epoch number     |
34/// | flags          | 16     | 8    | Behavior flags (see [`flags`]) |
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
36#[repr(C)]
37pub struct CheckpointBarrier {
38    /// Unique identifier for this checkpoint.
39    pub checkpoint_id: u64,
40    /// Monotonically increasing epoch number.
41    pub epoch: u64,
42    /// Behavior flags (see [`flags`] module constants).
43    pub flags: u64,
44}
45
46// Verify the struct is exactly 24 bytes as promised.
47const _: () = assert!(std::mem::size_of::<CheckpointBarrier>() == 24);
48
49impl CheckpointBarrier {
50    /// Create a new barrier with the given checkpoint ID and epoch.
51    #[must_use]
52    pub const fn new(checkpoint_id: u64, epoch: u64) -> Self {
53        Self {
54            checkpoint_id,
55            epoch,
56            flags: flags::NONE,
57        }
58    }
59
60    /// Create a barrier that requests a full snapshot.
61    #[must_use]
62    pub const fn full_snapshot(checkpoint_id: u64, epoch: u64) -> Self {
63        Self {
64            checkpoint_id,
65            epoch,
66            flags: flags::FULL_SNAPSHOT,
67        }
68    }
69
70    /// Whether the duplicated wire fields represent one nonzero durable checkpoint identity.
71    #[must_use]
72    pub const fn is_canonical(self) -> bool {
73        self.checkpoint_id != 0 && self.epoch == self.checkpoint_id
74    }
75
76    /// Check whether this barrier requests a full (non-incremental) snapshot.
77    #[must_use]
78    pub const fn is_full_snapshot(&self) -> bool {
79        self.flags & flags::FULL_SNAPSHOT != 0
80    }
81
82    /// Check whether this barrier signals drain/shutdown.
83    #[must_use]
84    pub const fn is_drain(&self) -> bool {
85        self.flags & flags::DRAIN != 0
86    }
87
88    /// Check whether this barrier cancels an in-progress checkpoint.
89    #[must_use]
90    pub const fn is_cancel(&self) -> bool {
91        self.flags & flags::CANCEL != 0
92    }
93}
94
95/// A message that flows through streaming channels.
96///
97/// Wraps user events with control messages (watermarks and barriers)
98/// in a single enum. Operators pattern-match on this to handle
99/// data vs. control flow.
100///
101/// ## Generic Parameter
102///
103/// `T` is the event payload type — typically `RecordBatch` or a
104/// domain-specific event struct.
105#[derive(Debug, Clone, PartialEq)]
106pub enum StreamMessage<T> {
107    /// A user data event.
108    Event(T),
109    /// A watermark indicating event-time progress (millis since epoch).
110    Watermark(i64),
111    /// A checkpoint barrier for consistent snapshots.
112    Barrier(CheckpointBarrier),
113}
114
115impl<T> StreamMessage<T> {
116    /// Returns `true` if this is a barrier message.
117    #[must_use]
118    pub const fn is_barrier(&self) -> bool {
119        matches!(self, Self::Barrier(_))
120    }
121
122    /// Returns `true` if this is a watermark message.
123    #[must_use]
124    pub const fn is_watermark(&self) -> bool {
125        matches!(self, Self::Watermark(_))
126    }
127
128    /// Returns `true` if this is a data event.
129    #[must_use]
130    pub const fn is_event(&self) -> bool {
131        matches!(self, Self::Event(_))
132    }
133
134    /// Extracts the barrier if this is a [`StreamMessage::Barrier`].
135    #[must_use]
136    pub const fn as_barrier(&self) -> Option<&CheckpointBarrier> {
137        match self {
138            Self::Barrier(b) => Some(b),
139            _ => None,
140        }
141    }
142}
143
144/// Cross-thread barrier injector for source operators.
145///
146/// The coordinator thread publishes a barrier command via
147/// [`trigger`](Self::trigger). Source operators poll via
148/// [`BarrierPollHandle::poll`] on each iteration of their event loop.
149///
150/// ## Fast Path
151///
152/// The poll path is a single `AtomicU64::load(Relaxed)` — typically < 10ns.
153/// Only when a barrier is pending does the source perform a compare-exchange
154/// to claim it.
155#[derive(Debug)]
156pub struct CheckpointBarrierInjector {
157    /// Command lifecycle. Payload fields are published before `PENDING`.
158    state: Arc<AtomicU64>,
159    /// Full-width checkpoint ID for the pending command.
160    checkpoint_id: Arc<AtomicU64>,
161    /// Full-width epoch for the pending command.
162    epoch: Arc<AtomicU64>,
163    /// Full-width flags for the pending command.
164    flags: Arc<AtomicU64>,
165}
166
167const STATE_IDLE: u64 = 0;
168const STATE_WRITING: u64 = 1;
169const STATE_PENDING: u64 = 2;
170const STATE_CONSUMING: u64 = 3;
171
172impl CheckpointBarrierInjector {
173    /// Create a new injector with no pending barrier.
174    #[must_use]
175    pub fn new() -> Self {
176        Self {
177            state: Arc::new(AtomicU64::new(STATE_IDLE)),
178            checkpoint_id: Arc::new(AtomicU64::new(0)),
179            epoch: Arc::new(AtomicU64::new(0)),
180            flags: Arc::new(AtomicU64::new(0)),
181        }
182    }
183
184    /// Get a handle that source operators use to poll for barriers.
185    #[must_use]
186    pub fn handle(&self) -> BarrierPollHandle {
187        BarrierPollHandle {
188            state: Arc::clone(&self.state),
189            checkpoint_id: Arc::clone(&self.checkpoint_id),
190            epoch: Arc::clone(&self.epoch),
191            flags: Arc::clone(&self.flags),
192        }
193    }
194
195    /// Whether no barrier command is currently being published, pending, or consumed.
196    /// The coordinator uses this as an all-source preflight before fan-out.
197    #[must_use]
198    pub fn can_trigger(&self) -> bool {
199        self.state.load(Ordering::Acquire) == STATE_IDLE
200    }
201
202    /// Trigger a new checkpoint barrier.
203    ///
204    /// The next [`BarrierPollHandle::poll`] call on any source will
205    /// observe this barrier and return it. Returns `false` without
206    /// modifying the pending command if the identity is zero or noncanonical, or
207    /// if another trigger is being published, is pending, or is being consumed.
208    ///
209    /// # Arguments
210    ///
211    /// * `barrier` - Exact barrier command to publish. `checkpoint_id` and `epoch`
212    ///   must name the same nonzero durable checkpoint.
213    #[must_use = "a rejected trigger leaves the existing barrier pending"]
214    pub fn trigger(&self, barrier: CheckpointBarrier) -> bool {
215        if !barrier.is_canonical()
216            || self
217                .state
218                .compare_exchange(
219                    STATE_IDLE,
220                    STATE_WRITING,
221                    Ordering::Acquire,
222                    Ordering::Relaxed,
223                )
224                .is_err()
225        {
226            return false;
227        }
228
229        // Writers own the payload while state is WRITING. A release-publish of
230        // PENDING makes every full-width field visible to the claiming poller.
231        self.checkpoint_id
232            .store(barrier.checkpoint_id, Ordering::Relaxed);
233        self.epoch.store(barrier.epoch, Ordering::Relaxed);
234        self.flags.store(barrier.flags, Ordering::Relaxed);
235        self.state.store(STATE_PENDING, Ordering::Release);
236        true
237    }
238
239    /// Cancel the pending barrier with this exact checkpoint identity.
240    ///
241    /// Returns `true` only when this call claims and cancels a pending command
242    /// whose checkpoint ID and epoch both match. A stale identity, an idle
243    /// injector, or a barrier already claimed by a poller returns `false`.
244    /// Mismatched pending commands remain available to pollers unchanged.
245    #[must_use = "a false result means the exact pending barrier was not cancelled"]
246    pub fn cancel_exact(&self, barrier: CheckpointBarrier) -> bool {
247        if !barrier.is_canonical()
248            || self.state.load(Ordering::Acquire) != STATE_PENDING
249            || self.checkpoint_id.load(Ordering::Relaxed) != barrier.checkpoint_id
250            || self.epoch.load(Ordering::Relaxed) != barrier.epoch
251        {
252            return false;
253        }
254
255        if self
256            .state
257            .compare_exchange(
258                STATE_PENDING,
259                STATE_CONSUMING,
260                Ordering::Acquire,
261                Ordering::Relaxed,
262            )
263            .is_err()
264        {
265            return false;
266        }
267
268        // Revalidate under exclusive ownership: a poll and a subsequent
269        // trigger may have changed the payload after the optimistic check.
270        let matches = self.checkpoint_id.load(Ordering::Relaxed) == barrier.checkpoint_id
271            && self.epoch.load(Ordering::Relaxed) == barrier.epoch;
272        self.state.store(
273            if matches { STATE_IDLE } else { STATE_PENDING },
274            Ordering::Release,
275        );
276        matches
277    }
278}
279
280impl Default for CheckpointBarrierInjector {
281    fn default() -> Self {
282        Self::new()
283    }
284}
285
286impl Clone for CheckpointBarrierInjector {
287    fn clone(&self) -> Self {
288        Self {
289            state: Arc::clone(&self.state),
290            checkpoint_id: Arc::clone(&self.checkpoint_id),
291            epoch: Arc::clone(&self.epoch),
292            flags: Arc::clone(&self.flags),
293        }
294    }
295}
296
297/// Handle used by source operators to poll for pending barriers.
298///
299/// Cloned from [`CheckpointBarrierInjector::handle`] and stored in the
300/// source operator. The fast path is a single atomic load.
301#[derive(Debug, Clone)]
302pub struct BarrierPollHandle {
303    /// Shared command lifecycle.
304    state: Arc<AtomicU64>,
305    /// Full-width checkpoint ID payload.
306    checkpoint_id: Arc<AtomicU64>,
307    /// Full-width epoch payload.
308    epoch: Arc<AtomicU64>,
309    /// Full-width flags payload.
310    flags: Arc<AtomicU64>,
311}
312
313impl BarrierPollHandle {
314    /// Poll for a pending barrier.
315    ///
316    /// Returns `Some(CheckpointBarrier)` if a barrier is pending and
317    /// this call successfully claimed it (exactly-once delivery across
318    /// handles sharing the same injector). Returns `None` if no barrier
319    /// is pending or another handle already claimed it.
320    ///
321    /// ## Performance
322    ///
323    /// Fast path (no barrier): single `load(Relaxed)` — < 10ns.
324    /// Slow path (barrier pending): one `compare_exchange`.
325    #[must_use]
326    pub fn poll(&self) -> Option<CheckpointBarrier> {
327        // Fast path: no barrier pending
328        if self.state.load(Ordering::Relaxed) != STATE_PENDING {
329            return None;
330        }
331
332        // Claim before reading. CONSUMING prevents the next trigger from
333        // overwriting payload fields until this poller has copied both.
334        if self
335            .state
336            .compare_exchange(
337                STATE_PENDING,
338                STATE_CONSUMING,
339                Ordering::Acquire,
340                Ordering::Relaxed,
341            )
342            .is_ok()
343        {
344            let barrier = CheckpointBarrier {
345                checkpoint_id: self.checkpoint_id.load(Ordering::Relaxed),
346                epoch: self.epoch.load(Ordering::Relaxed),
347                flags: self.flags.load(Ordering::Relaxed),
348            };
349            self.state.store(STATE_IDLE, Ordering::Release);
350            Some(barrier)
351        } else {
352            // Another thread claimed it first
353            None
354        }
355    }
356}
357
358#[cfg(test)]
359mod tests {
360    use super::*;
361
362    #[test]
363    fn test_barrier_size() {
364        assert_eq!(std::mem::size_of::<CheckpointBarrier>(), 24);
365    }
366
367    #[test]
368    fn test_barrier_flags() {
369        let barrier = CheckpointBarrier::new(1, 1);
370        assert!(barrier.is_canonical());
371        assert!(!barrier.is_full_snapshot());
372        assert!(!barrier.is_drain());
373        assert!(!barrier.is_cancel());
374
375        let full = CheckpointBarrier::full_snapshot(1, 1);
376        assert!(full.is_full_snapshot());
377        assert!(!full.is_drain());
378
379        let drain = CheckpointBarrier {
380            checkpoint_id: 1,
381            epoch: 1,
382            flags: flags::DRAIN,
383        };
384        assert!(drain.is_drain());
385        assert!(!CheckpointBarrier::new(0, 0).is_canonical());
386        assert!(!CheckpointBarrier::new(1, 2).is_canonical());
387    }
388
389    #[test]
390    fn test_barrier_roundtrip_via_injector() {
391        let injector = CheckpointBarrierInjector::new();
392        let handle = injector.handle();
393        let expected = CheckpointBarrier {
394            checkpoint_id: 42,
395            epoch: 42,
396            flags: flags::DRAIN,
397        };
398
399        assert!(injector.trigger(expected));
400        assert_eq!(handle.poll(), Some(expected));
401        assert!(handle.poll().is_none(), "cleared after one poll");
402    }
403
404    #[test]
405    fn test_stream_message_variants() {
406        let event: StreamMessage<String> = StreamMessage::Event("hello".into());
407        assert!(event.is_event());
408        assert!(!event.is_barrier());
409        assert!(!event.is_watermark());
410
411        let watermark: StreamMessage<String> = StreamMessage::Watermark(1000);
412        assert!(watermark.is_watermark());
413
414        let barrier: StreamMessage<String> = StreamMessage::Barrier(CheckpointBarrier::new(1, 1));
415        assert!(barrier.is_barrier());
416        assert_eq!(barrier.as_barrier().unwrap().checkpoint_id, 1);
417    }
418
419    #[test]
420    fn test_injector_poll_no_barrier() {
421        let injector = CheckpointBarrierInjector::new();
422        let handle = injector.handle();
423
424        // No barrier pending
425        assert!(handle.poll().is_none());
426    }
427
428    #[test]
429    fn test_injector_trigger_and_poll() {
430        let injector = CheckpointBarrierInjector::new();
431        let handle = injector.handle();
432
433        // Trigger barrier
434        let expected = CheckpointBarrier::full_snapshot(42, 42);
435        assert!(injector.trigger(expected));
436
437        // Poll should return the barrier
438        let barrier = handle.poll().unwrap();
439        assert_eq!(barrier, expected);
440        assert!(barrier.is_full_snapshot());
441
442        // Second poll should return None (already claimed)
443        assert!(handle.poll().is_none());
444    }
445
446    #[test]
447    fn test_injector_multiple_handles() {
448        let injector = CheckpointBarrierInjector::new();
449        let handle1 = injector.handle();
450        let handle2 = injector.handle();
451
452        assert!(injector.trigger(CheckpointBarrier::new(1, 1)));
453
454        // Only one handle should claim it
455        let r1 = handle1.poll();
456        let r2 = handle2.poll();
457
458        // Exactly one should succeed
459        assert!(r1.is_some() || r2.is_some());
460        if r1.is_some() {
461            assert!(r2.is_none());
462        }
463    }
464
465    #[test]
466    fn test_injector_supports_full_width_id_and_flags() {
467        let injector = CheckpointBarrierInjector::new();
468        let handle = injector.handle();
469        let checkpoint_id = u64::from(u32::MAX) + 17;
470        let epoch = checkpoint_id;
471        let barrier_flags = 1_u64 << 63;
472        let expected = CheckpointBarrier {
473            checkpoint_id,
474            epoch,
475            flags: barrier_flags,
476        };
477
478        assert!(injector.trigger(expected));
479        assert_eq!(handle.poll(), Some(expected));
480    }
481
482    #[test]
483    fn test_pending_trigger_is_rejected_without_overwrite() {
484        let injector = CheckpointBarrierInjector::new();
485        let handle = injector.handle();
486        let first = CheckpointBarrier {
487            checkpoint_id: u64::MAX,
488            epoch: u64::MAX,
489            flags: flags::DRAIN,
490        };
491        let second = CheckpointBarrier::full_snapshot(7, 7);
492
493        assert!(injector.trigger(first));
494        assert!(!injector.trigger(second));
495
496        assert_eq!(handle.poll(), Some(first));
497        assert!(injector.trigger(second));
498    }
499
500    #[test]
501    fn test_cancel_exact_removes_matching_pending_barrier() {
502        let injector = CheckpointBarrierInjector::new();
503        let handle = injector.handle();
504        let pending = CheckpointBarrier::full_snapshot(41, 41);
505
506        assert!(injector.trigger(pending));
507        assert!(injector.cancel_exact(pending));
508        assert!(handle.poll().is_none());
509        assert!(injector.can_trigger());
510    }
511
512    #[test]
513    fn test_cancel_exact_preserves_mismatched_pending_barrier() {
514        let injector = CheckpointBarrierInjector::new();
515        let handle = injector.handle();
516        let pending = CheckpointBarrier::full_snapshot(41, 41);
517
518        assert!(injector.trigger(pending));
519        assert!(!injector.cancel_exact(CheckpointBarrier {
520            epoch: pending.epoch + 1,
521            ..pending
522        }));
523        assert!(!injector.cancel_exact(CheckpointBarrier {
524            checkpoint_id: pending.checkpoint_id + 1,
525            ..pending
526        }));
527        assert_eq!(handle.poll(), Some(pending));
528    }
529
530    #[test]
531    fn test_cancel_exact_after_poll_does_not_claim_completed_barrier() {
532        let injector = CheckpointBarrierInjector::new();
533        let handle = injector.handle();
534        let pending = CheckpointBarrier::new(41, 41);
535
536        assert!(injector.trigger(pending));
537        assert_eq!(handle.poll(), Some(pending));
538        assert!(!injector.cancel_exact(pending));
539        assert!(injector.can_trigger());
540    }
541
542    #[test]
543    fn test_cancel_exact_and_poll_have_exactly_one_winner() {
544        let injector = Arc::new(CheckpointBarrierInjector::new());
545        let handle = injector.handle();
546        let pending = CheckpointBarrier::new(41, 41);
547        assert!(injector.trigger(pending));
548
549        let start = Arc::new(std::sync::Barrier::new(3));
550        let cancel_thread = {
551            let injector = Arc::clone(&injector);
552            let start = Arc::clone(&start);
553            std::thread::spawn(move || {
554                start.wait();
555                injector.cancel_exact(pending)
556            })
557        };
558        let poll_thread = {
559            let start = Arc::clone(&start);
560            std::thread::spawn(move || {
561                start.wait();
562                handle.poll()
563            })
564        };
565
566        start.wait();
567        let cancelled = cancel_thread.join().unwrap();
568        let polled = poll_thread.join().unwrap();
569        assert_eq!(cancelled, polled.is_none());
570        if !cancelled {
571            assert_eq!(polled, Some(pending));
572        }
573        assert!(injector.can_trigger());
574    }
575
576    #[test]
577    fn test_noncanonical_identity_is_rejected_without_mutation() {
578        let injector = CheckpointBarrierInjector::new();
579        let handle = injector.handle();
580
581        assert!(!injector.trigger(CheckpointBarrier::new(0, 1)));
582        assert!(!injector.trigger(CheckpointBarrier::new(1, 0)));
583        assert!(!injector.trigger(CheckpointBarrier::new(1, 2)));
584        assert!(injector.can_trigger());
585        assert!(handle.poll().is_none());
586
587        let expected = CheckpointBarrier::new(9, 9);
588        assert!(injector.trigger(expected));
589        assert!(!injector.trigger(CheckpointBarrier::new(0, 11)));
590        assert!(!injector.trigger(CheckpointBarrier::new(11, 0)));
591        assert!(!injector.trigger(CheckpointBarrier::new(11, 12)));
592        assert_eq!(handle.poll(), Some(expected));
593    }
594
595    #[test]
596    fn test_concurrent_triggers_admit_exactly_one() {
597        let injector = Arc::new(CheckpointBarrierInjector::new());
598        let handle = injector.handle();
599        let mut threads = Vec::new();
600
601        for checkpoint_id in 1..=16 {
602            let injector = Arc::clone(&injector);
603            threads.push(std::thread::spawn(move || {
604                let barrier = CheckpointBarrier {
605                    checkpoint_id,
606                    epoch: checkpoint_id,
607                    flags: checkpoint_id << 40,
608                };
609                (barrier, injector.trigger(barrier))
610            }));
611        }
612
613        let accepted = threads
614            .into_iter()
615            .map(|thread| thread.join().unwrap())
616            .filter_map(|(barrier, accepted)| accepted.then_some(barrier))
617            .collect::<Vec<_>>();
618        assert_eq!(accepted.len(), 1);
619
620        assert_eq!(handle.poll(), Some(accepted[0]));
621    }
622}
laminar_core/checkpoint/barrier.rs

laminar_core/checkpoint/
barrier.rs
