//! `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 ''`. 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> { 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, scores_dir: PathBuf, shutdown: Arc, ) { 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) { 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) { tokio::spawn(async move { let _ = tokio::signal::ctrl_c().await; shutdown.store(true, Ordering::Relaxed); }); }