polyc_runtime/

health.rs

//! The `axum` side-server (PRD §11): liveness, readiness, and metrics on a
//! separate port from a binary's main surface.

use std::{
    net::SocketAddr,
    sync::{
        Arc, LazyLock,
        atomic::{AtomicBool, Ordering},
    },
    time::Duration,
};

use axum::{Router, extract::State, http::StatusCode, response::IntoResponse, routing::get};
use prometheus::{Encoder, IntGaugeVec, TextEncoder, register_int_gauge_vec};
use tokio_util::sync::CancellationToken;

/// `1` while the process is up, labelled with the build's release version, so a
/// `/metrics` scrape reports which release every process runs. This is the
/// always-on, passive half of update detection: an operator scrapes it today
/// and, paired with `polychrome_update_available` (control plane only), learns
/// when a newer release exists. See docs/reference/upgrade-operator-decision.md.
///
/// The version is the lockstep workspace version (`version.workspace = true`),
/// so every binary built from this workspace reports the same `vX.Y.Z`.
/// Registered to the default registry so it is always present in a scrape.
static BUILD_INFO: LazyLock<IntGaugeVec> = LazyLock::new(|| {
    let gauge = register_int_gauge_vec!(
        "polychrome_build_info",
        "Always 1 for a live process; the `version` label carries the build's release version.",
        &["version"]
    )
    .expect("register polychrome_build_info");
    gauge.with_label_values(&[env!("CARGO_PKG_VERSION")]).set(1);
    gauge
});

/// Default delay `drain` sleeps after firing the drain hook, before
/// responding 200 — long enough that the endpoints controller / kube-proxy
/// observes the readiness flip (see [`Health::with_drain_propagation_delay`]
/// for why this exists) before a `preStop` caller's response returns and
/// SIGTERM follows.
const DEFAULT_DRAIN_PROPAGATION_DELAY: Duration = Duration::from_secs(5);

/// A hook a binary installs to run its own drain-specific work (e.g.
/// flipping a native gRPC health check to `NOT_SERVING`, refusing new
/// work admission) the first time `/drain` is hit. Boxed as `Arc<dyn Fn>`
/// so [`Health`] stays `Clone` and generic over whatever a binary needs to
/// do — this crate has no knowledge of gRPC health checks, lease
/// acquisition, or any other binary-specific concern.
type DrainHook = Arc<dyn Fn() + Send + Sync>;

/// Shared readiness flag, flipped on once the process is wired up and flipped
/// off as shutdown begins.
#[derive(Clone)]
pub struct Health {
    ready: Arc<AtomicBool>,
    /// Set once `/drain` has been hit; `/readyz` returns 503 unconditionally
    /// while this is `true`, regardless of `ready`.
    draining: Arc<AtomicBool>,
    /// Optional binary-provided hook, fired once (on the FIRST `/drain` hit
    /// only — see `drain`).
    on_drain: Option<DrainHook>,
    /// How long `/drain` sleeps before responding; see
    /// [`Health::with_drain_propagation_delay`].
    drain_propagation_delay: Duration,
}

impl Default for Health {
    fn default() -> Self {
        Self {
            ready: Arc::default(),
            draining: Arc::default(),
            on_drain: None,
            drain_propagation_delay: DEFAULT_DRAIN_PROPAGATION_DELAY,
        }
    }
}

impl Health {
    /// A fresh, not-yet-ready health handle.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Set readiness; `/readyz` returns 200 only while `true` AND the
    /// process is not draining (see [`Health::is_draining`]).
    pub fn set_ready(&self, ready: bool) {
        self.ready.store(ready, Ordering::Relaxed);
    }

    /// Install a hook `/drain` fires once, the first time it's hit — never
    /// on a repeat call, so a lifecycle-hook retry or a human re-`curl`ing
    /// `/drain` doesn't repeat binary-specific side effects (e.g. don't
    /// flip a gRPC health check twice). Builder-style.
    #[must_use]
    pub fn with_on_drain(mut self, hook: DrainHook) -> Self {
        self.on_drain = Some(hook);
        self
    }

    /// Override the delay `drain` sleeps before responding (production
    /// default: `DEFAULT_DRAIN_PROPAGATION_DELAY`, 5s). Tests set this to
    /// [`Duration::ZERO`] to keep the delay out of the critical path.
    /// Builder-style.
    #[must_use]
    pub const fn with_drain_propagation_delay(mut self, delay: Duration) -> Self {
        self.drain_propagation_delay = delay;
        self
    }

