embassy_sync/waitqueue/
atomic_waker.rs

1use core::cell::Cell;
2use core::task::Waker;
3
4use crate::blocking_mutex::raw::{CriticalSectionRawMutex, RawMutex};
5use crate::blocking_mutex::Mutex;
6
7/// Utility struct to register and wake a waker.
8pub struct GenericAtomicWaker<M: RawMutex> {
9    waker: Mutex<M, Cell<Option<Waker>>>,
10}
11
12impl<M: RawMutex> GenericAtomicWaker<M> {
13    /// Create a new `AtomicWaker`.
14    pub const fn new(mutex: M) -> Self {
15        Self {
16            waker: Mutex::const_new(mutex, Cell::new(None)),
17        }
18    }
19
20    /// Register a waker. Overwrites the previous waker, if any.
21    pub fn register(&self, w: &Waker) {
22        self.waker.lock(|cell| {
23            cell.set(match cell.replace(None) {
24                Some(w2) if (w2.will_wake(w)) => Some(w2),
25                _ => Some(w.clone()),
26            })
27        })
28    }
29
30    /// Wake the registered waker, if any.
31    pub fn wake(&self) {
32        self.waker.lock(|cell| {
33            if let Some(w) = cell.replace(None) {
34                w.wake_by_ref();
35                cell.set(Some(w));
36            }
37        })
38    }
39}
40
41/// Utility struct to register and wake a waker.
42pub struct AtomicWaker {
43    waker: GenericAtomicWaker<CriticalSectionRawMutex>,
44}
45
46impl AtomicWaker {
47    /// Create a new `AtomicWaker`.
48    pub const fn new() -> Self {
49        Self {
50            waker: GenericAtomicWaker::new(CriticalSectionRawMutex::new()),
51        }
52    }
53
54    /// Register a waker. Overwrites the previous waker, if any.
55    pub fn register(&self, w: &Waker) {
56        self.waker.register(w);
57    }
58
59    /// Wake the registered waker, if any.
60    pub fn wake(&self) {
61        self.waker.wake();
62    }
63}