conc_util/event/
onceevent.rs

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

conc_util/event/
onceevent.rs
