muffy/cache/
memory.rs

1use super::{Cache, CacheError};
2use async_trait::async_trait;
3use scc::{HashMap, hash_map::Entry};
4
5/// An in-memory cache.
6pub struct MemoryCache<T> {
7    map: HashMap<String, T>,
8}
9
10impl<T> MemoryCache<T> {
11    /// Creates an in-memory cache.
12    pub fn new(capacity: usize) -> Self {
13        Self {
14            map: HashMap::with_capacity(capacity),
15        }
16    }
17}
18
19#[async_trait]
20impl<T: Clone + Send + Sync> Cache<T> for MemoryCache<T> {
21    async fn get_with(
22        &self,
23        key: String,
24        future: Box<dyn Future<Output = T> + Send>,
25    ) -> Result<T, CacheError> {
26        Ok(match self.map.entry_async(key).await {
27            Entry::Occupied(entry) => entry.get().clone(),
28            Entry::Vacant(entry) => {
29                // TODO Avoid deadlocks.
30                let value = Box::into_pin(future).await;
31                entry.insert_entry(value.clone());
32                value
33            }
34        })
35    }
36
37    async fn remove(&self, key: &str) -> Result<(), CacheError> {
38        self.map.remove(key);
39
40        Ok(())
41    }
42}
43
44#[cfg(test)]
45mod tests {
46    use super::*;
47
48    #[tokio::test]
49    async fn get_or_set() {
50        let cache = MemoryCache::new(1 << 10);
51
52        assert_eq!(
53            cache
54                .get_with("key".into(), Box::new(async { 42 }))
55                .await
56                .unwrap(),
57            42,
58        );
59        assert_eq!(
60            cache
61                .get_with("key".into(), Box::new(async { 0 }))
62                .await
63                .unwrap(),
64            42,
65        );
66    }
67}