reth_network/
cache.rs

1//! Network cache support
2
3use core::hash::BuildHasher;
4use derive_more::{Deref, DerefMut};
5use itertools::Itertools;
6use schnellru::{ByLength, Limiter, RandomState, Unlimited};
7use std::{fmt, hash::Hash};
8
9/// A minimal LRU cache based on a [`LruMap`](schnellru::LruMap) with limited capacity.
10///
11/// If the length exceeds the set capacity, the oldest element will be removed
12/// In the limit, for each element inserted the oldest existing element will be removed.
13pub struct LruCache<T: Hash + Eq + fmt::Debug> {
14    limit: u32,
15    inner: LruMap<T, ()>,
16}
17
18impl<T: Hash + Eq + fmt::Debug> LruCache<T> {
19    /// Creates a new [`LruCache`] using the given limit
20    pub fn new(limit: u32) -> Self {
21        // limit of lru map is one element more, so can give eviction feedback, which isn't
22        // supported by LruMap
23        Self { inner: LruMap::new(limit + 1), limit }
24    }
25
26    /// Insert an element into the set.
27    ///
28    /// This operation uses `get_or_insert` from the underlying `schnellru::LruMap` which:
29    /// - Automatically evicts the least recently used item if capacity is exceeded
30    ///
31    /// This method is more efficient than [`insert_and_get_evicted`](Self::insert_and_get_evicted)
32    /// as it performs fewer checks. Use this method when you don't need information about
33    /// evicted values.
34    ///
35    /// If the set did not have this value present, true is returned.
36    /// If the set did have this value present, false is returned.
37    pub fn insert(&mut self, entry: T) -> bool {
38        let mut is_new = false;
39        self.inner.get_or_insert(entry, || {
40            is_new = true;
41        });
42        is_new
43    }
44
45    /// Same as [`insert`](Self::insert) but returns a tuple, where the second index is the evicted
46    /// value, if one was evicted.
47    pub fn insert_and_get_evicted(&mut self, entry: T) -> (bool, Option<T>) {
48        let new = self.inner.peek(&entry).is_none();
49        let evicted =
50            (new && (self.limit as usize) <= self.inner.len()).then(|| self.remove_lru()).flatten();
51        _ = self.inner.get_or_insert(entry, || ());
52
53        (new, evicted)
54    }
55
56    /// Gets the given element, if exists, and promotes to lru.
57    pub fn get(&mut self, entry: &T) -> Option<&T> {
58        let _ = self.inner.get(entry)?;
59        self.iter().next()
60    }
61
62    /// Iterates through entries and returns a reference to the given entry, if exists, without
63    /// promoting to lru.
64    ///
65    /// NOTE: Use this for type that have custom impl of [`PartialEq`] and [`Eq`], that aren't
66    /// unique by all fields. If `PartialEq` and `Eq` are derived for a type, it's more efficient to
67    /// call [`contains`](Self::contains).
68    pub fn find(&self, entry: &T) -> Option<&T> {
69        self.iter().find(|key| *key == entry)
70    }
71
72    /// Remove the least recently used entry and return it.
73    ///
74    /// If the `LruCache` is empty or if the eviction feedback is
75    /// configured, this will return None.
76    #[inline]
77    fn remove_lru(&mut self) -> Option<T> {
78        self.inner.pop_oldest().map(|(k, ())| k)
79    }
80
81    /// Expels the given value. Returns true if the value existed.
82    pub fn remove(&mut self, value: &T) -> bool {
83        self.inner.remove(value).is_some()
84    }
85
86    /// Returns `true` if the set contains a value.
87    pub fn contains(&self, value: &T) -> bool {
88        self.inner.peek(value).is_some()
89    }
90
91    /// Returns an iterator over all cached entries in lru order
92    pub fn iter(&self) -> impl Iterator<Item = &T> + '_ {
93        self.inner.iter().map(|(k, ())| k)
94    }
95
96    /// Returns number of elements currently in cache.
97    pub fn len(&self) -> usize {
98        self.inner.len()
99    }
100
101    /// Returns `true` if there are currently no elements in the cache.
102    pub fn is_empty(&self) -> bool {
103        self.inner.is_empty()
104    }
105}
106
107impl<T> Extend<T> for LruCache<T>
108where
109    T: Eq + Hash + fmt::Debug,
110{
111    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
112        for item in iter {
113            _ = self.insert(item);
114        }
115    }
116}
117
118impl<T> fmt::Debug for LruCache<T>
119where
120    T: fmt::Debug + Hash + Eq,
121{
122    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
123        let mut debug_struct = f.debug_struct("LruCache");
124
125        debug_struct.field("limit", &self.limit);
126
127        debug_struct.field(
128            "ret %iter",
129            &format_args!("Iter: {{{} }}", self.iter().map(|k| format!(" {k:?}")).format(",")),
130        );
131
132        debug_struct.finish()
133    }
134}
135
136/// Wrapper of [`schnellru::LruMap`] that implements [`fmt::Debug`].
137#[derive(Deref, DerefMut, Default)]
138pub struct LruMap<K, V, L = ByLength, S = RandomState>(schnellru::LruMap<K, V, L, S>)
139where
140    K: Hash + PartialEq,
141    L: Limiter<K, V>,
142    S: BuildHasher;
143
144impl<K, V, L, S> fmt::Debug for LruMap<K, V, L, S>
145where
146    K: Hash + PartialEq + fmt::Display,
147    V: fmt::Debug,
148    L: Limiter<K, V> + fmt::Debug,
149    S: BuildHasher,
150{
151    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
152        let mut debug_struct = f.debug_struct("LruMap");
153
154        debug_struct.field("limiter", self.limiter());
155
156        debug_struct.field(
157            "ret %iter",
158            &format_args!(
159                "Iter: {{{} }}",
160                self.iter().map(|(k, v)| format!(" {k}: {v:?}")).format(",")
161            ),
162        );
163
164        debug_struct.finish()
165    }
166}
167
168impl<K, V> LruMap<K, V>
169where
170    K: Hash + PartialEq,
171{
172    /// Returns a new cache with default limiter and hash builder.
173    pub fn new(max_length: u32) -> Self {
174        Self(schnellru::LruMap::new(ByLength::new(max_length)))
175    }
176}
177
178impl<K, V> LruMap<K, V, Unlimited>
179where
180    K: Hash + PartialEq,
181{
182    /// Returns a new cache with [`Unlimited`] limiter and default hash builder.
183    pub fn new_unlimited() -> Self {
184        Self(schnellru::LruMap::new(Unlimited))
185    }
186}
187
188#[cfg(test)]
189mod test {
190    use super::*;
191    use derive_more::{Constructor, Display};
192    use std::hash::Hasher;
193
194    #[derive(Debug, Hash, PartialEq, Eq, Display, Clone, Copy)]
195    struct Key(i8);
196
197    #[derive(Debug, Eq, Constructor, Clone, Copy)]
198    struct CompoundKey {
199        // type unique for id
200        id: i8,
201        other: i8,
202    }
203
204    impl PartialEq for CompoundKey {
205        fn eq(&self, other: &Self) -> bool {
206            self.id == other.id
207        }
208    }
209
210    impl Hash for CompoundKey {
211        fn hash<H: Hasher>(&self, state: &mut H) {
212            self.id.hash(state)
213        }
214    }
215
216    #[test]
217    fn test_cache_should_insert_into_empty_set() {
218        let mut cache = LruCache::new(5);
219        let entry = "entry";
220        assert!(cache.insert(entry));
221        assert!(cache.contains(&entry));
222    }
223
224    #[test]
225    fn test_cache_should_not_insert_same_value_twice() {
226        let mut cache = LruCache::new(5);
227        let entry = "entry";
228        assert!(cache.insert(entry));
229        assert!(!cache.insert(entry));
230    }
231
232    #[test]
233    fn test_cache_should_remove_oldest_element_when_exceeding_limit() {
234        let mut cache = LruCache::new(1); // LruCache limit will be 2, check LruCache::new
235        let old_entry = "old_entry";
236        let new_entry = "new_entry";
237        cache.insert(old_entry);
238        cache.insert("entry");
239        cache.insert(new_entry);
240        assert!(cache.contains(&new_entry));
241        assert!(!cache.contains(&old_entry));
242    }
243
244    #[test]
245    fn test_cache_should_extend_an_array() {
246        let mut cache = LruCache::new(5);
247        let entries = ["some_entry", "another_entry"];
248        cache.extend(entries);
249        for e in entries {
250            assert!(cache.contains(&e));
251        }
252    }
253
254    #[test]
255    #[expect(dead_code)]
256    fn test_debug_impl_lru_map() {
257        #[derive(Debug)]
258        struct Value(i8);
259
260        let mut cache = LruMap::new(2);
261        let key_1 = Key(1);
262        let value_1 = Value(11);
263        cache.insert(key_1, value_1);
264        let key_2 = Key(2);
265        let value_2 = Value(22);
266        cache.insert(key_2, value_2);
267
268        assert_eq!("LruMap { limiter: ByLength { max_length: 2 }, ret %iter: Iter: { 2: Value(22), 1: Value(11) } }", format!("{cache:?}"))
269    }
270
271    #[test]
272    fn test_debug_impl_lru_cache() {
273        let mut cache = LruCache::new(2);
274        let key_1 = Key(1);
275        cache.insert(key_1);
276        let key_2 = Key(2);
277        cache.insert(key_2);
278
279        assert_eq!(
280            "LruCache { limit: 2, ret %iter: Iter: { Key(2), Key(1) } }",
281            format!("{cache:?}")
282        )
283    }
284
285    #[test]
286    fn get() {
287        let mut cache = LruCache::new(2);
288        let key_1 = Key(1);
289        cache.insert(key_1);
290        let key_2 = Key(2);
291        cache.insert(key_2);
292
293        // promotes key 1 to lru
294        _ = cache.get(&key_1);
295
296        assert_eq!(
297            "LruCache { limit: 2, ret %iter: Iter: { Key(1), Key(2) } }",
298            format!("{cache:?}")
299        )
300    }
301
302    #[test]
303    fn get_ty_custom_eq_impl() {
304        let mut cache = LruCache::new(2);
305        let key_1 = CompoundKey::new(1, 11);
306        cache.insert(key_1);
307        let key_2 = CompoundKey::new(2, 22);
308        cache.insert(key_2);
309
310        let key = cache.get(&key_1);
311
312        assert_eq!(key_1.other, key.unwrap().other)
313    }
314
315    #[test]
316    fn peek() {
317        let mut cache = LruCache::new(2);
318        let key_1 = Key(1);
319        cache.insert(key_1);
320        let key_2 = Key(2);
321        cache.insert(key_2);
322
323        // doesn't promote key 1 to lru
324        _ = cache.find(&key_1);
325
326        assert_eq!(
327            "LruCache { limit: 2, ret %iter: Iter: { Key(2), Key(1) } }",
328            format!("{cache:?}")
329        )
330    }
331
332    #[test]
333    fn peek_ty_custom_eq_impl() {
334        let mut cache = LruCache::new(2);
335        let key_1 = CompoundKey::new(1, 11);
336        cache.insert(key_1);
337        let key_2 = CompoundKey::new(2, 22);
338        cache.insert(key_2);
339
340        let key = cache.find(&key_1);
341
342        assert_eq!(key_1.other, key.unwrap().other)
343    }
344
345    #[test]
346    fn test_insert_methods() {
347        let mut cache = LruCache::new(2);
348
349        // Test basic insert
350        assert!(cache.insert("first")); // new entry
351        assert!(!cache.insert("first")); // existing entry
352        assert!(cache.insert("second")); // new entry
353
354        // Test insert_and_get_evicted
355        let (is_new, evicted) = cache.insert_and_get_evicted("third");
356        assert!(is_new); // should be new entry
357        assert_eq!(evicted, Some("first")); // should evict
358
359        assert!(cache.contains(&"second"));
360        assert!(cache.contains(&"third"));
361        assert!(!cache.contains(&"first"));
362
363        // Test insert_and_get_evicted with existing entry
364        let (is_new, evicted) = cache.insert_and_get_evicted("second");
365        assert!(!is_new); // should not be new
366        assert_eq!(evicted, None); // should not evict anything
367    }
368}