Skip to main content

otelite_api/
cache.rs

1//! Simple LRU cache for query results
2
3use std::collections::{HashMap, VecDeque};
4use std::hash::Hash;
5use std::sync::{Arc, RwLock};
6use std::time::{Duration, Instant};
7
8/// Entry in the cache with expiration
9#[derive(Clone)]
10struct CacheEntry<V> {
11    value: V,
12    expires_at: Instant,
13}
14
15/// Simple LRU cache with time-based expiration
16pub struct LruCache<K, V> {
17    cache: Arc<RwLock<HashMap<K, CacheEntry<V>>>>,
18    access_order: Arc<RwLock<VecDeque<K>>>,
19    max_size: usize,
20    ttl: Duration,
21}
22
23impl<K, V> LruCache<K, V>
24where
25    K: Eq + Hash + Clone,
26    V: Clone,
27{
28    /// Create a new LRU cache
29    pub fn new(max_size: usize, ttl: Duration) -> Self {
30        Self {
31            cache: Arc::new(RwLock::new(HashMap::new())),
32            access_order: Arc::new(RwLock::new(VecDeque::new())),
33            max_size,
34            ttl,
35        }
36    }
37
38    /// Get a value from the cache
39    pub fn get(&self, key: &K) -> Option<V> {
40        let cache = self.cache.read().ok()?;
41
42        if let Some(entry) = cache.get(key) {
43            // Check if entry has expired
44            if Instant::now() < entry.expires_at {
45                // Update access order
46                if let Ok(mut order) = self.access_order.write() {
47                    // Remove old position
48                    order.retain(|k| k != key);
49                    // Add to front (most recently used)
50                    order.push_front(key.clone());
51                }
52                return Some(entry.value.clone());
53            }
54        }
55
56        None
57    }
58
59    /// Insert a value into the cache
60    pub fn insert(&self, key: K, value: V) {
61        let expires_at = Instant::now() + self.ttl;
62        let entry = CacheEntry { value, expires_at };
63
64        // Acquire write locks
65        if let (Ok(mut cache), Ok(mut order)) = (self.cache.write(), self.access_order.write()) {
66            // If cache is at capacity, remove least recently used
67            if cache.len() >= self.max_size && !cache.contains_key(&key) {
68                if let Some(lru_key) = order.pop_back() {
69                    cache.remove(&lru_key);
70                }
71            }
72
73            // Insert new entry
74            cache.insert(key.clone(), entry);
75
76            // Update access order
77            order.retain(|k| k != &key);
78            order.push_front(key);
79        }
80    }
81
82    /// Clear all entries from the cache
83    pub fn clear(&self) {
84        if let (Ok(mut cache), Ok(mut order)) = (self.cache.write(), self.access_order.write()) {
85            cache.clear();
86            order.clear();
87        }
88    }
89
90    /// Remove expired entries
91    pub fn cleanup_expired(&self) {
92        let now = Instant::now();
93
94        if let (Ok(mut cache), Ok(mut order)) = (self.cache.write(), self.access_order.write()) {
95            // Collect expired keys
96            let expired_keys: Vec<K> = cache
97                .iter()
98                .filter(|(_, entry)| now >= entry.expires_at)
99                .map(|(k, _)| k.clone())
100                .collect();
101
102            // Remove expired entries
103            for key in &expired_keys {
104                cache.remove(key);
105                order.retain(|k| k != key);
106            }
107        }
108    }
109
110    /// Get current cache size
111    pub fn len(&self) -> usize {
112        self.cache.read().map(|c| c.len()).unwrap_or(0)
113    }
114
115    /// Check if cache is empty
116    pub fn is_empty(&self) -> bool {
117        self.len() == 0
118    }
119}
120
121impl<K, V> Clone for LruCache<K, V>
122where
123    K: Eq + Hash + Clone,
124    V: Clone,
125{
126    fn clone(&self) -> Self {
127        Self {
128            cache: Arc::clone(&self.cache),
129            access_order: Arc::clone(&self.access_order),
130            max_size: self.max_size,
131            ttl: self.ttl,
132        }
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139    use std::thread;
140
141    #[test]
142    fn test_basic_operations() {
143        let cache = LruCache::new(3, Duration::from_secs(60));
144
145        cache.insert("key1", "value1");
146        cache.insert("key2", "value2");
147
148        assert_eq!(cache.get(&"key1"), Some("value1"));
149        assert_eq!(cache.get(&"key2"), Some("value2"));
150        assert_eq!(cache.get(&"key3"), None);
151    }
152
153    #[test]
154    fn test_lru_eviction() {
155        let cache = LruCache::new(2, Duration::from_secs(60));
156
157        cache.insert("key1", "value1");
158        cache.insert("key2", "value2");
159        cache.insert("key3", "value3"); // Should evict key1
160
161        assert_eq!(cache.get(&"key1"), None);
162        assert_eq!(cache.get(&"key2"), Some("value2"));
163        assert_eq!(cache.get(&"key3"), Some("value3"));
164    }
165
166    #[test]
167    fn test_expiration() {
168        let cache = LruCache::new(10, Duration::from_millis(50));
169
170        cache.insert("key1", "value1");
171        assert_eq!(cache.get(&"key1"), Some("value1"));
172
173        thread::sleep(Duration::from_millis(100));
174        assert_eq!(cache.get(&"key1"), None);
175    }
176
177    #[test]
178    fn test_cleanup_expired() {
179        let cache = LruCache::new(10, Duration::from_millis(50));
180
181        cache.insert("key1", "value1");
182        cache.insert("key2", "value2");
183
184        thread::sleep(Duration::from_millis(100));
185        cache.cleanup_expired();
186
187        assert_eq!(cache.len(), 0);
188    }
189}