nuclei/sync/
set_once.rs

1//! A cell set once, whose readers can wait for the value.
2
3use core::error::Error;
4use core::fmt;
5
6use super::event::Event;
7use super::once_cell::{OnceCell, SetError};
8
9/// Error returned from [`SetOnce::set`]: the cell was set already, or
10/// another `set` is setting it now. The value comes back.
11#[derive(Debug, PartialEq, Eq)]
12pub struct SetOnceError<T>(pub T);
13
14impl<T> fmt::Display for SetOnceError<T> {
15    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
16        f.write_str("SetOnceError")
17    }
18}
19
20impl<T: fmt::Debug> Error for SetOnceError<T> {}
21
22/// A cell that is set once, and that readers can wait on until it is.
23///
24/// It is a [`OnceCell`] that nobody initializes asynchronously, paired with
25/// an [`Event`] that [`wait`](Self::wait) sleeps on: [`set`](Self::set)
26/// publishes the value with one compare-and-swap, and then sets the event,
27/// which wakes every waiter.
28///
29/// ```
30/// use nuclei::sync::SetOnce;
31///
32/// static GREETING: SetOnce<&str> = SetOnce::const_new();
33/// assert!(GREETING.get().is_none());
34/// assert!(GREETING.set("hello").is_ok());
35/// assert_eq!(GREETING.get(), Some(&"hello"));
36/// assert!(GREETING.set("again").is_err());
37/// ```
38pub struct SetOnce<T> {
39    cell: OnceCell<T>,
40    ready: Event,
41}
42
43impl<T> SetOnce<T> {
44    loom_const_fn! {
45        /// An unset cell.
46        pub fn new() -> Self {
47            Self { cell: OnceCell::new(), ready: Event::new() }
48        }
49    }
50
51    loom_const_fn! {
52        /// An unset cell, usable in a `static`.
53        pub fn const_new() -> Self {
54            Self::new()
55        }
56    }
57
58    /// A cell that holds `value` if there is one, and is unset otherwise.
59    pub fn new_with(value: Option<T>) -> Self {
60        match value {
61            Some(value) => Self::from(value),
62            None => Self::new(),
63        }
64    }
65
66    loom_const_fn! {
67        /// A cell that is set to `value` already, usable in a `static`.
68        pub fn const_new_with(value: T) -> Self {
69            Self { cell: OnceCell::const_new_with(value), ready: Event::new_set(true) }
70        }
71    }
72
73    /// Whether the cell holds a value.
74    pub fn initialized(&self) -> bool {
75        self.cell.initialized()
76    }
77
78    /// The value, once set.
79    pub fn get(&self) -> Option<&T> {
80        self.cell.get()
81    }
82
83    /// Sets the value, and wakes every task waiting in
84    /// [`wait`](Self::wait). One compare-and-swap publishes; never waits.
85    ///
86    /// # Errors
87    ///
88    /// [`SetOnceError`] with the value when the cell is set already, or
89    /// another call is setting it at this moment.
90    pub fn set(&self, value: T) -> Result<(), SetOnceError<T>> {
91        match self.cell.set(value) {
92            Ok(()) => {
93                self.ready.set();
94                Ok(())
95            }
96            Err(SetError::AlreadyInitializedError(value) | SetError::InitializingError(value)) => {
97                Err(SetOnceError(value))
98            }
99        }
100    }
101
102    /// Consumes the cell and returns its value, if it has one.
103    pub fn into_inner(self) -> Option<T> {
104        self.cell.into_inner()
105    }
106
107    /// Waits until the cell is set, and returns the value: at once if it
108    /// is set already.
109    ///
110    /// Dropping the future withdraws the wait.
111    pub async fn wait(&self) -> &T {
112        loop {
113            if let Some(value) = self.cell.get() {
114                return value;
115            }
116            self.ready.wait().await;
117        }
118    }
119}
120
121impl<T> Default for SetOnce<T> {
122    fn default() -> Self {
123        Self::new()
124    }
125}
126
127impl<T: fmt::Debug> fmt::Debug for SetOnce<T> {
128    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
129        f.debug_struct("SetOnce")
130            .field("value", &self.get())
131            .finish()
132    }
133}
134
135impl<T: Clone> Clone for SetOnce<T> {
136    fn clone(&self) -> Self {
137        Self::new_with(self.get().cloned())
138    }
139}
140
141impl<T: PartialEq> PartialEq for SetOnce<T> {
142    fn eq(&self, other: &Self) -> bool {
143        self.get() == other.get()
144    }
145}
146
147impl<T: Eq> Eq for SetOnce<T> {}
148
149impl<T> From<T> for SetOnce<T> {
150    fn from(value: T) -> Self {
151        Self::const_new_with(value)
152    }
153}
154
155#[cfg(all(test, not(loom)))]
156mod tests {
157    use super::*;
158    use crate::sync::test_support::{Polled, block_on};
159    use std::sync::Arc;
160
161    #[test]
162    fn a_set_wakes_every_waiter_and_later_waits_pass() {
163        let once = SetOnce::new();
164        let mut first = Polled::new(once.wait());
165        first.pending();
166        let mut second = Polled::new(once.wait());
167        second.pending();
168        assert!(!once.initialized());
169        assert_eq!(once.set(7), Ok(()));
170        assert!(once.initialized());
171        assert_eq!((first.wakes(), second.wakes()), (1, 1));
172        assert_eq!((first.ready(), second.ready()), (&7, &7));
173        assert_eq!(Polled::new(once.wait()).ready(), &7);
174    }
175
176    #[test]
177    fn a_second_set_gives_its_value_back() {
178        let once = SetOnce::new();
179        assert_eq!(once.set(1), Ok(()));
180        let err = once.set(2).unwrap_err();
181        assert_eq!(err, SetOnceError(2));
182        assert_eq!(err.0, 2);
183        assert_eq!(err.to_string(), "SetOnceError");
184        assert_eq!(once.get(), Some(&1));
185        assert_eq!(once.into_inner(), Some(1));
186    }
187
188    #[test]
189    fn a_dropped_wait_leaves_the_rest_to_their_wake() {
190        let once = SetOnce::new();
191        let mut kept = Polled::new(once.wait());
192        kept.pending();
193        Polled::new(once.wait()).pending();
194        assert_eq!(once.set(5), Ok(()));
195        assert_eq!(kept.ready(), &5);
196    }
197
198    #[test]
199    fn constructors_conversions_and_traits() {
200        static FULL: SetOnce<u8> = SetOnce::const_new_with(3);
201        static EMPTY: SetOnce<u8> = SetOnce::const_new();
202        assert_eq!(FULL.get(), Some(&3));
203        assert!(FULL.set(4).is_err());
204        assert_eq!(*block_on(FULL.wait()), 3);
205        assert_eq!(EMPTY.get(), None);
206        assert_eq!(SetOnce::new_with(Some(1)).get(), Some(&1));
207        assert_eq!(SetOnce::<u8>::new_with(None).get(), None);
208        assert_eq!(SetOnce::<u8>::default().into_inner(), None);
209        assert_eq!(
210            format!("{:?}", SetOnce::from(2)),
211            "SetOnce { value: Some(2) }"
212        );
213        let once = SetOnce::from(String::from("x"));
214        assert_eq!(once.clone(), once);
215        assert_ne!(once, SetOnce::new());
216    }
217
218    #[test]
219    fn a_thread_waiting_for_the_value_gets_it() {
220        let once = Arc::new(SetOnce::new());
221        let waiter = {
222            let once = Arc::clone(&once);
223            std::thread::spawn(move || *block_on(once.wait()))
224        };
225        assert_eq!(once.set(9), Ok(()));
226        assert_eq!(waiter.join().expect("waiter"), 9);
227    }
228
229    #[test]
230    fn two_sets_at_once_leave_exactly_one_winner() {
231        for _ in 0..if cfg!(miri) { 3 } else { 200 } {
232            let once = Arc::new(SetOnce::new());
233            let racers: Vec<_> = [1, 2]
234                .into_iter()
235                .map(|n| {
236                    let once = Arc::clone(&once);
237                    std::thread::spawn(move || once.set(n).is_ok())
238                })
239                .collect();
240            let wins = racers
241                .into_iter()
242                .map(|racer| racer.join().expect("racer"))
243                .filter(|won| *won)
244                .count();
245            assert_eq!(wins, 1);
246            assert!(once.get().is_some());
247        }
248    }
249}
nuclei/sync/set_once.rs

nuclei/sync/
set_once.rs
