foukoapi 0.1.2-alpha.2

//! In-memory rate limiting for incoming updates.
//!
//! A [`RateLimiter`] answers one question: "has this user sent too many
//! requests too fast?". It keeps a short sliding window of hit timestamps
//! per key and nothing else - no storage, no background task, no locks held
//! across `.await`. That makes it cheap enough to check on every update and
//! safe to share across adapter tasks.
//!
//! The limiter is deliberately unopinionated. The [`Bot`](crate::Bot)
//! wires a default one in, but you can build your own policy with
//! [`RateLimit`] and hand it over, or turn limiting off entirely. Nothing
//! here is hidden behind private types - a downstream bot can read the
//! same building blocks and layer stricter rules of its own on top.

use std::{
    collections::HashMap,
    sync::Mutex,
    time::{Duration, Instant},
};

/// A rate-limit policy: at most `max_hits` events per `window`.
#[derive(Debug, Clone, Copy)]
pub struct RateLimit {
    /// How many events are allowed inside one window.
    pub max_hits: u32,
    /// The length of the sliding window.
    pub window: Duration,
}

impl RateLimit {
    /// `max_hits` events per `window`.
    pub fn new(max_hits: u32, window: Duration) -> Self {
        Self { max_hits, window }
    }

    /// A sensible default for chat bots: 5 actions every 3 seconds. Enough
    /// for normal use, tight enough that holding a key down or scripting a
    /// flood gets throttled fast.
    pub fn relaxed() -> Self {
        Self::new(5, Duration::from_secs(3))
    }

    /// A stricter policy for expensive actions (network calls, image
    /// generation): 1 action per 3 seconds.
    pub fn strict() -> Self {
        Self::new(1, Duration::from_secs(3))
    }
}

impl Default for RateLimit {
    fn default() -> Self {
        Self::relaxed()
    }
}

/// The verdict for a single check.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Decision {
    /// Under the limit; the caller may proceed.
    Allow,
    /// Over the limit. `retry_after` is a hint for how long until a slot
    /// frees up.
    Deny {
        /// Roughly how long until the oldest hit leaves the window.
        retry_after: Duration,
        /// `true` on the first denial in the current window, so a caller
        /// can notify the user once instead of on every dropped update.
        first: bool,
    },
}

impl Decision {
    /// `true` when the action is allowed.
    pub fn allowed(&self) -> bool {
        matches!(self, Decision::Allow)
    }
}

/// A sliding-window rate limiter keyed by arbitrary strings.
///
/// Clone it freely; clones share the same underlying state through an
/// `Arc`-free `Mutex` wrapped by the caller if needed. In practice the
/// [`Bot`](crate::Bot) stores it behind an `Arc`.
#[derive(Debug)]
pub struct RateLimiter {
    policy: RateLimit,
    /// key -> recent hit timestamps plus deny bookkeeping.
    hits: Mutex<HashMap<String, KeyState>>,
    /// Counts calls so we can run a full sweep every so often and keep the
    /// map from accumulating keys of users who never come back.
    calls: std::sync::atomic::AtomicU64,
}

/// Per-key window state.
#[derive(Debug, Default)]
struct KeyState {
    /// Recent hit timestamps, oldest first.
    hits: Vec<Instant>,
    /// Set after the first deny in the current window; cleared once a
    /// hit is allowed again.
    denied: bool,
}

/// Run a full sweep every this-many checks.
const SWEEP_EVERY: u64 = 1024;

/// Sweep immediately when the map outgrows this many keys, so a flood of
/// unique keys can't balloon memory between periodic sweeps.
const SWEEP_MAP_SIZE: usize = 4096;

impl RateLimiter {
    /// A limiter enforcing `policy`.
    pub fn new(policy: RateLimit) -> Self {
        Self {
            policy,
            hits: Mutex::new(HashMap::new()),
            calls: std::sync::atomic::AtomicU64::new(0),
        }
    }

    /// The policy this limiter enforces.
    pub fn policy(&self) -> RateLimit {
        self.policy
    }

