est 0.10.1

#![allow(rustdoc::private_intra_doc_links)]

//! Asynchronous tasks that can be gracefully shutdown.
//!
//! Graceful shutdown is not a hard termination, it is only responsible for notifying your `Future`
//! that it needs to exit. If your `Future` never quits, the whole shutdown process will also be
//! hung up.
//!
//! The main entry type of this module is [`GracefulTask`].

use super::TaskId;
use crate::{
    future::IntoFutureWithArgs,
    sync::once::{OnceTrigger, once_event},
};
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
use std::{
    marker::PhantomData,
    pin::Pin,
    sync::{Arc, Mutex},
    task::{Context, Poll},
};
use tokio::{
    sync::watch::{Receiver, channel},
    task::{JoinError, JoinHandle},
};

#[cfg(feature = "signal")]
use tokio::signal::ctrl_c;

/// Trigger kind of graceful shutdown (triggered by `ctrl-c` or explicit call).
#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum GracefulKind {
    /// Triggered by `ctrl-c` signal.
    CtrlC,
    /// Triggered by explicit call.
    Explicit,
}

/// Finish mode of the task (whether it is completed completely and exited, or terminated by
/// graceful shutdown).
#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum FinishMode {
    /// Task completed completely and exited.
    Complete,
    /// Task terminated by graceful shutdown.
    Shutdown(GracefulKind),
}

/// Output of the task.
#[derive(Debug)]
pub struct TaskOutput<T> {
    /// Finish mode of the task.
    pub finish_mode: FinishMode,
    /// Join result of the inner task (which spawns the `Future` specified by the user).
    ///
    /// The meaning of this field is similar to that of the `Output` of
    /// [`tokio::task::JoinHandle`].
    pub join_result: Result<T, JoinError>,
}

/// Trigger handle for graceful shutdown (can be `Clone`d).
#[derive(Debug, Clone)]
pub struct ShutdownTrigger(Arc<Mutex<Option<OnceTrigger>>>);

impl ShutdownTrigger {
    /// Trigger the graceful shutdown.
    ///
    /// Returns true if triggered successfully. If it has been triggered previously (including
    /// triggered by `ctrl-c`), or the task has been completed (whether naturally completed or
    /// terminated), false is returned.
    pub fn trigger(&self) -> bool {
        match self.0.try_lock() {
            Err(_) => false,
            Ok(mut guard) => match guard.take() {
                None => false,
                Some(trigger) => trigger.trigger(),
            },
        }
    }
}

/// Receiver handle for graceful shutdown (can be `Clone`d).
///
/// This handle will be passed to the `Future` provided by the user (see the
/// [`GracefulTaskBuilder::spawn`] method). This type provides an async [`recv()`] method and
/// implements [`IntoFuture`]. Therefore, users can call `.recv().await` or directly `.await` to
/// obtain the graceful shutdown signal and its kind. Calling the [`recv()`] method is a better
/// choice because it does not generate additional memory allocation, and its performance will be
/// slightly better than that of [`IntoFuture`]. The implementation of [`IntoFuture`] will
/// [`Box::pin`] the future generated by the compiler.
///
/// [`recv()`]: ShutdownReceiver::recv
///
/// If you call `.recv().await` multiple times, the second and subsequent calls will return
/// immediately and carry the last graceful shutdown signal (which means there will only be one
/// signal).
///
/// Cloning this receiver will be a more intuitive process, that is, the signal will be transmitted
/// to all receivers without loss or delay of the signal. This is useful when subtasks are spawned
/// in the `Future` you provide.
#[derive(Debug, Clone)]
pub struct ShutdownReceiver(RecvInner);

#[derive(Debug, Clone)]
enum RecvInner {
    Pending(Receiver<Option<GracefulKind>>),
    Shutdown(GracefulKind),
}

