governor/state/keyed/
hashmap.rs

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

governor/state/keyed/
hashmap.rs
