use std::{ collections::HashMap, sync::Arc, time::{Duration, Instant}, }; use tokio::{sync::RwLock, time::interval}; #[derive(Debug, Clone)] pub struct RateLimitConfig { pub read_requests_per_window: u32, pub write_requests_per_window: u32, pub window_secs: u64, } impl Default for RateLimitConfig { fn default() -> Self { Self { read_requests_per_window: 120, write_requests_per_window: 30, window_secs: 60, } } } #[derive(Debug)] struct RateLimitBucket { read_count: u32, write_count: u32, reset_time: Instant, } #[derive(Clone, Copy)] enum BucketOp { Read, Write, } pub struct RateLimiter { buckets: Arc>>, config: RateLimitConfig, } impl RateLimiter { pub fn new(config: RateLimitConfig) -> Self { Self { buckets: Arc::new(RwLock::new(HashMap::new())), config, } } pub async fn is_read_allowed(&self) -> bool { self.is_allowed( "global:read", BucketOp::Read, self.config.read_requests_per_window, ) .await } pub async fn is_write_allowed(&self) -> bool { self.is_allowed( "global:write", BucketOp::Write, self.config.write_requests_per_window, ) .await } pub async fn is_repo_write_allowed(&self, repo_path: &str) -> bool { let key = format!("repo:write:{}", repo_path); self.is_allowed( &key, BucketOp::Write, self.config.write_requests_per_window, ) .await } async fn is_allowed(&self, key: &str, op: BucketOp, limit: u32) -> bool { let now = Instant::now(); let mut buckets = self.buckets.write().await; let bucket = buckets .entry(key.to_string()) .or_insert_with(|| RateLimitBucket { read_count: 0, write_count: 0, reset_time: now + Duration::from_secs(self.config.window_secs), }); if now >= bucket.reset_time { bucket.read_count = 0; bucket.write_count = 0; bucket.reset_time = now + Duration::from_secs(self.config.window_secs); } let over_limit = match op { BucketOp::Read => bucket.read_count >= limit, BucketOp::Write => bucket.write_count >= limit, }; if over_limit { return false; } match op { BucketOp::Read => bucket.read_count += 1, BucketOp::Write => bucket.write_count += 1, } true } pub fn start_cleanup(self: Arc) -> tokio::task::JoinHandle<()> { tokio::spawn(async move { let mut ticker = interval(Duration::from_secs(300)); loop { ticker.tick().await; let now = Instant::now(); let mut buckets = self.buckets.write().await; buckets.retain(|_, bucket| now < bucket.reset_time); } }) } }