aion_server/

shutdown.rs

//! Graceful shutdown and single-node activity drain coordination.

use std::process::ExitCode;
use std::sync::Arc;
use std::time::Duration;

use tokio::sync::{Notify, watch};
use tracing::{error, info, warn};

use crate::ServerState;
use crate::error::ServerError;
use crate::worker::{AttemptKey, LostWorkerReport};

/// Process exit selected by the shutdown coordinator.
///
/// Serialized by name into the shutdown outcome record
/// ([`crate::control::outcome`]), so the face an operator reads from
/// `aion server stop` is this exact enum, not a re-derivation.
#[derive(Clone, Copy, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)]
pub enum ShutdownOutcome {
    /// Drain completed before the configured timeout.
    Clean,
    /// In-flight activities outlived the drain timeout and were parked for
    /// restart recovery (#207): nothing recorded, nothing delivered — the
    /// recoverable-by-design state, so a fully-parked drain is a SUCCESS. A
    /// long-running activity (an agent round runs hours) outliving any sane
    /// drain window is the expected case, and a non-zero exit on every routine
    /// deploy would train operators to ignore failures.
    Parked,
    /// The drain timed out AND the park itself failed (lock poison, sink
    /// error): in-flight state could not be handed to restart recovery.
    TimedOut,
    /// A second termination signal requested immediate process exit.
    Forced,
}

impl ShutdownOutcome {
    /// Convert the outcome to the process exit code required by the operations contract.
    #[must_use]
    pub fn exit_code(self) -> ExitCode {
        match self {
            Self::Clean | Self::Parked => ExitCode::SUCCESS,
            Self::TimedOut => ExitCode::FAILURE,
            Self::Forced => ExitCode::from(130),
        }
    }
}

/// Everything one graceful shutdown observed, for the outcome record the
/// dying server writes into the death note ([`crate::control::outcome`]).
/// The process exit code still collapses this to the #207 contract via
/// [`ShutdownOutcome::exit_code`]; the report is how the rest crosses the
/// process boundary.
#[derive(Debug)]
pub struct ShutdownReport {
    /// The drain's face.
    pub outcome: ShutdownOutcome,
    /// The window that governed the wait.
    pub drain_timeout: Duration,
    /// How many connected workers received the drain request.
    pub delivered_drain_requests: usize,
    /// Parked in-flight work at drain timeout, by worker. Empty on a clean
    /// drain and on a forced exit (a second signal skips the accounting).
    pub parked: Vec<LostWorkerReport>,
    /// Declared-body commands still executing when the drain ended (timeout or
    /// forced exit). A declared body is in-flight work no worker holds and no
    /// heartbeat tracks; these attempts are parked by silence exactly like the
    /// worker dispatches above — nothing is recorded, the engine teardown ends
    /// their process groups with the server, and post-restart replay
    /// re-dispatches each one. Empty on a clean drain: the drain gate waits for
    /// this census too.
    pub parked_declared: Vec<AttemptKey>,
    /// Managed workers proven stopped, by deployment name.
    pub managed_workers_stopped: Vec<String>,
    /// Managed workers that could not be proven stopped, each rendered with
    /// the observation that contradicted the stop.
    pub managed_workers_unstopped: Vec<String>,
}

/// Cloneable gate shared by transports, dispatchers, worker streams, and the
/// shutdown coordinator.
#[derive(Clone, Debug, Default)]
pub struct DrainState {
    inner: Arc<DrainStateInner>,
}

#[derive(Debug)]
struct DrainStateInner {
    /// The drain latch, held in a `watch` rather than an `AtomicBool` because
    /// one gated seam cannot poll a flag: the bridge's park for an arriving
    /// worker BLOCKS until a worker appears, so it needs this same latch in
    /// awaitable form to stop blocking when the server starts draining.
    ///
    /// One latch answering both questions is the point. Two independent notions
    /// of "we are shutting down" inside the dispatch path is exactly how a
    /// drain gate and a dispatch parked behind it come to disagree — and the
    /// disagreement is unobservable until a process refuses to exit.
    draining: watch::Sender<bool>,
    empty: Notify,
}

