use std::time::{Duration, Instant};
use dashmap::DashMap;
use once_cell::sync::Lazy;
pub static FIXED_LIMITER: Lazy<FixedRateLimiter> = Lazy::new(FixedRateLimiter::new);
#[derive(Debug, Clone)]
pub struct FixedRateLimiter {
buckets: DashMap<String, RateBucket>,
}
#[derive(Debug, Clone)]
struct RateBucket {
count: u64,
window_start: Instant,
}
impl FixedRateLimiter {
pub fn new() -> Self {
Self {
buckets: DashMap::new(),
}
}
pub fn check(&self, key: &str, max: u64, window: Duration) -> bool {
let now = Instant::now();
let mut bucket = self.buckets.entry(key.to_string()).or_insert(RateBucket {
count: 0,
window_start: now,
});
if now.duration_since(bucket.window_start) > window {
bucket.count = 1;
bucket.window_start = now;
true
} else if bucket.count < max {
bucket.count += 1;
true
} else {
false
}
}
pub fn remaining(&self, key: &str, max: u64, window: Duration) -> u64 {
let now = Instant::now();
if let Some(bucket) = self.buckets.get(key) {
if now.duration_since(bucket.window_start) > window {
max
} else {
max.saturating_sub(bucket.count)
}
} else {
max
}
}
pub fn reset_in(&self, key: &str, window: Duration) -> Duration {
let now = Instant::now();
if let Some(bucket) = self.buckets.get(key) {
let elapsed = now.duration_since(bucket.window_start);
if elapsed > window {
Duration::ZERO
} else {
window - elapsed
}
} else {
Duration::ZERO
}
}
}
impl Default for FixedRateLimiter {
fn default() -> Self {
Self::new()
}
}
pub fn parse_duration(unit: &str) -> Duration {
match unit.trim().to_lowercase().as_str() {
"second" | "seconds" => Duration::from_secs(1),
"minute" | "minutes" => Duration::from_secs(60),
"hour" | "hours" => Duration::from_secs(3600),
"day" | "days" => Duration::from_secs(86400),
_ => Duration::from_secs(1),
}
}