1 #![allow(unknown_lints, unexpected_cfgs)]
2 #![cfg(all(
3     tokio_unstable,
4     tokio_taskdump,
5     target_os = "linux",
6     any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
7 ))]
8 
9 use std::hint::black_box;
10 use tokio::runtime::{self, Handle};
11 
12 #[inline(never)]
a()13 async fn a() {
14     black_box(b()).await
15 }
16 
17 #[inline(never)]
b()18 async fn b() {
19     black_box(c()).await
20 }
21 
22 #[inline(never)]
c()23 async fn c() {
24     loop {
25         black_box(tokio::task::yield_now()).await
26     }
27 }
28 
29 #[test]
current_thread()30 fn current_thread() {
31     let rt = runtime::Builder::new_current_thread()
32         .enable_all()
33         .build()
34         .unwrap();
35 
36     async fn dump() {
37         let handle = Handle::current();
38         let dump = handle.dump().await;
39 
40         let tasks: Vec<_> = dump.tasks().iter().collect();
41 
42         assert_eq!(tasks.len(), 3);
43 
44         for task in tasks {
45             let id = task.id();
46             let trace = task.trace().to_string();
47             eprintln!("\n\n{id}:\n{trace}\n\n");
48             assert!(trace.contains("dump::a"));
49             assert!(trace.contains("dump::b"));
50             assert!(trace.contains("dump::c"));
51             assert!(trace.contains("tokio::task::yield_now"));
52         }
53     }
54 
55     rt.block_on(async {
56         tokio::select!(
57             biased;
58             _ = tokio::spawn(a()) => {},
59             _ = tokio::spawn(a()) => {},
60             _ = tokio::spawn(a()) => {},
61             _ = dump() => {},
62         );
63     });
64 }
65 
66 #[test]
multi_thread()67 fn multi_thread() {
68     let rt = runtime::Builder::new_multi_thread()
69         .enable_all()
70         .worker_threads(3)
71         .build()
72         .unwrap();
73 
74     async fn dump() {
75         let handle = Handle::current();
76         let dump = handle.dump().await;
77 
78         let tasks: Vec<_> = dump.tasks().iter().collect();
79 
80         assert_eq!(tasks.len(), 3);
81 
82         for task in tasks {
83             let id = task.id();
84             let trace = task.trace().to_string();
85             eprintln!("\n\n{id}:\n{trace}\n\n");
86             assert!(trace.contains("dump::a"));
87             assert!(trace.contains("dump::b"));
88             assert!(trace.contains("dump::c"));
89             assert!(trace.contains("tokio::task::yield_now"));
90         }
91     }
92 
93     rt.block_on(async {
94         tokio::select!(
95             biased;
96             _ = tokio::spawn(a()) => {},
97             _ = tokio::spawn(a()) => {},
98             _ = tokio::spawn(a()) => {},
99             _ = dump() => {},
100         );
101     });
102 }
103 
104 /// Regression tests for #6035.
105 ///
106 /// These tests ensure that dumping will not deadlock if a future completes
107 /// during a trace.
108 mod future_completes_during_trace {
109     use super::*;
110 
111     use core::future::{poll_fn, Future};
112 
113     /// A future that completes only during a trace.
complete_during_trace() -> impl Future<Output = ()> + Send114     fn complete_during_trace() -> impl Future<Output = ()> + Send {
115         use std::task::Poll;
116         poll_fn(|cx| {
117             if Handle::is_tracing() {
118                 Poll::Ready(())
119             } else {
120                 cx.waker().wake_by_ref();
121                 Poll::Pending
122             }
123         })
124     }
125 
126     #[test]
current_thread()127     fn current_thread() {
128         let rt = runtime::Builder::new_current_thread()
129             .enable_all()
130             .build()
131             .unwrap();
132 
133         async fn dump() {
134             let handle = Handle::current();
135             let _dump = handle.dump().await;
136         }
137 
138         rt.block_on(async {
139             let _ = tokio::join!(tokio::spawn(complete_during_trace()), dump());
140         });
141     }
142 
143     #[test]
multi_thread()144     fn multi_thread() {
145         let rt = runtime::Builder::new_multi_thread()
146             .enable_all()
147             .build()
148             .unwrap();
149 
150         async fn dump() {
151             let handle = Handle::current();
152             let _dump = handle.dump().await;
153             tokio::task::yield_now().await;
154         }
155 
156         rt.block_on(async {
157             let _ = tokio::join!(tokio::spawn(complete_during_trace()), dump());
158         });
159     }
160 }
161 
162 /// Regression test for #6051.
163 ///
164 /// This test ensures that tasks notified outside of a worker will not be
165 /// traced, since doing so will un-set their notified bit prior to them being
166 /// run and panic.
167 #[test]
notified_during_tracing()168 fn notified_during_tracing() {
169     let rt = runtime::Builder::new_multi_thread()
170         .enable_all()
171         .worker_threads(3)
172         .build()
173         .unwrap();
174 
175     let timeout = async {
176         tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
177     };
178 
179     let timer = rt.spawn(async {
180         loop {
181             tokio::time::sleep(tokio::time::Duration::from_nanos(1)).await;
182         }
183     });
184 
185     let dump = async {
186         loop {
187             let handle = Handle::current();
188             let _dump = handle.dump().await;
189         }
190     };
191 
192     rt.block_on(async {
193         tokio::select!(
194             biased;
195             _ = timeout => {},
196             _ = timer => {},
197             _ = dump => {},
198         );
199     });
200 }
201