embassy_sync/waitqueue/
waker_registration.rs

1use core::mem;
2use core::task::Waker;
3
4/// Utility struct to register and wake a waker.
5/// If a waker is registered, registering another waker will replace the previous one.
6/// The previous waker will be woken in this case, giving it a chance to reregister itself.
7/// Although it is possible to wake multiple tasks this way,
8/// this will cause them to wake each other in a loop registering themselves.
9#[derive(Debug, Default)]
10pub struct WakerRegistration {
11    waker: Option<Waker>,
12}
13
14impl WakerRegistration {
15    /// Create a new `WakerRegistration`.
16    pub const fn new() -> Self {
17        Self { waker: None }
18    }
19
20    /// Register a waker. Overwrites the previous waker, if any.
21    pub fn register(&mut self, w: &Waker) {
22        match self.waker {
23            // Optimization: If both the old and new Wakers wake the same task, we can simply
24            // keep the old waker, skipping the clone. (In most executor implementations,
25            // cloning a waker is somewhat expensive, comparable to cloning an Arc).
26            Some(ref w2) if (w2.will_wake(w)) => {}
27            _ => {
28                // clone the new waker and store it
29                if let Some(old_waker) = mem::replace(&mut self.waker, Some(w.clone())) {
30                    // We had a waker registered for another task. Wake it, so the other task can
31                    // reregister itself if it's still interested.
32                    //
33                    // If two tasks are waiting on the same thing concurrently, this will cause them
34                    // to wake each other in a loop fighting over this WakerRegistration. This wastes
35                    // CPU but things will still work.
36                    //
37                    // If the user wants to have two tasks waiting on the same thing they should use
38                    // a more appropriate primitive that can store multiple wakers.
39                    old_waker.wake()
40                }
41            }
42        }
43    }
44
45    /// Wake the registered waker, if any.
46    pub fn wake(&mut self) {
47        if let Some(w) = self.waker.take() {
48            w.wake()
49        }
50    }
51
52    /// Returns true if a waker is currently registered
53    pub fn occupied(&self) -> bool {
54        self.waker.is_some()
55    }
56}