join.rs

use crate::runtime::task::RawTask;

use std::fmt;
use std::future::Future;
use std::marker::PhantomData;
use std::pin::Pin;
use std::task::{Context, Poll};

doc_rt_core! {
    /// An owned permission to join on a task (await its termination).
    ///
    /// This can be thought of as the equivalent of [`std::thread::JoinHandle`] for
    /// a task rather than a thread.
    ///
    /// A `JoinHandle` *detaches* the associated task when it is dropped, which
    /// means that there is no longer any handle to the task, and no way to `join`
    /// on it.
    ///
    /// This `struct` is created by the [`task::spawn`] and [`task::spawn_blocking`]
    /// functions.
    ///
    /// # Examples
    ///
    /// Creation from [`task::spawn`]:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn doc() {
    /// let join_handle: task::JoinHandle<_> = task::spawn(async {
    ///     // some work here
    /// });
    /// # }
    /// ```
    ///
    /// Creation from [`task::spawn_blocking`]:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn doc() {
    /// let join_handle: task::JoinHandle<_> = task::spawn_blocking(|| {
    ///     // some blocking work here
    /// });
    /// # }
    /// ```
    ///
    /// Child being detached and outliving its parent:
    ///
    /// ```no_run
    /// use tokio::task;
    /// use tokio::time;
    /// use std::time::Duration;
    ///
    /// # #[tokio::main] async fn main() {
    /// let original_task = task::spawn(async {
    ///     let _detached_task = task::spawn(async {
    ///         // Here we sleep to make sure that the first task returns before.
    ///         time::delay_for(Duration::from_millis(10)).await;
    ///         // This will be called, even though the JoinHandle is dropped.
    ///         println!("♫ Still alive ♫");
    ///     });
    /// });
    ///
    /// original_task.await.expect("The task being joined has panicked");
    /// println!("Original task is joined.");
    ///
    /// // We make sure that the new task has time to run, before the main
    /// // task returns.
    ///
    /// time::delay_for(Duration::from_millis(1000)).await;
    /// # }
    /// ```
    ///
    /// [`task::spawn`]: crate::task::spawn()
    /// [`task::spawn_blocking`]: crate::task::spawn_blocking
    /// [`std::thread::JoinHandle`]: std::thread::JoinHandle
    pub struct JoinHandle<T> {
        raw: RawJoinHandle,
        _p: PhantomData<T>,
    }

    /// A type-erased variant of [`JoinHandle`], these are created by using
    /// [`JoinHandle::into_raw_handle`].
    ///
    /// Raw join handles erase the type information associated with a
    /// [`JoinHandle`], allowing them to easily be stored in containers for
    /// future cancellation.
    ///
    /// They behave exactly the same as regular [`JoinHandle`], except that
    /// instead of resolving to `Result<T, JoinError>`, they resolve to
    /// `Result<(), JoinError>`.
    ///
    /// [`JoinHandle`]: JoinHandle
    /// [`JoinHandle::into_raw_handle`]: JoinHandle::into_raw_handle
    pub struct RawJoinHandle {
        raw: Option<RawTask>,
    }
}

unsafe impl<T: Send> Send for JoinHandle<T> {}
unsafe impl<T: Send> Sync for JoinHandle<T> {}

impl<T> JoinHandle<T> {
    pub(super) fn new(raw: RawTask) -> JoinHandle<T> {
        JoinHandle {
            raw: RawJoinHandle { raw: Some(raw) },
            _p: PhantomData,
        }
    }

    /// Cancel the task associated with the handle.
    ///
    /// Awaiting a cancelled task might complete as usual if the task was
    /// already completed at the time it was cancelled, but most likely it
    /// will complete with a `Err(JoinError::Cancelled)`.
    ///
    /// ```rust
    /// use tokio::time;
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///    let mut handles = Vec::new();
    ///
    ///    handles.push(tokio::spawn(async {
    ///       time::delay_for(time::Duration::from_secs(10)).await;
    ///       true
    ///    }));
    ///
    ///    handles.push(tokio::spawn(async {
    ///       time::delay_for(time::Duration::from_secs(10)).await;
    ///       false
    ///    }));
    ///
    ///    for handle in &handles {
    ///        handle.cancel();
    ///    }
    ///
    ///    for handle in handles {
    ///        assert!(handle.await.unwrap_err().is_cancelled());
    ///    }
    /// }
    /// ```
    pub fn cancel(&self) {
        if let Some(raw) = self.raw.raw {
            raw.shutdown();
        }
    }

