ferron/util/
ttl_cache.rs

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 for 2 seconds to ensure the entry expires
67    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 for 2 seconds to ensure the entries expire
88    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}