Skip to main content

volga_rate_limiter/
rate_limiter.rs

1//! Generic rate limiter abstractions and shared utilities.
2//!
3//! This module defines the core traits and building blocks used by
4//! all rate limiting algorithms provided by this crate.
5//!
6//! The primary abstraction is [`RateLimiter`], which represents a
7//! stateful, thread-safe rate limiting algorithm operating on a
8//! partition key.
9//!
10//! ## Design principles
11//!
12//! - **Algorithm-agnostic interface** - higher-level frameworks can
13//!   work with any rate limiting strategy through a common API.
14//! - **Partition-based limiting** - each limiter operates on a `u64`
15//!   partition key representing a logical client or request group.
16//! - **Time abstraction** - all time-dependent logic is driven by a
17//!   pluggable [`TimeSource`] to allow deterministic testing.
18//!
19//! ## Thread safety
20//!
21//! All implementations of [`RateLimiter`] provided by this crate are:
22//!
23//! - Safe to use concurrently
24//! - Designed for high-contention scenarios
25//! - Intended to be shared between threads and async tasks
26//!
27//! ## Scope
28//!
29//! This module does **not** define how partition keys are created or
30//! how rate limiting is applied to HTTP requests.
31//! Those concerns are intentionally left to higher-level layers.
32
33use std::time::Instant;
34
35pub use fixed_window::{FixedWindowRateLimiter, InMemoryFixedWindowStore};
36pub use gcra::{GcraRateLimiter, InMemoryGcraStore};
37pub use sliding_window::{InMemorySlidingWindowStore, SlidingWindowRateLimiter};
38pub use token_bucket::{InMemoryTokenBucketStore, TokenBucketRateLimiter};
39
40mod fixed_window;
41mod gcra;
42mod sliding_window;
43mod store;
44mod token_bucket;
45
46pub use store::{
47    FixedWindowParams, FixedWindowStore, GcraParams, GcraStore, SlidingWindowParams,
48    SlidingWindowStore, TokenBucketParams, TokenBucketStore,
49};
50
51const MICROS_PER_SEC: u64 = 1_000_000;
52
53/// A generic rate limiter interface.
54///
55/// A rate limiter tracks request counts per **partition key** and
56/// determines whether new requests are allowed.
57///
58/// Implementations must:
59///
60/// - Be thread-safe
61/// - Handle concurrent access correctly
62/// - Execute the `check` operation efficiently, as it is typically
63///   called on every incoming request
64///
65/// The meaning of the partition key is defined by the caller
66/// (for example: IP address, user ID, tenant ID, or API key).
67pub trait RateLimiter {
68    /// Checks whether a request is allowed for the given partition key.
69    ///
70    /// # Parameters
71    ///
72    /// - `key`: A stable `u64` value identifying a logical client or
73    ///   request group.
74    ///
75    /// # Returns
76    ///
77    /// - `true` if the request is allowed and should proceed
78    /// - `false` if the rate limit has been exceeded
79    ///
80    /// # Notes
81    ///
82    /// - This method may mutate internal state.
83    /// - It is expected to be called on the hot path and should be fast.
84    fn check(&self, key: u64) -> bool;
85}
86
87/// A source of time used by rate-limiting algorithms.
88///
89/// This abstraction allows rate limiters to be decoupled from
90/// the system clock, enabling deterministic and fast unit tests.
91///
92/// Time is expressed in **microseconds** and must be **monotonic**
93/// (non-decreasing).
94pub trait TimeSource: Send + Sync {
95    /// Returns a monotonic timestamp in microseconds.
96    fn now_micros(&self) -> u64;
97
98    /// Returns the number of seconds elapsed since [`UNIX_EPOCH`](std::time::UNIX_EPOCH)
99    /// (`1970-01-01 00:00:00 UTC`).
100    ///
101    /// Implementations must ensure that the returned value is:
102    ///
103    /// - Monotonic (non-decreasing)
104    /// - Cheap to compute
105    #[inline(always)]
106    fn now_secs(&self) -> u64 {
107        self.now_micros() / MICROS_PER_SEC
108    }
109}
110
111/// Monotonic system time source backed by `Instant`.
112///
113/// Uses an internal start anchor and returns elapsed microseconds since that anchor.
114/// This avoids wall-clock jumps (NTP, manual adjustments, etc.).
115#[derive(Debug, Default, Clone, Copy)]
116pub struct SystemTimeSource;
117
118impl SystemTimeSource {
119    #[inline]
120    fn anchor() -> Instant {
121        // `Instant::now()` is cheap and monotonic.
122        // We want a stable anchor shared across calls.
123        // Using `OnceLock` gives us a process-wide start point.
124        static START: std::sync::OnceLock<Instant> = std::sync::OnceLock::new();
125        *START.get_or_init(Instant::now)
126    }
127}
128
129impl TimeSource for SystemTimeSource {
130    #[inline]
131    fn now_micros(&self) -> u64 {
132        let elapsed = Self::anchor().elapsed();
133        // Saturating conversion to be extra defensive (though practically safe).
134        elapsed.as_micros().try_into().unwrap_or(u64::MAX)
135    }
136}
137
138#[cfg(test)]
139pub(super) mod test_utils {
140    use super::{MICROS_PER_SEC, TimeSource};
141    use std::sync::{Arc, Mutex};
142
143    #[derive(Clone)]
144    pub(super) struct MockTimeSource {
145        current_time: Arc<Mutex<u64>>,
146    }
147
148    impl MockTimeSource {
149        pub(super) fn new(initial_time: u64) -> Self {
150            Self {
151                current_time: Arc::new(Mutex::new(initial_time * MICROS_PER_SEC)),
152            }
153        }
154
155        pub(super) fn advance(&self, seconds: u64) {
156            let mut time = self.current_time.lock().unwrap();
157            *time += seconds * MICROS_PER_SEC;
158        }
159    }
160
161    impl TimeSource for MockTimeSource {
162        fn now_micros(&self) -> u64 {
163            *self.current_time.lock().unwrap()
164        }
165    }
166}