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    #[allow(dead_code)]
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    #[allow(dead_code)]
104    pub fn is_empty(&self) -> bool {
105        self.inner.is_empty()
106    }
107}
108
109impl<T> Extend<T> for LruCache<T>
110where
111    T: Eq + Hash + fmt::Debug,
112{
113    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
114        for item in iter {
115            _ = self.insert(item);
116        }
117    }
118}
119
120impl<T> fmt::Debug for LruCache<T>
121where
122    T: fmt::Debug + Hash + Eq,
123{
124    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
125        let mut debug_struct = f.debug_struct("LruCache");
126
127        debug_struct.field("limit", &self.limit);
128
129        debug_struct.field(
130            "ret %iter",
131            &format_args!("Iter: {{{} }}", self.iter().map(|k| format!(" {k:?}")).format(",")),
132        );
133
134        debug_struct.finish()
135    }
136}
137
138/// Wrapper of [`schnellru::LruMap`] that implements [`fmt::Debug`].
139#[derive(Deref, DerefMut, Default)]
140pub struct LruMap<K, V, L = ByLength, S = RandomState>(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::new(ByLength::new(max_length)))
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::new(Unlimited))
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    #[allow(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!("LruMap { limiter: ByLength { max_length: 2 }, ret %iter: Iter: { 2: Value(22), 1: Value(11) } }", format!("{cache:?}"))
271    }
272
273    #[test]
274    #[allow(dead_code)]
275    fn test_debug_impl_lru_cache() {
276        let mut cache = LruCache::new(2);
277        let key_1 = Key(1);
278        cache.insert(key_1);
279        let key_2 = Key(2);
280        cache.insert(key_2);
281
282        assert_eq!(
283            "LruCache { limit: 2, ret %iter: Iter: { Key(2), Key(1) } }",
284            format!("{cache:?}")
285        )
286    }
287
288    #[test]
289    fn get() {
290        let mut cache = LruCache::new(2);
291        let key_1 = Key(1);
292        cache.insert(key_1);
293        let key_2 = Key(2);
294        cache.insert(key_2);
295
296        // promotes key 1 to lru
297        _ = cache.get(&key_1);
298
299        assert_eq!(
300            "LruCache { limit: 2, ret %iter: Iter: { Key(1), Key(2) } }",
301            format!("{cache:?}")
302        )
303    }
304
305    #[test]
306    fn get_ty_custom_eq_impl() {
307        let mut cache = LruCache::new(2);
308        let key_1 = CompoundKey::new(1, 11);
309        cache.insert(key_1);
310        let key_2 = CompoundKey::new(2, 22);
311        cache.insert(key_2);
312
313        let key = cache.get(&key_1);
314
315        assert_eq!(key_1.other, key.unwrap().other)
316    }
317
318    #[test]
319    fn peek() {
320        let mut cache = LruCache::new(2);
321        let key_1 = Key(1);
322        cache.insert(key_1);
323        let key_2 = Key(2);
324        cache.insert(key_2);
325
326        // doesn't promote key 1 to lru
327        _ = cache.find(&key_1);
328
329        assert_eq!(
330            "LruCache { limit: 2, ret %iter: Iter: { Key(2), Key(1) } }",
331            format!("{cache:?}")
332        )
333    }
334
335    #[test]
336    fn peek_ty_custom_eq_impl() {
337        let mut cache = LruCache::new(2);
338        let key_1 = CompoundKey::new(1, 11);
339        cache.insert(key_1);
340        let key_2 = CompoundKey::new(2, 22);
341        cache.insert(key_2);
342
343        let key = cache.find(&key_1);
344
345        assert_eq!(key_1.other, key.unwrap().other)
346    }
347
348    #[test]
349    fn test_insert_methods() {
350        let mut cache = LruCache::new(2);
351
352        // Test basic insert
353        assert!(cache.insert("first")); // new entry
354        assert!(!cache.insert("first")); // existing entry
355        assert!(cache.insert("second")); // new entry
356
357        // Test insert_and_get_evicted
358        let (is_new, evicted) = cache.insert_and_get_evicted("third");
359        assert!(is_new); // should be new entry
360        assert_eq!(evicted, Some("first")); // should evict
361
362        assert!(cache.contains(&"second"));
363        assert!(cache.contains(&"third"));
364        assert!(!cache.contains(&"first"));
365
366        // Test insert_and_get_evicted with existing entry
367        let (is_new, evicted) = cache.insert_and_get_evicted("second");
368        assert!(!is_new); // should not be new
369        assert_eq!(evicted, None); // should not evict anything
370    }
371}
reth_network/cache.rs

reth_network/
cache.rs
