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

etl_core/source/
barrier.rs
