ocd/src/fetch.rs

185 lines
6.0 KiB
Rust
Raw Normal View History

//! The paced workspace walk.
//!
//! Targets are flattened into one ordered list of (account, workspace). The
//! walk is sequential with 500 ms pacing. Per workspace two requests run at
//! once: Zen for balance and Go for limits. Fresh cache skips the network.
use std::time::Duration;
use chrono::Utc;
use crate::cache::{self, CachedRecord};
use crate::client::{self, HttpClient};
use crate::model::{AccountReport, Balance, Limits, ReportStatus};
use crate::scrape;
/// One fetch target: an account and one workspace of it, with the cookie.
#[derive(Clone)]
pub struct Target {
pub account: String,
pub workspace: String,
pub cookie: String,
}
/// Resolved pipeline options from the CLI.
#[derive(Clone)]
pub struct FetchOptions {
pub ttl_secs: u32,
pub timeout_secs: u64,
pub retries: u8,
pub force_refresh: bool,
pub no_cache: bool,
pub workspace_pace_ms: u64,
pub account_filter: Vec<String>,
pub workspace_filter: Vec<String>,
pub user_agent: String,
}
/// Build the target list from the config, applying CLI filters.
pub fn targets(cfg: &crate::config::Config, opts: &FetchOptions) -> Vec<Target> {
let mut out = Vec::new();
for acct in &cfg.account {
if !opts.account_filter.is_empty() && !opts.account_filter.contains(&acct.name) {
continue;
}
for ws in acct.workspace_ids() {
if !opts.workspace_filter.is_empty() && !opts.workspace_filter.contains(&ws) {
continue;
}
out.push(Target {
account: acct.name.clone(),
workspace: ws,
cookie: acct.cookie.clone(),
});
}
}
out
}
/// Walk all targets. Returns one report per target, in config order.
pub async fn run(cfg: &crate::config::Config, opts: &FetchOptions) -> anyhow::Result<Vec<AccountReport>> {
let http = HttpClient::new(opts.user_agent.clone(), opts.timeout_secs, opts.retries)?;
let targets = targets(cfg, opts);
let mut reports = Vec::with_capacity(targets.len());
for (i, tgt) in targets.iter().enumerate() {
let report = fetch_target(tgt, opts, &http).await;
reports.push(report);
// Pace before the next workspace. Do not pace after the last.
if i + 1 < targets.len() && opts.workspace_pace_ms > 0 {
tokio::time::sleep(Duration::from_millis(opts.workspace_pace_ms)).await;
}
}
Ok(reports)
}
/// Fetch one target. Cache short-circuits the network.
async fn fetch_target(tgt: &Target, opts: &FetchOptions, http: &HttpClient) -> AccountReport {
// Cache hit path.
if !opts.force_refresh && !opts.no_cache {
if let Some(rec) = cache::read(&tgt.account, &tgt.workspace, opts.ttl_secs) {
let mut r = rec.report;
r.status = ReportStatus::Cached;
return r;
}
}
// Network path. Fire Zen and Go at once.
let zen_url = client::zen_url(&tgt.workspace);
let go_url = client::go_url(&tgt.workspace);
let (zen_res, go_res) = tokio::join!(
http.get_text(&zen_url, &tgt.cookie),
http.get_text(&go_url, &tgt.cookie),
);
let fetched_at = Some(Utc::now().to_rfc3339());
let balance = match &zen_res {
Ok(html) => match scrape::parse_zen_balance(html) {
Ok(b) => Some(b),
Err(e) => {
tracing::warn!(account = %tgt.account, workspace = %tgt.workspace, error = %e, "zen parse failed");
None
}
},
Err(e) => {
tracing::warn!(account = %tgt.account, workspace = %tgt.workspace, error = %e, "zen fetch failed");
None
}
};
let limits = match &go_res {
Ok(html) => match scrape::parse_go_limits(html) {
Ok(l) => Some(l),
Err(e) => {
tracing::warn!(account = %tgt.account, workspace = %tgt.workspace, error = %e, "go parse failed");
None
}
},
Err(e) => {
tracing::warn!(account = %tgt.account, workspace = %tgt.workspace, error = %e, "go fetch failed");
None
}
};
// Decide the status. RateLimited and Timeout take priority over generic
// Error so callers can react.
let status = report_status(&zen_res, &go_res, limits.as_ref(), balance.as_ref());
let report = AccountReport {
account: tgt.account.clone(),
workspace: tgt.workspace.clone(),
fetched_at,
status: status.clone(),
limits: limits.clone(),
balance: balance.clone(),
};
// Write cache only on a clean Ok, so errors never poison the cache.
if status == ReportStatus::Ok {
let record = CachedRecord {
report: report.clone(),
go_etag: None,
go_last_modified: None,
zen_etag: None,
zen_last_modified: None,
};
if let Err(e) = cache::write(&tgt.account, &tgt.workspace, &record) {
tracing::warn!(error = %e, "cache write failed");
}
}
report
}
/// Fold the two fetch results into one status.
fn report_status(
zen: &Result<String, crate::error::FetchError>,
go: &Result<String, crate::error::FetchError>,
limits: Option<&Limits>,
balance: Option<&Balance>,
) -> ReportStatus {
use crate::error::FetchError;
for r in [zen, go] {
match r {
Err(FetchError::RateLimited) => return ReportStatus::RateLimited,
Err(FetchError::Timeout) => return ReportStatus::Timeout,
Err(FetchError::Status { status, .. }) if *status == 429 => {
return ReportStatus::RateLimited;
}
_ => {}
}
}
for r in [zen, go] {
if let Err(e) = r {
if limits.is_none() && balance.is_none() {
return ReportStatus::Error(e.to_string());
}
return ReportStatus::Error(e.to_string());
}
}
if limits.is_some() || balance.is_some() {
ReportStatus::Ok
} else {
ReportStatus::Error("no limits or balance parsed".to_string())
}
}