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

tokio/runtime/scheduler/
mod.rs
