Graham McIntire
e9a38623d8
Batched from a system-wide review pass.
**Rate-limit + transient-error log hygiene**
- WeatherFetchWorker: classify IEM HTTP 429 as {:snooze, 300} instead
of {:error}. Stops Oban.PerformError stack-trace spam on every
routine rate-limit during backfill, keeps the retry semantics.
- IonosphereFetchWorker: compress GIRO TLS :unknown_ca error from a
~400-char inspect blob to 'TLS unknown_ca (CA bundle missing)'.
- FreshnessMonitor: log only when an enqueue actually lands (the
Oban unique conflict case was silently dropping jobs but still
producing 'enqueuing grid worker' info lines every 5 minutes).
**Perf**
- Rust grid_level_keys(): pre-format 'TMP:{p} mb' / DPT / HGT keys
once via OnceLock instead of format!()'ing per-cell. Removes ~3.6M
String allocations per f01..f18 chain step across pipeline.rs,
hrrr_points.rs. Same for the wgrib2 :(TMP|DPT|HGT): pattern in
hrrr_points.process_batch (sfc_pattern / prs_pattern OnceLock).
- MapLive preload_forecast: Task.async_stream with max_concurrency=4
replaces the 18-wide Enum.map serial walk. Forecast cache warms
~4× faster after a band change, with ordering preserved.
- grid_tasks: partial composite index on (run_time DESC,
forecast_hour ASC) WHERE status='queued' AND kind='forecast',
matching the Rust claim query's ORDER BY. Drops the old
status-only partial that forced a sort per claim.
**Correctness**
- ContactImportWorker: add Oban unique:[keys: [:import_run_id,
:offset]]. Was missing on a worker whose perform() does a
non-idempotent atomic counter increment — a retry would double-
count imported rows.
- CommonVolumeRadarWorker: x_min..x_max default step is -1 when
x_min > x_max, triggering a runtime deprecation warning. Force
Range.new(_, _, 1) explicitly.
**Cleanup**
- Drop unused tmp_dir parameter from hrrr_point_worker + its
process_batch signature.
758 lines
24 KiB
Rust
758 lines
24 KiB
Rust
//! 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 futures::stream::{FuturesUnordered, StreamExt};
|
||
use reqwest::header::{HeaderMap, HeaderValue, RANGE};
|
||
|
||
#[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<Vec<LevelKeys>> = 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<IdxEntry> {
|
||
text.split('\n')
|
||
.filter(|l| !l.is_empty())
|
||
.filter_map(parse_idx_line)
|
||
.collect()
|
||
}
|
||
|
||
fn parse_idx_line(line: &str) -> Option<IdxEntry> {
|
||
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<ByteRange> {
|
||
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<ByteRange>) -> Vec<ByteRange> {
|
||
if ranges.is_empty() {
|
||
return ranges;
|
||
}
|
||
ranges.sort_by_key(|r| r.start);
|
||
let mut out: Vec<ByteRange> = 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
|
||
}
|
||
|
||
pub fn nearest_hrrr_hour(dt: DateTime<Utc>) -> DateTime<Utc> {
|
||
let total_seconds = dt.minute() as i64 * 60 + dt.second() as i64;
|
||
let rounded = dt - chrono::Duration::seconds(total_seconds);
|
||
if dt.minute() >= 30 {
|
||
rounded + chrono::Duration::hours(1)
|
||
} else {
|
||
rounded
|
||
}
|
||
}
|
||
|
||
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<Path>, 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<String> {
|
||
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<HashMap<String, IdxCacheEntry>>,
|
||
disk: Option<DiskIdxCache>,
|
||
}
|
||
|
||
impl IdxCache {
|
||
pub fn new() -> Self {
|
||
Self::default()
|
||
}
|
||
|
||
pub fn with_disk(dir: impl AsRef<Path>) -> Self {
|
||
Self {
|
||
inner: Mutex::new(HashMap::new()),
|
||
disk: Some(DiskIdxCache::new(dir, IDX_DISK_TTL)),
|
||
}
|
||
}
|
||
|
||
pub fn get(&self, url: &str) -> Option<String> {
|
||
if let Ok(map) = self.inner.lock() {
|
||
if let Some(entry) = map.get(url) {
|
||
if entry.inserted.elapsed() < IDX_TTL {
|
||
return Some(entry.body.clone());
|
||
}
|
||
}
|
||
}
|
||
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(),
|
||
},
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
// ── 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<String>,
|
||
idx_cache: IdxCache,
|
||
}
|
||
|
||
impl HrrrClient {
|
||
pub fn new(base: impl Into<String>) -> Result<Self, FetchError> {
|
||
Self::build(base.into(), None)
|
||
}
|
||
|
||
pub fn new_with_fallback(
|
||
base: impl Into<String>,
|
||
fallback: impl Into<String>,
|
||
) -> Result<Self, FetchError> {
|
||
Self::build(base.into(), Some(fallback.into()))
|
||
}
|
||
|
||
fn build(base: String, fallback: Option<String>) -> Result<Self, FetchError> {
|
||
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, FetchError> {
|
||
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<String, FetchError> {
|
||
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<ByteRange>,
|
||
) -> Result<Vec<u8>, FetchError> {
|
||
if ranges.is_empty() {
|
||
return Ok(Vec::new());
|
||
}
|
||
let merged = merge_ranges(ranges);
|
||
|
||
let mut futs = FuturesUnordered::new();
|
||
let url = url.to_string();
|
||
for r in merged.iter().copied() {
|
||
let client = self.http.clone();
|
||
let url = url.clone();
|
||
futs.push(async move {
|
||
let bytes = retry_get_range(&client, &url, r).await?;
|
||
Ok::<(u64, Vec<u8>), FetchError>((r.start, bytes))
|
||
});
|
||
}
|
||
|
||
// Cap concurrent in-flight requests at MAX_PARALLEL_RANGES.
|
||
let mut results: Vec<(u64, Vec<u8>)> = Vec::with_capacity(merged.len());
|
||
while let Some(item) = futs.next().await {
|
||
results.push(item?);
|
||
}
|
||
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<Vec<u8>, 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<Vec<u8>, 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<Vec<u8>, 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<Vec<u8>, 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<String, FetchError> {
|
||
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<Vec<u8>, 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:\n<html>bad</html>\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_rounds_properly() {
|
||
let base = Utc.with_ymd_and_hms(2026, 4, 19, 15, 29, 0).unwrap();
|
||
assert_eq!(nearest_hrrr_hour(base).hour(), 15);
|
||
let up = Utc.with_ymd_and_hms(2026, 4, 19, 15, 30, 0).unwrap();
|
||
assert_eq!(nearest_hrrr_hour(up).hour(), 16);
|
||
}
|
||
|
||
#[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 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"));
|
||
}
|
||
}
|