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::keyed::DefaultHasher;
7use crate::state::{InMemoryState, StateStore};
8use crate::{clock, Quota, RateLimiter};
9use crate::{middleware::NoOpMiddleware, state::keyed::ShrinkableKeyedStateStore};
10use core::hash::Hash;
11use dashmap::DashMap;
12
13/// A concurrent, thread-safe and fairly performant hashmap based on [`DashMap`].
14pub type DashMapStateStore<K, S = DefaultHasher> = DashMap<K, InMemoryState, S>;
15
16impl<K: Hash + Eq + Clone, S: core::hash::BuildHasher + Clone> StateStore
17    for DashMapStateStore<K, S>
18{
19    type Key = K;
20
21    fn measure_and_replace<T, F, E>(&self, key: &Self::Key, f: F) -> Result<T, E>
22    where
23        F: Fn(Option<Nanos>) -> Result<(T, Nanos), E>,
24    {
25        if let Some(v) = self.get(key) {
26            // fast path: measure existing entry
27            return v.measure_and_replace_one(f);
28        }
29        // make an entry and measure that:
30        let entry = self.entry(key.clone()).or_default();
31        (*entry).measure_and_replace_one(f)
32    }
33}
34
35/// # Keyed rate limiters - [`DashMap`]-backed with a default hasher
36impl<K, C> RateLimiter<K, DashMapStateStore<K>, C, NoOpMiddleware<C::Instant>>
37where
38    K: Hash + Eq + Clone,
39    C: clock::Clock,
40{
41    /// Constructs a new rate limiter with a custom clock, backed by a
42    /// [`DashMap`] with the default hasher.
43    pub fn dashmap_with_clock(quota: Quota, clock: C) -> Self {
44        let state: DashMapStateStore<K> = DashMap::default();
45        RateLimiter::new(quota, state, clock)
46    }
47}
48
49/// # Keyed rate limiters - [`DashMap`]-backed with a custom hasher
50impl<K, S, C> RateLimiter<K, DashMapStateStore<K, S>, C, NoOpMiddleware<C::Instant>>
51where
52    K: Hash + Eq + Clone,
53    S: core::hash::BuildHasher + Default + Clone,
54    C: clock::Clock,
55{
56    /// Constructs a new rate limiter with a custom clock and hasher, backed by a
57    /// [`DashMap`].
58    pub fn dashmap_with_clock_and_hasher(quota: Quota, clock: C, hasher: S) -> Self {
59        let state: DashMapStateStore<K, S> = DashMap::with_hasher(hasher);
60        RateLimiter::new(quota, state, clock)
61    }
62}
63
64impl<K: Hash + Eq + Clone, S: core::hash::BuildHasher + Clone> ShrinkableKeyedStateStore<K>
65    for DashMapStateStore<K, S>
66{
67    fn retain_recent(&self, drop_below: Nanos) {
68        self.retain(|_, v| !v.is_older_than(drop_below));
69    }
70
71    fn shrink_to_fit(&self) {
72        self.shrink_to_fit();
73    }
74
75    fn len(&self) -> usize {
76        self.len()
77    }
78
79    fn is_empty(&self) -> bool {
80        self.is_empty()
81    }
82}