tokio/runtime/scheduler/
mod.rs

1cfg_rt! {
2    pub(crate) mod current_thread;
3    pub(crate) use current_thread::CurrentThread;
4
5    mod defer;
6    use defer::Defer;
7
8    pub(crate) mod inject;
9    pub(crate) use inject::Inject;
10
11    use crate::runtime::TaskHooks;
12
13    use crate::runtime::WorkerMetrics;
14}
15
16cfg_rt_multi_thread! {
17    mod block_in_place;
18    pub(crate) use block_in_place::block_in_place;
19
20    mod lock;
21    use lock::Lock;
22
23    pub(crate) mod multi_thread;
24    pub(crate) use multi_thread::MultiThread;
25}
26
27pub(super) mod util;
28
29use crate::runtime::driver;
30
31#[derive(Debug, Clone)]
32pub(crate) enum Handle {
33    #[cfg(feature = "rt")]
34    CurrentThread(Arc<current_thread::Handle>),
35
36    #[cfg(feature = "rt-multi-thread")]
37    MultiThread(Arc<multi_thread::Handle>),
38
39    // TODO: This is to avoid triggering "dead code" warnings many other places
40    // in the codebase. Remove this during a later cleanup
41    #[cfg(not(feature = "rt"))]
42    #[allow(dead_code)]
43    Disabled,
44}
45
46#[cfg(feature = "rt")]
47pub(super) enum Context {
48    CurrentThread(current_thread::Context),
49
50    #[cfg(feature = "rt-multi-thread")]
51    MultiThread(multi_thread::Context),
52}
53
54impl Handle {
55    #[cfg_attr(not(feature = "full"), allow(dead_code))]
56    pub(crate) fn driver(&self) -> &driver::Handle {
57        match *self {
58            #[cfg(feature = "rt")]
59            Handle::CurrentThread(ref h) => &h.driver,
60
61            #[cfg(feature = "rt-multi-thread")]
62            Handle::MultiThread(ref h) => &h.driver,
63
64            #[cfg(not(feature = "rt"))]
65            Handle::Disabled => unreachable!(),
66        }
67    }
68}
69
70cfg_rt! {
71    use crate::future::Future;
72    use crate::loom::sync::Arc;
73    use crate::runtime::{blocking, task::{Id, SpawnLocation}};
74    use crate::runtime::context;
75    use crate::task::JoinHandle;
76    use crate::util::RngSeedGenerator;
77    use std::task::Waker;
78
79    macro_rules! match_flavor {
80        ($self:expr, $ty:ident($h:ident) => $e:expr) => {
81            match $self {
82                $ty::CurrentThread($h) => $e,
83
84                #[cfg(feature = "rt-multi-thread")]
85                $ty::MultiThread($h) => $e,
86            }
87        }
88    }
89
90    impl Handle {
91        #[track_caller]
92        pub(crate) fn current() -> Handle {
93            match context::with_current(Clone::clone) {
94                Ok(handle) => handle,
95                Err(e) => panic!("{}", e),
96            }
97        }
98
99        pub(crate) fn blocking_spawner(&self) -> &blocking::Spawner {
100            match_flavor!(self, Handle(h) => &h.blocking_spawner)
101        }
102
103        pub(crate) fn is_local(&self) -> bool {
104            match self {
105                Handle::CurrentThread(h) => h.local_tid.is_some(),
106
107                #[cfg(feature = "rt-multi-thread")]
108                Handle::MultiThread(_) => false,
109            }
110        }
111
112        #[cfg(feature = "time")]
113        pub(crate) fn timer_flavor(&self) -> crate::runtime::TimerFlavor {
114            match self {
115                Handle::CurrentThread(_) => crate::runtime::TimerFlavor::Traditional,
116
117                #[cfg(feature = "rt-multi-thread")]
118                Handle::MultiThread(h) => h.timer_flavor,
119            }
120        }
121
122        #[cfg(all(tokio_unstable, feature = "rt-multi-thread", feature = "time"))]
123        /// Returns true if the runtime is shutting down.
124        pub(crate) fn is_shutdown(&self) -> bool {
125            match self {
126                Handle::CurrentThread(_) => panic!("the alternative timer implementation is not supported on CurrentThread runtime"),
127                Handle::MultiThread(h) => h.is_shutdown(),
128            }
129        }
130
131        #[cfg(all(tokio_unstable, feature = "rt-multi-thread", feature = "time"))]
132        /// Push a timer entry that was created outside of this runtime
133        /// into the runtime-global queue. The pushed timer will be
134        /// processed by a random worker thread.
135        pub(crate) fn push_remote_timer(&self, entry_hdl: crate::runtime::time_alt::EntryHandle) {
136            match self {
137                Handle::CurrentThread(_) => panic!("the alternative timer implementation is not supported on CurrentThread runtime"),
138                Handle::MultiThread(h) => h.push_remote_timer(entry_hdl),
139            }
140        }
141
142        /// Returns true if this is a local runtime and the runtime is owned by the current thread.
143        pub(crate) fn can_spawn_local_on_local_runtime(&self) -> bool {
144            match self {
145                Handle::CurrentThread(h) => h.local_tid.map(|x| std::thread::current().id() == x).unwrap_or(false),
146
147                #[cfg(feature = "rt-multi-thread")]
148                Handle::MultiThread(_) => false,
149            }
150        }
151
152        pub(crate) fn spawn<F>(&self, future: F, id: Id, spawned_at: SpawnLocation) -> JoinHandle<F::Output>
153        where
154            F: Future + Send + 'static,
155            F::Output: Send + 'static,
156        {
157            match self {
158                Handle::CurrentThread(h) => current_thread::Handle::spawn(h, future, id, spawned_at),
159
160                #[cfg(feature = "rt-multi-thread")]
161                Handle::MultiThread(h) => multi_thread::Handle::spawn(h, future, id, spawned_at),
162            }
163        }
164
165        /// Spawn a local task
166        ///
167        /// # Safety
168        ///
169        /// This should only be called in `LocalRuntime` if the runtime has been verified to be owned
170        /// by the current thread.
171        #[allow(irrefutable_let_patterns)]
172        #[track_caller]
173        pub(crate) unsafe fn spawn_local<F>(&self, future: F, id: Id, spawned_at: SpawnLocation) -> JoinHandle<F::Output>
174        where
175            F: Future + 'static,
176            F::Output: 'static,
177        {
178            if let Handle::CurrentThread(h) = self {
179                // Safety: caller guarantees that this is a `LocalRuntime`.
180                unsafe { current_thread::Handle::spawn_local(h, future, id, spawned_at) }
181            } else {
182                panic!("Only current_thread and LocalSet have spawn_local internals implemented")
183            }
184        }
185
186        pub(crate) fn shutdown(&self) {
187            match *self {
188                Handle::CurrentThread(_) => {},
189
190                #[cfg(feature = "rt-multi-thread")]
191                Handle::MultiThread(ref h) => h.shutdown(),
192            }
193        }
194
195        pub(crate) fn seed_generator(&self) -> &RngSeedGenerator {
196            match_flavor!(self, Handle(h) => &h.seed_generator)
197        }
198
199        pub(crate) fn as_current_thread(&self) -> &Arc<current_thread::Handle> {
200            match self {
201                Handle::CurrentThread(handle) => handle,
202                #[cfg(feature = "rt-multi-thread")]
203                _ => panic!("not a CurrentThread handle"),
204            }
205        }
206
207        pub(crate) fn hooks(&self) -> &TaskHooks {
208            match self {
209                Handle::CurrentThread(h) => &h.task_hooks,
210                #[cfg(feature = "rt-multi-thread")]
211                Handle::MultiThread(h) => &h.task_hooks,
212            }
213        }
214    }
215
216    impl Handle {
217        pub(crate) fn num_workers(&self) -> usize {
218            match self {
219                Handle::CurrentThread(_) => 1,
220                #[cfg(feature = "rt-multi-thread")]
221                Handle::MultiThread(handle) => handle.num_workers(),
222            }
223        }
224
225        pub(crate) fn num_alive_tasks(&self) -> usize {
226            match_flavor!(self, Handle(handle) => handle.num_alive_tasks())
227        }
228
229        pub(crate) fn injection_queue_depth(&self) -> usize {
230            match_flavor!(self, Handle(handle) => handle.injection_queue_depth())
231        }
232
233        pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics {
234            match_flavor!(self, Handle(handle) => handle.worker_metrics(worker))
235        }
236    }
237
238    cfg_unstable_metrics! {
239        use crate::runtime::SchedulerMetrics;
240
241        impl Handle {
242            cfg_64bit_metrics! {
243                pub(crate) fn spawned_tasks_count(&self) -> u64 {
244                    match_flavor!(self, Handle(handle) => handle.spawned_tasks_count())
245                }
246            }
247
248            pub(crate) fn num_blocking_threads(&self) -> usize {
249                match_flavor!(self, Handle(handle) => handle.num_blocking_threads())
250            }
251
252            pub(crate) fn num_idle_blocking_threads(&self) -> usize {
253                match_flavor!(self, Handle(handle) => handle.num_idle_blocking_threads())
254            }
255
256            pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics {
257                match_flavor!(self, Handle(handle) => handle.scheduler_metrics())
258            }
259
260            pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize {
261                match_flavor!(self, Handle(handle) => handle.worker_local_queue_depth(worker))
262            }
263
264            pub(crate) fn blocking_queue_depth(&self) -> usize {
265                match_flavor!(self, Handle(handle) => handle.blocking_queue_depth())
266            }
267        }
268    }
269
270    impl Context {
271        #[track_caller]
272        pub(crate) fn expect_current_thread(&self) -> &current_thread::Context {
273            match self {
274                Context::CurrentThread(context) => context,
275                #[cfg(feature = "rt-multi-thread")]
276                _ => panic!("expected `CurrentThread::Context`")
277            }
278        }
279
280        pub(crate) fn defer(&self, waker: &Waker) {
281            match_flavor!(self, Context(context) => context.defer(waker));
282        }
283
284        #[cfg(tokio_unstable)]
285        pub(crate) fn worker_index(&self) -> Option<usize> {
286            match self {
287                Context::CurrentThread(_) => Some(0),
288                #[cfg(feature = "rt-multi-thread")]
289                Context::MultiThread(context) => Some(context.worker_index()),
290            }
291        }
292
293        cfg_rt_multi_thread! {
294            #[track_caller]
295            pub(crate) fn expect_multi_thread(&self) -> &multi_thread::Context {
296                match self {
297                    Context::MultiThread(context) => context,
298                    _ => panic!("expected `MultiThread::Context`")
299                }
300            }
301        }
302    }
303}
304
305cfg_not_rt! {
306    #[cfg(any(
307        feature = "net",
308        all(unix, feature = "process"),
309        all(unix, feature = "signal"),
310        feature = "time",
311    ))]
312    impl Handle {
313        #[track_caller]
314        pub(crate) fn current() -> Handle {
315            panic!("{}", crate::util::error::CONTEXT_MISSING_ERROR)
316        }
317
318        #[cfg_attr(not(feature = "time"), allow(dead_code))]
319        #[track_caller]
320        pub(crate) fn timer_flavor(&self) -> crate::runtime::TimerFlavor {
321            panic!("{}", crate::util::error::CONTEXT_MISSING_ERROR)
322        }
323    }
324}
tokio/runtime/scheduler/mod.rs

tokio/runtime/scheduler/
mod.rs
