1use super::{Cache, CacheError};
2use async_trait::async_trait;
3use scc::{HashMap, hash_map::Entry};
4
5pub struct MemoryCache<T> {
7 map: HashMap<String, T>,
8}
9
10impl<T> MemoryCache<T> {
11 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 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}