impl Default for DrainStateInner {
    fn default() -> Self {
        Self {
            draining: watch::Sender::new(false),
            empty: Notify::default(),
        }
    }
}

impl DrainState {
    /// Return whether drain has begun and new workflow/activity starts must be rejected.
    #[must_use]
    pub fn is_draining(&self) -> bool {
        *self.inner.draining.borrow()
    }

    /// Mark the server draining. Returns true for the first caller that changed the state.
    #[must_use]
    pub fn begin(&self) -> bool {
        // Sets the latch AND wakes every awaiting seam in one write, so a park
        // released by drain can never observe a latch that has not been set
        // yet.
        !self.inner.draining.send_replace(true)
    }

    /// Resolve as soon as drain has begun — [`Self::is_draining`] in awaitable
    /// form, over the same latch.
    ///
    /// For the seam that must block on an external arrival (a worker
    /// registering) and therefore cannot re-read a flag between iterations.
    /// Waking here decides nothing on its own: the woken caller re-runs its
    /// normal loop and meets [`Self::ensure_accepting`], which is still the only
    /// place a drain refusal is produced.
    pub async fn wait_for_drain(&self) {
        let mut draining = self.inner.draining.subscribe();
        while !*draining.borrow_and_update() {
            if draining.changed().await.is_err() {
                // Unreachable while this handle lives — it owns the `Arc` the
                // sender sits in — but a closed latch must read as "drained"
                // rather than block forever on a state that cannot recover.
                break;
            }
        }
    }

    /// Reject a new unit of work if drain has already begun.
    ///
    /// # Errors
    ///
    /// Returns [`ServerError::WorkerDispatch`] with a stable drain message when work is closed.
    pub fn ensure_accepting(
        &self,
        namespace: &str,
        activity_type: &str,
    ) -> Result<(), ServerError> {
        if self.is_draining() {
            Err(ServerError::worker_dispatch(
                namespace.to_owned(),
                activity_type.to_owned(),
                "server is draining and not accepting new activity tasks",
            ))
        } else {
            Ok(())
        }
    }

    /// Wake waiters after in-flight accounting may have reached zero.
    pub fn notify_activity_drained(&self) {
        self.inner.empty.notify_waiters();
    }

    /// Whether both in-flight censuses read empty: the heartbeat tracker's
    /// (worker-held activities) AND the declared-body registry's (commands this
    /// server is executing itself, which no worker holds and no heartbeat
    /// tracks). Draining on the first census alone is how a server once
    /// reported `Clean` over a live declared command.
    fn nothing_in_flight(state: &ServerState) -> Result<bool, ServerError> {
        Ok(state.heartbeat_tracker().in_flight_count()? == 0
            && state.declared_attempts().executing()?.is_empty())
    }

    async fn wait_for_empty(&self, state: &ServerState) -> Result<(), ServerError> {
        loop {
            if Self::nothing_in_flight(state)? {
                return Ok(());
            }
            let notified = self.inner.empty.notified();
            if Self::nothing_in_flight(state)? {
                return Ok(());
            }
            notified.await;
        }
    }
}

