use std::io::Read;
use std::num::NonZeroU32;
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use thiserror::Error;
const MAX_BACKOFF_MS: u64 = 30_000;
#[derive(Debug, Error)]
pub enum HttpError {
#[error("http {code} for {url}")]
Status { code: u16, url: String },
#[error("transport error: {0}")]
Transport(String),
#[error("io error: {0}")]
Io(#[from] std::io::Error),
}
impl HttpError {
fn is_transient(&self) -> bool {
match self {
HttpError::Status { code, .. } => *code == 429 || (500..=599).contains(code),
HttpError::Transport(_) => true,
HttpError::Io(_) => false,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FetchMode {
Always,
OnlyIfMissing,
}
#[derive(Debug, Clone)]
pub struct DatasetClientConfig {
pub user_agent: String,
pub connect_timeout: Duration,
pub overall_timeout: Option<Duration>,
pub rate_per_sec: Option<NonZeroU32>,
pub retry_count: usize,
pub base_backoff_ms: u64,
}
#[derive(Clone)]
pub struct DatasetClient {
agent: ureq::Agent,
last_request: Arc<Mutex<Option<Instant>>>,
min_gap: Option<Duration>,
retry_count: usize,
base_backoff_ms: u64,
}
impl DatasetClient {
pub fn new(config: DatasetClientConfig) -> Self {
let mut builder = ureq::AgentBuilder::new()
.user_agent(&config.user_agent)
.timeout_connect(config.connect_timeout);
if let Some(timeout) = config.overall_timeout {
builder = builder.timeout(timeout);
}
let min_gap = config
.rate_per_sec
.map(|r| Duration::from_secs_f64(1.0 / r.get() as f64));
DatasetClient {
agent: builder.build(),
last_request: Arc::new(Mutex::new(None)),
min_gap,
retry_count: config.retry_count,
base_backoff_ms: config.base_backoff_ms,
}
}
pub(crate) fn agent(&self) -> &ureq::Agent {
&self.agent
}
fn rate_gate(&self) {
let Some(gap) = self.min_gap else {
return;
};
let mut last = self.last_request.lock().unwrap();
if let Some(prev) = *last {
let elapsed = prev.elapsed();
if elapsed < gap {
std::thread::sleep(gap - elapsed);
}
}
*last = Some(Instant::now());
}
fn request_once(&self, url: &str) -> Result<Vec<u8>, HttpError> {
match self.agent.get(url).call() {
Ok(resp) => {
let mut buf = Vec::new();
resp.into_reader()
.read_to_end(&mut buf)
.map_err(|e| HttpError::Transport(e.to_string()))?;
Ok(buf)
}
Err(ureq::Error::Status(code, _resp)) => Err(HttpError::Status {
code,
url: url.to_string(),
}),
Err(ureq::Error::Transport(t)) => Err(HttpError::Transport(t.to_string())),
}
}
pub fn fetch_bytes(&self, url: &str) -> Result<Vec<u8>, HttpError> {
let mut delay_ms = self.base_backoff_ms;
for attempt in 0..=self.retry_count {
self.rate_gate();
match self.request_once(url) {
Ok(bytes) => return Ok(bytes),
Err(e) => {
if !e.is_transient() || attempt == self.retry_count {
return Err(e);
}
eprintln!(
"kglite-http: {e} (attempt {}/{}); backing off {delay_ms}ms",
attempt + 1,
self.retry_count + 1,
);
std::thread::sleep(Duration::from_millis(delay_ms));
delay_ms = (delay_ms * 2).min(MAX_BACKOFF_MS);
}
}
}
unreachable!("fetch_bytes loop returns or errors before completing")
}
pub fn fetch_json(&self, url: &str) -> Result<serde_json::Value, HttpError> {
let bytes = self.fetch_bytes(url)?;
serde_json::from_slice(&bytes)
.map_err(|e| HttpError::Transport(format!("json decode from {url}: {e}")))
}
pub fn fetch_to_file(
&self,
url: &str,
path: &Path,
mode: FetchMode,
) -> Result<bool, HttpError> {
if mode == FetchMode::OnlyIfMissing && path.is_file() {
return Ok(false);
}
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let bytes = self.fetch_bytes(url)?;
let tmp = path.with_extension("tmp");
std::fs::write(&tmp, &bytes)?;
std::fs::rename(&tmp, path)?;
Ok(true)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn client_is_send_sync_clone() {
fn assert_bounds<T: Send + Sync + Clone>() {}
assert_bounds::<DatasetClient>();
}
#[test]
fn transient_classification() {
assert!(HttpError::Status {
code: 429,
url: String::new()
}
.is_transient());
assert!(HttpError::Status {
code: 503,
url: String::new()
}
.is_transient());
assert!(!HttpError::Status {
code: 404,
url: String::new()
}
.is_transient());
assert!(!HttpError::Status {
code: 403,
url: String::new()
}
.is_transient());
assert!(HttpError::Transport("boom".into()).is_transient());
}
#[test]
fn rate_gate_enforces_min_gap() {
let client = DatasetClient::new(DatasetClientConfig {
user_agent: "test".into(),
connect_timeout: Duration::from_secs(5),
overall_timeout: Some(Duration::from_secs(5)),
rate_per_sec: NonZeroU32::new(20),
retry_count: 0,
base_backoff_ms: 100,
});
client.rate_gate();
let start = Instant::now();
client.rate_gate();
assert!(
start.elapsed() >= Duration::from_millis(40),
"second gate should have waited ~50ms, waited {:?}",
start.elapsed()
);
}
#[test]
fn no_rate_gate_when_unset() {
let client = DatasetClient::new(DatasetClientConfig {
user_agent: "test".into(),
connect_timeout: Duration::from_secs(5),
overall_timeout: None,
rate_per_sec: None,
retry_count: 0,
base_backoff_ms: 100,
});
let start = Instant::now();
client.rate_gate();
client.rate_gate();
assert!(start.elapsed() < Duration::from_millis(20));
}
}