impl ShutdownReceiver {
    /// Await and receive the graceful shutdown signal.
    pub async fn recv(&mut self) -> GracefulKind {
        match &mut self.0 {
            RecvInner::Pending(receiver) => {
                let init = *receiver.borrow_and_update();
                let kind = match init {
                    Some(kind) => kind,
                    None => {
                        // The `Sender` will never drop before the `Receiver` drops, so
                        // calling `changed()` here will always resolve to `Ok(())`. Therefore,
                        // the next `borrow_and_update()` call must return `Some`, so it can be
                        // unwrapped safely.
                        receiver.changed().await.ok();
                        receiver.borrow_and_update().unwrap()
                    }
                };

                self.0 = RecvInner::Shutdown(kind);
                kind
            }
            RecvInner::Shutdown(kind) => *kind,
        }
    }
}

impl IntoFuture for ShutdownReceiver {
    type Output = GracefulKind;
    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send>>;

    fn into_future(mut self) -> Self::IntoFuture {
        Box::pin(async move { self.recv().await })
    }
}

/// The builder of [`GracefulTask`] (currently only affects whether to respond to `ctrl-c`
/// signals).
///
/// The default value is not to respond to the `ctrl-c` signal.
#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
pub struct GracefulTaskBuilder<T> {
    ctrlc_shutdown: bool,
    _phantom: PhantomData<T>,
}

impl<T> Default for GracefulTaskBuilder<T> {
    fn default() -> Self {
        Self {
            ctrlc_shutdown: false,
            _phantom: PhantomData,
        }
    }
}

impl<T> GracefulTaskBuilder<T> {
    /// Enable `ctrl-c` shutdown.
    #[cfg(feature = "signal")]
    pub fn ctrlc_shutdown(self) -> Self {
        Self {
            ctrlc_shutdown: true,
            ..self
        }
    }

    /// Spawn an asynchronous task that can be gracefully shutdown.
    ///
    /// The parameter `ifwa` can be a closure that returns `Future`, an async closure, an async
    /// function, or a type that implements the [`IntoFutureWithArgs`] trait -- as long as they
    /// hold a parameter of type [`ShutdownReceiver`].
    ///
    /// The returned handle can be used to obtain the ID of the created tasks, trigger graceful
    /// shutdown, and await for the output of the task.
    ///
    /// For more details, see:
    /// - [`GracefulTask`]
    /// - [`IntoFutureWithArgs`]
    /// - [`ShutdownReceiver`]
    /// - [`tokio::spawn`]
    pub fn spawn<I, F>(self, ifwa: I) -> GracefulTask<T>
    where
        I: IntoFutureWithArgs<ShutdownReceiver, F>,
        F: Future<Output = T> + Send + 'static,
        T: Send + 'static,
    {
        self.spawn_ctrlc_mocked(ifwa, async move {
            #[cfg(feature = "signal")]
            ctrl_c().await.ok();
        })
    }

    fn spawn_ctrlc_mocked<I, F, C>(self, ifwa: I, ctrlc: C) -> GracefulTask<T>
    where
        I: IntoFutureWithArgs<ShutdownReceiver, F>,
        F: Future<Output = T> + Send + 'static,
        C: Future<Output = ()> + Send + 'static,
        T: Send + 'static,
    {
        let ctrlc_shutdown = self.ctrlc_shutdown;
        let ctrlc = if ctrlc_shutdown { Some(ctrlc) } else { None };
        let (sender, recver) = channel(None);
        let (trigger, waiter) = once_event();
        let trigger = ShutdownTrigger(Arc::new(Mutex::new(Some(trigger))));
        let mut inner_task =
            tokio::spawn(ifwa.into_future_with_args(ShutdownReceiver(RecvInner::Pending(recver))));

        let inner = inner_task.id().into();
        let graceful = trigger.clone();
        let task = tokio::spawn(async move {
            // This async block will only be polled when `ctrlc` is `Some`, so the unwrap here is
            // safe.
            let ctrlc = async move { ctrlc.unwrap().await };
            let (finish_mode, join_result) = tokio::select! {
                _ = ctrlc, if ctrlc_shutdown => {
                    trigger.trigger();
                    let kind = GracefulKind::CtrlC;
                    sender.send(Some(kind)).ok();
                    (FinishMode::Shutdown(kind), inner_task.await)
                },
                _ = waiter => {
                    let kind = GracefulKind::Explicit;
                    sender.send(Some(kind)).ok();
                    (FinishMode::Shutdown(kind), inner_task.await)
                },
                join_result = &mut inner_task => (FinishMode::Complete, join_result),
            };

            TaskOutput {
                finish_mode,
                join_result,
            }
        });
        let outer = task.id().into();

        GracefulTask {
            inner,
            outer,
            graceful,
            task,
        }
    }
}

