1use std::collections::HashMap;
15use std::hash::Hash;
16use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
17use std::time::Duration;
18
19use tokio::time::Instant;
20use tracing::{debug, warn};
21
22#[derive(Debug, Clone)]
24struct CacheEntry<V> {
25 value: V,
26 inserted_at: Instant,
27 ttl: Duration,
28}
29
30impl<V> CacheEntry<V> {
31 fn new(value: V, ttl: Duration) -> Self {
33 Self {
34 value,
35 inserted_at: Instant::now(),
36 ttl,
37 }
38 }
39
40 fn is_expired(&self) -> bool {
42 self.inserted_at.elapsed() > self.ttl
43 }
44
45 fn age(&self) -> Duration {
47 self.inserted_at.elapsed()
48 }
49}
50
51pub struct TtlCache<K, V> {
76 entries: RwLock<HashMap<K, CacheEntry<V>>>,
77 default_ttl: Duration,
78 max_capacity: usize,
81}
82
83const DEFAULT_MAX_CAPACITY: usize = 1024;
85
86impl<K, V> TtlCache<K, V>
87where
88 K: Eq + Hash + Clone + std::fmt::Debug,
89 V: Clone,
90{
91 pub fn new(default_ttl: Duration) -> Self {
93 Self {
94 entries: RwLock::new(HashMap::new()),
95 default_ttl,
96 max_capacity: DEFAULT_MAX_CAPACITY,
97 }
98 }
99
100 pub fn with_max_capacity(default_ttl: Duration, max_capacity: usize) -> Self {
102 Self {
103 entries: RwLock::new(HashMap::new()),
104 default_ttl,
105 max_capacity,
106 }
107 }
108
109 fn read(&self) -> RwLockReadGuard<'_, HashMap<K, CacheEntry<V>>> {
111 self.entries.read().unwrap_or_else(|poisoned| {
112 warn!("Cache read lock poisoned, recovering");
113 poisoned.into_inner()
114 })
115 }
116
117 fn write(&self) -> RwLockWriteGuard<'_, HashMap<K, CacheEntry<V>>> {
119 self.entries.write().unwrap_or_else(|poisoned| {
120 warn!("Cache write lock poisoned, recovering");
121 poisoned.into_inner()
122 })
123 }
124
125 pub fn get(&self, key: &K) -> Option<V> {
129 let entries = self.read();
130 let entry = entries.get(key)?;
131
132 if entry.is_expired() {
133 debug!(
134 hit = false,
135 ?key,
136 age_secs = entry.age().as_secs(),
137 "cache lookup (expired)"
138 );
139 None
140 } else {
141 debug!(hit = true, ?key, "cache lookup");
142 Some(entry.value.clone())
143 }
144 }
145
146 pub fn insert(&self, key: K, value: V) {
148 self.insert_with_ttl(key, value, self.default_ttl);
149 }
150
151 pub fn insert_with_ttl(&self, key: K, value: V, ttl: Duration) {
156 let mut entries = self.write();
157
158 if entries.len() >= self.max_capacity && !entries.contains_key(&key) {
160 let before = entries.len();
162 entries.retain(|_, entry| !entry.is_expired());
163 let removed = before - entries.len();
164 if removed > 0 {
165 debug!(removed, "Evicted expired entries to make room");
166 }
167
168 if entries.len() >= self.max_capacity {
177 let low_water = self.max_capacity.saturating_mul(9) / 10;
178 let to_remove = entries.len().saturating_sub(low_water).max(1);
179 let mut aged: Vec<(K, Duration)> =
180 entries.iter().map(|(k, e)| (k.clone(), e.age())).collect();
181 let take = to_remove.min(aged.len());
182 if take < aged.len() {
183 aged.select_nth_unstable_by(take - 1, |a, b| b.1.cmp(&a.1));
186 }
187 for (k, _) in aged.into_iter().take(take) {
188 entries.remove(&k);
189 }
190 debug!(evicted = take, "Batch-evicted oldest entries to make room");
191 }
192 }
193
194 debug!(?key, ttl_secs = ttl.as_secs(), "Inserting cache entry");
195 entries.insert(key, CacheEntry::new(value, ttl));
196 }
197
198 pub fn remove(&self, key: &K) -> Option<V> {
200 self.write().remove(key).map(|e| e.value)
201 }
202
203 pub fn len(&self) -> usize {
205 self.read().len()
206 }
207
208 pub fn is_empty(&self) -> bool {
210 self.len() == 0
211 }
212
213 pub fn clear(&self) {
215 self.write().clear();
216 }
217}
218
219#[cfg(test)]
220mod tests {
221 use super::*;
222
223 #[test]
224 fn test_cache_insert_and_get() {
225 let cache: TtlCache<String, String> = TtlCache::new(Duration::from_secs(3600));
226
227 cache.insert("key".to_string(), "value".to_string());
228
229 assert_eq!(cache.get(&"key".to_string()), Some("value".to_string()));
230 }
231
232 #[test]
233 fn test_cache_get_missing_key() {
234 let cache: TtlCache<String, String> = TtlCache::new(Duration::from_secs(3600));
235
236 assert_eq!(cache.get(&"missing".to_string()), None);
237 }
238
239 #[test]
240 fn test_cache_expiration() {
241 let cache: TtlCache<String, String> = TtlCache::new(Duration::from_millis(10));
242
243 cache.insert("key".to_string(), "value".to_string());
244 assert_eq!(cache.get(&"key".to_string()), Some("value".to_string()));
245
246 std::thread::sleep(Duration::from_millis(20));
248
249 assert_eq!(cache.get(&"key".to_string()), None);
250 }
251
252 #[test]
253 fn test_cache_remove() {
254 let cache: TtlCache<String, String> = TtlCache::new(Duration::from_secs(3600));
255
256 cache.insert("key".to_string(), "value".to_string());
257 assert!(cache.get(&"key".to_string()).is_some());
258
259 cache.remove(&"key".to_string());
260 assert!(cache.get(&"key".to_string()).is_none());
261 }
262
263 #[tokio::test(start_paused = true)]
264 async fn capacity_eviction_batches_to_low_water_mark() {
265 let cap = 100;
275 let cache: TtlCache<u32, u32> = TtlCache::with_max_capacity(Duration::from_secs(3600), cap);
276 for i in 0..=cap as u32 {
277 cache.insert(i, i);
279 tokio::time::advance(Duration::from_millis(1)).await;
280 }
281 assert!(
282 cache.len() <= (cap * 9 / 10) + 1,
283 "batch eviction should drop to ~90% low-water, got len {}",
284 cache.len()
285 );
286 assert!(
287 cache.len() < cap,
288 "must be below capacity after a batch evict"
289 );
290 assert_eq!(cache.get(&(cap as u32)), Some(cap as u32));
292 assert_eq!(cache.get(&0), None);
293 }
294
295 #[test]
296 fn capacity_eviction_never_exceeds_capacity_under_churn() {
297 let cap = 50;
298 let cache: TtlCache<u32, u32> = TtlCache::with_max_capacity(Duration::from_secs(3600), cap);
299 for i in 0..1000u32 {
300 cache.insert(i, i);
301 assert!(cache.len() <= cap, "len {} exceeded cap {cap}", cache.len());
302 }
303 }
304
305 #[test]
306 fn test_cache_clear() {
307 let cache: TtlCache<String, String> = TtlCache::new(Duration::from_secs(3600));
308
309 cache.insert("key1".to_string(), "value1".to_string());
310 cache.insert("key2".to_string(), "value2".to_string());
311
312 assert_eq!(cache.len(), 2);
313
314 cache.clear();
315
316 assert_eq!(cache.len(), 0);
317 assert!(cache.is_empty());
318 }
319
320 #[tokio::test(start_paused = true)]
321 async fn test_entry_expires_after_ttl() {
322 let cache: TtlCache<String, String> = TtlCache::new(Duration::from_secs(1));
325 cache.insert("key".to_string(), "value".to_string());
326
327 tokio::time::advance(Duration::from_millis(800)).await;
328 assert!(
329 cache.get(&"key".to_string()).is_some(),
330 "entry must not be expired at t=800ms (< 1000ms TTL)"
331 );
332
333 tokio::time::advance(Duration::from_millis(300)).await;
334 assert!(
335 cache.get(&"key".to_string()).is_none(),
336 "entry must be expired at t=1100ms (> 1000ms TTL)"
337 );
338 }
339}