use std::time::Duration;
#[derive(Debug, Clone, Default)]
pub struct RateLimitInfo {
pub limit: Option<u32>,
pub remaining: Option<u32>,
pub reset_after: Option<Duration>,
pub retry_after: Option<Duration>,
pub is_rate_limited: bool,
}
impl RateLimitInfo {
pub fn from_headers<'a>(headers: impl Iterator<Item = (&'a str, &'a str)>) -> Self {
let mut info = Self::default();
for (name, value) in headers {
let lower = name.to_lowercase();
match lower.as_str() {
"x-ratelimit-limit" | "x-rate-limit-limit" => {
info.limit = value.parse().ok();
}
"x-ratelimit-remaining" | "x-rate-limit-remaining" => {
info.remaining = value.parse().ok();
}
"x-ratelimit-reset" | "x-rate-limit-reset" => {
if let Ok(secs) = value.parse::<u64>() {
info.reset_after = Some(Duration::from_secs(secs));
}
}
"x-rate-limit-cu-second" => {
info.limit = value.parse().ok();
}
"x-rate-limit-request-second" => {
if info.limit.is_none() {
info.limit = value.parse().ok();
}
}
"retry-after" => {
info.retry_after = parse_retry_after(value);
info.is_rate_limited = true;
}
_ => {}
}
}
info
}
pub fn should_backoff(&self) -> bool {
self.is_rate_limited || self.remaining == Some(0)
}
pub fn suggested_wait(&self) -> Option<Duration> {
self.retry_after.or(self.reset_after).or_else(|| {
if self.should_backoff() {
Some(Duration::from_secs(1))
} else {
None
}
})
}
}
fn parse_retry_after(value: &str) -> Option<Duration> {
if let Ok(secs) = value.parse::<u64>() {
return Some(Duration::from_secs(secs));
}
if let Ok(secs) = value.parse::<f64>() {
return Some(Duration::from_secs_f64(secs));
}
Some(Duration::from_secs(1))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn standard_headers() {
let headers = vec![
("X-RateLimit-Limit", "100"),
("X-RateLimit-Remaining", "42"),
("X-RateLimit-Reset", "30"),
];
let info = RateLimitInfo::from_headers(headers.iter().map(|(k, v)| (*k, *v)));
assert_eq!(info.limit, Some(100));
assert_eq!(info.remaining, Some(42));
assert_eq!(info.reset_after, Some(Duration::from_secs(30)));
assert!(!info.is_rate_limited);
assert!(!info.should_backoff());
}
#[test]
fn retry_after_seconds() {
let headers = vec![("Retry-After", "5")];
let info = RateLimitInfo::from_headers(headers.iter().map(|(k, v)| (*k, *v)));
assert!(info.is_rate_limited);
assert_eq!(info.retry_after, Some(Duration::from_secs(5)));
assert!(info.should_backoff());
}
#[test]
fn remaining_zero_triggers_backoff() {
let headers = vec![("X-RateLimit-Remaining", "0")];
let info = RateLimitInfo::from_headers(headers.iter().map(|(k, v)| (*k, *v)));
assert!(info.should_backoff());
assert!(info.suggested_wait().is_some());
}
#[test]
fn alchemy_cu_headers() {
let headers = vec![("x-rate-limit-cu-second", "330")];
let info = RateLimitInfo::from_headers(headers.iter().map(|(k, v)| (*k, *v)));
assert_eq!(info.limit, Some(330));
}
#[test]
fn case_insensitive() {
let headers = vec![
("x-ratelimit-limit", "200"),
("X-RATELIMIT-REMAINING", "50"),
];
let info = RateLimitInfo::from_headers(headers.iter().map(|(k, v)| (*k, *v)));
assert_eq!(info.limit, Some(200));
assert_eq!(info.remaining, Some(50));
}
#[test]
fn empty_headers() {
let info = RateLimitInfo::from_headers(std::iter::empty());
assert!(!info.should_backoff());
assert!(info.suggested_wait().is_none());
}
#[test]
fn suggested_wait_prefers_retry_after() {
let headers = vec![("Retry-After", "10"), ("X-RateLimit-Reset", "30")];
let info = RateLimitInfo::from_headers(headers.iter().map(|(k, v)| (*k, *v)));
assert_eq!(info.suggested_wait(), Some(Duration::from_secs(10)));
}
}