polyc_runtime/

supervise.rs

//! Server-task supervision: run a binary's spawned servers until a shutdown
//! signal, treating any earlier exit as fatal.
//!
//! A process whose serving task has died must flip readiness off and exit
//! (so the orchestrator restarts it) rather than keep answering `/readyz`
//! over a dead surface.

use std::time::Duration;

use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;

use crate::health::Health;

/// Default grace before remaining tasks are hard-aborted, when
/// `POLYCHROME_DRAIN_GRACE_SECS` is unset. Comfortably under the Kubernetes
/// default `terminationGracePeriodSeconds` (30s) so a wedged task can't drag the
/// drain past the grace window and force a SIGKILL.
const DEFAULT_DRAIN_GRACE: Duration = Duration::from_secs(25);

/// Resolve the drain grace from `POLYCHROME_DRAIN_GRACE_SECS`, falling back to
/// [`DEFAULT_DRAIN_GRACE`].
///
/// The right budget is **per-deployment**: a binary running long in-flight work
/// (the harness drains LLM turns and sets `terminationGracePeriodSeconds: 150`)
/// must drain for most of its grace, while edges/control-plane keep the 30s
/// default. A single compile-time const would either truncate the harness's
/// drain or overrun an edge's grace, so it is read from the environment and the
/// deployment sets it to sit just under its own `terminationGracePeriodSeconds`.
fn drain_grace() -> Duration {
    std::env::var("POLYCHROME_DRAIN_GRACE_SECS")
        .ok()
        .and_then(|v| v.parse::<u64>().ok())
        .map_or(DEFAULT_DRAIN_GRACE, Duration::from_secs)
}

/// Block until `shutdown` is cancelled or any task in `servers` finishes.
///
/// Either way, readiness is flipped off, `shutdown` is cancelled, and the
/// remaining tasks are drained. A task finishing before the shutdown signal —
/// even cleanly — is an error: servers run until told to stop.
///
/// # Errors
///
/// Returns the first server failure: an exit before shutdown, a serve error
/// surfaced during drain, or a panic.
pub async fn until_shutdown(
    mut servers: JoinSet<anyhow::Result<()>>,
    health: &Health,
    shutdown: &CancellationToken,
) -> anyhow::Result<()> {
    let early = tokio::select! {
        // Biased so an already-delivered shutdown signal reads as graceful
        // even when a server task has (consequently) already finished.
        biased;
        () = shutdown.cancelled() => None,
        res = servers.join_next() => res,
    };
    health.set_ready(false);
    shutdown.cancel();

    let mut failure = early.map(|res| {
        task_failure(res).map_or_else(
            || anyhow::anyhow!("server exited cleanly before shutdown"),
            |err| err.context("server exited before shutdown"),
        )
    });
    if failure.is_some() {
        tracing::error!("server task exited before shutdown; draining and exiting");
    } else {
        tracing::info!("shutdown signal received; draining in-flight work");
    }

    drain_bounded(&mut servers, &mut failure, drain_grace()).await;
    failure.map_or(Ok(()), Err)
}

/// Drain `servers`, recording the first failure into `failure`, but no longer
/// than `grace`. If the grace window expires with tasks still running, they are
/// hard-aborted and reaped so a wedged task can't block shutdown indefinitely
/// (an abort is our own doing, so cancelled tasks are not counted as failures).
async fn drain_bounded(
    servers: &mut JoinSet<anyhow::Result<()>>,
    failure: &mut Option<anyhow::Error>,
    grace: Duration,
) {
    let reap = async {
        while let Some(res) = servers.join_next().await {
            if let Some(err) = task_failure(res) {
                tracing::error!(error = ?err, "server task failed during drain");
                failure.get_or_insert(err);
            }
        }
    };
    if tokio::time::timeout(grace, reap).await.is_err() {
        tracing::error!(
            grace_secs = grace.as_secs(),
            "drain exceeded grace window; aborting remaining server tasks"
        );
        servers.abort_all();
        // Reap the remaining tasks. A genuine failure that surfaced right at the
        // grace boundary (a serve error or panic) must still be recorded — only
        // the cancellations from our own `abort_all` are not failures.
        while let Some(res) = servers.join_next().await {
            match res {
                Ok(Ok(())) => {}
                Ok(Err(err)) => {
                    tracing::error!(error = ?err, "server task failed during bounded drain");
                    failure.get_or_insert(err);
                }
                // Our own abort — expected, not a failure.
                Err(join) if join.is_cancelled() => {}
                Err(join) => {
                    let err = anyhow::Error::new(join).context("server task panicked");
                    tracing::error!(error = ?err, "server task panicked during bounded drain");
                    failure.get_or_insert(err);
                }
            }
        }
    }
}

