sf_rate_limiter/
storage.rs

1use hashbrown::HashMap;
2use parking_lot::Mutex;
3use std::any::Any;
4use std::marker::PhantomData;
5
6pub trait Storage<Inner, S: State<Inner>> {
7    fn fetch(&self, key: &str) -> Option<S>;
8
9    fn save<IntoString: Into<String>>(&mut self, key: IntoString, value: S);
10}
11
12pub trait State<Body>: Clone {
13    fn get_id(&self) -> String;
14
15    fn get_expiration_time(&self) -> usize;
16}
17
18pub struct InMemoryStorage<A: Sized, S: State<A>> {
19    store: HashMap<String, Mutex<S>>,
20    _phantom_data: PhantomData<A>,
21}
22
23impl<A: Sized, S: State<A>> InMemoryStorage<A, S> {
24    pub fn new() -> Self {
25        Self {
26            store: HashMap::new(),
27            _phantom_data: Default::default(),
28        }
29    }
30}
31
32impl<A: Sized, S: State<A>> Storage<A, S> for InMemoryStorage<A, S> {
33    fn fetch(&self, key: &str) -> Option<S> {
34        if let Some(value) = self.store.get(key) {
35            return Some(value.lock().clone());
36        }
37
38        None
39    }
40
41    fn save<IntoString: Into<String>>(&mut self, key: IntoString, value: S) {
42        let key = key.into();
43
44        if self.store.contains_key(&key) {
45            let mut state = self.store.get(&key).unwrap().lock();
46            *state = value;
47        } else {
48            self.store.insert(key, Mutex::new(value));
49        }
50    }
51}