Skip to main content

conc_util/event/
onceevent.rs

1#![cfg(feature = "alloc")]
2
3use core::{
4    mem::ManuallyDrop,
5    pin::Pin,
6    task::{self, Waker},
7};
8
9use alloc::{sync::Arc, vec::Vec};
10
11use crate::sync::{OnceFlag, SpinLock};
12
13#[derive(Debug)]
14struct OnceInner {
15    used: OnceFlag,
16    wakers: SpinLock<ManuallyDrop<Vec<Waker>>>,
17}
18
19impl Drop for OnceInner {
20    fn drop(&mut self) {
21        if !self.used.fired_mut() {
22            // SAFETY: this is only dropped here and taken in `fire()`,
23            // but it also sets the `used` flag which is checked above.
24            unsafe { ManuallyDrop::drop(self.wakers.get()) };
25        }
26    }
27}
28
29#[derive(Clone, Debug)]
30pub struct OnceEvent(Arc<OnceInner>);
31
32impl OnceEvent {
33    pub fn new() -> Self {
34        Self(Arc::new(OnceInner {
35            used: OnceFlag::new(),
36            wakers: SpinLock::new(ManuallyDrop::new(Vec::new())),
37        }))
38    }
39
40    pub fn fire(&self) {
41        if self.0.used.fire() {
42            return;
43        }
44
45        let mut lock = self.0.wakers.lock();
46        // SAFETY:
47        // 1) Value is taken once - `used` is checked at the beginning
48        // 2) Value is never used again - `poll` checks the same `used`
49        let wakers = unsafe { ManuallyDrop::take(&mut lock) };
50        drop(lock);
51
52        for waker in wakers {
53            waker.wake();
54        }
55    }
56}
57
58impl Future for OnceEvent {
59    type Output = ();
60
61    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
62        if self.0.used.fired() {
63            return task::Poll::Ready(());
64        }
65
66        let mut wakers = self.0.wakers.lock();
67
68        if self.0.used.fired() {
69            return task::Poll::Ready(());
70        }
71
72        wakers.push(cx.waker().clone());
73        return task::Poll::Pending;
74    }
75}