initial commit

Signed-off-by: NotAShelf <raf@notashelf.dev>
Change-Id: Ife1391ed23a1e7f388b1b5eca90b9ea76a6a6964

src/rate_limiter.rs

@@ -0,0 +1,104 @@
+use std::{
+  collections::HashMap,
+  sync::Arc,
+  time::{Duration, Instant},
+};
+
+use tokio::sync::Mutex;
+
+use crate::error::Result;
+
+#[derive(Clone)]
+pub struct RateLimiter {
+  inner: Arc<Mutex<RateLimiterInner>>,
+}
+
+struct RateLimiterInner {
+  requests: HashMap<String, Vec<Instant>>,
+  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<RateLimitConfig>) -> 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 {
+      if 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());
+
+    Ok(())
+  }
+
+  pub async fn wait_for(&self, platform: &str) {
+    let _ = self.acquire(platform).await;
+  }
+}