#[derive(Debug, Clone, Default)]
pub struct RateLimitInfo {
pub limit: Option<u64>,
pub remaining: Option<u64>,
pub reset_secs: Option<u64>,
pub cost: Option<u64>,
}
impl std::fmt::Display for RateLimitInfo {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let mut parts = Vec::new();
if let (Some(remaining), Some(limit)) = (self.remaining, self.limit) {
let cost = self.cost.filter(|cost| *cost > 0).unwrap_or(1);
parts.push(format!(
"remaining={}/{} reqs",
remaining / cost,
limit / cost,
));
} else {
if let Some(remaining) = self.remaining {
parts.push(format!("remaining={remaining}"));
}
if let Some(limit) = self.limit {
parts.push(format!("limit={limit}"));
}
}
if let Some(reset) = self.reset_secs {
parts.push(format!("resets_in={reset}s"));
}
write!(f, "{}", parts.join(", "))
}
}
impl RateLimitInfo {
pub(crate) fn from_response(res: &reqwest::Response) -> Self {
Self {
limit: Self::parse_limit_header(res),
remaining: Self::parse_u64_header(res, "x-ratelimit-remaining"),
reset_secs: Self::parse_u64_header(res, "x-ratelimit-reset"),
cost: Self::parse_u64_header(res, "x-ratelimit-cost"),
}
}
pub fn is_rate_limited(&self) -> bool {
self.remaining == Some(0)
}
pub fn suggested_wait_secs(&self) -> Option<u64> {
self.reset_secs
}
fn parse_limit_header(res: &reqwest::Response) -> Option<u64> {
let value = res.headers().get("x-ratelimit-limit")?.to_str().ok()?;
let first = value.split(',').next()?.trim();
first.parse().ok()
}
fn parse_u64_header(res: &reqwest::Response, name: &str) -> Option<u64> {
res.headers().get(name)?.to_str().ok()?.trim().parse().ok()
}
}
#[derive(Debug, Clone)]
pub struct QueryResponseWithRateLimit<T> {
pub response: T,
pub rate_limit: RateLimitInfo,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_is_rate_limited() {
let info = RateLimitInfo {
remaining: Some(0),
..Default::default()
};
assert!(info.is_rate_limited());
let info = RateLimitInfo {
remaining: Some(5),
..Default::default()
};
assert!(!info.is_rate_limited());
let info = RateLimitInfo::default();
assert!(!info.is_rate_limited());
}
#[test]
fn test_from_response_header_case_insensitive() {
let http_resp = http::Response::builder()
.header("X-RateLimit-Remaining", "42")
.header("X-RATELIMIT-RESET", "30")
.header("X-Ratelimit-Limit", "100, 100;w=60")
.header("X-Ratelimit-Cost", "10")
.body("")
.unwrap();
let resp: reqwest::Response = http_resp.into();
let info = RateLimitInfo::from_response(&resp);
assert_eq!(
(info.limit, info.remaining, info.reset_secs, info.cost),
(Some(100), Some(42), Some(30), Some(10))
);
}
#[test]
fn test_from_response_ignores_malformed_headers() {
let http_resp = http::Response::builder()
.header("x-ratelimit-limit", "not-a-number, 50;w=60")
.header("x-ratelimit-remaining", "-1")
.header("x-ratelimit-reset", "tomorrow")
.header("x-ratelimit-cost", "1.5")
.body("")
.unwrap();
let resp: reqwest::Response = http_resp.into();
let info = RateLimitInfo::from_response(&resp);
assert_eq!(
(info.limit, info.remaining, info.reset_secs, info.cost),
(None, None, None, None)
);
}
#[test]
fn test_suggested_wait_secs() {
let info = RateLimitInfo {
reset_secs: Some(30),
..Default::default()
};
assert_eq!(info.suggested_wait_secs(), Some(30));
let info = RateLimitInfo::default();
assert_eq!(info.suggested_wait_secs(), None);
}
#[test]
fn test_display_full() {
let info = RateLimitInfo {
limit: Some(50),
remaining: Some(0),
reset_secs: Some(59),
cost: Some(10),
};
assert_eq!(info.to_string(), "remaining=0/5 reqs, resets_in=59s");
}
#[test]
fn test_display_partial() {
let info = RateLimitInfo {
remaining: Some(3),
reset_secs: Some(30),
..Default::default()
};
assert_eq!(info.to_string(), "remaining=3, resets_in=30s");
}
#[test]
fn test_display_empty() {
let info = RateLimitInfo::default();
assert_eq!(info.to_string(), "");
}
#[test]
fn test_display_zero_cost_does_not_panic() {
let info = RateLimitInfo {
limit: Some(50),
remaining: Some(0),
cost: Some(0),
..Default::default()
};
assert_eq!(info.to_string(), "remaining=0/50 reqs");
}
}