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/boinkor-net/governor).
17//!
18//! The governor does not support different quota for different key. It is an open [issue](https://github.com/boinkor-net/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 an error 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 an error 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    #[case(1)]
474    #[case(3)]
475    #[case(8)]
476    fn test_burst_capacity_and_pacing_after_idle(#[case] burst: u32) {
477        let interval = Duration::from_nanos(100);
478        let quota = Quota::with_period(interval)
479            .unwrap()
480            .allow_burst(NonZeroU32::new(burst).unwrap());
481        let clock = FakeRelativeClock::default();
482        let limiter = RateLimiter::new_with_clock(Some(quota), vec![], clock.clone());
483        let key = "requests";
484
485        for idle in [
486            Duration::ZERO,
487            Duration::from_secs(1),
488            Duration::from_secs(1),
489        ] {
490            clock.advance(idle);
491
492            for _ in 0..burst {
493                assert_eq!(limiter.check_key(&key), Ok(()));
494            }
495
496            let denied = limiter.check_key(&key).unwrap_err();
497            assert_eq!(denied.quota(), quota);
498            assert_eq!(denied.wait_time_from(clock.now()), interval);
499            clock.advance(Duration::from_nanos(99));
500            assert_eq!(
501                limiter
502                    .check_key(&key)
503                    .unwrap_err()
504                    .wait_time_from(clock.now()),
505                Duration::from_nanos(1),
506            );
507            clock.advance(Duration::from_nanos(1));
508            assert_eq!(limiter.check_key(&key), Ok(()));
509            assert_eq!(
510                limiter
511                    .check_key(&key)
512                    .unwrap_err()
513                    .wait_time_from(clock.now()),
514                interval,
515            );
516        }
517    }
518
519    #[rstest]
520    #[case(1)]
521    #[case(3)]
522    #[case(8)]
523    #[tokio::test]
524    async fn test_multi_key_admission_after_idle_preserves_capacity_and_atomicity(
525        #[case] burst: u32,
526    ) {
527        let interval = Duration::from_nanos(100);
528        let quota = Quota::with_period(interval)
529            .unwrap()
530            .allow_burst(NonZeroU32::new(burst).unwrap());
531        let clock = FakeRelativeClock::default();
532        let limiter = RateLimiter::new_with_clock(Some(quota), vec![], clock.clone());
533        let keys = ["first", "second"];
534
535        for _ in 0..burst {
536            limiter.await_keys_ready(Some(&keys)).await;
537        }
538
539        clock.advance(Duration::from_secs(1));
540        let now = clock.now();
541
542        for _ in 0..burst {
543            limiter.await_keys_ready(Some(&keys)).await;
544            assert_eq!(clock.now(), now);
545        }
546
547        assert_eq!(limiter.plan_keys(&["spare", "first"], now), Err(interval));
548        assert!(!limiter.state.contains_key(&"spare"));
549        limiter.await_keys_ready(Some(&["spare", "first"])).await;
550        assert_eq!(clock.now(), now + Nanos::from(interval));
551        assert_eq!(
552            limiter
553                .check_key(&"first")
554                .unwrap_err()
555                .wait_time_from(clock.now()),
556            interval
557        );
558
559        for _ in 1..burst {
560            assert_eq!(limiter.check_key(&"spare"), Ok(()));
561        }
562
563        assert_eq!(
564            limiter
565                .check_key(&"spare")
566                .unwrap_err()
567                .wait_time_from(clock.now()),
568            interval
569        );
570    }
571
572    #[rstest]
573    fn test_custom_key_quota() {
574        let mock_limiter = initialize_mock_rate_limiter();
575
576        // Add new key quota pair
577        mock_limiter.add_quota_for_key(
578            "custom".to_string(),
579            Quota::per_second(NonZeroU32::new(1).unwrap()).unwrap(),
580        );
581
582        // Check custom quota
583        assert!(mock_limiter.check_key(&"custom".to_string()).is_ok());
584        assert!(mock_limiter.check_key(&"custom".to_string()).is_err());
585
586        // Check that default quota still applies to other keys
587        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
588        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
589        assert!(mock_limiter.check_key(&"default".to_string()).is_err());
590    }
591
592    #[rstest]
593    fn test_multiple_keys() {
594        let mock_limiter = initialize_mock_rate_limiter();
595
596        mock_limiter.add_quota_for_key(
597            "key1".to_string(),
598            Quota::per_second(NonZeroU32::new(1).unwrap()).unwrap(),
599        );
600        mock_limiter.add_quota_for_key(
601            "key2".to_string(),
602            Quota::per_second(NonZeroU32::new(3).unwrap()).unwrap(),
603        );
604
605        // Test key1
606        assert!(mock_limiter.check_key(&"key1".to_string()).is_ok());
607        assert!(mock_limiter.check_key(&"key1".to_string()).is_err());
608
609        // Test key2
610        assert!(mock_limiter.check_key(&"key2".to_string()).is_ok());
611        assert!(mock_limiter.check_key(&"key2".to_string()).is_ok());
612        assert!(mock_limiter.check_key(&"key2".to_string()).is_ok());
613        assert!(mock_limiter.check_key(&"key2".to_string()).is_err());
614    }
615
616    #[rstest]
617    fn test_quota_reset() {
618        let mock_limiter = initialize_mock_rate_limiter();
619
620        // Exhaust quota
621        assert!(mock_limiter.check_key(&"reset".to_string()).is_ok());
622        assert!(mock_limiter.check_key(&"reset".to_string()).is_ok());
623        assert!(mock_limiter.check_key(&"reset".to_string()).is_err());
624
625        // Advance clock by less than a second
626        mock_limiter.advance_clock(Duration::from_millis(499));
627        assert!(mock_limiter.check_key(&"reset".to_string()).is_err());
628
629        // Advance clock to reset
630        mock_limiter.advance_clock(Duration::from_millis(501));
631        assert!(mock_limiter.check_key(&"reset".to_string()).is_ok());
632    }
633
634    #[rstest]
635    fn test_different_quotas() {
636        let mock_limiter = initialize_mock_rate_limiter();
637
638        mock_limiter.add_quota_for_key(
639            "per_second".to_string(),
640            Quota::per_second(NonZeroU32::new(2).unwrap()).unwrap(),
641        );
642        mock_limiter.add_quota_for_key(
643            "per_minute".to_string(),
644            Quota::per_minute(NonZeroU32::new(3).unwrap()),
645        );
646
647        // Test per_second quota
648        assert!(mock_limiter.check_key(&"per_second".to_string()).is_ok());
649        assert!(mock_limiter.check_key(&"per_second".to_string()).is_ok());
650        assert!(mock_limiter.check_key(&"per_second".to_string()).is_err());
651
652        // Test per_minute quota
653        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_ok());
654        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_ok());
655        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_ok());
656        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_err());
657
658        // Advance clock and check reset
659        mock_limiter.advance_clock(Duration::from_secs(1));
660        assert!(mock_limiter.check_key(&"per_second".to_string()).is_ok());
661        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_err());
662    }
663
664    #[tokio::test]
665    async fn test_await_keys_ready() {
666        let mock_limiter = initialize_mock_rate_limiter();
667
668        // Check base quota is not exceeded
669        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
670        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
671
672        // Check base quota is exceeded
673        assert!(mock_limiter.check_key(&"default".to_string()).is_err());
674
675        // Wait keys to be ready and check base quota is reset
676        mock_limiter.advance_clock(Duration::from_secs(1));
677        let keys = ["default".to_string()];
678        mock_limiter.await_keys_ready(Some(keys.as_slice())).await;
679        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
680    }
681
682    #[cfg_attr(
683        not(all(feature = "simulation", madsim)),
684        tokio::test(start_paused = true)
685    )]
686    #[cfg_attr(all(feature = "simulation", madsim), madsim::test)]
687    async fn test_await_keys_ready_reserves_keys_together() {
688        let fast = "fast".to_string();
689        let slow = "slow".to_string();
690        let limiter = Arc::new(RateLimiter::new_with_quota(
691            None,
692            vec![
693                (
694                    fast.clone(),
695                    Quota::with_period(Duration::from_secs(1)).unwrap(),
696                ),
697                (
698                    slow.clone(),
699                    Quota::with_period(Duration::from_secs(10)).unwrap(),
700                ),
701            ],
702        ));
703        limiter.check_key(&slow).unwrap();
704
705        let waiting_limiter = Arc::clone(&limiter);
706        let waiting_fast = fast.clone();
707        let waiting_slow = slow.clone();
708
709        let request = test_task::spawn(async move {
710            waiting_limiter
711                .await_keys_ready(Some(&[waiting_fast, waiting_slow]))
712                .await;
713        });
714        test_task::yield_now().await;
715
716        limiter.check_key(&fast).unwrap();
717        assert!(!request.is_finished());
718
719        advance_test_clock(Duration::from_millis(9_999)).await;
720        limiter.until_key_ready(&fast).await;
721        advance_test_clock(Duration::from_millis(1)).await;
722        test_task::yield_now().await;
723        assert!(!request.is_finished());
724
725        advance_test_clock(Duration::from_millis(998)).await;
726        test_task::yield_now().await;
727        assert!(!request.is_finished());
728
729        advance_test_clock(Duration::from_millis(1)).await;
730        request.await.unwrap();
731
732        assert!(limiter.check_key(&fast).is_err());
733        assert!(limiter.check_key(&slow).is_err());
734    }
735
736    #[cfg(all(feature = "simulation", madsim))]
737    async fn advance_test_clock(duration: Duration) {
738        madsim::time::advance(duration);
739        test_task::yield_now().await;
740    }
741
742    #[cfg(not(all(feature = "simulation", madsim)))]
743    async fn advance_test_clock(duration: Duration) {
744        tokio::time::advance(duration).await;
745    }
746
747    #[rstest]
748    fn test_per_second_returns_none_on_zero_replenish_interval() {
749        assert!(Quota::per_second(NonZeroU32::new(u32::MAX).unwrap()).is_none());
750    }
751
752    #[rstest]
753    fn test_per_minute_accepts_max_burst() {
754        let quota = Quota::per_minute(NonZeroU32::new(u32::MAX).unwrap());
755        assert_eq!(quota.replenish_interval(), Duration::from_nanos(13));
756    }
757
758    #[rstest]
759    fn test_per_hour_accepts_max_burst() {
760        let quota = Quota::per_hour(NonZeroU32::new(u32::MAX).unwrap());
761        assert_eq!(quota.replenish_interval(), Duration::from_nanos(838));
762    }
763
764    mod property_tests {
765        use proptest::prelude::*;
766        use rstest::rstest;
767
768        use crate::ratelimiter::nanos::Nanos;
769
770        proptest! {
771            #![proptest_config(ProptestConfig {
772                failure_persistence: Some(Box::new(
773                    proptest::test_runner::FileFailurePersistence::WithSource("ratelimiter")
774                )),
775                ..ProptestConfig::default()
776            })]
777
778            // Operators must saturate across the full u64 domain (a wrapped TAT admits everything)
779            #[rstest]
780            fn nanos_operators_never_panic(a in proptest::num::u64::ANY, b in proptest::num::u64::ANY) {
781                let na = Nanos::from(a);
782                let nb = Nanos::from(b);
783
784                prop_assert_eq!((na + nb).as_u64(), a.saturating_add(b));
785                prop_assert_eq!((na * b).as_u64(), a.saturating_mul(b));
786                prop_assert_eq!(na.saturating_sub(nb).as_u64(), a.saturating_sub(b));
787            }
788        }
789    }
790
791    #[rstest]
792    fn test_gcra_boundary_exact_replenishment() {
793        // Test GCRA boundary condition where t0 equals earliest_time exactly.
794        // This exercises the saturating_sub edge case deterministically without sleeps.
795        let mock_limiter = initialize_mock_rate_limiter();
796        let key = "boundary_test".to_string();
797
798        assert!(mock_limiter.check_key(&key).is_ok());
799        assert!(mock_limiter.check_key(&key).is_ok());
800        assert!(mock_limiter.check_key(&key).is_err());
801
802        // Advance clock by exactly one replenish interval (500ms for 2 req/sec)
803        let quota = Quota::per_second(NonZeroU32::new(2).unwrap()).unwrap();
804        let replenish_interval = quota.replenish_interval();
805        mock_limiter.advance_clock(replenish_interval);
806
807        assert!(
808            mock_limiter.check_key(&key).is_ok(),
809            "Request at exact replenish boundary should be allowed"
810        );
811        assert!(
812            mock_limiter.check_key(&key).is_err(),
813            "Immediate follow-up should be rate-limited"
814        );
815    }
816
817    #[rstest]
818    fn test_per_second_boundary_exact_limit() {
819        // 1_000_000_000ns / 1_000_000_000 = 1ns per replenish, the exact boundary
820        let quota = Quota::per_second(NonZeroU32::new(1_000_000_000).unwrap()).unwrap();
821        assert_eq!(quota.replenish_interval().as_nanos(), 1);
822    }
823
824    #[rstest]
825    fn test_per_second_returns_none_above_one_billion() {
826        // 1_000_000_000ns / 1_000_000_001 rounds to 0ns
827        assert!(Quota::per_second(NonZeroU32::new(1_000_000_001).unwrap()).is_none());
828    }
829
830    #[rstest]
831    #[case::large(Duration::from_secs(100), u32::MAX, Duration::from_mins(7_158_278_825))]
832    #[case::saturated(Duration::MAX, 2, Duration::MAX)]
833    fn test_burst_size_replenished_in(
834        #[case] replenish_interval: Duration,
835        #[case] burst_size: u32,
836        #[case] expected: Duration,
837    ) {
838        let quota = Quota::with_period(replenish_interval)
839            .unwrap()
840            .allow_burst(NonZeroU32::new(burst_size).unwrap());
841
842        assert_eq!(quota.burst_size_replenished_in(), expected);
843    }
844
845    #[rstest]
846    #[should_panic(expected = "t cannot be zero")]
847    fn test_from_gcra_parameters_panics_on_zero_t() {
848        let _ = Quota::from_gcra_parameters(Nanos::from(0u64), Nanos::from(100u64));
849    }
850
851    #[rstest]
852    #[should_panic(expected = "tau/t results in zero burst capacity")]
853    fn test_from_gcra_parameters_panics_on_zero_division() {
854        // tau=1, t=2 → integer division yields 0
855        let _ = Quota::from_gcra_parameters(Nanos::from(2u64), Nanos::from(1u64));
856    }
857
858    #[rstest]
859    #[should_panic(expected = "tau/t exceeds u32::MAX")]
860    fn test_from_gcra_parameters_panics_on_overflow() {
861        let _ = Quota::from_gcra_parameters(Nanos::from(1u64), Nanos::from(u64::MAX));
862    }
863
864    #[rstest]
865    fn test_concurrent_check_key_respects_burst() {
866        let rate = 10u32;
867        let clock = FakeRelativeClock::default();
868        let start = clock.now();
869        let limiter = RateLimiter {
870            default_gcra: Some(Gcra::new(
871                Quota::per_second(NonZeroU32::new(rate).unwrap()).unwrap(),
872            )),
873            state: DashMapStateStore::new(),
874            gcra: DashMap::new(),
875            clock,
876            start,
877            decision_lock: parking_lot::Mutex::new(()),
878        };
879
880        let accepted = AtomicU32::new(0);
881        let num_threads = 50;
882
883        // Clock is frozen: no replenishment occurs
884        std::thread::scope(|s| {
885            for _ in 0..num_threads {
886                s.spawn(|| {
887                    if limiter.check_key(&"hot_key".to_string()).is_ok() {
888                        accepted.fetch_add(1, Ordering::Relaxed);
889                    }
890                });
891            }
892        });
893
894        let total = accepted.load(Ordering::Relaxed);
895        assert!(total >= 1, "At least one request should be accepted");
896        assert!(
897            total <= rate,
898            "Accepted {total} but burst capacity is {rate}"
899        );
900    }
901}
nautilus_network/ratelimiter/mod.rs

nautilus_network/ratelimiter/
mod.rs
