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}