/// The error inside a finished task's join outcome, if any: a serve error or
/// a panic. A clean `Ok(())` exit returns `None`.
fn task_failure(res: Result<anyhow::Result<()>, tokio::task::JoinError>) -> Option<anyhow::Error> {
    match res {
        Ok(Ok(())) => None,
        Ok(Err(err)) => Some(err),
        Err(join) => Some(anyhow::Error::new(join).context("server task panicked")),
    }
}

#[cfg(test)]
mod tests {
    #![allow(clippy::pedantic, clippy::nursery, missing_docs)]

    use std::time::Duration;

    use tokio::task::JoinSet;
    use tokio_util::sync::CancellationToken;

    use super::{drain_bounded, until_shutdown};
    use crate::health::Health;

    #[tokio::test]
    async fn graceful_shutdown_drains_and_returns_ok() {
        let health = Health::new();
        let shutdown = CancellationToken::new();
        let mut servers = JoinSet::new();
        for _ in 0..2 {
            let shutdown = shutdown.clone();
            servers.spawn(async move {
                shutdown.cancelled().await;
                Ok(())
            });
        }
        health.set_ready(true);

        shutdown.cancel();
        until_shutdown(servers, &health, &shutdown)
            .await
            .expect("graceful shutdown is not an error");
    }

    #[tokio::test]
    async fn failing_server_is_fatal_and_cancels_the_rest() {
        let health = Health::new();
        let shutdown = CancellationToken::new();
        let mut servers = JoinSet::new();
        servers.spawn(async { Err(anyhow::anyhow!("bind lost")) });
        // The healthy peer must be drained via the cancelled token, proving
        // the early exit propagates shutdown to the rest.
        let peer = shutdown.clone();
        servers.spawn(async move {
            peer.cancelled().await;
            Ok(())
        });
        health.set_ready(true);

        let err = until_shutdown(servers, &health, &shutdown)
            .await
            .expect_err("a dead server is fatal");
        assert!(err.to_string().contains("server exited before shutdown"));
        assert!(
            shutdown.is_cancelled(),
            "remaining servers are told to stop"
        );
    }

    #[tokio::test]
    async fn clean_early_exit_is_still_fatal() {
        let health = Health::new();
        let shutdown = CancellationToken::new();
        let mut servers = JoinSet::new();
        servers.spawn(async { Ok(()) });

        let err = until_shutdown(servers, &health, &shutdown)
            .await
            .expect_err("servers run until told to stop");
        assert!(err.to_string().contains("exited cleanly before shutdown"));
    }

    #[tokio::test]
    async fn drain_aborts_a_wedged_task_after_grace() {
        let mut servers = JoinSet::new();
        // A task that never completes (ignores the shutdown signal) — without a
        // bounded drain this would hang until the orchestrator SIGKILLs the pod.
        servers.spawn(async {
            std::future::pending::<()>().await;
            Ok(())
        });
        let mut failure = None;

        let start = std::time::Instant::now();
        drain_bounded(&mut servers, &mut failure, Duration::from_millis(50)).await;

        assert!(
            start.elapsed() < Duration::from_secs(5),
            "drain is bounded by the grace window, not the wedged task"
        );
        assert!(
            failure.is_none(),
            "a task we aborted ourselves is not counted as a failure"
        );
        assert!(
            servers.is_empty(),
            "remaining tasks were aborted and reaped"
        );
    }

    #[tokio::test]
    async fn drain_records_a_failure_even_when_grace_expires() {
        let mut servers = JoinSet::new();
        // A wedged task forces the grace-timeout / abort path…
        servers.spawn(async {
            std::future::pending::<()>().await;
            Ok(())
        });
        // …while a real serve error must still be surfaced, not swallowed as
        // "our own abort".
        servers.spawn(async { Err(anyhow::anyhow!("serve error during drain")) });
        let mut failure = None;

        drain_bounded(&mut servers, &mut failure, Duration::from_millis(50)).await;

        let err = failure.expect("a genuine failure during the bounded drain must be recorded");
        assert!(err.to_string().contains("serve error during drain"));
    }
}