Skip to main content

ntex_rt/
handle.rs

1use std::{fmt, panic::Location, pin::Pin, task::Context, task::Poll, task::ready};
2
3type Task<T> = crate::rt::task::Task<T>;
4
5/// The location a task was spawned at.
6pub(crate) type SpawnLocation = &'static Location<'static>;
7
8#[derive(Debug)]
9/// A spawned task.
10pub struct JoinHandle<T> {
11    task: Option<Task<T>>,
12}
13
14impl<T> JoinHandle<T> {
15    pub(crate) fn new(task: Task<T>) -> Self {
16        JoinHandle { task: Some(task) }
17    }
18
19    /// Cancels the task.
20    pub fn cancel(mut self) {
21        if let Some(t) = self.task.take() {
22            t.cancel();
23        }
24    }
25
26    /// Detaches the task so it continues running independently.
27    pub fn detach(mut self) {
28        if let Some(t) = self.task.take() {
29            t.detach();
30        }
31    }
32
33    /// Returns `true` if the task has finished.
34    pub fn is_finished(&self) -> bool {
35        match &self.task {
36            Some(fut) => fut.is_finished(),
37            None => true,
38        }
39    }
40}
41
42impl<T> Drop for JoinHandle<T> {
43    fn drop(&mut self) {
44        if let Some(fut) = self.task.take() {
45            fut.detach();
46        }
47    }
48}
49
50impl<T> Future for JoinHandle<T> {
51    type Output = Result<T, JoinError>;
52
53    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
54        Poll::Ready(match self.task.as_mut() {
55            Some(fut) => ready!(fut.poll_join(cx)).ok_or(JoinError),
56            None => Err(JoinError),
57        })
58    }
59}
60
61/// Error returned when a task handle no longer contains a task.
62#[derive(Debug, Copy, Clone)]
63pub struct JoinError;
64
65impl fmt::Display for JoinError {
66    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
67        write!(f, "JoinError")
68    }
69}
70
71impl std::error::Error for JoinError {}
72
73#[cfg(all(test, not(feature = "tokio"), not(feature = "compio")))]
74mod tests {
75    use std::future::{pending, poll_fn};
76
77    use crate::{Handle, System, testing::TestRunner};
78
79    #[test]
80    fn join_handle() {
81        System::new("test", TestRunner).block_on(async {
82            let hnd = crate::spawn(pending::<()>());
83            assert!(!hnd.is_finished());
84            hnd.cancel();
85
86            assert_eq!(crate::spawn(async { 1 }).await.unwrap(), 1);
87
88            let hnd = crate::spawn(async {});
89            poll_fn(|cx| {
90                if hnd.is_finished() {
91                    std::task::Poll::Ready(())
92                } else {
93                    cx.waker().wake_by_ref();
94                    std::task::Poll::Pending
95                }
96            })
97            .await;
98
99            let hnd = Handle::current().clone();
100            hnd.notify().unwrap();
101        });
102        assert_eq!(super::JoinError.to_string(), "JoinError");
103        assert_eq!(crate::rt_default::JoinError.to_string(), "JoinError");
104    }
105}