governor/state/
keyed.rs

Help
1//! Keyed rate limiters (those that can hold one state per key).
2//!
3//! These are rate limiters that have one set of parameters (burst capacity per time period) but
4//! apply those to several sets of actual rate-limiting states, e.g. to enforce one API call rate
5//! limit per API key.
6//!
7//! Rate limiters based on these types are constructed with
8//! [the `RateLimiter` constructors](../struct.RateLimiter.html#keyed-rate-limiters---default-constructors)
9
10use std::hash::Hash;
11use std::num::NonZeroU32;
12use std::prelude::v1::*;
13
14use crate::state::StateStore;
15use crate::{
16    clock::{self, Reference},
17    errors::InsufficientCapacity,
18    middleware::RateLimitingMiddleware,
19    nanos::Nanos,
20    Quota, RateLimiter,
21};
22
23/// A trait for state stores with one rate limiting state per key.
24///
25/// This is blanket-implemented by all [`StateStore`]s with hashable (`Eq + Hash + Clone`) key
26/// associated types.
27pub trait KeyedStateStore<K: Hash>: StateStore<Key = K> {}
28
29impl<T, K: Hash> KeyedStateStore<K> for T
30where
31    T: StateStore<Key = K>,
32    K: Eq + Clone + Hash,
33{
34}
35
36/// # Keyed rate limiters - default constructors
37impl<K> RateLimiter<K, DefaultKeyedStateStore<K>, clock::DefaultClock>
38where
39    K: Clone + Hash + Eq,
40{
41    /// Constructs a new keyed rate limiter backed by
42    /// the [`DefaultKeyedStateStore`].
43    pub fn keyed(quota: Quota) -> Self {
44        let state = DefaultKeyedStateStore::default();
45        let clock = clock::DefaultClock::default();
46        RateLimiter::new(quota, state, &clock)
47    }
48
49    #[cfg(all(feature = "std", feature = "dashmap"))]
50    /// Constructs a new keyed rate limiter explicitly backed by a [`DashMap`][::dashmap::DashMap].
51    pub fn dashmap(quota: Quota) -> Self {
52        let state = DashMapStateStore::default();
53        let clock = clock::DefaultClock::default();
54        RateLimiter::new(quota, state, &clock)
55    }
56
57    #[cfg(any(all(feature = "std", not(feature = "dashmap")), not(feature = "std")))]
58    /// Constructs a new keyed rate limiter explicitly backed by a
59    /// [`HashMap`][std::collections::HashMap].
60    pub fn hashmap(quota: Quota) -> Self {
61        let state = HashMapStateStore::default();
62        let clock = clock::DefaultClock::default();
63        RateLimiter::new(quota, state, &clock)
64    }
65}
66
67#[cfg(all(feature = "std", feature = "dashmap"))]
68impl<K> RateLimiter<K, HashMapStateStore<K>, clock::DefaultClock>
69where
70    K: Clone + Hash + Eq,
71{
72    /// Constructs a new keyed rate limiter explicitly backed by a
73    /// [`HashMap`][std::collections::HashMap].
74    pub fn hashmap(quota: Quota) -> Self {
75        let state = HashMapStateStore::default();
76        let clock = clock::DefaultClock::default();
77        RateLimiter::new(quota, state, &clock)
78    }
79}
80
81/// # Keyed rate limiters - Manually checking cells
82impl<K, S, C, MW> RateLimiter<K, S, C, MW>
83where
84    S: KeyedStateStore<K>,
85    K: Hash,
86    C: clock::Clock,
87    MW: RateLimitingMiddleware<C::Instant>,
88{
89    /// Allow a single cell through the rate limiter for the given key.
90    ///
91    /// If the rate limit is reached, `check_key` returns information about the earliest
92    /// time that a cell might be allowed through again under that key.
93    pub fn check_key(&self, key: &K) -> Result<MW::PositiveOutcome, MW::NegativeOutcome> {
94        self.gcra.test_and_update::<K, C::Instant, S, MW>(
95            self.start,
96            key,
97            &self.state,
98            self.clock.now(),
99        )
100    }
101
102    /// Allow *only all* `n` cells through the rate limiter for the given key.
103    ///
104    /// This method can succeed in only one way and fail in two ways:
105    /// * Success: If all `n` cells can be accommodated, it returns `Ok(Ok(()))`.
106    /// * Failure (but ok): Not all cells can make it through at the current time.
107    ///   The result is `Ok(Err(NotUntil))`, which can
108    ///   be interrogated about when the batch might next conform.
109    /// * Failure (the batch can never go through): The rate limit is too low for the given number
110    ///   of cells. The result is `Err(InsufficientCapacity)`
111    ///
112    /// ### Performance
113    /// This method diverges a little from the GCRA algorithm, using
114    /// multiplication to determine the next theoretical arrival time, and so
115    /// is not as fast as checking a single cell.
116    pub fn check_key_n(
117        &self,
118        key: &K,
119        n: NonZeroU32,
120    ) -> Result<Result<MW::PositiveOutcome, MW::NegativeOutcome>, InsufficientCapacity> {
121        self.gcra.test_n_all_and_update::<K, C::Instant, S, MW>(
122            self.start,
123            key,
124            n,
125            &self.state,
126            self.clock.now(),
127        )
128    }
129}
130
131/// Keyed rate limiters that can be "cleaned up".
132///
133/// Any keyed state store implementing this trait allows users to evict elements that are
134/// indistinguishable from fresh rate-limiting states (that is, if a key hasn't been used for
135/// rate-limiting decisions for as long as the bucket capacity).
136///
137/// As this does not make sense for not all keyed state stores (e.g. stores that auto-expire like
138/// memcache), this is an optional trait. All the keyed state stores in this crate implement
139/// shrinking.
140pub trait ShrinkableKeyedStateStore<K: Hash>: KeyedStateStore<K> {
141    /// Remove those keys with state older than `drop_below`.
142    fn retain_recent(&self, drop_below: Nanos);
143
144    /// Shrinks the capacity of the state store, if possible.
145    ///
146    /// If the state store does not support shrinking, this method is a no-op.
147    fn shrink_to_fit(&self) {}
148
149    /// Returns the number of "live" keys stored in the state store.
150    ///
151    /// Depending on how the state store is implemented, this may
152    /// return an estimate or an out-of-date result.
153    fn len(&self) -> usize;
154
155    /// Returns `true` if `self` has no keys stored in it.
156    ///
157    /// As with [`len`](#tymethod.len), this method may return
158    /// imprecise results (indicating that the state store is empty
159    /// while a concurrent rate-limiting operation is taking place).
160    fn is_empty(&self) -> bool;
161}
162
163/// # Keyed rate limiters - Housekeeping
164///
165/// As the inputs to a keyed rate-limiter can be arbitrary keys, the set of retained keys retained
166/// grows, while the number of active keys may stay smaller. To save on space, a keyed rate-limiter
167/// allows removing those keys that are "stale", i.e., whose values are no different from keys' that
168/// aren't present in the rate limiter state store.
169impl<K, S, C, MW> RateLimiter<K, S, C, MW>
170where
171    S: ShrinkableKeyedStateStore<K>,
172    K: Hash,
173    C: clock::Clock,
174    MW: RateLimitingMiddleware<C::Instant>,
175{
176    /// Retains all keys in the rate limiter that were used recently enough.
177    ///
178    /// Any key whose rate limiting state is indistinguishable from a "fresh" state (i.e., the
179    /// theoretical arrival time lies in the past).
180    pub fn retain_recent(&self) {
181        // calculate the minimum retention parameter: Any key whose state store's theoretical
182        // arrival time is larger than a starting state for the bucket gets to stay, everything
183        // else (that's indistinguishable from a starting state) goes.
184        let now = self.clock.now();
185        let drop_below = now.duration_since(self.start);
186
187        self.state.retain_recent(drop_below);
188    }
189
190    /// Shrinks the capacity of the rate limiter's state store, if possible.
191    pub fn shrink_to_fit(&self) {
192        self.state.shrink_to_fit();
193    }
194
195    /// Returns the number of "live" keys in the rate limiter's state store.
196    ///
197    /// Depending on how the state store is implemented, this may
198    /// return an estimate or an out-of-date result.
199    pub fn len(&self) -> usize {
200        self.state.len()
201    }
202
203    /// Returns `true` if the rate limiter has no keys in it.
204    ///
205    /// As with [`len`](#method.len), this method may return
206    /// imprecise results (indicating that the state store is empty
207    /// while a concurrent rate-limiting operation is taking place).
208    pub fn is_empty(&self) -> bool {
209        self.state.is_empty()
210    }
211}
212
213mod hashmap;
214
215pub use hashmap::HashMapStateStore;
216
217#[cfg(all(feature = "std", feature = "dashmap"))]
218mod dashmap;
219
220#[cfg(all(feature = "std", feature = "dashmap"))]
221pub use self::dashmap::DashMapStateStore;
222
223#[cfg(feature = "std")]
224mod future;
225
226#[cfg(any(all(feature = "std", not(feature = "dashmap")), not(feature = "std")))]
227/// The default keyed rate limiter type: a mutex-wrapped [`HashMap`][std::collections::HashMap].
228pub type DefaultKeyedStateStore<K> = HashMapStateStore<K>;
229
230#[cfg(all(feature = "std", feature = "dashmap"))]
231/// The default keyed rate limiter type: the concurrent [`DashMap`][::dashmap::DashMap].
232pub type DefaultKeyedStateStore<K> = DashMapStateStore<K>;
233
234#[cfg(test)]
235mod test {
236    use std::marker::PhantomData;
237
238    use nonzero_ext::nonzero;
239
240    use crate::{
241        clock::{Clock, FakeRelativeClock},
242        middleware::NoOpMiddleware,
243    };
244
245    use super::*;
246
247    #[test]
248    fn default_nonshrinkable_state_store_coverage() {
249        #[derive(Default)]
250        struct NaiveKeyedStateStore<K>(PhantomData<K>);
251
252        impl<K: Hash + Eq + Clone> StateStore for NaiveKeyedStateStore<K> {
253            type Key = K;
254
255            fn measure_and_replace<T, F, E>(&self, _key: &Self::Key, f: F) -> Result<T, E>
256            where
257                F: Fn(Option<Nanos>) -> Result<(T, Nanos), E>,
258            {
259                f(None).map(|(res, _)| res)
260            }
261        }
262
263        impl<K: Hash + Eq + Clone> ShrinkableKeyedStateStore<K> for NaiveKeyedStateStore<K> {
264            fn retain_recent(&self, _drop_below: Nanos) {
265                // nothing to do
266            }
267
268            fn len(&self) -> usize {
269                0
270            }
271            fn is_empty(&self) -> bool {
272                true
273            }
274        }
275
276        let lim: RateLimiter<
277            u32,
278            NaiveKeyedStateStore<u32>,
279            FakeRelativeClock,
280            NoOpMiddleware<<FakeRelativeClock as Clock>::Instant>,
281        > = RateLimiter::new(
282            Quota::per_second(nonzero!(1_u32)),
283            NaiveKeyedStateStore::default(),
284            &FakeRelativeClock::default(),
285        );
286        assert_eq!(lim.check_key(&1u32), Ok(()));
287        assert!(lim.is_empty());
288        assert_eq!(lim.len(), 0);
289        lim.retain_recent();
290        lim.shrink_to_fit();
291    }
292}
governor/state/keyed.rs

governor/state/
keyed.rs
