ruststream 0.5.3

//! Running the service: startup sequence, signal handling and graceful shutdown.

use std::fmt;
use std::{future::Future, sync::Arc, time::Duration};

#[cfg(unix)]
use tokio::signal::unix::{SignalKind, signal};

use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
use tracing::{debug, info, warn};

use crate::runtime::failure::ErrorShutdown;
use crate::runtime::lifecycle::{BoxError, BoxFuture};

use super::service::RegisteredBroker;
use super::{LifecycleHook, RustStream, RustStreamError};

/// A lifecycle hook with the state already bound, so [`RunningApp`] stays non-generic.
type BoundHook = Box<dyn FnOnce() -> BoxFuture<'static, Result<(), BoxError>> + Send>;

/// Binds the shared state into each hook, erasing the state type from the hook list.
fn bind_hooks<St: Send + Sync + 'static>(
    hooks: Vec<LifecycleHook<St>>,
    state: &Arc<St>,
) -> Vec<BoundHook> {
    hooks
        .into_iter()
        .map(|hook| {
            let state = Arc::clone(state);
            Box::new(move || hook(state)) as BoundHook
        })
        .collect()
}

// `run`/`run_until` are routinely driven from a multi-thread runtime (`tokio::spawn`, the CLI's
// `block_on`), so their futures must be `Send`: the shared state is held as `Arc<St>` across the
// startup awaits (needs `St: Sync`) and the global stack `L` is carried in `self` (needs `L: Send`).
// `St: 'static` is what every constructible app already satisfies (the `on_startup` producer
// returns the state from a `'static` boxed future); naming it here lets `start` box the shutdown
// hooks with the state bound in.
impl<L: Send, St: Send + Sync + 'static, PP> RustStream<L, St, PP> {
    /// Runs the service until an interrupt (`SIGINT` / `SIGTERM`) is received, then shuts down
    /// gracefully.
    ///
    /// # Errors
    ///
    /// Returns [`RustStreamError`] if a broker fails to connect, a subscription fails to open, a
    /// dispatch task panics, or a broker fails to shut down.
    pub async fn run(self) -> Result<(), RustStreamError> {
        self.run_until(wait_for_signal()).await
    }

    /// Runs the service until `shutdown` resolves, then shuts down gracefully.
    ///
    /// Use this instead of [`run`](Self::run) to drive shutdown from a caller-owned future (a
    /// name, a timeout, a test signal) rather than from process signals.
    ///
    /// # Errors
    ///
    /// Returns [`RustStreamError`] if a broker fails to connect, a subscription fails to open, a
    /// dispatch task panics, or a broker fails to shut down.
    pub async fn run_until<F>(self, shutdown: F) -> Result<(), RustStreamError>
    where
        F: Future<Output = ()> + Send,
    {
        let running = self.start().await?;
        tokio::select! {
            () = shutdown => info!(target: "ruststream::lifecycle", "shutdown signal received"),
            () = running.stopping() => {
                info!(target: "ruststream::lifecycle", "fail-fast shutdown triggered");
            }
        }
        running.shutdown().await
    }

    /// Starts the service in the background and hands back a [`RunningApp`] handle.
    ///
    /// Performs the same startup sequence as [`run`](Self::run) - the `on_startup` state
    /// producer, broker connects, subscription opens, `after_startup` hooks - and resolves once
    /// the service is running, so a startup failure surfaces here, before the caller starts
    /// accepting its own traffic. Installs no signal handlers: the caller decides what stops the
    /// service, by calling [`RunningApp::shutdown`].
    ///
    /// Use this to run the service beside another foreground server (an HTTP framework) in the
    /// same process; when the service is the whole process, [`run`](Self::run) /
    /// [`run_until`](Self::run_until) stay the simpler form.
    ///
    /// # Errors
    ///
    /// Returns [`RustStreamError`] if the state producer or an `after_startup` hook fails, a
    /// broker fails to connect, or a subscription fails to open.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # #[cfg(feature = "memory")]
    /// # async fn run() -> Result<(), ruststream::runtime::RustStreamError> {
    /// use ruststream::memory::MemoryBroker;
    /// use ruststream::runtime::{AppInfo, RustStream};
    ///
    /// let app = RustStream::new(AppInfo::new("svc", "0.1.0")).register_broker(MemoryBroker::new());
    /// let running = app.start().await?;
    /// // ... serve HTTP in the foreground ...
    /// running.shutdown().await
    /// # }
    /// ```
    pub async fn start(self) -> Result<RunningApp, RustStreamError> {
        let Self {
            info,
            brokers,
            starters,
            handlers,
            state_init,
            after_startup,
            on_shutdown,
            after_shutdown,
            shutdown_timeout,
            continuations,
            ..
        } = self;

        info!(
            target: "ruststream::lifecycle",
            service = %info.title,
            version = %info.version,
            brokers = brokers.len(),
            subscribers = starters.len(),
            "starting service",
        );

        debug!(target: "ruststream::lifecycle", "producing application state");
        let state = state_init().await.map_err(RustStreamError::Startup)?;
        let state = Arc::new(state);

        for broker in &brokers {
            broker
                .lifecycle
                .connect()
                .await
                .map_err(RustStreamError::Connect)?;
            info!(
                target: "ruststream::lifecycle",
                broker = broker.label.as_deref().unwrap_or_else(|| broker.lifecycle.name()),
                "broker connected",
            );
        }

        let token = CancellationToken::new();
        // Shared with every dispatch task: a fail-fast failure records its reason here and cancels
        // the token, which both stops the loops and resolves `stopping()`.
        let error_shutdown = ErrorShutdown::new(token.clone());
        let mut handles = Vec::with_capacity(starters.len());
        for (starter, meta) in starters.into_iter().zip(handlers) {
            let handle = starter(state.clone(), error_shutdown.clone(), token.clone())
                .await
                .map_err(RustStreamError::Subscribe)?;
            info!(
                target: "ruststream::dispatch",
                subscriber = %meta.name,
                input = meta.input_type,
                "subscriber started",
            );
            handles.push(handle);
        }

        if !after_startup.is_empty() {
            debug!(target: "ruststream::lifecycle", count = after_startup.len(), "running after_startup hooks");
        }
        for hook in after_startup {
            hook(Arc::clone(&state))
                .await
                .map_err(RustStreamError::Startup)?;
        }

        info!(target: "ruststream::lifecycle", subscribers = handles.len(), "service running");

        Ok(RunningApp {
            token,
            error_shutdown,
            handles,
            on_shutdown: bind_hooks(on_shutdown, &state),
            after_shutdown: bind_hooks(after_shutdown, &state),
            brokers,
            shutdown_timeout,
            continuations,
        })
    }
}

