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}