atomic_wake/
atomic-wake.rs

1use option_lock::OptionLock;
2use std::task::Waker;
3
4enum ResultState<T, E> {
5    Ready(Result<T, E>),
6    Wake(Waker),
7}
8
9pub struct AsyncResult<T, E> {
10    state: OptionLock<ResultState<T, E>>,
11}
12
13impl<T, E> AsyncResult<T, E> {
14    pub const fn new() -> Self {
15        Self {
16            state: OptionLock::empty(),
17        }
18    }
19
20    pub fn poll(&self, waker: Option<&Waker>) -> Option<Result<T, E>> {
21        match self.state.try_lock() {
22            Ok(mut guard) => match guard.take() {
23                Some(ResultState::Ready(result)) => Some(result),
24                Some(ResultState::Wake(_)) | None => {
25                    waker.map(|waker| guard.replace(ResultState::Wake(waker.clone())));
26                    None
27                }
28            },
29            _ => {
30                // result is currently being stored
31                waker.map(Waker::wake_by_ref);
32                None
33            }
34        }
35    }
36
37    pub fn fulfill(&self, result: Result<T, E>) -> Result<(), Result<T, E>> {
38        // retry method is left up to the caller (spin, yield thread, etc)
39        if let Ok(mut guard) = self.state.try_lock() {
40            let prev = guard.replace(ResultState::Ready(result));
41            drop(guard);
42            if let Some(ResultState::Wake(waker)) = prev {
43                waker.wake();
44            }
45            Ok(())
46        } else {
47            Err(result)
48        }
49    }
50}
51
52fn main() {
53    let fut = AsyncResult::<u32, ()>::new();
54    fut.fulfill(Ok(100)).unwrap();
55    assert_eq!(fut.poll(None), Some(Ok(100)));
56}