kapot_cache/backend/policy/lru/
lru_cache.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::backend::policy::lru::ResourceCounter;
19use crate::backend::policy::{lru::LruCachePolicy, CachePolicy, CachePolicyPutResult};
20use hashbrown::hash_map::DefaultHashBuilder;
21use hashlink::linked_hash_map::{self, IntoIter, Iter, IterMut, LinkedHashMap};
22use std::any::Any;
23use std::fmt;
24use std::fmt::Debug;
25use std::hash::{BuildHasher, Hash};
26
27pub struct LruCache<K, V, H = DefaultHashBuilder>
28where
29    K: Clone + Eq + Hash + Ord + Debug + Send + 'static,
30    V: Clone + Debug + Send + 'static,
31{
32    map: LinkedHashMap<K, V, H>,
33    resource_counter: Box<dyn ResourceCounter<K = K, V = V>>,
34}
35
36impl<K, V> LruCache<K, V>
37where
38    K: Clone + Eq + Hash + Ord + Debug + Send + 'static,
39    V: Clone + Debug + Send + 'static,
40{
41    pub fn with_resource_counter<R>(resource_counter: R) -> Self
42    where
43        R: ResourceCounter<K = K, V = V>,
44    {
45        LruCache {
46            map: LinkedHashMap::new(),
47            resource_counter: Box::new(resource_counter),
48        }
49    }
50}
51
52impl<K, V, H> LruCache<K, V, H>
53where
54    K: Clone + Eq + Hash + Ord + Debug + Send + 'static,
55    V: Clone + Debug + Send + 'static,
56{
57    pub fn with_resource_counter_and_hasher<R>(
58        resource_counter: R,
59        hash_builder: H,
60    ) -> Self
61    where
62        R: ResourceCounter<K = K, V = V>,
63    {
64        LruCache {
65            map: LinkedHashMap::with_hasher(hash_builder),
66            resource_counter: Box::new(resource_counter),
67        }
68    }
69
70    pub fn len(&self) -> usize {
71        self.map.len()
72    }
73
74    pub fn is_empty(&self) -> bool {
75        self.map.is_empty()
76    }
77
78    pub fn iter(&self) -> Iter<K, V> {
79        self.map.iter()
80    }
81
82    pub fn iter_mut(&mut self) -> IterMut<K, V> {
83        self.map.iter_mut()
84    }
85}
86
87impl<K, V, H> LruCachePolicy for LruCache<K, V, H>
88where
89    K: 'static + Clone + Debug + Eq + Hash + Ord + Send,
90    H: 'static + BuildHasher + Debug + Send,
91    V: 'static + Clone + Debug + Send,
92{
93    fn get_lru(&mut self, k: &Self::K) -> Option<Self::V> {
94        match self.map.raw_entry_mut().from_key(k) {
95            linked_hash_map::RawEntryMut::Occupied(mut occupied) => {
96                occupied.to_back();
97                Some(occupied.into_mut().clone())
98            }
99            linked_hash_map::RawEntryMut::Vacant(_) => None,
100        }
101    }
102
103    fn put_lru(
104        &mut self,
105        k: Self::K,
106        v: Self::V,
107    ) -> CachePolicyPutResult<Self::K, Self::V> {
108        let old_val = self.map.insert(k.clone(), v.clone());
109        // Consume resources for (k, v)
110        self.resource_counter.consume(&k, &v);
111        // Restore resources for old (k, old_val)
112        if let Some(old_val) = &old_val {
113            self.resource_counter.restore(&k, old_val);
114        }
115
116        let mut popped_entries = vec![];
117        while self.resource_counter.exceed_capacity() {
118            if let Some(entry) = self.pop_lru() {
119                popped_entries.push(entry);
120            }
121        }
122        (old_val, popped_entries)
123    }
124
125    fn pop_lru(&mut self) -> Option<(Self::K, Self::V)> {
126        if let Some(entry) = self.map.pop_front() {
127            self.resource_counter.restore(&entry.0, &entry.1);
128            Some(entry)
129        } else {
130            None
131        }
132    }
133}
134
135impl<K, V, H> CachePolicy for LruCache<K, V, H>
136where
137    K: 'static + Clone + Debug + Eq + Hash + Ord + Send,
138    H: 'static + BuildHasher + Debug + Send,
139    V: 'static + Clone + Debug + Send,
140{
141    type K = K;
142    type V = V;
143
144    fn get(&mut self, k: &Self::K) -> Option<Self::V> {
145        self.get_lru(k)
146    }
147
148    fn peek(&mut self, k: &Self::K) -> Option<Self::V> {
149        self.map.get(k).cloned()
150    }
151
152    fn put(&mut self, k: Self::K, v: Self::V) -> CachePolicyPutResult<Self::K, Self::V> {
153        self.put_lru(k, v)
154    }
155
156    fn remove(&mut self, k: &Self::K) -> Option<Self::V> {
157        if let Some(v) = self.map.remove(k) {
158            self.resource_counter.restore(k, &v);
159            Some(v)
160        } else {
161            None
162        }
163    }
164
165    fn pop(&mut self) -> Option<(Self::K, Self::V)> {
166        self.pop_lru()
167    }
168
169    fn as_any(&self) -> &dyn Any {
170        self
171    }
172}
173
174impl<K, V, H> IntoIterator for LruCache<K, V, H>
175where
176    K: 'static + Clone + Debug + Eq + Hash + Ord + Send,
177    V: 'static + Clone + Debug + Send,
178{
179    type Item = (K, V);
180    type IntoIter = IntoIter<K, V>;
181
182    fn into_iter(self) -> IntoIter<K, V> {
183        self.map.into_iter()
184    }
185}
186
187impl<'a, K, V, H> IntoIterator for &'a LruCache<K, V, H>
188where
189    K: 'static + Clone + Debug + Eq + Hash + Ord + Send,
190    V: 'static + Clone + Debug + Send,
191{
192    type Item = (&'a K, &'a V);
193    type IntoIter = Iter<'a, K, V>;
194
195    fn into_iter(self) -> Iter<'a, K, V> {
196        self.iter()
197    }
198}
199
200impl<'a, K, V, H> IntoIterator for &'a mut LruCache<K, V, H>
201where
202    K: 'static + Clone + Debug + Eq + Hash + Ord + Send,
203    V: 'static + Clone + Debug + Send,
204{
205    type Item = (&'a K, &'a mut V);
206    type IntoIter = IterMut<'a, K, V>;
207
208    fn into_iter(self) -> IterMut<'a, K, V> {
209        self.iter_mut()
210    }
211}
212
213impl<K, V, H> Debug for LruCache<K, V, H>
214where
215    K: 'static + Clone + Debug + Eq + Hash + Ord + Send,
216    V: 'static + Clone + Debug + Send,
217{
218    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
219        f.debug_map().entries(self.iter().rev()).finish()
220    }
221}
222
223#[cfg(test)]
224mod tests {
225    use crate::backend::policy::lru::lru_cache::LruCache;
226    use crate::backend::policy::lru::{DefaultResourceCounter, ResourceCounter};
227    use crate::backend::policy::CachePolicy;
228    use hashbrown::HashMap;
229
230    #[test]
231    fn test_cache_with_lru_policy() {
232        let mut cache = LruCache::with_resource_counter(DefaultResourceCounter::new(3));
233
234        cache.put("1".to_string(), "file1".to_string());
235        cache.put("2".to_string(), "file2".to_string());
236        cache.put("3".to_string(), "file3".to_string());
237        assert_eq!(3, cache.len());
238
239        cache.put("4".to_string(), "file4".to_string());
240        assert_eq!(3, cache.len());
241        assert!(cache.peek(&"1".to_string()).is_none());
242
243        assert!(cache.peek(&"2".to_string()).is_some());
244        let mut iter = cache.iter();
245        assert_eq!("2", iter.next().unwrap().0);
246        assert_eq!("3", iter.next().unwrap().0);
247        assert_eq!("4", iter.next().unwrap().0);
248
249        assert!(cache.get(&"2".to_string()).is_some());
250        let mut iter = cache.iter();
251        assert_eq!("3", iter.next().unwrap().0);
252        assert_eq!("4", iter.next().unwrap().0);
253        assert_eq!("2", iter.next().unwrap().0);
254
255        assert_eq!(Some("file4".to_string()), cache.remove(&"4".to_string()));
256
257        assert_eq!("3".to_string(), cache.pop().unwrap().0);
258        assert_eq!("2".to_string(), cache.pop().unwrap().0);
259        assert!(cache.pop().is_none());
260    }
261
262    #[test]
263    fn test_cache_with_size_resource_counter() {
264        let mut cache =
265            LruCache::with_resource_counter(get_test_size_resource_counter(50));
266
267        cache.put("1".to_string(), "file1".to_string());
268        cache.put("2".to_string(), "file2".to_string());
269        cache.put("3".to_string(), "file3".to_string());
270        assert_eq!(3, cache.len());
271
272        cache.put("4".to_string(), "file4".to_string());
273        assert_eq!(2, cache.len());
274        assert!(cache.peek(&"1".to_string()).is_none());
275        assert!(cache.peek(&"2".to_string()).is_none());
276
277        assert!(cache.peek(&"3".to_string()).is_some());
278        let mut iter = cache.iter();
279        assert_eq!("3", iter.next().unwrap().0);
280        assert_eq!("4", iter.next().unwrap().0);
281
282        assert!(cache.get(&"3".to_string()).is_some());
283        let mut iter = cache.iter();
284        assert_eq!("4", iter.next().unwrap().0);
285        assert_eq!("3", iter.next().unwrap().0);
286
287        cache.put("5".to_string(), "file5".to_string());
288        cache.put("3".to_string(), "file3-bak".to_string());
289        cache.put("1".to_string(), "file1".to_string());
290        let mut iter = cache.iter();
291        assert_eq!("3", iter.next().unwrap().0);
292        assert_eq!("1", iter.next().unwrap().0);
293        assert!(iter.next().is_none());
294    }
295
296    fn get_test_size_resource_counter(max_size: usize) -> TestSizeResourceCounter {
297        let mut size_map = HashMap::new();
298        size_map.insert(("1".to_string(), "file1".to_string()), 10);
299        size_map.insert(("2".to_string(), "file2".to_string()), 20);
300        size_map.insert(("3".to_string(), "file3".to_string()), 15);
301        size_map.insert(("3".to_string(), "file3-bak".to_string()), 30);
302        size_map.insert(("4".to_string(), "file4".to_string()), 35);
303        size_map.insert(("5".to_string(), "file5".to_string()), 25);
304
305        TestSizeResourceCounter {
306            size_map,
307            max_size,
308            current_size: 0,
309        }
310    }
311
312    #[derive(Debug)]
313    struct TestSizeResourceCounter {
314        size_map: HashMap<(String, String), usize>,
315        max_size: usize,
316        current_size: usize,
317    }
318
319    impl ResourceCounter for TestSizeResourceCounter {
320        type K = String;
321        type V = String;
322
323        fn consume(&mut self, k: &Self::K, v: &Self::V) {
324            let s = self.size_map.get(&(k.clone(), v.clone())).unwrap();
325            self.current_size += s;
326        }
327
328        fn restore(&mut self, k: &Self::K, v: &Self::V) {
329            let s = self.size_map.get(&(k.clone(), v.clone())).unwrap();
330            self.current_size -= s;
331        }
332
333        fn exceed_capacity(&self) -> bool {
334            self.current_size > self.max_size
335        }
336    }
337}
kapot_cache/backend/policy/lru/lru_cache.rs

kapot_cache/backend/policy/lru/
lru_cache.rs
