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