reth_network/
cache.rs

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