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
12pub 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 return v.measure_and_replace_one(f);
25 }
26 let entry = self.entry(key.clone()).or_default();
28 (*entry).measure_and_replace_one(f)
29 }
30}
31
32impl<K, C> RateLimiter<K, DashMapStateStore<K>, C, NoOpMiddleware<C::Instant>>
34where
35 K: Hash + Eq + Clone,
36 C: clock::Clock,
37{
38 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}