use std::collections::HashMap;
use std::sync::RwLock;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};
use serde::{Deserialize, Serialize};
use crate::pipeline::traits::RateLimitable;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct RateLimiterMetrics {
pub window_seconds: u32,
pub global_rate_limit: u32,
pub global_request_count: u64,
pub tracked_senders: usize,
pub max_tracked_users: usize,
pub global_utilization_pct: Option<f64>,
}
struct SlidingWindow {
timestamps: Vec<Instant>,
last_access: u64,
}
pub struct RateLimiter {
windows: RwLock<HashMap<String, SlidingWindow>>,
window_seconds: u32,
window_duration: Duration,
global_counter: AtomicU64,
global_rate_limit: u32,
global_reset: RwLock<Instant>,
max_tracked_users: usize,
access_counter: AtomicU64,
}
impl RateLimiter {
pub fn new(window_seconds: u32, global_rate_limit: u32) -> Self {
Self {
windows: RwLock::new(HashMap::new()),
window_seconds,
window_duration: Duration::from_secs(u64::from(window_seconds)),
global_counter: AtomicU64::new(0),
global_rate_limit,
global_reset: RwLock::new(Instant::now()),
max_tracked_users: 10_000,
access_counter: AtomicU64::new(0),
}
}
pub fn with_max_tracked_users(mut self, max: usize) -> Self {
self.max_tracked_users = max;
self
}
pub fn check(&self, sender_id: &str, limit: u32) -> bool {
if !self.check_global() {
return false;
}
if limit == 0 {
return true;
}
let now = Instant::now();
let order = self.access_counter.fetch_add(1, Ordering::Relaxed);
let mut windows = self.windows.write().unwrap();
let entry = windows
.entry(sender_id.to_string())
.or_insert_with(|| SlidingWindow {
timestamps: Vec::new(),
last_access: order,
});
entry.last_access = order;
entry
.timestamps
.retain(|ts| now.duration_since(*ts) < self.window_duration);
if entry.timestamps.len() >= limit as usize {
return false;
}
entry.timestamps.push(now);
let needs_eviction = windows.len() > self.max_tracked_users;
if needs_eviction {
self.evict_oldest(&mut windows);
}
true
}
pub fn check_global(&self) -> bool {
if self.global_rate_limit == 0 {
return true;
}
let now = Instant::now();
{
let mut window_start = self.global_reset.write().unwrap();
if now.duration_since(*window_start) >= self.window_duration {
self.global_counter.store(0, Ordering::Relaxed);
*window_start = now;
}
}
let prev = self.global_counter.fetch_add(1, Ordering::Relaxed);
if prev >= u64::from(self.global_rate_limit) {
self.global_counter.fetch_sub(1, Ordering::Relaxed);
return false;
}
true
}
pub fn get_count(&self, sender_id: &str) -> u32 {
let now = Instant::now();
let windows = self.windows.read().unwrap();
windows.get(sender_id).map_or(0, |entry| {
entry
.timestamps
.iter()
.filter(|ts| now.duration_since(**ts) < self.window_duration)
.count() as u32
})
}
pub fn tracked_senders(&self) -> usize {
self.windows.read().unwrap().len()
}
pub fn global_request_count(&self) -> u64 {
self.global_counter.load(Ordering::Relaxed)
}
pub fn global_rate_limit(&self) -> u32 {
self.global_rate_limit
}
pub fn window_duration(&self) -> Duration {
self.window_duration
}
pub fn window_seconds(&self) -> u32 {
self.window_seconds
}
pub fn max_tracked_users(&self) -> usize {
self.max_tracked_users
}
pub fn metrics(&self) -> RateLimiterMetrics {
let global_request_count = self.global_request_count();
let global_rate_limit = self.global_rate_limit;
let global_utilization_pct = if global_rate_limit > 0 {
Some((global_request_count as f64 / f64::from(global_rate_limit)) * 100.0)
} else {
None
};
RateLimiterMetrics {
window_seconds: self.window_seconds,
global_rate_limit,
global_request_count,
tracked_senders: self.tracked_senders(),
max_tracked_users: self.max_tracked_users,
global_utilization_pct,
}
}
pub fn clear(&self) {
let mut windows = self.windows.write().unwrap();
windows.clear();
self.access_counter.store(0, Ordering::Relaxed);
self.global_counter.store(0, Ordering::Relaxed);
let mut reset = self.global_reset.write().unwrap();
*reset = Instant::now();
}
pub fn flush_lru(&self) {
self.clear();
}
fn evict_oldest(&self, windows: &mut HashMap<String, SlidingWindow>) {
if windows.len() <= self.max_tracked_users {
return;
}
let mut oldest_key: Option<String> = None;
let mut oldest_access = u64::MAX;
for (key, entry) in windows.iter() {
if entry.last_access < oldest_access {
oldest_access = entry.last_access;
oldest_key = Some(key.clone());
}
}
if let Some(key) = oldest_key {
windows.remove(&key);
}
}
pub fn evict_if_needed(&self) {
let mut windows = self.windows.write().unwrap();
self.evict_oldest(&mut windows);
}
}
impl RateLimitable for RateLimiter {
fn check(&self, sender_id: &str, limit: u32) -> bool {
RateLimiter::check(self, sender_id, limit)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use std::thread;
use std::time::Duration;
#[test]
fn test_allows_under_limit() {
let limiter = RateLimiter::new(60, 0);
for _ in 0..10 {
assert!(limiter.check("user_1", 10));
}
}
#[test]
fn test_rejects_at_limit() {
let limiter = RateLimiter::new(60, 0);
for _ in 0..5 {
assert!(limiter.check("user_1", 5));
}
assert!(!limiter.check("user_1", 5));
}
#[test]
fn test_window_expiry() {
let limiter = RateLimiter::new(1, 0);
for _ in 0..3 {
assert!(limiter.check("user_1", 3));
}
assert!(!limiter.check("user_1", 3));
thread::sleep(Duration::from_millis(1100));
assert!(limiter.check("user_1", 3));
}
#[test]
fn test_global_rate_limit_rejects() {
let limiter = RateLimiter::new(60, 5);
for i in 0..5 {
let sender = format!("user_{}", i);
assert!(limiter.check(&sender, 0), "request {} should be allowed", i);
}
assert!(!limiter.check("user_new", 0));
assert_eq!(limiter.global_request_count(), 5);
}
#[test]
fn test_global_window_resets() {
let limiter = RateLimiter::new(1, 2);
assert!(limiter.check("user_1", 0));
assert!(limiter.check("user_2", 0));
assert!(!limiter.check("user_3", 0));
thread::sleep(Duration::from_millis(1100));
assert!(limiter.check("user_4", 0));
assert_eq!(limiter.global_request_count(), 1);
}
#[test]
fn test_unlimited_per_user() {
let limiter = RateLimiter::new(60, 0);
for _ in 0..1000 {
assert!(limiter.check("user_1", 0));
}
assert_eq!(limiter.tracked_senders(), 0);
}
#[test]
fn test_global_limit_zero_is_unlimited() {
let limiter = RateLimiter::new(60, 0);
for i in 0..1000 {
let sender = format!("user_{}", i);
assert!(limiter.check(&sender, 0));
}
}
#[test]
fn test_independent_senders() {
let limiter = RateLimiter::new(60, 0);
for _ in 0..5 {
limiter.check("user_1", 5);
}
assert!(!limiter.check("user_1", 5));
assert!(limiter.check("user_2", 5));
}
#[test]
fn test_lru_eviction() {
let limiter = RateLimiter::new(60, 0).with_max_tracked_users(3);
limiter.check("user_a", 10);
limiter.check("user_b", 10);
limiter.check("user_c", 10);
limiter.check("user_d", 10);
assert!(limiter.tracked_senders() <= 3);
}
#[test]
fn test_get_count() {
let limiter = RateLimiter::new(60, 0);
assert_eq!(limiter.get_count("user_1"), 0);
limiter.check("user_1", 10);
limiter.check("user_1", 10);
limiter.check("user_1", 10);
assert_eq!(limiter.get_count("user_1"), 3);
}
#[test]
fn test_concurrent_access() {
let limiter = Arc::new(RateLimiter::new(60, 0));
let mut handles = vec![];
for i in 0..10 {
let limiter = Arc::clone(&limiter);
handles.push(thread::spawn(move || {
let sender = format!("user_{}", i);
for _ in 0..100 {
limiter.check(&sender, 200);
}
}));
}
for handle in handles {
handle.join().unwrap();
}
assert_eq!(limiter.tracked_senders(), 10);
}
#[test]
fn test_empty_sender_id() {
let limiter = RateLimiter::new(60, 0);
assert!(limiter.check("", 2));
assert!(limiter.check("", 2));
assert!(!limiter.check("", 2));
}
#[test]
fn test_global_limit_before_per_user() {
let limiter = RateLimiter::new(60, 3);
assert!(limiter.check("user_1", 100));
assert!(limiter.check("user_2", 100));
assert!(limiter.check("user_3", 100));
assert!(!limiter.check("user_4", 100));
}
#[test]
fn test_clear() {
let limiter = RateLimiter::new(60, 100);
for _ in 0..5 {
limiter.check("user_1", 5);
}
assert!(!limiter.check("user_1", 5));
assert!(limiter.global_request_count() > 0);
limiter.clear();
assert!(limiter.check("user_1", 5));
assert_eq!(limiter.tracked_senders(), 1);
assert_eq!(limiter.global_request_count(), 1);
}
#[test]
fn test_rate_limitable_trait_impl() {
let limiter = RateLimiter::new(60, 0);
let trait_obj: &dyn RateLimitable = &limiter;
assert!(trait_obj.check("user_1", 2));
assert!(trait_obj.check("user_1", 2));
assert!(!trait_obj.check("user_1", 2));
}
#[test]
fn test_different_limits_per_call() {
let limiter = RateLimiter::new(60, 0);
assert!(limiter.check("user_1", 5));
assert!(limiter.check("user_1", 5));
assert!(!limiter.check("user_1", 2));
assert!(limiter.check("user_1", 5));
}
#[test]
fn test_global_limit_concurrent() {
let limiter = Arc::new(RateLimiter::new(60, 50));
let mut handles = vec![];
for i in 0..10 {
let limiter = Arc::clone(&limiter);
handles.push(thread::spawn(move || {
let sender = format!("user_{}", i);
let mut allowed = 0u32;
for _ in 0..20 {
if limiter.check(&sender, 100) {
allowed += 1;
}
}
allowed
}));
}
let total_allowed: u32 = handles.into_iter().map(|h| h.join().unwrap()).sum();
assert!(
total_allowed <= 50,
"total_allowed={} exceeded global limit 50",
total_allowed
);
}
#[test]
fn test_lru_evicts_oldest() {
let limiter = RateLimiter::new(60, 0).with_max_tracked_users(2);
limiter.check("user_a", 10);
limiter.check("user_b", 10);
limiter.check("user_c", 10);
assert_eq!(limiter.tracked_senders(), 2);
assert_eq!(limiter.get_count("user_a"), 0);
assert_eq!(limiter.get_count("user_b"), 1);
assert_eq!(limiter.get_count("user_c"), 1);
}
#[test]
fn test_zero_window_allows_all() {
let limiter = RateLimiter::new(0, 0);
for _ in 0..100 {
limiter.check("user_1", 1);
}
}
#[test]
fn test_builder_max_tracked_users() {
let limiter = RateLimiter::new(60, 0).with_max_tracked_users(5);
for i in 0..10 {
let sender = format!("user_{}", i);
limiter.check(&sender, 10);
}
assert!(limiter.tracked_senders() <= 5);
}
#[test]
fn test_rejected_request_no_timestamp() {
let limiter = RateLimiter::new(60, 0);
assert!(limiter.check("user_1", 2));
assert!(limiter.check("user_1", 2));
assert_eq!(limiter.get_count("user_1"), 2);
assert!(!limiter.check("user_1", 2));
assert_eq!(limiter.get_count("user_1"), 2);
}
#[test]
fn test_metrics_snapshot() {
let limiter = RateLimiter::new(60, 10).with_max_tracked_users(100);
limiter.check("alice", 5);
limiter.check("bob", 5);
let m = limiter.metrics();
assert_eq!(m.window_seconds, 60);
assert_eq!(m.global_rate_limit, 10);
assert_eq!(m.global_request_count, 2);
assert_eq!(m.tracked_senders, 2);
assert_eq!(m.max_tracked_users, 100);
assert!((m.global_utilization_pct.unwrap() - 20.0).abs() < f64::EPSILON);
let unlimited = RateLimiter::new(30, 0);
assert!(unlimited.metrics().global_utilization_pct.is_none());
}
#[test]
fn test_flush_lru() {
let limiter = RateLimiter::new(60, 100);
limiter.check("u1", 10);
limiter.check("u2", 10);
assert_eq!(limiter.tracked_senders(), 2);
assert!(limiter.global_request_count() > 0);
limiter.flush_lru();
assert_eq!(limiter.tracked_senders(), 0);
assert_eq!(limiter.global_request_count(), 0);
assert_eq!(limiter.get_count("u1"), 0);
}
#[test]
fn test_evict_if_needed_public() {
let limiter = RateLimiter::new(60, 0).with_max_tracked_users(2);
limiter.check("a", 10);
limiter.check("b", 10);
limiter.evict_if_needed();
assert_eq!(limiter.tracked_senders(), 2);
limiter.check("c", 10);
assert_eq!(limiter.tracked_senders(), 2);
limiter.evict_if_needed();
assert!(limiter.tracked_senders() <= 2);
}
}