/// Run the graceful drain after the first termination signal.
///
/// The caller is responsible for stopping transports as soon as drain begins.
///
/// # Errors
///
/// Returns [`ServerError`] if worker-drain broadcast, in-flight accounting, timeout failure
/// surfacing, or engine shutdown fails.
pub async fn drain_after_first_signal(
    state: ServerState,
    second_signal: impl std::future::Future<Output = ()>,
) -> Result<ShutdownReport, ServerError> {
    let drain = state.drain_state().clone();
    let first = drain.begin();
    if first {
        info!("shutdown signal received; beginning graceful drain");
    }

    let delivered_workers = state.worker_registry().broadcast_drain()?;
    info!(delivered_workers, "sent drain request to connected workers");

    let timeout = state.runtime_config().drain_timeout;
    tokio::pin!(second_signal);

    let (outcome, parked, parked_declared) = tokio::select! {
        () = &mut second_signal => {
            warn!("second shutdown signal received; forcing immediate exit");
            // A forced exit skips the worker-park accounting, but the
            // declared-command census is a read this process can still afford:
            // these commands die with the server either way, and a report that
            // cannot name them is the report the next operator debugs against.
            let declared = state.declared_attempts().executing().unwrap_or_else(|census_error| {
                error!(
                    %census_error,
                    "forced exit could not read the declared-command census; the \
                     outcome record will not name the commands that died with the server"
                );
                Vec::new()
            });
            (ShutdownOutcome::Forced, Vec::new(), declared)
        }
        result = wait_for_drain_or_timeout(&state, &drain, timeout) => result?,
    };

    // W-4 containment: managed worker PROCESSES stop with the server — AFTER
    // the drain window, never before it. Draining exists to let in-flight work
    // finish, and killing the workers doing that work would invert it; by the
    // time this runs the activities have either completed or been parked for
    // restart recovery (#207), and the next boot reconciles the fleet back up.
    //
    // It runs on the FORCED path too. A second signal asks for an immediate
    // exit, and this does bound that by the operator's own `stop_grace` — but a
    // worker outliving the server that owns it is worse than a bounded moment,
    // and the alternative (relying on the drop guard as the process unwinds)
    // reaps without ever verifying that it did.
    //
    // Deliberately NOT allowed to change the process exit contract (#72/#207):
    // a failure here is reported in full, loudly, and shutdown proceeds.
    // Failing the exit on a worker that would not die would turn a routine
    // deploy red, which is how operators learn to ignore failures.
    let managed = stop_managed_workers(&state).await;

    let report = ShutdownReport {
        outcome,
        drain_timeout: timeout,
        delivered_drain_requests: delivered_workers,
        parked,
        parked_declared,
        managed_workers_stopped: managed.stopped,
        managed_workers_unstopped: managed.failures.iter().map(ToString::to_string).collect(),
    };
    if matches!(outcome, ShutdownOutcome::Forced) {
        return Ok(report);
    }

    state.shutdown()?;
    Ok(report)
}

async fn wait_for_drain_or_timeout(
    state: &ServerState,
    drain: &DrainState,
    timeout: Duration,
) -> Result<(ShutdownOutcome, Vec<LostWorkerReport>, Vec<AttemptKey>), ServerError> {
    match tokio::time::timeout(timeout, drain.wait_for_empty(state)).await {
        Ok(result) => {
            result?;
            info!("activity drain completed cleanly");
            Ok((ShutdownOutcome::Clean, Vec::new(), Vec::new()))
        }
        Err(_elapsed) => {
            // #207 drain-timeout backstop: PARK the remaining in-flight
            // dispatches for restart recovery instead of synthesizing
            // transport-loss failures. Nothing is recorded, so the
            // durable log converges on the kill -9 shape and post-restart
            // replay re-dispatches every parked ordinal. A park that itself
            // fails leaves in-flight state unhanded — the one remaining
            // FAILURE-worthy drain outcome.
            //
            // Declared-body attempts park the same way but need no sweep: the
            // engine teardown that follows ends their process groups with the
            // server and their unrecorded attempts re-dispatch on restart. What
            // the backstop owes them is the CENSUS — a park that cannot NAME
            // the in-flight commands it is abandoning is the same unhanded
            // state as a failed worker park, and is reported the same way.
            let declared = match state.declared_attempts().executing() {
                Ok(executing) => executing,
                Err(census_error) => {
                    error!(
                        %census_error,
                        "activity drain timed out and the declared-command census could \
                         not be read; exiting with the failure drain outcome"
                    );
                    return Ok((ShutdownOutcome::TimedOut, Vec::new(), Vec::new()));
                }
            };
            match state
                .heartbeat_tracker()
                .park_all_in_flight_workers(state.worker_registry(), state.pending_activities())
            {
                Ok(reports) => {
                    log_parked_workers(&reports);
                    log_parked_declared(&declared);
                    Ok((ShutdownOutcome::Parked, reports, declared))
                }
                Err(park_error) => {
                    error!(
                        %park_error,
                        "activity drain timed out and parking the remaining in-flight \
                         activities failed; exiting with the failure drain outcome"
                    );
                    Ok((ShutdownOutcome::TimedOut, Vec::new(), declared))
                }
            }
        }
    }
}

