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;
20pub mod quota;
21
22mod gcra;
23mod nanos;
24
25use std::{
26    collections::HashMap,
27    fmt::Debug,
28    hash::Hash,
29    num::NonZeroU64,
30    sync::{
31        Arc,
32        atomic::{AtomicU64, Ordering},
33    },
34    time::Duration,
35};
36
37use dashmap::DashMap;
38use parking_lot::Mutex;
39
40use self::{
41    clock::{Clock, FakeRelativeClock, MonotonicClock},
42    gcra::{Gcra, NotUntil},
43    nanos::Nanos,
44    quota::Quota,
45};
46
47/// An in-memory representation of a GCRA's rate-limiting state.
48///
49/// Implemented using [`AtomicU64`] operations, this state representation can be used to
50/// construct rate limiting states for other in-memory states: e.g., this crate uses
51/// `InMemoryState` as the states it tracks in the keyed rate limiters it implements.
52///
53/// Internally, the number tracked here is the theoretical arrival time (a GCRA term) in number of
54/// nanoseconds since the rate limiter was created.
55#[derive(Debug, Default)]
56pub struct InMemoryState(AtomicU64);
57
58impl InMemoryState {
59    fn load(&self) -> Option<Nanos> {
60        NonZeroU64::new(self.0.load(Ordering::Acquire)).map(|n| n.get().into())
61    }
62
63    fn store(&self, value: Nanos) {
64        self.0.store(value.into(), Ordering::Release);
65    }
66
67    /// Measures and updates the GCRA's state atomically, retrying on concurrent modifications.
68    ///
69    /// # Errors
70    ///
71    /// Returns an error if the provided closure returns an error.
72    pub(crate) fn measure_and_replace_one<T, F, E>(&self, mut f: F) -> Result<T, E>
73    where
74        F: FnMut(Option<Nanos>) -> Result<(T, Nanos), E>,
75    {
76        let mut prev = self.0.load(Ordering::Acquire);
77        loop {
78            let (result, new_data) = f(NonZeroU64::new(prev).map(|n| n.get().into()))?;
79
80            // Lock-free CAS loop: retry with current value if another thread modified it,
81            // uses weak variant (faster) since spurious failures are fine in a retry loop.
82            match self.0.compare_exchange_weak(
83                prev,
84                new_data.into(),
85                Ordering::Release,
86                Ordering::Relaxed,
87            ) {
88                Ok(_) => return Ok(result),
89                Err(e) => prev = e, // Retry with value written by another thread
90            }
91        }
92    }
93}
94
95/// A concurrent, thread-safe and fairly performant hashmap based on [`DashMap`].
96pub type DashMapStateStore<K> = DashMap<K, InMemoryState>;
97
98/// A way for rate limiters to keep state.
99///
100/// There are two important kinds of state stores: Direct and keyed. The direct kind have only
101/// one state, and are useful for "global" rate limit enforcement (e.g. a process should never
102/// do more than N tasks a day). The keyed kind allows one rate limit per key (e.g. an API
103/// call budget per client API key).
104///
105/// A direct state store is expressed as [`StateStore::Key`] = `NotKeyed`.
106/// Keyed state stores have a
107/// type parameter for the key and set their key to that.
108pub trait StateStore {
109    /// The type of key that the state store can represent.
110    type Key;
111
112    /// Updates a state store's rate limiting state for a given key, using the given closure.
113    ///
114    /// The closure parameter takes the old value (`None` if this is the first measurement) of the
115    /// state store at the key's location, checks if the request an be accommodated and:
116    ///
117    /// - If the request is rate-limited, returns `Err(E)`.
118    /// - If the request can make it through, returns `Ok(T)` (an arbitrary positive return
119    ///   value) and the updated state.
120    ///
121    /// It is `measure_and_replace`'s job then to safely replace the value at the key - it must
122    /// only update the value if the value hasn't changed. The implementations in this
123    /// crate use `AtomicU64` operations for this.
124    ///
125    /// # Errors
126    ///
127    /// Returns `Err(E)` if the closure returns an error or the request is rate-limited.
128    fn measure_and_replace<T, F, E>(&self, key: &Self::Key, f: F) -> Result<T, E>
129    where
130        F: Fn(Option<Nanos>) -> Result<(T, Nanos), E>;
131}
132
133impl<K: Hash + Eq + Clone> StateStore for DashMapStateStore<K> {
134    type Key = K;
135
136    fn measure_and_replace<T, F, E>(&self, key: &Self::Key, f: F) -> Result<T, E>
137    where
138        F: Fn(Option<Nanos>) -> Result<(T, Nanos), E>,
139    {
140        if let Some(v) = self.get(key) {
141            // fast path: measure existing entry
142            return v.measure_and_replace_one(f);
143        }
144        // make an entry and measure that:
145        let entry = self.entry(key.clone()).or_default();
146        (*entry).measure_and_replace_one(f)
147    }
148}
149
150/// A rate limiter that enforces different quotas per key using the GCRA algorithm.
151///
152/// This implementation allows setting different rate limits for different keys,
153/// with an optional default quota for keys that don't have specific quotas.
154pub struct RateLimiter<K, C>
155where
156    C: Clock,
157{
158    default_gcra: Option<Gcra>,
159    state: DashMapStateStore<K>,
160    gcra: DashMap<K, Gcra>,
161    clock: C,
162    start: C::Instant,
163    decision_lock: Mutex<()>,
164}
165
166impl<K, C> Debug for RateLimiter<K, C>
167where
168    K: Debug,
169    C: Clock,
170{
171    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
172        f.debug_struct(stringify!(RateLimiter)).finish()
173    }
174}
175
176impl<K> RateLimiter<K, MonotonicClock>
177where
178    K: Eq + Hash,
179{
180    /// Creates a new rate limiter with a base quota and keyed quotas.
181    ///
182    /// The base quota applies to all keys that don't have specific quotas.
183    /// Keyed quotas override the base quota for specific keys.
184    #[must_use]
185    pub fn new_with_quota(base_quota: Option<Quota>, keyed_quotas: Vec<(K, Quota)>) -> Self {
186        Self::new_with_clock(base_quota, keyed_quotas, MonotonicClock)
187    }
188}
189
190impl<K, C> RateLimiter<K, C>
191where
192    K: Eq + Hash,
193    C: Clock,
194{
195    /// Creates a new rate limiter with an explicit clock.
196    ///
197    /// The base quota applies to all keys that do not have specific quotas.
198    /// Keyed quotas override the base quota for specific keys.
199    #[must_use]
200    pub fn new_with_clock(
201        base_quota: Option<Quota>,
202        keyed_quotas: Vec<(K, Quota)>,
203        clock: C,
204    ) -> Self {
205        let start = clock.now();
206        let gcra: DashMap<_, _> = keyed_quotas
207            .into_iter()
208            .map(|(k, q)| (k, Gcra::new(q)))
209            .collect();
210        Self {
211            default_gcra: base_quota.map(Gcra::new),
212            state: DashMapStateStore::new(),
213            gcra,
214            clock,
215            start,
216            decision_lock: Mutex::new(()),
217        }
218    }
219}
220
221impl<K> RateLimiter<K, FakeRelativeClock>
222where
223    K: Hash + Eq + Clone,
224{
225    /// Advances the fake clock by the specified duration.
226    ///
227    /// This is only available for testing with `FakeRelativeClock`.
228    pub fn advance_clock(&self, by: Duration) {
229        self.clock.advance(by);
230    }
231}
232
233impl<K, C> RateLimiter<K, C>
234where
235    K: Hash + Eq + Clone,
236    C: Clock,
237{
238    /// Adds or updates a quota for a specific key.
239    pub fn add_quota_for_key(&self, key: K, value: Quota) {
240        let _guard = self.decision_lock.lock();
241        self.gcra.insert(key, Gcra::new(value));
242    }
243
244    /// Checks if the given key is allowed under the rate limit.
245    ///
246    /// # Errors
247    ///
248    /// Returns `Err(NotUntil)` if the key is rate-limited, indicating when it will be allowed.
249    pub fn check_key(&self, key: &K) -> Result<(), NotUntil<C::Instant>> {
250        let _guard = self.decision_lock.lock();
251
252        match self.gcra.get(key) {
253            Some(quota) => quota.test_and_update(self.start, key, &self.state, self.clock.now()),
254            None => self.default_gcra.as_ref().map_or(Ok(()), |gcra| {
255                gcra.test_and_update(self.start, key, &self.state, self.clock.now())
256            }),
257        }
258    }
259
260    /// Waits until the specified key is ready (not rate-limited).
261    pub async fn until_key_ready(&self, key: &K) {
262        loop {
263            match self.check_key(key) {
264                Ok(()) => {
265                    break;
266                }
267                Err(e) => {
268                    self.clock.sleep(e.wait_time_from(self.clock.now())).await;
269                }
270            }
271        }
272    }
273
274    /// Waits until all specified keys are ready (not rate-limited).
275    ///
276    /// If no keys are provided, this function returns immediately.
277    pub async fn await_keys_ready(&self, keys: Option<&[K]>) {
278        let Some(keys) = keys else {
279            return;
280        };
281
282        loop {
283            let wait = {
284                let _guard = self.decision_lock.lock();
285
286                match self.plan_keys(keys, self.clock.now()) {
287                    Ok(planned) => {
288                        self.commit_keys(planned);
289                        None
290                    }
291                    Err(wait) => Some(wait),
292                }
293            };
294
295            match wait {
296                Some(wait) => self.clock.sleep(wait).await,
297                None => return,
298            }
299        }
300    }
301
302    fn plan_keys<'a>(
303        &self,
304        keys: &'a [K],
305        now: C::Instant,
306    ) -> Result<HashMap<&'a K, Nanos>, Duration> {
307        let mut planned = HashMap::with_capacity(keys.len());
308        let mut wait: Option<Duration> = None;
309
310        for key in keys {
311            let tat = planned
312                .get(key)
313                .copied()
314                .or_else(|| self.state.get(key).and_then(|state| state.load()));
315            let decision = match self.gcra.get(key) {
316                Some(quota) => Some(quota.test(self.start, tat, now)),
317                None => self
318                    .default_gcra
319                    .as_ref()
320                    .map(|gcra| gcra.test(self.start, tat, now)),
321            };
322
323            match decision {
324                Some(Ok(next)) => {
325                    planned.insert(key, next);
326                }
327                Some(Err(denied)) => {
328                    let duration = denied.wait_time_from(now);
329                    wait = Some(wait.map_or(duration, |current| current.max(duration)));
330                }
331                None => {}
332            }
333        }
334
335        match wait {
336            Some(wait) => Err(wait),
337            None => Ok(planned),
338        }
339    }
340
341    fn commit_keys(&self, planned: HashMap<&K, Nanos>) {
342        for (key, tat) in planned {
343            self.state.entry(key.clone()).or_default().store(tat);
344        }
345    }
346}
347
348impl<K> RateLimiter<K, MonotonicClock>
349where
350    K: Hash + Eq + Clone,
351{
352    pub(crate) async fn await_limiters_ready(rate_limiters: &[Arc<Self>], keys: Option<&[K]>) {
353        let Some(keys) = keys else {
354            return;
355        };
356
357        if rate_limiters.is_empty() || keys.is_empty() {
358            return;
359        }
360
361        let mut ordered = rate_limiters.iter().map(Arc::as_ref).collect::<Vec<_>>();
362        ordered.sort_unstable_by_key(|limiter| std::ptr::from_ref(*limiter) as usize);
363        ordered.dedup_by(|a, b| std::ptr::eq(*a, *b));
364
365        loop {
366            let wait = {
367                let _guards = ordered
368                    .iter()
369                    .map(|limiter| limiter.decision_lock.lock())
370                    .collect::<Vec<_>>();
371                let mut plans = Vec::with_capacity(ordered.len());
372                let mut wait: Option<Duration> = None;
373
374                for limiter in &ordered {
375                    match limiter.plan_keys(keys, limiter.clock.now()) {
376                        Ok(planned) => plans.push((*limiter, planned)),
377                        Err(duration) => {
378                            wait = Some(wait.map_or(duration, |current| current.max(duration)));
379                        }
380                    }
381                }
382
383                if wait.is_none() {
384                    for (limiter, planned) in plans {
385                        limiter.commit_keys(planned);
386                    }
387                }
388                wait
389            };
390
391            match wait {
392                Some(wait) => ordered[0].clock.sleep(wait).await,
393                None => return,
394            }
395        }
396    }
397}
398
399#[cfg(test)]
400mod tests {
401    use std::{
402        num::NonZeroU32,
403        sync::{
404            Arc,
405            atomic::{AtomicU32, Ordering},
406        },
407        time::Duration,
408    };
409
410    use dashmap::DashMap;
411    #[cfg(all(feature = "simulation", madsim))]
412    use madsim::task as test_task;
413    use rstest::rstest;
414    #[cfg(not(all(feature = "simulation", madsim)))]
415    use tokio::task as test_task;
416
417    use super::{
418        DashMapStateStore, RateLimiter,
419        clock::{Clock, FakeRelativeClock},
420        gcra::Gcra,
421        nanos::Nanos,
422        quota::Quota,
423    };
424
425    fn initialize_mock_rate_limiter() -> RateLimiter<String, FakeRelativeClock> {
426        let clock = FakeRelativeClock::default();
427        let start = clock.now();
428        let gcra = DashMap::new();
429        let base_quota = Quota::per_second(NonZeroU32::new(2).unwrap()).unwrap();
430        RateLimiter {
431            default_gcra: Some(Gcra::new(base_quota)),
432            state: DashMapStateStore::new(),
433            gcra,
434            clock,
435            start,
436            decision_lock: parking_lot::Mutex::new(()),
437        }
438    }
439
440    #[rstest]
441    fn test_enormous_quota_denies_after_burst() {
442        // Regression: a period beyond ~584 years panicked in Gcra::new; with
443        // clamping it must admit the burst and then deny, not admit everything
444        let quota = Quota::with_period(Duration::MAX)
445            .unwrap()
446            .allow_burst(NonZeroU32::new(u32::MAX).unwrap());
447        let clock = FakeRelativeClock::default();
448        let limiter: RateLimiter<String, FakeRelativeClock> =
449            RateLimiter::new_with_clock(Some(quota), vec![], clock);
450
451        let key = "key".to_string();
452        assert!(limiter.check_key(&key).is_ok());
453        assert!(limiter.check_key(&key).is_err());
454    }
455
456    #[rstest]
457    fn test_default_quota() {
458        let mock_limiter = initialize_mock_rate_limiter();
459
460        // Check base quota is not exceeded
461        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
462        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
463
464        // Check base quota is exceeded
465        assert!(mock_limiter.check_key(&"default".to_string()).is_err());
466
467        // Increment clock and check base quota is reset
468        mock_limiter.advance_clock(Duration::from_secs(1));
469        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
470    }
471
472    #[rstest]
473    fn test_custom_key_quota() {
474        let mock_limiter = initialize_mock_rate_limiter();
475
476        // Add new key quota pair
477        mock_limiter.add_quota_for_key(
478            "custom".to_string(),
479            Quota::per_second(NonZeroU32::new(1).unwrap()).unwrap(),
480        );
481
482        // Check custom quota
483        assert!(mock_limiter.check_key(&"custom".to_string()).is_ok());
484        assert!(mock_limiter.check_key(&"custom".to_string()).is_err());
485
486        // Check that default quota still applies to other keys
487        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
488        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
489        assert!(mock_limiter.check_key(&"default".to_string()).is_err());
490    }
491
492    #[rstest]
493    fn test_multiple_keys() {
494        let mock_limiter = initialize_mock_rate_limiter();
495
496        mock_limiter.add_quota_for_key(
497            "key1".to_string(),
498            Quota::per_second(NonZeroU32::new(1).unwrap()).unwrap(),
499        );
500        mock_limiter.add_quota_for_key(
501            "key2".to_string(),
502            Quota::per_second(NonZeroU32::new(3).unwrap()).unwrap(),
503        );
504
505        // Test key1
506        assert!(mock_limiter.check_key(&"key1".to_string()).is_ok());
507        assert!(mock_limiter.check_key(&"key1".to_string()).is_err());
508
509        // Test key2
510        assert!(mock_limiter.check_key(&"key2".to_string()).is_ok());
511        assert!(mock_limiter.check_key(&"key2".to_string()).is_ok());
512        assert!(mock_limiter.check_key(&"key2".to_string()).is_ok());
513        assert!(mock_limiter.check_key(&"key2".to_string()).is_err());
514    }
515
516    #[rstest]
517    fn test_quota_reset() {
518        let mock_limiter = initialize_mock_rate_limiter();
519
520        // Exhaust quota
521        assert!(mock_limiter.check_key(&"reset".to_string()).is_ok());
522        assert!(mock_limiter.check_key(&"reset".to_string()).is_ok());
523        assert!(mock_limiter.check_key(&"reset".to_string()).is_err());
524
525        // Advance clock by less than a second
526        mock_limiter.advance_clock(Duration::from_millis(499));
527        assert!(mock_limiter.check_key(&"reset".to_string()).is_err());
528
529        // Advance clock to reset
530        mock_limiter.advance_clock(Duration::from_millis(501));
531        assert!(mock_limiter.check_key(&"reset".to_string()).is_ok());
532    }
533
534    #[rstest]
535    fn test_different_quotas() {
536        let mock_limiter = initialize_mock_rate_limiter();
537
538        mock_limiter.add_quota_for_key(
539            "per_second".to_string(),
540            Quota::per_second(NonZeroU32::new(2).unwrap()).unwrap(),
541        );
542        mock_limiter.add_quota_for_key(
543            "per_minute".to_string(),
544            Quota::per_minute(NonZeroU32::new(3).unwrap()),
545        );
546
547        // Test per_second quota
548        assert!(mock_limiter.check_key(&"per_second".to_string()).is_ok());
549        assert!(mock_limiter.check_key(&"per_second".to_string()).is_ok());
550        assert!(mock_limiter.check_key(&"per_second".to_string()).is_err());
551
552        // Test per_minute quota
553        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_ok());
554        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_ok());
555        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_ok());
556        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_err());
557
558        // Advance clock and check reset
559        mock_limiter.advance_clock(Duration::from_secs(1));
560        assert!(mock_limiter.check_key(&"per_second".to_string()).is_ok());
561        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_err());
562    }
563
564    #[tokio::test]
565    async fn test_await_keys_ready() {
566        let mock_limiter = initialize_mock_rate_limiter();
567
568        // Check base quota is not exceeded
569        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
570        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
571
572        // Check base quota is exceeded
573        assert!(mock_limiter.check_key(&"default".to_string()).is_err());
574
575        // Wait keys to be ready and check base quota is reset
576        mock_limiter.advance_clock(Duration::from_secs(1));
577        let keys = ["default".to_string()];
578        mock_limiter.await_keys_ready(Some(keys.as_slice())).await;
579        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
580    }
581
582    #[cfg_attr(
583        not(all(feature = "simulation", madsim)),
584        tokio::test(start_paused = true)
585    )]
586    #[cfg_attr(all(feature = "simulation", madsim), madsim::test)]
587    async fn test_await_keys_ready_reserves_keys_together() {
588        let fast = "fast".to_string();
589        let slow = "slow".to_string();
590        let limiter = Arc::new(RateLimiter::new_with_quota(
591            None,
592            vec![
593                (
594                    fast.clone(),
595                    Quota::with_period(Duration::from_secs(1)).unwrap(),
596                ),
597                (
598                    slow.clone(),
599                    Quota::with_period(Duration::from_secs(10)).unwrap(),
600                ),
601            ],
602        ));
603        limiter.check_key(&slow).unwrap();
604
605        let waiting_limiter = Arc::clone(&limiter);
606        let waiting_fast = fast.clone();
607        let waiting_slow = slow.clone();
608
609        let request = test_task::spawn(async move {
610            waiting_limiter
611                .await_keys_ready(Some(&[waiting_fast, waiting_slow]))
612                .await;
613        });
614        test_task::yield_now().await;
615
616        limiter.check_key(&fast).unwrap();
617        assert!(!request.is_finished());
618
619        advance_test_clock(Duration::from_millis(9_999)).await;
620        limiter.until_key_ready(&fast).await;
621        limiter.until_key_ready(&fast).await;
622        advance_test_clock(Duration::from_millis(1)).await;
623        test_task::yield_now().await;
624        assert!(!request.is_finished());
625
626        advance_test_clock(Duration::from_millis(998)).await;
627        test_task::yield_now().await;
628        assert!(!request.is_finished());
629
630        advance_test_clock(Duration::from_millis(1)).await;
631        request.await.unwrap();
632
633        assert!(limiter.check_key(&fast).is_err());
634        assert!(limiter.check_key(&slow).is_err());
635    }
636
637    #[cfg(all(feature = "simulation", madsim))]
638    async fn advance_test_clock(duration: Duration) {
639        madsim::time::advance(duration);
640        test_task::yield_now().await;
641    }
642
643    #[cfg(not(all(feature = "simulation", madsim)))]
644    async fn advance_test_clock(duration: Duration) {
645        tokio::time::advance(duration).await;
646    }
647
648    #[rstest]
649    fn test_per_second_returns_none_on_zero_replenish_interval() {
650        assert!(Quota::per_second(NonZeroU32::new(u32::MAX).unwrap()).is_none());
651    }
652
653    #[rstest]
654    fn test_per_minute_accepts_max_burst() {
655        let quota = Quota::per_minute(NonZeroU32::new(u32::MAX).unwrap());
656        assert!(quota.replenish_interval().as_nanos() > 0);
657    }
658
659    #[rstest]
660    fn test_per_hour_accepts_max_burst() {
661        let quota = Quota::per_hour(NonZeroU32::new(u32::MAX).unwrap());
662        assert!(quota.replenish_interval().as_nanos() > 0);
663    }
664
665    mod property_tests {
666        use proptest::prelude::*;
667        use rstest::rstest;
668
669        use crate::ratelimiter::nanos::Nanos;
670
671        proptest! {
672            #![proptest_config(ProptestConfig {
673                failure_persistence: Some(Box::new(
674                    proptest::test_runner::FileFailurePersistence::WithSource("ratelimiter")
675                )),
676                ..ProptestConfig::default()
677            })]
678
679            // Operators must saturate across the full u64 domain (a wrapped TAT admits everything)
680            #[rstest]
681            fn nanos_operators_never_panic(a in proptest::num::u64::ANY, b in proptest::num::u64::ANY) {
682                let na = Nanos::from(a);
683                let nb = Nanos::from(b);
684
685                prop_assert_eq!((na + nb).as_u64(), a.saturating_add(b));
686                prop_assert_eq!((na * b).as_u64(), a.saturating_mul(b));
687                prop_assert_eq!(na.saturating_sub(nb).as_u64(), a.saturating_sub(b));
688            }
689        }
690    }
691
692    #[rstest]
693    fn test_gcra_boundary_exact_replenishment() {
694        // Test GCRA boundary condition where t0 equals earliest_time exactly.
695        // This exercises the saturating_sub edge case deterministically without sleeps.
696        let mock_limiter = initialize_mock_rate_limiter();
697        let key = "boundary_test".to_string();
698
699        assert!(mock_limiter.check_key(&key).is_ok());
700        assert!(mock_limiter.check_key(&key).is_ok());
701        assert!(mock_limiter.check_key(&key).is_err());
702
703        // Advance clock by exactly one replenish interval (500ms for 2 req/sec)
704        let quota = Quota::per_second(NonZeroU32::new(2).unwrap()).unwrap();
705        let replenish_interval = quota.replenish_interval();
706        mock_limiter.advance_clock(replenish_interval);
707
708        assert!(
709            mock_limiter.check_key(&key).is_ok(),
710            "Request at exact replenish boundary should be allowed"
711        );
712        assert!(
713            mock_limiter.check_key(&key).is_err(),
714            "Immediate follow-up should be rate-limited"
715        );
716    }
717
718    #[rstest]
719    fn test_per_second_boundary_exact_limit() {
720        // 1_000_000_000ns / 1_000_000_000 = 1ns per replenish, the exact boundary
721        let quota = Quota::per_second(NonZeroU32::new(1_000_000_000).unwrap()).unwrap();
722        assert_eq!(quota.replenish_interval().as_nanos(), 1);
723    }
724
725    #[rstest]
726    fn test_per_second_returns_none_above_one_billion() {
727        // 1_000_000_000ns / 1_000_000_001 rounds to 0ns
728        assert!(Quota::per_second(NonZeroU32::new(1_000_000_001).unwrap()).is_none());
729    }
730
731    #[rstest]
732    #[case::large(Duration::from_secs(100), u32::MAX, Duration::from_mins(7_158_278_825))]
733    #[case::saturated(Duration::MAX, 2, Duration::MAX)]
734    fn test_burst_size_replenished_in(
735        #[case] replenish_interval: Duration,
736        #[case] burst_size: u32,
737        #[case] expected: Duration,
738    ) {
739        let quota = Quota::with_period(replenish_interval)
740            .unwrap()
741            .allow_burst(NonZeroU32::new(burst_size).unwrap());
742
743        assert_eq!(quota.burst_size_replenished_in(), expected);
744    }
745
746    #[rstest]
747    #[should_panic(expected = "t cannot be zero")]
748    fn test_from_gcra_parameters_panics_on_zero_t() {
749        let _ = Quota::from_gcra_parameters(Nanos::from(0u64), Nanos::from(100u64));
750    }
751
752    #[rstest]
753    #[should_panic(expected = "tau/t results in zero burst capacity")]
754    fn test_from_gcra_parameters_panics_on_zero_division() {
755        // tau=1, t=2 → integer division yields 0
756        let _ = Quota::from_gcra_parameters(Nanos::from(2u64), Nanos::from(1u64));
757    }
758
759    #[rstest]
760    #[should_panic(expected = "tau/t exceeds u32::MAX")]
761    fn test_from_gcra_parameters_panics_on_overflow() {
762        let _ = Quota::from_gcra_parameters(Nanos::from(1u64), Nanos::from(u64::MAX));
763    }
764
765    #[rstest]
766    fn test_concurrent_check_key_respects_burst() {
767        let rate = 10u32;
768        let clock = FakeRelativeClock::default();
769        let start = clock.now();
770        let limiter = RateLimiter {
771            default_gcra: Some(Gcra::new(
772                Quota::per_second(NonZeroU32::new(rate).unwrap()).unwrap(),
773            )),
774            state: DashMapStateStore::new(),
775            gcra: DashMap::new(),
776            clock,
777            start,
778            decision_lock: parking_lot::Mutex::new(()),
779        };
780
781        let accepted = AtomicU32::new(0);
782        let num_threads = 50;
783
784        // Clock is frozen: no replenishment occurs
785        std::thread::scope(|s| {
786            for _ in 0..num_threads {
787                s.spawn(|| {
788                    if limiter.check_key(&"hot_key".to_string()).is_ok() {
789                        accepted.fetch_add(1, Ordering::Relaxed);
790                    }
791                });
792            }
793        });
794
795        let total = accepted.load(Ordering::Relaxed);
796        assert!(total >= 1, "At least one request should be accepted");
797        assert!(
798            total <= rate,
799            "Accepted {total} but burst capacity is {rate}"
800        );
801    }
802}
nautilus_network/ratelimiter/mod.rs

nautilus_network/ratelimiter/
mod.rs
