//! HRRR fetcher. 1:1 port of the grid-fetch path of //! `lib/microwaveprop/weather/hrrr_client.ex`: //! * URL derivation (`hrrr.tHHz.wrfsfcfFH.grib2` / `wrfprsfFH.grib2`) //! * `.idx` parse + per-`(var, level)` byte-range lookup //! * range merge + parallel `Range:` downloads //! * in-process 1h TTL `.idx` memoization (idx files are immutable per run) //! * retry on HTTP 429/5xx with jittered exponential backoff use std::collections::HashMap; use std::fs; use std::io::Write; use std::path::{Path, PathBuf}; use std::sync::Mutex; use std::time::{Duration, Instant}; use chrono::{DateTime, NaiveDate, Timelike, Utc}; use reqwest::header::{HeaderMap, HeaderValue, RANGE}; use tokio::task::JoinSet; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Product { Surface, Pressure, } #[derive(Debug, Clone)] pub struct IdxEntry { pub msg: u32, pub offset: u64, pub var: String, pub level: String, } #[derive(Debug, Clone, Copy)] pub struct ByteRange { pub start: u64, pub end: u64, } #[derive(Debug, thiserror::Error)] pub enum FetchError { #[error("http: {0}")] Http(#[from] reqwest::Error), #[error("idx HTTP {status}")] IdxStatus { status: u16 }, #[error("grib HTTP {status}")] GribStatus { status: u16 }, #[error("idx text was empty")] IdxEmpty, } pub const HRRR_BASE_DEFAULT: &str = "https://noaa-hrrr-bdp-pds.s3.amazonaws.com"; pub const IDX_TTL: Duration = Duration::from_secs(3_600); pub const MAX_PARALLEL_RANGES: usize = 8; pub const REQUEST_TIMEOUT: Duration = Duration::from_secs(120); pub const GRID_PRESSURE_LEVELS: &[u16] = &[ 1000, 975, 950, 925, 900, 875, 850, 825, 800, 775, 750, 725, 700, ]; /// Pre-formatted lookup keys for per-cell pressure-level access. /// /// The hot decode → score loop runs over 92k cells × 13 levels and /// would otherwise `format!("TMP:{p} mb")` on every cell, every level, /// every variable — ~3.6M String allocations per chain step. This /// table is built once (at first access), stored as owned Strings, /// and borrowed as `&str` at every call site. pub struct LevelKeys { pub pres_mb: f64, pub tmp: String, pub dpt: String, pub hgt: String, } pub fn grid_level_keys() -> &'static [LevelKeys] { static KEYS: std::sync::OnceLock> = std::sync::OnceLock::new(); KEYS.get_or_init(|| { GRID_PRESSURE_LEVELS .iter() .map(|&p| LevelKeys { pres_mb: p as f64, tmp: format!("TMP:{p} mb"), dpt: format!("DPT:{p} mb"), hgt: format!("HGT:{p} mb"), }) .collect() }) } pub const SURFACE_MESSAGES: &[(&str, &str)] = &[ ("TMP", "2 m above ground"), ("DPT", "2 m above ground"), ("PRES", "surface"), ("HPBL", "surface"), ("PWAT", "entire atmosphere (considered as a single layer)"), ("UGRD", "10 m above ground"), ("VGRD", "10 m above ground"), ("TCDC", "entire atmosphere"), ("APCP", "surface"), ]; pub fn hrrr_url( base: &str, date: NaiveDate, hour: u8, product: Product, forecast_hour: u8, ) -> String { let date_str = date.format("%Y%m%d").to_string(); let hh = format!("{hour:02}"); let fh = format!("{forecast_hour:02}"); let file = match product { Product::Surface => format!("hrrr.t{hh}z.wrfsfcf{fh}.grib2"), Product::Pressure => format!("hrrr.t{hh}z.wrfprsf{fh}.grib2"), }; let trimmed = base.trim_end_matches('/'); format!("{trimmed}/hrrr.{date_str}/conus/{file}") } pub fn hrrr_native_url(base: &str, date: NaiveDate, hour: u8, forecast_hour: u8) -> String { let date_str = date.format("%Y%m%d").to_string(); let hh = format!("{hour:02}"); let fh = format!("{forecast_hour:02}"); let trimmed = base.trim_end_matches('/'); format!("{trimmed}/hrrr.{date_str}/conus/hrrr.t{hh}z.wrfnatf{fh}.grib2") } pub fn parse_idx(text: &str) -> Vec { text.split('\n') .filter(|l| !l.is_empty()) .filter_map(parse_idx_line) .collect() } fn parse_idx_line(line: &str) -> Option { let mut parts = line.splitn(8, ':'); let msg: u32 = parts.next()?.parse().ok()?; let offset: u64 = parts.next()?.parse().ok()?; parts.next()?; // date let var = parts.next()?.to_string(); let level = parts.next()?.to_string(); Some(IdxEntry { msg, offset, var, level, }) } pub fn byte_ranges_for_messages(idx: &[IdxEntry], wanted: &[(String, String)]) -> Vec { let mut out = Vec::new(); for (var, level) in wanted { for (i, entry) in idx.iter().enumerate() { if entry.var == *var && entry.level == *level { let end = idx .get(i + 1) .map(|nxt| nxt.offset.saturating_sub(1)) .unwrap_or(entry.offset + 10_000_000); out.push(ByteRange { start: entry.offset, end, }); } } } out } pub fn merge_ranges(mut ranges: Vec) -> Vec { if ranges.is_empty() { return ranges; } ranges.sort_by_key(|r| r.start); let mut out: Vec = Vec::with_capacity(ranges.len()); for r in ranges { match out.last_mut() { Some(prev) if r.start <= prev.end.saturating_add(1) => { prev.end = prev.end.max(r.end); } _ => out.push(r), } } out } /// Truncates a DateTime down to the nearest published HRRR cycle hour. /// /// Always rounds DOWN — rounding up past `:30` used to point at the next /// cycle, which produced 404/403 errors when that cycle had not yet been /// published by NOAA. Callers that need a later valid_time should reach /// it via `forecast_hour` instead. pub fn nearest_hrrr_hour(dt: DateTime) -> DateTime { let total_seconds = dt.minute() as i64 * 60 + dt.second() as i64; dt - chrono::Duration::seconds(total_seconds) } pub fn pressure_messages_grid() -> Vec<(String, String)> { let mut out = Vec::with_capacity(GRID_PRESSURE_LEVELS.len() * 3); for &level in GRID_PRESSURE_LEVELS { for var in ["TMP", "DPT", "HGT"] { out.push((var.to_string(), format!("{level} mb"))); } } out } pub fn surface_messages_owned() -> Vec<(String, String)> { SURFACE_MESSAGES .iter() .map(|(v, l)| ((*v).to_string(), (*l).to_string())) .collect() } // ── Idx cache ──────────────────────────────────────────────────────── // Disk TTL is longer than the in-memory TTL because idx files are // immutable per HRRR run — cross-replica sharing on NFS benefits from // keeping hits alive across pod restarts and the next hourly cycle. pub const IDX_DISK_TTL: Duration = Duration::from_secs(24 * 3_600); #[derive(Debug, Clone)] struct IdxCacheEntry { body: String, inserted: Instant, } #[derive(Debug)] pub struct DiskIdxCache { dir: PathBuf, ttl: Duration, } impl DiskIdxCache { pub fn new(dir: impl AsRef, ttl: Duration) -> Self { let dir = dir.as_ref().to_path_buf(); let _ = fs::create_dir_all(&dir); Self { dir, ttl } } fn path_for(&self, url: &str) -> PathBuf { // Sanitize URL into a safe filename. Collisions would require // two distinct URLs to differ only in unsafe characters mapped // to '_' — not possible for HRRR S3 URLs. let mut sanitized = String::with_capacity(url.len()); for c in url.chars() { if c.is_ascii_alphanumeric() || c == '.' || c == '-' || c == '_' { sanitized.push(c); } else { sanitized.push('_'); } } self.dir.join(sanitized) } pub fn get(&self, url: &str) -> Option { let path = self.path_for(url); let meta = fs::metadata(&path).ok()?; let mtime = meta.modified().ok()?; if mtime.elapsed().ok()? > self.ttl { return None; } fs::read_to_string(&path).ok() } pub fn put(&self, url: &str, body: &str) { // Atomic write: tmp file + rename so concurrent readers across // replicas never observe a partial body. let path = self.path_for(url); let tmp = path.with_extension("tmp"); if let Ok(mut f) = fs::File::create(&tmp) { if f.write_all(body.as_bytes()).is_ok() { drop(f); let _ = fs::rename(&tmp, &path); } } } } #[derive(Debug, Default)] pub struct IdxCache { inner: Mutex>, disk: Option, } impl IdxCache { pub fn new() -> Self { Self::default() } pub fn with_disk(dir: impl AsRef) -> Self { Self { inner: Mutex::new(HashMap::new()), disk: Some(DiskIdxCache::new(dir, IDX_DISK_TTL)), } } pub fn get(&self, url: &str) -> Option { if let Ok(mut map) = self.inner.lock() { if let Some(entry) = map.get(url) { if entry.inserted.elapsed() < IDX_TTL { return Some(entry.body.clone()); } // Expired entry — evict so unique idx URLs don't // accumulate across hourly HRRR cycles. map.remove(url); } } if let Some(disk) = &self.disk { if let Some(body) = disk.get(url) { if let Ok(mut map) = self.inner.lock() { map.insert( url.to_string(), IdxCacheEntry { body: body.clone(), inserted: Instant::now(), }, ); } return Some(body); } } None } pub fn put(&self, url: String, body: String) { if let Some(disk) = &self.disk { disk.put(&url, &body); } if let Ok(mut map) = self.inner.lock() { map.insert( url, IdxCacheEntry { body, inserted: Instant::now(), }, ); } } #[cfg(test)] fn len(&self) -> usize { self.inner.lock().map(|m| m.len()).unwrap_or(0) } #[cfg(test)] fn put_with_age(&self, url: String, body: String, age: Duration) { if let Ok(mut map) = self.inner.lock() { map.insert( url, IdxCacheEntry { body, inserted: Instant::now() - age, }, ); } } } // ── HTTP client ────────────────────────────────────────────────────── pub struct HrrrClient { http: reqwest::Client, base: String, // When the primary base fails every retry (e.g. an on-cluster // caching proxy is unreachable), fetches retry once against this // base. Production sets it to the NOAA S3 bucket so a proxy outage // degrades to direct S3 reads instead of a hung pipeline. fallback_base: Option, idx_cache: IdxCache, } impl HrrrClient { pub fn new(base: impl Into) -> Result { Self::build(base.into(), None) } pub fn new_with_fallback( base: impl Into, fallback: impl Into, ) -> Result { Self::build(base.into(), Some(fallback.into())) } fn build(base: String, fallback: Option) -> Result { let http = reqwest::Client::builder() .timeout(REQUEST_TIMEOUT) .pool_max_idle_per_host(MAX_PARALLEL_RANGES) .build()?; let idx_cache = match std::env::var("HRRR_IDX_CACHE_DIR") { Ok(dir) if !dir.is_empty() => IdxCache::with_disk(dir), _ => IdxCache::new(), }; let fallback = fallback.filter(|f| f != &base); Ok(Self { http, base, fallback_base: fallback, idx_cache, }) } pub fn default_base() -> Result { Self::new(HRRR_BASE_DEFAULT) } pub fn base(&self) -> &str { &self.base } pub fn fallback_base(&self) -> Option<&str> { self.fallback_base.as_deref() } pub async fn fetch_idx(&self, url: &str) -> Result { if let Some(body) = self.idx_cache.get(url) { return Ok(body); } let body = retry_get_text(&self.http, url).await?; if body.is_empty() { return Err(FetchError::IdxEmpty); } self.idx_cache.put(url.to_string(), body.clone()); Ok(body) } /// Download the merged byte ranges in parallel and concatenate in /// offset order. Mirrors `download_grib_ranges_remote/2` in Elixir. pub async fn download_grib_ranges( &self, url: &str, ranges: Vec, ) -> Result, FetchError> { if ranges.is_empty() { return Ok(Vec::new()); } let merged = merge_ranges(ranges); let mut futs: JoinSet), FetchError>> = JoinSet::new(); let url = url.to_string(); for r in merged.iter().copied() { let client = self.http.clone(); let url = url.clone(); futs.spawn(async move { let bytes = retry_get_range(&client, &url, r).await?; Ok((r.start, bytes)) }); } // Cap concurrent in-flight requests at MAX_PARALLEL_RANGES. let mut results: Vec<(u64, Vec)> = Vec::with_capacity(merged.len()); while let Some(item) = futs.join_next().await { results.push(item.expect("range task")?); } results.sort_by_key(|(off, _)| *off); let total: usize = results.iter().map(|(_, b)| b.len()).sum(); let mut out = Vec::with_capacity(total); for (_, b) in results { out.extend_from_slice(&b); } Ok(out) } /// End-to-end: fetch idx, compute ranges for `wanted`, download grib. /// Tries the primary base first, then the fallback base (if configured) /// when the primary exhausts its retries. pub async fn fetch_product_blob( &self, date: NaiveDate, hour: u8, product: Product, forecast_hour: u8, wanted: &[(String, String)], ) -> Result, FetchError> { let primary = hrrr_url(&self.base, date, hour, product, forecast_hour); let fallback = self .fallback_base .as_deref() .map(|fb| hrrr_url(fb, date, hour, product, forecast_hour)); self.fetch_blob_with_fallback(&primary, fallback.as_deref(), wanted) .await } /// End-to-end native-level fetch (hrrr.tHHz.wrfnatfFH.grib2) with the /// same primary-then-fallback logic as `fetch_product_blob`. Callers /// (native_duct) supply the duct-variable `wanted` list. pub async fn fetch_native_blob( &self, date: NaiveDate, hour: u8, forecast_hour: u8, wanted: &[(String, String)], ) -> Result, FetchError> { let primary = hrrr_native_url(&self.base, date, hour, forecast_hour); let fallback = self .fallback_base .as_deref() .map(|fb| hrrr_native_url(fb, date, hour, forecast_hour)); self.fetch_blob_with_fallback(&primary, fallback.as_deref(), wanted) .await } async fn fetch_blob_with_fallback( &self, primary_url: &str, fallback_url: Option<&str>, wanted: &[(String, String)], ) -> Result, FetchError> { match self.fetch_blob_at(primary_url, wanted).await { Ok(blob) => Ok(blob), Err(primary_err) => match fallback_url { Some(fb) => { tracing::warn!( error = %primary_err, primary = primary_url, fallback = fb, "hrrr primary failed; retrying against fallback" ); self.fetch_blob_at(fb, wanted).await } None => Err(primary_err), }, } } async fn fetch_blob_at( &self, url: &str, wanted: &[(String, String)], ) -> Result, FetchError> { let idx_url = format!("{url}.idx"); let idx_body = self.fetch_idx(&idx_url).await?; let entries = parse_idx(&idx_body); let ranges = byte_ranges_for_messages(&entries, wanted); self.download_grib_ranges(url, ranges).await } } async fn retry_get_text(client: &reqwest::Client, url: &str) -> Result { for attempt in 0..5 { match client.get(url).send().await { Ok(resp) => { let status = resp.status().as_u16(); if status == 200 { return Ok(resp.text().await?); } if !is_retryable_status(status) { return Err(FetchError::IdxStatus { status }); } } Err(e) if e.is_connect() || e.is_timeout() => { tracing::warn!(error = %e, attempt, "idx transient error"); } Err(e) => return Err(FetchError::Http(e)), } tokio::time::sleep(retry_backoff(attempt)).await; } Err(FetchError::IdxStatus { status: 0 }) } async fn retry_get_range( client: &reqwest::Client, url: &str, range: ByteRange, ) -> Result, FetchError> { let value = HeaderValue::from_str(&format!("bytes={}-{}", range.start, range.end)) .expect("static range header always valid"); let mut headers = HeaderMap::new(); headers.insert(RANGE, value); for attempt in 0..5 { match client.get(url).headers(headers.clone()).send().await { Ok(resp) => { let status = resp.status().as_u16(); if status == 206 || status == 200 { return Ok(resp.bytes().await?.to_vec()); } if !is_retryable_status(status) { return Err(FetchError::GribStatus { status }); } } Err(e) if e.is_connect() || e.is_timeout() => { tracing::warn!(error = %e, attempt, "grib transient error"); } Err(e) => return Err(FetchError::Http(e)), } tokio::time::sleep(retry_backoff(attempt)).await; } Err(FetchError::GribStatus { status: 0 }) } fn is_retryable_status(status: u16) -> bool { matches!(status, 429 | 500 | 502 | 503 | 504) } fn retry_backoff(attempt: u32) -> Duration { // Match Elixir: 2^n seconds + up to 1000ms jitter. let base = 2u64.pow(attempt).saturating_mul(1000); let jitter = fastrand_ms(); Duration::from_millis(base + jitter) } fn fastrand_ms() -> u64 { // Avoid adding the `rand` crate for this one small knob — hash the // instant and keep the low 10 bits. let nanos = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .map(|d| d.as_nanos() as u64) .unwrap_or(0); nanos & 0x3ff // 0..1023 ms } #[cfg(test)] mod tests { use super::*; use chrono::TimeZone; #[test] fn url_matches_elixir() { let date = NaiveDate::from_ymd_opt(2026, 4, 19).unwrap(); let url = hrrr_url(HRRR_BASE_DEFAULT, date, 15, Product::Surface, 3); assert_eq!( url, "https://noaa-hrrr-bdp-pds.s3.amazonaws.com/hrrr.20260419/conus/hrrr.t15z.wrfsfcf03.grib2" ); } #[test] fn hrrr_url_trims_trailing_slash() { let d = NaiveDate::from_ymd_opt(2026, 4, 19).unwrap(); assert_eq!( hrrr_url("https://x.y/", d, 14, Product::Surface, 0), "https://x.y/hrrr.20260419/conus/hrrr.t14z.wrfsfcf00.grib2" ); } #[test] fn hrrr_native_url_format() { let d = NaiveDate::from_ymd_opt(2026, 4, 19).unwrap(); assert_eq!( hrrr_native_url("https://x.y", d, 14, 0), "https://x.y/hrrr.20260419/conus/hrrr.t14z.wrfnatf00.grib2" ); assert_eq!( hrrr_native_url("https://x.y/", d, 4, 7), "https://x.y/hrrr.20260419/conus/hrrr.t04z.wrfnatf07.grib2" ); } #[test] fn new_stores_no_fallback() { let c = HrrrClient::new("https://primary").unwrap(); assert_eq!(c.base(), "https://primary"); assert!(c.fallback_base().is_none()); } #[test] fn new_with_fallback_stores_fallback() { let c = HrrrClient::new_with_fallback("http://primary", "https://fallback").unwrap(); assert_eq!(c.base(), "http://primary"); assert_eq!(c.fallback_base(), Some("https://fallback")); } #[test] fn fallback_equal_to_primary_is_dropped() { // Pointing fallback at the same base adds no resilience and // would double the wait on every failure. Treat it as "no // fallback" so retries stop after the primary exhausts. let c = HrrrClient::new_with_fallback("https://same", "https://same").unwrap(); assert!(c.fallback_base().is_none()); } #[test] fn idx_parsing_drops_bad_lines() { let idx = "1:0:d=2024:TMP:2 m above ground:anl:\nbad\n2:12345:d=2024:DPT:2 m above ground:anl:"; let parsed = parse_idx(idx); assert_eq!(parsed.len(), 2); assert_eq!(parsed[0].var, "TMP"); assert_eq!(parsed[1].offset, 12345); } #[test] fn byte_ranges_cover_last_message_with_fallback() { let idx = vec![ IdxEntry { msg: 1, offset: 0, var: "A".into(), level: "x".into(), }, IdxEntry { msg: 2, offset: 100, var: "B".into(), level: "y".into(), }, ]; let wanted = vec![("A".into(), "x".into()), ("B".into(), "y".into())]; let r = byte_ranges_for_messages(&idx, &wanted); assert_eq!(r.len(), 2); assert_eq!(r[0].start, 0); assert_eq!(r[0].end, 99); assert_eq!(r[1].start, 100); assert_eq!(r[1].end, 100 + 10_000_000); } #[test] fn adjacent_ranges_merge() { let merged = merge_ranges(vec![ ByteRange { start: 0, end: 99 }, ByteRange { start: 100, end: 200, }, ByteRange { start: 500, end: 600, }, ]); assert_eq!(merged.len(), 2); assert_eq!(merged[0].start, 0); assert_eq!(merged[0].end, 200); } #[test] fn nearest_hour_always_rounds_down() { // Round-down keeps us from probing a cycle that NOAA has not // published yet. Callers reach later valid_times via forecast_hour. let early = Utc.with_ymd_and_hms(2026, 4, 19, 15, 29, 0).unwrap(); assert_eq!(nearest_hrrr_hour(early).hour(), 15); let mid = Utc.with_ymd_and_hms(2026, 4, 19, 15, 30, 0).unwrap(); assert_eq!(nearest_hrrr_hour(mid).hour(), 15); let late = Utc.with_ymd_and_hms(2026, 4, 19, 15, 59, 0).unwrap(); assert_eq!(nearest_hrrr_hour(late).hour(), 15); } #[test] fn pressure_messages_sized_correctly() { let msgs = pressure_messages_grid(); assert_eq!(msgs.len(), GRID_PRESSURE_LEVELS.len() * 3); } #[test] fn idx_cache_returns_within_ttl() { let cache = IdxCache::new(); cache.put("u".into(), "body".into()); assert_eq!(cache.get("u").as_deref(), Some("body")); assert!(cache.get("missing").is_none()); } #[test] fn idx_cache_removes_expired_entries_on_get() { // Without eviction the HashMap grows forever as unique idx URLs // accumulate each hour. Calling get() on an expired entry must // remove it so the in-memory map bounds itself to the hot set. let cache = IdxCache::new(); cache.put_with_age( "stale".into(), "old".into(), IDX_TTL + Duration::from_secs(1), ); cache.put("fresh".into(), "new".into()); assert_eq!(cache.len(), 2); assert!(cache.get("stale").is_none()); assert_eq!(cache.len(), 1, "expired entry should have been evicted"); assert_eq!(cache.get("fresh").as_deref(), Some("new")); } #[test] fn disk_idx_cache_round_trips() { let tmp = tempfile::TempDir::new().unwrap(); let cache = DiskIdxCache::new(tmp.path(), Duration::from_secs(3_600)); cache.put("https://example.com/a.grib2.idx", "body1"); assert_eq!( cache.get("https://example.com/a.grib2.idx").as_deref(), Some("body1") ); assert!(cache.get("https://example.com/missing.idx").is_none()); } #[test] fn disk_idx_cache_respects_ttl() { let tmp = tempfile::TempDir::new().unwrap(); let cache = DiskIdxCache::new(tmp.path(), Duration::from_millis(10)); cache.put("u", "body"); std::thread::sleep(Duration::from_millis(30)); assert!(cache.get("u").is_none()); } #[test] fn idx_cache_disk_fallback_rewarms_memory() { let tmp = tempfile::TempDir::new().unwrap(); // Seed disk via one cache instance... let c1 = IdxCache::with_disk(tmp.path()); c1.put("u".into(), "body".into()); // A fresh cache simulates a new pod — in-memory map is empty // but the disk entry must still resolve. let c2 = IdxCache::with_disk(tmp.path()); assert_eq!(c2.get("u").as_deref(), Some("body")); } }