1#![warn(rust_2018_idioms)]
2#![allow(clippy::declare_interior_mutable_const)]
3#![cfg(all(feature = "full", tokio_unstable))]
4
5#[cfg(not(target_os = "wasi"))]
6use std::error::Error;
7use std::future::Future;
8use std::pin::Pin;
9use std::task::{Context, Poll};
10#[cfg(not(target_os = "wasi"))]
11use tokio::runtime::{Builder, Runtime};
12use tokio::sync::oneshot;
13use tokio::task::{self, Id, LocalSet};
14
15#[cfg(not(target_os = "wasi"))]
16mod support {
17 pub mod panic;
18}
19#[cfg(not(target_os = "wasi"))]
20use support::panic::test_panic;
21
22#[tokio::test(flavor = "current_thread")]
23async fn task_id_spawn() {
24 tokio::spawn(async { println!("task id: {}", task::id()) })
25 .await
26 .unwrap();
27}
28
29#[cfg(not(target_os = "wasi"))]
30#[tokio::test(flavor = "current_thread")]
31async fn task_id_spawn_blocking() {
32 task::spawn_blocking(|| println!("task id: {}", task::id()))
33 .await
34 .unwrap();
35}
36
37#[tokio::test(flavor = "current_thread")]
38async fn task_id_collision_current_thread() {
39 let handle1 = tokio::spawn(async { task::id() });
40 let handle2 = tokio::spawn(async { task::id() });
41
42 let (id1, id2) = tokio::join!(handle1, handle2);
43 assert_ne!(id1.unwrap(), id2.unwrap());
44}
45
46#[cfg(not(target_os = "wasi"))]
47#[tokio::test(flavor = "multi_thread")]
48async fn task_id_collision_multi_thread() {
49 let handle1 = tokio::spawn(async { task::id() });
50 let handle2 = tokio::spawn(async { task::id() });
51
52 let (id1, id2) = tokio::join!(handle1, handle2);
53 assert_ne!(id1.unwrap(), id2.unwrap());
54}
55
56#[tokio::test(flavor = "current_thread")]
57async fn task_ids_match_current_thread() {
58 let (tx, rx) = oneshot::channel();
59 let handle = tokio::spawn(async {
60 let id = rx.await.unwrap();
61 assert_eq!(id, task::id());
62 });
63 tx.send(handle.id()).unwrap();
64 handle.await.unwrap();
65}
66
67#[cfg(not(target_os = "wasi"))]
68#[tokio::test(flavor = "multi_thread")]
69async fn task_ids_match_multi_thread() {
70 let (tx, rx) = oneshot::channel();
71 let handle = tokio::spawn(async {
72 let id = rx.await.unwrap();
73 assert_eq!(id, task::id());
74 });
75 tx.send(handle.id()).unwrap();
76 handle.await.unwrap();
77}
78
79#[cfg(not(target_os = "wasi"))]
80#[tokio::test(flavor = "multi_thread")]
81async fn task_id_future_destructor_completion() {
82 struct MyFuture {
83 tx: Option<oneshot::Sender<Id>>,
84 }
85
86 impl Future for MyFuture {
87 type Output = ();
88
89 fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
90 Poll::Ready(())
91 }
92 }
93
94 impl Drop for MyFuture {
95 fn drop(&mut self) {
96 let _ = self.tx.take().unwrap().send(task::id());
97 }
98 }
99
100 let (tx, rx) = oneshot::channel();
101 let handle = tokio::spawn(MyFuture { tx: Some(tx) });
102 let id = handle.id();
103 handle.await.unwrap();
104 assert_eq!(rx.await.unwrap(), id);
105}
106
107#[cfg(not(target_os = "wasi"))]
108#[tokio::test(flavor = "multi_thread")]
109async fn task_id_future_destructor_abort() {
110 struct MyFuture {
111 tx: Option<oneshot::Sender<Id>>,
112 }
113
114 impl Future for MyFuture {
115 type Output = ();
116
117 fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
118 Poll::Pending
119 }
120 }
121 impl Drop for MyFuture {
122 fn drop(&mut self) {
123 let _ = self.tx.take().unwrap().send(task::id());
124 }
125 }
126
127 let (tx, rx) = oneshot::channel();
128 let handle = tokio::spawn(MyFuture { tx: Some(tx) });
129 let id = handle.id();
130 handle.abort();
131 assert!(handle.await.unwrap_err().is_cancelled());
132 assert_eq!(rx.await.unwrap(), id);
133}
134
135#[tokio::test(flavor = "current_thread")]
136async fn task_id_output_destructor_handle_dropped_before_completion() {
137 struct MyOutput {
138 tx: Option<oneshot::Sender<Id>>,
139 }
140
141 impl Drop for MyOutput {
142 fn drop(&mut self) {
143 let _ = self.tx.take().unwrap().send(task::id());
144 }
145 }
146
147 struct MyFuture {
148 tx: Option<oneshot::Sender<Id>>,
149 }
150
151 impl Future for MyFuture {
152 type Output = MyOutput;
153
154 fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
155 Poll::Ready(MyOutput { tx: self.tx.take() })
156 }
157 }
158
159 let (tx, mut rx) = oneshot::channel();
160 let handle = tokio::spawn(MyFuture { tx: Some(tx) });
161 let id = handle.id();
162 drop(handle);
163 assert!(rx.try_recv().is_err());
164 assert_eq!(rx.await.unwrap(), id);
165}
166
167#[tokio::test(flavor = "current_thread")]
168async fn task_id_output_destructor_handle_dropped_after_completion() {
169 struct MyOutput {
170 tx: Option<oneshot::Sender<Id>>,
171 }
172
173 impl Drop for MyOutput {
174 fn drop(&mut self) {
175 let _ = self.tx.take().unwrap().send(task::id());
176 }
177 }
178
179 struct MyFuture {
180 tx_output: Option<oneshot::Sender<Id>>,
181 tx_future: Option<oneshot::Sender<()>>,
182 }
183
184 impl Future for MyFuture {
185 type Output = MyOutput;
186
187 fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
188 let _ = self.tx_future.take().unwrap().send(());
189 Poll::Ready(MyOutput {
190 tx: self.tx_output.take(),
191 })
192 }
193 }
194
195 let (tx_output, mut rx_output) = oneshot::channel();
196 let (tx_future, rx_future) = oneshot::channel();
197 let handle = tokio::spawn(MyFuture {
198 tx_output: Some(tx_output),
199 tx_future: Some(tx_future),
200 });
201 let id = handle.id();
202 rx_future.await.unwrap();
203 assert!(rx_output.try_recv().is_err());
204 drop(handle);
205 assert_eq!(rx_output.await.unwrap(), id);
206}
207
208#[test]
209fn task_try_id_outside_task() {
210 assert_eq!(None, task::try_id());
211}
212
213#[cfg(not(target_os = "wasi"))]
214#[test]
215fn task_try_id_inside_block_on() {
216 let rt = Runtime::new().unwrap();
217 rt.block_on(async {
218 assert_eq!(None, task::try_id());
219 });
220}
221
222#[tokio::test(flavor = "current_thread")]
223async fn task_id_spawn_local() {
224 LocalSet::new()
225 .run_until(async {
226 task::spawn_local(async { println!("task id: {}", task::id()) })
227 .await
228 .unwrap();
229 })
230 .await
231}
232
233#[tokio::test(flavor = "current_thread")]
234async fn task_id_nested_spawn_local() {
235 LocalSet::new()
236 .run_until(async {
237 task::spawn_local(async {
238 let parent_id = task::id();
239 LocalSet::new()
240 .run_until(async {
241 task::spawn_local(async move {
242 assert_ne!(parent_id, task::id());
243 })
244 .await
245 .unwrap();
246 })
247 .await;
248 assert_eq!(parent_id, task::id());
249 })
250 .await
251 .unwrap();
252 })
253 .await;
254}
255
256#[cfg(not(target_os = "wasi"))]
257#[tokio::test(flavor = "multi_thread")]
258async fn task_id_block_in_place_block_on_spawn() {
259 task::spawn(async {
260 let parent_id = task::id();
261
262 task::block_in_place(move || {
263 let rt = Builder::new_current_thread().build().unwrap();
264 rt.block_on(rt.spawn(async move {
265 assert_ne!(parent_id, task::id());
266 }))
267 .unwrap();
268 });
269
270 assert_eq!(parent_id, task::id());
271 })
272 .await
273 .unwrap();
274}
275
276#[cfg(not(target_os = "wasi"))]
277#[test]
278fn task_id_outside_task_panic_caller() -> Result<(), Box<dyn Error>> {
279 let panic_location_file = test_panic(|| {
280 let _ = task::id();
281 });
282
283 // The panic location should be in this file
284 assert_eq!(&panic_location_file.unwrap(), file!());
285
286 Ok(())
287}
288
289#[cfg(not(target_os = "wasi"))]
290#[test]
291fn task_id_inside_block_on_panic_caller() -> Result<(), Box<dyn Error>> {
292 let panic_location_file = test_panic(|| {
293 let rt = Runtime::new().unwrap();
294 rt.block_on(async {
295 task::id();
296 });
297 });
298
299 // The panic location should be in this file
300 assert_eq!(&panic_location_file.unwrap(), file!());
301
302 Ok(())
303}
304