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// }