/// The handle and entry type of asynchronous tasks that can be shutdown gracefully.
///
/// This type is very similar to [`tokio::task::JoinHandle`], but the difference is that it can
/// gracefully shutdown instead of directly abort.
///
/// Awaiting on it can obtain the [`TaskOutput`].
///
/// # Examples
///
/// ```rust
/// use est::task::{
///     graceful::{FinishMode, GracefulKind, ShutdownReceiver},
///     GracefulTask
/// };
///
/// #[tokio::main]
/// async fn main() {
///     // Task completed normally.
///     let task_output = GracefulTask::builder_default()
///         .spawn(async |_| 42)
///         .await;
///     assert_eq!(task_output.finish_mode, FinishMode::Complete);
///     assert_eq!(task_output.join_result.unwrap(), 42);
///
///     // Task terminated by explicit shutdown.
///     let task_output = GracefulTask::builder_default()
///         .spawn(async |mut shutdown: ShutdownReceiver| shutdown.recv().await)
///         .graceful_shutdown()
///         .await;
///     assert_eq!(task_output.finish_mode, FinishMode::Shutdown(GracefulKind::Explicit));
///     assert_eq!(task_output.join_result.unwrap(), GracefulKind::Explicit);
///
///     // You can separate triggering graceful shutdown and awaiting.
///     let graceful_task = GracefulTask::builder_default()
///         .spawn(async |mut shutdown: ShutdownReceiver| shutdown.recv().await);
///     graceful_task.trigger_graceful_shutdown();
///     let task_output = graceful_task.await;
///     assert_eq!(task_output.finish_mode, FinishMode::Shutdown(GracefulKind::Explicit));
///     assert_eq!(task_output.join_result.unwrap(), GracefulKind::Explicit);
/// }
/// ```
///
/// If you want the `ctrl-c` signal to trigger a graceful shutdown, you need to enable
/// [`ctrlc_shutdown`] in the [`Builder`]:
///
/// [`ctrlc_shutdown`]: GracefulTaskBuilder::ctrlc_shutdown
/// [`Builder`]: GracefulTaskBuilder
///
/// ```ignore
/// use est::task::{
///     graceful::{FinishMode, GracefulKind, ShutdownReceiver},
///     GracefulTask
/// };
///
/// #[tokio::main]
/// async fn main() {
///     // Task terminated by `ctrl-c` signal.
///     // Note: This example will block until `ctrl-c` is pressed.
///     // Note2: You need to enable the `signal` feature flag to run this example.
///     let task_output = GracefulTask::builder_default()
///         .ctrlc_shutdown()
///         .spawn(async |mut shutdown: ShutdownReceiver| shutdown.recv().await)
///         .await;
///     assert_eq!(task_output.finish_mode, FinishMode::Shutdown(GracefulKind::CtrlC));
///     assert_eq!(task_output.join_result.unwrap(), GracefulKind::CtrlC);
/// }
/// ```
#[derive(Debug)]
pub struct GracefulTask<T> {
    inner: TaskId,
    outer: TaskId,
    graceful: ShutdownTrigger,
    task: JoinHandle<TaskOutput<T>>,
}