/// A started service, handed out by [`RustStream::start`].
///
/// The handle owns the graceful teardown: dropping it without calling
/// [`shutdown`](Self::shutdown) detaches the service (dispatch tasks keep running, nothing is
/// drained), per the crate rule that destructors never block. The intended shape when running
/// beside a foreground server:
///
/// ```no_run
/// # #[cfg(feature = "memory")]
/// # async fn run() -> Result<(), ruststream::runtime::RustStreamError> {
/// use ruststream::memory::MemoryBroker;
/// use ruststream::runtime::{AppInfo, RustStream};
///
/// let app = RustStream::new(AppInfo::new("svc", "0.1.0")).register_broker(MemoryBroker::new());
/// let running = app.start().await?;
/// // e.g. axum::serve(listener, router)
/// //     .with_graceful_shutdown(running.stopping())
/// //     .await?;
/// running.shutdown().await
/// # }
/// ```
#[must_use = "dropping the handle detaches the service without graceful shutdown"]
pub struct RunningApp {
    token: CancellationToken,
    error_shutdown: ErrorShutdown,
    handles: Vec<JoinHandle<()>>,
    on_shutdown: Vec<BoundHook>,
    after_shutdown: Vec<BoundHook>,
    brokers: Vec<RegisteredBroker>,
    shutdown_timeout: Option<Duration>,
    continuations: TaskTracker,
}

impl fmt::Debug for RunningApp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("RunningApp")
            .field("subscribers", &self.handles.len())
            .field("brokers", &self.brokers.len())
            .field("shutdown_timeout", &self.shutdown_timeout)
            .finish_non_exhaustive()
    }
}

