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)]
47#[cfg_attr(
48 feature = "python",
49 pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.network")
50)]
51pub struct Quota {
52 pub(crate) max_burst: NonZeroU32,
53 pub(crate) replenish_1_per: Duration,
54}
55
56/// Constructors for Quotas
57impl Quota {
58 /// Construct a quota for a number of cells per second. The given number of cells is also
59 /// assumed to be the maximum burst size.
60 ///
61 /// Returns `None` if `max_burst` is so large that the replenish interval rounds to zero
62 /// nanoseconds (i.e. `max_burst > 1_000_000_000`).
63 #[must_use]
64 pub const fn per_second(max_burst: NonZeroU32) -> Option<Self> {
65 let replenish_interval_ns = Duration::from_secs(1).as_nanos() / (max_burst.get() as u128);
66 if replenish_interval_ns == 0 {
67 return None;
68 }
69 Some(Self {
70 max_burst,
71 replenish_1_per: Duration::from_nanos(replenish_interval_ns as u64),
72 })
73 }
74
75 /// Construct a quota for a number of cells per 60-second period. The given number of cells is
76 /// also assumed to be the maximum burst size.
77 #[must_use]
78 pub const fn per_minute(max_burst: NonZeroU32) -> Self {
79 let replenish_interval_ns = Duration::from_mins(1).as_nanos() / (max_burst.get() as u128);
80 Self {
81 max_burst,
82 replenish_1_per: Duration::from_nanos(replenish_interval_ns as u64),
83 }
84 }
85
86 /// Construct a quota for a number of cells per 60-minute (3600-second) period. The given number
87 /// of cells is also assumed to be the maximum burst size.
88 #[must_use]
89 pub const fn per_hour(max_burst: NonZeroU32) -> Self {
90 let replenish_interval_ns = Duration::from_hours(1).as_nanos() / (max_burst.get() as u128);
91 Self {
92 max_burst,
93 replenish_1_per: Duration::from_nanos(replenish_interval_ns as u64),
94 }
95 }
96
97 /// Construct a quota that replenishes one cell in a given interval.
98 ///
99 /// If the time interval is zero, returns `None`.
100 #[must_use]
101 pub const fn with_period(replenish_1_per: Duration) -> Option<Self> {
102 if replenish_1_per.as_nanos() == 0 {
103 None
104 } else {
105 Some(Self {
106 max_burst: NonZeroU32::MIN,
107 replenish_1_per,
108 })
109 }
110 }
111
112 /// Adjusts the maximum burst size for a quota to construct a rate limiter with a capacity
113 /// for at most the given number of cells.
114 #[must_use]
115 pub const fn allow_burst(self, max_burst: NonZeroU32) -> Self {
116 Self { max_burst, ..self }
117 }
118}
119
120/// Retrieving information about a quota
121impl Quota {
122 /// The time it takes for a rate limiter with an exhausted burst budget to replenish
123 /// a single element.
124 #[must_use]
125 pub const fn replenish_interval(&self) -> Duration {
126 self.replenish_1_per
127 }
128
129 /// The maximum number of cells that can be allowed in one burst.
130 #[must_use]
131 pub const fn burst_size(&self) -> NonZeroU32 {
132 self.max_burst
133 }
134
135 /// The time it takes to replenish the entire maximum burst size.
136 #[must_use]
137 pub const fn burst_size_replenished_in(&self) -> Duration {
138 let fill_in_ns = self.replenish_1_per.as_nanos() * self.max_burst.get() as u128;
139 Duration::from_nanos(fill_in_ns as u64)
140 }
141}
142
143impl Quota {
144 /// A way to reconstruct a Quota from an in-use Gcra.
145 ///
146 /// This is useful mainly for [`crate::middleware::RateLimitingMiddleware`]
147 /// where custom code may want to construct information based on
148 /// the amount of burst balance remaining.
149 ///
150 /// # Panics
151 ///
152 /// Panics if the division result is 0 or exceeds `u32::MAX`.
153 pub(crate) fn from_gcra_parameters(t: Nanos, tau: Nanos) -> Self {
154 let t_u64 = t.as_u64();
155 let tau_u64 = tau.as_u64();
156
157 // Validate division won't be zero or overflow
158 assert!(t_u64 != 0, "Invalid GCRA parameter: t cannot be zero");
159
160 let division_result = tau_u64 / t_u64;
161 assert!(
162 division_result != 0,
163 "Invalid GCRA parameters: tau/t results in zero burst capacity"
164 );
165 assert!(
166 u32::try_from(division_result).is_ok(),
167 "Invalid GCRA parameters: tau/t exceeds u32::MAX"
168 );
169
170 // We've verified the result is non-zero and fits in u32
171 let max_burst = NonZeroU32::new(division_result as u32)
172 .expect("Division result should be non-zero after validation");
173 let replenish_1_per = t.into();
174 Self {
175 max_burst,
176 replenish_1_per,
177 }
178 }
179}
180
181// #[cfg(test)]
182// mod test {
183// use nonzero_ext::nonzero;
184
185// use super::*;
186// use rstest::rstest;
187
188// #[rstest]
189// fn time_multiples() {
190// let hourly = Quota::per_hour(nonzero!(1u32));
191// let minutely = Quota::per_minute(nonzero!(1u32));
192// let secondly = Quota::per_second(nonzero!(1u32));
193
194// assert_eq!(
195// hourly.replenish_interval() / 60,
196// minutely.replenish_interval()
197// );
198// assert_eq!(
199// minutely.replenish_interval() / 60,
200// secondly.replenish_interval()
201// );
202// }
203
204// #[rstest]
205// fn period_error_cases() {
206// assert!(Quota::with_period(Duration::from_secs(0)).is_none());
207// }
208// }