1use std::collections::HashMap;
2use std::time::{Duration, Instant};
3
4pub struct TtlCache<K, V> {
5 cache: HashMap<K, (V, Instant)>,
6 ttl: Duration,
7}
8
9impl<K, V> TtlCache<K, V>
10where
11 K: std::cmp::Eq + std::hash::Hash + Clone,
12 V: Clone,
13{
14 pub fn new(ttl: Duration) -> Self {
15 Self {
16 cache: HashMap::new(),
17 ttl,
18 }
19 }
20
21 pub fn insert(&mut self, key: K, value: V) {
22 self.cache.insert(key, (value, Instant::now()));
23 }
24
25 pub fn get(&self, key: &K) -> Option<V> {
26 self.cache.get(key).and_then(|(value, timestamp)| {
27 if timestamp.elapsed() < self.ttl {
28 Some(value.clone())
29 } else {
30 None
31 }
32 })
33 }
34
35 #[allow(dead_code)]
36 pub fn remove(&mut self, key: &K) -> Option<V> {
37 self.cache.remove(key).map(|(value, _)| value)
38 }
39
40 pub fn cleanup(&mut self) {
41 self
42 .cache
43 .retain(|_, (_, timestamp)| timestamp.elapsed() < self.ttl);
44 }
45}
46
47#[cfg(test)]
48mod tests {
49 use super::*;
50 use std::thread::sleep;
51 use std::time::Duration;
52
53 #[test]
54 fn test_insert_and_get() {
55 let mut cache = TtlCache::new(Duration::new(5, 0));
56 cache.insert("key1", "value1");
57
58 assert_eq!(cache.get(&"key1"), Some("value1"));
59 }
60
61 #[test]
62 fn test_get_expired() {
63 let mut cache = TtlCache::new(Duration::new(1, 0));
64 cache.insert("key1", "value1");
65
66 sleep(Duration::new(2, 0));
68
69 assert_eq!(cache.get(&"key1"), None);
70 }
71
72 #[test]
73 fn test_remove() {
74 let mut cache = TtlCache::new(Duration::new(5, 0));
75 cache.insert("key1", "value1");
76 cache.remove(&"key1");
77
78 assert_eq!(cache.get(&"key1"), None);
79 }
80
81 #[test]
82 fn test_cleanup() {
83 let mut cache = TtlCache::new(Duration::new(1, 0));
84 cache.insert("key1", "value1");
85 cache.insert("key2", "value2");
86
87 sleep(Duration::new(2, 0));
89
90 cache.cleanup();
91
92 assert_eq!(cache.get(&"key1"), None);
93 assert_eq!(cache.get(&"key2"), None);
94 }
95
96 #[test]
97 fn test_get_non_existent() {
98 let cache: TtlCache<&str, &str> = TtlCache::new(Duration::new(5, 0));
99 assert_eq!(cache.get(&"key1"), None);
100 }
101
102 #[test]
103 fn test_insert_and_get_multiple() {
104 let mut cache = TtlCache::new(Duration::new(5, 0));
105 cache.insert("key1", "value1");
106 cache.insert("key2", "value2");
107
108 assert_eq!(cache.get(&"key1"), Some("value1"));
109 assert_eq!(cache.get(&"key2"), Some("value2"));
110 }
111}