impl RunningApp {
    /// Resolves when the service begins stopping on its own: a subscriber hit a fail-fast
    /// failure and tore the service down.
    ///
    /// The future is owned (`'static`), so it plugs directly into another server's graceful
    /// shutdown (for example axum's `with_graceful_shutdown`), stopping the host when the
    /// messaging side dies. It does not resolve on an orderly [`shutdown`](Self::shutdown) call -
    /// the caller drives that path itself.
    ///
    /// # Cancel safety
    ///
    /// Cancel-safe: dropping the future loses nothing; a fresh call observes the same state.
    pub fn stopping(&self) -> impl Future<Output = ()> + Send + 'static {
        self.token.clone().cancelled_owned()
    }

    /// Shuts the service down gracefully.
    ///
    /// Runs the `on_shutdown` hooks, stops the dispatch loops, drains in-flight handlers and
    /// post-settle continuations (bounded by the configured shutdown timeout), shuts the brokers
    /// down in reverse registration order, then runs the `after_shutdown` hooks.
    ///
    /// # Errors
    ///
    /// Returns [`RustStreamError`] if a dispatch task panicked, a broker failed to shut down, or
    /// the service had already torn itself down on a fail-fast failure (surfaced as
    /// [`RustStreamError::Dispatch`] so the operator sees a non-zero exit, not a silent stop).
    pub async fn shutdown(self) -> Result<(), RustStreamError> {
        let Self {
            token,
            error_shutdown,
            handles,
            on_shutdown,
            after_shutdown,
            brokers,
            shutdown_timeout,
            continuations,
        } = self;

        for hook in on_shutdown {
            if let Err(err) = hook().await {
                warn!(target: "ruststream::lifecycle", error = %err, "on_shutdown hook failed");
            }
        }

        token.cancel();
        debug!(target: "ruststream::lifecycle", "draining in-flight handlers");
        drain_handles(handles, shutdown_timeout).await?;

        // Handlers have stopped, so no new post-settle continuations can be spawned: close the
        // tracker and drain the in-flight ones, bounded by the same shutdown timeout. They are
        // at-most-once, so timing one out only abandons follow-up work, never a settlement.
        drain_continuations(continuations, shutdown_timeout).await;

        for broker in brokers.iter().rev() {
            broker
                .lifecycle
                .shutdown()
                .await
                .map_err(RustStreamError::Shutdown)?;
            debug!(
                target: "ruststream::lifecycle",
                broker = broker.label.as_deref().unwrap_or_else(|| broker.lifecycle.name()),
                "broker shut down",
            );
        }

        for hook in after_shutdown {
            if let Err(err) = hook().await {
                warn!(target: "ruststream::lifecycle", error = %err, "after_shutdown hook failed");
            }
        }
        info!(target: "ruststream::lifecycle", "service stopped");

        // A fail-fast failure tore the service down: surface it so an orchestrator restarts the
        // service and the operator sees a non-zero exit, not a silent stop.
        if let Some(reason) = error_shutdown.taken_failure() {
            return Err(RustStreamError::Dispatch(reason));
        }
        Ok(())
    }
}

/// Awaits all handler tasks, bounded by `timeout` if set. On timeout the remaining tasks are
/// aborted; without a timeout, a panicking task surfaces as [`RustStreamError::Join`].
async fn drain_handles(
    handles: Vec<JoinHandle<()>>,
    timeout: Option<Duration>,
) -> Result<(), RustStreamError> {
    let Some(timeout) = timeout else {
        for handle in handles {
            handle.await.map_err(RustStreamError::Join)?;
        }
        return Ok(());
    };

    let aborts: Vec<_> = handles.iter().map(JoinHandle::abort_handle).collect();
    if tokio::time::timeout(timeout, futures::future::join_all(handles))
        .await
        .is_err()
    {
        warn!(
            target: "ruststream::lifecycle",
            "graceful shutdown timed out; aborting in-flight handlers",
        );
        for abort in aborts {
            abort.abort();
        }
    }
    Ok(())
}

/// Closes the post-settle continuation tracker and waits for the in-flight continuations to finish,
/// bounded by `timeout` when set. On timeout the remaining continuations keep running detached (the
/// tracker does not own abort handles); they are at-most-once side effects, so abandoning them is
/// safe.
async fn drain_continuations(continuations: TaskTracker, timeout: Option<Duration>) {
    continuations.close();
    if continuations.is_empty() {
        return;
    }
    debug!(target: "ruststream::lifecycle", "draining post-settle continuations");
    match timeout {
        Some(timeout) => {
            if tokio::time::timeout(timeout, continuations.wait())
                .await
                .is_err()
            {
                warn!(
                    target: "ruststream::lifecycle",
                    "graceful shutdown timed out; abandoning in-flight continuations",
                );
            }
        }
        None => continuations.wait().await,
    }
}

async fn wait_for_signal() {
    #[cfg(unix)]
    {
        let Ok(mut term) = signal(SignalKind::terminate()) else {
            let _ = tokio::signal::ctrl_c().await;
            return;
        };
        tokio::select! {
            _ = tokio::signal::ctrl_c() => {}
            _ = term.recv() => {}
        }
    }
    #[cfg(not(unix))]
    {
        let _ = tokio::signal::ctrl_c().await;
    }
}