salvo_cache/
moka_store.rs

1//! Memory store module.
2use std::borrow::Borrow;
3use std::convert::Infallible;
4use std::hash::Hash;
5use std::sync::Arc;
6use std::time::Duration;
7
8use moka::future::Cache as MokaCache;
9use moka::future::CacheBuilder as MokaCacheBuilder;
10use moka::notification::RemovalCause;
11
12use super::{CacheStore, CachedEntry};
13
14/// A builder for [`MokaStore`].
15pub struct Builder<K> {
16    inner: MokaCacheBuilder<K, CachedEntry, MokaCache<K, CachedEntry>>,
17}
18impl<K> Builder<K>
19where
20    K: Hash + Eq + Send + Sync + Clone + 'static,
21{
22    /// Sets the initial capacity (number of entries) of the cache.
23    pub fn initial_capacity(mut self, capacity: usize) -> Self {
24        self.inner = self.inner.initial_capacity(capacity);
25        self
26    }
27
28    /// Sets the max capacity of the cache.
29    pub fn max_capacity(mut self, capacity: u64) -> Self {
30        self.inner = self.inner.max_capacity(capacity);
31        self
32    }
33
34    /// Sets the time to idle of the cache.
35    ///
36    /// A cached entry will expire after the specified duration has passed since `get`
37    /// or `insert`.
38    ///
39    /// # Panics
40    ///
41    /// `CacheBuilder::build*` methods will panic if the given `duration` is longer
42    /// than 1000 years. This is done to protect against overflow when computing key
43    /// expiration.
44    pub fn time_to_idle(mut self, duration: Duration) -> Self {
45        self.inner = self.inner.time_to_idle(duration);
46        self
47    }
48
49    /// Sets the time to live of the cache.
50    ///
51    /// A cached entry will expire after the specified duration has passed since
52    /// `insert`.
53    ///
54    /// # Panics
55    ///
56    /// `CacheBuilder::build*` methods will panic if the given `duration` is longer
57    /// than 1000 years. This is done to protect against overflow when computing key
58    /// expiration.
59    pub fn time_to_live(mut self, duration: Duration) -> Self {
60        self.inner = self.inner.time_to_live(duration);
61        self
62    }
63
64    /// Sets the eviction listener closure to the cache.
65    ///
66    /// # Panics
67    ///
68    /// It is very important to ensure the listener closure does not panic. Otherwise,
69    /// the cache will stop calling the listener after a panic. This is intended
70    /// behavior because the cache cannot know whether it is memory safe to
71    /// call the panicked listener again.
72    pub fn eviction_listener(
73        mut self,
74        listener: impl Fn(Arc<K>, CachedEntry, RemovalCause) + Send + Sync + 'static,
75    ) -> Self {
76        self.inner = self.inner.eviction_listener(listener);
77        self
78    }
79
80    /// Build a [`MokaStore`].
81    ///
82    /// # Panics
83    ///
84    /// Panics if configured with either `time_to_live` or `time_to_idle` higher than
85    /// 1000 years. This is done to protect against overflow when computing key
86    /// expiration.
87    pub fn build(self) -> MokaStore<K> {
88        MokaStore {
89            inner: self.inner.build(),
90        }
91    }
92}
93/// A simple in-memory store for rate limiter.
94pub struct MokaStore<K> {
95    inner: MokaCache<K, CachedEntry>,
96}
97impl<K> MokaStore<K>
98where
99    K: Hash + Eq + Send + Sync + Clone + 'static,
100{
101    /// Create a new `MokaStore`.
102    pub fn new(max_capacity: u64) -> Self {
103        Self {
104            inner: MokaCache::new(max_capacity),
105        }
106    }
107
108    /// Returns a [`Builder`], which can build a `MokaStore`.
109    pub fn builder() -> Builder<K> {
110        Builder {
111            inner: MokaCache::builder(),
112        }
113    }
114}
115
116impl<K> CacheStore for MokaStore<K>
117where
118    K: Hash + Eq + Send + Sync + Clone + 'static,
119{
120    type Error = Infallible;
121    type Key = K;
122
123    async fn load_entry<Q>(&self, key: &Q) -> Option<CachedEntry>
124    where
125        Self::Key: Borrow<Q>,
126        Q: Hash + Eq + Sync,
127    {
128        self.inner.get(key).await
129    }
130
131    async fn save_entry(&self, key: Self::Key, entry: CachedEntry) -> Result<(), Self::Error> {
132        self.inner.insert(key, entry).await;
133        Ok(())
134    }
135}