use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::time::Duration;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RateLimitInfo {
pub limit: u64,
pub remaining: u64,
pub reset_at: DateTime<Utc>,
pub is_limited: bool,
}
impl RateLimitInfo {
pub fn from_headers(
limit: Option<&str>,
remaining: Option<&str>,
reset: Option<&str>,
) -> Option<Self> {
let limit = limit?.parse::<u64>().ok()?;
let remaining = remaining?.parse::<u64>().ok()?;
let reset_timestamp = reset?.parse::<i64>().ok()?;
let reset_at = DateTime::from_timestamp(reset_timestamp, 0)?;
let is_limited = remaining == 0;
Some(RateLimitInfo {
limit,
remaining,
reset_at,
is_limited,
})
}
pub fn is_near_limit(&self, threshold_pct: f64) -> bool {
let threshold = (self.limit as f64 * threshold_pct) as u64;
self.remaining < threshold
}
pub fn time_until_reset(&self) -> Duration {
let now = Utc::now();
if self.reset_at > now {
Duration::from_secs((self.reset_at - now).num_seconds() as u64)
} else {
Duration::from_secs(0)
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RetryPolicy {
pub max_retries: u32,
pub initial_delay: Duration,
pub max_delay: Duration,
pub backoff_multiplier: f64,
pub use_jitter: bool,
}
impl Default for RetryPolicy {
fn default() -> Self {
Self {
max_retries: 3,
initial_delay: Duration::from_millis(100),
max_delay: Duration::from_secs(60),
backoff_multiplier: 2.0,
use_jitter: true,
}
}
}
impl RetryPolicy {
pub fn new(max_retries: u32, initial_delay: Duration, max_delay: Duration) -> Self {
Self {
max_retries,
initial_delay,
max_delay,
backoff_multiplier: 2.0,
use_jitter: true,
}
}
pub fn with_jitter(mut self) -> Self {
self.use_jitter = true;
self
}
pub fn without_jitter(mut self) -> Self {
self.use_jitter = false;
self
}
pub fn calculate_delay(&self, attempt: u32) -> Duration {
if attempt == 0 {
return Duration::from_secs(0);
}
let multiplier = self.backoff_multiplier.powi(attempt as i32 - 1);
let delay_ms = (self.initial_delay.as_millis() as f64 * multiplier) as u64;
let mut delay = Duration::from_millis(delay_ms);
if delay > self.max_delay {
delay = self.max_delay;
}
if self.use_jitter {
use rand::RngExt;
let mut rng = rand::rng();
let jitter_factor = rng.random_range(0.75..=1.25);
delay = Duration::from_millis((delay.as_millis() as f64 * jitter_factor) as u64);
}
delay
}
pub fn should_retry(&self, attempt: u32) -> bool {
attempt < self.max_retries
}
}
pub fn parse_retry_after(retry_after: &str) -> Option<Duration> {
if let Ok(seconds) = retry_after.parse::<u64>() {
return Some(Duration::from_secs(seconds));
}
if let Ok(date_time) = chrono::DateTime::parse_from_rfc2822(retry_after) {
let now = Utc::now();
let retry_time = date_time.with_timezone(&Utc);
if retry_time > now {
let duration = (retry_time - now).num_seconds();
if duration > 0 {
return Some(Duration::from_secs(duration as u64));
}
}
}
None
}
pub fn calculate_rate_limit_delay(
retry_after: Option<&str>,
rate_limit_reset: Option<&str>,
) -> Duration {
if let Some(retry_after_value) = retry_after {
if let Some(delay) = parse_retry_after(retry_after_value) {
return delay;
}
}
if let Some(reset_value) = rate_limit_reset {
if let Ok(reset_timestamp) = reset_value.parse::<i64>() {
if let Some(reset_time) = DateTime::from_timestamp(reset_timestamp, 0) {
let now = Utc::now();
if reset_time > now {
let duration = (reset_time - now).num_seconds();
if duration > 0 {
return Duration::from_secs(duration as u64);
}
}
}
}
}
Duration::from_secs(60)
}
pub fn detect_secondary_rate_limit(status: u16, body: &str) -> bool {
if status != 403 {
return false;
}
let body_lower = body.to_lowercase();
body_lower.contains("rate limit")
|| body_lower.contains("rate_limit")
|| body_lower.contains("abuse")
|| body_lower.contains("too many requests")
}
#[cfg(test)]
#[path = "retry_tests.rs"]
mod tests;