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