    /// `true` once `/drain` has been hit at least once.
    #[must_use]
    pub fn is_draining(&self) -> bool {
        self.draining.load(Ordering::Relaxed)
    }
}

/// Serve the side-server until `shutdown` is cancelled, binding `addr` here.
///
/// Callers that gate readiness on their listeners actually binding should
/// bind up front and use [`serve_on`] instead.
///
/// # Errors
///
/// Returns an error if the listener cannot bind or `axum` serving fails.
pub async fn serve(
    addr: SocketAddr,
    health: Health,
    shutdown: CancellationToken,
) -> anyhow::Result<()> {
    let listener = tokio::net::TcpListener::bind(addr).await?;
    serve_on(listener, health, shutdown).await
}

/// Serve the side-server on an already-bound listener until `shutdown` is
/// cancelled. Binding is the caller's job, so readiness can be flipped on
/// only once the socket actually exists.
///
/// # Errors
///
/// Returns an error if `axum` serving fails.
pub async fn serve_on(
    listener: tokio::net::TcpListener,
    health: Health,
    shutdown: CancellationToken,
) -> anyhow::Result<()> {
    let app = Router::new()
        .route("/healthz", get(|| async { StatusCode::OK }))
        .route("/livez", get(|| async { StatusCode::OK }))
        .route("/readyz", get(readyz))
        .route("/metrics", get(metrics))
        // Both verbs: Kubernetes' `preStop.httpGet` lifecycle hook issues a
        // GET (`httpGet` has no method override); POST is for a human or
        // script draining a pod out-of-band ahead of a manual restart.
        .route("/drain", get(drain).post(drain))
        .with_state(health);

    let addr = listener.local_addr()?;
    tracing::info!(%addr, "side-server listening");
    axum::serve(listener, app)
        .with_graceful_shutdown(async move { shutdown.cancelled().await })
        .await?;
    Ok(())
}

/// 200 once the process is ready, else 503. Draining always wins: once
/// `/drain` has fired, `/readyz` reports 503 unconditionally, even if
/// `set_ready(true)` is called afterward — there is no coming back from
/// draining within a process's lifetime.
async fn readyz(State(health): State<Health>) -> impl IntoResponse {
    if health.draining.load(Ordering::Relaxed) {
        return StatusCode::SERVICE_UNAVAILABLE;
    }
    if health.ready.load(Ordering::Relaxed) {
        StatusCode::OK
    } else {
        StatusCode::SERVICE_UNAVAILABLE
    }
}

/// Begin draining: flip the draining flag (idempotent — only the first call
/// fires the hook), fire the binary-provided [`DrainHook`] if one was
/// installed, then sleep [`Health::drain_propagation_delay`] before
/// responding.
///
/// The sleep matters for the `preStop` use case specifically: Kubernetes
/// only sends SIGTERM to the container AFTER the `preStop` hook's HTTP call
/// returns, but the endpoints controller learns the pod is `NotReady`
/// asynchronously (it has to observe the `/readyz` flip via its own probe
/// cycle, then reprogram kube-proxy / the Service's endpoint slice). Without
/// this delay, SIGTERM — and the bounded shutdown drain it starts — can
/// arrive before traffic has actually stopped being routed to this pod,
/// so a very short window of new connections could still land on a pod
/// that's about to stop accepting them.
async fn drain(State(health): State<Health>) -> impl IntoResponse {
    let already_draining = health.draining.swap(true, Ordering::SeqCst);
    if !already_draining && let Some(hook) = &health.on_drain {
        hook();
    }
    if !health.drain_propagation_delay.is_zero() {
        tokio::time::sleep(health.drain_propagation_delay).await;
    }
    (StatusCode::OK, "draining\n")
}

