forest/utils/cache/
lru.rs

1// Copyright 2019-2025 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use std::{
5    borrow::{Borrow, Cow},
6    fmt::Debug,
7    hash::Hash,
8    num::NonZeroUsize,
9    sync::{
10        Arc,
11        atomic::{AtomicUsize, Ordering},
12    },
13};
14
15use get_size2::GetSize;
16use hashlink::LruCache;
17use parking_lot::RwLock;
18use prometheus_client::{
19    collector::Collector,
20    encoding::{DescriptorEncoder, EncodeMetric},
21    metrics::gauge::Gauge,
22    registry::Unit,
23};
24
25pub trait KeyConstraints:
26    GetSize + Debug + Send + Sync + Hash + PartialEq + Eq + Clone + 'static
27{
28}
29
30impl<T> KeyConstraints for T where
31    T: GetSize + Debug + Send + Sync + Hash + PartialEq + Eq + Clone + 'static
32{
33}
34
35pub trait LruValueConstraints: GetSize + Debug + Send + Sync + Clone + 'static {}
36
37impl<T> LruValueConstraints for T where T: GetSize + Debug + Send + Sync + Clone + 'static {}
38
39#[derive(Debug, Clone)]
40pub struct SizeTrackingLruCache<K, V>
41where
42    K: KeyConstraints,
43    V: LruValueConstraints,
44{
45    cache_id: usize,
46    cache_name: Cow<'static, str>,
47    cache: Arc<RwLock<LruCache<K, V>>>,
48}
49
50impl<K, V> SizeTrackingLruCache<K, V>
51where
52    K: KeyConstraints,
53    V: LruValueConstraints,
54{
55    fn register_metrics(&self) {
56        crate::metrics::register_collector(Box::new(self.clone()));
57    }
58
59    fn new_inner(cache_name: Cow<'static, str>, capacity: Option<NonZeroUsize>) -> Self {
60        static ID_GENERATOR: AtomicUsize = AtomicUsize::new(0);
61
62        Self {
63            cache_id: ID_GENERATOR.fetch_add(1, Ordering::Relaxed),
64            cache_name,
65            #[allow(clippy::disallowed_methods)]
66            cache: Arc::new(RwLock::new(
67                capacity
68                    .map(From::from)
69                    .map(LruCache::new)
70                    // For constructing lru cache that is bounded by memory usage instead of length
71                    .unwrap_or_else(LruCache::new_unbounded),
72            )),
73        }
74    }
75
76    pub fn new_without_metrics_registry(
77        cache_name: Cow<'static, str>,
78        capacity: NonZeroUsize,
79    ) -> Self {
80        Self::new_inner(cache_name, Some(capacity))
81    }
82
83    pub fn new_with_metrics(cache_name: Cow<'static, str>, capacity: NonZeroUsize) -> Self {
84        let c = Self::new_without_metrics_registry(cache_name, capacity);
85        c.register_metrics();
86        c
87    }
88
89    pub fn unbounded_without_metrics_registry(cache_name: Cow<'static, str>) -> Self {
90        Self::new_inner(cache_name, None)
91    }
92
93    pub fn unbounded_with_metrics(cache_name: Cow<'static, str>) -> Self {
94        let c = Self::unbounded_without_metrics_registry(cache_name);
95        c.register_metrics();
96        c
97    }
98
99    pub fn cache(&self) -> &Arc<RwLock<LruCache<K, V>>> {
100        &self.cache
101    }
102
103    pub fn push(&self, k: K, v: V) -> Option<V> {
104        self.cache.write().insert(k, v)
105    }
106
107    pub fn contains<Q>(&self, k: &Q) -> bool
108    where
109        K: Borrow<Q>,
110        Q: Hash + Eq + ?Sized,
111    {
112        self.cache.read().contains_key(k)
113    }
114
115    pub fn get_cloned<Q>(&self, k: &Q) -> Option<V>
116    where
117        K: Borrow<Q>,
118        Q: Hash + Eq + ?Sized,
119    {
120        self.cache.write().get(k).cloned()
121    }
122
123    pub fn peek_cloned<Q>(&self, k: &Q) -> Option<V>
124    where
125        K: Borrow<Q>,
126        Q: Hash + Eq + ?Sized,
127    {
128        self.cache.read().peek(k).cloned()
129    }
130
131    pub fn pop_lru(&self) -> Option<(K, V)> {
132        self.cache.write().remove_lru()
133    }
134
135    pub fn len(&self) -> usize {
136        self.cache.read().len()
137    }
138
139    pub fn cap(&self) -> usize {
140        self.cache.read().capacity()
141    }
142
143    pub(crate) fn size_in_bytes(&self) -> usize {
144        let mut size = 0_usize;
145        for (k, v) in self.cache.read().iter() {
146            size = size
147                .saturating_add(k.get_size())
148                .saturating_add(v.get_size());
149        }
150        size
151    }
152}
153
154impl<K, V> Collector for SizeTrackingLruCache<K, V>
155where
156    K: KeyConstraints,
157    V: LruValueConstraints,
158{
159    fn encode(&self, mut encoder: DescriptorEncoder) -> Result<(), std::fmt::Error> {
160        {
161            let size_in_bytes = {
162                let g: Gauge = Default::default();
163                g.set(self.size_in_bytes() as _);
164                g
165            };
166            let size_metric_name = format!("cache_{}_{}_size", self.cache_name, self.cache_id);
167            let size_metric_help = format!(
168                "Size of LruCache {}_{} in bytes",
169                self.cache_name, self.cache_id
170            );
171            let size_metric_encoder = encoder.encode_descriptor(
172                &size_metric_name,
173                &size_metric_help,
174                Some(&Unit::Bytes),
175                size_in_bytes.metric_type(),
176            )?;
177            size_in_bytes.encode(size_metric_encoder)?;
178        }
179        {
180            let len_metric_name = format!("{}_{}_len", self.cache_name, self.cache_id);
181            let len_metric_help =
182                format!("Length of LruCache {}_{}", self.cache_name, self.cache_id);
183            let len: Gauge = Default::default();
184            len.set(self.len() as _);
185            let len_metric_encoder = encoder.encode_descriptor(
186                &len_metric_name,
187                &len_metric_help,
188                None,
189                len.metric_type(),
190            )?;
191            len.encode(len_metric_encoder)?;
192        }
193        {
194            let cap_metric_name = format!("{}_{}_cap", self.cache_name, self.cache_id);
195            let cap_metric_help =
196                format!("Capacity of LruCache {}_{}", self.cache_name, self.cache_id);
197            let cap: Gauge = Default::default();
198            cap.set(self.cap() as _);
199            let cap_metric_encoder = encoder.encode_descriptor(
200                &cap_metric_name,
201                &cap_metric_help,
202                None,
203                cap.metric_type(),
204            )?;
205            cap.encode(cap_metric_encoder)?;
206        }
207
208        Ok(())
209    }
210}
forest/utils/cache/lru.rs

forest/utils/cache/
lru.rs
