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
145        // make an entry and measure that:
146        let entry = self.entry(key.clone()).or_default();
147        (*entry).measure_and_replace_one(f)
148    }
149}
150
151/// A rate limiter that enforces different quotas per key using the GCRA algorithm.
152///
153/// This implementation allows setting different rate limits for different keys,
154/// with an optional default quota for keys that don't have specific quotas.
155pub struct RateLimiter<K, C>
156where
157    C: Clock,
158{
159    default_gcra: Option<Gcra>,
160    state: DashMapStateStore<K>,
161    gcra: DashMap<K, Gcra>,
162    clock: C,
163    start: C::Instant,
164    decision_lock: Mutex<()>,
165}
166
167impl<K, C> Debug for RateLimiter<K, C>
168where
169    K: Debug,
170    C: Clock,
171{
172    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
173        f.debug_struct(stringify!(RateLimiter)).finish()
174    }
175}
176
177impl<K> RateLimiter<K, MonotonicClock>
178where
179    K: Eq + Hash,
180{
181    /// Creates a new rate limiter with a base quota and keyed quotas.
182    ///
183    /// The base quota applies to all keys that don't have specific quotas.
184    /// Keyed quotas override the base quota for specific keys.
185    #[must_use]
186    pub fn new_with_quota(base_quota: Option<Quota>, keyed_quotas: Vec<(K, Quota)>) -> Self {
187        Self::new_with_clock(base_quota, keyed_quotas, MonotonicClock)
188    }
189}
190
191impl<K, C> RateLimiter<K, C>
192where
193    K: Eq + Hash,
194    C: Clock,
195{
196    /// Creates a new rate limiter with an explicit clock.
197    ///
198    /// The base quota applies to all keys that do not have specific quotas.
199    /// Keyed quotas override the base quota for specific keys.
200    #[must_use]
201    pub fn new_with_clock(
202        base_quota: Option<Quota>,
203        keyed_quotas: Vec<(K, Quota)>,
204        clock: C,
205    ) -> Self {
206        let start = clock.now();
207        let gcra: DashMap<_, _> = keyed_quotas
208            .into_iter()
209            .map(|(k, q)| (k, Gcra::new(q)))
210            .collect();
211
212        Self {
213            default_gcra: base_quota.map(Gcra::new),
214            state: DashMapStateStore::new(),
215            gcra,
216            clock,
217            start,
218            decision_lock: Mutex::new(()),
219        }
220    }
221}
222
223impl<K> RateLimiter<K, FakeRelativeClock>
224where
225    K: Hash + Eq + Clone,
226{
227    /// Advances the fake clock by the specified duration.
228    ///
229    /// This is only available for testing with `FakeRelativeClock`.
230    pub fn advance_clock(&self, by: Duration) {
231        self.clock.advance(by);
232    }
233}
234
235impl<K, C> RateLimiter<K, C>
236where
237    K: Hash + Eq + Clone,
238    C: Clock,
239{
240    /// Adds or updates a quota for a specific key.
241    pub fn add_quota_for_key(&self, key: K, value: Quota) {
242        let _guard = self.decision_lock.lock();
243        self.gcra.insert(key, Gcra::new(value));
244    }
245
246    /// Checks if the given key is allowed under the rate limit.
247    ///
248    /// # Errors
249    ///
250    /// Returns an error if the key is rate-limited, indicating when it will be allowed.
251    pub fn check_key(&self, key: &K) -> Result<(), NotUntil<C::Instant>> {
252        let _guard = self.decision_lock.lock();
253
254        match self.gcra.get(key) {
255            Some(quota) => quota.test_and_update(self.start, key, &self.state, self.clock.now()),
256            None => self.default_gcra.as_ref().map_or(Ok(()), |gcra| {
257                gcra.test_and_update(self.start, key, &self.state, self.clock.now())
258            }),
259        }
260    }
261
262    /// Waits until the specified key is ready (not rate-limited).
263    pub async fn until_key_ready(&self, key: &K) {
264        loop {
265            match self.check_key(key) {
266                Ok(()) => {
267                    break;
268                }
269                Err(e) => {
270                    self.clock.sleep(e.wait_time_from(self.clock.now())).await;
271                }
272            }
273        }
274    }
275
276    /// Waits until all specified keys are ready (not rate-limited).
277    ///
278    /// If no keys are provided, this function returns immediately.
279    pub async fn await_keys_ready(&self, keys: Option<&[K]>) {
280        let Some(keys) = keys else {
281            return;
282        };
283
284        loop {
285            let wait = {
286                let _guard = self.decision_lock.lock();
287
288                match self.plan_keys(keys, self.clock.now()) {
289                    Ok(planned) => {
290                        self.commit_keys(planned);
291                        None
292                    }
293                    Err(wait) => Some(wait),
294                }
295            };
296
297            match wait {
298                Some(wait) => self.clock.sleep(wait).await,
299                None => return,
300            }
301        }
302    }
303
304    fn plan_keys<'a>(
305        &self,
306        keys: &'a [K],
307        now: C::Instant,
308    ) -> Result<HashMap<&'a K, Nanos>, Duration> {
309        let mut planned = HashMap::with_capacity(keys.len());
310        let mut wait: Option<Duration> = None;
311
312        for key in keys {
313            let tat = planned
314                .get(key)
315                .copied()
316                .or_else(|| self.state.get(key).and_then(|state| state.load()));
317
318            let decision = match self.gcra.get(key) {
319                Some(quota) => Some(quota.test(self.start, tat, now)),
320                None => self
321                    .default_gcra
322                    .as_ref()
323                    .map(|gcra| gcra.test(self.start, tat, now)),
324            };
325
326            match decision {
327                Some(Ok(next)) => {
328                    planned.insert(key, next);
329                }
330                Some(Err(denied)) => {
331                    let duration = denied.wait_time_from(now);
332                    wait = Some(wait.map_or(duration, |current| current.max(duration)));
333                }
334                None => {}
335            }
336        }
337
338        match wait {
339            Some(wait) => Err(wait),
340            None => Ok(planned),
341        }
342    }
343
344    fn commit_keys(&self, planned: HashMap<&K, Nanos>) {
345        for (key, tat) in planned {
346            self.state.entry(key.clone()).or_default().store(tat);
347        }
348    }
349}
350
351impl<K> RateLimiter<K, MonotonicClock>
352where
353    K: Hash + Eq + Clone,
354{
355    pub(crate) async fn await_limiters_ready(rate_limiters: &[Arc<Self>], keys: Option<&[K]>) {
356        let Some(keys) = keys else {
357            return;
358        };
359
360        if rate_limiters.is_empty() || keys.is_empty() {
361            return;
362        }
363
364        let mut ordered = rate_limiters.iter().map(Arc::as_ref).collect::<Vec<_>>();
365        ordered.sort_unstable_by_key(|limiter| std::ptr::from_ref(*limiter) as usize);
366        ordered.dedup_by(|a, b| std::ptr::eq(*a, *b));
367
368        loop {
369            let wait = {
370                let _guards = ordered
371                    .iter()
372                    .map(|limiter| limiter.decision_lock.lock())
373                    .collect::<Vec<_>>();
374                let mut plans = Vec::with_capacity(ordered.len());
375                let mut wait: Option<Duration> = None;
376
377                for limiter in &ordered {
378                    match limiter.plan_keys(keys, limiter.clock.now()) {
379                        Ok(planned) => plans.push((*limiter, planned)),
380                        Err(duration) => {
381                            wait = Some(wait.map_or(duration, |current| current.max(duration)));
382                        }
383                    }
384                }
385
386                if wait.is_none() {
387                    for (limiter, planned) in plans {
388                        limiter.commit_keys(planned);
389                    }
390                }
391
392                wait
393            };
394
395            match wait {
396                Some(wait) => ordered[0].clock.sleep(wait).await,
397                None => return,
398            }
399        }
400    }
401}
402
403#[cfg(test)]
404mod tests {
405    use std::{
406        num::NonZeroU32,
407        sync::{
408            Arc,
409            atomic::{AtomicU32, Ordering},
410        },
411        time::Duration,
412    };
413
414    use dashmap::DashMap;
415    #[cfg(all(feature = "simulation", madsim))]
416    use madsim::task as test_task;
417    use rstest::rstest;
418    #[cfg(not(all(feature = "simulation", madsim)))]
419    use tokio::task as test_task;
420
421    use super::{
422        DashMapStateStore, RateLimiter,
423        clock::{Clock, FakeRelativeClock},
424        gcra::Gcra,
425        nanos::Nanos,
426        quota::Quota,
427    };
428
429    fn initialize_mock_rate_limiter() -> RateLimiter<String, FakeRelativeClock> {
430        let clock = FakeRelativeClock::default();
431        let start = clock.now();
432        let gcra = DashMap::new();
433        let base_quota = Quota::per_second(NonZeroU32::new(2).unwrap()).unwrap();
434
435        RateLimiter {
436            default_gcra: Some(Gcra::new(base_quota)),
437            state: DashMapStateStore::new(),
438            gcra,
439            clock,
440            start,
441            decision_lock: parking_lot::Mutex::new(()),
442        }
443    }
444
445    #[rstest]
446    fn test_enormous_quota_denies_after_burst() {
447        // Regression: a period beyond ~584 years panicked in Gcra::new; with
448        // clamping it must admit the burst and then deny, not admit everything
449        let quota = Quota::with_period(Duration::MAX)
450            .unwrap()
451            .allow_burst(NonZeroU32::new(u32::MAX).unwrap());
452        let clock = FakeRelativeClock::default();
453        let limiter: RateLimiter<String, FakeRelativeClock> =
454            RateLimiter::new_with_clock(Some(quota), vec![], clock);
455
456        let key = "key".to_string();
457        assert!(limiter.check_key(&key).is_ok());
458        assert!(limiter.check_key(&key).is_err());
459    }
460
461    #[rstest]
462    fn test_default_quota() {
463        let mock_limiter = initialize_mock_rate_limiter();
464
465        // Check base quota is not exceeded
466        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
467        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
468
469        // Check base quota is exceeded
470        assert!(mock_limiter.check_key(&"default".to_string()).is_err());
471
472        // Increment clock and check base quota is reset
473        mock_limiter.advance_clock(Duration::from_secs(1));
474        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
475    }
476
477    #[rstest]
478    #[case(1)]
479    #[case(3)]
480    #[case(8)]
481    fn test_burst_capacity_and_pacing_after_idle(#[case] burst: u32) {
482        let interval = Duration::from_nanos(100);
483        let quota = Quota::with_period(interval)
484            .unwrap()
485            .allow_burst(NonZeroU32::new(burst).unwrap());
486        let clock = FakeRelativeClock::default();
487        let limiter = RateLimiter::new_with_clock(Some(quota), vec![], clock.clone());
488        let key = "requests";
489
490        for idle in [
491            Duration::ZERO,
492            Duration::from_secs(1),
493            Duration::from_secs(1),
494        ] {
495            clock.advance(idle);
496
497            for _ in 0..burst {
498                assert_eq!(limiter.check_key(&key), Ok(()));
499            }
500
501            let denied = limiter.check_key(&key).unwrap_err();
502            assert_eq!(denied.quota(), quota);
503            assert_eq!(denied.wait_time_from(clock.now()), interval);
504            clock.advance(Duration::from_nanos(99));
505            assert_eq!(
506                limiter
507                    .check_key(&key)
508                    .unwrap_err()
509                    .wait_time_from(clock.now()),
510                Duration::from_nanos(1),
511            );
512            clock.advance(Duration::from_nanos(1));
513            assert_eq!(limiter.check_key(&key), Ok(()));
514            assert_eq!(
515                limiter
516                    .check_key(&key)
517                    .unwrap_err()
518                    .wait_time_from(clock.now()),
519                interval,
520            );
521        }
522    }
523
524    #[rstest]
525    #[case(1)]
526    #[case(3)]
527    #[case(8)]
528    #[tokio::test]
529    async fn test_multi_key_admission_after_idle_preserves_capacity_and_atomicity(
530        #[case] burst: u32,
531    ) {
532        let interval = Duration::from_nanos(100);
533        let quota = Quota::with_period(interval)
534            .unwrap()
535            .allow_burst(NonZeroU32::new(burst).unwrap());
536        let clock = FakeRelativeClock::default();
537        let limiter = RateLimiter::new_with_clock(Some(quota), vec![], clock.clone());
538        let keys = ["first", "second"];
539        for _ in 0..burst {
540            limiter.await_keys_ready(Some(&keys)).await;
541        }
542
543        clock.advance(Duration::from_secs(1));
544        let now = clock.now();
545
546        for _ in 0..burst {
547            limiter.await_keys_ready(Some(&keys)).await;
548            assert_eq!(clock.now(), now);
549        }
550
551        assert_eq!(limiter.plan_keys(&["spare", "first"], now), Err(interval));
552        assert!(!limiter.state.contains_key(&"spare"));
553        limiter.await_keys_ready(Some(&["spare", "first"])).await;
554        assert_eq!(clock.now(), now + Nanos::from(interval));
555        assert_eq!(
556            limiter
557                .check_key(&"first")
558                .unwrap_err()
559                .wait_time_from(clock.now()),
560            interval
561        );
562
563        for _ in 1..burst {
564            assert_eq!(limiter.check_key(&"spare"), Ok(()));
565        }
566
567        assert_eq!(
568            limiter
569                .check_key(&"spare")
570                .unwrap_err()
571                .wait_time_from(clock.now()),
572            interval
573        );
574    }
575
576    #[rstest]
577    fn test_custom_key_quota() {
578        let mock_limiter = initialize_mock_rate_limiter();
579
580        // Add new key quota pair
581        mock_limiter.add_quota_for_key(
582            "custom".to_string(),
583            Quota::per_second(NonZeroU32::new(1).unwrap()).unwrap(),
584        );
585
586        // Check custom quota
587        assert!(mock_limiter.check_key(&"custom".to_string()).is_ok());
588        assert!(mock_limiter.check_key(&"custom".to_string()).is_err());
589
590        // Check that default quota still applies to other keys
591        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
592        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
593        assert!(mock_limiter.check_key(&"default".to_string()).is_err());
594    }
595
596    #[rstest]
597    fn test_multiple_keys() {
598        let mock_limiter = initialize_mock_rate_limiter();
599
600        mock_limiter.add_quota_for_key(
601            "key1".to_string(),
602            Quota::per_second(NonZeroU32::new(1).unwrap()).unwrap(),
603        );
604        mock_limiter.add_quota_for_key(
605            "key2".to_string(),
606            Quota::per_second(NonZeroU32::new(3).unwrap()).unwrap(),
607        );
608
609        // Test key1
610        assert!(mock_limiter.check_key(&"key1".to_string()).is_ok());
611        assert!(mock_limiter.check_key(&"key1".to_string()).is_err());
612
613        // Test key2
614        assert!(mock_limiter.check_key(&"key2".to_string()).is_ok());
615        assert!(mock_limiter.check_key(&"key2".to_string()).is_ok());
616        assert!(mock_limiter.check_key(&"key2".to_string()).is_ok());
617        assert!(mock_limiter.check_key(&"key2".to_string()).is_err());
618    }
619
620    #[rstest]
621    fn test_quota_reset() {
622        let mock_limiter = initialize_mock_rate_limiter();
623
624        // Exhaust quota
625        assert!(mock_limiter.check_key(&"reset".to_string()).is_ok());
626        assert!(mock_limiter.check_key(&"reset".to_string()).is_ok());
627        assert!(mock_limiter.check_key(&"reset".to_string()).is_err());
628
629        // Advance clock by less than a second
630        mock_limiter.advance_clock(Duration::from_millis(499));
631        assert!(mock_limiter.check_key(&"reset".to_string()).is_err());
632
633        // Advance clock to reset
634        mock_limiter.advance_clock(Duration::from_millis(501));
635        assert!(mock_limiter.check_key(&"reset".to_string()).is_ok());
636    }
637
638    #[rstest]
639    fn test_different_quotas() {
640        let mock_limiter = initialize_mock_rate_limiter();
641
642        mock_limiter.add_quota_for_key(
643            "per_second".to_string(),
644            Quota::per_second(NonZeroU32::new(2).unwrap()).unwrap(),
645        );
646        mock_limiter.add_quota_for_key(
647            "per_minute".to_string(),
648            Quota::per_minute(NonZeroU32::new(3).unwrap()),
649        );
650
651        // Test per_second quota
652        assert!(mock_limiter.check_key(&"per_second".to_string()).is_ok());
653        assert!(mock_limiter.check_key(&"per_second".to_string()).is_ok());
654        assert!(mock_limiter.check_key(&"per_second".to_string()).is_err());
655
656        // Test per_minute quota
657        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_ok());
658        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_ok());
659        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_ok());
660        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_err());
661
662        // Advance clock and check reset
663        mock_limiter.advance_clock(Duration::from_secs(1));
664        assert!(mock_limiter.check_key(&"per_second".to_string()).is_ok());
665        assert!(mock_limiter.check_key(&"per_minute".to_string()).is_err());
666    }
667
668    #[tokio::test]
669    async fn test_await_keys_ready() {
670        let mock_limiter = initialize_mock_rate_limiter();
671
672        // Check base quota is not exceeded
673        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
674        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
675
676        // Check base quota is exceeded
677        assert!(mock_limiter.check_key(&"default".to_string()).is_err());
678
679        // Wait keys to be ready and check base quota is reset
680        mock_limiter.advance_clock(Duration::from_secs(1));
681        let keys = ["default".to_string()];
682        mock_limiter.await_keys_ready(Some(keys.as_slice())).await;
683        assert!(mock_limiter.check_key(&"default".to_string()).is_ok());
684    }
685
686    #[cfg_attr(
687        not(all(feature = "simulation", madsim)),
688        tokio::test(start_paused = true)
689    )]
690    #[cfg_attr(all(feature = "simulation", madsim), madsim::test)]
691    async fn test_await_keys_ready_reserves_keys_together() {
692        let fast = "fast".to_string();
693        let slow = "slow".to_string();
694
695        let limiter = Arc::new(RateLimiter::new_with_quota(
696            None,
697            vec![
698                (
699                    fast.clone(),
700                    Quota::with_period(Duration::from_secs(1)).unwrap(),
701                ),
702                (
703                    slow.clone(),
704                    Quota::with_period(Duration::from_secs(10)).unwrap(),
705                ),
706            ],
707        ));
708        limiter.check_key(&slow).unwrap();
709
710        let waiting_limiter = Arc::clone(&limiter);
711        let waiting_fast = fast.clone();
712        let waiting_slow = slow.clone();
713
714        let request = test_task::spawn(async move {
715            waiting_limiter
716                .await_keys_ready(Some(&[waiting_fast, waiting_slow]))
717                .await;
718        });
719
720        test_task::yield_now().await;
721
722        limiter.check_key(&fast).unwrap();
723        assert!(!request.is_finished());
724
725        advance_test_clock(Duration::from_millis(9_999)).await;
726        limiter.until_key_ready(&fast).await;
727        advance_test_clock(Duration::from_millis(1)).await;
728        test_task::yield_now().await;
729        assert!(!request.is_finished());
730
731        advance_test_clock(Duration::from_millis(998)).await;
732        test_task::yield_now().await;
733        assert!(!request.is_finished());
734
735        advance_test_clock(Duration::from_millis(1)).await;
736        request.await.unwrap();
737
738        assert!(limiter.check_key(&fast).is_err());
739        assert!(limiter.check_key(&slow).is_err());
740    }
741
742    #[cfg(all(feature = "simulation", madsim))]
743    async fn advance_test_clock(duration: Duration) {
744        madsim::time::advance(duration);
745        test_task::yield_now().await;
746    }
747
748    #[cfg(not(all(feature = "simulation", madsim)))]
749    async fn advance_test_clock(duration: Duration) {
750        tokio::time::advance(duration).await;
751    }
752
753    #[rstest]
754    fn test_per_second_returns_none_on_zero_replenish_interval() {
755        assert!(Quota::per_second(NonZeroU32::new(u32::MAX).unwrap()).is_none());
756    }
757
758    #[rstest]
759    fn test_per_minute_accepts_max_burst() {
760        let quota = Quota::per_minute(NonZeroU32::new(u32::MAX).unwrap());
761        assert_eq!(quota.replenish_interval(), Duration::from_nanos(13));
762    }
763
764    #[rstest]
765    fn test_per_hour_accepts_max_burst() {
766        let quota = Quota::per_hour(NonZeroU32::new(u32::MAX).unwrap());
767        assert_eq!(quota.replenish_interval(), Duration::from_nanos(838));
768    }
769
770    mod property_tests {
771        use proptest::prelude::*;
772        use rstest::rstest;
773
774        use crate::ratelimiter::nanos::Nanos;
775
776        proptest! {
777            #![proptest_config(ProptestConfig {
778                failure_persistence: Some(Box::new(
779                    proptest::test_runner::FileFailurePersistence::WithSource("ratelimiter")
780                )),
781                ..ProptestConfig::default()
782            })]
783
784            // Operators must saturate across the full u64 domain (a wrapped TAT admits everything)
785            #[rstest]
786            fn nanos_operators_never_panic(a in proptest::num::u64::ANY, b in proptest::num::u64::ANY) {
787                let na = Nanos::from(a);
788                let nb = Nanos::from(b);
789
790                prop_assert_eq!((na + nb).as_u64(), a.saturating_add(b));
791                prop_assert_eq!((na * b).as_u64(), a.saturating_mul(b));
792                prop_assert_eq!(na.saturating_sub(nb).as_u64(), a.saturating_sub(b));
793            }
794        }
795    }
796
797    #[rstest]
798    fn test_gcra_boundary_exact_replenishment() {
799        // Test GCRA boundary condition where t0 equals earliest_time exactly.
800        // This exercises the saturating_sub edge case deterministically without sleeps.
801        let mock_limiter = initialize_mock_rate_limiter();
802        let key = "boundary_test".to_string();
803
804        assert!(mock_limiter.check_key(&key).is_ok());
805        assert!(mock_limiter.check_key(&key).is_ok());
806        assert!(mock_limiter.check_key(&key).is_err());
807
808        // Advance clock by exactly one replenish interval (500ms for 2 req/sec)
809        let quota = Quota::per_second(NonZeroU32::new(2).unwrap()).unwrap();
810        let replenish_interval = quota.replenish_interval();
811        mock_limiter.advance_clock(replenish_interval);
812
813        assert!(
814            mock_limiter.check_key(&key).is_ok(),
815            "Request at exact replenish boundary should be allowed"
816        );
817        assert!(
818            mock_limiter.check_key(&key).is_err(),
819            "Immediate follow-up should be rate-limited"
820        );
821    }
822
823    #[rstest]
824    fn test_per_second_boundary_exact_limit() {
825        // 1_000_000_000ns / 1_000_000_000 = 1ns per replenish, the exact boundary
826        let quota = Quota::per_second(NonZeroU32::new(1_000_000_000).unwrap()).unwrap();
827        assert_eq!(quota.replenish_interval().as_nanos(), 1);
828    }
829
830    #[rstest]
831    fn test_per_second_returns_none_above_one_billion() {
832        // 1_000_000_000ns / 1_000_000_001 rounds to 0ns
833        assert!(Quota::per_second(NonZeroU32::new(1_000_000_001).unwrap()).is_none());
834    }
835
836    #[rstest]
837    #[case::large(Duration::from_secs(100), u32::MAX, Duration::from_mins(7_158_278_825))]
838    #[case::saturated(Duration::MAX, 2, Duration::MAX)]
839    fn test_burst_size_replenished_in(
840        #[case] replenish_interval: Duration,
841        #[case] burst_size: u32,
842        #[case] expected: Duration,
843    ) {
844        let quota = Quota::with_period(replenish_interval)
845            .unwrap()
846            .allow_burst(NonZeroU32::new(burst_size).unwrap());
847
848        assert_eq!(quota.burst_size_replenished_in(), expected);
849    }
850
851    #[rstest]
852    #[should_panic(expected = "t cannot be zero")]
853    fn test_from_gcra_parameters_panics_on_zero_t() {
854        let _ = Quota::from_gcra_parameters(Nanos::from(0u64), Nanos::from(100u64));
855    }
856
857    #[rstest]
858    #[should_panic(expected = "tau/t results in zero burst capacity")]
859    fn test_from_gcra_parameters_panics_on_zero_division() {
860        // tau=1, t=2 → integer division yields 0
861        let _ = Quota::from_gcra_parameters(Nanos::from(2u64), Nanos::from(1u64));
862    }
863
864    #[rstest]
865    #[should_panic(expected = "tau/t exceeds u32::MAX")]
866    fn test_from_gcra_parameters_panics_on_overflow() {
867        let _ = Quota::from_gcra_parameters(Nanos::from(1u64), Nanos::from(u64::MAX));
868    }
869
870    #[rstest]
871    fn test_concurrent_check_key_respects_burst() {
872        let rate = 10u32;
873        let clock = FakeRelativeClock::default();
874        let start = clock.now();
875
876        let limiter = RateLimiter {
877            default_gcra: Some(Gcra::new(
878                Quota::per_second(NonZeroU32::new(rate).unwrap()).unwrap(),
879            )),
880            state: DashMapStateStore::new(),
881            gcra: DashMap::new(),
882            clock,
883            start,
884            decision_lock: parking_lot::Mutex::new(()),
885        };
886
887        let accepted = AtomicU32::new(0);
888        let num_threads = 50;
889
890        // Clock is frozen: no replenishment occurs
891        std::thread::scope(|s| {
892            for _ in 0..num_threads {
893                s.spawn(|| {
894                    if limiter.check_key(&"hot_key".to_string()).is_ok() {
895                        accepted.fetch_add(1, Ordering::Relaxed);
896                    }
897                });
898            }
899        });
900
901        let total = accepted.load(Ordering::Relaxed);
902        assert!(total >= 1, "At least one request should be accepted");
903        assert!(
904            total <= rate,
905            "Accepted {total} but burst capacity is {rate}"
906        );
907    }
908}
nautilus_network/ratelimiter/mod.rs

nautilus_network/ratelimiter/
mod.rs
