tycho_util/futures/
join_task.rs

1use std::pin::Pin;
2use std::task::{Context, Poll};
3
4use futures_util::{Future, FutureExt};
5use tokio::task::JoinHandle;
6
7#[must_use = "futures do nothing unless you `.await` or poll them"]
8pub struct JoinTask<T> {
9    handle: JoinHandle<T>,
10    completed: bool,
11}
12
13impl<T> JoinTask<T> {
14    #[inline]
15    pub fn new<F>(f: F) -> Self
16    where
17        F: Future<Output = T> + Send + 'static,
18        T: Send + 'static,
19    {
20        Self {
21            handle: tokio::spawn(f),
22            completed: false,
23        }
24    }
25
26    pub fn is_finished(&self) -> bool {
27        self.handle.is_finished()
28    }
29}
30
31impl<T> Drop for JoinTask<T> {
32    fn drop(&mut self) {
33        if !self.completed {
34            self.handle.abort();
35        }
36    }
37}
38
39impl<T> Future for JoinTask<T> {
40    type Output = T;
41
42    #[inline]
43    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
44        let res = futures_util::ready!(self.handle.poll_unpin(cx));
45        match res {
46            Ok(value) => {
47                self.completed = true;
48                Poll::Ready(value)
49            }
50            Err(e) => {
51                if e.is_panic() {
52                    std::panic::resume_unwind(e.into_panic());
53                }
54
55                // Can only happen on program termination
56                Poll::Pending
57            }
58        }
59    }
60}
tycho_util/futures/join_task.rs

tycho_util/futures/
join_task.rs
