+ ShallowCopy,
149 | I: clock::Reference,
150 | {
151 | }
152 | }
153 |
154 | #[cfg(feature = "std")]
155 | pub use self::std::*;
156 |
--------------------------------------------------------------------------------
/src/algorithms/gcra.rs:
--------------------------------------------------------------------------------
1 | //! The Generic Cell Rate Algorithm
2 |
3 | use crate::lib::*;
4 |
5 | use crate::{
6 | algorithms::{Algorithm, NonConformance, RateLimitState},
7 | clock,
8 | thread_safety::ThreadsafeWrapper,
9 | InconsistentCapacity, NegativeMultiDecision,
10 | };
11 |
12 | #[cfg(feature = "std")]
13 | mod std {
14 | use crate::clock;
15 | use evmap::ShallowCopy;
16 |
17 | impl ShallowCopy for super::State {
18 | unsafe fn shallow_copy(&mut self) -> Self {
19 | super::State(self.0.shallow_copy())
20 | }
21 | }
22 | }
23 |
24 | #[derive(Debug, Eq, PartialEq, Clone)]
25 | struct Tat(Option);
26 |
27 | impl Default for Tat {
28 | fn default() -> Self {
29 | Tat(None)
30 | }
31 | }
32 |
33 | /// The GCRA's state about a single rate limiting history.
34 | #[derive(Debug, Eq, PartialEq, Clone)]
35 | pub struct State(ThreadsafeWrapper>);
36 |
37 | impl Default for State {
38 | fn default() -> Self {
39 | State(Default::default())
40 | }
41 | }
42 |
43 | impl RateLimitState, P> for State {
44 | fn last_touched(&self, params: &GCRA
) -> Option
{
45 | let data = self.0.snapshot();
46 | Some(data.0? + params.tau)
47 | }
48 | }
49 |
50 | /// Returned in case of a negative rate-limiting decision. Indicates
51 | /// the earliest instant that a cell might get accepted again.
52 | ///
53 | /// To avoid thundering herd effects, client code should always add a
54 | /// random amount of jitter to wait time estimates.
55 | #[derive(Debug, PartialEq)]
56 | pub struct NotUntil(P);
57 |
58 | impl fmt::Display for NotUntil {
59 | fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
60 | write!(f, "rate-limited until {:?}", self.0)
61 | }
62 | }
63 |
64 | impl NonConformance for NotUntil
{
65 | #[inline]
66 | fn earliest_possible(&self) -> P {
67 | self.0
68 | }
69 | }
70 |
71 | /// Implements the virtual scheduling description of the Generic Cell
72 | /// Rate Algorithm, attributed to ITU-T in recommendation I.371
73 | /// Traffic control and congestion control in B-ISDN; from
74 | /// [Wikipedia](https://en.wikipedia.org/wiki/Generic_cell_rate_algorithm).
75 | ///
76 | ///
77 | /// While algorithms like leaky-bucket rate limiters allow cells to be
78 | /// distributed across time in any way, GCRA is a rate-limiting *and*
79 | /// traffic-shaping algorithm. It mandates that a minimum amount of
80 | /// time passes between cells being measured. For example, if your API
81 | /// mandates that only 20 requests can be made per second, GCRA will
82 | /// ensure that each request is at least 50ms apart from the previous
83 | /// request. This makes GCRA suitable for shaping traffic in
84 | /// networking and telecom equipment (it was initially made for
85 | /// asynchronous transfer mode networks), or for outgoing workloads on
86 | /// *consumers* of attention, e.g. distributing outgoing emails across
87 | /// a day.
88 | ///
89 | /// # A note about batch decisions
90 | /// In a blatant side-stepping of the above traffic-shaping criteria,
91 | /// this implementation of GCRA comes with an extension that allows
92 | /// measuring multiple cells at once, assuming that if a pause of
93 | /// `n*(the minimum time between cells)` has passed, we can allow a
94 | /// single big batch of `n` cells through. This assumption may not be
95 | /// correct for your application, but if you depend on GCRA's
96 | /// traffic-shaping properties, it's better to not use the `_n`
97 | /// suffixed check functions.
98 | ///
99 | /// # Example
100 | /// In this example, we construct a rate-limiter with the GCR
101 | /// algorithm that can accommodate 20 cells per second. This translates
102 | /// to the GCRA parameters τ=1s, T=50ms (that's 1s / 20 cells).
103 | ///
104 | /// ```
105 | /// # use ratelimit_meter::{DirectRateLimiter, GCRA};
106 | /// # use std::num::NonZeroU32;
107 | /// # use std::time::{Instant, Duration};
108 | /// # #[macro_use] extern crate nonzero_ext;
109 | /// # extern crate ratelimit_meter;
110 | /// # #[cfg(feature = "std")]
111 | /// # fn main () {
112 | /// let mut limiter = DirectRateLimiter::::per_second(nonzero!(20u32));
113 | /// let now = Instant::now();
114 | /// let ms = Duration::from_millis(1);
115 | /// assert_eq!(Ok(()), limiter.check_at(now)); // the first cell is free
116 | /// for i in 0..20 {
117 | /// // Spam a lot:
118 | /// assert!(limiter.check_at(now).is_ok(), "at {}", i);
119 | /// }
120 | /// // We have exceeded the bucket capacity:
121 | /// assert!(limiter.check_at(now).is_err());
122 | ///
123 | /// // After a sufficient time period, cells are allowed again:
124 | /// assert_eq!(Ok(()), limiter.check_at(now + ms*50));
125 | /// # }
126 | /// # #[cfg(not(feature = "std"))] fn main() {}
127 | /// ```
128 | #[derive(Debug, Clone)]
129 | pub struct GCRA::Instant> {
130 | // The "weight" of a single packet in units of time.
131 | t: Duration,
132 |
133 | // The "capacity" of the bucket.
134 | tau: Duration,
135 |
136 | point: PhantomData,
137 | }
138 |
139 | impl Algorithm for GCRA
{
140 | type BucketState = State
;
141 |
142 | type NegativeDecision = NotUntil
;
143 |
144 | fn construct(
145 | capacity: NonZeroU32,
146 | cell_weight: NonZeroU32,
147 | per_time_unit: Duration,
148 | ) -> Result {
149 | if capacity < cell_weight {
150 | return Err(InconsistentCapacity::new(capacity, cell_weight));
151 | }
152 | Ok(GCRA {
153 | t: (per_time_unit / capacity.get()) * cell_weight.get(),
154 | tau: per_time_unit,
155 | point: PhantomData,
156 | })
157 | }
158 |
159 | /// Tests if a single cell can be accommodated by the
160 | /// rate-limiter and updates the state, if so.
161 | fn test_and_update(
162 | &self,
163 | state: &Self::BucketState,
164 | t0: P,
165 | ) -> Result<(), Self::NegativeDecision> {
166 | let tau = self.tau;
167 | let t = self.t;
168 | state.0.measure_and_replace(|tat| {
169 | // the "theoretical arrival time" of the next cell:
170 | let tat = tat.0.unwrap_or(t0);
171 | if t0 < tat.saturating_sub(tau) {
172 | (Err(NotUntil(tat)), None)
173 | } else {
174 | (Ok(()), Some(Tat(Some(cmp::max(tat, t0) + t))))
175 | }
176 | })
177 | }
178 |
179 | /// Tests if `n` cells can be accommodated by the rate-limiter
180 | /// and updates rate limiter state iff they can be.
181 | ///
182 | /// As this method is an extension of GCRA (using multiplication),
183 | /// it is likely not as fast (and not as obviously "right") as the
184 | /// single-cell variant.
185 | fn test_n_and_update(
186 | &self,
187 | state: &Self::BucketState,
188 | n: u32,
189 | t0: P,
190 | ) -> Result<(), NegativeMultiDecision> {
191 | let tau = self.tau;
192 | let t = self.t;
193 | state.0.measure_and_replace(|tat| {
194 | let tat = tat.0.unwrap_or(t0);
195 | let tat = match n {
196 | 0 => t0,
197 | 1 => tat,
198 | _ => {
199 | let weight = t * (n - 1);
200 | if (weight + t) > tau {
201 | // The bucket capacity can never accommodate this request
202 | return (Err(NegativeMultiDecision::InsufficientCapacity(n)), None);
203 | }
204 | tat + weight
205 | }
206 | };
207 |
208 | let additional_weight = match n {
209 | 0 => Duration::new(0, 0),
210 | 1 => t,
211 | _ => t * n,
212 | };
213 | if t0 < tat.saturating_sub(tau) {
214 | (
215 | Err(NegativeMultiDecision::BatchNonConforming(n, NotUntil(tat))),
216 | None,
217 | )
218 | } else {
219 | (
220 | Ok(()),
221 | Some(Tat(Some(cmp::max(tat, t0) + additional_weight))),
222 | )
223 | }
224 | })
225 | }
226 | }
227 |
--------------------------------------------------------------------------------
/src/algorithms/leaky_bucket.rs:
--------------------------------------------------------------------------------
1 | //! A classic leaky bucket algorithm
2 |
3 | use crate::lib::*;
4 | use crate::thread_safety::ThreadsafeWrapper;
5 | use crate::{
6 | algorithms::{Algorithm, RateLimitState},
7 | clock, InconsistentCapacity, NegativeMultiDecision, NonConformance,
8 | };
9 |
10 | /// Implements the industry-standard leaky bucket rate-limiting
11 | /// as-a-meter. The bucket keeps a "fill height", pretending to drip
12 | /// steadily (which reduces the fill height), and increases the fill
13 | /// height with every cell that is found conforming. If cells would
14 | /// make the bucket overflow, they count as non-conforming.
15 | ///
16 | /// # Drip implementation
17 | ///
18 | /// Instead of having a background task update the bucket's fill
19 | /// level, this implementation re-computes the fill level of the
20 | /// bucket on every call to [`check`](#method.check) and related
21 | /// methods.
22 | ///
23 | /// # Wait time calculation
24 | ///
25 | /// If the cell does not fit, this implementation computes the minimum
26 | /// wait time until the cell can be accommodated. This minimum wait
27 | /// time does not account for thundering herd effects or other
28 | /// problems in concurrent resource acquisition, so users of this
29 | /// library must take care to apply positive jitter to these wait
30 | /// times.
31 | ///
32 | /// # Example
33 | /// ``` rust
34 | /// # use ratelimit_meter::{DirectRateLimiter, LeakyBucket};
35 | /// # #[macro_use] extern crate nonzero_ext;
36 | /// # extern crate ratelimit_meter;
37 | /// # #[cfg(feature = "std")]
38 | /// # fn main () {
39 | /// let mut lb = DirectRateLimiter::::per_second(nonzero!(2u32));
40 | /// assert_eq!(Ok(()), lb.check());
41 | /// # }
42 | /// # #[cfg(not(feature = "std"))] fn main() {}
43 | /// ```
44 | #[derive(Debug, Clone, Eq, PartialEq)]
45 | pub struct LeakyBucket::Instant> {
46 | full: Duration,
47 | token_interval: Duration,
48 | point: PhantomData,
49 | }
50 |
51 | /// Represents the state of a single history of decisions.
52 | #[derive(Debug, Eq, PartialEq, Clone)]
53 | pub struct State(ThreadsafeWrapper>);
54 |
55 | impl Default for State {
56 | fn default() -> Self {
57 | State(Default::default())
58 | }
59 | }
60 |
61 | impl RateLimitState, P> for State {
62 | fn last_touched(&self, _params: &LeakyBucket
) -> Option
{
63 | let data = self.0.snapshot();
64 | Some(data.last_update? + data.level)
65 | }
66 | }
67 |
68 | #[cfg(feature = "std")]
69 | mod std {
70 | use crate::clock;
71 | use evmap::ShallowCopy;
72 |
73 | impl ShallowCopy for super::State {
74 | unsafe fn shallow_copy(&mut self) -> Self {
75 | super::State(self.0.shallow_copy())
76 | }
77 | }
78 | }
79 |
80 | /// Returned in case of a negative rate-limiting decision.
81 | ///
82 | /// To avoid the thundering herd effect, client code should always add
83 | /// some jitter to the wait time.
84 | #[derive(Debug, PartialEq)]
85 | pub struct TooEarly(P, Duration);
86 |
87 | impl fmt::Display for TooEarly {
88 | fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
89 | write!(f, "rate-limited until {:?}", self.0 + self.1)
90 | }
91 | }
92 |
93 | impl NonConformance for TooEarly
{
94 | #[inline]
95 | fn earliest_possible(&self) -> P {
96 | self.0 + self.1
97 | }
98 | }
99 |
100 | #[derive(Debug, Clone, PartialEq, Eq)]
101 | struct BucketState {
102 | level: Duration,
103 | last_update: Option,
104 | }
105 |
106 | impl Default for BucketState {
107 | fn default() -> Self {
108 | BucketState {
109 | level: Duration::new(0, 0),
110 | last_update: None,
111 | }
112 | }
113 | }
114 |
115 | impl Algorithm for LeakyBucket
{
116 | type BucketState = State
;
117 |
118 | type NegativeDecision = TooEarly
;
119 |
120 | fn construct(
121 | capacity: NonZeroU32,
122 | cell_weight: NonZeroU32,
123 | per_time_unit: Duration,
124 | ) -> Result {
125 | if capacity < cell_weight {
126 | return Err(InconsistentCapacity::new(capacity, cell_weight));
127 | }
128 | let token_interval = (per_time_unit * cell_weight.get()) / capacity.get();
129 | Ok(LeakyBucket {
130 | full: per_time_unit,
131 | token_interval,
132 | point: PhantomData,
133 | })
134 | }
135 |
136 | fn test_n_and_update(
137 | &self,
138 | state: &Self::BucketState,
139 | n: u32,
140 | t0: P,
141 | ) -> Result<(), NegativeMultiDecision>> {
142 | let full = self.full;
143 | let weight = self.token_interval * n;
144 | if weight > self.full {
145 | return Err(NegativeMultiDecision::InsufficientCapacity(n));
146 | }
147 | state.0.measure_and_replace(|state| {
148 | let mut new = BucketState {
149 | last_update: Some(t0),
150 | level: Duration::new(0, 0),
151 | };
152 | let last = state.last_update.unwrap_or(t0);
153 | // Prevent time travel: If any parallel calls get re-ordered,
154 | // or any tests attempt silly things, make sure to answer from
155 | // the last query onwards instead.
156 | let t0 = cmp::max(t0, last);
157 | // Decrement the level by the amount the bucket
158 | // has dripped in the meantime:
159 | new.level = state.level - cmp::min(t0.duration_since(last), state.level);
160 | if weight + new.level <= full {
161 | new.level += weight;
162 | (Ok(()), Some(new))
163 | } else {
164 | let wait_period = (weight + new.level) - full;
165 | (
166 | Err(NegativeMultiDecision::BatchNonConforming(
167 | n,
168 | TooEarly(t0, wait_period),
169 | )),
170 | None,
171 | )
172 | }
173 | })
174 | }
175 | }
176 |
--------------------------------------------------------------------------------
/src/clock.rs:
--------------------------------------------------------------------------------
1 | //! Time sources for the rate limiter.
2 | //!
3 | //! The time sources contained in this module allow the rate limiter
4 | //! to be (optionally) independent of std, and should additionally
5 | //! allow mocking the passage of time.
6 |
7 | use crate::lib::*;
8 |
9 | /// A measurement from a clock.
10 | pub trait Reference:
11 | Sized + Add + PartialEq + Eq + Ord + Copy + Clone + Send + Sync + Debug
12 | {
13 | /// Determines the time that separates two measurements of a
14 | /// clock. Implementations of this must perform a saturating
15 | /// subtraction - if the `earlier` timestamp should be later,
16 | /// `duration_since` must return the zero duration.
17 | fn duration_since(&self, earlier: Self) -> Duration;
18 |
19 | /// Returns a reference point that lies at most `duration` in the
20 | /// past from the current reference. If an underflow should occur,
21 | /// returns the current reference.
22 | fn saturating_sub(&self, duration: Duration) -> Self;
23 | }
24 |
25 | /// A time source used by rate limiters.
26 | pub trait Clock: Default + Clone {
27 | /// A measurement of a monotonically increasing clock.
28 | type Instant: Reference;
29 |
30 | /// Returns a measurement of the clock.
31 | fn now(&self) -> Self::Instant;
32 | }
33 |
34 | impl Reference for Duration {
35 | fn duration_since(&self, earlier: Self) -> Duration {
36 | self.checked_sub(earlier)
37 | .unwrap_or_else(|| Duration::new(0, 0))
38 | }
39 |
40 | fn saturating_sub(&self, duration: Duration) -> Self {
41 | self.checked_sub(duration).unwrap_or(*self)
42 | }
43 | }
44 |
45 | /// A mock implementation of a clock. All it does is keep track of
46 | /// what "now" is (relative to some point meaningful to the program),
47 | /// and returns that.
48 | #[derive(Debug, PartialEq, Clone, Default)]
49 | pub struct FakeRelativeClock {
50 | now: Duration,
51 | }
52 |
53 | impl FakeRelativeClock {
54 | /// Advances the fake clock by the given amount.
55 | pub fn advance(&mut self, by: Duration) {
56 | self.now += by
57 | }
58 | }
59 |
60 | impl Clock for FakeRelativeClock {
61 | type Instant = Duration;
62 |
63 | fn now(&self) -> Self::Instant {
64 | self.now
65 | }
66 | }
67 |
68 | #[cfg(not(feature = "std"))]
69 | mod no_std;
70 | #[cfg(not(feature = "std"))]
71 | pub use no_std::*;
72 |
73 | #[cfg(feature = "std")]
74 | mod with_std;
75 | #[cfg(feature = "std")]
76 | pub use with_std::*;
77 |
--------------------------------------------------------------------------------
/src/clock/no_std.rs:
--------------------------------------------------------------------------------
1 | use super::FakeRelativeClock;
2 |
3 | /// The default `no_std` clock that reports [`Durations`] must be advanced by the program.
4 | pub type DefaultClock = FakeRelativeClock;
5 |
--------------------------------------------------------------------------------
/src/clock/with_std.rs:
--------------------------------------------------------------------------------
1 | use super::{Clock, Reference};
2 | use crate::lib::*;
3 | use parking_lot::Mutex;
4 | use std::time::SystemTime;
5 |
6 | /// The default clock that reports [`Instant`]s.
7 | pub type DefaultClock = MonotonicClock;
8 |
9 | /// A mock implementation of a clock tracking [`Instant`]s. All it
10 | /// does is keep track of what "now" is by allowing the program to
11 | /// increment the current time (taken at time of construction) by some
12 | /// arbitrary [`Duration`].
13 | #[derive(Debug, Clone)]
14 | pub struct FakeAbsoluteClock {
15 | now: Arc>,
16 | }
17 |
18 | impl Default for FakeAbsoluteClock {
19 | fn default() -> Self {
20 | FakeAbsoluteClock {
21 | now: Arc::new(Mutex::new(Instant::now())),
22 | }
23 | }
24 | }
25 |
26 | impl FakeAbsoluteClock {
27 | /// Advances the fake clock by the given amount.
28 | pub fn advance(&mut self, by: Duration) {
29 | *(self.now.lock()) += by
30 | }
31 | }
32 |
33 | impl Clock for FakeAbsoluteClock {
34 | type Instant = Instant;
35 |
36 | fn now(&self) -> Self::Instant {
37 | *self.now.lock()
38 | }
39 | }
40 |
41 | /// The monotonic clock implemented by [`Instant`].
42 | #[derive(Clone, Debug, Default)]
43 | pub struct MonotonicClock();
44 |
45 | impl Reference for Instant {
46 | fn duration_since(&self, earlier: Self) -> Duration {
47 | if earlier < *self {
48 | *self - earlier
49 | } else {
50 | Duration::new(0, 0)
51 | }
52 | }
53 |
54 | fn saturating_sub(&self, duration: Duration) -> Self {
55 | self.checked_sub(duration).unwrap_or(*self)
56 | }
57 | }
58 |
59 | impl Clock for MonotonicClock {
60 | type Instant = Instant;
61 |
62 | fn now(&self) -> Self::Instant {
63 | Instant::now()
64 | }
65 | }
66 |
67 | /// The non-monotonic clock implemented by [`SystemTime`].
68 | #[derive(Clone, Debug, Default)]
69 | pub struct SystemClock();
70 |
71 | impl Reference for SystemTime {
72 | /// Returns the difference in times between the two
73 | /// SystemTimes. Due to the fallible nature of SystemTimes,
74 | /// returns the zero duration if a negative duration would
75 | /// result (e.g. due to system clock adjustments).
76 | fn duration_since(&self, earlier: Self) -> Duration {
77 | self.duration_since(earlier)
78 | .unwrap_or_else(|_| Duration::new(0, 0))
79 | }
80 |
81 | fn saturating_sub(&self, duration: Duration) -> Self {
82 | self.checked_sub(duration).unwrap_or(*self)
83 | }
84 | }
85 |
86 | impl Clock for SystemClock {
87 | type Instant = SystemTime;
88 |
89 | fn now(&self) -> Self::Instant {
90 | SystemTime::now()
91 | }
92 | }
93 |
--------------------------------------------------------------------------------
/src/errors.rs:
--------------------------------------------------------------------------------
1 | use crate::lib::*;
2 |
3 | /// An error that is returned when initializing a rate limiter that is
4 | /// too small to let a single cell through.
5 | #[derive(Debug)]
6 | pub struct InconsistentCapacity {
7 | capacity: NonZeroU32,
8 | cell_weight: NonZeroU32,
9 | }
10 |
11 | impl InconsistentCapacity {
12 | pub(crate) fn new(capacity: NonZeroU32, cell_weight: NonZeroU32) -> InconsistentCapacity {
13 | InconsistentCapacity {
14 | capacity,
15 | cell_weight,
16 | }
17 | }
18 | }
19 |
20 | impl fmt::Display for InconsistentCapacity {
21 | fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
22 | write!(
23 | f,
24 | "bucket capacity {} too small for a single cell with weight {}",
25 | self.capacity, self.cell_weight
26 | )
27 | }
28 | }
29 |
30 | /// Gives additional information about the negative outcome of a batch
31 | /// cell decision.
32 | ///
33 | /// Since batch queries can be made for batch sizes bigger than the
34 | /// rate limiter parameter could accomodate, there are now two
35 | /// possible negative outcomes:
36 | ///
37 | /// * `BatchNonConforming` - the query is valid but the Decider can
38 | /// not accomodate them.
39 | ///
40 | /// * `InsufficientCapacity` - the query was invalid as the rate
41 | /// limite parameters can never accomodate the number of cells
42 | /// queried for.
43 | #[derive(Debug, PartialEq)]
44 | pub enum NegativeMultiDecision {
45 | /// A batch of cells (the first argument) is non-conforming and
46 | /// can not be let through at this time. The second argument gives
47 | /// information about when that batch of cells might be let
48 | /// through again (not accounting for thundering herds and other,
49 | /// simultaneous decisions).
50 | BatchNonConforming(u32, E),
51 |
52 | /// The number of cells tested (the first argument) is larger than
53 | /// the bucket's capacity, which means the decision can never have
54 | /// a conforming result.
55 | InsufficientCapacity(u32),
56 | }
57 |
58 | impl fmt::Display for NegativeMultiDecision
59 | where
60 | E: fmt::Display,
61 | {
62 | fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
63 | match self {
64 | NegativeMultiDecision::BatchNonConforming(n, err) => write!(f, "{} cells: {}", n, err),
65 | NegativeMultiDecision::InsufficientCapacity(n) => write!(
66 | f,
67 | "bucket does not have enough capacity to accomodate {} cells",
68 | n
69 | ),
70 | }
71 | }
72 | }
73 |
--------------------------------------------------------------------------------
/src/example_algorithms.rs:
--------------------------------------------------------------------------------
1 | //! No-op examples.
2 |
3 | use crate::lib::*;
4 | use crate::{
5 | algorithms::{Algorithm, RateLimitState},
6 | clock, DirectRateLimiter, InconsistentCapacity, NegativeMultiDecision,
7 | };
8 |
9 | /// The most naive implementation of a rate-limiter ever: Always
10 | /// allows every cell through.
11 | /// # Example
12 | /// ```
13 | /// use ratelimit_meter::DirectRateLimiter;
14 | /// use ratelimit_meter::example_algorithms::Allower;
15 | /// let mut allower = Allower::ratelimiter();
16 | /// assert!(allower.check().is_ok());
17 | /// ```
18 | #[derive(Default, Copy, Clone, Debug)]
19 | pub struct Allower {}
20 |
21 | impl Allower {
22 | /// Return a rate-limiter that lies, i.e. that allows all requests
23 | /// through.
24 | pub fn ratelimiter() -> DirectRateLimiter {
25 | // These numbers are fake, but we make them up for convenience:
26 | DirectRateLimiter::per_second(nonzero!(1u32))
27 | }
28 | }
29 |
30 | impl RateLimitState for () {
31 | fn last_touched(&self, _params: &Allower) -> Option {
32 | None
33 | }
34 | }
35 |
36 | /// A non-error - the Allower example rate-limiter always returns a
37 | /// positive result, so this error is never returned.
38 | #[derive(Debug, PartialEq)]
39 | pub enum Impossible {}
40 |
41 | impl fmt::Display for Impossible {
42 | fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
43 | write!(f, "can't happen")
44 | }
45 | }
46 |
47 | impl Algorithm for Allower {
48 | type BucketState = ();
49 | type NegativeDecision = Impossible;
50 |
51 | fn construct(
52 | _capacity: NonZeroU32,
53 | _cell_weight: NonZeroU32,
54 | _per_time_unit: Duration,
55 | ) -> Result {
56 | Ok(Allower {})
57 | }
58 |
59 | /// Allows all cells through unconditionally.
60 | fn test_n_and_update(
61 | &self,
62 | _state: &Self::BucketState,
63 | _n: u32,
64 | _t0: Always,
65 | ) -> Result<(), NegativeMultiDecision> {
66 | Ok(())
67 | }
68 | }
69 |
70 | /// A pseudo-instant that never changes.
71 | ///
72 | /// It is used to implement the `Allower` rate-limiter type, which
73 | /// never denies any requests.
74 | #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
75 | pub struct Always();
76 | impl clock::Reference for Always {
77 | fn duration_since(&self, _other: Self) -> Duration {
78 | Duration::new(0, 0)
79 | }
80 |
81 | fn saturating_sub(&self, _: Duration) -> Self {
82 | *self
83 | }
84 | }
85 |
86 | impl Add for Always {
87 | type Output = Always;
88 | fn add(self, _rhs: Duration) -> Always {
89 | Always()
90 | }
91 | }
92 |
93 | impl Sub for Always {
94 | type Output = Always;
95 | fn sub(self, _rhs: Duration) -> Always {
96 | Always()
97 | }
98 | }
99 |
100 | #[derive(Default, Debug, Clone)]
101 | pub struct ForeverClock();
102 |
103 | impl ForeverClock {
104 | pub fn now() -> Always {
105 | Always()
106 | }
107 | }
108 |
109 | impl clock::Clock for ForeverClock {
110 | type Instant = Always;
111 |
112 | fn now(&self) -> Self::Instant {
113 | Always()
114 | }
115 | }
116 |
--------------------------------------------------------------------------------
/src/lib.rs:
--------------------------------------------------------------------------------
1 | //! # Leaky Bucket Rate-Limiting (as a meter) in Rust
2 | //!
3 | //! This crate provides generic rate-limiting interfaces and
4 | //! implements a few rate-limiting algorithms for programs that need
5 | //! to regulate the rate of their outgoing requests.
6 | //!
7 | //! This crate currently provides in-memory implementations of a by-key
8 | //! (limits enforced per key, e.g. an IP address or a customer ID) and a
9 | //! simple (one limit per object) state tracker.
10 | //!
11 | //! The simple (one limit per object) state tracker can be used in
12 | //! `no_std` environments, such as embedded systems.
13 | //!
14 | //! ## Interface
15 | //!
16 | //! This crate implements two "serious" rate-limiting/traffic-shaping
17 | //! algorithms:
18 | //! [GCRA](https://en.wikipedia.org/wiki/Generic_cell_rate_algorithm)
19 | //! and a [Leaky
20 | //! Bucket](https://en.wikipedia.org/wiki/Leaky_bucket#As_a_meter). An
21 | //! "unserious" implementation is provided also: The
22 | //! [`Allower`](example_algorithms/struct.Allower.html), which returns
23 | //! "Yes" to all rate-limiting queries.
24 | //!
25 | //! The Generic Cell Rate Algorithm can be used by in an in-memory
26 | //! rate limiter like so:
27 | //!
28 | //! ``` rust
29 | //! use std::num::NonZeroU32;
30 | //! use ratelimit_meter::{DirectRateLimiter, GCRA};
31 | //!
32 | //! # #[macro_use] extern crate nonzero_ext;
33 | //! # extern crate ratelimit_meter;
34 | //! # #[cfg(feature = "std")]
35 | //! # fn main () {
36 | //! let mut lim = DirectRateLimiter::::per_second(nonzero!(50u32)); // Allow 50 units per second
37 | //! assert_eq!(Ok(()), lim.check());
38 | //! # }
39 | //! # #[cfg(not(feature = "std"))]
40 | //! # fn main() {}
41 | //! ```
42 | //!
43 | //! The rate-limiter interface is intentionally geared towards only
44 | //! providing callers with the information they need to make decisions
45 | //! about what to do with each cell. Deciders return additional
46 | //! information about why a cell should be denied alongside the
47 | //! decision. This allows callers to e.g. provide better error
48 | //! messages to users.
49 | //!
50 | //! As a consequence, the `ratelimit_meter` crate does not provide any
51 | //! facility to wait until a cell would be allowed - if you require
52 | //! this, you should use the
53 | //! [`NonConformance`](struct.NonConformance.html) returned with
54 | //! negative decisions and have the program wait using the method best
55 | //! suited for this, e.g. an event loop.
56 | //!
57 | //! ## Using this crate effectively
58 | //!
59 | //! Many of the parameters in use by this crate are `NonZeroU32` -
60 | //! since they are not very ergonomic to construct from constants
61 | //! using stdlib means, I recommend using the
62 | //! [nonzero_ext](https://crates.io/crates/nonzero_ext) crate, which
63 | //! comes with a macro `nonzero!()`. This macro makes it far easier to
64 | //! construct rate limiters without cluttering your code.
65 | //!
66 | //! ## Rate-limiting Algorithms
67 | //!
68 | //! ### Design and implementation of GCRA
69 | //!
70 | //! The GCRA limits the rate of cells by determining when the "next"
71 | //! cell is expected to arrive; any cells that arrive before that time
72 | //! are classified as non-conforming; the methods for checking cells
73 | //! also return an expected arrival time for these cells, so that
74 | //! callers can choose to wait (adding jitter), or reject the cell.
75 | //!
76 | //! Since using the GCRA results in a much smoother usage pattern, it
77 | //! appears to be very useful for "outgoing" traffic behaviors,
78 | //! e.g. throttling API call rates, or emails sent to a person in a
79 | //! period of time.
80 | //!
81 | //! Unlike token or leaky bucket algorithms, the GCRA assumes that all
82 | //! units of work are of the same "weight", and so allows some
83 | //! optimizations which result in much more concise and fast code (it
84 | //! does not even use multiplication or division in the "hot" path).
85 | //!
86 | //! See [the documentation of the GCRA type](algorithms/gcra/struct.GCRA.html) for
87 | //! more details on its implementation and on trade-offs that apply to
88 | //! it.
89 | //!
90 | //! ### Design and implementation of the leaky bucket
91 | //!
92 | //! In contrast to the GCRA, the leaky bucket algorithm does not place
93 | //! any constraints on the next cell's arrival time: Whenever there is
94 | //! capacity left in the bucket, it can be used. This means that the
95 | //! distribution of "yes" decisions from heavy usage on the leaky
96 | //! bucket rate-limiter will be clustered together. On average, the
97 | //! cell rates of both the GCRA and the leaky bucket will be the same,
98 | //! but in terms of observable behavior, the leaky bucket will appear
99 | //! to allow requests at a more predictable rate.
100 | //!
101 | //! This kind of behavior is usually what people of online APIs expect
102 | //! these days, which makes the leaky bucket a very popular technique
103 | //! for rate-limiting on these kinds of services.
104 | //!
105 | //! The leaky bucket algorithm implemented in this crate is fairly
106 | //! standard: It only updates the bucket fill gauge when a cell is
107 | //! checked, and supports checking "batches" of cells in a single call
108 | //! with no problems.
109 | //!
110 | //! ## Thread-safe operation
111 | //!
112 | //! The in-memory implementations in this crate use parking_lot
113 | //! mutexes to ensure rate-limiting operations can happen safely
114 | //! across threads.
115 | //!
116 | //! Example:
117 | //!
118 | //! ```
119 | //! use std::thread;
120 | //! use std::num::NonZeroU32;
121 | //! use std::time::Duration;
122 | //! use ratelimit_meter::{DirectRateLimiter, GCRA};
123 | //!
124 | //! # #[macro_use] extern crate nonzero_ext;
125 | //! # extern crate ratelimit_meter;
126 | //! # #[cfg(feature = "std")]
127 | //! # fn main () {
128 | //! // Allow 50 units/second across all threads:
129 | //! let mut lim = DirectRateLimiter::::per_second(nonzero!(50u32));
130 | //! let mut thread_lim = lim.clone();
131 | //! thread::spawn(move || { assert_eq!(Ok(()), thread_lim.check());});
132 | //! assert_eq!(Ok(()), lim.check());
133 | //! # }
134 | //! # #[cfg(not(feature = "std"))]
135 | //! # fn main() {}
136 | //! ```
137 | //!
138 | //! ## Usage with `no_std`
139 | //!
140 | //! `ratelimit_meter` can be used in `no_std` crates, with a reduced
141 | //! feature set. These features are available:
142 | //!
143 | //! * [`DirectRateLimiter`](state/direct/struct.DirectRateLimiter.html)
144 | //! for a single rate-limiting history per limit,
145 | //! * measurements using relative timestamps (`Duration`) by default,
146 | //! * extensibility for integrating a custom time source.
147 | //!
148 | //! The following things are not available in `no_std` builds by default:
149 | //!
150 | //! * `check` and `check_n` - unless you implement a custom time
151 | //! source, you have to pass a timestamp to check the rate-limit
152 | //! against.
153 | //! * [`KeyedRateLimiter`](state/keyed/struct.KeyedRateLimiter.html) -
154 | //! the keyed state representation requires too much of `std` right
155 | //! now to be feasible to implement.
156 | //!
157 | //! To use the crate, turn off default features and enable the
158 | //! `"no_std"` feature, like so:
159 | //!
160 | //! ``` toml
161 | //! [dependencies.ratelimit_meter]
162 | //! version = "..."
163 | //! default-features = false
164 | //! features = ["no_std"]
165 | //! ```
166 | //!
167 | //! ### Implementing your own custom time source in `no_std`
168 | //!
169 | //! On platforms that do have a clock or other time source, you can
170 | //! use that time source to implement a trait provided by
171 | //! `ratelimit_meter`, which will enable the `check` and `check_n`
172 | //! methods on rate limiters. Here is an example:
173 | //!
174 | //! ```rust,ignore
175 | //! // MyTimeSource is what provides your timestamps. Since it probably
176 | //! // doesn't live in your crate, we make a newtype:
177 | //! use ratelimit_meter::instant;
178 | //! struct MyInstant(MyTimeSource);
179 | //!
180 | //! impl instant::Relative for MyInstant {
181 | //! fn duration_since(&self, other: Self) -> Duration {
182 | //! self.duration_since(other)
183 | //! }
184 | //! }
185 | //!
186 | //! impl instant::Absolute for MyInstant {
187 | //! fn now() -> Self {
188 | //! MyTimeSource::now()
189 | //! }
190 | //! }
191 | //!
192 | //! impl Add for MyInstant {
193 | //! type Output = MyInstant;
194 | //! fn add(self, rhs: Duration) -> Always {
195 | //! self.0 + rhs
196 | //! }
197 | //! }
198 | //!
199 | //! impl Sub for MyInstant {
200 | //! type Output = MyInstant;
201 | //! fn sub(self, rhs: Duration) -> Always {
202 | //! self.0 - rhs
203 | //! }
204 | //! }
205 | //! ```
206 | //!
207 | //! Then, using that type to create a rate limiter with that time
208 | //! source is a little more verbose. It looks like this:
209 | //!
210 | //! ```rust,ignore
211 | //! let mut lim = DirectRateLimiter::,MyInstant>::per_second(nonzero!(50u32));
212 | //! lim.check().ok();
213 | //! ```
214 |
215 | // Allow using ratelimit_meter without std
216 | #![cfg_attr(not(feature = "std"), no_std)]
217 | // Deny warnings
218 | #![cfg_attr(feature = "cargo-clippy", deny(warnings))]
219 |
220 | pub mod algorithms;
221 | pub mod clock;
222 | mod errors;
223 | pub mod example_algorithms;
224 | pub mod state;
225 | pub mod test_utilities;
226 | mod thread_safety;
227 |
228 | #[macro_use]
229 | extern crate nonzero_ext;
230 |
231 | #[cfg(not(feature = "std"))]
232 | extern crate alloc;
233 |
234 | pub use self::algorithms::LeakyBucket;
235 | pub use self::algorithms::NonConformance;
236 | pub use self::algorithms::GCRA;
237 |
238 | pub use self::state::DirectRateLimiter;
239 |
240 | #[cfg(feature = "std")]
241 | pub use self::state::KeyedRateLimiter;
242 |
243 | pub use self::errors::*;
244 |
245 | /// A facade around all the types we need from std/core crates, to
246 | /// avoid unnecessary cfg-conditionalization everywhere.
247 | mod lib {
248 | mod core {
249 | #[cfg(not(feature = "std"))]
250 | pub use core::*;
251 |
252 | #[cfg(feature = "std")]
253 | pub use std::*;
254 | }
255 |
256 | pub use self::core::clone::Clone;
257 | pub use self::core::cmp::{Eq, Ord, PartialEq};
258 | pub use self::core::default::Default;
259 | pub use self::core::fmt::Debug;
260 | pub use self::core::marker::{Copy, PhantomData, Send, Sized, Sync};
261 | pub use self::core::num::NonZeroU32;
262 | pub use self::core::ops::{Add, Sub};
263 | pub use self::core::time::Duration;
264 |
265 | pub use self::core::cmp;
266 | pub use self::core::fmt;
267 |
268 | /// Imports that are only available on std.
269 | #[cfg(feature = "std")]
270 | mod std {
271 | pub use std::collections::hash_map::RandomState;
272 | pub use std::hash::{BuildHasher, Hash};
273 | pub use std::sync::Arc;
274 | pub use std::time::Instant;
275 | }
276 |
277 | #[cfg(feature = "no_std")]
278 | mod no_std {
279 | pub use alloc::sync::Arc;
280 | }
281 |
282 | #[cfg(feature = "std")]
283 | pub use self::std::*;
284 |
285 | #[cfg(not(feature = "std"))]
286 | pub use self::no_std::*;
287 | }
288 |
--------------------------------------------------------------------------------
/src/state.rs:
--------------------------------------------------------------------------------
1 | //! Data structures that keep rate-limiting state.
2 |
3 | pub mod direct;
4 |
5 | #[cfg(feature = "std")]
6 | pub mod keyed;
7 |
8 | pub use self::direct::DirectRateLimiter;
9 |
10 | #[cfg(feature = "std")]
11 | pub use self::keyed::KeyedRateLimiter;
12 |
--------------------------------------------------------------------------------
/src/state/direct.rs:
--------------------------------------------------------------------------------
1 | //! An in-memory rate limiter that can make decisions for a single
2 | //! situation.
3 |
4 | use crate::lib::*;
5 |
6 | use crate::{
7 | algorithms::{Algorithm, DefaultAlgorithm},
8 | clock, InconsistentCapacity, NegativeMultiDecision,
9 | };
10 |
11 | /// An in-memory rate limiter that makes direct (un-keyed)
12 | /// rate-limiting decisions. Direct rate limiters can be used to
13 | /// e.g. regulate the transmission of packets on a single connection,
14 | /// or to ensure that an API client stays within a server's rate
15 | /// limit.
16 | #[derive(Debug, Clone)]
17 | pub struct DirectRateLimiter<
18 | A: Algorithm = DefaultAlgorithm,
19 | C: clock::Clock = clock::DefaultClock,
20 | > {
21 | state: A::BucketState,
22 | algorithm: A,
23 | clock: C,
24 | }
25 |
26 | impl DirectRateLimiter
27 | where
28 | C: clock::Clock,
29 | A: Algorithm,
30 | {
31 | /// Construct a new rate limiter that allows `capacity` cells per
32 | /// time unit through.
33 | /// # Examples
34 | /// You can construct a GCRA rate limiter like so:
35 | /// ```
36 | /// # use std::num::NonZeroU32;
37 | /// # use std::time::Duration;
38 | /// use ratelimit_meter::{DirectRateLimiter, GCRA};
39 | /// # #[macro_use] extern crate nonzero_ext;
40 | /// # extern crate ratelimit_meter;
41 | /// # fn main () {
42 | /// let _gcra = DirectRateLimiter::::new(nonzero!(100u32), Duration::from_secs(5));
43 | /// # }
44 | /// ```
45 | ///
46 | /// and similarly, for a leaky bucket:
47 | /// ```
48 | /// # use std::time::Duration;
49 | /// use ratelimit_meter::{DirectRateLimiter, LeakyBucket};
50 | /// # #[macro_use] extern crate nonzero_ext;
51 | /// # extern crate ratelimit_meter;
52 | /// # fn main () {
53 | /// let _lb = DirectRateLimiter::::new(nonzero!(100u32), Duration::from_secs(5));
54 | /// # }
55 | /// ```
56 | pub fn new(capacity: NonZeroU32, per_time_unit: Duration) -> Self {
57 | DirectRateLimiter {
58 | state: >::BucketState::default(),
59 | algorithm: >::construct(
60 | capacity,
61 | nonzero!(1u32),
62 | per_time_unit,
63 | )
64 | .unwrap(),
65 | clock: Default::default(),
66 | }
67 | }
68 |
69 | /// Construct a new rate limiter that allows `capacity` cells per
70 | /// second.
71 | /// # Examples
72 | /// Constructing a GCRA rate limiter that lets through 100 cells per second:
73 | /// ```
74 | /// # use std::time::Duration;
75 | /// use ratelimit_meter::{DirectRateLimiter, GCRA};
76 | /// # #[macro_use] extern crate nonzero_ext;
77 | /// # extern crate ratelimit_meter;
78 | /// # fn main () {
79 | /// let _gcra = DirectRateLimiter::::per_second(nonzero!(100u32));
80 | /// # }
81 | /// ```
82 | ///
83 | /// and a leaky bucket:
84 | /// ```
85 | /// # use std::time::Duration;
86 | /// use ratelimit_meter::{DirectRateLimiter, LeakyBucket};
87 | /// # #[macro_use] extern crate nonzero_ext;
88 | /// # extern crate ratelimit_meter;
89 | /// # fn main () {
90 | /// let _gcra = DirectRateLimiter::::per_second(nonzero!(100u32));
91 | /// # }
92 | /// ```
93 | pub fn per_second(capacity: NonZeroU32) -> Self {
94 | Self::new(capacity, Duration::from_secs(1))
95 | }
96 |
97 | /// Return a builder that can be used to construct a rate limiter using
98 | /// the parameters passed to the Builder.
99 | pub fn build_with_capacity(capacity: NonZeroU32) -> Builder {
100 | Builder {
101 | capacity,
102 | cell_weight: nonzero!(1u32),
103 | time_unit: Duration::from_secs(1),
104 | end_result: PhantomData,
105 | clock: Default::default(),
106 | }
107 | }
108 |
109 | /// Tests whether a single cell can be accommodated at the given
110 | /// time stamp. See [`check`](#method.check).
111 | pub fn check_at(
112 | &mut self,
113 | at: C::Instant,
114 | ) -> Result<(), >::NegativeDecision> {
115 | self.algorithm.test_and_update(&self.state, at)
116 | }
117 |
118 | /// Tests if `n` cells can be accommodated at the given time
119 | /// (`Instant::now()`), using [`check_n`](#method.check_n)
120 | pub fn check_n_at(
121 | &mut self,
122 | n: u32,
123 | at: C::Instant,
124 | ) -> Result<(), NegativeMultiDecision<>::NegativeDecision>> {
125 | self.algorithm.test_n_and_update(&self.state, n, at)
126 | }
127 |
128 | /// Tests if a single cell can be accommodated at the clock's
129 | /// current reading. If it can be, `check` updates the rate
130 | /// limiter state to account for the conforming cell and returns
131 | /// `Ok(())`.
132 | ///
133 | /// If the cell is non-conforming (i.e., it can't be accomodated
134 | /// at this time stamp), `check_at` returns `Err` with information
135 | /// about the earliest time at which a cell could be considered
136 | /// conforming.
137 | pub fn check(&mut self) -> Result<(), >::NegativeDecision> {
138 | self.algorithm
139 | .test_and_update(&self.state, self.clock.now())
140 | }
141 |
142 | /// Tests if `n` cells can be accommodated at the clock's current
143 | /// reading. If (and only if) all cells in the batch can be
144 | /// accomodated, the `MultiDecider` updates the internal state to
145 | /// account for all cells and returns `Ok(())`.
146 | ///
147 | /// If the entire batch of cells would not be conforming but the
148 | /// rate limiter has the capacity to accomodate the cells at any
149 | /// point in time, `check_n_at` returns error
150 | /// [`NegativeMultiDecision::BatchNonConforming`](../../enum.NegativeMultiDecision.html#variant.BatchNonConforming),
151 | /// holding the number of cells the rate limiter's negative
152 | /// outcome result.
153 | ///
154 | /// If `n` exceeds the bucket capacity, `check_n_at` returns
155 | /// [`NegativeMultiDecision::InsufficientCapacity`](../../enum.NegativeMultiDecision.html#variant.InsufficientCapacity),
156 | /// indicating that a batch of this many cells can never succeed.
157 | pub fn check_n(
158 | &mut self,
159 | n: u32,
160 | ) -> Result<(), NegativeMultiDecision<>::NegativeDecision>> {
161 | self.algorithm
162 | .test_n_and_update(&self.state, n, self.clock.now())
163 | }
164 | }
165 |
166 | /// An object that allows incrementally constructing rate Limiter
167 | /// objects.
168 | pub struct Builder
169 | where
170 | C: clock::Clock,
171 | A: Algorithm + Sized,
172 | {
173 | capacity: NonZeroU32,
174 | cell_weight: NonZeroU32,
175 | time_unit: Duration,
176 | end_result: PhantomData,
177 | clock: C,
178 | }
179 |
180 | impl Builder
181 | where
182 | C: clock::Clock,
183 | A: Algorithm + Sized,
184 | {
185 | /// Sets the "weight" of each cell being checked against the
186 | /// bucket. Each cell fills the bucket by this much.
187 | pub fn cell_weight(
188 | &mut self,
189 | weight: NonZeroU32,
190 | ) -> Result<&mut Builder, InconsistentCapacity> {
191 | if self.cell_weight > self.capacity {
192 | return Err(InconsistentCapacity::new(self.capacity, self.cell_weight));
193 | }
194 | self.cell_weight = weight;
195 | Ok(self)
196 | }
197 |
198 | /// Sets the "unit of time" within which the bucket drains.
199 | ///
200 | /// The assumption is that in a period of `time_unit` (if no cells
201 | /// are being checked), the bucket is fully drained.
202 | pub fn per(&mut self, time_unit: Duration) -> &mut Builder {
203 | self.time_unit = time_unit;
204 | self
205 | }
206 |
207 | /// Sets the clock used by the bucket.
208 | pub fn using_clock(&mut self, clock: C) -> &mut Builder {
209 | self.clock = clock;
210 | self
211 | }
212 |
213 | /// Builds a rate limiter of the specified type.
214 | pub fn build(&self) -> Result, InconsistentCapacity> {
215 | Ok(DirectRateLimiter {
216 | state: >::BucketState::default(),
217 | algorithm: >::construct(
218 | self.capacity,
219 | self.cell_weight,
220 | self.time_unit,
221 | )?,
222 | clock: self.clock.clone(),
223 | })
224 | }
225 | }
226 |
--------------------------------------------------------------------------------
/src/state/keyed.rs:
--------------------------------------------------------------------------------
1 | #![cfg(feature = "std")]
2 | //! An in-memory rate limiter that can keep track of rates for
3 | //! multiple keys, e.g. per-customer or per-IC rates.
4 |
5 | use crate::lib::*;
6 |
7 | use evmap::{self, ReadHandle, WriteHandle};
8 | use parking_lot::Mutex;
9 |
10 | use crate::{
11 | algorithms::{Algorithm, DefaultAlgorithm, KeyableRateLimitState, RateLimitState},
12 | clock,
13 | clock::Reference,
14 | InconsistentCapacity, NegativeMultiDecision,
15 | };
16 |
17 | type MapWriteHandle =
18 | Arc::Instant>>::BucketState, (), H>>>;
19 |
20 | /// An in-memory rate limiter that regulates a single rate limit for
21 | /// multiple keys.
22 | ///
23 | /// Keyed rate limiters can be used to e.g. enforce a per-IC address
24 | /// or a per-customer request limit on the server side.
25 | ///
26 | /// This implementation of the keyed rate limiter uses
27 | /// [`evmap`](../../../evmap/index.html), a read lock-free, concurrent
28 | /// hash map. Addition of new keys (e.g. a new customer making their
29 | /// first request) is synchronized and happens one at a time (it
30 | /// synchronizes writes to minimize the effects from `evmap`'s
31 | /// eventually consistent behavior on key addition), while reads of
32 | /// existing keys all happen simultaneously, then get synchronized by
33 | /// the rate limiting algorithm itself.
34 | ///
35 | /// ```
36 | /// # use std::num::NonZeroU32;
37 | /// # use std::time::Duration;
38 | /// use ratelimit_meter::{KeyedRateLimiter};
39 | /// # #[macro_use] extern crate nonzero_ext;
40 | /// # extern crate ratelimit_meter;
41 | /// # fn main () {
42 | /// let mut limiter = KeyedRateLimiter::<&str>::new(nonzero!(1u32), Duration::from_secs(5));
43 | /// assert_eq!(Ok(()), limiter.check("customer1")); // allowed!
44 | /// assert_ne!(Ok(()), limiter.check("customer1")); // ...but now customer1 must wait 5 seconds.
45 | ///
46 | /// assert_eq!(Ok(()), limiter.check("customer2")); // it's customer2's first request!
47 | /// # }
48 | /// ```
49 | ///
50 | /// # Expiring old keys
51 | /// If a key has not been checked in a long time, that key can be
52 | /// expired safely (the next rate limit check for that key would
53 | /// behave as if the key was not present in the map, after all). To
54 | /// remove the unused keys and free up space, use the
55 | /// [`cleanup`](method.cleanup) method:
56 | ///
57 | /// ```
58 | /// # use std::num::NonZeroU32;
59 | /// # use std::time::Duration;
60 | /// use ratelimit_meter::{KeyedRateLimiter};
61 | /// # #[macro_use] extern crate nonzero_ext;
62 | /// # extern crate ratelimit_meter;
63 | /// # fn main () {
64 | /// let mut limiter = KeyedRateLimiter::<&str>::new(nonzero!(100u32), Duration::from_secs(5));
65 | /// limiter.check("hi there");
66 | /// // time passes...
67 | ///
68 | /// // remove all keys that have been expireable for 10 minutes:
69 | /// limiter.cleanup(Duration::from_secs(600));
70 | /// # }
71 | /// ```
72 | #[derive(Clone)]
73 | pub struct KeyedRateLimiter<
74 | K: Eq + Hash + Clone,
75 | A: Algorithm = DefaultAlgorithm,
76 | C: clock::Clock = clock::DefaultClock,
77 | H: BuildHasher + Clone = RandomState,
78 | > where
79 | A::BucketState: KeyableRateLimitState,
80 | {
81 | algorithm: A,
82 | map_reader: ReadHandle,
83 | map_writer: MapWriteHandle,
84 | clock: C,
85 | }
86 |
87 | impl fmt::Debug for KeyedRateLimiter
88 | where
89 | A: Algorithm,
90 | A::BucketState: KeyableRateLimitState,
91 | K: Eq + Hash + Clone,
92 | {
93 | fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
94 | write!(f, "KeyedRateLimiter{{{params:?}}}", params = self.algorithm)
95 | }
96 | }
97 |
98 | impl KeyedRateLimiter
99 | where
100 | C: clock::Clock,
101 | A: Algorithm,
102 | A::BucketState: KeyableRateLimitState,
103 | K: Eq + Hash + Clone,
104 | {
105 | /// Construct a new rate limiter that allows `capacity` cells per
106 | /// time unit through.
107 | /// # Examples
108 | /// ```
109 | /// # use std::num::NonZeroU32;
110 | /// # use std::time::Duration;
111 | /// use ratelimit_meter::{KeyedRateLimiter};
112 | /// # #[macro_use] extern crate nonzero_ext;
113 | /// # extern crate ratelimit_meter;
114 | /// # fn main () {
115 | /// let _limiter = KeyedRateLimiter::<&str>::new(nonzero!(100u32), Duration::from_secs(5));
116 | /// # }
117 | /// ```
118 | pub fn new(capacity: NonZeroU32, per_time_unit: Duration) -> Self {
119 | let (r, mut w): (
120 | ReadHandle,
121 | WriteHandle,
122 | ) = evmap::new();
123 | w.refresh();
124 | KeyedRateLimiter {
125 | algorithm: >::construct(
126 | capacity,
127 | nonzero!(1u32),
128 | per_time_unit,
129 | )
130 | .unwrap(),
131 | map_reader: r,
132 | map_writer: Arc::new(Mutex::new(w)),
133 | clock: Default::default(),
134 | }
135 | }
136 |
137 | /// Returns the number of non-empty keys present in the map.
138 | pub fn len(&self) -> usize {
139 | self.map_reader.len()
140 | }
141 |
142 | /// Returns `true` if `self` has no keys stored in it.
143 | pub fn is_empty(&self) -> bool {
144 | self.map_reader.is_empty()
145 | }
146 |
147 | /// Construct a new keyed rate limiter that allows `capacity`
148 | /// cells per second.
149 | ///
150 | /// # Examples
151 | /// Constructing a rate limiter keyed by `&str` that lets through
152 | /// 100 cells per second:
153 | ///
154 | /// ```
155 | /// # use std::time::Duration;
156 | /// use ratelimit_meter::{KeyedRateLimiter, GCRA};
157 | /// # #[macro_use] extern crate nonzero_ext;
158 | /// # extern crate ratelimit_meter;
159 | /// # fn main () {
160 | /// let _limiter = KeyedRateLimiter::<&str, GCRA>::per_second(nonzero!(100u32));
161 | /// # }
162 | /// ```
163 | pub fn per_second(capacity: NonZeroU32) -> Self {
164 | Self::new(capacity, Duration::from_secs(1))
165 | }
166 |
167 | /// Return a constructor that can be used to construct a keyed
168 | /// rate limiter with the builder pattern.
169 | pub fn build_with_capacity(capacity: NonZeroU32) -> Builder {
170 | Builder {
171 | capacity,
172 | ..Default::default()
173 | }
174 | }
175 |
176 | fn check_and_update_key(&self, key: K, update: F) -> Result<(), E>
177 | where
178 | F: Fn(&A::BucketState) -> Result<(), E>,
179 | {
180 | self.map_reader
181 | .get_and(&key, |v| {
182 | // we have at least one element (owing to the nature of
183 | // the evmap, it says there could be >1
184 | // entries, but we'll only ever add one):
185 | let state = &v[0];
186 | update(state)
187 | })
188 | .unwrap_or_else(|| {
189 | // entry does not exist, let's add one.
190 | let mut w = self.map_writer.lock();
191 | let state: A::BucketState = Default::default();
192 | let result = update(&state);
193 | w.update(key, state);
194 | w.flush();
195 | result
196 | })
197 | }
198 |
199 | /// Tests if a single cell for the given key can be accommodated
200 | /// at `Instant::now()`. If it can be, `check` updates the rate
201 | /// limiter state on that key to account for the conforming cell
202 | /// and returns `Ok(())`.
203 | ///
204 | /// If the cell is non-conforming (i.e., it can't be accomodated
205 | /// at this time stamp), `check_at` returns `Err` with information
206 | /// about the earliest time at which a cell could be considered
207 | /// conforming under that key.
208 | pub fn check(&mut self, key: K) -> Result<(), >::NegativeDecision> {
209 | self.check_at(key, self.clock.now())
210 | }
211 |
212 | /// Tests if `n` cells for the given key can be accommodated at
213 | /// the current time stamp. If (and only if) all cells in the
214 | /// batch can be accomodated, the `MultiDecider` updates the rate
215 | /// limiter state on that key to account for all cells and returns
216 | /// `Ok(())`.
217 | ///
218 | /// If the entire batch of cells would not be conforming but the
219 | /// rate limiter has the capacity to accomodate the cells at any
220 | /// point in time, `check_n_at` returns error
221 | /// [`NegativeMultiDecision::BatchNonConforming`](../../enum.NegativeMultiDecision.html#variant.BatchNonConforming),
222 | /// holding the number of cells and the rate limiter's negative
223 | /// outcome result.
224 | ///
225 | /// If `n` exceeds the bucket capacity, `check_n_at` returns
226 | /// [`NegativeMultiDecision::InsufficientCapacity`](../../enum.NegativeMultiDecision.html#variant.InsufficientCapacity),
227 | /// indicating that a batch of this many cells can never succeed.
228 | pub fn check_n(
229 | &mut self,
230 | key: K,
231 | n: u32,
232 | ) -> Result<(), NegativeMultiDecision<>::NegativeDecision>> {
233 | self.check_n_at(key, n, self.clock.now())
234 | }
235 |
236 | /// Tests whether a single cell for the given key can be
237 | /// accommodated at the given time stamp. See
238 | /// [`check`](#method.check).
239 | pub fn check_at(
240 | &mut self,
241 | key: K,
242 | at: C::Instant,
243 | ) -> Result<(), >::NegativeDecision> {
244 | self.check_and_update_key(key, |state| self.algorithm.test_and_update(state, at))
245 | }
246 |
247 | /// Tests if `n` cells for the given key can be accommodated at
248 | /// the given time (`Instant::now()`), using
249 | /// [`check_n`](#method.check_n)
250 | pub fn check_n_at(
251 | &mut self,
252 | key: K,
253 | n: u32,
254 | at: C::Instant,
255 | ) -> Result<(), NegativeMultiDecision<>::NegativeDecision>> {
256 | self.check_and_update_key(key, |state| self.algorithm.test_n_and_update(state, n, at))
257 | }
258 |
259 | /// Removes the keys from this rate limiter that can be expired
260 | /// safely and returns the keys that were removed.
261 | ///
262 | /// To be eligible for expiration, a key's rate limiter state must
263 | /// be at least `min_age` past its last relevance (see
264 | /// [`RateLimitState.last_touched`](../../algorithms/trait.RateLimitState.html#method.last_touched)).
265 | ///
266 | /// This method works in two parts, but both parts block new keys
267 | /// from getting added while they're running:
268 | /// * First, it collects the keys that are eligible for expiration.
269 | /// * Then, it expires these keys.
270 | ///
271 | /// Note that this only affects new keys that need to be
272 | /// added. Rate-limiting operations on existing keys continue
273 | /// concurrently.
274 | ///
275 | /// # Race conditions
276 | /// Since this is happening concurrently with other operations,
277 | /// race conditions can & will occur. It's possible that cells are
278 | /// accounted between the time `cleanup_at` is called and their
279 | /// expiry. These cells will be lost.
280 | ///
281 | /// The time window in which this can occur is hopefully short
282 | /// enough that this is an acceptable risk of loss in accuracy.
283 | pub fn cleanup>>(&mut self, min_age: D) -> Vec {
284 | self.cleanup_at(min_age, self.clock.now())
285 | }
286 |
287 | /// Removes the keys from this rate limiter that can be expired
288 | /// safely at the given time stamp. See
289 | /// [`cleanup`](#method.cleanup). It returns the list of expired
290 | /// keys.
291 | pub fn cleanup_at>, I: Into>>(
292 | &mut self,
293 | min_age: D,
294 | at: I,
295 | ) -> Vec {
296 | let params = &self.algorithm;
297 | let min_age = min_age.into().unwrap_or_else(|| Duration::new(0, 0));
298 | let at = at.into().unwrap_or_else(|| self.clock.now());
299 |
300 | let mut expireable: Vec = vec![];
301 | self.map_reader.for_each(|k, v| {
302 | if let Some(state) = v.get(0) {
303 | if state
304 | .last_touched(params)
305 | .unwrap_or_else(|| self.clock.now())
306 | < at.saturating_sub(min_age)
307 | {
308 | expireable.push(k.clone());
309 | }
310 | }
311 | });
312 |
313 | // Now take the map write lock and remove all the keys that we
314 | // collected:
315 | let mut w = self.map_writer.lock();
316 | for key in expireable.iter().cloned() {
317 | w.empty(key);
318 | }
319 | w.refresh();
320 | expireable
321 | }
322 | }
323 |
324 | /// A constructor for keyed rate limiters.
325 | pub struct Builder, H: BuildHasher>
326 | {
327 | end_result: PhantomData<(K, A)>,
328 | clock: C,
329 | capacity: NonZeroU32,
330 | cell_weight: NonZeroU32,
331 | per_time_unit: Duration,
332 | hasher: H,
333 | map_capacity: Option,
334 | }
335 |
336 | impl Default for Builder
337 | where
338 | K: Eq + Hash + Clone,
339 | C: clock::Clock,
340 | A: Algorithm,
341 | A::BucketState: KeyableRateLimitState,
342 | {
343 | fn default() -> Builder {
344 | Builder {
345 | end_result: PhantomData,
346 | clock: Default::default(),
347 | map_capacity: None,
348 | capacity: nonzero!(1u32),
349 | cell_weight: nonzero!(1u32),
350 | per_time_unit: Duration::from_secs(1),
351 | hasher: RandomState::new(),
352 | }
353 | }
354 | }
355 |
356 | impl Builder
357 | where
358 | K: Eq + Hash + Clone,
359 | C: clock::Clock,
360 | A: Algorithm,
361 | A::BucketState: KeyableRateLimitState,
362 | H: BuildHasher,
363 | {
364 | /// Sets the hashing method used for the map.
365 | pub fn with_hasher(self, hash_builder: H2) -> Builder {
366 | Builder {
367 | hasher: hash_builder,
368 | clock: Default::default(),
369 | end_result: self.end_result,
370 | capacity: self.capacity,
371 | cell_weight: self.cell_weight,
372 | per_time_unit: self.per_time_unit,
373 | map_capacity: self.map_capacity,
374 | }
375 | }
376 |
377 | /// Sets the "weight" of each cell that is checked against the
378 | /// bucket.
379 | pub fn with_cell_weight(self, cell_weight: NonZeroU32) -> Result {
380 | if self.cell_weight > self.capacity {
381 | return Err(InconsistentCapacity::new(self.capacity, cell_weight));
382 | }
383 | Ok(Builder {
384 | cell_weight,
385 | ..self
386 | })
387 | }
388 |
389 | /// Sets the initial number of keys that the map can hold before
390 | /// rehashing.
391 | pub fn with_map_capacity(self, map_capacity: usize) -> Self {
392 | Builder {
393 | map_capacity: Some(map_capacity),
394 | ..self
395 | }
396 | }
397 |
398 | /// Sets the clock used by the bucket.
399 | pub fn using_clock(self, clock: C) -> Self {
400 | Builder { clock, ..self }
401 | }
402 |
403 | /// Constructs a keyed rate limiter with the given options.
404 | pub fn build(self) -> Result, InconsistentCapacity>
405 | where
406 | H: Clone,
407 | {
408 | let map_opts = evmap::Options::default().with_hasher(self.hasher);
409 | let (r, mut w) = if self.map_capacity.is_some() {
410 | map_opts
411 | .with_capacity(self.map_capacity.unwrap())
412 | .construct()
413 | } else {
414 | map_opts.construct()
415 | };
416 |
417 | w.refresh();
418 | Ok(KeyedRateLimiter {
419 | algorithm: >::construct(
420 | self.capacity,
421 | self.cell_weight,
422 | self.per_time_unit,
423 | )?,
424 | clock: self.clock,
425 | map_reader: r,
426 | map_writer: Arc::new(Mutex::new(w)),
427 | })
428 | }
429 | }
430 |
--------------------------------------------------------------------------------
/src/test_utilities.rs:
--------------------------------------------------------------------------------
1 | #![doc(hidden)]
2 | //! A module for code shared between integration tests & benchmarks in this crate.
3 |
4 | pub mod algorithms;
5 | pub mod variants;
6 |
7 | use crate::lib::*;
8 |
9 | use crate::clock;
10 | use crate::clock::Clock;
11 |
12 | /// Returns a "current" moment that's suitable for tests.
13 | pub fn current_moment() -> ::Instant {
14 | let c = clock::DefaultClock::default();
15 | c.now()
16 | }
17 |
--------------------------------------------------------------------------------
/src/test_utilities/algorithms.rs:
--------------------------------------------------------------------------------
1 | use crate::lib::*;
2 | use crate::{algorithms::Algorithm, clock, NegativeMultiDecision};
3 |
4 | /// A representation of a bare in-memory algorithm, without any bucket
5 | /// attached.
6 | #[derive(Debug)]
7 | pub struct AlgorithmForTest, C: clock::Clock>(A, C);
8 |
9 | impl<'a, A, C> AlgorithmForTest
10 | where
11 | A: Algorithm,
12 | C: clock::Clock,
13 | {
14 | pub fn new>, D: Into>>(
15 | cap: NonZeroU32,
16 | weight: U,
17 | duration: D,
18 | ) -> Self {
19 | AlgorithmForTest(
20 | A::construct(
21 | cap,
22 | weight.into().unwrap_or(nonzero!(1u32)),
23 | duration
24 | .into()
25 | .unwrap_or(crate::lib::Duration::from_secs(1)),
26 | )
27 | .unwrap(),
28 | Default::default(),
29 | )
30 | }
31 |
32 | pub fn algorithm(&'a self) -> &'a A {
33 | &self.0
34 | }
35 |
36 | pub fn state(&self) -> A::BucketState {
37 | A::BucketState::default()
38 | }
39 |
40 | pub fn check(&self, state: &A::BucketState, t0: C::Instant) -> Result<(), A::NegativeDecision> {
41 | self.0.test_and_update(state, t0)
42 | }
43 |
44 | pub fn check_n(
45 | &self,
46 | state: &A::BucketState,
47 | n: u32,
48 | t0: C::Instant,
49 | ) -> Result<(), NegativeMultiDecision> {
50 | self.0.test_n_and_update(state, n, t0)
51 | }
52 | }
53 |
54 | impl Default for AlgorithmForTest
55 | where
56 | A: Algorithm,
57 | C: clock::Clock,
58 | {
59 | fn default() -> Self {
60 | Self::new(nonzero!(1u32), None, None)
61 | }
62 | }
63 |
64 | #[macro_export]
65 | #[doc(hidden)]
66 | macro_rules! bench_with_algorithm_variants {
67 | ($variant:expr, $var:ident, $code:block) => {
68 | match $variant {
69 | $crate::test_utilities::variants::Variant::GCRA => {
70 | let mut $var = $crate::test_utilities::algorithms::AlgorithmForTest::<
71 | $crate::GCRA,
72 | $crate::clock::DefaultClock,
73 | >::default();
74 | $code
75 | }
76 | $crate::test_utilities::variants::Variant::LeakyBucket => {
77 | let mut $var = $crate::test_utilities::algorithms::AlgorithmForTest::<
78 | $crate::LeakyBucket,
79 | $crate::clock::DefaultClock,
80 | >::default();
81 | $code
82 | }
83 | }
84 | };
85 | }
86 |
--------------------------------------------------------------------------------
/src/test_utilities/variants.rs:
--------------------------------------------------------------------------------
1 | use crate::algorithms::Algorithm;
2 | use crate::clock;
3 | use crate::state::DirectRateLimiter;
4 |
5 | #[derive(Debug)]
6 | pub enum Variant {
7 | GCRA,
8 | LeakyBucket,
9 | }
10 |
11 | impl Variant {
12 | pub const ALL: &'static [Variant; 2] = &[Variant::GCRA, Variant::LeakyBucket];
13 | }
14 |
15 | pub struct DirectBucket, C: clock::Clock>(DirectRateLimiter);
16 | impl Default for DirectBucket
17 | where
18 | C: clock::Clock,
19 | A: Algorithm,
20 | {
21 | fn default() -> Self {
22 | DirectBucket(DirectRateLimiter::per_second(nonzero!(50u32)))
23 | }
24 | }
25 | impl DirectBucket
26 | where
27 | C: clock::Clock,
28 | A: Algorithm,
29 | {
30 | pub fn limiter(self) -> DirectRateLimiter {
31 | self.0
32 | }
33 | }
34 |
35 | #[cfg(feature = "std")]
36 | mod std {
37 | use super::*;
38 | use crate::{algorithms::KeyableRateLimitState, clock, KeyedRateLimiter};
39 |
40 | pub struct KeyedBucket, C: clock::Clock>(KeyedRateLimiter)
41 | where
42 | A::BucketState: KeyableRateLimitState;
43 |
44 | impl Default for KeyedBucket
45 | where
46 | A: Algorithm,
47 | A::BucketState: KeyableRateLimitState,
48 | C: clock::Clock,
49 | {
50 | fn default() -> Self {
51 | KeyedBucket(KeyedRateLimiter::per_second(nonzero!(50u32)))
52 | }
53 | }
54 | impl KeyedBucket
55 | where
56 | A: Algorithm,
57 | A::BucketState: KeyableRateLimitState,
58 | C: clock::Clock,
59 | {
60 | pub fn limiter(self) -> KeyedRateLimiter {
61 | self.0
62 | }
63 | }
64 | }
65 | #[cfg(feature = "std")]
66 | pub use self::std::*;
67 |
68 | // I really wish I could just have a function that returns an impl
69 | // Trait that was usable in all the benchmarks, but alas it should not
70 | // be so.
71 | #[doc(hidden)]
72 | #[macro_export]
73 | macro_rules! bench_with_variants {
74 | ($variant:expr, $var:ident : $bucket:tt, $code:block) => {
75 | match $variant {
76 | $crate::test_utilities::variants::Variant::GCRA => {
77 | let mut $var = $bucket::<
78 | ::ratelimit_meter::GCRA<::Instant>,
79 | clock::DefaultClock,
80 | >::default()
81 | .limiter();
82 | $code
83 | }
84 | $crate::test_utilities::variants::Variant::LeakyBucket => {
85 | let mut $var = $bucket::<
86 | ::ratelimit_meter::LeakyBucket<::Instant>,
87 | clock::DefaultClock,
88 | >::default()
89 | .limiter();
90 | $code
91 | }
92 | }
93 | };
94 | }
95 |
--------------------------------------------------------------------------------
/src/thread_safety.rs:
--------------------------------------------------------------------------------
1 | use crate::lib::*;
2 |
3 | #[cfg(feature = "std")]
4 | use parking_lot::Mutex;
5 |
6 | #[cfg(not(feature = "std"))]
7 | use spin::Mutex;
8 |
9 | #[derive(Clone)]
10 | /// Wraps the atomic operations on a Decider's state in a threadsafe
11 | /// fashion.
12 | pub(crate) struct ThreadsafeWrapper
13 | where
14 | T: fmt::Debug + Default + Clone + PartialEq + Eq,
15 | {
16 | data: Arc>,
17 | }
18 |
19 | impl Default for ThreadsafeWrapper
20 | where
21 | T: fmt::Debug + Default + Clone + PartialEq + Eq,
22 | {
23 | fn default() -> Self {
24 | ThreadsafeWrapper {
25 | data: Arc::new(Mutex::new(T::default())),
26 | }
27 | }
28 | }
29 |
30 | impl PartialEq for ThreadsafeWrapper
31 | where
32 | T: fmt::Debug + Default + Clone + PartialEq + Eq,
33 | {
34 | fn eq(&self, other: &Self) -> bool {
35 | if self as *const _ == other as *const _ {
36 | return true;
37 | }
38 | let mine = self.data.lock();
39 | let other = other.data.lock();
40 | *other == *mine
41 | }
42 | }
43 |
44 | impl Eq for ThreadsafeWrapper where T: fmt::Debug + Default + Clone + PartialEq + Eq {}
45 |
46 | impl fmt::Debug for ThreadsafeWrapper
47 | where
48 | T: fmt::Debug + Default + Clone + PartialEq + Eq,
49 | {
50 | fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
51 | let data = self.data.lock();
52 | data.fmt(f)
53 | }
54 | }
55 |
56 | #[cfg(feature = "std")]
57 | mod std {
58 | use super::*;
59 | use evmap::ShallowCopy;
60 |
61 | impl ShallowCopy for ThreadsafeWrapper
62 | where
63 | T: fmt::Debug + Default + Clone + PartialEq + Eq,
64 | {
65 | unsafe fn shallow_copy(&mut self) -> Self {
66 | ThreadsafeWrapper {
67 | data: self.data.shallow_copy(),
68 | }
69 | }
70 | }
71 | }
72 |
73 | impl ThreadsafeWrapper
74 | where
75 | T: fmt::Debug + Default + Clone + PartialEq + Eq,
76 | {
77 | #[inline]
78 | /// Wraps retrieving a bucket's data, calls a function to make a
79 | /// decision and return a new state, and then tries to set the
80 | /// state on the bucket.
81 | ///
82 | /// This function can loop and call the decision closure again if
83 | /// the bucket state couldn't be set.
84 | ///
85 | /// # Panics
86 | /// Panics if an error occurs in acquiring any locks.
87 | pub(crate) fn measure_and_replace(&self, f: F) -> Result<(), E>
88 | where
89 | F: Fn(&T) -> (Result<(), E>, Option),
90 | {
91 | let mut data = self.data.lock();
92 | let (decision, new_data) = f(&*data);
93 | if let Some(new_data) = new_data {
94 | *data = new_data;
95 | }
96 | decision
97 | }
98 |
99 | /// Retrieves and returns a snapshot of the bucket state. This
100 | /// isn't thread safe, but can be used to restore an old copy of
101 | /// the bucket if necessary.
102 | ///
103 | /// # Thread safety
104 | /// This function operates threadsafely, but you're literally
105 | /// taking a copy of data that will change. Relying on the data
106 | /// that is returned *will* race.
107 | pub(crate) fn snapshot(&self) -> T {
108 | let data = self.data.lock();
109 | data.clone()
110 | }
111 | }
112 |
113 | #[cfg(test)]
114 | mod test {
115 | use super::*;
116 |
117 | #[test]
118 | fn no_deadlock_on_eq() {
119 | let wrapper: ThreadsafeWrapper = ThreadsafeWrapper::default();
120 | assert_eq!(wrapper, wrapper);
121 | }
122 | }
123 |
--------------------------------------------------------------------------------
/tests/gcra.rs:
--------------------------------------------------------------------------------
1 | extern crate ratelimit_meter;
2 | #[macro_use]
3 | extern crate nonzero_ext;
4 |
5 | use ratelimit_meter::{
6 | algorithms::Algorithm, test_utilities::current_moment, NegativeMultiDecision, NonConformance,
7 | GCRA,
8 | };
9 | use std::thread;
10 | use std::time::Duration;
11 |
12 | #[test]
13 | fn accepts_first_cell() {
14 | let gcra = GCRA::construct(nonzero!(5u32), nonzero!(1u32), Duration::from_secs(1)).unwrap();
15 | let state = ::BucketState::default();
16 | let now = current_moment();
17 | assert_eq!(Ok(()), gcra.test_and_update(&state, now));
18 | }
19 |
20 | #[test]
21 | fn rejects_too_many() {
22 | let gcra = GCRA::construct(nonzero!(1u32), nonzero!(1u32), Duration::from_secs(1)).unwrap();
23 | let state = ::BucketState::default();
24 | let now = current_moment();
25 | gcra.test_and_update(&state, now).unwrap();
26 | gcra.test_and_update(&state, now).unwrap();
27 | assert_ne!(
28 | Ok(()),
29 | gcra.test_and_update(&state, now),
30 | "{:?} {:?}",
31 | &state,
32 | &gcra
33 | );
34 | }
35 |
36 | #[test]
37 | fn allows_after_interval() {
38 | let gcra = GCRA::construct(nonzero!(1u32), nonzero!(1u32), Duration::from_secs(1)).unwrap();
39 | let state = ::BucketState::default();
40 | let now = current_moment();
41 | let ms = Duration::from_millis(1);
42 | gcra.test_and_update(&state, now).unwrap();
43 | assert_eq!(Ok(()), gcra.test_and_update(&state, now + ms));
44 | assert_ne!(Ok(()), gcra.test_and_update(&state, now + ms * 2));
45 | // should be ok again in 1s:
46 | let next = now + Duration::from_secs(1);
47 | assert_eq!(Ok(()), gcra.test_and_update(&state, next));
48 | }
49 |
50 | #[test]
51 | fn allows_n_after_interval() {
52 | let gcra = GCRA::construct(nonzero!(2u32), nonzero!(1u32), Duration::from_secs(1)).unwrap();
53 | let state = ::BucketState::default();
54 | let now = current_moment() + Duration::from_secs(1);
55 | let ms = Duration::from_millis(1);
56 | assert_eq!(Ok(()), gcra.test_n_and_update(&state, 2, now));
57 | assert!(!gcra.test_n_and_update(&state, 2, now + ms).is_ok());
58 | // should be ok again in 1.5s:
59 | let next = now + Duration::from_secs(1);
60 | assert_eq!(
61 | Ok(()),
62 | gcra.test_n_and_update(&state, 2, next),
63 | "now: {:?}",
64 | next
65 | );
66 |
67 | // should always accommodate 0 cells:
68 | assert_eq!(Ok(()), gcra.test_n_and_update(&state, 0, next));
69 | }
70 |
71 | #[test]
72 | fn correctly_handles_per() {
73 | let ms = Duration::from_millis(1);
74 | let gcra = GCRA::construct(nonzero!(1u32), nonzero!(1u32), ms * 20).unwrap();
75 | let state = ::BucketState::default();
76 | let now = current_moment();
77 |
78 | assert_eq!(Ok(()), gcra.test_and_update(&state, now));
79 | assert_eq!(Ok(()), gcra.test_and_update(&state, now + ms));
80 | assert!(!gcra.test_and_update(&state, now + ms * 10).is_ok());
81 | assert_eq!(Ok(()), gcra.test_and_update(&state, now + ms * 20));
82 | }
83 |
84 | #[test]
85 | fn never_allows_more_than_capacity() {
86 | let ms = Duration::from_millis(1);
87 | let gcra = GCRA::construct(nonzero!(5u32), nonzero!(1u32), Duration::from_secs(1)).unwrap();
88 | let state = ::BucketState::default();
89 | let now = current_moment();
90 |
91 | // Should not allow the first 15 cells on a capacity 5 bucket:
92 | assert!(gcra.test_n_and_update(&state, 15, now).is_err());
93 |
94 | // After 3 and 20 seconds, it should not allow 15 on that bucket either:
95 | assert!(gcra
96 | .test_n_and_update(&state, 15, now + (ms * 3 * 1000))
97 | .is_err());
98 |
99 | let result = gcra.test_n_and_update(&state, 15, now + (ms * 20 * 1000));
100 | match result {
101 | Err(NegativeMultiDecision::InsufficientCapacity(n)) => assert_eq!(n, 15),
102 | _ => panic!("Did not expect {:?}", result),
103 | }
104 | }
105 |
106 | #[test]
107 | fn correct_wait_time() {
108 | // Bucket adding a new element per 200ms:
109 | let gcra = GCRA::construct(nonzero!(5u32), nonzero!(1u32), Duration::from_secs(1)).unwrap();
110 | let state = ::BucketState::default();
111 | let mut now = current_moment();
112 | let ms = Duration::from_millis(1);
113 | let mut conforming = 0;
114 | for _i in 0..20 {
115 | now += ms;
116 | let res = gcra.test_and_update(&state, now);
117 | match res {
118 | Ok(()) => {
119 | conforming += 1;
120 | }
121 | Err(wait) => {
122 | now += wait.wait_time_from(now);
123 | assert_eq!(Ok(()), gcra.test_and_update(&state, now));
124 | conforming += 1;
125 | }
126 | }
127 | }
128 | assert_eq!(20, conforming);
129 | }
130 |
131 | #[test]
132 | fn actual_threadsafety() {
133 | let gcra = GCRA::construct(nonzero!(20u32), nonzero!(1u32), Duration::from_secs(1))
134 | .expect("can't build GCRA");
135 | let state = ::BucketState::default();
136 |
137 | let now = current_moment() + Duration::from_secs(1);
138 | let ms = Duration::from_millis(1);
139 | let mut children = vec![];
140 |
141 | gcra.test_and_update(&state, now)
142 | .expect("first update should work");
143 | for _i in 0..20 {
144 | let state = state.clone();
145 | let gcra = gcra.clone();
146 | children.push(thread::spawn(move || gcra.test_and_update(&state, now)));
147 | }
148 | let results: Vec::NegativeDecision>> = children
149 | .into_iter()
150 | .enumerate()
151 | .map(|(n, c)| c.join().expect(&format!("thread {} panicked", n)))
152 | .collect();
153 | let expected: Vec::NegativeDecision>> =
154 | results.iter().map(|_| Ok(())).collect();
155 | assert_eq!(expected, results);
156 |
157 | assert_ne!(Ok(()), gcra.test_and_update(&state, now + ms * 2));
158 | assert_eq!(Ok(()), gcra.test_and_update(&state, now + ms * 1000));
159 | }
160 |
161 | #[test]
162 | fn nonconformance_wait_time_from() {
163 | let gcra = GCRA::construct(nonzero!(1u32), nonzero!(1u32), Duration::from_secs(1)).unwrap();
164 | let state = ::BucketState::default();
165 | let now = current_moment();
166 | let ms = Duration::from_millis(1);
167 | gcra.test_and_update(&state, now).unwrap();
168 | gcra.test_and_update(&state, now).unwrap();
169 | if let Err(failure) = gcra.test_and_update(&state, now) {
170 | assert_eq!(ms * 2000, failure.wait_time_from(now));
171 | assert_eq!(Duration::new(0, 0), failure.wait_time_from(now + ms * 2000));
172 | assert_eq!(Duration::new(0, 0), failure.wait_time_from(now + ms * 2001));
173 | } else {
174 | assert!(false, "Second attempt should fail");
175 | }
176 | }
177 |
--------------------------------------------------------------------------------
/tests/keyed.rs:
--------------------------------------------------------------------------------
1 | #![cfg(feature = "std")]
2 |
3 | extern crate ratelimit_meter;
4 | #[macro_use]
5 | extern crate nonzero_ext;
6 |
7 | use ratelimit_meter::{KeyedRateLimiter, GCRA};
8 | use std::thread;
9 | use std::time::{Duration, Instant};
10 |
11 | #[test]
12 | fn different_states_per_key() {
13 | let mut lim = KeyedRateLimiter::<&str>::new(nonzero!(1u32), Duration::from_secs(1));
14 | let ms = Duration::from_millis(1);
15 | let now = Instant::now();
16 | assert_eq!(Ok(()), lim.check_at("foo", now + ms));
17 | assert_eq!(Ok(()), lim.check_at("bar", now + ms));
18 | assert_eq!(Ok(()), lim.check_at("baz", now + ms));
19 |
20 | assert_ne!(Ok(()), lim.check_at("foo", now + ms * 3), "{:?}", lim);
21 | assert_ne!(Ok(()), lim.check_at("bar", now + ms * 3), "{:?}", lim);
22 | assert_ne!(Ok(()), lim.check_at("baz", now + ms * 3), "{:?}", lim);
23 | }
24 |
25 | #[test]
26 | fn expiration() {
27 | let ms = Duration::from_millis(1);
28 | let now = Instant::now();
29 | let then = now + ms * 2000; // two seconds later
30 |
31 | fn make_bucket<'a>() -> KeyedRateLimiter<&'a str> {
32 | let ms = Duration::from_millis(1);
33 | let now = Instant::now();
34 | let mut lim = KeyedRateLimiter::<&str>::new(nonzero!(1u32), Duration::from_secs(1));
35 | lim.check_at("foo", now).unwrap();
36 | lim.check_at("bar", now + ms * 200).unwrap();
37 | lim.check_at("baz", now + ms * 800).unwrap();
38 | lim
39 | }
40 |
41 | // clean up all keys that are indistinguishable from unoccupied keys:
42 | let mut lim = make_bucket();
43 | let mut removed = lim.cleanup_at(None, then);
44 | removed.sort();
45 | assert_eq!(vec!["bar", "baz", "foo"], removed);
46 | assert_eq!(0, lim.len());
47 | assert!(lim.is_empty());
48 |
49 | // clean up all keys that have been so for 300ms:
50 | let mut lim = make_bucket();
51 | let mut removed = lim.cleanup_at(Some(Duration::from_millis(300)), then);
52 | removed.sort();
53 | assert_eq!(vec!["bar", "foo"], removed);
54 | assert_eq!(1, lim.len());
55 | assert!(!lim.is_empty());
56 |
57 | // clean up 2 seconds plus change later:
58 | let mut lim = make_bucket();
59 | let mut removed = lim.cleanup_at(Some(Duration::from_secs(1)), now + ms * 2100);
60 | removed.sort();
61 | assert_eq!(vec!["foo"], removed);
62 | assert_eq!(2, lim.len());
63 | assert!(!lim.is_empty());
64 | }
65 |
66 | #[test]
67 | fn actual_threadsafety() {
68 | let mut lim = KeyedRateLimiter::<&str, GCRA>::new(nonzero!(20u32), Duration::from_secs(1));
69 | let now = Instant::now();
70 | let ms = Duration::from_millis(1);
71 | let mut children = vec![];
72 |
73 | lim.check_at("foo", now).unwrap();
74 | for _i in 0..20 {
75 | let mut lim = lim.clone();
76 | children.push(thread::spawn(move || {
77 | lim.check_at("foo", now).unwrap();
78 | }));
79 | }
80 | for child in children {
81 | child.join().unwrap();
82 | }
83 | assert!(!lim.check_at("foo", now + ms * 2).is_ok());
84 | assert_eq!(Ok(()), lim.check_at("foo", now + ms * 1000));
85 | }
86 |
--------------------------------------------------------------------------------
/tests/leaky_bucket.rs:
--------------------------------------------------------------------------------
1 | extern crate ratelimit_meter;
2 | #[macro_use]
3 | extern crate nonzero_ext;
4 |
5 | use ratelimit_meter::{
6 | algorithms::Algorithm, test_utilities::current_moment, DirectRateLimiter, LeakyBucket,
7 | NegativeMultiDecision, NonConformance,
8 | };
9 | use std::thread;
10 | use std::time::Duration;
11 |
12 | #[test]
13 | fn accepts_first_cell() {
14 | let mut lb = DirectRateLimiter::::per_second(nonzero!(5u32));
15 | assert_eq!(Ok(()), lb.check_at(current_moment()));
16 | }
17 |
18 | #[test]
19 | fn rejects_too_many() {
20 | let mut lb = DirectRateLimiter::::per_second(nonzero!(2u32));
21 | let now = current_moment();
22 | let ms = Duration::from_millis(1);
23 | assert_eq!(Ok(()), lb.check_at(now));
24 | assert_eq!(Ok(()), lb.check_at(now));
25 |
26 | assert_ne!(Ok(()), lb.check_at(now + ms * 2));
27 |
28 | // should be ok again in 1s:
29 | let next = now + Duration::from_millis(1002);
30 | assert_eq!(Ok(()), lb.check_at(next));
31 | assert_eq!(Ok(()), lb.check_at(next + ms));
32 |
33 | assert_ne!(Ok(()), lb.check_at(next + ms * 2), "{:?}", lb);
34 | }
35 |
36 | #[test]
37 | fn never_allows_more_than_capacity() {
38 | let mut lb = DirectRateLimiter::::per_second(nonzero!(5u32));
39 | let now = current_moment();
40 | let ms = Duration::from_millis(1);
41 |
42 | // Should not allow the first 15 cells on a capacity 5 bucket:
43 | assert_ne!(Ok(()), lb.check_n_at(15, now));
44 |
45 | // After 3 and 20 seconds, it should not allow 15 on that bucket either:
46 | assert_ne!(Ok(()), lb.check_n_at(15, now + (ms * 3 * 1000)));
47 | let result = lb.check_n_at(15, now + (ms * 20 * 1000));
48 | match result {
49 | Err(NegativeMultiDecision::InsufficientCapacity(n)) => assert_eq!(n, 15),
50 | _ => panic!("Did not expect {:?}", result),
51 | }
52 | }
53 |
54 | #[test]
55 | fn correct_wait_time() {
56 | // Bucket adding a new element per 200ms:
57 | let mut lb = DirectRateLimiter::::per_second(nonzero!(5u32));
58 | let mut now = current_moment();
59 | let ms = Duration::from_millis(1);
60 | let mut conforming = 0;
61 | for _i in 0..20 {
62 | now += ms;
63 | let res = lb.check_at(now);
64 | match res {
65 | Ok(()) => {
66 | conforming += 1;
67 | }
68 | Err(wait) => {
69 | now += wait.wait_time_from(now);
70 | assert_eq!(Ok(()), lb.check_at(now));
71 | conforming += 1;
72 | }
73 | }
74 | }
75 | assert_eq!(20, conforming);
76 | }
77 |
78 | #[test]
79 | fn prevents_time_travel() {
80 | let mut lb = DirectRateLimiter::::per_second(nonzero!(5u32));
81 | let now = current_moment() + Duration::from_secs(1);
82 | let ms = Duration::from_millis(1);
83 |
84 | assert!(lb.check_at(now).is_ok());
85 | assert!(lb.check_at(now - ms).is_ok());
86 | assert!(lb.check_at(now - ms * 500).is_ok());
87 | }
88 |
89 | #[test]
90 | fn actual_threadsafety() {
91 | let mut lim = DirectRateLimiter::::per_second(nonzero!(20u32));
92 | let now = current_moment();
93 | let ms = Duration::from_millis(1);
94 | let mut children = vec![];
95 |
96 | lim.check_at(now).unwrap();
97 | for _i in 0..20 {
98 | let mut lim = lim.clone();
99 | children.push(thread::spawn(move || lim.check_at(now).is_ok()));
100 | }
101 | for child in children {
102 | child.join().unwrap();
103 | }
104 | assert!(!lim.check_at(now + ms * 2).is_ok());
105 | assert_eq!(Ok(()), lim.check_at(now + ms * 1000));
106 | }
107 |
108 | #[test]
109 | fn tooearly_wait_time_from() {
110 | let lim =
111 | LeakyBucket::construct(nonzero!(1u32), nonzero!(1u32), Duration::from_secs(1)).unwrap();
112 | let state = ::BucketState::default();
113 | let now = current_moment();
114 | let ms = Duration::from_millis(1);
115 | lim.test_and_update(&state, now).unwrap();
116 | if let Err(failure) = lim.test_and_update(&state, now) {
117 | assert_eq!(ms * 1000, failure.wait_time_from(now));
118 | assert_eq!(Duration::new(0, 0), failure.wait_time_from(now + ms * 1000));
119 | assert_eq!(Duration::new(0, 0), failure.wait_time_from(now + ms * 2001));
120 | } else {
121 | assert!(false, "Second attempt should fail");
122 | }
123 | }
124 |
--------------------------------------------------------------------------------
/tests/memory.rs:
--------------------------------------------------------------------------------
1 | #![cfg(feature = "std")]
2 |
3 | // This test uses procinfo, so can only be run on Linux.
4 | extern crate libc;
5 | extern crate ratelimit_meter;
6 | #[macro_use]
7 | extern crate nonzero_ext;
8 |
9 | use ratelimit_meter::{DirectRateLimiter, LeakyBucket, GCRA};
10 | use std::thread;
11 |
12 | fn resident_memsize() -> i64 {
13 | let mut out: libc::rusage = unsafe { std::mem::zeroed() };
14 | assert!(unsafe { libc::getrusage(libc::RUSAGE_SELF, &mut out) } == 0);
15 | out.ru_maxrss
16 | }
17 |
18 | const LEAK_TOLERANCE: i64 = 1024 * 1024 * 10;
19 |
20 | struct LeakCheck {
21 | usage_before: i64,
22 | n_iter: usize,
23 | }
24 |
25 | impl Drop for LeakCheck {
26 | fn drop(&mut self) {
27 | let usage_after = resident_memsize();
28 | assert!(
29 | usage_after <= self.usage_before + LEAK_TOLERANCE,
30 | "Plausible memory leak!\nAfter {} iterations, usage before: {}, usage after: {}",
31 | self.n_iter,
32 | self.usage_before,
33 | usage_after
34 | );
35 | }
36 | }
37 |
38 | impl LeakCheck {
39 | fn new(n_iter: usize) -> Self {
40 | LeakCheck {
41 | n_iter,
42 | usage_before: resident_memsize(),
43 | }
44 | }
45 | }
46 |
47 | #[test]
48 | fn memleak_gcra() {
49 | let mut bucket = DirectRateLimiter::::build_with_capacity(nonzero!(1_000_000u32))
50 | .build()
51 | .unwrap();
52 | let leak_check = LeakCheck::new(500_000);
53 |
54 | for _i in 0..leak_check.n_iter {
55 | drop(bucket.check());
56 | }
57 | }
58 |
59 | #[test]
60 | fn memleak_gcra_multi() {
61 | let mut bucket = DirectRateLimiter::::build_with_capacity(nonzero!(1_000_000u32))
62 | .build()
63 | .unwrap();
64 | let leak_check = LeakCheck::new(500_000);
65 |
66 | for _i in 0..leak_check.n_iter {
67 | drop(bucket.check_n(2));
68 | }
69 | }
70 |
71 | #[test]
72 | fn memleak_gcra_threaded() {
73 | let bucket = DirectRateLimiter::::build_with_capacity(nonzero!(1_000_000u32))
74 | .build()
75 | .unwrap();
76 | let leak_check = LeakCheck::new(5_000);
77 |
78 | for _i in 0..leak_check.n_iter {
79 | let mut bucket = bucket.clone();
80 | thread::spawn(move || drop(bucket.check())).join().unwrap();
81 | }
82 | }
83 |
84 | #[test]
85 | fn memleak_leakybucket() {
86 | let mut bucket = DirectRateLimiter::::per_second(nonzero!(1_000_000u32));
87 | let leak_check = LeakCheck::new(500_000);
88 |
89 | for _i in 0..leak_check.n_iter {
90 | drop(bucket.check());
91 | }
92 | }
93 |
94 | #[test]
95 | fn memleak_leakybucket_threaded() {
96 | let bucket = DirectRateLimiter::::per_second(nonzero!(1_000_000u32));
97 | let leak_check = LeakCheck::new(5_000);
98 |
99 | for _i in 0..leak_check.n_iter {
100 | let mut bucket = bucket.clone();
101 | thread::spawn(move || drop(bucket.check())).join().unwrap();
102 | }
103 | }
104 |
--------------------------------------------------------------------------------