rama_utils/rate/
mod.rs

1//! Rate primitives: express "N units per period" and enforce it.
2//!
3//! The core of this module is sans-IO and `no_std`: [`TokenBucket`] never
4//! reads a clock, the caller provides time as monotonic nanoseconds. This
5//! keeps it trivially testable and embeddable in other sans-IO state
6//! machines (e.g. a send scheduler pacing datagrams).
7//!
8//! On top (behind the `std` feature) [`RateLimiter`] wraps a [`TokenBucket`]
9//! in a cheap-to-clone shared handle that can wait (pace) or reject
10//! (rate limit) using tokio timers.
11//!
12//! Units are whatever you want them to be: requests, connections,
13//! messages, bytes.
14
15mod bucket;
16#[doc(inline)]
17pub use bucket::TokenBucket;
18
19#[cfg(feature = "std")]
20mod limiter;
21#[cfg(feature = "std")]
22#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
23#[doc(inline)]
24pub use limiter::{RateLimiter, RefundWait};
25
26use core::time::Duration;
27
28/// A rate: `units` per `per` period.
29///
30/// Units are dimensionless: use it for requests, connections, messages
31/// or bytes alike. For byte rates [`Rate::per_sec`] *is* bytes/sec;
32/// compose with [`crate::octets`] helpers for larger units.
33#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
34pub struct Rate {
35    units: u64,
36    per: Duration,
37}
38
39impl Rate {
40    /// Create a new [`Rate`] of `units` per `per` period.
41    ///
42    /// # Panics
43    ///
44    /// Panics if `units` is zero, `per` is zero, or `per` exceeds
45    /// `u64::MAX` nanoseconds (~584 years). In const contexts this is
46    /// a compile-time error. Use [`Rate::try_new`] for fallible
47    /// construction from dynamic input.
48    ///
49    /// # Examples
50    ///
51    /// ```
52    /// use std::time::Duration;
53    /// use rama_utils::rate::Rate;
54    ///
55    /// const RATE: Rate = Rate::new(100, Duration::from_millis(250));
56    /// assert_eq!(RATE.units(), 100);
57    /// assert_eq!(RATE.per(), Duration::from_millis(250));
58    /// ```
59    #[must_use]
60    pub const fn new(units: u64, per: Duration) -> Self {
61        assert!(units != 0, "Rate: units must be non-zero");
62        assert!(
63            per.as_nanos() != 0 && per.as_nanos() <= u64::MAX as u128,
64            "Rate: per must be non-zero and at most u64::MAX nanos"
65        );
66        Self { units, per }
67    }
68
69    /// Create a new [`Rate`] of `units` per `per` period,
70    /// or `None` for the invalid inputs that make [`Rate::new`] panic.
71    #[must_use]
72    pub const fn try_new(units: u64, per: Duration) -> Option<Self> {
73        if units == 0 || per.as_nanos() == 0 || per.as_nanos() > u64::MAX as u128 {
74            return None;
75        }
76        Some(Self { units, per })
77    }
78
79    /// Create a new [`Rate`] of `units` per second.
80    ///
81    /// For byte rates this *is* bytes per second; combine with
82    /// [`crate::octets::kib_u64`], [`crate::octets::mib_u64`] and
83    /// [`crate::octets::gib_u64`] for KiB/MiB/GiB rates, and with
84    /// [`crate::octets::from_bits_u64`] for link rates expressed in bits.
85    ///
86    /// # Panics
87    ///
88    /// Panics if `units` is zero (compile-time error in const contexts).
89    ///
90    /// # Examples
91    ///
92    /// ```
93    /// use rama_utils::{octets, rate::Rate};
94    ///
95    /// // 100 requests per second
96    /// const REQ_RATE: Rate = Rate::per_sec(100);
97    /// # let _ = REQ_RATE;
98    /// // 10 MiB per second
99    /// const BYTE_RATE: Rate = Rate::per_sec(octets::mib_u64(10));
100    /// # let _ = BYTE_RATE;
101    /// // a 100 Mbit/s link rate
102    /// const LINK_RATE: Rate = Rate::per_sec(octets::from_bits_u64(100_000_000));
103    /// # let _ = LINK_RATE;
104    /// ```
105    #[must_use]
106    pub const fn per_sec(units: u64) -> Self {
107        Self::new(units, Duration::from_secs(1))
108    }
109
110    /// The number of units per [`Rate::per`] period.
111    #[must_use]
112    pub const fn units(&self) -> u64 {
113        self.units
114    }
115
116    /// The period over which [`Rate::units`] units are allowed.
117    #[must_use]
118    pub const fn per(&self) -> Duration {
119        self.per
120    }
121
122    pub(crate) const fn per_nanos(&self) -> u64 {
123        // validated by the constructors to fit
124        self.per.as_nanos() as u64
125    }
126}
127
128/// The result of [`TokenBucket::try_acquire`] (and
129/// [`RateLimiter::try_acquire`]).
130#[derive(Debug, Clone, Copy, PartialEq, Eq)]
131pub enum Acquire {
132    /// The requested units were spent; the caller may proceed.
133    Granted,
134    /// Not enough budget: the earliest monotonic instant (in the caller's
135    /// clock, nanoseconds) at which the requested units can be granted.
136    RetryAt(u64),
137    /// The request exceeds the burst capacity and can never be granted
138    /// as a single acquisition. Split it up (e.g. in burst-sized chunks,
139    /// as [`RateLimiter::acquire`] does) or reject it.
140    Never,
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146
147    #[test]
148    fn rate_valid() {
149        let rate = Rate::per_sec(10);
150        assert_eq!(rate.units(), 10);
151        assert_eq!(rate.per(), Duration::from_secs(1));
152        assert_eq!(rate.per_nanos(), 1_000_000_000);
153    }
154
155    #[test]
156    fn rate_try_new_invalid() {
157        assert!(Rate::try_new(0, Duration::from_secs(1)).is_none());
158        assert!(Rate::try_new(1, Duration::ZERO).is_none());
159        assert!(Rate::try_new(1, Duration::MAX).is_none());
160        assert!(Rate::try_new(1, Duration::from_nanos(1)).is_some());
161    }
162
163    #[test]
164    #[should_panic(expected = "non-zero")]
165    fn rate_new_zero_units_panics() {
166        assert_eq!(Rate::new(0, Duration::from_secs(1)).units(), 0);
167    }
168}
rama_utils/rate/mod.rs

rama_utils/rate/
mod.rs
