arty 0.5.0

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

use std::pin::Pin;
use std::task::{self, Poll};

use arty_executor::remote::{JoinHandle as RemoteJoinHandle, Pool as RemotePool};
use performables::arc::Arc;
use pin_project::pin_project;

use super::JoinError;
use crate::runtime::blocking_worker::is_current_blocking_pool;

/// A handle for receiving an async or blocking task's result.
///
/// Await the handle inside async code. Completion produces `Ok(result)`. A task
/// panic or shutdown cancellation produces [`JoinError`] without unwinding the
/// joining task.
///
/// Dropping the handle does not cancel its task or rethrow a task panic.
/// The runtime must remain running for pending async work to complete.
/// If the task returns its own `Result<T, E>`, joining it produces
/// `Result<Result<T, E>, JoinError>`.
///
/// # Panics
///
/// Panics if polled again after its result has been received. A blocking
/// handle also panics if it is polled directly from a callback running in the
/// same blocking pool.
///
/// # Examples
///
/// ```
/// # #[cfg(all(feature = "macros", feature = "rt"))]
/// # #[arty::main]
/// # async fn main(cx: arty::task::Builtins) -> Result<(), arty::task::JoinError> {
/// let task = cx.scheduler().spawn(async |_| 42);
/// assert_eq!(task.await?, 42);
/// # Ok(())
/// # }
/// # #[cfg(not(all(feature = "macros", feature = "rt")))] fn main() {}
/// ```
#[derive(derive_more::Debug)]
#[pin_project]
pub struct JoinHandle<R>
where
    R: Send + 'static,
{
    #[debug(ignore)]
    #[pin]
    inner: RemoteJoinHandle<R>,
    blocking_pool: Option<Arc<()>>,
    completed: bool,
}

impl<R> JoinHandle<R>
where
    R: Send + 'static,
{
    pub(crate) fn new(inner: RemoteJoinHandle<R>) -> Self {
        Self {
            inner,
            blocking_pool: None,
            completed: false,
        }
    }

    pub(crate) fn shutdown(pool: &RemotePool) -> Self {
        let (sender, inner) = pool.channel();
        drop(sender);
        Self::new(inner)
    }

    pub(crate) fn blocking_pool(mut self, pool: Arc<()>) -> Self {
        self.blocking_pool = Some(pool);
        self
    }

    #[cfg(test)]
    pub(crate) fn join(self) -> Result<R, JoinError> {
        futures::executor::block_on(self)
    }
}

#[cfg_attr(test, mutants::skip)] // Removing this guard makes the black-box blocking-pool regression deadlock.
fn assert_not_current_blocking_pool(blocking_pool: Option<&Arc<()>>) {
    assert!(
        !blocking_pool.is_some_and(is_current_blocking_pool),
        "blocking JoinHandle cannot be polled while the current blocking pool waits for it"
    );
}

impl<R> Future for JoinHandle<R>
where
    R: Send + 'static,
{
    type Output = Result<R, JoinError>;

    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        assert!(!*this.completed, "JoinHandle polled after completion");
        assert_not_current_blocking_pool(this.blocking_pool.as_ref());
        match this.inner.poll(cx) {
            Poll::Ready(result) => {
                *this.completed = true;
                Poll::Ready(result.map_err(JoinError::from_executor))
            }
            Poll::Pending => Poll::Pending,
        }
    }
}