embassy_sync/waitqueue/
waker_registration.rs

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