pakker/src/rate_limiter.rs

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
      && 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;
  }
}