    /// Record a hit for `key` and decide whether it's allowed. Prunes
    /// timestamps that have aged out of the window as it goes, so memory
    /// stays proportional to *active* keys, not lifetime keys. Every
    /// `SWEEP_EVERY` calls the whole map is swept too, and a map that
    /// grew past `SWEEP_MAP_SIZE` keys is swept on the spot, so keys
    /// from one-off users don't pile up over months of uptime.
    pub fn check(&self, key: &str) -> Decision {
        let n = self
            .calls
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        if n % SWEEP_EVERY == SWEEP_EVERY - 1 {
            self.sweep();
        }

        let now = Instant::now();
        let window = self.policy.window;
        let mut map = match self.hits.lock() {
            Ok(m) => m,
            // A poisoned lock shouldn't take the bot down; fail open so a
            // panic elsewhere can't wedge every command behind a deny.
            Err(poisoned) => poisoned.into_inner(),
        };
        // Emergency sweep: too many keys means someone is spraying unique
        // ids; prune stale entries now instead of waiting for the counter.
        if map.len() > SWEEP_MAP_SIZE {
            map.retain(|_, state| {
                state.hits.retain(|t| now.duration_since(*t) < window);
                !state.hits.is_empty()
            });
        }
        let entry = map.entry(key.to_owned()).or_default();
        entry.hits.retain(|t| now.duration_since(*t) < window);

        if entry.hits.len() as u32 >= self.policy.max_hits {
            let first = !entry.denied;
            entry.denied = true;
            let oldest = entry.hits.first().copied().unwrap_or(now);
            let retry_after = window.saturating_sub(now.duration_since(oldest));
            return Decision::Deny { retry_after, first };
        }
        entry.denied = false;
        entry.hits.push(now);
        Decision::Allow
    }

    /// Drop stale keys entirely. Optional housekeeping for very long-lived
    /// processes; `check` already prunes per key, so this is only worth
    /// calling occasionally if you have a huge churn of one-off keys.
    pub fn sweep(&self) {
        let now = Instant::now();
        let window = self.policy.window;
        if let Ok(mut map) = self.hits.lock() {
            map.retain(|_, state| {
                state.hits.retain(|t| now.duration_since(*t) < window);
                !state.hits.is_empty()
            });
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn allows_up_to_the_limit_then_denies() {
        let rl = RateLimiter::new(RateLimit::new(3, Duration::from_secs(10)));
        assert!(rl.check("u").allowed());
        assert!(rl.check("u").allowed());
        assert!(rl.check("u").allowed());
        assert!(!rl.check("u").allowed()); // 4th in the window
    }

    #[test]
    fn keys_are_independent() {
        let rl = RateLimiter::new(RateLimit::new(1, Duration::from_secs(10)));
        assert!(rl.check("a").allowed());
        assert!(rl.check("b").allowed());
        assert!(!rl.check("a").allowed());
    }

    #[test]
    fn window_frees_up_over_time() {
        let rl = RateLimiter::new(RateLimit::new(1, Duration::from_millis(30)));
        assert!(rl.check("u").allowed());
        assert!(!rl.check("u").allowed());
        std::thread::sleep(Duration::from_millis(40));
        assert!(rl.check("u").allowed());
    }

    #[test]
    fn deny_reports_retry_after() {
        let rl = RateLimiter::new(RateLimit::new(1, Duration::from_secs(10)));
        rl.check("u");
        match rl.check("u") {
            Decision::Deny { retry_after, .. } => assert!(retry_after <= Duration::from_secs(10)),
            Decision::Allow => panic!("expected a deny"),
        }
    }

    #[test]
    fn only_first_deny_in_a_window_is_flagged() {
        let rl = RateLimiter::new(RateLimit::new(1, Duration::from_millis(30)));
        assert!(rl.check("u").allowed());
        match rl.check("u") {
            Decision::Deny { first, .. } => assert!(first),
            Decision::Allow => panic!("expected a deny"),
        }
        match rl.check("u") {
            Decision::Deny { first, .. } => assert!(!first),
            Decision::Allow => panic!("expected a deny"),
        }
        // A fresh window resets the flag.
        std::thread::sleep(Duration::from_millis(40));
        assert!(rl.check("u").allowed());
        match rl.check("u") {
            Decision::Deny { first, .. } => assert!(first),
            Decision::Allow => panic!("expected a deny"),
        }
    }

    #[test]
    fn oversized_map_is_swept() {
        let rl = RateLimiter::new(RateLimit::new(1, Duration::from_millis(1)));
        // Stuff the map with expired keys directly, so the periodic
        // call-count sweep can't kick in first.
        {
            let now = Instant::now();
            let stale = now.checked_sub(Duration::from_secs(60)).unwrap_or(now);
            let mut map = rl.hits.lock().unwrap();
            for i in 0..(SWEEP_MAP_SIZE + 10) {
                map.insert(
                    format!("k{i}"),
                    KeyState {
                        hits: vec![stale],
                        denied: false,
                    },
                );
            }
        }
        // The next check sees the oversized map and prunes it in place.
        rl.check("fresh");
        let len = rl.hits.lock().unwrap().len();
        assert!(len <= 2, "map should have been swept, len={len}");
    }
}