Skip to main content

nautilus_network/ratelimiter/
quota.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
16use std::{num::NonZeroU32, prelude::v1::*, time::Duration};
17
18use super::nanos::Nanos;
19
20/// A rate-limiting quota.
21///
22/// Quotas are expressed in a positive number of "cells" (the maximum number of positive decisions /
23/// allowed items until the rate limiter needs to replenish) and the amount of time for the rate
24/// limiter to replenish a single cell.
25///
26/// Neither the number of cells nor the replenishment unit of time may be zero.
27///
28/// # Burst Sizes
29/// There are multiple ways of expressing the same quota: a quota given as `Quota::per_second(1)`
30/// allows, on average, the same number of cells through as a quota given as `Quota::per_minute(60)`.
31/// The quota of `Quota::per_minute(60)` has a burst size of 60 cells, meaning it is
32/// possible to accommodate 60 cells in one go, after which the equivalent of a minute of inactivity
33/// is required for the burst allowance to be fully restored.
34///
35/// Burst size gets really important when you construct a rate limiter that should allow multiple
36/// elements through at one time (using [`RateLimiter.check_n`](struct.RateLimiter.html#method.check_n)
37/// and its related functions): Only
38/// at most as many cells can be let through in one call as are given as the burst size.
39///
40/// In other words, the burst size is the maximum number of cells that the rate limiter will ever
41/// allow through without replenishing them.
42#[derive(Debug, PartialEq, Eq, Clone, Copy)]
43#[cfg_attr(
44    feature = "python",
45    pyo3::pyclass(module = "nautilus_trader.core.nautilus_pyo3.network", from_py_object)
46)]
47pub struct Quota {
48    pub(crate) max_burst: NonZeroU32,
49    pub(crate) replenish_1_per: Duration,
50}
51
52/// Constructors for Quotas
53impl Quota {
54    /// Construct a quota for a number of cells per second. The given number of cells is also
55    /// assumed to be the maximum burst size.
56    #[must_use]
57    pub const fn per_second(max_burst: NonZeroU32) -> Self {
58        let replenish_interval_ns = Duration::from_secs(1).as_nanos() / (max_burst.get() as u128);
59        Self {
60            max_burst,
61            replenish_1_per: Duration::from_nanos(replenish_interval_ns as u64),
62        }
63    }
64
65    /// Construct a quota for a number of cells per 60-second period. The given number of cells is
66    /// also assumed to be the maximum burst size.
67    #[must_use]
68    pub const fn per_minute(max_burst: NonZeroU32) -> Self {
69        let replenish_interval_ns = Duration::from_secs(60).as_nanos() / (max_burst.get() as u128);
70        Self {
71            max_burst,
72            replenish_1_per: Duration::from_nanos(replenish_interval_ns as u64),
73        }
74    }
75
76    /// Construct a quota for a number of cells per 60-minute (3600-second) period. The given number
77    /// of cells is also assumed to be the maximum burst size.
78    #[must_use]
79    pub const fn per_hour(max_burst: NonZeroU32) -> Self {
80        let replenish_interval_ns = Duration::from_hours(1).as_nanos() / (max_burst.get() as u128);
81        Self {
82            max_burst,
83            replenish_1_per: Duration::from_nanos(replenish_interval_ns as u64),
84        }
85    }
86
87    /// Construct a quota that replenishes one cell in a given
88    /// interval.
89    ///
90    /// This constructor is meant to replace [`::new`](#method.new),
91    /// in cases where a longer refresh period than 1 cell/hour is
92    /// necessary.
93    ///
94    /// If the time interval is zero, returns `None`.
95    #[must_use]
96    pub const fn with_period(replenish_1_per: Duration) -> Option<Self> {
97        if replenish_1_per.as_nanos() == 0 {
98            None
99        } else {
100            // SAFETY: Unwrap is safe because 1 is always non-zero
101            #[allow(clippy::missing_panics_doc)]
102            Some(Self {
103                max_burst: NonZeroU32::new(1).unwrap(),
104                replenish_1_per,
105            })
106        }
107    }
108
109    /// Adjusts the maximum burst size for a quota to construct a rate limiter with a capacity
110    /// for at most the given number of cells.
111    #[must_use]
112    pub const fn allow_burst(self, max_burst: NonZeroU32) -> Self {
113        Self { max_burst, ..self }
114    }
115
116    /// Construct a quota for a given burst size, replenishing the entire burst size in that
117    /// given unit of time.
118    ///
119    /// Returns `None` if the duration is zero.
120    ///
121    /// This constructor allows greater control over the resulting
122    /// quota, but doesn't make as much intuitive sense as other
123    /// methods of constructing the same quotas. Unless your quotas
124    /// are given as "max burst size, and time it takes to replenish
125    /// that burst size", you are better served by the
126    /// [`Quota::per_second`](#method.per_second) (and similar)
127    /// constructors with the [`allow_burst`](#method.allow_burst)
128    /// modifier.
129    #[deprecated(
130        since = "0.2.0",
131        note = "This constructor is often confusing and non-intuitive. \
132    Use the `per_(interval)` / `with_period` and `max_burst` constructors instead."
133    )]
134    #[must_use]
135    pub fn new(max_burst: NonZeroU32, replenish_all_per: Duration) -> Option<Self> {
136        if replenish_all_per.as_nanos() == 0 {
137            None
138        } else {
139            Some(Self {
140                max_burst,
141                replenish_1_per: replenish_all_per / max_burst.get(),
142            })
143        }
144    }
145}
146
147/// Retrieving information about a quota
148impl Quota {
149    /// The time it takes for a rate limiter with an exhausted burst budget to replenish
150    /// a single element.
151    #[must_use]
152    pub const fn replenish_interval(&self) -> Duration {
153        self.replenish_1_per
154    }
155
156    /// The maximum number of cells that can be allowed in one burst.
157    #[must_use]
158    pub const fn burst_size(&self) -> NonZeroU32 {
159        self.max_burst
160    }
161
162    /// The time it takes to replenish the entire maximum burst size.
163    #[must_use]
164    pub const fn burst_size_replenished_in(&self) -> Duration {
165        let fill_in_ns = self.replenish_1_per.as_nanos() * self.max_burst.get() as u128;
166        Duration::from_nanos(fill_in_ns as u64)
167    }
168}
169
170impl Quota {
171    /// A way to reconstruct a Quota from an in-use Gcra.
172    ///
173    /// This is useful mainly for [`crate::middleware::RateLimitingMiddleware`]
174    /// where custom code may want to construct information based on
175    /// the amount of burst balance remaining.
176    ///
177    /// # Panics
178    ///
179    /// Panics if the division result is 0 or exceeds `u32::MAX`.
180    pub(crate) fn from_gcra_parameters(t: Nanos, tau: Nanos) -> Self {
181        let t_u64 = t.as_u64();
182        let tau_u64 = tau.as_u64();
183
184        // Validate division won't be zero or overflow
185        assert!(t_u64 != 0, "Invalid GCRA parameter: t cannot be zero");
186
187        let division_result = tau_u64 / t_u64;
188        assert!(
189            division_result != 0,
190            "Invalid GCRA parameters: tau/t results in zero burst capacity"
191        );
192        assert!(
193            u32::try_from(division_result).is_ok(),
194            "Invalid GCRA parameters: tau/t exceeds u32::MAX"
195        );
196
197        // We've verified the result is non-zero and fits in u32
198        let max_burst = NonZeroU32::new(division_result as u32)
199            .expect("Division result should be non-zero after validation");
200        let replenish_1_per = t.into();
201        Self {
202            max_burst,
203            replenish_1_per,
204        }
205    }
206}
207
208// #[cfg(test)]
209// mod test {
210//     use nonzero_ext::nonzero;
211
212//     use super::*;
213//     use rstest::rstest;
214
215//     #[rstest]
216//     fn time_multiples() {
217//         let hourly = Quota::per_hour(nonzero!(1u32));
218//         let minutely = Quota::per_minute(nonzero!(1u32));
219//         let secondly = Quota::per_second(nonzero!(1u32));
220
221//         assert_eq!(
222//             hourly.replenish_interval() / 60,
223//             minutely.replenish_interval()
224//         );
225//         assert_eq!(
226//             minutely.replenish_interval() / 60,
227//             secondly.replenish_interval()
228//         );
229//     }
230
231//     #[rstest]
232//     fn period_error_cases() {
233//         assert!(Quota::with_period(Duration::from_secs(0)).is_none());
234
235//         #[allow(deprecated)]
236//         {
237//             assert!(Quota::new(nonzero!(1u32), Duration::from_secs(0)).is_none());
238//         }
239//     }
240// }