compio_executor/
join_handle.rs

1use std::{
2    error::Error,
3    fmt::Display,
4    io,
5    marker::PhantomData,
6    mem::ManuallyDrop,
7    panic::resume_unwind,
8    pin::Pin,
9    ptr,
10    task::{Context, Poll},
11};
12
13use compio_log::{instrument, trace};
14
15use crate::{Panic, task::Task};
16
17/// A handle that awaits the result of a task.
18///
19/// Dropping a [`JoinHandle`] will cancel the task. To run the task in the
20/// background, use [`JoinHandle::detach`].
21#[must_use = "Drop `JoinHandle` will cancel the task. Use `detach` to run it in background."]
22#[derive(Debug)]
23#[repr(transparent)]
24pub struct JoinHandle<T> {
25    task: Option<Task>,
26    _marker: PhantomData<T>,
27}
28
29/// If T is send, we can poll result from other thread
30unsafe impl<T: Send> Send for JoinHandle<T> {}
31
32/// JoinHandle does not expose any &self interface, so it's unconditionally
33/// Sync.
34unsafe impl<T> Sync for JoinHandle<T> {}
35
36impl<T> Unpin for JoinHandle<T> {}
37
38impl<T> JoinHandle<T> {
39    pub(crate) fn new(task: Task) -> Self {
40        Self {
41            task: Some(task),
42            _marker: PhantomData,
43        }
44    }
45
46    /// Cancel the task and wait for the result, if any.
47    pub async fn cancel(self) -> Option<T> {
48        self.task.as_ref()?.cancel(false);
49        self.await.ok()
50    }
51
52    /// Detach the task to let it run in the background.
53    pub fn detach(self) {
54        unsafe { ptr::drop_in_place(&raw mut ManuallyDrop::new(self).task) };
55    }
56
57    /// Returns true if this JoinHandle has been finalized, either because
58    /// it's been completed or canceled.
59    pub fn is_finished(&self) -> bool {
60        match &self.task {
61            Some(task) => task.is_finished(),
62            None => true,
63        }
64    }
65}
66
67/// Task failed to execute to completion.
68#[derive(Debug)]
69pub enum JoinError {
70    /// The task was cancelled.
71    Cancelled,
72    /// The task panicked.
73    Panicked(Panic),
74}
75
76/// Trait to resume unwind from a [`JoinError`].
77pub trait ResumeUnwind {
78    /// The output type.
79    type Output;
80
81    /// Resume the panic if the task panicked.
82    fn resume_unwind(self) -> Self::Output;
83}
84
85impl<T> ResumeUnwind for Result<T, JoinError> {
86    type Output = Option<T>;
87
88    fn resume_unwind(self) -> Self::Output {
89        match self {
90            Ok(res) => Some(res),
91            Err(JoinError::Cancelled) => None,
92            Err(JoinError::Panicked(e)) => resume_unwind(e),
93        }
94    }
95}
96
97impl JoinError {
98    /// Resume unwind if the task panicked, otherwise do nothing.
99    pub fn resume_unwind(self) {
100        if let JoinError::Panicked(e) = self {
101            resume_unwind(e)
102        }
103    }
104}
105
106impl Display for JoinError {
107    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
108        match self {
109            JoinError::Cancelled => write!(f, "Task was cancelled"),
110            JoinError::Panicked(_) => write!(f, "Task has panicked"),
111        }
112    }
113}
114
115impl Error for JoinError {}
116
117impl From<JoinError> for io::Error {
118    fn from(e: JoinError) -> Self {
119        match e {
120            JoinError::Cancelled => io::Error::other("Task was cancelled"),
121            JoinError::Panicked(_) => io::Error::other("Task has panicked"),
122        }
123    }
124}
125
126impl<T> Future for JoinHandle<T> {
127    type Output = Result<T, JoinError>;
128
129    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
130        instrument!(compio_log::Level::TRACE, "JoinHandle::poll");
131
132        let task = self.task.as_ref().expect("Cannot poll after completion");
133
134        unsafe { task.poll(cx) }.map(|res| {
135            trace!("Poll ready");
136
137            self.task = None;
138
139            match res {
140                Some(Ok(res)) => Ok(res),
141                Some(Err(e)) => Err(JoinError::Panicked(e)),
142                None => Err(JoinError::Cancelled),
143            }
144        })
145    }
146}
147
148impl<T> Drop for JoinHandle<T> {
149    fn drop(&mut self) {
150        instrument!(compio_log::Level::TRACE, "JoinHandle::drop");
151
152        if let Some(task) = self.task.as_ref() {
153            task.cancel(true);
154        }
155    }
156}
compio_executor/join_handle.rs

compio_executor/
join_handle.rs
