perf(prop-grid-rs): claim N forecast hours concurrently

Replaces the serial claim→process loop with a JoinSet of N worker tasks,
each running the same claim→fetch→decode→score→complete flow against
grid_tasks. FOR UPDATE SKIP LOCKED already prevents contention between
workers so no additional coordination is needed.

talos5 is 4c/4t (i5-7260U); PROP_GRID_RS_PARALLELISM=4 turns the full
f01..f18 chain from ~18 min serial into roughly ~5 min. Per-task peak
memory after the streamed-bands refactor is ~250 Mi, so 4 concurrent is
~1 Gi — well under the 2 Gi container limit.

Postgres pool sized to parallelism + 2 so NOTIFY and an opportunistic
claim retry share the pool without blocking a worker transaction.
This commit is contained in:
Graham McIntire 2026-04-19 17:31:12 -05:00
parent 7b1de3f995
commit 022e6928bc
No known key found for this signature in database
GPG key ID: F4ABF488E6029E59
2 changed files with 61 additions and 14 deletions

View file

@ -61,8 +61,17 @@ spec:
value: "/data/scores"
- name: RUST_LOG
value: "info"
- name: PROP_GRID_RS_PG_CONNS
# talos5 is 4c/4t (i5-7260U). Claiming 4 forecast hours in
# parallel drops full-chain wall time from ~18 min serial to
# ~5 min. Each concurrent task peaks around 250 Mi after the
# streamed-bands refactor; 4×250 ≈ 1 Gi is comfortably under
# the 2 Gi memory limit.
- name: PROP_GRID_RS_PARALLELISM
value: "4"
# Pool size = parallelism + 2 (NOTIFY + retry slack). Kept in
# sync with PROP_GRID_RS_PARALLELISM when either changes.
- name: PROP_GRID_RS_PG_CONNS
value: "6"
envFrom:
- secretRef:
name: prop-secrets

View file

@ -2,11 +2,13 @@
//!
//! Flow:
//! 1. Connect to Postgres (DATABASE_URL env).
//! 2. Loop: claim one `grid_tasks` row (fh>0, status='queued'),
//! run the fetch → decode → score → write-scores-file chain, mark done +
//! `NOTIFY propagation_ready '<valid_time iso>'`. On error mark
//! 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: requeue current row and exit.
//! 3. Graceful shutdown on SIGTERM: workers exit after their current task.
use std::path::PathBuf;
use std::sync::Arc;
@ -14,6 +16,7 @@ use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use prop_grid_rs::{db, fetcher::HrrrClient, pipeline};
use tokio::task::JoinSet;
use tracing::{error, info, warn};
use tracing_subscriber::EnvFilter;
@ -33,10 +36,18 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
.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(4);
.unwrap_or((parallelism as u32) + 2);
let pool = db::connect(&database_url, max_conn).await?;
let client = Arc::new(HrrrClient::new(hrrr_base)?);
@ -44,8 +55,37 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let shutdown = Arc::new(AtomicBool::new(false));
spawn_signal_handler(shutdown.clone());
info!(scores_dir = %scores_dir.display(), "prop-grid-rs started");
info!(
scores_dir = %scores_dir.display(),
parallelism,
pg_conns = max_conn,
"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) {
match db::claim_next(&pool).await {
Ok(Some(task)) => {
@ -62,6 +102,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
{
Ok(stats) => {
info!(
worker_id,
task_id = %task.id,
run_time = %task.run_time,
forecast_hour = task.forecast_hour,
@ -70,13 +111,13 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
"chain step done"
);
if let Err(e) = db::complete(&pool, &task).await {
error!(%task_id, error = %e, "failed to mark task done");
error!(worker_id, %task_id, error = %e, "failed to mark task done");
}
}
Err(e) => {
warn!(%task_id, error = %e, "chain step failed");
warn!(worker_id, %task_id, error = %e, "chain step failed");
if let Err(dbe) = db::fail(&pool, task_id, &e.to_string()).await {
error!(%task_id, error = %dbe, "failed to mark task failed");
error!(worker_id, %task_id, error = %dbe, "failed to mark task failed");
}
}
}
@ -85,14 +126,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
tokio::time::sleep(IDLE_SLEEP).await;
}
Err(e) => {
error!(error = %e, "claim_next failed; backing off");
error!(worker_id, error = %e, "claim_next failed; backing off");
tokio::time::sleep(IDLE_SLEEP).await;
}
}
}
info!("prop-grid-rs shutting down");
Ok(())
}
#[cfg(unix)]