use chrono::{DateTime, Utc};
use reqwest::header::HeaderMap;
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use crate::auth::InstallationId;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum RateLimitContext {
App,
Installation(InstallationId),
}
#[derive(Debug, Clone)]
pub struct RateLimit {
limit: u32,
remaining: u32,
reset_at: DateTime<Utc>,
resource: String,
}
impl RateLimit {
pub fn new(
limit: u32,
remaining: u32,
reset_at: DateTime<Utc>,
resource: impl Into<String>,
) -> Self {
Self {
limit,
remaining,
reset_at,
resource: resource.into(),
}
}
pub fn limit(&self) -> u32 {
self.limit
}
pub fn remaining(&self) -> u32 {
self.remaining
}
pub fn reset_at(&self) -> DateTime<Utc> {
self.reset_at
}
pub fn resource(&self) -> &str {
&self.resource
}
pub fn is_exhausted(&self) -> bool {
self.remaining == 0
}
pub fn is_near_exhaustion(&self, margin: f64) -> bool {
let threshold = (self.limit as f64 * margin) as u32;
self.remaining <= threshold
}
pub fn has_reset(&self) -> bool {
Utc::now() >= self.reset_at
}
}
pub fn parse_rate_limit_from_headers(headers: &HeaderMap) -> Option<RateLimit> {
let limit_str = headers.get("x-ratelimit-limit")?.to_str().ok()?;
let remaining_str = headers.get("x-ratelimit-remaining")?.to_str().ok()?;
let reset_str = headers.get("x-ratelimit-reset")?.to_str().ok()?;
let limit = limit_str.parse::<u32>().ok()?;
let remaining = remaining_str.parse::<u32>().ok()?;
let reset_timestamp = reset_str.parse::<i64>().ok()?;
let reset_at = DateTime::from_timestamp(reset_timestamp, 0)?;
let resource = headers
.get("x-ratelimit-resource")
.and_then(|v| v.to_str().ok())
.unwrap_or("core")
.to_string();
Some(RateLimit::new(limit, remaining, reset_at, resource))
}
#[derive(Debug, Clone)]
pub struct RateLimiter {
limits: Arc<RwLock<HashMap<(RateLimitContext, String), RateLimit>>>,
margin: f64,
}
impl RateLimiter {
pub fn new(margin: f64) -> Self {
Self {
limits: Arc::new(RwLock::new(HashMap::new())),
margin: margin.clamp(0.0, 1.0),
}
}
pub fn update_from_headers(&self, context: &RateLimitContext, headers: &HeaderMap) {
if let Some(rate_limit) = parse_rate_limit_from_headers(headers) {
let resource = rate_limit.resource().to_string();
if let Ok(mut limits) = self.limits.write() {
limits.insert((context.clone(), resource), rate_limit);
}
}
}
pub fn can_proceed(&self, context: &RateLimitContext, resource: &str) -> bool {
let limits = match self.limits.read() {
Ok(limits) => limits,
Err(_) => return true, };
let key = (context.clone(), resource.to_string());
let rate_limit = match limits.get(&key) {
Some(limit) => limit,
None => return true,
};
if rate_limit.has_reset() {
return true;
}
!rate_limit.is_exhausted() && !rate_limit.is_near_exhaustion(self.margin)
}
pub fn get_limit(&self, context: &RateLimitContext, resource: &str) -> Option<RateLimit> {
let key = (context.clone(), resource.to_string());
self.limits.read().ok()?.get(&key).cloned()
}
}
impl Default for RateLimiter {
fn default() -> Self {
Self::new(0.1)
}
}
#[cfg(test)]
#[path = "rate_limit_tests.rs"]
mod tests;