lru_cache_rs/
lib.rs

1use std::collections::HashMap;
2use std::hash::Hash;
3use std::sync::{Arc, Mutex};
4use std::time::{Duration, Instant};
5
6pub struct SafeLRUCache<K, V> {
7    inner: Arc<Mutex<InnerLRUCache<K, V>>>,
8}
9
10struct InnerLRUCache<K, V> {
11    capacity: usize,
12    max_size: Option<usize>,
13    current_size: usize,
14    map: HashMap<Arc<K>, CacheItem<V>>,
15    order: Vec<Arc<K>>,
16}
17
18#[derive(Clone)]
19struct CacheItem<V> {
20    value: V,
21    size: usize,
22    expires_at: Option<Instant>,
23}
24
25impl<K, V> SafeLRUCache<K, V>
26where
27    K: Eq + Hash + Clone + 'static,
28    V: Clone + 'static,
29{
30    pub fn new(capacity: usize, max_size: Option<usize>) -> Self {
31        Self {
32            inner: Arc::new(Mutex::new(InnerLRUCache {
33                capacity,
34                max_size,
35                current_size: 0,
36                map: HashMap::with_capacity(capacity),
37                order: Vec::with_capacity(capacity),
38            })),
39        }
40    }
41
42    pub fn put(&self, key: K, value: V, ttl: Option<Duration>, size: usize) {
43        let mut inner = self.inner.lock().unwrap();
44        let expires_at = ttl.map(|d| Instant::now() + d);
45
46        let item = CacheItem {
47            value,
48            size,
49            expires_at,
50        };
51
52        // Check size limit
53        if let Some(max) = inner.max_size {
54            while inner.current_size + item.size > max && !inner.map.is_empty() {
55                inner.remove_oldest();
56            }
57        }
58
59        // Check capacity limit
60        while inner.map.len() >= inner.capacity && !inner.map.is_empty() {
61            inner.remove_oldest();
62        }
63
64        let key_arc = Arc::new(key);
65        if inner.map.insert(Arc::clone(&key_arc), item).is_none() {
66            inner.current_size += size;
67        }
68        inner.order.push(key_arc);
69    }
70
71    pub fn get(&self, key: &K) -> Option<V> {
72        let mut inner = self.inner.lock().unwrap();
73        inner.clear_expired();
74
75        // First check if key exists to avoid cloning unnecessarily
76        if !inner.map.contains_key(key) {
77            return None;
78        }
79
80        // Now we know key exists, we can modify order safely
81        if let Some(pos) = inner.order.iter().position(|k| k.as_ref() == key) {
82            let key_clone = Arc::new(key.clone());
83            inner.order.remove(pos);
84            inner.order.push(key_clone);
85        }
86
87        inner.map.get(key).map(|item| item.value.clone())
88    }
89
90    pub fn clear_expired(&self) {
91        let mut inner = self.inner.lock().unwrap();
92        inner.clear_expired();
93    }
94}
95
96impl<K, V> InnerLRUCache<K, V>
97where
98    K: Eq + Hash,
99{
100    fn remove_oldest(&mut self) {
101        if let Some(oldest_key) = self.order.first() {
102            if let Some(item) = self.map.remove(oldest_key) {
103                self.current_size -= item.size;
104            }
105            self.order.remove(0);
106        }
107    }
108
109    fn clear_expired(&mut self) {
110        let now = Instant::now();
111        let mut i = 0;
112        while i < self.order.len() {
113            let key = &self.order[i];
114            if let Some(item) = self.map.get(key) {
115                if let Some(expiry) = item.expires_at {
116                    if expiry <= now {
117                        if let Some(removed_item) = self.map.remove(key) {
118                            self.current_size -= removed_item.size;
119                        }
120                        self.order.remove(i);
121                        continue;
122                    }
123                }
124            }
125            i += 1;
126        }
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133    use std::thread;
134
135    #[test]
136    fn test_basic_operations() {
137        let cache = SafeLRUCache::new(2, None);
138
139        cache.put("key1", "value1", None, 1);
140        cache.put("key2", "value2", None, 1);
141        assert_eq!(cache.get(&"key1"), Some("value1"));
142
143        // LRU eviction
144        cache.put("key3", "value3", None, 1);
145        assert_eq!(cache.get(&"key2"), None);
146    }
147
148    #[test]
149    fn test_ttl() {
150        let cache = SafeLRUCache::new(10, None);
151
152        cache.put("key1", "value1", Some(Duration::from_millis(100)), 1);
153        assert_eq!(cache.get(&"key1"), Some("value1"));
154
155        thread::sleep(Duration::from_millis(150));
156        assert_eq!(cache.get(&"key1"), None);
157    }
158
159    #[test]
160    fn test_size_limit() {
161        let cache = SafeLRUCache::new(10, Some(100));
162
163        cache.put("key1", vec![0u8, 1, 2], None, 50);
164        cache.put("key2", vec![3u8, 4, 5], None, 60);
165
166        // Both items should be evicted when adding a large one
167        cache.put("key3", vec![6u8, 7, 8, 9], None, 90);
168        assert_eq!(cache.get(&"key1"), None);
169        assert_eq!(cache.get(&"key2"), None);
170        assert_eq!(cache.get(&"key3"), Some(vec![6u8, 7, 8, 9]));
171    }
172
173    #[test]
174    fn test_update_existing() {
175        let cache = SafeLRUCache::new(2, None);
176
177        cache.put("key1", "value1", None, 1);
178        cache.put("key1", "value1_new", None, 1);
179        assert_eq!(cache.get(&"key1"), Some("value1_new"));
180    }
181}