forest/utils/cache/
size_tracking.rs

1// Copyright 2019-2026 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use std::{borrow::Cow, fmt::Debug, hash::Hash, num::NonZeroUsize, sync::atomic::Ordering};
5
6use get_size2::GetSize;
7use prometheus_client::{
8    collector::Collector,
9    encoding::{DescriptorEncoder, EncodeMetric},
10    metrics::{counter::Counter, gauge::Gauge},
11    registry::Unit,
12};
13use quick_cache::sync::Cache;
14
15use crate::prelude::*;
16
17pub trait CacheKeyConstraints:
18    GetSize + Debug + Send + Sync + Hash + PartialEq + Eq + Clone + 'static
19{
20}
21
22impl<T> CacheKeyConstraints for T where
23    T: GetSize + Debug + Send + Sync + Hash + PartialEq + Eq + Clone + 'static
24{
25}
26
27pub trait CacheValueConstraints: GetSize + Debug + Send + Sync + Clone + 'static {}
28
29impl<T> CacheValueConstraints for T where T: GetSize + Debug + Send + Sync + Clone + 'static {}
30
31/// A concurrent cache with Prometheus instrumentation.
32///
33/// Backed by [`quick_cache::sync::Cache`], which uses the scan-resistant
34/// CLOCK-PRO eviction policy. Tracks total entry size in bytes for
35/// observability.
36#[derive(Debug, derive_more::Deref)]
37pub struct SizeTrackingCache<K, V>
38where
39    K: CacheKeyConstraints,
40    V: CacheValueConstraints,
41{
42    cache_name: Cow<'static, str>,
43    #[deref]
44    cache: Arc<Cache<K, V>>,
45}
46
47impl<K, V> ShallowClone for SizeTrackingCache<K, V>
48where
49    K: CacheKeyConstraints,
50    V: CacheValueConstraints,
51{
52    fn shallow_clone(&self) -> Self {
53        Self {
54            cache_name: self.cache_name.clone(),
55            cache: self.cache.shallow_clone(),
56        }
57    }
58}
59
60impl<K, V> SizeTrackingCache<K, V>
61where
62    K: CacheKeyConstraints,
63    V: CacheValueConstraints,
64{
65    fn register_metrics(&self) {
66        crate::metrics::register_collector(Box::new(self.shallow_clone()));
67    }
68
69    fn new_inner(cache_name: impl Into<Cow<'static, str>>, capacity: NonZeroUsize) -> Self {
70        Self {
71            cache_name: cache_name.into(),
72            cache: Arc::new(Cache::new(capacity.get())),
73        }
74    }
75
76    pub fn new_without_metrics_registry(
77        cache_name: impl Into<Cow<'static, str>>,
78        capacity: NonZeroUsize,
79    ) -> Self {
80        Self::new_inner(cache_name, capacity)
81    }
82
83    pub fn new_with_metrics(
84        cache_name: impl Into<Cow<'static, str>>,
85        capacity: NonZeroUsize,
86    ) -> Self {
87        let c = Self::new_without_metrics_registry(cache_name, capacity);
88        c.register_metrics();
89        c
90    }
91
92    /// Insert `k`/`v`. If a previous entry existed for `k`, return it.
93    ///
94    /// `quick_cache::sync::Cache::insert` does not return the displaced
95    /// value, so this is a peek-then-insert. The two steps are not atomic;
96    /// concurrent callers for the same key may both observe `None`. None of
97    /// the existing callers depend on atomicity here.
98    #[inline]
99    pub fn push_and_get_prev(&self, k: K, v: V) -> Option<V> {
100        let prev = self.cache.peek(&k);
101        self.cache.insert(k, v);
102        prev
103    }
104
105    pub(crate) fn size_in_bytes(&self) -> usize {
106        let mut size = 0_usize;
107        for (k, v) in self.cache.iter() {
108            size = size
109                .saturating_add(k.get_size())
110                .saturating_add(v.get_size());
111        }
112        size
113    }
114}
115
116impl<K, V> Collector for SizeTrackingCache<K, V>
117where
118    K: CacheKeyConstraints,
119    V: CacheValueConstraints,
120{
121    fn encode(&self, mut encoder: DescriptorEncoder) -> Result<(), std::fmt::Error> {
122        {
123            let size_in_bytes = {
124                let g: Gauge = Default::default();
125                g.set(self.size_in_bytes() as _);
126                g
127            };
128            let size_metric_name = format!("cache_{}_size", self.cache_name);
129            let size_metric_help = format!("Size of cache {} in bytes", self.cache_name);
130            let size_metric_encoder = encoder.encode_descriptor(
131                &size_metric_name,
132                &size_metric_help,
133                Some(&Unit::Bytes),
134                size_in_bytes.metric_type(),
135            )?;
136            size_in_bytes.encode(size_metric_encoder)?;
137        }
138        {
139            let len_metric_name = format!("cache_{}_len", self.cache_name);
140            let len_metric_help = format!("Length of cache {}", self.cache_name);
141            let len: Gauge = Default::default();
142            len.set(self.len() as _);
143            let len_metric_encoder = encoder.encode_descriptor(
144                &len_metric_name,
145                &len_metric_help,
146                None,
147                len.metric_type(),
148            )?;
149            len.encode(len_metric_encoder)?;
150        }
151        {
152            let cap_metric_name = format!("cache_{}_cap", self.cache_name);
153            let cap_metric_help = format!("Capacity of cache {}", self.cache_name);
154            let cap: Gauge = Default::default();
155            cap.set(self.capacity() as _);
156            let cap_metric_encoder = encoder.encode_descriptor(
157                &cap_metric_name,
158                &cap_metric_help,
159                None,
160                cap.metric_type(),
161            )?;
162            cap.encode(cap_metric_encoder)?;
163        }
164        {
165            let hits_metric_name = format!("cache_{}_hits", self.cache_name);
166            let hits_metric_help = format!("Cache hits of {}", self.cache_name);
167            let hits: Counter = Default::default();
168            hits.inner().store(self.cache.hits(), Ordering::Relaxed);
169            let hits_metric_encoder = encoder.encode_descriptor(
170                &hits_metric_name,
171                &hits_metric_help,
172                None,
173                hits.metric_type(),
174            )?;
175            hits.encode(hits_metric_encoder)?;
176        }
177        {
178            let misses_metric_name = format!("cache_{}_misses", self.cache_name);
179            let misses_metric_help = format!("Cache misses of {}", self.cache_name);
180            let misses: Counter = Default::default();
181            misses.inner().store(self.cache.misses(), Ordering::Relaxed);
182            let misses_metric_encoder = encoder.encode_descriptor(
183                &misses_metric_name,
184                &misses_metric_help,
185                None,
186                misses.metric_type(),
187            )?;
188            misses.encode(misses_metric_encoder)?;
189        }
190
191        Ok(())
192    }
193}
forest/utils/cache/size_tracking.rs

forest/utils/cache/
size_tracking.rs
