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 (rate, burst) = {
      let inner = self.inner.lock().await;
      match platform.to_lowercase().as_str() {
        "modrinth" => {
          (
            inner.config.modrinth_requests_per_min,
            inner.config.modrinth_burst,
          )
        },
        "curseforge" => {
          (
            inner.config.curseforge_requests_per_min,
            inner.config.curseforge_burst,
          )
        },
        "github" => {
          (
            inner.config.github_requests_per_min,
            inner.config.github_burst,
          )
        },
        _ => {
          (
            inner.config.default_requests_per_min,
            inner.config.default_burst,
          )
        },
      }
    };

    let interval = Duration::from_secs(60) / rate.max(1);

    loop {
      let mut inner = self.inner.lock().await;
      let now = Instant::now();
      let platform_requests =
        inner.requests.entry(platform.to_owned()).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;
          continue;
        }
      }

      platform_requests.push(Instant::now());
      return Ok(());
    }
  }

  pub async fn wait_for(&self, platform: &str) {
    let _ = self.acquire(platform).await;
  }
}