webterm_core/
simple_cache.rs

1use crate::random::random_in_range;
2use std::collections::HashMap;
3use std::fmt::{Debug, Display};
4use std::hash::Hash;
5use std::sync::Arc;
6use tokio::sync::RwLock;
7use tokio::task::JoinHandle;
8use tokio::time;
9use tokio::time::{Duration, Instant};
10use tracing::debug;
11
12const MAX_CLEANUP_DURATION: Duration = Duration::from_millis(200);
13const CLEANUP_EVERY: Duration = Duration::from_secs(10);
14
15#[derive(Debug)]
16pub enum CacheError {
17    ReadError,
18    WriteError,
19    AtCapacity,
20    KeyNotFound,
21}
22
23impl std::error::Error for CacheError {}
24
25impl std::fmt::Display for CacheError {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        match self {
28            CacheError::ReadError => write!(f, "Read error"),
29            CacheError::WriteError => write!(f, "Write error"),
30            CacheError::AtCapacity => write!(f, "Cache at capacity"),
31            CacheError::KeyNotFound => write!(f, "Key not found"),
32        }
33    }
34}
35
36pub struct SimpleCache<K, V> {
37    map: Arc<RwLock<HashMap<K, (V, Instant)>>>,
38    max_size: usize, // since HashMap#len() returns usize
39    cleanup_handle: JoinHandle<()>,
40}
41
42impl<K, V> SimpleCache<K, V>
43where
44    K: Hash + Eq + Send + Sync + Display + Debug + Clone + 'static,
45    V: Send + Sync + Clone + 'static,
46{
47    pub fn new(max_size: usize) -> Self {
48        let map = Arc::new(RwLock::new(HashMap::new()));
49        let map_clone = map.clone();
50
51        let cleanup_handle = tokio::spawn(async move {
52            let mut interval = time::interval(CLEANUP_EVERY);
53            loop {
54                interval.tick().await;
55                let _ = Self::remove_expired(&map_clone).await;
56            }
57        });
58
59        SimpleCache {
60            map,
61            max_size,
62            cleanup_handle,
63        }
64    }
65
66    pub async fn len(&self) -> Result<usize, CacheError> {
67        Ok(self.map.read().await.len())
68    }
69
70    pub async fn is_empty(&self) -> Result<bool, CacheError> {
71        Ok(self.map.read().await.is_empty())
72    }
73
74    pub async fn insert(&self, key: K, value: V, duration: Duration) -> Result<(), CacheError> {
75        if self.len().await? >= self.max_size {
76            return Err(CacheError::AtCapacity);
77        }
78        let expire_at = Instant::now() + duration;
79        self.map.write().await.insert(key, (value, expire_at));
80        Ok(())
81    }
82
83    pub async fn get(&self, key: &K) -> Result<V, CacheError> {
84        debug!("starting simple_cache/get for key {:?}", key);
85        let map = self.map.read().await;
86        debug!("simple_cache/get map read lock acquired");
87        let result = map.get(key).map(|(value, expires_at)| {
88            debug!("simple_cache/get map.get() result: {:?}", expires_at);
89            if &Instant::now() <= expires_at {
90                Some(value)
91            } else {
92                None
93            }
94        });
95
96        debug!("simple_cache/get map.get() result loop finished");
97
98        if let Some(Some(value)) = result {
99            debug!("simple_cache/get returning value");
100            Ok(value.clone())
101        } else {
102            debug!("simple_cache/get returning key not found");
103            Err(CacheError::KeyNotFound)
104        }
105    }
106
107    pub async fn remove(&self, key: &K) -> Result<V, CacheError> {
108        if let Some((value, _expires_at)) = self.map.write().await.remove(key) {
109            Ok(value)
110        } else {
111            Err(CacheError::KeyNotFound)
112        }
113    }
114
115    pub async fn reset_expiration(&self, key: K, duration: Duration) -> Result<(), CacheError> {
116        let existing = self.get(&key).await?;
117        self.insert(key, existing, duration).await
118    }
119
120    pub async fn remove_expired(
121        map: &Arc<RwLock<HashMap<K, (V, Instant)>>>,
122    ) -> Result<(), CacheError> {
123        let start_time = Instant::now();
124        let mut keys_to_remove: Vec<K> = Vec::new();
125
126        {
127            for (key, (_value, expires_at)) in map.write().await.iter() {
128                if start_time.elapsed() > MAX_CLEANUP_DURATION {
129                    break;
130                }
131                if expires_at < &start_time {
132                    keys_to_remove.push(key.clone());
133                }
134            }
135        }
136
137        {
138            let mut write_guard = map.write().await;
139            for key in keys_to_remove.clone() {
140                write_guard.remove(&key);
141            }
142        }
143
144        // shrink to fit at randomly every 100th iteration
145        if random_in_range(0, 100) == 0 {
146            map.write().await.shrink_to_fit();
147        }
148
149        Ok(())
150    }
151}
152
153impl<K, V> Drop for SimpleCache<K, V> {
154    fn drop(&mut self) {
155        self.cleanup_handle.abort()
156    }
157}
webterm_core/simple_cache.rs

webterm_core/
simple_cache.rs
