rtic_common/
waker_registration.rs

Help
1//! Waker registration utility.
2
3use core::cell::UnsafeCell;
4use core::task::Waker;
5
6/// A critical section based waker handler.
7pub struct CriticalSectionWakerRegistration {
8    waker: UnsafeCell<Option<Waker>>,
9}
10
11unsafe impl Send for CriticalSectionWakerRegistration {}
12unsafe impl Sync for CriticalSectionWakerRegistration {}
13
14impl CriticalSectionWakerRegistration {
15    /// Create a new waker registration.
16    pub const fn new() -> Self {
17        Self {
18            waker: UnsafeCell::new(None),
19        }
20    }
21
22    /// Register a waker.
23    /// This will overwrite the previous waker if there was one.
24    pub fn register(&self, new_waker: &Waker) {
25        critical_section::with(|_| {
26            // SAFETY: This access is protected by the critical section.
27            let self_waker = unsafe { &mut *self.waker.get() };
28
29            // From embassy
30            // https://github.com/embassy-rs/embassy/blob/b99533607ceed225dd12ae73aaa9a0d969a7365e/embassy-sync/src/waitqueue/waker.rs#L59-L61
31            match self_waker {
32                // Optimization: If both the old and new Wakers wake the same task, we can simply
33                // keep the old waker, skipping the clone. (In most executor implementations,
34                // cloning a waker is somewhat expensive, comparable to cloning an Arc).
35                Some(ref w2) if (w2.will_wake(new_waker)) => {}
36                _ => {
37                    // clone the new waker and store it
38                    if let Some(old_waker) = core::mem::replace(self_waker, Some(new_waker.clone()))
39                    {
40                        // We had a waker registered for another task. Wake it, so the other task can
41                        // reregister itself if it's still interested.
42                        //
43                        // If two tasks are waiting on the same thing concurrently, this will cause them
44                        // to wake each other in a loop fighting over this WakerRegistration. This wastes
45                        // CPU but things will still work.
46                        //
47                        // If the user wants to have two tasks waiting on the same thing they should use
48                        // a more appropriate primitive that can store multiple wakers.
49                        old_waker.wake()
50                    }
51                }
52            }
53        });
54    }
55
56    /// Wake the waker.
57    pub fn wake(&self) {
58        critical_section::with(|_| {
59            // SAFETY: This access is protected by the critical section.
60            let self_waker = unsafe { &mut *self.waker.get() };
61            if let Some(waker) = self_waker.take() {
62                waker.wake()
63            }
64        });
65    }
66}
67
68impl Default for CriticalSectionWakerRegistration {
69    fn default() -> Self {
70        Self::new()
71    }
72}
rtic_common/waker_registration.rs

rtic_common/
waker_registration.rs