/// Stop every supervised managed worker, and say exactly what happened.
///
/// An empty failure list is the no-orphan claim: each stop returned only after
/// a signal-zero probe found the worker's process group empty. Anything else is
/// logged per worker with the observation that contradicted it — never summed
/// into a single count that could read as calm.
async fn stop_managed_workers(state: &ServerState) -> crate::worker::FleetShutdownReport {
    let report = state.worker_supervisor().shutdown().await;
    if report.failures.is_empty() {
        info!("managed workers stopped; every process group confirmed empty");
        return report;
    }
    for failure in &report.failures {
        error!(%failure, "a managed worker could not be confirmed stopped at shutdown");
    }
    error!(
        unstopped = report.failures.len(),
        "shutdown could not prove every managed worker stopped; check for orphaned processes"
    );
    report
}

fn log_parked_declared(declared: &[AttemptKey]) {
    for key in declared {
        info!(
            workflow_id = %key.workflow_id,
            run_id = %key.run_id,
            activity_id = %key.activity_id,
            attempt = key.attempt,
            "drain timed out with this declared command still executing; its process \
             group ends with the server and the unrecorded attempt re-dispatches on \
             the next boot"
        );
    }
}

fn log_parked_workers(reports: &[LostWorkerReport]) {
    let parked_tasks: usize = reports.iter().map(|report| report.tasks.len()).sum();
    if parked_tasks == 0 {
        info!("activity drain timed out with no tracked in-flight activities to park");
    } else {
        info!(
            parked_workers = reports.len(),
            parked_tasks,
            "activity drain timed out; remaining activities parked for restart recovery"
        );
    }
}

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

    use super::{DrainState, ShutdownOutcome};

    type TestResult = Result<(), Box<dyn std::error::Error>>;

    /// How long a woken waiter is allowed to take. Generous, and not a
    /// behavioural bound: a correct latch resolves in microseconds and a broken
    /// one never resolves, so this only decides how long a failure takes to
    /// report.
    const WAKE_BUDGET: Duration = Duration::from_secs(5);

    #[test]
    fn begin_is_idempotent_and_sets_draining() {
        let drain = DrainState::default();

        assert!(!drain.is_draining());
        assert!(drain.begin());
        assert!(drain.is_draining());
        assert!(!drain.begin());
    }

    /// #72: a waiter already parked on the latch is woken by `begin`.
    ///
    /// This is the ordering the bridge's park depends on and the one a
    /// notification-only signal gets wrong: the waiter registers first and the
    /// latch flips afterwards, so nothing it could poll has changed yet. If the
    /// wake is ever lost here, a dispatch parked for a worker becomes
    /// unwakeable and the process cannot exit.
    #[tokio::test]
    async fn begin_wakes_a_waiter_that_registered_before_the_latch_flipped() {
        let drain = DrainState::default();
        let waiting = drain.clone();
        let waiter = tokio::spawn(async move { waiting.wait_for_drain().await });
        // Let the waiter reach its await before the latch is touched.
        tokio::task::yield_now().await;
        assert!(!drain.is_draining());
        assert!(drain.begin());

        let woken = tokio::time::timeout(WAKE_BUDGET, waiter).await;
        assert!(
            matches!(woken, Ok(Ok(()))),
            "a waiter registered before `begin` was not woken: {woken:?}"
        );
    }

    /// The other half of the same race: a waiter arriving AFTER the latch
    /// flipped must not wait for a notification that has already been sent.
    #[tokio::test]
    async fn wait_for_drain_resolves_at_once_once_drain_has_begun() {
        let drain = DrainState::default();
        assert!(drain.begin());

        let resolved = tokio::time::timeout(WAKE_BUDGET, drain.wait_for_drain()).await;
        assert!(
            resolved.is_ok(),
            "a waiter arriving after `begin` blocked instead of resolving"
        );
    }

    /// #207 exit contract: a fully-parked drain is a SUCCESS (parked state is
    /// recoverable by design); FAILURE is reserved for a park that itself
    /// failed; a forced exit keeps 130. `ExitCode` carries no `PartialEq`, so
    /// the mapping is asserted through its debug representation.
    #[test]
    fn exit_codes_map_parked_to_success_and_timed_out_to_failure() {
        let debug = |code: ExitCode| format!("{code:?}");
        assert_eq!(
            debug(ShutdownOutcome::Clean.exit_code()),
            debug(ExitCode::SUCCESS)
        );
        assert_eq!(
            debug(ShutdownOutcome::Parked.exit_code()),
            debug(ExitCode::SUCCESS)
        );
        assert_eq!(
            debug(ShutdownOutcome::TimedOut.exit_code()),
            debug(ExitCode::FAILURE)
        );
        assert_eq!(
            debug(ShutdownOutcome::Forced.exit_code()),
            debug(ExitCode::from(130))
        );
    }

    /// Red first — the drain gate must wait for an executing declared-body
    /// command: in-flight work no worker holds and no heartbeat tracks. Before
    /// the declared census joined the gate, a drain over a live declared
    /// command resolved immediately and the server reported `Clean` while the
    /// command still ran (found by `a_parked_drain_is_visible_in_the_stop_report`).
    #[tokio::test]
    async fn the_drain_gate_waits_for_an_executing_declared_command() -> TestResult {
        let (engine, _store, _visibility) = crate::api::http::test_support::shared_engine().await?;
        let resolver = crate::NamespaceResolver::from_config(
            crate::config::NamespaceConfig {
                mode: crate::config::NamespaceMode::SharedEngine,
            },
            engine.handle(),
        );
        let state = crate::api::http::test_support::server_state(
            engine,
            resolver,
            crate::api::http::test_support::runtime_config(),
        )
        .await?;

        let key = crate::worker::AttemptKey::new(
            aion_core::WorkflowId::new_v4(),
            aion_core::RunId::new_v4(),
            aion_core::ActivityId::from_sequence_position(1),
            1,
        );
        let (_context, cancellation) = aion_worker::ActivityContext::new(
            key.workflow_id.clone(),
            key.run_id.clone(),
            key.activity_id.clone(),
            key.attempt,
        );
        let registration = state.declared_attempts().register(key, cancellation)?;

        let drain = state.drain_state().clone();
        let waiter_state = state.clone();
        let waiter = tokio::spawn(async move { drain.wait_for_empty(&waiter_state).await });
        for _ in 0..32 {
            tokio::task::yield_now().await;
        }
        assert!(
            !waiter.is_finished(),
            "the drain gate resolved over a live declared command — the in-flight \
             census must include the declared-body registry"
        );

        // Finishing the command (the guard's drop) is what completes the drain,
        // through the registry's own wake of the drain latch.
        drop(registration);
        let joined = tokio::time::timeout(WAKE_BUDGET, waiter).await;
        joined
            .map_err(|_elapsed| "the drain gate never woke after the declared command finished")?
            .map_err(|join_error| format!("the drain waiter panicked: {join_error}"))??;
        Ok(())
    }
}