nautilus_network/ratelimiter/
gcra.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2025 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::{cmp, fmt::Display, time::Duration};
17
18use super::{clock, nanos::Nanos, quota::Quota, StateStore};
19
20/// Information about the rate-limiting state used to reach a decision.
21#[derive(Clone, PartialEq, Eq, Debug)]
22pub struct StateSnapshot {
23    /// The "weight" of a single packet in units of time.
24    t: Nanos,
25
26    /// The "burst capacity" of the bucket.
27    tau: Nanos,
28
29    /// The time at which the measurement was taken.
30    pub(crate) time_of_measurement: Nanos,
31
32    /// The next time a cell is expected to arrive
33    pub(crate) tat: Nanos,
34}
35
36impl StateSnapshot {
37    /// Creates a new [`StateSnapshot`] instance.
38    #[inline]
39    pub(crate) const fn new(t: Nanos, tau: Nanos, time_of_measurement: Nanos, tat: Nanos) -> Self {
40        Self {
41            t,
42            tau,
43            time_of_measurement,
44            tat,
45        }
46    }
47
48    /// Returns the quota used to make the rate limiting decision.
49    pub fn quota(&self) -> Quota {
50        Quota::from_gcra_parameters(self.t, self.tau)
51    }
52
53    /// Returns the number of cells that can be let through in
54    /// addition to a (possible) positive outcome.
55    ///
56    /// If this state snapshot is based on a negative rate limiting
57    /// outcome, this method returns 0.
58    pub fn remaining_burst_capacity(&self) -> u32 {
59        let t0 = self.time_of_measurement + self.t;
60        (cmp::min(
61            (t0 + self.tau).saturating_sub(self.tat).as_u64(),
62            self.tau.as_u64(),
63        ) / self.t.as_u64()) as u32
64    }
65}
66
67/// A negative rate-limiting outcome.
68///
69/// `NotUntil`'s methods indicate when a caller can expect the next positive
70/// rate-limiting result.
71#[derive(Debug, PartialEq, Eq)]
72pub struct NotUntil<P: clock::Reference> {
73    state: StateSnapshot,
74    start: P,
75}
76
77impl<P: clock::Reference> NotUntil<P> {
78    /// Create a `NotUntil` as a negative rate-limiting result.
79    #[inline]
80    pub(crate) const fn new(state: StateSnapshot, start: P) -> Self {
81        Self { state, start }
82    }
83
84    /// Returns the earliest time at which a decision could be
85    /// conforming (excluding conforming decisions made by the Decider
86    /// that are made in the meantime).
87    #[inline]
88    pub fn earliest_possible(&self) -> P {
89        let tat: Nanos = self.state.tat;
90        self.start + tat
91    }
92
93    /// Returns the minimum amount of time from the time that the
94    /// decision was made that must pass before a
95    /// decision can be conforming.
96    ///
97    /// If the time of the next expected positive result is in the past,
98    /// `wait_time_from` returns a zero `Duration`.
99    #[inline]
100    pub fn wait_time_from(&self, from: P) -> Duration {
101        let earliest = self.earliest_possible();
102        earliest.duration_since(earliest.min(from)).into()
103    }
104
105    /// Returns the rate limiting [`Quota`] used to reach the decision.
106    #[inline]
107    pub fn quota(&self) -> Quota {
108        self.state.quota()
109    }
110}
111
112impl<P: clock::Reference> Display for NotUntil<P> {
113    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
114        write!(f, "rate-limited until {:?}", self.start + self.state.tat)
115    }
116}
117
118#[derive(Debug, PartialEq, Eq)]
119pub struct Gcra {
120    /// The "weight" of a single packet in units of time.
121    t: Nanos,
122
123    /// The "burst capacity" of the bucket.
124    tau: Nanos,
125}
126
127impl Gcra {
128    pub(crate) fn new(quota: Quota) -> Self {
129        let tau: Nanos = (quota.replenish_1_per * quota.max_burst.get()).into();
130        let t: Nanos = quota.replenish_1_per.into();
131        Self { t, tau }
132    }
133
134    /// Computes and returns a new ratelimiter state if none exists yet.
135    fn starting_state(&self, t0: Nanos) -> Nanos {
136        t0 + self.t
137    }
138
139    /// Tests a single cell against the rate limiter state and updates it at the given key.
140    pub(crate) fn test_and_update<K, S: StateStore<Key = K>, P: clock::Reference>(
141        &self,
142        start: P,
143        key: &K,
144        state: &S,
145        t0: P,
146    ) -> Result<(), NotUntil<P>> {
147        let t0 = t0.duration_since(start);
148        let tau = self.tau;
149        let t = self.t;
150        state.measure_and_replace(key, |tat| {
151            let tat = tat.unwrap_or_else(|| self.starting_state(t0));
152            let earliest_time = tat.saturating_sub(tau);
153            if t0 < earliest_time {
154                Err(NotUntil::new(
155                    StateSnapshot::new(self.t, self.tau, earliest_time, earliest_time),
156                    start,
157                ))
158            } else {
159                let next = cmp::max(tat, t0) + t;
160                Ok(((), next))
161            }
162        })
163    }
164}
nautilus_network/ratelimiter/gcra.rs

nautilus_network/ratelimiter/
gcra.rs
