nautilus_network/ratelimiter/
gcra.rs

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

nautilus_network/ratelimiter/
gcra.rs
