Skip to main content

gatel_core/
ttl_cache.rs

1use std::collections::HashMap;
2use std::hash::Hash;
3use std::sync::Mutex;
4use std::time::{Duration, Instant};
5
6/// A simple thread-safe TTL cache.
7///
8/// Entries are lazily evicted: expired entries are removed on the next
9/// `get` or `cleanup` call, not by a background timer.
10pub 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    /// Create a new TTL cache.
27    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    /// Insert a value with the default TTL.
38    pub fn insert(&self, key: K, value: V) {
39        self.insert_with_ttl(key, value, self.default_ttl);
40    }
41
42    /// Insert a value with a custom TTL.
43    pub fn insert_with_ttl(&self, key: K, value: V, ttl: Duration) {
44        let mut inner = self.inner.lock().unwrap();
45        // Evict expired entries if we're at capacity.
46        if inner.entries.len() >= self.max_entries {
47            let now = Instant::now();
48            inner.entries.retain(|_, e| e.expires_at > now);
49        }
50        // If still at capacity after eviction, drop the oldest entry.
51        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    /// Get a value if it exists and hasn't expired.
70    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    /// Remove a value.
82    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    /// Remove all expired entries.
88    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    /// Number of entries (including potentially expired ones).
95    pub fn len(&self) -> usize {
96        self.inner.lock().unwrap().entries.len()
97    }
98
99    /// Whether the cache is empty.
100    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); // should evict oldest
130        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}