actix_rl/store/
mem_store.rs

1use std::collections::HashMap;
2use std::ops::{Deref, DerefMut};
3use std::sync::Arc;
4use chrono::{DateTime, Utc};
5use tokio::sync::Mutex;
6use crate::store::{Store, Value};
7
8pub const DEFAULT_STORE_CAPACITY: usize = 4096;
9
10/// [DateCount] stores the creation time and the current count.
11#[derive(Debug, Clone, Copy)]
12pub struct DateCount {
13    pub create_date: DateTime<Utc>,
14    pub count: u32,
15}
16
17impl Default for DateCount {
18    fn default() -> Self {
19        Self {
20            create_date: Utc::now(),
21            count: 0u32,
22        }
23    }
24}
25
26impl DateCount {
27    /// Check if [DateCount] has expired.
28    pub fn expired(&self, ttl: chrono::Duration) -> bool {
29        self.expired_at(ttl, Utc::now())
30    }
31
32    /// Check if [DateCount] has expired at instant.
33    pub fn expired_at(&self, ttl: chrono::Duration, instant: DateTime<Utc>) -> bool {
34        self.create_date + ttl < instant
35    }
36}
37
38#[derive(Debug, Clone)]
39pub struct DateCountUntil {
40    pub date_count: DateCount,
41    pub until: DateTime<Utc>,
42}
43
44impl Value for DateCountUntil {
45    type Count = u32;
46
47    fn count(&self) -> Self::Count {
48        self.date_count.count
49    }
50
51    fn create_date(&self) -> Option<DateTime<Utc>> {
52        Some(self.date_count.create_date)
53    }
54
55    fn expire_date(&self) -> Option<DateTime<Utc>> {
56        Some(self.until)
57    }
58}
59
60/// [MemStore] stores data in memory.
61#[derive(Debug, Clone)]
62pub struct MemStore {
63    pub(crate) inner: Arc<Mutex<MemStoreInner>>,
64}
65
66impl MemStore {
67    pub fn new(capacity: usize, ttl: chrono::Duration) -> Self {
68        Self {
69            inner: Arc::new(Mutex::new(MemStoreInner::new(capacity, ttl))),
70        }
71    }
72}
73
74impl Default for MemStore {
75    fn default() -> Self {
76        Self::new(DEFAULT_STORE_CAPACITY, chrono::Duration::seconds(60))
77    }
78}
79
80#[async_trait::async_trait]
81impl Store for MemStore {
82    type Error = ();
83    type Key = String;
84    type Value = DateCountUntil;
85    type Count = u32;
86
87    async fn incr_by(&self, key: Self::Key, val: u32) -> Result<Self::Value, Self::Error> {
88        Ok(self.inner.lock().await.incr_by(key, val))
89    }
90
91    async fn incr(&self, key: Self::Key) -> Result<Self::Value, Self::Error> {
92        self.incr_by(key, 1).await
93    }
94
95    async fn del(&self, key: Self::Key) -> Result<Option<Self::Value>, Self::Error> {
96        Ok(self.inner.lock().await.del(key))
97    }
98
99    async fn clear(&self) -> Result<(), Self::Error> {
100        self.inner.lock().await.clear();
101        Ok(())
102    }
103}
104
105#[derive(Debug, Clone)]
106pub(crate) struct MemStoreInner {
107    pub(crate) data: HashMap<String, DateCount>,
108    /// The [ttl] field indicates the time-to-live (TTL)
109    /// of data from its creation. Once this TTL expires,
110    /// the data in the cache is considered empty or expired.
111    pub(crate) ttl: chrono::Duration,
112}
113
114impl MemStoreInner {
115    pub fn new(capacity: usize, ttl: chrono::Duration) -> Self {
116        Self {
117            data: HashMap::with_capacity(capacity),
118            ttl,
119        }
120    }
121
122    pub fn incr_by(&mut self, key: String, val: u32) -> DateCountUntil {
123        let entry = self.data.entry(key).or_default();
124
125        if entry.expired(self.ttl) {
126            *entry = DateCount::default()
127        }
128
129        entry.count += val;
130
131        DateCountUntil {
132            date_count: *entry,
133            until: entry.create_date + self.ttl,
134        }
135    }
136
137    pub fn del(&mut self, key: String) -> Option<DateCountUntil> {
138        self.data.remove(&key)
139            .map(|entry| DateCountUntil {
140                date_count: entry,
141                until: entry.create_date + self.ttl,
142            })
143    }
144
145    pub fn clear(&mut self) {
146        self.data.clear()
147    }
148}
149
150impl Deref for MemStoreInner {
151    type Target = HashMap<String, DateCount>;
152
153    fn deref(&self) -> &Self::Target {
154        &self.data
155    }
156}
157
158impl DerefMut for MemStoreInner {
159    fn deref_mut(&mut self) -> &mut Self::Target {
160        &mut self.data
161    }
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167
168    #[tokio::test]
169    async fn incr_del() -> Result<(), ()> {
170        let store = MemStore::new(8, chrono::Duration::seconds(100000));
171
172        assert_eq!(store.incr("John".to_string()).await?.date_count.count, 1);
173        assert_eq!(store.incr("John".to_string()).await?.date_count.count, 2);
174        assert_eq!(store.incr("John".to_string()).await?.date_count.count, 3);
175        assert_eq!(store.incr_by("John".to_string(), 2).await?.date_count.count, 5);
176
177        assert_eq!(store.incr_by("Meg".to_string(), 3).await?.date_count.count, 3);
178        assert_eq!(store.incr_by("Meg".to_string(), 3).await?.date_count.count, 6);
179        assert_eq!(store.incr_by("Meg".to_string(), 3).await?.date_count.count, 9);
180        assert_eq!(store.incr_by("Meg".to_string(), 4).await?.date_count.count, 13);
181        assert_eq!(store.incr_by("Meg".to_string(), 4).await?.date_count.count, 17);
182        assert_eq!(store.incr("Meg".to_string()).await?.date_count.count, 18);
183
184        let cloned = store.clone();
185        assert_eq!(cloned.incr("John".to_string()).await?.date_count.count, 6);
186        assert_eq!(store.incr("John".to_string()).await?.date_count.count, 7);
187
188        store.del("Meg".to_string()).await?;
189        assert_eq!(store.incr_by("Meg".to_string(), 3).await?.date_count.count, 3);
190        assert_eq!(cloned.incr_by("Meg".to_string(), 3).await?.date_count.count, 6);
191        assert_eq!(store.incr_by("Meg".to_string(), 3).await?.date_count.count, 9);
192        assert_eq!(store.incr("John".to_string()).await?.date_count.count, 8);
193
194        Ok(())
195    }
196
197    #[tokio::test]
198    async fn clear() -> Result<(), ()> {
199        let store = MemStore::new(8, chrono::Duration::seconds(100000));
200
201        assert_eq!(store.incr("John".to_string()).await?.date_count.count, 1);
202        assert_eq!(store.incr_by("Meg".to_string(), 3).await?.date_count.count, 3);
203
204        store.clear().await?;
205        assert_eq!(store.incr("John".to_string()).await?.date_count.count, 1);
206        assert_eq!(store.incr_by("Meg".to_string(), 3).await?.date_count.count, 3);
207
208        Ok(())
209    }
210
211    #[tokio::test]
212    async fn ttl() -> Result<(), ()> {
213        let store = MemStore::new(8, chrono::Duration::seconds(5));
214        assert_eq!(store.incr("John".to_string()).await?.date_count.count, 1);
215
216        // wait 2 seconds to add a new one...
217        tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
218        assert_eq!(store.incr_by("Meg".to_string(), 3).await?.date_count.count, 3);
219
220        // wait 4 seconds, "John" should be expired, while "Meg" should still exist.
221        tokio::time::sleep(tokio::time::Duration::from_secs(4)).await;
222        assert_eq!(store.incr("John".to_string()).await?.date_count.count, 1);
223        assert_eq!(store.incr_by("Meg".to_string(), 3).await?.date_count.count, 6);
224
225        Ok(())
226    }
227}
actix_rl/store/mem_store.rs

actix_rl/store/
mem_store.rs
