1use std::collections::HashMap;
2use std::hash::Hash;
3use std::sync::Mutex;
4use std::time::{Duration, Instant};
5
6pub struct TtlCache<K, V> {
11 inner: Mutex<CacheInner<K, V>>,
12 default_ttl: Duration,
13 max_entries: usize,
14}
15
16struct CacheInner<K, V> {
17 entries: HashMap<K, CacheEntry<V>>,
18}
19
20struct CacheEntry<V> {
21 value: V,
22 expires_at: Instant,
23}
24
25impl<K: Eq + Hash + Clone, V: Clone> TtlCache<K, V> {
26 pub fn new(default_ttl: Duration, max_entries: usize) -> Self {
28 Self {
29 inner: Mutex::new(CacheInner {
30 entries: HashMap::new(),
31 }),
32 default_ttl,
33 max_entries,
34 }
35 }
36
37 pub fn insert(&self, key: K, value: V) {
39 self.insert_with_ttl(key, value, self.default_ttl);
40 }
41
42 pub fn insert_with_ttl(&self, key: K, value: V, ttl: Duration) {
44 let mut inner = self.inner.lock().unwrap();
45 if inner.entries.len() >= self.max_entries {
47 let now = Instant::now();
48 inner.entries.retain(|_, e| e.expires_at > now);
49 }
50 if inner.entries.len() >= self.max_entries
52 && let Some(oldest_key) = inner
53 .entries
54 .iter()
55 .min_by_key(|(_, e)| e.expires_at)
56 .map(|(k, _)| k.clone())
57 {
58 inner.entries.remove(&oldest_key);
59 }
60 inner.entries.insert(
61 key,
62 CacheEntry {
63 value,
64 expires_at: Instant::now() + ttl,
65 },
66 );
67 }
68
69 pub fn get(&self, key: &K) -> Option<V> {
71 let mut inner = self.inner.lock().unwrap();
72 let entry = inner.entries.get(key)?;
73 if entry.expires_at <= Instant::now() {
74 inner.entries.remove(key);
75 None
76 } else {
77 Some(entry.value.clone())
78 }
79 }
80
81 pub fn remove(&self, key: &K) -> Option<V> {
83 let mut inner = self.inner.lock().unwrap();
84 inner.entries.remove(key).map(|e| e.value)
85 }
86
87 pub fn cleanup(&self) {
89 let mut inner = self.inner.lock().unwrap();
90 let now = Instant::now();
91 inner.entries.retain(|_, e| e.expires_at > now);
92 }
93
94 pub fn len(&self) -> usize {
96 self.inner.lock().unwrap().entries.len()
97 }
98
99 pub fn is_empty(&self) -> bool {
101 self.inner.lock().unwrap().entries.is_empty()
102 }
103}
104
105#[cfg(test)]
106mod tests {
107 use super::*;
108
109 #[test]
110 fn insert_and_get() {
111 let cache = TtlCache::new(Duration::from_secs(60), 100);
112 cache.insert("key1", "value1");
113 assert_eq!(cache.get(&"key1"), Some("value1"));
114 }
115
116 #[test]
117 fn expired_entry_returns_none() {
118 let cache = TtlCache::new(Duration::from_millis(1), 100);
119 cache.insert("key1", "value1");
120 std::thread::sleep(Duration::from_millis(10));
121 assert_eq!(cache.get(&"key1"), None);
122 }
123
124 #[test]
125 fn max_entries_eviction() {
126 let cache = TtlCache::new(Duration::from_secs(60), 2);
127 cache.insert("a", 1);
128 cache.insert("b", 2);
129 cache.insert("c", 3); assert_eq!(cache.len(), 2);
131 assert!(cache.get(&"c").is_some());
132 }
133
134 #[test]
135 fn remove() {
136 let cache = TtlCache::new(Duration::from_secs(60), 100);
137 cache.insert("key1", "value1");
138 assert_eq!(cache.remove(&"key1"), Some("value1"));
139 assert_eq!(cache.get(&"key1"), None);
140 }
141
142 #[test]
143 fn cleanup_removes_expired() {
144 let cache = TtlCache::new(Duration::from_millis(1), 100);
145 cache.insert("a", 1);
146 cache.insert("b", 2);
147 std::thread::sleep(Duration::from_millis(10));
148 cache.cleanup();
149 assert_eq!(cache.len(), 0);
150 }
151}