impl<T> GracefulTask<T> {
    /// Create a default [`GracefulTaskBuilder`].
    pub fn builder_default() -> GracefulTaskBuilder<T> {
        GracefulTaskBuilder::default()
    }

    /// Get the [`TaskId`]s of the two tasks created when calling [`spawn`].
    ///
    /// [`spawn`]: GracefulTaskBuilder::spawn
    ///
    /// The first is the outer task, which is a "monitoring layer" auxiliary task. The second is
    /// the inner task, which is actually the task of executing the `Future` provided by the user.
    pub fn ids(&self) -> (TaskId, TaskId) {
        (self.outer, self.inner)
    }

    /// Trigger the graceful shutdown.
    ///
    /// Returns true if triggered successfully. If it has been triggered previously (including
    /// triggered by `ctrl-c`), or the task has been completed (whether naturally completed or
    /// terminated), false is returned.
    pub fn trigger_graceful_shutdown(&self) -> bool {
        self.graceful.trigger()
    }

    /// Trigger graceful shutdown, and await for the task to complete and obtain the
    /// [`TaskOutput`].
    pub async fn graceful_shutdown(self) -> TaskOutput<T> {
        self.trigger_graceful_shutdown();
        self.await
    }

    /// Returns a new [`ShutdownTrigger`] that can be used to remotely graceful shutdown this task.
    pub fn shutdown_handle(&self) -> ShutdownTrigger {
        self.graceful.clone()
    }

    /// Checks if the tasks associated with this `GracefulTask` have finished.
    pub fn is_finished(&self) -> bool {
        self.task.is_finished()
    }
}

impl<T> Future for GracefulTask<T> {
    type Output = TaskOutput<T>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // `task` will never panic or be aborted, so it can be unwrapped safely.
        Pin::new(&mut self.task).poll(cx).map(Result::unwrap)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    fn ctrlc_mocked() -> (OnceTrigger, impl Future<Output = ()> + Send + 'static) {
        let (trigger, waiter) = once_event();
        (trigger, async move {
            waiter.await;
        })
    }

    async fn sleep() {
        tokio::time::sleep(Duration::from_millis(100)).await;
    }

