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 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 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 if !inner.map.contains_key(key) {
77 return None;
78 }
79
80 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 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 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}