nautilus_network/ratelimiter/
mod.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! A rate limiter implementation heavily inspired by [governor](https://github.com/antifuchs/governor).
17//!
18//! The governor does not support different quota for different key. It is an open [issue](https://github.com/antifuchs/governor/issues/193).
19pub mod clock;
20mod gcra;
21mod nanos;
22pub mod quota;
23
24use std::{
25    fmt::Debug,
26    hash::Hash,
27    num::NonZeroU64,
28    sync::atomic::{AtomicU64, Ordering},
29    time::Duration,
30};
31
32use dashmap::DashMap;
33use futures_util::StreamExt;
34
35use self::{
36    clock::{Clock, FakeRelativeClock, MonotonicClock},
37    gcra::{Gcra, NotUntil},
38    nanos::Nanos,
39    quota::Quota,
40};
41
42/// An in-memory representation of a GCRA's rate-limiting state.
43///
44/// Implemented using [`AtomicU64`] operations, this state representation can be used to
45/// construct rate limiting states for other in-memory states: e.g., this crate uses
46/// `InMemoryState` as the states it tracks in the keyed rate limiters it implements.
47///
48/// Internally, the number tracked here is the theoretical arrival time (a GCRA term) in number of
49/// nanoseconds since the rate limiter was created.
50#[derive(Debug, Default)]
51pub struct InMemoryState(AtomicU64);
52
53impl InMemoryState {
54    /// Measures and updates the GCRA's state atomically, retrying on concurrent modifications.
55    ///
56    /// # Errors
57    ///
58    /// Returns an error if the provided closure returns an error.
59    pub(crate) fn measure_and_replace_one<T, F, E>(&self, mut f: F) -> Result<T, E>
60    where
61        F: FnMut(Option<Nanos>) -> Result<(T, Nanos), E>,
62    {
63        let mut prev = self.0.load(Ordering::Acquire);
64        let mut decision = f(NonZeroU64::new(prev).map(|n| n.get().into()));
65        while let Ok((result, new_data)) = decision {
66            // Lock-free CAS loop: retry with current value if another thread modified it,
67            // uses weak variant (faster) since spurious failures are fine in a retry loop.
68            match self.0.compare_exchange_weak(
69                prev,
70                new_data.into(),
71                Ordering::Release,
72                Ordering::Relaxed,
73            ) {
74                Ok(_) => return Ok(result),
75                Err(e) => prev = e, // Retry with value written by another thread
76            }
77            decision = f(NonZeroU64::new(prev).map(|n| n.get().into()));
78        }
79        // This map shouldn't be needed, as we only get here in the error case, but the compiler
80        // can't see it.
81        decision.map(|(result, _)| result)
82    }
83}
84
85/// A concurrent, thread-safe and fairly performant hashmap based on [`DashMap`].
86pub type DashMapStateStore<K> = DashMap<K, InMemoryState>;
87
88/// A way for rate limiters to keep state.
89///
90/// There are two important kinds of state stores: Direct and keyed. The direct kind have only
91/// one state, and are useful for "global" rate limit enforcement (e.g. a process should never
92/// do more than N tasks a day). The keyed kind allows one rate limit per key (e.g. an API
93/// call budget per client API key).
94///
95/// A direct state store is expressed as [`StateStore::Key`] = `NotKeyed`.
96/// Keyed state stores have a
97/// type parameter for the key and set their key to that.
98pub trait StateStore {
99    /// The type of key that the state store can represent.
100    type Key;
101
102    /// Updates a state store's rate limiting state for a given key, using the given closure.
103    ///
104    /// The closure parameter takes the old value (`None` if this is the first measurement) of the
105    /// state store at the key's location, checks if the request an be accommodated and:
106    ///
107    /// - If the request is rate-limited, returns `Err(E)`.
108    /// - If the request can make it through, returns `Ok(T)` (an arbitrary positive return
109    ///   value) and the updated state.
110    ///
111    /// It is `measure_and_replace`'s job then to safely replace the value at the key - it must
112    /// only update the value if the value hasn't changed. The implementations in this
113    /// crate use `AtomicU64` operations for this.
114    ///
115    /// # Errors
116    ///
117    /// Returns `Err(E)` if the closure returns an error or the request is rate-limited.
118    fn measure_and_replace<T, F, E>(&self, key: &Self::Key, f: F) -> Result<T, E>
119    where
120        F: Fn(Option<Nanos>) -> Result<(T, Nanos), E>;
121}
122
123impl<K: Hash + Eq + Clone> StateStore for DashMapStateStore<K> {
124    type Key = K;
125
126    fn measure_and_replace<T, F, E>(&self, key: &Self::Key, f: F) -> Result<T, E>
127    where
128        F: Fn(Option<Nanos>) -> Result<(T, Nanos), E>,
129    {
130        if let Some(v) = self.get(key) {
131            // fast path: measure existing entry
132            return v.measure_and_replace_one(f);
133        }
134        // make an entry and measure that:
135        let entry = self.entry(key.clone()).or_default();
136        (*entry).measure_and_replace_one(f)
137    }
138}
139
140/// A rate limiter that enforces different quotas per key using the GCRA algorithm.
141///
142/// This implementation allows setting different rate limits for different keys,
143/// with an optional default quota for keys that don't have specific quotas.
144pub struct RateLimiter<K, C>
145where
146    C: Clock,
147{
148    default_gcra: Option<Gcra>,
149    state: DashMapStateStore<K>,
150    gcra: DashMap<K, Gcra>,
151    clock: C,
152    start: C::Instant,
153}
154
155impl<K, C> Debug for RateLimiter<K, C>
156where
157    K: Debug,
158    C: Clock,
159{
160    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
161        f.debug_struct(stringify!(RateLimiter)).finish()
162    }
163}
164
165impl<K> RateLimiter<K, MonotonicClock>
166where
167    K: Eq + Hash,
168{
169    /// Creates a new rate limiter with a base quota and keyed quotas.
170    ///
171    /// The base quota applies to all keys that don't have specific quotas.
172    /// Keyed quotas override the base quota for specific keys.
173    #[must_use]
174    pub fn new_with_quota(base_quota: Option<Quota>, keyed_quotas: Vec<(K, Quota)>) -> Self {
175        let clock = MonotonicClock {};
176        Self::new_with_clock(base_quota, keyed_quotas, clock)
177    }
178}
179
180impl<K, C> RateLimiter<K, C>
181where
182    K: Eq + Hash,
183    C: Clock,
184{
185    /// Creates a new rate limiter with an explicit clock.
186    ///
187    /// The base quota applies to all keys that do not have specific quotas.
188    /// Keyed quotas override the base quota for specific keys.
189    #[must_use]
190    pub fn new_with_clock(
191        base_quota: Option<Quota>,
192        keyed_quotas: Vec<(K, Quota)>,
193        clock: C,
194    ) -> Self {
195        let start = clock.now();
196        let gcra: DashMap<_, _> = keyed_quotas
197            .into_iter()
198            .map(|(k, q)| (k, Gcra::new(q)))
199            .collect();
200        Self {
201            default_gcra: base_quota.map(Gcra::new),
202            state: DashMapStateStore::new(),
203            gcra,
204            clock,
205            start,
206        }
207    }
208}
209
210impl<K> RateLimiter<K, FakeRelativeClock>
211where
212    K: Hash + Eq + Clone,
213{
214    /// Advances the fake clock by the specified duration.
215    ///
216    /// This is only available for testing with `FakeRelativeClock`.
217    pub fn advance_clock(&self, by: Duration) {
218        self.clock.advance(by);
219    }
220}
221
222impl<K, C> RateLimiter<K, C>
223where
224    K: Hash + Eq + Clone,
225    C: Clock,
226{
227    /// Adds or updates a quota for a specific key.
228    pub fn add_quota_for_key(&self, key: K, value: Quota) {
229        self.gcra.insert(key, Gcra::new(value));
230    }
231
232    /// Checks if the given key is allowed under the rate limit.
233    ///
234    /// # Errors
235    ///
236    /// Returns `Err(NotUntil)` if the key is rate-limited, indicating when it will be allowed.
237    pub fn check_key(&self, key: &K) -> Result<(), NotUntil<C::Instant>> {
238        match self.gcra.get(key) {
239            Some(quota) => quota.test_and_update(self.start, key, &self.state, self.clock.now()),
240            None => self.default_gcra.as_ref().map_or(Ok(()), |gcra| {
241                gcra.test_and_update(self.start, key, &self.state, self.clock.now())
242            }),
243        }
244    }
245
246    /// Waits until the specified key is ready (not rate-limited).
247    pub async fn until_key_ready(&self, key: &K) {
248        loop {
249            match self.check_key(key) {
250                Ok(()) => {
251                    break;
252                }
253                Err(e) => {
254                    self.clock.sleep(e.wait_time_from(self.clock.now())).await;
255                }
256            }
257        }
258    }
259
260    /// Waits until all specified keys are ready (not rate-limited).
261    ///
262    /// If no keys are provided, this function returns immediately.
263    /// Uses fast paths for 0-2 keys to avoid stream scheduling overhead.
264    pub async fn await_keys_ready(&self, keys: Option<&[K]>) {
265        let Some(keys) = keys else {
266            return;
267        };
268
269        match keys.len() {
270            0 => {}
271            1 => self.until_key_ready(&keys[0]).await,
272            2 => {
273                tokio::join!(
274                    self.until_key_ready(&keys[0]),
275                    self.until_key_ready(&keys[1]),
276                );
277            }
278            _ => {
279                let tasks = keys.iter().map(|key| self.until_key_ready(key));
280                futures::stream::iter(tasks)
281                    .for_each_concurrent(None, |key_future| async move {
282                        key_future.await;
283                    })
284                    .await;
285            }
286        }
287    }
288}
289
290#[cfg(test)]
291mod tests {
292    use std::{
293        num::NonZeroU32,
294        sync::atomic::{AtomicU32, Ordering},
295        time::Duration,
296    };
297
298    use dashmap::DashMap;
299    use rstest::rstest;
300
301    use super::{
302        DashMapStateStore, RateLimiter,
303        clock::{Clock, FakeRelativeClock},
304        gcra::{Gcra, StateSnapshot},
305        nanos::Nanos,
306        quota::Quota,
307    };
308
309    fn initialize_mock_rate_limiter() -> RateLimiter<String, FakeRelativeClock> {
310        let clock = FakeRelativeClock::default();
311        let start = clock.now();
312        let gcra = DashMap::new();
313        let base_quota = Quota::per_second(NonZeroU32::new(2).unwrap()).unwrap();
314        RateLimiter {
315            default_gcra: Some(Gcra::new(base_quota)),
316            state: DashMapStateStore::new(),
317            gcra,
318            clock,
319            start,
320        }
321    }
322
323    #[rstest]
324    fn test_enormous_quota_denies_after_burst() {
325        // Regression: a period beyond ~584 years panicked in Gcra::new; with
326        // clamping it must admit the burst and then deny, not admit everything
327        let quota = Quota::with_period(Duration::MAX)
328            .unwrap()
329            .allow_burst(NonZeroU32::new(u32::MAX).unwrap());
330        let clock = FakeRelativeClock::default();
331        let limiter: RateLimiter<String, FakeRelativeClock> =
332            RateLimiter::new_with_clock(Some(quota), vec![], clock);
333
334        let key = "key".to_string();
335        assert!(limiter.check_key(&key).is_ok());
336        assert!(limiter.check_key(&key).is_err());
337    }
338
339    #[rstest]
340    fn test_default_quota() {
341        let mock_limiter = initialize_mock_rate_limiter();
342
343        // Check base quota is not exceeded
344        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
345        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
346
347        // Check base quota is exceeded
348        assert!(mock_limiter.check_key(&"default".to_string()).is_err());
349
350        // Increment clock and check base quota is reset
351        mock_limiter.advance_clock(Duration::from_secs(1));
352        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
353    }
354
355    #[rstest]
356    fn test_custom_key_quota() {
357        let mock_limiter = initialize_mock_rate_limiter();
358
359        // Add new key quota pair
360        mock_limiter.add_quota_for_key(
361            "custom".to_string(),
362            Quota::per_second(NonZeroU32::new(1).unwrap()).unwrap(),
363        );
364
365        // Check custom quota
366        assert!(mock_limiter.check_key(&"custom".to_string()).is_ok());
367        assert!(mock_limiter.check_key(&"custom".to_string()).is_err());
368
369        // Check that default quota still applies to other keys
370        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
371        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
372        assert!(mock_limiter.check_key(&"default".to_string()).is_err());
373    }
374
375    #[rstest]
376    fn test_multiple_keys() {
377        let mock_limiter = initialize_mock_rate_limiter();
378
379        mock_limiter.add_quota_for_key(
380            "key1".to_string(),
381            Quota::per_second(NonZeroU32::new(1).unwrap()).unwrap(),
382        );
383        mock_limiter.add_quota_for_key(
384            "key2".to_string(),
385            Quota::per_second(NonZeroU32::new(3).unwrap()).unwrap(),
386        );
387
388        // Test key1
389        assert!(mock_limiter.check_key(&"key1".to_string()).is_ok());
390        assert!(mock_limiter.check_key(&"key1".to_string()).is_err());
391
392        // Test key2
393        assert!(mock_limiter.check_key(&"key2".to_string()).is_ok());
394        assert!(mock_limiter.check_key(&"key2".to_string()).is_ok());
395        assert!(mock_limiter.check_key(&"key2".to_string()).is_ok());
396        assert!(mock_limiter.check_key(&"key2".to_string()).is_err());
397    }
398
399    #[rstest]
400    fn test_quota_reset() {
401        let mock_limiter = initialize_mock_rate_limiter();
402
403        // Exhaust quota
404        assert!(mock_limiter.check_key(&"reset".to_string()).is_ok());
405        assert!(mock_limiter.check_key(&"reset".to_string()).is_ok());
406        assert!(mock_limiter.check_key(&"reset".to_string()).is_err());
407
408        // Advance clock by less than a second
409        mock_limiter.advance_clock(Duration::from_millis(499));
410        assert!(mock_limiter.check_key(&"reset".to_string()).is_err());
411
412        // Advance clock to reset
413        mock_limiter.advance_clock(Duration::from_millis(501));
414        assert!(mock_limiter.check_key(&"reset".to_string()).is_ok());
415    }
416
417    #[rstest]
418    fn test_different_quotas() {
419        let mock_limiter = initialize_mock_rate_limiter();
420
421        mock_limiter.add_quota_for_key(
422            "per_second".to_string(),
423            Quota::per_second(NonZeroU32::new(2).unwrap()).unwrap(),
424        );
425        mock_limiter.add_quota_for_key(
426            "per_minute".to_string(),
427            Quota::per_minute(NonZeroU32::new(3).unwrap()),
428        );
429
430        // Test per_second quota
431        assert!(mock_limiter.check_key(&"per_second".to_string()).is_ok());
432        assert!(mock_limiter.check_key(&"per_second".to_string()).is_ok());
433        assert!(mock_limiter.check_key(&"per_second".to_string()).is_err());
434
435        // Test per_minute quota
436        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_ok());
437        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_ok());
438        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_ok());
439        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_err());
440
441        // Advance clock and check reset
442        mock_limiter.advance_clock(Duration::from_secs(1));
443        assert!(mock_limiter.check_key(&"per_second".to_string()).is_ok());
444        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_err());
445    }
446
447    #[tokio::test]
448    async fn test_await_keys_ready() {
449        let mock_limiter = initialize_mock_rate_limiter();
450
451        // Check base quota is not exceeded
452        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
453        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
454
455        // Check base quota is exceeded
456        assert!(mock_limiter.check_key(&"default".to_string()).is_err());
457
458        // Wait keys to be ready and check base quota is reset
459        mock_limiter.advance_clock(Duration::from_secs(1));
460        let keys = ["default".to_string()];
461        mock_limiter.await_keys_ready(Some(keys.as_slice())).await;
462        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
463    }
464
465    #[rstest]
466    fn test_remaining_burst_capacity_zero_t() {
467        let snapshot = StateSnapshot::new(
468            Nanos::from(0u64),
469            Nanos::from(1_000_000u64),
470            Nanos::from(0u64),
471            Nanos::from(0u64),
472        );
473        assert_eq!(snapshot.remaining_burst_capacity(), 0);
474    }
475
476    #[rstest]
477    fn test_per_second_returns_none_on_zero_replenish_interval() {
478        assert!(Quota::per_second(NonZeroU32::new(u32::MAX).unwrap()).is_none());
479    }
480
481    #[rstest]
482    fn test_per_minute_accepts_max_burst() {
483        let quota = Quota::per_minute(NonZeroU32::new(u32::MAX).unwrap());
484        assert!(quota.replenish_interval().as_nanos() > 0);
485    }
486
487    #[rstest]
488    fn test_per_hour_accepts_max_burst() {
489        let quota = Quota::per_hour(NonZeroU32::new(u32::MAX).unwrap());
490        assert!(quota.replenish_interval().as_nanos() > 0);
491    }
492
493    mod property_tests {
494        use proptest::prelude::*;
495        use rstest::rstest;
496
497        use crate::ratelimiter::{gcra::StateSnapshot, nanos::Nanos};
498
499        proptest! {
500            #![proptest_config(ProptestConfig {
501                failure_persistence: Some(Box::new(
502                    proptest::test_runner::FileFailurePersistence::WithSource("ratelimiter")
503                )),
504                ..ProptestConfig::default()
505            })]
506
507            // Full u64 domain: the historical overflow lived above the narrowed one-hour range
508            #[rstest]
509            fn remaining_burst_capacity_never_panics(
510                t in proptest::num::u64::ANY,
511                tau in proptest::num::u64::ANY,
512                time_of_measurement in proptest::num::u64::ANY,
513                tat in proptest::num::u64::ANY,
514            ) {
515                let snapshot = StateSnapshot::new(
516                    Nanos::from(t),
517                    Nanos::from(tau),
518                    Nanos::from(time_of_measurement),
519                    Nanos::from(tat),
520                );
521
522                let _ = snapshot.remaining_burst_capacity();
523            }
524
525            // Operators must saturate across the full u64 domain (a wrapped TAT admits everything)
526            #[rstest]
527            fn nanos_operators_never_panic(a in proptest::num::u64::ANY, b in proptest::num::u64::ANY) {
528                let na = Nanos::from(a);
529                let nb = Nanos::from(b);
530
531                prop_assert_eq!((na + nb).as_u64(), a.saturating_add(b));
532                prop_assert_eq!((na * b).as_u64(), a.saturating_mul(b));
533                prop_assert_eq!(na.saturating_sub(nb).as_u64(), a.saturating_sub(b));
534            }
535        }
536    }
537
538    #[rstest]
539    fn test_gcra_boundary_exact_replenishment() {
540        // Test GCRA boundary condition where t0 equals earliest_time exactly.
541        // This exercises the saturating_sub edge case deterministically without sleeps.
542        let mock_limiter = initialize_mock_rate_limiter();
543        let key = "boundary_test".to_string();
544
545        assert!(mock_limiter.check_key(&key).is_ok());
546        assert!(mock_limiter.check_key(&key).is_ok());
547        assert!(mock_limiter.check_key(&key).is_err());
548
549        // Advance clock by exactly one replenish interval (500ms for 2 req/sec)
550        let quota = Quota::per_second(NonZeroU32::new(2).unwrap()).unwrap();
551        let replenish_interval = quota.replenish_interval();
552        mock_limiter.advance_clock(replenish_interval);
553
554        assert!(
555            mock_limiter.check_key(&key).is_ok(),
556            "Request at exact replenish boundary should be allowed"
557        );
558        assert!(
559            mock_limiter.check_key(&key).is_err(),
560            "Immediate follow-up should be rate-limited"
561        );
562    }
563
564    #[rstest]
565    fn test_per_second_boundary_exact_limit() {
566        // 1_000_000_000ns / 1_000_000_000 = 1ns per replenish, the exact boundary
567        let quota = Quota::per_second(NonZeroU32::new(1_000_000_000).unwrap()).unwrap();
568        assert_eq!(quota.replenish_interval().as_nanos(), 1);
569    }
570
571    #[rstest]
572    fn test_per_second_returns_none_above_one_billion() {
573        // 1_000_000_000ns / 1_000_000_001 rounds to 0ns
574        assert!(Quota::per_second(NonZeroU32::new(1_000_000_001).unwrap()).is_none());
575    }
576
577    #[rstest]
578    fn test_burst_size_replenished_in_truncation() {
579        // 100_000_000_000ns * u32::MAX overflows u64, `as u64` silently truncates
580        let quota = Quota::with_period(Duration::from_secs(100))
581            .unwrap()
582            .allow_burst(NonZeroU32::new(u32::MAX).unwrap());
583
584        let replenished_in = quota.burst_size_replenished_in();
585        let full: u128 = 100_000_000_000u128 * u128::from(u32::MAX);
586        let truncated = full as u64;
587
588        assert_eq!(replenished_in, Duration::from_nanos(truncated));
589        assert_ne!(
590            full,
591            u128::from(truncated),
592            "Truncation should have occurred"
593        );
594    }
595
596    #[rstest]
597    #[should_panic(expected = "t cannot be zero")]
598    fn test_from_gcra_parameters_panics_on_zero_t() {
599        let _ = Quota::from_gcra_parameters(Nanos::from(0u64), Nanos::from(100u64));
600    }
601
602    #[rstest]
603    #[should_panic(expected = "tau/t results in zero burst capacity")]
604    fn test_from_gcra_parameters_panics_on_zero_division() {
605        // tau=1, t=2 → integer division yields 0
606        let _ = Quota::from_gcra_parameters(Nanos::from(2u64), Nanos::from(1u64));
607    }
608
609    #[rstest]
610    #[should_panic(expected = "tau/t exceeds u32::MAX")]
611    fn test_from_gcra_parameters_panics_on_overflow() {
612        let _ = Quota::from_gcra_parameters(Nanos::from(1u64), Nanos::from(u64::MAX));
613    }
614
615    #[rstest]
616    fn test_concurrent_check_key_respects_burst() {
617        let rate = 10u32;
618        let clock = FakeRelativeClock::default();
619        let start = clock.now();
620        let limiter = RateLimiter {
621            default_gcra: Some(Gcra::new(
622                Quota::per_second(NonZeroU32::new(rate).unwrap()).unwrap(),
623            )),
624            state: DashMapStateStore::new(),
625            gcra: DashMap::new(),
626            clock,
627            start,
628        };
629
630        let accepted = AtomicU32::new(0);
631        let num_threads = 50;
632
633        // Clock is frozen: no replenishment occurs
634        std::thread::scope(|s| {
635            for _ in 0..num_threads {
636                s.spawn(|| {
637                    if limiter.check_key(&"hot_key".to_string()).is_ok() {
638                        accepted.fetch_add(1, Ordering::Relaxed);
639                    }
640                });
641            }
642        });
643
644        let total = accepted.load(Ordering::Relaxed);
645        assert!(total >= 1, "At least one request should be accepted");
646        assert!(
647            total <= rate,
648            "Accepted {total} but burst capacity is {rate}"
649        );
650    }
651}
nautilus_network/ratelimiter/mod.rs

nautilus_network/ratelimiter/
mod.rs
