prop/rust/prop_grid_rs/src/bin/worker.rs

//! `prop-grid-rs` main loop.
//!
//! Flow:
//! 1. Connect to Postgres (DATABASE_URL env).
//! 2. Spawn N concurrent worker tasks (`PROP_GRID_RS_PARALLELISM`, default
//!    = num_cpus). Each worker independently claims one `grid_tasks` row
//!    (fh>0, status='queued') via `FOR UPDATE SKIP LOCKED`, runs the
//!    fetch → decode → score → write-scores-file chain, and marks done +
//!    `NOTIFY propagation_ready '<valid_time iso>'`. On error marks
//!    failed; hourly cron reseeds.
//! 3. Graceful shutdown on SIGTERM: workers exit after their current task.

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

use prop_grid_rs::{db, fetcher::HrrrClient, metrics, pipeline, telemetry};
use tokio::task::JoinSet;
use tracing::{error, info, warn};

const IDLE_SLEEP: Duration = Duration::from_secs(5);

enum ChainOutcome {
    Forecast(pipeline::ChainStepStats),
    Analysis(pipeline::AnalysisStepStats),
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let _telemetry = telemetry::init("prop-grid-rs");

    let database_url = std::env::var("DATABASE_URL")
        .map_err(|_| "DATABASE_URL environment variable is required")?;
    let scores_dir = std::env::var("PROP_SCORES_DIR")
        .map(PathBuf::from)
        .unwrap_or_else(|_| PathBuf::from("/data/scores"));
    let hrrr_base = std::env::var("HRRR_BASE_URL")
        .unwrap_or_else(|_| prop_grid_rs::fetcher::HRRR_BASE_DEFAULT.to_string());
    let parallelism: usize = std::env::var("PROP_GRID_RS_PARALLELISM")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or_else(num_cpus::get)
        .max(1);
    // +2 slots over parallelism: each worker holds one connection during
    // claim/complete transactions; the spare lets NOTIFY and an opportunistic
    // retry share the pool without blocking.
    let max_conn: u32 = std::env::var("PROP_GRID_RS_PG_CONNS")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or((parallelism as u32) + 2);

    let pool = db::connect(&database_url, max_conn).await?;
    let client = Arc::new(HrrrClient::new(hrrr_base)?);

    let shutdown = Arc::new(AtomicBool::new(false));
    spawn_signal_handler(shutdown.clone());

    // Metrics endpoint binds on a separate port so scrape traffic
    // never interferes with the worker's NFS writes or Postgres calls.
    let metrics_addr: SocketAddr = std::env::var("METRICS_ADDR")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or_else(|| "0.0.0.0:9100".parse().expect("valid default addr"));
    tokio::spawn(async move {
        metrics::serve(metrics_addr).await;
    });

    info!(
        scores_dir = %scores_dir.display(),
        parallelism,
        pg_conns = max_conn,
        metrics_addr = %metrics_addr,
        "prop-grid-rs started"
    );

    let mut set: JoinSet<()> = JoinSet::new();
    for worker_id in 0..parallelism {
        let pool = pool.clone();
        let client = client.clone();
        let scores_dir = scores_dir.clone();
        let shutdown = shutdown.clone();
        set.spawn(async move {
            worker_loop(worker_id, pool, client, scores_dir, shutdown).await;
        });
    }

    while set.join_next().await.is_some() {}

    info!("prop-grid-rs shutting down");
    Ok(())
}

async fn worker_loop(
    worker_id: usize,
    pool: sqlx::PgPool,
    client: Arc<HrrrClient>,
    scores_dir: PathBuf,
    shutdown: Arc<AtomicBool>,
) {
    while !shutdown.load(Ordering::Relaxed) {
        // Analysis tasks (kind='analysis', f00) get priority — a /map
        // user values the current hour more than +1..+18h. Forecast is
        // only attempted when there is no analysis row ready.
        let claimed = match db::claim_next_analysis(&pool).await {
            Ok(Some(task)) => Ok(Some(task)),
            Ok(None) => db::claim_next(&pool).await,
            Err(e) => Err(e),
        };

        match claimed {
            Ok(Some(task)) => {
                let task_id = task.id;
                let _guard = metrics::InFlightGuard::new();
                let started = Instant::now();
                let result = match task.kind {
                    db::TaskKind::Forecast => pipeline::run_chain_step(
                        &client,
                        &scores_dir,
                        &pipeline::ChainStepInput {
                            run_time: task.run_time,
                            forecast_hour: task.forecast_hour as u8,
                        },
                    )
                    .await
                    .map(ChainOutcome::Forecast),
                    db::TaskKind::Analysis => pipeline::run_analysis_step(
                        &client,
                        &pool,
                        &scores_dir,
                        &pipeline::AnalysisStepInput {
                            run_time: task.run_time,
                        },
                    )
                    .await
                    .map(ChainOutcome::Analysis),
                };

                match result {
                    Ok(outcome) => {
                        let elapsed = started.elapsed();
                        metrics::record_chain_step(elapsed, true);
                        match &outcome {
                            ChainOutcome::Forecast(stats) => info!(
                                worker_id,
                                task_id = %task.id,
                                kind = "forecast",
                                run_time = %task.run_time,
                                forecast_hour = task.forecast_hour,
                                files = stats.score_files_written,
                                points = stats.point_count,
                                duration_s = elapsed.as_secs_f64(),
                                "chain step done"
                            ),
                            ChainOutcome::Analysis(stats) => info!(
                                worker_id,
                                task_id = %task.id,
                                kind = "analysis",
                                run_time = %task.run_time,
                                files = stats.score_files_written,
                                points = stats.point_count,
                                profile_cells = stats.profile_cells_written,
                                duct_cells = stats.duct_cells,
                                nexrad_cells = stats.nexrad_cells_with_echo,
                                commercial_cells = stats.commercial_cells_boosted,
                                duration_s = elapsed.as_secs_f64(),
                                "analysis step done"
                            ),
                        }
                        if let Err(e) = db::complete(&pool, &task).await {
                            error!(worker_id, %task_id, error = %e, "failed to mark task done");
                        }
                    }
                    Err(e) => {
                        metrics::record_chain_step(started.elapsed(), false);
                        warn!(worker_id, %task_id, error = %e, "chain step failed");
                        if let Err(dbe) = db::fail(&pool, task_id, &e.to_string()).await {
                            error!(worker_id, %task_id, error = %dbe, "failed to mark task failed");
                        }
                    }
                }
            }
            Ok(None) => {
                tokio::time::sleep(IDLE_SLEEP).await;
            }
            Err(e) => {
                error!(worker_id, error = %e, "claim_next failed; backing off");
                tokio::time::sleep(IDLE_SLEEP).await;
            }
        }
    }
}

#[cfg(unix)]
fn spawn_signal_handler(shutdown: Arc<AtomicBool>) {
    use tokio::signal::unix::{signal, SignalKind};
    tokio::spawn(async move {
        let mut term = signal(SignalKind::terminate()).expect("install SIGTERM handler");
        let mut intr = signal(SignalKind::interrupt()).expect("install SIGINT handler");
        tokio::select! {
            _ = term.recv() => info!("SIGTERM received"),
            _ = intr.recv() => info!("SIGINT received"),
        };
        shutdown.store(true, Ordering::Relaxed);
    });
}

#[cfg(not(unix))]
fn spawn_signal_handler(shutdown: Arc<AtomicBool>) {
    tokio::spawn(async move {
        let _ = tokio::signal::ctrl_c().await;
        shutdown.store(true, Ordering::Relaxed);
    });
}