Skip to main content

reth_rpc_eth_types/cache/
multi_consumer.rs

1//! Metered cache, which also provides storage for senders in order to queue queries that result in
2//! a cache miss.
3
4use super::metrics::CacheMetrics;
5use reth_primitives_traits::InMemorySize;
6use schnellru::{ByLength, Limiter, LruMap};
7use std::{
8    collections::{hash_map::Entry, HashMap},
9    fmt::{self, Debug, Formatter},
10    hash::Hash,
11};
12
13/// A multi-consumer LRU cache.
14pub struct MultiConsumerLruCache<K, V, L, S>
15where
16    K: Hash + Eq,
17    L: Limiter<K, V>,
18{
19    /// The LRU cache.
20    cache: LruMap<K, V, L>,
21    /// All queued consumers.
22    queued: HashMap<K, Vec<S>>,
23    /// Cache metrics
24    metrics: CacheMetrics,
25    // Tracked heap usage
26    memory_usage: usize,
27}
28
29impl<K, V, L, S> Debug for MultiConsumerLruCache<K, V, L, S>
30where
31    K: Hash + Eq,
32    L: Limiter<K, V>,
33{
34    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
35        f.debug_struct("MultiConsumerLruCache")
36            .field("cache_length", &self.cache.len())
37            .field("cache_memory_usage", &self.cache.memory_usage())
38            .field("queued_length", &self.queued.len())
39            .field("memory_usage", &self.memory_usage)
40            .finish()
41    }
42}
43
44impl<K, V, L, S> MultiConsumerLruCache<K, V, L, S>
45where
46    K: Hash + Eq + Debug,
47    L: Limiter<K, V>,
48{
49    /// Adds the sender to the queue for the given key.
50    ///
51    /// Returns true if this is the first queued sender for the key
52    pub fn queue(&mut self, key: K, sender: S) -> bool {
53        self.metrics.queued_consumers_count.increment(1.0);
54        match self.queued.entry(key) {
55            Entry::Occupied(mut entry) => {
56                entry.get_mut().push(sender);
57                false
58            }
59            Entry::Vacant(entry) => {
60                entry.insert(vec![sender]);
61                true
62            }
63        }
64    }
65
66    /// Remove consumers for a given key, this will also remove the key from the cache.
67    pub fn remove(&mut self, key: &K) -> Option<Vec<S>>
68    where
69        V: InMemorySize,
70    {
71        self.cache
72            .remove(key)
73            .inspect(|value| self.memory_usage = self.memory_usage.saturating_sub(value.size()));
74        self.queued
75            .remove(key)
76            .inspect(|removed| self.metrics.queued_consumers_count.decrement(removed.len() as f64))
77    }
78
79    /// Returns a reference to the value for a given key and promotes that element to be the most
80    /// recently used.
81    pub fn get(&mut self, key: &K) -> Option<&mut V> {
82        let entry = self.cache.get(key);
83        if entry.is_some() {
84            self.metrics.hits_total.increment(1);
85        } else {
86            self.metrics.misses_total.increment(1);
87        }
88        entry
89    }
90
91    /// Inserts a new element into the map.
92    ///
93    /// Can fail if the element is rejected by the limiter or if we fail to grow an empty map.
94    ///
95    /// See [`LruMap::insert`] for more info.
96    pub fn insert<'a>(&mut self, key: L::KeyToInsert<'a>, value: V) -> bool
97    where
98        L::KeyToInsert<'a>: Hash + PartialEq<K>,
99        V: InMemorySize,
100    {
101        let size = value.size();
102
103        if self.cache.limiter().is_over_the_limit(self.cache.len() + 1) &&
104            let Some((_, evicted)) = self.cache.pop_oldest()
105        {
106            // update tracked memory with the evicted value
107            self.memory_usage = self.memory_usage.saturating_sub(evicted.size());
108        }
109
110        if self.cache.insert(key, value) {
111            self.memory_usage = self.memory_usage.saturating_add(size);
112            true
113        } else {
114            false
115        }
116    }
117
118    /// Shrinks the capacity of the queue with a lower limit.
119    #[inline]
120    pub fn shrink_to(&mut self, min_capacity: usize) {
121        self.queued.shrink_to(min_capacity);
122    }
123
124    /// Update metrics for the inner cache.
125    #[inline]
126    pub fn update_cached_metrics(&self) {
127        self.metrics.cached_count.set(self.cache.len() as f64);
128        self.metrics.memory_usage.set(self.memory_usage as f64);
129    }
130}
131
132impl<K, V, S> MultiConsumerLruCache<K, V, ByLength, S>
133where
134    K: Hash + Eq,
135{
136    /// Creates a new empty map with a given `max_len` and metric label.
137    pub fn new(max_len: u32, cache_id: &str) -> Self {
138        Self {
139            cache: LruMap::new(ByLength::new(max_len)),
140            queued: Default::default(),
141            metrics: CacheMetrics::new_with_labels(&[("cache", cache_id.to_string())]),
142            memory_usage: 0,
143        }
144    }
145}