Skip to main content

conc_util/event/
onceevent.rs

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