use std::{
borrow::{Borrow, Cow},
fmt::Debug,
hash::Hash,
num::NonZeroUsize,
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
},
};
use get_size2::GetSize;
use hashlink::LruCache;
use parking_lot::RwLock;
use prometheus_client::{
collector::Collector,
encoding::{DescriptorEncoder, EncodeMetric},
metrics::gauge::Gauge,
registry::Unit,
};
use crate::utils::ShallowClone;
pub trait KeyConstraints:
GetSize + Debug + Send + Sync + Hash + PartialEq + Eq + Clone + 'static
{
}
impl<T> KeyConstraints for T where
T: GetSize + Debug + Send + Sync + Hash + PartialEq + Eq + Clone + 'static
{
}
pub trait LruValueConstraints: GetSize + Debug + Send + Sync + Clone + 'static {}
impl<T> LruValueConstraints for T where T: GetSize + Debug + Send + Sync + Clone + 'static {}
#[derive(Debug)]
pub struct SizeTrackingLruCache<K, V>
where
K: KeyConstraints,
V: LruValueConstraints,
{
cache_id: usize,
cache_name: Cow<'static, str>,
cache: Arc<RwLock<LruCache<K, V>>>,
}
impl<K, V> ShallowClone for SizeTrackingLruCache<K, V>
where
K: KeyConstraints,
V: LruValueConstraints,
{
fn shallow_clone(&self) -> Self {
Self {
cache_id: self.cache_id,
cache_name: self.cache_name.clone(),
cache: self.cache.shallow_clone(),
}
}
}
impl<K, V> SizeTrackingLruCache<K, V>
where
K: KeyConstraints,
V: LruValueConstraints,
{
fn register_metrics(&self) {
crate::metrics::register_collector(Box::new(self.shallow_clone()));
}
fn new_inner(cache_name: Cow<'static, str>, capacity: Option<NonZeroUsize>) -> Self {
static ID_GENERATOR: AtomicUsize = AtomicUsize::new(0);
Self {
cache_id: ID_GENERATOR.fetch_add(1, Ordering::Relaxed),
cache_name,
#[allow(clippy::disallowed_methods)]
cache: Arc::new(RwLock::new(
capacity
.map(From::from)
.map(LruCache::new)
.unwrap_or_else(LruCache::new_unbounded),
)),
}
}
pub fn new_without_metrics_registry(
cache_name: Cow<'static, str>,
capacity: NonZeroUsize,
) -> Self {
Self::new_inner(cache_name, Some(capacity))
}
pub fn new_with_metrics(cache_name: Cow<'static, str>, capacity: NonZeroUsize) -> Self {
let c = Self::new_without_metrics_registry(cache_name, capacity);
c.register_metrics();
c
}
pub fn unbounded_without_metrics_registry(cache_name: Cow<'static, str>) -> Self {
Self::new_inner(cache_name, None)
}
pub fn unbounded_with_metrics(cache_name: Cow<'static, str>) -> Self {
let c = Self::unbounded_without_metrics_registry(cache_name);
c.register_metrics();
c
}
pub fn cache(&self) -> &Arc<RwLock<LruCache<K, V>>> {
&self.cache
}
pub fn remove<Q>(&self, k: &Q) -> Option<V>
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
self.cache.write().remove(k)
}
pub fn push(&self, k: K, v: V) -> Option<V> {
self.cache.write().insert(k, v)
}
pub fn get_map<Q, T>(&self, k: &Q, mapper: impl Fn(&V) -> T) -> Option<T>
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
self.cache.write().get(k).map(mapper)
}
pub fn get_cloned<Q>(&self, k: &Q) -> Option<V>
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
self.get_map(k, Clone::clone)
}
pub fn peek_cloned<Q>(&self, k: &Q) -> Option<V>
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
self.cache.read().peek(k).cloned()
}
pub fn pop_lru(&self) -> Option<(K, V)> {
self.cache.write().remove_lru()
}
pub fn len(&self) -> usize {
self.cache.read().len()
}
pub fn cap(&self) -> usize {
self.cache.read().capacity()
}
pub fn clear(&self) {
self.cache.write().clear()
}
pub(crate) fn size_in_bytes(&self) -> usize {
let mut size = 0_usize;
for (k, v) in self.cache.read().iter() {
size = size
.saturating_add(k.get_size())
.saturating_add(v.get_size());
}
size
}
#[cfg(test)]
pub(crate) fn new_mocked() -> Self {
Self::new_inner(Cow::Borrowed("mocked_cache"), NonZeroUsize::new(1))
}
}
impl<K, V> Collector for SizeTrackingLruCache<K, V>
where
K: KeyConstraints,
V: LruValueConstraints,
{
fn encode(&self, mut encoder: DescriptorEncoder) -> Result<(), std::fmt::Error> {
{
let size_in_bytes = {
let g: Gauge = Default::default();
g.set(self.size_in_bytes() as _);
g
};
let size_metric_name = format!("cache_{}_{}_size", self.cache_name, self.cache_id);
let size_metric_help = format!(
"Size of LruCache {}_{} in bytes",
self.cache_name, self.cache_id
);
let size_metric_encoder = encoder.encode_descriptor(
&size_metric_name,
&size_metric_help,
Some(&Unit::Bytes),
size_in_bytes.metric_type(),
)?;
size_in_bytes.encode(size_metric_encoder)?;
}
{
let len_metric_name = format!("{}_{}_len", self.cache_name, self.cache_id);
let len_metric_help =
format!("Length of LruCache {}_{}", self.cache_name, self.cache_id);
let len: Gauge = Default::default();
len.set(self.len() as _);
let len_metric_encoder = encoder.encode_descriptor(
&len_metric_name,
&len_metric_help,
None,
len.metric_type(),
)?;
len.encode(len_metric_encoder)?;
}
{
let cap_metric_name = format!("{}_{}_cap", self.cache_name, self.cache_id);
let cap_metric_help =
format!("Capacity of LruCache {}_{}", self.cache_name, self.cache_id);
let cap: Gauge = Default::default();
cap.set(self.cap() as _);
let cap_metric_encoder = encoder.encode_descriptor(
&cap_metric_name,
&cap_metric_help,
None,
cap.metric_type(),
)?;
cap.encode(cap_metric_encoder)?;
}
Ok(())
}
}