Skip to main content

klauthed_data/rate_limit/
mod.rs

1//! Rate limiting with a pluggable backend and algorithm.
2//!
3//! A [`RateLimiter`] records requests against a string key and reports whether
4//! each is [`Allowed`](RateLimitOutcome::Allowed) or
5//! [`Limited`](RateLimitOutcome::Limited), permitting up to `max` per `window`.
6//! All implementations share that `(max, window)` API, so they're
7//! interchangeable behind `Arc<dyn RateLimiter>`.
8//!
9//! **Fixed-window** (hard reset each window):
10//!
11//! * [`InMemoryRateLimiter`] — a clock-injected `Mutex<HashMap>`, per-process
12//!   (each replica counts independently). Ideal for single-node deployments and
13//!   tests (drive it with a `FixedClock`).
14//! * [`RedisRateLimiter`] (`redis` feature) — a shared counter in Redis, so a
15//!   fleet of replicas enforces one global budget per key.
16//!
17//! **Token-bucket** (continuous refill — smooths traffic, allows short bursts up
18//! to `max`): [`InMemoryTokenBucket`] and [`RedisTokenBucket`] (`redis`),
19//! with the same `(max, window)` parameters (`max` = burst capacity, refilled at
20//! `max / window`).
21//!
22//! ```
23//! use std::sync::Arc;
24//! use std::time::Duration;
25//! use klauthed_core::time::FixedClock;
26//! use klauthed_data::rate_limit::{InMemoryRateLimiter, RateLimiter, RateLimitOutcome};
27//!
28//! # #[tokio::main]
29//! # async fn main() -> Result<(), klauthed_data::DataError> {
30//! let clock = Arc::new(FixedClock::at_unix_millis(0));
31//! let limiter = InMemoryRateLimiter::new(clock.clone());
32//! let window = Duration::from_secs(60);
33//!
34//! // First request of two is allowed.
35//! assert!(matches!(limiter.check("ip:1.2.3.4", 2, window).await?, RateLimitOutcome::Allowed { .. }));
36//! limiter.check("ip:1.2.3.4", 2, window).await?; // second
37//! // Third exceeds the budget.
38//! assert!(matches!(limiter.check("ip:1.2.3.4", 2, window).await?, RateLimitOutcome::Limited { .. }));
39//! # Ok(())
40//! # }
41//! ```
42
43use std::time::Duration;
44
45use async_trait::async_trait;
46
47use crate::error::DataError;
48
49pub mod memory;
50#[cfg(feature = "redis")]
51pub mod redis;
52
53pub use memory::{InMemoryRateLimiter, InMemoryTokenBucket};
54#[cfg(feature = "redis")]
55pub use redis::{RedisRateLimiter, RedisTokenBucket};
56
57/// The result of recording one request against a key in its current window.
58#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59pub enum RateLimitOutcome {
60    /// The request is within budget. `remaining` is how many more are allowed
61    /// before the window resets.
62    Allowed {
63        /// Requests still permitted in the current window.
64        remaining: u32,
65    },
66    /// The budget is exhausted; the caller should retry after `retry_after`.
67    Limited {
68        /// Time until the current window resets.
69        retry_after: Duration,
70    },
71}
72
73impl RateLimitOutcome {
74    /// Whether the request was allowed.
75    #[must_use]
76    pub fn is_allowed(&self) -> bool {
77        matches!(self, RateLimitOutcome::Allowed { .. })
78    }
79}
80
81/// A fixed-window rate limiter keyed by an arbitrary string.
82///
83/// Implementations are `Send + Sync` so a limiter can be shared as
84/// `Arc<dyn RateLimiter>` across worker threads.
85#[async_trait]
86pub trait RateLimiter: Send + Sync {
87    /// Record one request for `key`, permitting up to `max` per `window`.
88    ///
89    /// `max` is clamped to at least 1. Returns the [`RateLimitOutcome`].
90    ///
91    /// # Errors
92    /// Returns [`DataError`] only on a backend failure (e.g. a Redis error); an
93    /// over-budget request is a successful [`Limited`](RateLimitOutcome::Limited)
94    /// outcome, not an error.
95    async fn check(
96        &self,
97        key: &str,
98        max: u32,
99        window: Duration,
100    ) -> Result<RateLimitOutcome, DataError>;
101}