    async fn sleep_double() {
        tokio::time::sleep(Duration::from_millis(200)).await;
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn graceful_shutdown() {
        let task_output = GracefulTask::builder_default()
            .spawn(async |shutdown| shutdown.await)
            .graceful_shutdown()
            .await;
        assert_eq!(
            task_output.finish_mode,
            FinishMode::Shutdown(GracefulKind::Explicit)
        );
        assert_eq!(task_output.join_result.unwrap(), GracefulKind::Explicit);
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn trigger_graceful_shutdown() {
        let mut graceful_task =
            GracefulTask::builder_default().spawn(async |shutdown| shutdown.await);
        assert!(!graceful_task.is_finished());
        assert!(graceful_task.trigger_graceful_shutdown());
        let task_output = (&mut graceful_task).await;
        assert!(graceful_task.is_finished());
        assert_eq!(
            task_output.finish_mode,
            FinishMode::Shutdown(GracefulKind::Explicit)
        );
        assert_eq!(task_output.join_result.unwrap(), GracefulKind::Explicit);

        let mut graceful_task =
            GracefulTask::builder_default().spawn(async |shutdown| shutdown.await);
        assert!(!graceful_task.is_finished());
        assert!(graceful_task.trigger_graceful_shutdown());
        assert!(!graceful_task.trigger_graceful_shutdown());
        let task_output = (&mut graceful_task).await;
        assert!(graceful_task.is_finished());
        assert_eq!(
            task_output.finish_mode,
            FinishMode::Shutdown(GracefulKind::Explicit)
        );
        assert_eq!(task_output.join_result.unwrap(), GracefulKind::Explicit);

        #[cfg(feature = "signal")]
        {
            let (trigger, ctrlc) = ctrlc_mocked();
            let mut graceful_task = GracefulTask::builder_default()
                .ctrlc_shutdown()
                .spawn_ctrlc_mocked(async |shutdown| shutdown.await, ctrlc);
            assert!(!graceful_task.is_finished());
            assert!(graceful_task.trigger_graceful_shutdown());
            sleep().await;
            assert!(!trigger.trigger());
            let task_output = (&mut graceful_task).await;
            assert!(graceful_task.is_finished());
            assert_eq!(
                task_output.finish_mode,
                FinishMode::Shutdown(GracefulKind::Explicit)
            );
            assert_eq!(task_output.join_result.unwrap(), GracefulKind::Explicit);

            let (trigger, ctrlc) = ctrlc_mocked();
            let mut graceful_task = GracefulTask::builder_default()
                .ctrlc_shutdown()
                .spawn_ctrlc_mocked(async |shutdown| shutdown.await, ctrlc);
            assert!(!graceful_task.is_finished());
            assert!(trigger.trigger());
            sleep().await;
            assert!(!graceful_task.trigger_graceful_shutdown());
            let task_output = (&mut graceful_task).await;
            assert!(graceful_task.is_finished());
            assert_eq!(
                task_output.finish_mode,
                FinishMode::Shutdown(GracefulKind::CtrlC)
            );
            assert_eq!(task_output.join_result.unwrap(), GracefulKind::CtrlC);
        }

        let (trigger, ctrlc) = ctrlc_mocked();
        let mut graceful_task = GracefulTask::builder_default()
            .spawn_ctrlc_mocked(async |shutdown| shutdown.await, ctrlc);
        assert!(!graceful_task.is_finished());
        assert!(!trigger.trigger());
        sleep().await;
        assert!(graceful_task.trigger_graceful_shutdown());
        let task_output = (&mut graceful_task).await;
        assert!(graceful_task.is_finished());
        assert_eq!(
            task_output.finish_mode,
            FinishMode::Shutdown(GracefulKind::Explicit)
        );
        assert_eq!(task_output.join_result.unwrap(), GracefulKind::Explicit);

        let mut graceful_task = GracefulTask::builder_default().spawn(async |_| 42);
        sleep().await;
        assert!(!graceful_task.trigger_graceful_shutdown());
        let task_output = (&mut graceful_task).await;
        assert!(graceful_task.is_finished());
        assert_eq!(task_output.finish_mode, FinishMode::Complete);
        assert_eq!(task_output.join_result.unwrap(), 42);

        let mut graceful_task =
            GracefulTask::builder_default().spawn(async |_| sleep_double().await);
        assert!(!graceful_task.is_finished());
        sleep().await;
        assert!(graceful_task.trigger_graceful_shutdown());
        let task_output = (&mut graceful_task).await;
        assert!(graceful_task.is_finished());
        assert_eq!(
            task_output.finish_mode,
            FinishMode::Shutdown(GracefulKind::Explicit)
        );
        assert_eq!(task_output.join_result.unwrap(), ());
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn shutdown_trigger() {
        let mut graceful_task =
            GracefulTask::builder_default().spawn(async |shutdown| shutdown.await);
        assert!(!graceful_task.is_finished());
        assert!(graceful_task.shutdown_handle().trigger());
        let task_output = (&mut graceful_task).await;
        assert!(graceful_task.is_finished());
        assert_eq!(
            task_output.finish_mode,
            FinishMode::Shutdown(GracefulKind::Explicit)
        );
        assert_eq!(task_output.join_result.unwrap(), GracefulKind::Explicit);

        let mut graceful_task =
            GracefulTask::builder_default().spawn(async |shutdown| shutdown.await);
        assert!(!graceful_task.is_finished());
        let trigger1 = graceful_task.shutdown_handle();
        let trigger2 = graceful_task.shutdown_handle();
        let trigger1 = tokio::spawn(async move { trigger1.trigger() });
        let trigger2 = tokio::spawn(async move { trigger2.trigger() });
        assert!(trigger1.await.unwrap() ^ trigger2.await.unwrap());
        let task_output = (&mut graceful_task).await;
        assert!(graceful_task.is_finished());
        assert_eq!(
            task_output.finish_mode,
            FinishMode::Shutdown(GracefulKind::Explicit)
        );
        assert_eq!(task_output.join_result.unwrap(), GracefulKind::Explicit);

        let mut graceful_task =
            GracefulTask::builder_default().spawn(async |shutdown| shutdown.await);
        assert!(!graceful_task.is_finished());
        assert!(graceful_task.shutdown_handle().trigger());
        assert!(!graceful_task.shutdown_handle().trigger());
        let task_output = (&mut graceful_task).await;
        assert!(graceful_task.is_finished());
        assert_eq!(
            task_output.finish_mode,
            FinishMode::Shutdown(GracefulKind::Explicit)
        );
        assert_eq!(task_output.join_result.unwrap(), GracefulKind::Explicit);

        #[cfg(feature = "signal")]
        {
            let (trigger, ctrlc) = ctrlc_mocked();
            let mut graceful_task = GracefulTask::builder_default()
                .ctrlc_shutdown()
                .spawn_ctrlc_mocked(async |shutdown| shutdown.await, ctrlc);
            assert!(!graceful_task.is_finished());
            assert!(graceful_task.shutdown_handle().trigger());
            sleep().await;
            assert!(!trigger.trigger());
            let task_output = (&mut graceful_task).await;
            assert!(graceful_task.is_finished());
            assert_eq!(
                task_output.finish_mode,
                FinishMode::Shutdown(GracefulKind::Explicit)
            );
            assert_eq!(task_output.join_result.unwrap(), GracefulKind::Explicit);

            let (trigger, ctrlc) = ctrlc_mocked();
            let mut graceful_task = GracefulTask::builder_default()
                .ctrlc_shutdown()
                .spawn_ctrlc_mocked(async |shutdown| shutdown.await, ctrlc);
            assert!(!graceful_task.is_finished());
            assert!(trigger.trigger());
            sleep().await;
            assert!(!graceful_task.shutdown_handle().trigger());
            let task_output = (&mut graceful_task).await;
            assert!(graceful_task.is_finished());
            assert_eq!(
                task_output.finish_mode,
                FinishMode::Shutdown(GracefulKind::CtrlC)
            );
            assert_eq!(task_output.join_result.unwrap(), GracefulKind::CtrlC);
        }

        let (trigger, ctrlc) = ctrlc_mocked();
        let mut graceful_task = GracefulTask::builder_default()
            .spawn_ctrlc_mocked(async |shutdown| shutdown.await, ctrlc);
        assert!(!graceful_task.is_finished());
        assert!(!trigger.trigger());
        sleep().await;
        assert!(graceful_task.shutdown_handle().trigger());
        let task_output = (&mut graceful_task).await;
        assert!(graceful_task.is_finished());
        assert_eq!(
            task_output.finish_mode,
            FinishMode::Shutdown(GracefulKind::Explicit)
        );
        assert_eq!(task_output.join_result.unwrap(), GracefulKind::Explicit);

        let mut graceful_task = GracefulTask::builder_default().spawn(async |_| 42);
        sleep().await;
        assert!(!graceful_task.shutdown_handle().trigger());
        let task_output = (&mut graceful_task).await;
        assert!(graceful_task.is_finished());
        assert_eq!(task_output.finish_mode, FinishMode::Complete);
        assert_eq!(task_output.join_result.unwrap(), 42);

        let mut graceful_task =
            GracefulTask::builder_default().spawn(async |_| sleep_double().await);
        assert!(!graceful_task.is_finished());
        sleep().await;
        assert!(graceful_task.shutdown_handle().trigger());
        let task_output = (&mut graceful_task).await;
        assert!(graceful_task.is_finished());
        assert_eq!(
            task_output.finish_mode,
            FinishMode::Shutdown(GracefulKind::Explicit)
        );
        assert_eq!(task_output.join_result.unwrap(), ());
    }
}