compio_runtime/runtime/
opt_waker.rs

1use std::{
2    sync::{
3        Arc,
4        atomic::{AtomicBool, Ordering},
5    },
6    task::{Wake, Waker},
7};
8
9#[cfg(not(feature = "notify-always"))]
10use crate::runtime::send_wrapper::SendWrapper;
11
12/// An optimized waker that avoids unnecessary wake-ups on the same thread.
13pub struct OptWaker {
14    waker: Waker,
15    #[cfg(not(feature = "notify-always"))]
16    current_thread: SendWrapper<()>,
17    is_woke: AtomicBool,
18}
19
20impl OptWaker {
21    pub(crate) fn new(waker: Waker) -> Arc<Self> {
22        Arc::new(Self {
23            waker,
24            #[cfg(not(feature = "notify-always"))]
25            current_thread: SendWrapper::new(()),
26            is_woke: AtomicBool::new(false),
27        })
28    }
29
30    /// Returns `true` if the waker has been woke, and resets the state to not
31    /// woke.
32    pub fn reset(&self) -> bool {
33        self.is_woke.swap(false, Ordering::AcqRel)
34    }
35
36    #[inline]
37    fn should_wake(&self) -> bool {
38        #[cfg(feature = "notify-always")]
39        {
40            true
41        }
42        #[cfg(not(feature = "notify-always"))]
43        {
44            !self.current_thread.valid()
45        }
46    }
47}
48
49impl Wake for OptWaker {
50    fn wake(self: Arc<Self>) {
51        self.wake_by_ref();
52    }
53
54    fn wake_by_ref(self: &Arc<Self>) {
55        if !self.is_woke.swap(true, Ordering::AcqRel) && self.should_wake() {
56            self.waker.wake_by_ref();
57        }
58    }
59}