Skip to main content

tokio/runtime/scheduler/
inject.rs

1//! Inject queue used to send wakeups to a work-stealing scheduler
2
3use crate::loom::sync::Mutex;
4use crate::runtime::task;
5
6mod pop;
7pub(crate) use pop::Pop;
8
9mod shared;
10pub(crate) use shared::Shared;
11
12mod synced;
13pub(crate) use synced::Synced;
14
15cfg_rt_multi_thread! {
16    mod rt_multi_thread;
17    pub(crate) use rt_multi_thread::InjectQueue;
18}
19
20mod metrics;
21
22/// Growable, MPMC queue used to inject new tasks into the scheduler and as an
23/// overflow queue when the local, fixed-size, array queue overflows.
24pub(crate) struct Inject<T: 'static> {
25    shared: Shared<T>,
26    synced: Mutex<Synced>,
27}
28
29impl<T: 'static> Inject<T> {
30    pub(crate) fn new() -> Inject<T> {
31        let (shared, synced) = Shared::new();
32
33        Inject {
34            shared,
35            synced: Mutex::new(synced),
36        }
37    }
38
39    // Kind of annoying to have to include the cfg here
40    #[cfg(any(all(tokio_unstable, feature = "taskdump"), feature = "rt-multi-thread"))]
41    pub(crate) fn is_closed(&self) -> bool {
42        let synced = self.synced.lock();
43        self.shared.is_closed(&synced)
44    }
45
46    /// Closes the injection queue, returns `true` if the queue is open when the
47    /// transition is made.
48    pub(crate) fn close(&self) -> bool {
49        let mut synced = self.synced.lock();
50        self.shared.close(&mut synced)
51    }
52
53    /// Pushes a value into the queue.
54    ///
55    /// This does nothing if the queue is closed.
56    pub(crate) fn push(&self, task: task::Notified<T>) {
57        let mut synced = self.synced.lock();
58        // safety: passing correct `Synced`
59        unsafe { self.shared.push(&mut synced, task) }
60    }
61
62    pub(crate) fn pop(&self) -> Option<task::Notified<T>> {
63        if self.shared.is_empty() {
64            return None;
65        }
66
67        let mut synced = self.synced.lock();
68        // safety: passing correct `Synced`
69        unsafe { self.shared.pop(&mut synced) }
70    }
71}