    /// Convert the join handle into a raw join handle.
    ///
    /// Raw join handles erase the type information associated with a
    /// [`JoinHandle`], allowing them to easily be stored in containers for
    /// future cancellation.
    ///
    /// [`JoinHandle`]: JoinHandle
    ///
    /// # Examples
    ///
    /// ```rust
    /// use tokio::time;
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///    let mut handles = Vec::new();
    ///
    ///    let handle = tokio::spawn(async {
    ///       time::delay_for(time::Duration::from_secs(10)).await;
    ///       true
    ///    });
    ///
    ///    handles.push(handle.into_raw_handle());
    ///
    ///    let handle = tokio::spawn(async {
    ///       time::delay_for(time::Duration::from_secs(10)).await;
    ///       1u32
    ///    });
    ///
    ///    handles.push(handle.into_raw_handle());
    ///
    ///    for handle in &handles {
    ///        handle.cancel();
    ///    }
    ///
    ///    for handle in handles {
    ///        assert!(handle.await.unwrap_err().is_cancelled());
    ///    }
    /// }
    /// ```
    pub fn into_raw_handle(self) -> RawJoinHandle {
        self.raw
    }
}

impl RawJoinHandle {
    /// Cancel the task associated with the handle.
    ///
    /// Awaiting a cancelled task might complete as usual if the task was
    /// already completed at the time it was cancelled, but most likely it
    /// will complete with a `Err(JoinError::Cancelled)`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use tokio::time;
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///    let mut handles = Vec::new();
    ///
    ///    let handle = tokio::spawn(async {
    ///       time::delay_for(time::Duration::from_secs(10)).await;
    ///       true
    ///    });
    ///
    ///    handles.push(handle.into_raw_handle());
    ///
    ///    let handle = tokio::spawn(async {
    ///       time::delay_for(time::Duration::from_secs(10)).await;
    ///       1u32
    ///    });
    ///
    ///    handles.push(handle.into_raw_handle());
    ///
    ///    for handle in &handles {
    ///        handle.cancel();
    ///    }
    ///
    ///    for handle in handles {
    ///        assert!(handle.await.unwrap_err().is_cancelled());
    ///    }
    /// }
    /// ```
    pub fn cancel(&self) {
        if let Some(raw) = self.raw {
            raw.shutdown();
        }
    }
}

impl Unpin for RawJoinHandle {}

impl Future for RawJoinHandle {
    type Output = super::Result<()>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut ret = Poll::Pending;

        // Keep track of task budget
        ready!(crate::coop::poll_proceed(cx));

        // Raw should always be set. If it is not, this is due to polling after
        // completion
        let raw = self
            .raw
            .as_ref()
            .expect("polling after `JoinHandle` already completed");

        // Try to read the task completion. If the task is not yet complete, the
        // waker is stored and is notified once the task does complete.
        //
        // Safety:
        //
        // `ret` must be of type `Poll<()>`.
        unsafe {
            raw.try_read_completion(&mut ret as *mut _ as *mut (), cx.waker());
        }

        ret
    }
}

impl<T> Future for JoinHandle<T> {
    type Output = super::Result<T>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut ret = Poll::Pending;

        // Keep track of task budget
        let coop = ready!(crate::coop::poll_proceed(cx));

        // Raw should always be set. If it is not, this is due to polling after
        // completion
        let raw = self
            .raw
            .raw
            .as_ref()
            .expect("polling after `JoinHandle` already completed");

        // Try to read the task output. If the task is not yet complete, the
        // waker is stored and is notified once the task does complete.
        //
        // The function must go via the vtable, which requires erasing generic
        // types. To do this, the function "return" is placed on the stack
        // **before** calling the function and is passed into the function using
        // `*mut ()`.
        //
        // Safety:
        //
        // The type of `T` must match the task's output type.
        unsafe {
            raw.try_read_output(&mut ret as *mut _ as *mut (), cx.waker());
        }

        if ret.is_ready() {
            coop.made_progress();
        }

        ret
    }
}

impl Drop for RawJoinHandle {
    fn drop(&mut self) {
        if let Some(raw) = self.raw.take() {
            if raw.header().state.drop_join_handle_fast().is_ok() {
                return;
            }

            raw.drop_join_handle_slow();
        }
    }
}

impl<T> fmt::Debug for JoinHandle<T>
where
    T: fmt::Debug,
{
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_struct("JoinHandle").finish()
    }
}

impl fmt::Debug for RawJoinHandle {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_struct("RawJoinHandle").finish()
    }
}