governor/state/keyed/
dashmap.rs

1#![cfg(all(feature = "std", feature = "dashmap"))]
2
3use std::prelude::v1::*;
4
5use crate::nanos::Nanos;
6use crate::state::{InMemoryState, StateStore};
7use crate::{clock, Quota, RateLimiter};
8use crate::{middleware::NoOpMiddleware, state::keyed::ShrinkableKeyedStateStore};
9use dashmap::DashMap;
10use std::hash::Hash;
11
12/// A concurrent, thread-safe and fairly performant hashmap based on [`DashMap`].
13pub type DashMapStateStore<K> = DashMap<K, InMemoryState>;
14
15impl<K: Hash + Eq + Clone> StateStore for DashMapStateStore<K> {
16    type Key = K;
17
18    fn measure_and_replace<T, F, E>(&self, key: &Self::Key, f: F) -> Result<T, E>
19    where
20        F: Fn(Option<Nanos>) -> Result<(T, Nanos), E>,
21    {
22        if let Some(v) = self.get(key) {
23            // fast path: measure existing entry
24            return v.measure_and_replace_one(f);
25        }
26        // make an entry and measure that:
27        let entry = self.entry(key.clone()).or_default();
28        (*entry).measure_and_replace_one(f)
29    }
30}
31
32/// # Keyed rate limiters - [`DashMap`]-backed
33impl<K, C> RateLimiter<K, DashMapStateStore<K>, C, NoOpMiddleware<C::Instant>>
34where
35    K: Hash + Eq + Clone,
36    C: clock::Clock,
37{
38    /// Constructs a new rate limiter with a custom clock, backed by a
39    /// [`DashMap`].
40    pub fn dashmap_with_clock(quota: Quota, clock: &C) -> Self {
41        let state: DashMapStateStore<K> = DashMap::default();
42        RateLimiter::new(quota, state, clock)
43    }
44}
45
46impl<K: Hash + Eq + Clone> ShrinkableKeyedStateStore<K> for DashMapStateStore<K> {
47    fn retain_recent(&self, drop_below: Nanos) {
48        self.retain(|_, v| !v.is_older_than(drop_below));
49    }
50
51    fn shrink_to_fit(&self) {
52        self.shrink_to_fit();
53    }
54
55    fn len(&self) -> usize {
56        self.len()
57    }
58
59    fn is_empty(&self) -> bool {
60        self.is_empty()
61    }
62}