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
/// a default of 25 seconds, comfortably under the Kubernetes default
/// `terminationGracePeriodSeconds` (30s).
///
/// 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`.
///
/// `pub`, not module-private, so a caller with a task OUTSIDE the `servers`
/// `JoinSet` this module otherwise bounds — `polyc-control-plane`'s
/// post-drain awaits, #1370 review finding 3 — can bound those awaits against
/// this budget instead of inventing one.
///
/// Such a caller must spend what is LEFT of this budget. It must not take a
/// fresh copy. The post-drain awaits run after `until_shutdown` has already
/// spent part of the budget. They also run one after another. A fresh copy per
/// phase lets the total reach a multiple of what the deployment sized. The
/// kubelet then kills the pod mid-drain. `polyc-control-plane` stamps an
/// `Instant` before `until_shutdown`. It subtracts the elapsed time at each
/// later phase.
#[must_use]
pub 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)
}

/// Run `future` inside what is LEFT of the drain budget, measured from
/// `started`.
///
/// Post-drain work runs after [`until_shutdown`]. That call has already spent
/// part of the budget. Each later phase also runs after the last one. A phase
/// that takes its own copy of [`drain_grace`] lets the total reach a multiple
/// of what the deployment sized. The kubelet then kills the pod mid-drain.
///
/// Stamp one `Instant` before [`until_shutdown`]. Pass it to every phase. This
/// function computes the budget. A caller passes a future and a start time. It
/// never passes a duration. There is no budget argument to get wrong.
///
/// This does not make the rule unbreakable. A caller can still reach for
/// `tokio::time::timeout` with a duration of its own. The compiler accepts
/// that. What this removes is the easy mistake. A reviewer sees one call shape
/// in the shutdown path. Anything else stands out.
///
/// # Errors
///
/// Returns [`tokio::time::error::Elapsed`] when the remaining budget runs out
/// before `future` finishes.
///
/// The budget can be zero. `tokio::time::timeout` polls the future once before
/// it reads the delay. A future that is ready on that first poll still returns
/// `Ok`. A zero budget denies further time to become ready. It does not deny
/// the first poll.
pub async fn within_drain_budget<F: std::future::Future>(
    started: std::time::Instant,
    future: F,
) -> Result<F::Output, tokio::time::error::Elapsed> {
    tokio::time::timeout(drain_grace().saturating_sub(started.elapsed()), future).await
}

/// 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, within_drain_budget};
    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 a_spent_budget_denies_further_time_but_not_the_first_poll() {
        // `tokio::time::timeout` polls the future before it reads the delay, so
        // a zero budget still admits one poll. The doc on
        // `within_drain_budget` states this, and the shutdown path relies on
        // it: a lease release that can settle immediately still settles.
        let spent = std::time::Instant::now()
            .checked_sub(Duration::from_secs(3600))
            .expect("an hour before now is representable");

        let ready = within_drain_budget(spent, async { 7 })
            .await
            .expect("a future that is ready on the first poll survives a spent budget");
        assert_eq!(ready, 7);

        let pending = within_drain_budget(spent, std::future::pending::<()>()).await;
        assert!(
            pending.is_err(),
            "a pending future gets no further time once the budget is spent"
        );
    }

    #[tokio::test]
    async fn an_unspent_budget_lets_a_slow_future_finish() {
        let started = std::time::Instant::now();
        let done = within_drain_budget(started, async {
            tokio::time::sleep(Duration::from_millis(20)).await;
            "settled"
        })
        .await
        .expect("a short future fits inside a fresh budget");
        assert_eq!(done, "settled");
    }

    #[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"));
    }
}