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    ///
37    /// By default this counts the **number of entries**, not their total size in
38    /// bytes. Since cached entries hold full response bodies, a cache bounded
39    /// only by entry count can still grow unbounded in memory when individual
40    /// responses are large. To bound by bytes instead, set a [`weigher`] that
41    /// returns each entry's size; `max_capacity` is then interpreted as the
42    /// maximum total weight.
43    ///
44    /// [`weigher`]: Self::weigher
45    #[must_use] pub fn max_capacity(mut self, capacity: u64) -> Self {
46        self.inner = self.inner.max_capacity(capacity);
47        self
48    }
49
50    /// Sets a weigher that returns the size of each entry, so that
51    /// [`max_capacity`](Self::max_capacity) bounds the cache by total weight
52    /// (e.g. bytes) rather than entry count.
53    ///
54    /// A common choice is the cached body's byte length, which keeps memory use
55    /// bounded even when a few responses are very large.
56    #[must_use]
57    pub fn weigher(
58        mut self,
59        weigher: impl Fn(&K, &CachedEntry) -> u32 + Send + Sync + 'static,
60    ) -> Self {
61        self.inner = self.inner.weigher(weigher);
62        self
63    }
64
65    /// Sets the time to idle of the cache.
66    ///
67    /// A cached entry will expire after the specified duration has passed since `get`
68    /// or `insert`.
69    ///
70    /// # Panics
71    ///
72    /// `CacheBuilder::build*` methods will panic if the given `duration` is longer
73    /// than 1000 years. This is done to protect against overflow when computing key
74    /// expiration.
75    #[must_use] pub fn time_to_idle(mut self, duration: Duration) -> Self {
76        self.inner = self.inner.time_to_idle(duration);
77        self
78    }
79
80    /// Sets the time to live of the cache.
81    ///
82    /// A cached entry will expire after the specified duration has passed since
83    /// `insert`.
84    ///
85    /// # Panics
86    ///
87    /// `CacheBuilder::build*` methods will panic if the given `duration` is longer
88    /// than 1000 years. This is done to protect against overflow when computing key
89    /// expiration.
90    #[must_use] pub fn time_to_live(mut self, duration: Duration) -> Self {
91        self.inner = self.inner.time_to_live(duration);
92        self
93    }
94
95    /// Sets the eviction listener closure to the cache.
96    ///
97    /// # Panics
98    ///
99    /// It is very important to ensure the listener closure does not panic. Otherwise,
100    /// the cache will stop calling the listener after a panic. This is intended
101    /// behavior because the cache cannot know whether it is memory safe to
102    /// call the panicked listener again.
103    #[must_use]
104    pub fn eviction_listener(
105        mut self,
106        listener: impl Fn(Arc<K>, CachedEntry, RemovalCause) + Send + Sync + 'static,
107    ) -> Self {
108        self.inner = self.inner.eviction_listener(listener);
109        self
110    }
111
112    /// Build a [`MokaStore`].
113    ///
114    /// # Panics
115    ///
116    /// Panics if configured with either `time_to_live` or `time_to_idle` higher than
117    /// 1000 years. This is done to protect against overflow when computing key
118    /// expiration.
119    #[must_use] pub fn build(self) -> MokaStore<K> {
120        MokaStore {
121            inner: self.inner.build(),
122        }
123    }
124}
125/// A simple in-memory store for the cache.
126pub struct MokaStore<K> {
127    inner: MokaCache<K, CachedEntry>,
128}
129
130impl<K> Debug for MokaStore<K> {
131    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
132        f.debug_struct("MokaStore").finish()
133    }
134}
135impl<K> MokaStore<K>
136where
137    K: Hash + Eq + Send + Sync + Clone + 'static,
138{
139    /// Create a new `MokaStore`.
140    #[must_use] pub fn new(max_capacity: u64) -> Self {
141        Self {
142            inner: MokaCache::new(max_capacity),
143        }
144    }
145
146    /// Returns a [`Builder`], which can build a `MokaStore`.
147    #[must_use] pub fn builder() -> Builder<K> {
148        Builder {
149            inner: MokaCache::builder(),
150        }
151    }
152}
153
154impl<K> CacheStore for MokaStore<K>
155where
156    K: Hash + Eq + Send + Sync + Clone + 'static,
157{
158    type Error = Infallible;
159    type Key = K;
160
161    async fn load_entry<Q>(&self, key: &Q) -> Option<CachedEntry>
162    where
163        Self::Key: Borrow<Q>,
164        Q: Hash + Eq + Sync,
165    {
166        self.inner.get(key).await
167    }
168
169    async fn save_entry(&self, key: Self::Key, entry: CachedEntry) -> Result<(), Self::Error> {
170        self.inner.insert(key, entry).await;
171        Ok(())
172    }
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178    use std::time::Duration;
179    use salvo_core::http::{HeaderMap, StatusCode};
180    use crate::{CachedBody, CachedEntry};
181
182    #[tokio::test]
183    async fn test_moka_store() {
184        let store = MokaStore::new(100);
185        let key = "test_key".to_owned();
186        let entry = CachedEntry {
187            status: Some(StatusCode::OK),
188            headers: HeaderMap::new(),
189            body: CachedBody::Once("test_body".into()),
190        };
191        store.save_entry(key.clone(), entry.clone()).await.unwrap();
192        let loaded_entry = store.load_entry(&key).await.unwrap();
193        assert_eq!(loaded_entry.status, entry.status);
194        assert_eq!(loaded_entry.body, entry.body);
195    }
196
197    #[tokio::test]
198    async fn test_moka_store_builder() {
199        let store = MokaStore::<String>::builder()
200            .initial_capacity(50)
201            .max_capacity(100)
202            .time_to_live(Duration::from_secs(1))
203            .time_to_idle(Duration::from_secs(1))
204            .build();
205        let key = "test_key".to_owned();
206        let entry = CachedEntry {
207            status: Some(StatusCode::OK),
208            headers: HeaderMap::new(),
209            body: CachedBody::Once("test_body".into()),
210        };
211        store.save_entry(key.clone(), entry.clone()).await.unwrap();
212        let loaded_entry = store.load_entry(&key).await.unwrap();
213        assert_eq!(loaded_entry.status, entry.status);
214        assert_eq!(loaded_entry.body, entry.body);
215
216        tokio::time::sleep(Duration::from_secs(2)).await;
217        let loaded_entry = store.load_entry(&key).await;
218        assert!(loaded_entry.is_none());
219    }
220    
221    #[test]
222    fn test_builder_debug() {
223        let builder = MokaStore::<String>::builder();
224        let dbg_str = format!("{builder:?}");
225        assert_eq!(dbg_str, "Builder");
226    }
227
228    #[test]
229    fn test_moka_store_debug() {
230        let store = MokaStore::<String>::new(100);
231        let dbg_str = format!("{store:?}");
232        assert_eq!(dbg_str, "MokaStore");
233    }
234    
235    #[tokio::test]
236    async fn test_eviction_listener() {
237        use std::sync::atomic::{AtomicBool, Ordering};
238        let evicted = Arc::new(AtomicBool::new(false));
239        let evicted_clone = evicted.clone();
240        let store = MokaStore::<String>::builder()
241            .max_capacity(1)
242            .eviction_listener(move |_, _, _| {
243                evicted_clone.store(true, Ordering::SeqCst);
244            })
245            .build();
246        let entry = CachedEntry {
247            status: None,
248            headers: HeaderMap::new(),
249            body: CachedBody::Once("test_body".into()),
250        };
251        store.save_entry("key1".to_owned(), entry.clone()).await.unwrap();
252        store.save_entry("key2".to_owned(), entry.clone()).await.unwrap();
253        
254        // Try to get the key to give time to the eviction listener to run.
255        for _ in 0..10 {
256            store.load_entry(&"key1".to_owned()).await;
257            store.load_entry(&"key2".to_owned()).await;
258            if evicted.load(Ordering::SeqCst) {
259                break;
260            }
261            tokio::time::sleep(Duration::from_millis(50)).await;
262        }
263
264        assert!(evicted.load(Ordering::SeqCst));
265    }
266
267    #[tokio::test]
268    async fn test_weigher_bounds_by_bytes() {
269        use std::sync::atomic::{AtomicBool, Ordering};
270
271        fn body_len(entry: &CachedEntry) -> u32 {
272            match &entry.body {
273                CachedBody::None => 0,
274                CachedBody::Once(bytes) => bytes.len() as u32,
275                CachedBody::Chunks(chunks) => {
276                    chunks.iter().map(|c| c.len()).sum::<usize>() as u32
277                }
278            }
279        }
280
281        let evicted = Arc::new(AtomicBool::new(false));
282        let evicted_clone = evicted.clone();
283        // Weigh by body bytes; cap total weight so two 9-byte bodies can't coexist.
284        let store = MokaStore::<String>::builder()
285            .weigher(|_k, entry| body_len(entry))
286            .max_capacity(12)
287            .eviction_listener(move |_, _, _| {
288                evicted_clone.store(true, Ordering::SeqCst);
289            })
290            .build();
291        let entry = CachedEntry {
292            status: None,
293            headers: HeaderMap::new(),
294            body: CachedBody::Once("test_body".into()), // 9 bytes
295        };
296        store.save_entry("key1".to_owned(), entry.clone()).await.unwrap();
297        store.save_entry("key2".to_owned(), entry.clone()).await.unwrap();
298
299        for _ in 0..10 {
300            store.load_entry(&"key1".to_owned()).await;
301            store.load_entry(&"key2".to_owned()).await;
302            if evicted.load(Ordering::SeqCst) {
303                break;
304            }
305            tokio::time::sleep(Duration::from_millis(50)).await;
306        }
307
308        assert!(evicted.load(Ordering::SeqCst));
309    }
310}
salvo_cache/moka_store.rs

salvo_cache/
moka_store.rs
