use std::{ collections::HashMap, sync::Arc, time::{Duration, Instant}, }; use tokio::sync::Mutex; use crate::error::Result; #[derive(Clone)] pub struct RateLimiter { inner: Arc>, } struct RateLimiterInner { requests: HashMap>, config: RateLimitConfig, } #[derive(Clone, Debug)] pub struct RateLimitConfig { pub modrinth_requests_per_min: u32, pub modrinth_burst: u32, pub curseforge_requests_per_min: u32, pub curseforge_burst: u32, pub github_requests_per_min: u32, pub github_burst: u32, pub default_requests_per_min: u32, pub default_burst: u32, } impl Default for RateLimitConfig { fn default() -> Self { Self { modrinth_requests_per_min: 100, modrinth_burst: 10, curseforge_requests_per_min: 60, curseforge_burst: 5, github_requests_per_min: 50, github_burst: 5, default_requests_per_min: 30, default_burst: 3, } } } impl RateLimiter { pub fn new(config: Option) -> Self { Self { inner: Arc::new(Mutex::new(RateLimiterInner { requests: HashMap::new(), config: config.unwrap_or_default(), })), } } pub async fn acquire(&self, platform: &str) -> Result<()> { let config = { let inner = self.inner.lock().await; inner.config.clone() }; let (rate, burst) = match platform.to_lowercase().as_str() { "modrinth" => (config.modrinth_requests_per_min, config.modrinth_burst), "curseforge" => { (config.curseforge_requests_per_min, config.curseforge_burst) }, "github" => (config.github_requests_per_min, config.github_burst), _ => (config.default_requests_per_min, config.default_burst), }; let interval = Duration::from_secs(60) / rate.max(1); let mut inner = self.inner.lock().await; let now = Instant::now(); let platform_requests = inner.requests.entry(platform.to_string()).or_default(); platform_requests .retain(|t| now.duration_since(*t) < Duration::from_secs(60)); if platform_requests.len() >= burst as usize && let Some(oldest) = platform_requests.first() { let wait_time = interval.saturating_sub(now.duration_since(*oldest)); if wait_time > Duration::ZERO { drop(inner); tokio::time::sleep(wait_time).await; let mut inner = self.inner.lock().await; let platform_requests = inner.requests.entry(platform.to_string()).or_default(); platform_requests.push(Instant::now()); return Ok(()); } } platform_requests.push(Instant::now()); Ok(()) } pub async fn wait_for(&self, platform: &str) { let _ = self.acquire(platform).await; } }