use std::time::{Duration, Instant};
use tokio::sync::RwLock;
#[derive(Debug)]
pub struct EnhancedRateLimiter {
pub max_requests: u32,
pub window: Duration,
requests: RwLock<Vec<Instant>>,
burst_allowance: u32,
}
impl EnhancedRateLimiter {
pub fn new(max_requests: u32, window_secs: u64, burst_allowance: u32) -> Self {
Self {
max_requests,
window: Duration::from_secs(window_secs),
requests: RwLock::new(Vec::new()),
burst_allowance,
}
}
pub async fn check_limit(&self) -> bool {
let now = Instant::now();
let mut requests = self.requests.write().await;
requests.retain(|&time| now.duration_since(time) < self.window);
let current_count = requests.len();
let effective_limit = self.max_requests + self.burst_allowance;
if current_count < effective_limit as usize {
requests.push(now);
true
} else {
false
}
}
pub async fn get_remaining(&self) -> u32 {
let requests = self.requests.read().await;
self.max_requests.saturating_sub(requests.len() as u32)
}
}