monoio/src/task/waker_fn.rs


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

use core::task::{RawWaker, RawWakerVTable, Waker};
use std::cell::Cell;

/// Creates a waker that does nothing.
///
/// This `Waker` is useful for polling a `Future` to check whether it is
/// `Ready`, without doing any additional work.
pub(crate) fn dummy_waker() -> Waker {
    fn raw_waker() -> RawWaker {
        // the pointer is never dereferenced, so null is ok
        RawWaker::new(std::ptr::null::<()>(), vtable())
    }

    fn vtable() -> &'static RawWakerVTable {
        &RawWakerVTable::new(
            |_| raw_waker(),
            |_| {
                set_poll();
            },
            |_| {
                set_poll();
            },
            |_| {},
        )
    }

    unsafe { Waker::from_raw(raw_waker()) }
}

thread_local! {
    pub(crate) static SHOULD_POLL: Cell<bool> = Cell::new(true);
}

#[inline]
pub(crate) fn should_poll() -> bool {
    SHOULD_POLL.with(|b| b.replace(false))
}

#[inline]
pub(crate) fn set_poll() {
    SHOULD_POLL.with(|b| {
        b.set(true);
    })
}