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