rs-zero 0.2.7

use std::{future::Future, pin::Pin, sync::Arc};

use tokio::sync::Notify;

/// Cloneable shutdown token used by [`crate::core::ServiceGroup`].
#[derive(Debug, Clone)]
pub struct ShutdownToken {
    inner: Arc<ShutdownState>,
}

#[derive(Debug)]
struct ShutdownState {
    notify: Notify,
    cancelled: std::sync::atomic::AtomicBool,
}

impl Default for ShutdownToken {
    fn default() -> Self {
        Self::new()
    }
}

impl ShutdownToken {
    /// Creates a token in the running state.
    pub fn new() -> Self {
        Self {
            inner: Arc::new(ShutdownState {
                notify: Notify::new(),
                cancelled: std::sync::atomic::AtomicBool::new(false),
            }),
        }
    }

    /// Requests shutdown. Calling this method multiple times is safe.
    pub fn cancel(&self) {
        if !self
            .inner
            .cancelled
            .swap(true, std::sync::atomic::Ordering::SeqCst)
        {
            self.inner.notify.notify_waiters();
        }
    }

    /// Returns whether shutdown has been requested.
    pub fn is_cancelled(&self) -> bool {
        self.inner
            .cancelled
            .load(std::sync::atomic::Ordering::SeqCst)
    }

    /// Waits until shutdown is requested.
    pub fn cancelled(&self) -> ShutdownFuture<'_> {
        Box::pin(async move {
            loop {
                let notified = self.inner.notify.notified();
                if self.is_cancelled() {
                    return;
                }
                notified.await;
            }
        })
    }
}

/// Future returned by [`ShutdownToken::cancelled`].
pub type ShutdownFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;

/// Waits until the process receives Ctrl-C.
///
/// Services can pass this future to server helpers that support graceful
/// shutdown.
pub async fn shutdown_signal() {
    if let Err(error) = tokio::signal::ctrl_c().await {
        tracing::warn!(%error, "failed to listen for ctrl-c");
    }
}