spate_core/source/
barrier.rs

1//! Drain barrier: revocation and shutdown share one choreography.
2
3use std::sync::atomic::{AtomicUsize, Ordering};
4use std::sync::{Arc, Condvar, Mutex};
5use std::time::Duration;
6
7use crate::clock::{Clock, SystemClock};
8
9#[derive(Debug)]
10struct Inner {
11    remaining: AtomicUsize,
12    done: Mutex<bool>,
13    cv: Condvar,
14}
15
16/// A countdown barrier the controller waits on while pipeline threads
17/// drain. Created with the number of arriving parties (the threads that own
18/// the lanes being stopped); each thread calls [`DrainBarrier::arrive`]
19/// exactly once after it has stopped the lanes and flushed its in-flight
20/// records; the controller's [`DrainBarrier::wait`] returns `true` once all
21/// parties arrived, or `false` when the drain deadline elapsed. On timeout the
22/// caller proceeds and unflushed data replays after restart; at-least-once
23/// holds either way.
24#[derive(Clone, Debug)]
25pub struct DrainBarrier {
26    inner: Arc<Inner>,
27}
28
29impl DrainBarrier {
30    /// A barrier expecting `parties` arrivals. `parties == 0` is complete
31    /// immediately.
32    #[must_use]
33    pub fn new(parties: usize) -> Self {
34        DrainBarrier {
35            inner: Arc::new(Inner {
36                remaining: AtomicUsize::new(parties),
37                done: Mutex::new(parties == 0),
38                cv: Condvar::new(),
39            }),
40        }
41    }
42
43    /// Record this party's drain as complete. Idempotence is the caller's
44    /// responsibility. Call exactly once per party.
45    pub fn arrive(&self) {
46        let prev = self.inner.remaining.fetch_sub(1, Ordering::AcqRel);
47        debug_assert!(prev > 0, "more arrivals than barrier parties");
48        if prev == 1 {
49            let mut done = self.inner.done.lock().expect("barrier lock");
50            *done = true;
51            self.inner.cv.notify_all();
52        }
53    }
54
55    /// Parties still outstanding.
56    #[must_use]
57    pub fn remaining(&self) -> usize {
58        self.inner.remaining.load(Ordering::Acquire)
59    }
60
61    /// Block until every party arrived or `timeout` elapsed. Returns
62    /// whether the drain completed.
63    #[must_use]
64    pub fn wait(&self, timeout: Duration) -> bool {
65        // Real time: a test clock cannot wake the condvar. Callers own the
66        // deadline and pass what is left of it.
67        let deadline = SystemClock.now() + timeout;
68        let mut done = self.inner.done.lock().expect("barrier lock");
69        while !*done {
70            let now = SystemClock.now();
71            if now >= deadline {
72                return false;
73            }
74            let (guard, _timeout) = self
75                .inner
76                .cv
77                .wait_timeout(done, deadline - now)
78                .expect("barrier lock");
79            done = guard;
80        }
81        true
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88
89    #[test]
90    #[expect(clippy::disallowed_methods, reason = "the test joins the thread")]
91    fn completes_when_all_arrive() {
92        let barrier = DrainBarrier::new(2);
93        let b2 = barrier.clone();
94        let t = std::thread::spawn(move || {
95            b2.arrive();
96            b2.arrive();
97        });
98        assert!(barrier.wait(Duration::from_secs(5)));
99        assert_eq!(barrier.remaining(), 0);
100        t.join().unwrap();
101    }
102
103    #[test]
104    fn times_out_when_a_party_never_arrives() {
105        let barrier = DrainBarrier::new(2);
106        barrier.arrive();
107        assert!(!barrier.wait(Duration::from_millis(50)));
108        assert_eq!(barrier.remaining(), 1);
109    }
110
111    #[test]
112    fn zero_parties_is_immediately_complete() {
113        let barrier = DrainBarrier::new(0);
114        assert!(barrier.wait(Duration::ZERO));
115    }
116
117    #[test]
118    fn wait_after_completion_returns_immediately() {
119        let barrier = DrainBarrier::new(1);
120        barrier.arrive();
121        assert!(barrier.wait(Duration::ZERO));
122        assert!(barrier.wait(Duration::from_secs(1)));
123    }
124}
spate_core/source/barrier.rs

spate_core/source/
barrier.rs
