rok-rate-limit 0.3.0

Rate limiting Tower middleware and programmatic Limiter API for the rok ecosystem
Documentation
//! Fixed-window in-memory rate limiter.
//!
//! Simpler than the builder-based [`Limiter`](crate::Limiter) — good for
//! attribute macros like `#[throttle]`.

use std::time::{Duration, Instant};

use dashmap::DashMap;
use once_cell::sync::Lazy;

/// Global default fixed-window limiter (10 req/s — must be configured per-endpoint).
pub static FIXED_LIMITER: Lazy<FixedRateLimiter> = Lazy::new(FixedRateLimiter::new);

/// Fixed-window in-memory rate limiter backed by a [`DashMap`].
///
/// Each key gets its own bucket that tracks the count and window start.
/// Windows are fixed (aligned to wall clock for `reset_in` reporting), but
/// internally [`Instant`] is used for expiry checks.
#[derive(Debug, Clone)]
pub struct FixedRateLimiter {
    buckets: DashMap<String, RateBucket>,
}

#[derive(Debug, Clone)]
struct RateBucket {
    count: u64,
    window_start: Instant,
}

impl FixedRateLimiter {
    /// Create an empty rate limiter.
    pub fn new() -> Self {
        Self {
            buckets: DashMap::new(),
        }
    }

    /// Check if `key` is allowed.  Returns `true` if under limit.
    ///
    /// Consumes one unit from the bucket.
    pub fn check(&self, key: &str, max: u64, window: Duration) -> bool {
        let now = Instant::now();
        let mut bucket = self.buckets.entry(key.to_string()).or_insert(RateBucket {
            count: 0,
            window_start: now,
        });

        if now.duration_since(bucket.window_start) > window {
            bucket.count = 1;
            bucket.window_start = now;
            true
        } else if bucket.count < max {
            bucket.count += 1;
            true
        } else {
            false
        }
    }

    /// Returns the number of remaining allowed attempts for `key`.
    pub fn remaining(&self, key: &str, max: u64, window: Duration) -> u64 {
        let now = Instant::now();
        if let Some(bucket) = self.buckets.get(key) {
            if now.duration_since(bucket.window_start) > window {
                max
            } else {
                max.saturating_sub(bucket.count)
            }
        } else {
            max
        }
    }

    /// Returns the time until the current window resets, or `Duration::ZERO`
    /// if no window has been started for `key`.
    pub fn reset_in(&self, key: &str, window: Duration) -> Duration {
        let now = Instant::now();
        if let Some(bucket) = self.buckets.get(key) {
            let elapsed = now.duration_since(bucket.window_start);
            if elapsed > window {
                Duration::ZERO
            } else {
                window - elapsed
            }
        } else {
            Duration::ZERO
        }
    }
}

impl Default for FixedRateLimiter {
    fn default() -> Self {
        Self::new()
    }
}

/// Parse a human-readable duration string into a [`Duration`].
///
/// Supported units: `"second"`, `"minute"`, `"hour"`, `"day"`.
/// Plural forms (`"seconds"`, `"minutes"`, `"hours"`, `"days"`) also work.
/// Returns 1 second for unknown units.
pub fn parse_duration(unit: &str) -> Duration {
    match unit.trim().to_lowercase().as_str() {
        "second" | "seconds" => Duration::from_secs(1),
        "minute" | "minutes" => Duration::from_secs(60),
        "hour" | "hours" => Duration::from_secs(3600),
        "day" | "days" => Duration::from_secs(86400),
        _ => Duration::from_secs(1),
    }
}