nautilus_network/ratelimiter/
quota.rs

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