Skip to main content

rate_net/
quota.rs

1//! How much a key may do, and how fast it recovers.
2
3use core::time::Duration;
4
5use crate::error::RateLimiterError;
6
7/// A rate limit: `limit` requests per `period`, per key, with a `burst` ceiling.
8///
9/// A quota describes the sustained rate a key is allowed and how much it may
10/// spend at once. Under the default token-bucket algorithm each key starts with
11/// a full allowance of `burst`, spends one unit per admitted request, and
12/// accrues `limit` units back over `period` — so a key may burst up to `burst`
13/// immediately and then sustain `limit` per `period` thereafter. `burst`
14/// defaults to `limit` (the classic "burst equals rate" bucket); raise it with
15/// [`with_burst`](Self::with_burst) to allow larger spikes, or lower it to shape
16/// traffic more tightly.
17///
18/// The convenience constructors [`per_second`](Self::per_second) and
19/// [`per_minute`](Self::per_minute) are infallible (a `limit` of `0` yields a
20/// quota that admits nothing). The general [`rate`](Self::rate) constructor
21/// validates its inputs and returns a [`RateLimiterError`] for values that
22/// cannot describe a working limit.
23///
24/// The window algorithms (fixed and sliding window) admit at most `limit` per
25/// `period` and ignore `burst`; it applies to the token and leaky buckets.
26///
27/// # Examples
28///
29/// ```
30/// use rate_net::Quota;
31/// use std::time::Duration;
32///
33/// let per_sec = Quota::per_second(100);
34/// assert_eq!(per_sec.limit(), 100);
35/// assert_eq!(per_sec.period(), Duration::from_secs(1));
36/// assert_eq!(per_sec.burst(), 100); // defaults to the limit
37///
38/// // 1000 requests per minute, but bursts capped at 50.
39/// let shaped = Quota::rate(1000, Duration::from_secs(60))?.with_burst(50);
40/// assert_eq!(shaped.limit(), 1000);
41/// assert_eq!(shaped.burst(), 50);
42/// # Ok::<(), rate_net::RateLimiterError>(())
43/// ```
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
45pub struct Quota {
46    limit: u32,
47    period: Duration,
48    burst: u32,
49}
50
51impl Quota {
52    /// A quota of `limit` requests per second, per key.
53    ///
54    /// Infallible: a `limit` of `0` produces a quota that admits nothing, which
55    /// is well-defined. Use [`rate`](Self::rate) when you want a zero limit
56    /// rejected as an error.
57    ///
58    /// # Examples
59    ///
60    /// ```
61    /// use rate_net::Quota;
62    /// use std::time::Duration;
63    ///
64    /// let quota = Quota::per_second(50);
65    /// assert_eq!(quota.limit(), 50);
66    /// assert_eq!(quota.period(), Duration::from_secs(1));
67    /// ```
68    #[must_use]
69    pub const fn per_second(limit: u32) -> Self {
70        Self {
71            limit,
72            period: Duration::from_secs(1),
73            burst: limit,
74        }
75    }
76
77    /// A quota of `limit` requests per minute, per key.
78    ///
79    /// Infallible, with the same zero-limit semantics as
80    /// [`per_second`](Self::per_second).
81    ///
82    /// # Examples
83    ///
84    /// ```
85    /// use rate_net::Quota;
86    /// use std::time::Duration;
87    ///
88    /// let quota = Quota::per_minute(600);
89    /// assert_eq!(quota.period(), Duration::from_secs(60));
90    /// ```
91    #[must_use]
92    pub const fn per_minute(limit: u32) -> Self {
93        Self {
94            limit,
95            period: Duration::from_secs(60),
96            burst: limit,
97        }
98    }
99
100    /// A quota of `limit` requests per arbitrary `period`, validated.
101    ///
102    /// Use this when the natural window is neither a second nor a minute — for
103    /// example 5 requests per 100 milliseconds, or 10 000 per hour.
104    ///
105    /// # Errors
106    ///
107    /// - [`RateLimiterError::ZeroQuota`] if `limit` is `0`.
108    /// - [`RateLimiterError::ZeroPeriod`] if `period` is zero.
109    ///
110    /// # Examples
111    ///
112    /// ```
113    /// use rate_net::{Quota, RateLimiterError};
114    /// use std::time::Duration;
115    ///
116    /// let quota = Quota::rate(5, Duration::from_millis(100))?;
117    /// assert_eq!(quota.limit(), 5);
118    ///
119    /// // A zero limit is rejected.
120    /// assert_eq!(
121    ///     Quota::rate(0, Duration::from_secs(1)),
122    ///     Err(RateLimiterError::ZeroQuota),
123    /// );
124    /// # Ok::<(), RateLimiterError>(())
125    /// ```
126    pub const fn rate(limit: u32, period: Duration) -> Result<Self, RateLimiterError> {
127        if limit == 0 {
128            return Err(RateLimiterError::ZeroQuota);
129        }
130        if period.is_zero() {
131            return Err(RateLimiterError::ZeroPeriod);
132        }
133        Ok(Self {
134            limit,
135            period,
136            burst: limit,
137        })
138    }
139
140    /// The number of requests admitted per [`period`](Self::period).
141    ///
142    /// # Examples
143    ///
144    /// ```
145    /// use rate_net::Quota;
146    ///
147    /// assert_eq!(Quota::per_second(100).limit(), 100);
148    /// ```
149    #[must_use]
150    pub const fn limit(&self) -> u32 {
151        self.limit
152    }
153
154    /// The window over which [`limit`](Self::limit) requests accrue.
155    ///
156    /// # Examples
157    ///
158    /// ```
159    /// use rate_net::Quota;
160    /// use std::time::Duration;
161    ///
162    /// assert_eq!(Quota::per_minute(60).period(), Duration::from_secs(60));
163    /// ```
164    #[must_use]
165    pub const fn period(&self) -> Duration {
166        self.period
167    }
168
169    /// The burst ceiling: the most a key may spend at once before it must wait
170    /// for the rate to refill. Defaults to [`limit`](Self::limit).
171    ///
172    /// Applies to the token and leaky buckets; the window algorithms admit at
173    /// most `limit` per `period` regardless of `burst`.
174    ///
175    /// # Examples
176    ///
177    /// ```
178    /// use rate_net::Quota;
179    ///
180    /// assert_eq!(Quota::per_second(100).burst(), 100);
181    /// assert_eq!(Quota::per_second(100).with_burst(250).burst(), 250);
182    /// ```
183    #[must_use]
184    pub const fn burst(&self) -> u32 {
185        self.burst
186    }
187
188    /// Returns a copy with the burst ceiling set to `burst`.
189    ///
190    /// # Examples
191    ///
192    /// ```
193    /// use rate_net::Quota;
194    ///
195    /// // Sustain 1000/min but never let a key spend more than 50 at once.
196    /// let quota = Quota::per_minute(1000).with_burst(50);
197    /// assert_eq!(quota.limit(), 1000);
198    /// assert_eq!(quota.burst(), 50);
199    /// ```
200    #[must_use]
201    pub const fn with_burst(mut self, burst: u32) -> Self {
202        self.burst = burst;
203        self
204    }
205
206    /// Assembles a quota from raw parts without validation, for the infallible
207    /// [`Builder`](crate::Builder) path. A zero `limit` admits nothing; a zero
208    /// `period` yields a degenerate limiter that never refills.
209    pub(crate) const fn from_parts(limit: u32, period: Duration, burst: u32) -> Self {
210        Self {
211            limit,
212            period,
213            burst,
214        }
215    }
216}
217
218#[cfg(test)]
219mod tests {
220    #![allow(clippy::unwrap_used)]
221
222    use super::Quota;
223    use crate::error::RateLimiterError;
224    use core::time::Duration;
225
226    #[test]
227    fn test_per_second_sets_one_second_period() {
228        let quota = Quota::per_second(10);
229        assert_eq!(quota.limit(), 10);
230        assert_eq!(quota.period(), Duration::from_secs(1));
231    }
232
233    #[test]
234    fn test_per_minute_sets_sixty_second_period() {
235        assert_eq!(Quota::per_minute(10).period(), Duration::from_secs(60));
236    }
237
238    #[test]
239    fn test_burst_defaults_to_limit_and_overrides() {
240        assert_eq!(Quota::per_second(10).burst(), 10);
241        assert_eq!(Quota::per_minute(10).burst(), 10);
242        let q = Quota::rate(10, Duration::from_secs(1)).unwrap();
243        assert_eq!(q.burst(), 10);
244        assert_eq!(q.with_burst(25).burst(), 25);
245        // Overriding burst leaves the sustained limit unchanged.
246        assert_eq!(q.with_burst(25).limit(), 10);
247    }
248
249    #[test]
250    fn test_rate_accepts_valid_values() {
251        let quota = Quota::rate(5, Duration::from_millis(100)).unwrap();
252        assert_eq!(quota.limit(), 5);
253        assert_eq!(quota.period(), Duration::from_millis(100));
254    }
255
256    #[test]
257    fn test_rate_rejects_zero_limit() {
258        assert_eq!(
259            Quota::rate(0, Duration::from_secs(1)),
260            Err(RateLimiterError::ZeroQuota)
261        );
262    }
263
264    #[test]
265    fn test_rate_rejects_zero_period() {
266        assert_eq!(
267            Quota::rate(10, Duration::ZERO),
268            Err(RateLimiterError::ZeroPeriod)
269        );
270    }
271
272    #[test]
273    fn test_per_second_zero_limit_is_allowed() {
274        // Infallible constructors accept zero (admits nothing), unlike `rate`.
275        assert_eq!(Quota::per_second(0).limit(), 0);
276    }
277}