/// Prometheus scrape target: encodes the default registry (process build info
/// plus whatever the binary has registered).
async fn metrics() -> impl IntoResponse {
    LazyLock::force(&BUILD_INFO);
    let families = prometheus::gather();
    let mut buf = Vec::new();
    let encoder = TextEncoder::new();
    if encoder.encode(&families, &mut buf).is_err() {
        return (StatusCode::INTERNAL_SERVER_ERROR, Vec::new()).into_response();
    }
    (
        [(axum::http::header::CONTENT_TYPE, encoder.format_type())],
        buf,
    )
        .into_response()
}

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

    use std::sync::atomic::{AtomicUsize, Ordering};

    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tokio_util::sync::CancellationToken;

    use super::{Duration, Health, serve_on};

    /// One HTTP/1.1 request over a raw socket; returns the status line.
    async fn status_line(addr: std::net::SocketAddr, path: &str) -> String {
        method_status_line(addr, "GET", path).await
    }

    /// Like [`status_line`], but with an explicit HTTP method — used to
    /// prove `/drain` answers both GET (the Kubernetes `preStop.httpGet`
    /// shape) and POST (a human/script hitting it out-of-band).
    async fn method_status_line(addr: std::net::SocketAddr, method: &str, path: &str) -> String {
        let mut stream = tokio::net::TcpStream::connect(addr).await.unwrap();
        stream
            .write_all(
                format!(
                    "{method} {path} HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\nContent-Length: 0\r\n\r\n"
                )
                .as_bytes(),
            )
            .await
            .unwrap();
        let mut buf = String::new();
        stream.read_to_string(&mut buf).await.unwrap();
        buf.lines().next().unwrap_or_default().to_owned()
    }

    #[tokio::test]
    async fn serve_on_takes_a_prebound_listener_and_readyz_tracks_the_flag() {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let health = Health::new();
        let shutdown = CancellationToken::new();
        let srv = tokio::spawn(serve_on(listener, health.clone(), shutdown.clone()));

        // The socket exists before the serving task is even polled (the
        // bind-then-ready contract); readiness stays the flag's job.
        assert!(status_line(addr, "/readyz").await.contains("503"));
        assert!(status_line(addr, "/livez").await.contains("200"));
        health.set_ready(true);
        assert!(status_line(addr, "/readyz").await.contains("200"));
        health.set_ready(false);
        assert!(status_line(addr, "/readyz").await.contains("503"));

        shutdown.cancel();
        srv.await.unwrap().unwrap();
    }

    /// GET `/drain` flips readiness to 503 (even though `set_ready(true)`
    /// was never undone) and fires the installed hook exactly once — a
    /// second call (GET or POST) must not fire it again. The propagation
    /// delay is zeroed so the test doesn't pay it.
    #[tokio::test]
    async fn drain_flips_readiness_and_fires_hook_once() {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let hook_calls = std::sync::Arc::new(AtomicUsize::new(0));
        let health = Health::new()
            .with_drain_propagation_delay(Duration::ZERO)
            .with_on_drain({
                let hook_calls = hook_calls.clone();
                std::sync::Arc::new(move || {
                    hook_calls.fetch_add(1, Ordering::SeqCst);
                })
            });
        health.set_ready(true);
        let shutdown = CancellationToken::new();
        let srv = tokio::spawn(serve_on(listener, health.clone(), shutdown.clone()));

        assert!(status_line(addr, "/readyz").await.contains("200"));
        assert!(!health.is_draining());

        assert!(
            method_status_line(addr, "GET", "/drain")
                .await
                .contains("200")
        );
        assert_eq!(hook_calls.load(Ordering::SeqCst), 1);
        assert!(health.is_draining());
        assert!(
            status_line(addr, "/readyz").await.contains("503"),
            "readyz must report 503 while draining even though set_ready(true) was never undone"
        );

        // A second hit — via POST this time — must be idempotent: no second
        // hook call, still draining.
        assert!(
            method_status_line(addr, "POST", "/drain")
                .await
                .contains("200")
        );
        assert_eq!(
            hook_calls.load(Ordering::SeqCst),
            1,
            "the drain hook must fire only once, on the first /drain hit"
        );
        assert!(status_line(addr, "/readyz").await.contains("503"));

        shutdown.cancel();
        srv.await.unwrap().unwrap();
    }

    /// `/metrics` exposes `polychrome_build_info` carrying the build version as
    /// a label — the surface an operator scrapes to learn which release runs.
    #[tokio::test]
    async fn metrics_reports_build_info_with_version_label() {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let shutdown = CancellationToken::new();
        let srv = tokio::spawn(serve_on(listener, Health::new(), shutdown.clone()));

        let mut stream = tokio::net::TcpStream::connect(addr).await.unwrap();
        stream
            .write_all(b"GET /metrics HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n")
            .await
            .unwrap();
        let mut buf = String::new();
        stream.read_to_string(&mut buf).await.unwrap();

        let expected = format!(
            "polychrome_build_info{{version=\"{}\"}} 1",
            env!("CARGO_PKG_VERSION")
        );
        assert!(
            buf.contains(&expected),
            "scrape missing labelled build_info:\n{buf}"
        );

        shutdown.cancel();
        srv.await.unwrap().unwrap();
    }
}