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