governor/state/keyed/
hashmap.rs

1use crate::nanos::Nanos;
2use crate::{clock, Quota, RateLimiter};
3use crate::{
4    middleware::NoOpMiddleware,
5    state::{InMemoryState, StateStore},
6};
7use core::hash::Hash;
8
9#[cfg(feature = "no_std")]
10pub use hashbrown::HashMap;
11#[cfg(not(feature = "no_std"))]
12pub use std::collections::HashMap;
13
14use crate::state::keyed::{DefaultHasher, ShrinkableKeyedStateStore};
15
16#[cfg(feature = "std")]
17type Mutex<T> = parking_lot::Mutex<T>;
18
19#[cfg(not(feature = "std"))]
20type Mutex<T> = spinning_top::Spinlock<T>;
21
22/// A thread-safe (but not very performant) implementation of a keyed rate limiter state
23/// store using [`HashMap`].
24///
25/// The `HashMapStateStore` is the default state store in `std` when no other thread-safe
26/// features are enabled.
27pub type HashMapStateStore<K, S = DefaultHasher> = Mutex<HashMap<K, InMemoryState, S>>;
28
29impl<K: Hash + Eq + Clone, S: core::hash::BuildHasher> StateStore for HashMapStateStore<K, S> {
30    type Key = K;
31
32    fn measure_and_replace<T, F, E>(&self, key: &Self::Key, f: F) -> Result<T, E>
33    where
34        F: Fn(Option<Nanos>) -> Result<(T, Nanos), E>,
35    {
36        let mut map = self.lock();
37        if let Some(v) = (*map).get(key) {
38            // fast path: a rate limiter is already present for the key.
39            return v.measure_and_replace_one(f);
40        }
41        // not-so-fast path: make a new entry and measure it.
42        let entry = (*map).entry(key.clone()).or_default();
43        entry.measure_and_replace_one(f)
44    }
45}
46
47impl<K: Hash + Eq + Clone, S: core::hash::BuildHasher> ShrinkableKeyedStateStore<K>
48    for HashMapStateStore<K, S>
49{
50    fn retain_recent(&self, drop_below: Nanos) {
51        let mut map = self.lock();
52        map.retain(|_, v| !v.is_older_than(drop_below));
53    }
54
55    fn shrink_to_fit(&self) {
56        let mut map = self.lock();
57        map.shrink_to_fit();
58    }
59
60    fn len(&self) -> usize {
61        let map = self.lock();
62        (*map).len()
63    }
64    fn is_empty(&self) -> bool {
65        let map = self.lock();
66        (*map).is_empty()
67    }
68}
69
70/// # Keyed rate limiters - [`HashMap`]-backed with a default hasher
71impl<K, C> RateLimiter<K, HashMapStateStore<K>, C, NoOpMiddleware<C::Instant>>
72where
73    K: Hash + Eq + Clone,
74    C: clock::Clock,
75{
76    /// Constructs a new rate limiter with a custom clock, backed by a [`HashMap`] with the default hasher.
77    pub fn hashmap_with_clock(quota: Quota, clock: C) -> Self {
78        let state: HashMapStateStore<K> = HashMapStateStore::new(HashMap::default());
79        RateLimiter::new(quota, state, clock)
80    }
81}
82
83/// # Keyed rate limiters - [`HashMap`]-backed with a custom hasher
84impl<K, S, C> RateLimiter<K, HashMapStateStore<K, S>, C, NoOpMiddleware<C::Instant>>
85where
86    K: Hash + Eq + Clone,
87    S: core::hash::BuildHasher,
88    C: clock::Clock,
89{
90    /// Constructs a new rate limiter with a custom clock and hasher, backed by a [`HashMap`].
91    pub fn hashmap_with_clock_and_hasher(quota: Quota, clock: C, hasher: S) -> Self {
92        let state: HashMapStateStore<K, S> = HashMapStateStore::new(HashMap::with_hasher(hasher));
93        RateLimiter::new(quota, state, clock)
94    }
95}
governor/state/keyed/hashmap.rs

governor/state/keyed/
hashmap.rs
