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
12cfg_rt_multi_thread! {
13 mod block_in_place;
14 pub(crate) use block_in_place::block_in_place;
15
16 mod lock;
17 use lock::Lock;
18
19 pub(crate) mod multi_thread;
20 pub(crate) use multi_thread::MultiThread;
21
22 cfg_unstable! {
23 pub(crate) mod multi_thread_alt;
24 pub(crate) use multi_thread_alt::MultiThread as MultiThreadAlt;
25 }
26}
27
28use crate::runtime::driver;
29
30#[derive(Debug, Clone)]
31pub(crate) enum Handle {
32 #[cfg(feature = "rt")]
33 CurrentThread(Arc<current_thread::Handle>),
34
35 #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))]
36 MultiThread(Arc<multi_thread::Handle>),
37
38 #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))]
39 MultiThreadAlt(Arc<multi_thread_alt::Handle>),
40
41 // TODO: This is to avoid triggering "dead code" warnings many other places
42 // in the codebase. Remove this during a later cleanup
43 #[cfg(not(feature = "rt"))]
44 #[allow(dead_code)]
45 Disabled,
46}
47
48#[cfg(feature = "rt")]
49pub(super) enum Context {
50 CurrentThread(current_thread::Context),
51
52 #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))]
53 MultiThread(multi_thread::Context),
54
55 #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))]
56 MultiThreadAlt(multi_thread_alt::Context),
57}
58
59impl Handle {
60 #[cfg_attr(not(feature = "full"), allow(dead_code))]
61 pub(crate) fn driver(&self) -> &driver::Handle {
62 match *self {
63 #[cfg(feature = "rt")]
64 Handle::CurrentThread(ref h: &Arc) => &h.driver,
65
66 #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))]
67 Handle::MultiThread(ref h: &Arc) => &h.driver,
68
69 #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))]
70 Handle::MultiThreadAlt(ref h) => &h.driver,
71
72 #[cfg(not(feature = "rt"))]
73 Handle::Disabled => unreachable!(),
74 }
75 }
76}
77
78cfg_rt! {
79 use crate::future::Future;
80 use crate::loom::sync::Arc;
81 use crate::runtime::{blocking, task::Id};
82 use crate::runtime::context;
83 use crate::task::JoinHandle;
84 use crate::util::RngSeedGenerator;
85 use std::task::Waker;
86
87 macro_rules! match_flavor {
88 ($self:expr, $ty:ident($h:ident) => $e:expr) => {
89 match $self {
90 $ty::CurrentThread($h) => $e,
91
92 #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))]
93 $ty::MultiThread($h) => $e,
94
95 #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))]
96 $ty::MultiThreadAlt($h) => $e,
97 }
98 }
99 }
100
101 impl Handle {
102 #[track_caller]
103 pub(crate) fn current() -> Handle {
104 match context::with_current(Clone::clone) {
105 Ok(handle) => handle,
106 Err(e) => panic!("{}", e),
107 }
108 }
109
110 pub(crate) fn blocking_spawner(&self) -> &blocking::Spawner {
111 match_flavor!(self, Handle(h) => &h.blocking_spawner)
112 }
113
114 pub(crate) fn spawn<F>(&self, future: F, id: Id) -> JoinHandle<F::Output>
115 where
116 F: Future + Send + 'static,
117 F::Output: Send + 'static,
118 {
119 match self {
120 Handle::CurrentThread(h) => current_thread::Handle::spawn(h, future, id),
121
122 #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))]
123 Handle::MultiThread(h) => multi_thread::Handle::spawn(h, future, id),
124
125 #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))]
126 Handle::MultiThreadAlt(h) => multi_thread_alt::Handle::spawn(h, future, id),
127 }
128 }
129
130 pub(crate) fn shutdown(&self) {
131 match *self {
132 Handle::CurrentThread(_) => {},
133
134 #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))]
135 Handle::MultiThread(ref h) => h.shutdown(),
136
137 #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))]
138 Handle::MultiThreadAlt(ref h) => h.shutdown(),
139 }
140 }
141
142 pub(crate) fn seed_generator(&self) -> &RngSeedGenerator {
143 match_flavor!(self, Handle(h) => &h.seed_generator)
144 }
145
146 pub(crate) fn as_current_thread(&self) -> &Arc<current_thread::Handle> {
147 match self {
148 Handle::CurrentThread(handle) => handle,
149 #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))]
150 _ => panic!("not a CurrentThread handle"),
151 }
152 }
153
154 cfg_rt_multi_thread! {
155 cfg_unstable! {
156 pub(crate) fn expect_multi_thread_alt(&self) -> &Arc<multi_thread_alt::Handle> {
157 match self {
158 Handle::MultiThreadAlt(handle) => handle,
159 _ => panic!("not a `MultiThreadAlt` handle"),
160 }
161 }
162 }
163 }
164 }
165
166 cfg_metrics! {
167 use crate::runtime::{SchedulerMetrics, WorkerMetrics};
168
169 impl Handle {
170 pub(crate) fn num_workers(&self) -> usize {
171 match self {
172 Handle::CurrentThread(_) => 1,
173 #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))]
174 Handle::MultiThread(handle) => handle.num_workers(),
175 #[cfg(all(tokio_unstable, feature = "rt-multi-thread", not(target_os = "wasi")))]
176 Handle::MultiThreadAlt(handle) => handle.num_workers(),
177 }
178 }
179
180 pub(crate) fn num_blocking_threads(&self) -> usize {
181 match_flavor!(self, Handle(handle) => handle.num_blocking_threads())
182 }
183
184 pub(crate) fn num_idle_blocking_threads(&self) -> usize {
185 match_flavor!(self, Handle(handle) => handle.num_idle_blocking_threads())
186 }
187
188 pub(crate) fn active_tasks_count(&self) -> usize {
189 match_flavor!(self, Handle(handle) => handle.active_tasks_count())
190 }
191
192 pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics {
193 match_flavor!(self, Handle(handle) => handle.scheduler_metrics())
194 }
195
196 pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics {
197 match_flavor!(self, Handle(handle) => handle.worker_metrics(worker))
198 }
199
200 pub(crate) fn injection_queue_depth(&self) -> usize {
201 match_flavor!(self, Handle(handle) => handle.injection_queue_depth())
202 }
203
204 pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize {
205 match_flavor!(self, Handle(handle) => handle.worker_local_queue_depth(worker))
206 }
207
208 pub(crate) fn blocking_queue_depth(&self) -> usize {
209 match_flavor!(self, Handle(handle) => handle.blocking_queue_depth())
210 }
211 }
212 }
213
214 impl Context {
215 #[track_caller]
216 pub(crate) fn expect_current_thread(&self) -> &current_thread::Context {
217 match self {
218 Context::CurrentThread(context) => context,
219 #[cfg(all(feature = "rt-multi-thread", not(target_os = "wasi")))]
220 _ => panic!("expected `CurrentThread::Context`")
221 }
222 }
223
224 pub(crate) fn defer(&self, waker: &Waker) {
225 match_flavor!(self, Context(context) => context.defer(waker));
226 }
227
228 cfg_rt_multi_thread! {
229 #[track_caller]
230 pub(crate) fn expect_multi_thread(&self) -> &multi_thread::Context {
231 match self {
232 Context::MultiThread(context) => context,
233 _ => panic!("expected `MultiThread::Context`")
234 }
235 }
236
237 cfg_unstable! {
238 #[track_caller]
239 pub(crate) fn expect_multi_thread_alt(&self) -> &multi_thread_alt::Context {
240 match self {
241 Context::MultiThreadAlt(context) => context,
242 _ => panic!("expected `MultiThreadAlt::Context`")
243 }
244 }
245 }
246 }
247 }
248}
249
250cfg_not_rt! {
251 #[cfg(any(
252 feature = "net",
253 all(unix, feature = "process"),
254 all(unix, feature = "signal"),
255 feature = "time",
256 ))]
257 impl Handle {
258 #[track_caller]
259 pub(crate) fn current() -> Handle {
260 panic!("{}", crate::util::error::CONTEXT_MISSING_ERROR)
261 }
262 }
263}
264