salvo_cache/
moka_store.rs

1//! Memory store module.
2use std::borrow::Borrow;
3use std::convert::Infallible;
4use std::fmt::{self, Debug, Formatter};
5use std::hash::Hash;
6use std::sync::Arc;
7use std::time::Duration;
8
9use moka::future::Cache as MokaCache;
10use moka::future::CacheBuilder as MokaCacheBuilder;
11use moka::notification::RemovalCause;
12
13use super::{CacheStore, CachedEntry};
14
15/// A builder for [`MokaStore`].
16pub struct Builder<K> {
17    inner: MokaCacheBuilder<K, CachedEntry, MokaCache<K, CachedEntry>>,
18}
19
20impl<K> Debug for Builder<K> {
21    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
22        f.debug_struct("Builder").finish()
23    }
24}
25impl<K> Builder<K>
26where
27    K: Hash + Eq + Send + Sync + Clone + 'static,
28{
29    /// Sets the initial capacity (number of entries) of the cache.
30    #[must_use] pub fn initial_capacity(mut self, capacity: usize) -> Self {
31        self.inner = self.inner.initial_capacity(capacity);
32        self
33    }
34
35    /// Sets the max capacity of the cache.
36    #[must_use] pub fn max_capacity(mut self, capacity: u64) -> Self {
37        self.inner = self.inner.max_capacity(capacity);
38        self
39    }
40
41    /// Sets the time to idle of the cache.
42    ///
43    /// A cached entry will expire after the specified duration has passed since `get`
44    /// or `insert`.
45    ///
46    /// # Panics
47    ///
48    /// `CacheBuilder::build*` methods will panic if the given `duration` is longer
49    /// than 1000 years. This is done to protect against overflow when computing key
50    /// expiration.
51    #[must_use] pub fn time_to_idle(mut self, duration: Duration) -> Self {
52        self.inner = self.inner.time_to_idle(duration);
53        self
54    }
55
56    /// Sets the time to live of the cache.
57    ///
58    /// A cached entry will expire after the specified duration has passed since
59    /// `insert`.
60    ///
61    /// # Panics
62    ///
63    /// `CacheBuilder::build*` methods will panic if the given `duration` is longer
64    /// than 1000 years. This is done to protect against overflow when computing key
65    /// expiration.
66    #[must_use] pub fn time_to_live(mut self, duration: Duration) -> Self {
67        self.inner = self.inner.time_to_live(duration);
68        self
69    }
70
71    /// Sets the eviction listener closure to the cache.
72    ///
73    /// # Panics
74    ///
75    /// It is very important to ensure the listener closure does not panic. Otherwise,
76    /// the cache will stop calling the listener after a panic. This is intended
77    /// behavior because the cache cannot know whether it is memory safe to
78    /// call the panicked listener again.
79    #[must_use]
80    pub fn eviction_listener(
81        mut self,
82        listener: impl Fn(Arc<K>, CachedEntry, RemovalCause) + Send + Sync + 'static,
83    ) -> Self {
84        self.inner = self.inner.eviction_listener(listener);
85        self
86    }
87
88    /// Build a [`MokaStore`].
89    ///
90    /// # Panics
91    ///
92    /// Panics if configured with either `time_to_live` or `time_to_idle` higher than
93    /// 1000 years. This is done to protect against overflow when computing key
94    /// expiration.
95    #[must_use] pub fn build(self) -> MokaStore<K> {
96        MokaStore {
97            inner: self.inner.build(),
98        }
99    }
100}
101/// A simple in-memory store for rate limiter.
102pub struct MokaStore<K> {
103    inner: MokaCache<K, CachedEntry>,
104}
105
106impl<K> Debug for MokaStore<K> {
107    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
108        f.debug_struct("MokaStore").finish()
109    }
110}
111impl<K> MokaStore<K>
112where
113    K: Hash + Eq + Send + Sync + Clone + 'static,
114{
115    /// Create a new `MokaStore`.
116    #[must_use] pub fn new(max_capacity: u64) -> Self {
117        Self {
118            inner: MokaCache::new(max_capacity),
119        }
120    }
121
122    /// Returns a [`Builder`], which can build a `MokaStore`.
123    #[must_use] pub fn builder() -> Builder<K> {
124        Builder {
125            inner: MokaCache::builder(),
126        }
127    }
128}
129
130impl<K> CacheStore for MokaStore<K>
131where
132    K: Hash + Eq + Send + Sync + Clone + 'static,
133{
134    type Error = Infallible;
135    type Key = K;
136
137    async fn load_entry<Q>(&self, key: &Q) -> Option<CachedEntry>
138    where
139        Self::Key: Borrow<Q>,
140        Q: Hash + Eq + Sync,
141    {
142        self.inner.get(key).await
143    }
144
145    async fn save_entry(&self, key: Self::Key, entry: CachedEntry) -> Result<(), Self::Error> {
146        self.inner.insert(key, entry).await;
147        Ok(())
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154    use std::time::Duration;
155    use salvo_core::http::{HeaderMap, StatusCode};
156    use crate::{CachedBody, CachedEntry};
157
158    #[tokio::test]
159    async fn test_moka_store() {
160        let store = MokaStore::new(100);
161        let key = "test_key".to_string();
162        let entry = CachedEntry {
163            status: Some(StatusCode::OK),
164            headers: HeaderMap::new(),
165            body: CachedBody::Once("test_body".into()),
166        };
167        store.save_entry(key.clone(), entry.clone()).await.unwrap();
168        let loaded_entry = store.load_entry(&key).await.unwrap();
169        assert_eq!(loaded_entry.status, entry.status);
170        assert_eq!(loaded_entry.body, entry.body);
171    }
172
173    #[tokio::test]
174    async fn test_moka_store_builder() {
175        let store = MokaStore::<String>::builder()
176            .initial_capacity(50)
177            .max_capacity(100)
178            .time_to_live(Duration::from_secs(1))
179            .time_to_idle(Duration::from_secs(1))
180            .build();
181        let key = "test_key".to_string();
182        let entry = CachedEntry {
183            status: Some(StatusCode::OK),
184            headers: HeaderMap::new(),
185            body: CachedBody::Once("test_body".into()),
186        };
187        store.save_entry(key.clone(), entry.clone()).await.unwrap();
188        let loaded_entry = store.load_entry(&key).await.unwrap();
189        assert_eq!(loaded_entry.status, entry.status);
190        assert_eq!(loaded_entry.body, entry.body);
191
192        tokio::time::sleep(Duration::from_secs(2)).await;
193        let loaded_entry = store.load_entry(&key).await;
194        assert!(loaded_entry.is_none());
195    }
196    
197    #[test]
198    fn test_builder_debug() {
199        let builder = MokaStore::<String>::builder();
200        let dbg_str = format!("{:?}", builder);
201        assert_eq!(dbg_str, "Builder");
202    }
203
204    #[test]
205    fn test_moka_store_debug() {
206        let store = MokaStore::<String>::new(100);
207        let dbg_str = format!("{:?}", store);
208        assert_eq!(dbg_str, "MokaStore");
209    }
210    
211    #[tokio::test]
212    async fn test_eviction_listener() {
213        use std::sync::atomic::{AtomicBool, Ordering};
214        let evicted = Arc::new(AtomicBool::new(false));
215        let evicted_clone = evicted.clone();
216        let store = MokaStore::<String>::builder()
217            .max_capacity(1)
218            .eviction_listener(move |_, _, _| {
219                evicted_clone.store(true, Ordering::SeqCst);
220            })
221            .build();
222        let entry = CachedEntry {
223            status: None,
224            headers: HeaderMap::new(),
225            body: CachedBody::Once("test_body".into()),
226        };
227        store.save_entry("key1".to_string(), entry.clone()).await.unwrap();
228        store.save_entry("key2".to_string(), entry.clone()).await.unwrap();
229        
230        // Try to get the key to give time to the eviction listener to run.
231        for _ in 0..10 {
232            store.load_entry(&"key1".to_string()).await;
233            store.load_entry(&"key2".to_string()).await;
234            if evicted.load(Ordering::SeqCst) {
235                break;
236            }
237            tokio::time::sleep(Duration::from_millis(50)).await;
238        }
239
240        assert!(evicted.load(Ordering::SeqCst));
241    }
242}
salvo_cache/moka_store.rs

salvo_cache/
moka_store.rs
