prop/rust/prop_grid_rs/src/metrics.rs

//! Prometheus metrics + `/metrics` HTTP endpoint.
//!
//! Scraped by the existing Prometheus server at 10.0.15.25. Exposed on
//! `METRICS_ADDR` (default `0.0.0.0:9100`).

use std::sync::LazyLock;
use std::time::Duration;

use axum::{extract::State, http::StatusCode, response::IntoResponse, routing::get, Router};
use prometheus::{
    register_histogram, register_histogram_vec, register_int_counter_vec, register_int_gauge,
    Encoder, Histogram, HistogramVec, IntCounterVec, IntGauge, TextEncoder,
};
use tracing::{info, warn};

/// End-to-end duration of one forecast-hour chain step (fetch → decode
/// → score → write). Labelled by outcome so p99 of successful steps can
/// be tracked independently from failure tails.
pub static CHAIN_STEP_DURATION: LazyLock<HistogramVec> = LazyLock::new(|| {
    register_histogram_vec!(
        "prop_grid_rs_chain_step_duration_seconds",
        "Wall time per grid_tasks chain step",
        &["outcome"],
        vec![5.0, 10.0, 20.0, 30.0, 45.0, 60.0, 90.0, 120.0, 180.0, 300.0, 600.0]
    )
    .expect("register histogram")
});

/// Count of chain steps by outcome. Makes the success/failure ratio
/// visible without having to derive it from the histogram.
pub static CHAIN_STEPS_TOTAL: LazyLock<IntCounterVec> = LazyLock::new(|| {
    register_int_counter_vec!(
        "prop_grid_rs_chain_steps_total",
        "Total grid_tasks chain steps processed",
        &["outcome"]
    )
    .expect("register counter")
});

/// Tasks currently in flight. A gauge so both "idle pod" (0) and
/// "fully saturated" (== PROP_GRID_RS_PARALLELISM) are distinguishable.
pub static TASKS_IN_FLIGHT: LazyLock<IntGauge> = LazyLock::new(|| {
    register_int_gauge!(
        "prop_grid_rs_tasks_in_flight",
        "Number of grid_tasks rows currently being processed"
    )
    .expect("register gauge")
});

/// Decode-only duration (the `spawn_blocking(wgrib2)` portion). Helps
/// answer "is wgrib2 fork overhead dominant?" without instrumenting the
/// full step.
pub static DECODE_DURATION: LazyLock<Histogram> = LazyLock::new(|| {
    register_histogram!(
        "prop_grid_rs_decode_duration_seconds",
        "wgrib2 subprocess duration per product (surface or pressure)",
        vec![0.5, 1.0, 2.0, 5.0, 10.0, 20.0, 30.0, 60.0]
    )
    .expect("register histogram")
});

/// Install the lazy statics so the `/metrics` endpoint lists every
/// series from the first scrape, even before any chain step runs.
pub fn init() {
    LazyLock::force(&CHAIN_STEP_DURATION);
    LazyLock::force(&CHAIN_STEPS_TOTAL);
    LazyLock::force(&TASKS_IN_FLIGHT);
    LazyLock::force(&DECODE_DURATION);
}

/// RAII helper: increment on construction, decrement on drop. Guarantees
/// the in-flight gauge returns to 0 even if the worker panics.
pub struct InFlightGuard;

impl InFlightGuard {
    pub fn new() -> Self {
        TASKS_IN_FLIGHT.inc();
        Self
    }
}

impl Drop for InFlightGuard {
    fn drop(&mut self) {
        TASKS_IN_FLIGHT.dec();
    }
}

impl Default for InFlightGuard {
    fn default() -> Self {
        Self::new()
    }
}

pub fn record_chain_step(duration: Duration, success: bool) {
    let outcome = if success { "success" } else { "failure" };
    CHAIN_STEP_DURATION
        .with_label_values(&[outcome])
        .observe(duration.as_secs_f64());
    CHAIN_STEPS_TOTAL.with_label_values(&[outcome]).inc();
}

async fn metrics_handler(State(()): State<()>) -> impl IntoResponse {
    let mut buf = Vec::new();
    let encoder = TextEncoder::new();
    let metric_families = prometheus::gather();
    match encoder.encode(&metric_families, &mut buf) {
        Ok(()) => (
            StatusCode::OK,
            [("content-type", TextEncoder::new().format_type().to_string())],
            buf,
        ),
        Err(e) => {
            warn!(error = %e, "metrics encode failed");
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                [("content-type", "text/plain".to_string())],
                Vec::new(),
            )
        }
    }
}

pub async fn serve(addr: std::net::SocketAddr) {
    init();
    let app = Router::new()
        .route("/metrics", get(metrics_handler))
        .route("/health", get(|| async { "ok" }))
        .with_state(());
    match tokio::net::TcpListener::bind(addr).await {
        Ok(listener) => {
            info!(%addr, "metrics server listening");
            if let Err(e) = axum::serve(listener, app).await {
                warn!(error = %e, "metrics server exited");
            }
        }
        Err(e) => {
            warn!(error = %e, %addr, "failed to bind metrics listener; continuing without /metrics");
        }
    }
}