nuclei/sync/
latch.rs

1//! A count-down latch: waits complete once a count reaches zero.
2
3use core::fmt;
4use core::future::Future;
5use core::pin::Pin;
6use core::task::{Context, Poll};
7
8use super::event::{Event, EventWait};
9use super::internal::{AtomicUsize, Ordering};
10
11/// A latch that opens when its count reaches zero, and stays open.
12///
13/// ```
14/// use nuclei::sync::Latch;
15///
16/// let done = Latch::new(2);
17/// done.count_down();
18/// assert!(!done.is_set());
19/// done.count_down();
20/// assert!(done.is_set());
21/// ```
22pub struct Latch {
23    count: AtomicUsize,
24    open: Event,
25}
26
27impl Latch {
28    loom_const_fn! {
29        /// A latch that opens after `count` calls to
30        /// [`count_down`](Self::count_down); open at once when `count` is
31        /// zero.
32        pub fn new(count: usize) -> Self {
33            Self { count: AtomicUsize::new(count), open: Event::new_set(count == 0) }
34        }
35    }
36
37    /// Lowers the count by one, opening the latch when it reaches zero.
38    /// Does nothing once the latch is open.
39    pub fn count_down(&self) {
40        let mut count = self.count.load(Ordering::Relaxed);
41        while count != 0 {
42            match self.count.compare_exchange_weak(
43                count,
44                count - 1,
45                Ordering::AcqRel,
46                Ordering::Relaxed,
47            ) {
48                Ok(1) => return self.open.set(),
49                Ok(_) => return,
50                Err(actual) => count = actual,
51            }
52        }
53    }
54
55    /// The count still to go.
56    pub fn count(&self) -> usize {
57        self.count.load(Ordering::Acquire)
58    }
59
60    /// Whether the latch is open.
61    pub fn is_set(&self) -> bool {
62        self.open.is_set()
63    }
64
65    /// A future that completes once the latch is open.
66    pub fn wait(&self) -> LatchWait<'_> {
67        LatchWait(self.open.wait())
68    }
69}
70
71impl fmt::Debug for Latch {
72    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
73        f.debug_struct("Latch")
74            .field("count", &self.count())
75            .finish()
76    }
77}
78
79/// The future [`Latch::wait`] returns.
80#[must_use = "futures do nothing unless polled"]
81#[derive(Debug)]
82pub struct LatchWait<'a>(EventWait<'a>);
83
84impl Future for LatchWait<'_> {
85    type Output = ();
86
87    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
88        Pin::new(&mut self.get_mut().0).poll(cx)
89    }
90}
91
92#[cfg(all(test, not(loom)))]
93mod tests {
94    use super::*;
95    use crate::sync::test_support::Polled;
96
97    #[test]
98    fn the_latch_opens_at_zero_and_stays_open() {
99        let latch = Latch::new(2);
100        let mut waiting = Polled::new(latch.wait());
101        waiting.pending();
102        latch.count_down();
103        assert_eq!(waiting.wakes(), 0);
104        latch.count_down();
105        latch.count_down();
106        assert_eq!(latch.count(), 0);
107        waiting.ready();
108        Polled::new(latch.wait()).ready();
109        Polled::new(Latch::new(0).wait()).ready();
110    }
111}
nuclei/sync/latch.rs

nuclei/sync/
latch.rs
