monoio/task/
join.rs

1use std::{
2    future::Future,
3    marker::PhantomData,
4    pin::Pin,
5    task::{Context, Poll},
6};
7
8use super::raw::RawTask;
9
10/// JoinHandle can be used to wait task finished.
11/// Note if you drop it directly, task will not be terminated.
12pub struct JoinHandle<T> {
13    raw: RawTask,
14    _p: PhantomData<T>,
15}
16
17unsafe impl<T: Send> Send for JoinHandle<T> {}
18unsafe impl<T: Send> Sync for JoinHandle<T> {}
19
20impl<T> JoinHandle<T> {
21    pub(super) fn new(raw: RawTask) -> JoinHandle<T> {
22        JoinHandle {
23            raw,
24            _p: PhantomData,
25        }
26    }
27
28    /// Checks if the task associated with this `JoinHandle` has finished.
29    pub fn is_finished(&self) -> bool {
30        let state = self.raw.header().state.load();
31        state.is_complete()
32    }
33}
34
35impl<T> Unpin for JoinHandle<T> {}
36
37impl<T> Future for JoinHandle<T> {
38    type Output = T;
39
40    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
41        let mut ret = Poll::Pending;
42
43        // Try to read the task output. If the task is not yet complete, the
44        // waker is stored and is notified once the task does complete.
45        //
46        // The function must go via the vtable, which requires erasing generic
47        // types. To do this, the function "return" is placed on the stack
48        // **before** calling the function and is passed into the function using
49        // `*mut ()`.
50        //
51        // Safety:
52        //
53        // The type of `T` must match the task's output type.
54        unsafe {
55            self.raw
56                .try_read_output(&mut ret as *mut _ as *mut (), cx.waker());
57        }
58        ret
59    }
60}
61
62impl<T> Drop for JoinHandle<T> {
63    fn drop(&mut self) {
64        if self.raw.header().state.drop_join_handle_fast().is_ok() {
65            return;
66        }
67
68        self.raw.drop_join_handle_slow();
69    }
70}