1use std::future::Future;
10use std::pin::Pin;
11use std::sync::LazyLock;
12use std::task::{Context, Poll, ready};
13
14use tokio::runtime::{Builder, Runtime};
15use tokio::task::JoinHandle;
16
17#[cfg(unix)]
18pub mod asyncfd;
19mod bufpool;
20pub mod immortal;
21pub mod reaper;
22mod timeout;
23
24pub use bufpool::{pooled_read, pooled_read_callback};
25pub use immortal::{Immortal, RespawnStrategy};
26pub use reaper::TaskReaper;
27pub use timeout::TimeoutExt;
28
29pub use tokio::time::{sleep, sleep_until};
31
32static RUNTIME: LazyLock<Runtime> = LazyLock::new(|| {
38 Builder::new_multi_thread()
39 .enable_all()
40 .build()
41 .expect("could not build the global tokio runtime")
42});
43
44pub fn handle() -> tokio::runtime::Handle {
46 RUNTIME.handle().clone()
47}
48
49pub fn spawn<F>(future: F) -> Task<F::Output>
51where
52 F: Future + Send + 'static,
53 F::Output: Send + 'static,
54{
55 Task(Some(RUNTIME.spawn(future)))
56}
57
58pub fn spawn_blocking<F, R>(f: F) -> Task<R>
60where
61 F: FnOnce() -> R + Send + 'static,
62 R: Send + 'static,
63{
64 Task(Some(RUNTIME.spawn_blocking(f)))
65}
66
67pub fn block_on<F: Future>(future: F) -> F::Output {
73 RUNTIME.block_on(future)
74}
75
76#[must_use = "dropping a Task cancels the task; use .detach() to let it run in the background"]
83pub struct Task<T>(Option<JoinHandle<T>>);
84
85impl<T> Task<T> {
86 pub fn detach(mut self) {
89 self.0.take();
90 }
91
92 pub async fn cancel(mut self) {
97 if let Some(handle) = self.0.take() {
98 handle.abort();
99 let _ = handle.await;
100 }
101 }
102}
103
104impl<T> Drop for Task<T> {
105 fn drop(&mut self) {
106 if let Some(handle) = &self.0 {
107 handle.abort();
108 }
109 }
110}
111
112impl<T> Future for Task<T> {
113 type Output = T;
114 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
115 let handle = self
116 .0
117 .as_mut()
118 .expect("Task polled after being detached or cancelled");
119 match ready!(Pin::new(handle).poll(cx)) {
120 Ok(value) => Poll::Ready(value),
121 Err(err) if err.is_panic() => std::panic::resume_unwind(err.into_panic()),
122 Err(_) => Poll::Pending,
127 }
128 }
129}
130
131#[cfg(test)]
132mod tests {
133 use super::*;
134 use std::sync::Arc;
135 use std::sync::atomic::{AtomicBool, Ordering};
136 use std::time::Duration;
137
138 #[test]
139 fn drop_cancels_the_task() {
140 block_on(async {
141 let flag = Arc::new(AtomicBool::new(false));
142 let f2 = flag.clone();
143 let task = spawn(async move {
144 tokio::time::sleep(Duration::from_millis(100)).await;
145 f2.store(true, Ordering::SeqCst);
146 });
147 drop(task);
148 tokio::time::sleep(Duration::from_millis(300)).await;
149 assert!(
150 !flag.load(Ordering::SeqCst),
151 "a dropped Task must be cancelled"
152 );
153 });
154 }
155
156 #[test]
157 fn detach_runs_to_completion() {
158 block_on(async {
159 let flag = Arc::new(AtomicBool::new(false));
160 let f2 = flag.clone();
161 spawn(async move {
162 tokio::time::sleep(Duration::from_millis(50)).await;
163 f2.store(true, Ordering::SeqCst);
164 })
165 .detach();
166 tokio::time::sleep(Duration::from_millis(300)).await;
167 assert!(
168 flag.load(Ordering::SeqCst),
169 "a detached Task must run to completion"
170 );
171 });
172 }
173
174 #[test]
175 fn await_yields_output() {
176 block_on(async {
177 let task = spawn(async { 42u32 });
178 assert_eq!(task.await, 42);
179 });
180 }
181
182 #[test]
183 fn timeout_returns_none_on_elapse() {
184 block_on(async {
185 let r = tokio::time::sleep(Duration::from_secs(10))
186 .timeout(Duration::from_millis(50))
187 .await;
188 assert!(r.is_none());
189 });
190 }
191
192 #[test]
193 fn reaper_cancels_on_drop() {
194 block_on(async {
195 let flag = Arc::new(AtomicBool::new(false));
196 let f2 = flag.clone();
197 let reaper = TaskReaper::new();
198 reaper.attach(spawn(async move {
199 tokio::time::sleep(Duration::from_millis(100)).await;
200 f2.store(true, Ordering::SeqCst);
201 }));
202 tokio::time::sleep(Duration::from_millis(10)).await;
204 drop(reaper);
205 tokio::time::sleep(Duration::from_millis(300)).await;
206 assert!(
207 !flag.load(Ordering::SeqCst),
208 "dropping the reaper must cancel attached tasks"
209 );
210 });
211 }
212}