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}
53