//! The long-running collection loop and the one-shot commands. use crate::config::{Config, USER_AGENT}; use crate::dams::DamsClient; use crate::db; use crate::forecasting::{self, Forecaster, RunInfo}; use crate::hii::HiiClient; use crate::notify; use crate::rain::OpenMeteoClient; use crate::rid::{self, CycleOutcome, RidClient}; use crate::stations::Stations; use crate::util::{bangkok_now, floor_hour}; use anyhow::{Result, bail}; use chrono::{Duration, NaiveDateTime}; use ping_forecast::ModelStore; use ping_forecast::rain::RainSeries; use serde_json::json; use sqlx::PgPool; use std::sync::{Arc, RwLock}; pub struct Collector { pub cfg: Config, pub pool: PgPool, /// Shared with the website (the station CRUD API edits it). pub stations: Arc>, /// The last status written, shared with /scraping/status. pub status: Arc>, /// POST /scrape/trigger: poll now. pub trigger: Arc, pub rid: RidClient, hii: HiiClient, dams: DamsClient, rain: OpenMeteoClient, http: reqwest::Client, } #[derive(Default)] struct Stats { total_runs: u64, successful_runs: u64, failed_runs: u64, last_run: Option, next_run: Option, feed_current: Option, latest_reading: Option, readings_saved: usize, forecast: Option, } fn iso(t: Option) -> serde_json::Value { t.map_or(serde_json::Value::Null, |t| { json!(t.format("%Y-%m-%dT%H:%M:%S").to_string()) }) } impl Collector { pub fn new(cfg: Config, pool: PgPool, stations: Stations) -> Result { let http = reqwest::Client::builder() .user_agent(USER_AGENT) .connect_timeout(std::time::Duration::from_secs(20)) .build()?; Ok(Self { rid: RidClient::new(http.clone(), cfg.rid_api_url.clone()), hii: HiiClient::new(http.clone(), cfg.hii_api_base.clone()), dams: DamsClient::new(http.clone(), cfg.dams_url.clone()), rain: OpenMeteoClient::new(http.clone(), cfg.openmeteo_url.clone()), http, cfg, pool, stations: Arc::new(RwLock::new(stations)), status: Arc::new(RwLock::new(json!({"is_running": false}))), trigger: Arc::new(tokio::sync::Notify::new()), }) } /// The station mapping as of now (the API may edit it between cycles). pub fn stations(&self) -> Stations { self.stations.read().unwrap().clone() } /// The website/API state over this collector's pool, stations and status. pub fn web_state(&self, in_process: bool) -> crate::web::Shared { crate::web::AppState::new( self.cfg.clone(), self.pool.clone(), self.http.clone(), self.stations.clone(), in_process.then(|| self.status.clone()), in_process.then(|| self.trigger.clone()), ) } /// `ping-collector serve`: the website and API only; /scraping/status /// reads the status file the collector service writes. pub async fn serve(&self, listen: String) -> Result<()> { db::ensure_schema(&self.pool).await?; crate::web::serve(self.web_state(false), listen, shutdown_signal()).await } /// HII rain + water level, RID dams (today and yesterday), Open-Meteo. /// Each source is failure-isolated. pub async fn collect_hourly(&self, now: NaiveDateTime) { if self.cfg.hii_enabled { match self.hii.fetch_rain(self.cfg.hii_basin_code).await { Ok(r) => match db::save_hii_rain(&self.pool, &r, now).await { Ok(n) => tracing::info!("HII rainfall: {n} rows saved"), Err(e) => tracing::error!("HII rainfall save failed: {e:#}"), }, Err(e) => tracing::error!("HII rainfall fetch failed: {e:#}"), } match self.hii.fetch_waterlevel(self.cfg.hii_basin_code).await { Ok(r) => match db::save_hii_waterlevel(&self.pool, &r, now).await { Ok(n) => tracing::info!("HII water level: {n} rows saved"), Err(e) => tracing::error!("HII water level save failed: {e:#}"), }, Err(e) => tracing::error!("HII water level fetch failed: {e:#}"), } } if self.cfg.dams_enabled { for date in [None, Some(now.date() - Duration::days(1))] { match self.dams.fetch_day(date).await { Ok(r) => match db::save_dams(&self.pool, &r, now).await { Ok(n) => tracing::info!("RID dams {date:?}: {n} dam-day rows saved"), Err(e) => tracing::error!("RID dams save failed: {e:#}"), }, Err(e) => tracing::error!("RID dams fetch {date:?} failed: {e:#}"), } } } if self.cfg.rain_enabled { match self.rain.fetch_forecast(5, 2).await { Ok(r) => match db::save_rain(&self.pool, &r).await { Ok(n) => tracing::info!("Open-Meteo rain: {n} hourly rows upserted"), Err(e) => tracing::error!("Open-Meteo save failed: {e:#}"), }, Err(e) => tracing::error!("Open-Meteo fetch failed: {e:#}"), } } } fn write_status(&self, stats: &Stats, running: bool) { let body = json!({ "is_running": running, "total_runs": stats.total_runs, "successful_runs": stats.successful_runs, "failed_runs": stats.failed_runs, "last_run": iso(stats.last_run), "next_run": iso(stats.next_run), "feed_current": stats.feed_current, "latest_reading": iso(stats.latest_reading), "readings_saved": stats.readings_saved, "collector": concat!("ping-collector ", env!("CARGO_PKG_VERSION")), "forecast": stats.forecast.as_ref().map(|f| json!({ "last_run": iso(Some(f.at)), "rows": f.rows, "model_rows": f.model_rows, "models": f.models, "notifications_sent": f.notified, })), }); *self.status.write().unwrap() = body.clone(); let Some(path) = &self.cfg.status_file else { return; }; let tmp = path.with_extension("json.tmp"); let result = std::fs::write(&tmp, body.to_string()).and_then(|_| std::fs::rename(&tmp, path)); if let Err(e) = result { tracing::warn!("could not write status file {}: {e}", path.display()); } } /// The service: startup gap fill, then RID polls on the hour (+offset) /// with retries while RID is behind, and the hourly sources once per /// clock hour. Returns when a shutdown signal arrives. pub async fn run(&self) -> Result<()> { db::ensure_schema(&self.pool).await?; let shutdown = shutdown_signal(); tokio::pin!(shutdown); // the website, when WEB_LISTEN is set: stopped with the collector let (stop_web, web_stopped) = tokio::sync::watch::channel(false); let web = self.cfg.web.listen.clone().map(|listen| { let state = self.web_state(true); let mut rx = web_stopped.clone(); tokio::spawn(async move { let stop = async move { let _ = rx.wait_for(|stop| *stop).await; }; if let Err(e) = crate::web::serve(state, listen, stop).await { tracing::error!("website stopped: {e:#}"); } }) }); let result = self.collect_loop(&mut shutdown).await; let _ = stop_web.send(true); if let Some(handle) = web { let _ = tokio::time::timeout(std::time::Duration::from_secs(10), handle).await; } result } async fn collect_loop( &self, shutdown: &mut std::pin::Pin<&mut impl std::future::Future>, ) -> Result<()> { if self.cfg.gap_fill_days_on_start > 0 { let days = Some(self.cfg.gap_fill_days_on_start); let stations = self.stations(); tokio::select! { r = rid::fill_gaps(&self.rid, &self.pool, &stations, days, bangkok_now()) => { if let Err(e) = r { tracing::error!("startup gap fill failed: {e:#}") } } _ = &mut *shutdown => return Ok(()), } } let mut stats = Stats::default(); let mut last_hourly_slot: Option = None; let mut forecaster = self.cfg.forecast_enabled.then(|| { tracing::info!( "FORECAST_IN_COLLECTOR: forecasting from {} and notifying here", self.cfg.models_dir.display() ); Forecaster::new(&self.cfg, self.http.clone()) }); loop { let start = bangkok_now(); self.write_status(&stats, true); let outcome: CycleOutcome = rid::run_cycle(&self.rid, &self.pool, &self.stations(), start).await; stats.total_runs += 1; stats.last_run = Some(start); if outcome.current { stats.successful_runs += 1; } else { stats.failed_runs += 1; } stats.feed_current = Some(outcome.current); stats.readings_saved = outcome.saved; if outcome.latest.is_some() { stats.latest_reading = outcome.latest.max(stats.latest_reading); } let slot = floor_hour(bangkok_now()); if last_hourly_slot != Some(slot) { last_hourly_slot = Some(slot); self.collect_hourly(bangkok_now()).await; } if let Some(f) = forecaster.as_mut() && let Some(info) = f.maybe_run(&self.pool, &self.rain, bangkok_now()).await { stats.forecast = Some(info); } let now = bangkok_now(); let next = rid::next_poll_time( now, outcome.current, self.cfg.poll_offset_minutes, self.cfg.retry_minutes, ); stats.next_run = Some(next); self.write_status(&stats, false); tracing::info!( "next RID poll at {}{}", next.format("%H:%M"), if outcome.current { "" } else { " (current hour not published yet)" } ); let wait = (next - now) .to_std() .unwrap_or_default() .max(std::time::Duration::from_secs(30)); tokio::select! { _ = tokio::time::sleep(wait) => {} _ = self.trigger.notified() => tracing::info!("collection triggered via the API"), _ = &mut *shutdown => { tracing::info!("shutdown signal received"); return Ok(()); } } } } /// One cycle of every source; with `dry_run` nothing is written and the /// RID fetch is compared against what the database already holds. pub async fn once(&self, dry_run: bool) -> Result<()> { let now = bangkok_now(); if !dry_run { db::ensure_schema(&self.pool).await?; let o = rid::run_cycle(&self.rid, &self.pool, &self.stations(), now).await; println!( "RID: reports {:?}, {} fetched, {} saved, latest {:?}, current hour in: {}{}", o.dates, o.fetched, o.saved, o.latest, o.current, o.error.map(|e| format!(", error: {e}")).unwrap_or_default() ); self.collect_hourly(now).await; return Ok(()); } let dates = rid::cycle_dates(&self.pool, now).await; let (readings, error) = rid::fetch_dates(&self.rid, &dates).await; let stored = db::count_stored_rid(&self.pool, &readings).await?; let latest = readings.iter().map(|r| r.timestamp).max(); let db_latest = db::latest_rid_timestamp(&self.pool).await?; println!( "RID (dry run): reports {dates:?}: {} readings fetched, {stored} already stored, {} would be added; latest fetched {latest:?}, latest stored {db_latest:?}{}", readings.len(), readings.len() as i64 - stored, error.map(|e| format!(", error: {e}")).unwrap_or_default() ); if self.cfg.hii_enabled { match self.hii.fetch_rain(self.cfg.hii_basin_code).await { Ok(r) => println!("HII rainfall (dry run): {} Ping-basin stations", r.len()), Err(e) => println!("HII rainfall (dry run): {e:#}"), } match self.hii.fetch_waterlevel(self.cfg.hii_basin_code).await { Ok(r) => println!("HII water level (dry run): {} Ping-basin stations", r.len()), Err(e) => println!("HII water level (dry run): {e:#}"), } } if self.cfg.dams_enabled { match self.dams.fetch_day(None).await { Ok(r) => println!("RID dams (dry run): {} dams today", r.len()), Err(e) => println!("RID dams (dry run): {e:#}"), } } if self.cfg.rain_enabled { match self.rain.fetch_forecast(5, 2).await { Ok(r) => println!("Open-Meteo (dry run): {} hourly rows", r.len()), Err(e) => println!("Open-Meteo (dry run): {e:#}"), } } Ok(()) } } impl Collector { /// `ping-collector forecast`: compute from the database now (or at /// `now`), print a table or the /forecast JSON, store only with `save`. pub async fn forecast( &self, json_out: bool, save: bool, now: Option, rain: Option, ) -> Result<()> { let mut store = ModelStore::new(&self.cfg.models_dir); let models = store.refresh()?; if models.is_empty() { bail!( "no exported models in {} (run `python -m src.ml.export`)", self.cfg.models_dir.display() ); } let now = now.unwrap_or_else(bangkok_now); let snap = forecasting::compute(&self.pool, &self.rain, &models, now, rain).await?; if json_out { println!("{}", serde_json::to_string(&snap.forecasts)?); } else { println!( "{:<6} {:>3} {:<9} {:<16} {:>7} {:>7} {:>9} {:>8}", "code", "h", "source", "as_of", "level", "peak", "p_warning", "p_danger" ); for f in &snap.forecasts { println!( "{:<6} {:>3} {:<9} {:<16} {:>7.2} {:>7.2} {:>9.3} {:>8.3}", f.station_code, f.horizon_hours, f.source, f.as_of.format("%Y-%m-%d %H:%M"), f.current_level, f.predicted_max_level, f.p_warning, f.p_danger ); } eprintln!( "{} models from {}; {}", models.len(), self.cfg.models_dir.display(), forecasting::summarize(&snap) ); } if save { db::ensure_schema(&self.pool).await?; let n = db::save_forecast_history(&self.pool, &snap.forecasts).await?; db::save_forecast_current(&self.pool, now, &snap.forecasts).await?; eprintln!("stored: {n} forecast_history rows, forecast_current"); } Ok(()) } /// `ping-collector notify`: what would be (or, without `dry_run`, is) /// published right now against the stored notification state. pub async fn notify(&self, dry_run: bool) -> Result<()> { if !dry_run && !self.cfg.forecast_enabled { bail!( "the Python app notifies while FORECAST_IN_COLLECTOR is not true; \ use --dry-run" ); } let Some(notifier) = forecasting::notifier(&self.cfg, self.http.clone()) else { bail!( "notifications are off: neither NTFY_SERVER nor MATRIX_ACCESS_TOKEN/MATRIX_ROOM_ID is set" ); }; let mut store = ModelStore::new(&self.cfg.models_dir); let models = store.refresh()?; let now = bangkok_now(); let snap = forecasting::compute(&self.pool, &self.rain, &models, now, None).await?; let readings = notify::readings_with_hii(&db::latest_rid_readings(&self.pool).await?, &snap.fills); for r in &readings { println!( "{:<6} {} {:>7} {:?}{}", r.station_code, r.timestamp.format("%Y-%m-%d %H:%M"), r.water_level.map_or("-".into(), |l| format!("{l:.2}")), r.source, r.discharge_percent .map_or(String::new(), |p| format!(" {p:.0}%")) ); } db::ensure_schema(&self.pool).await?; if dry_run { let states = db::notification_states(&self.pool).await?; let mut plan = vec![]; if notifier.ntfy.is_some() { plan.extend( notify::evaluate(&readings, &snap.forecasts, &states, ¬ifier.settings, now) .into_iter() .map(|p| (p.notification.topic.clone(), p)), ); } if notifier.matrix.is_some() { let mstates = notifier .matrix_states(&self.pool, &states, now, false) .await?; plan.extend( notifier .matrix_plan(&readings, &snap.forecasts, &mstates, now) .into_iter() .map(|p| ("matrix".to_string(), p)), ); } if plan.is_empty() { println!("no transitions: nothing would be published"); } for (to, p) in plan { println!( "would publish {to}: {} -- {}", p.notification.title, p.notification.message ); } } else { for n in notifier .run(&self.pool, &readings, &snap.forecasts, now) .await? { println!("published {}: {}", n.topic, n.title); } } Ok(()) } } async fn shutdown_signal() { #[cfg(unix)] { use tokio::signal::unix::{SignalKind, signal}; let mut term = signal(SignalKind::terminate()).expect("SIGTERM handler"); tokio::select! { _ = tokio::signal::ctrl_c() => {} _ = term.recv() => {} } } #[cfg(not(unix))] { let _ = tokio::signal::ctrl_c().await; } }