#[cfg(with_metrics)]
use std::any::type_name;
use std::{
borrow::Cow,
hash::Hash,
sync::{Arc, Weak},
};
use linera_base::{crypto::CryptoHash, hashed::Hashed};
use papaya::{Compute, Operation};
use quick_cache::sync::Cache;
pub const DEFAULT_CLEANUP_INTERVAL_SECS: u64 = 30;
pub struct ValueCache<K, V> {
cache: Cache<K, crate::Arc<V>>,
weak_index: Arc<papaya::HashMap<K, Weak<V>>>,
}
impl<K, V> ValueCache<K, V>
where
K: Hash + Eq + Clone + Send + Sync + 'static,
V: Send + Sync + 'static,
{
#[cfg(not(web))]
pub fn new(size: usize, cleanup_interval_secs: u64) -> Self {
let weak_index = Arc::new(papaya::HashMap::new());
Self::spawn_cleanup_task(
Arc::clone(&weak_index),
std::time::Duration::from_secs(cleanup_interval_secs),
);
ValueCache {
cache: Cache::new(size),
weak_index,
}
}
#[cfg(web)]
pub fn new(size: usize, _cleanup_interval_secs: u64) -> Self {
ValueCache {
cache: Cache::new(size),
weak_index: Arc::new(papaya::HashMap::new()),
}
}
pub fn insert(&self, key: &K, value: V) -> crate::Arc<V> {
self.dedup_insert(key, crate::Arc(Arc::new(value)))
}
pub fn remove(&self, key: &K) -> Option<crate::Arc<V>> {
let value = self.cache.peek(key);
if value.is_some() {
self.cache.remove(key);
}
Self::track_cache_usage(value)
}
pub fn get(&self, key: &K) -> Option<crate::Arc<V>> {
if let Some(arc) = self.cache.get(key) {
return Self::track_cache_usage(Some(arc));
}
let guard = self.weak_index.guard();
if let Some(weak) = self.weak_index.get(key, &guard) {
if let Some(arc) = weak.upgrade() {
let arc = crate::Arc(arc);
self.cache.insert(key.clone(), arc.clone());
return Self::track_cache_usage(Some(arc));
}
}
Self::track_cache_usage(None)
}
pub fn contains(&self, key: &K) -> bool {
if self.cache.peek(key).is_some() {
return true;
}
let guard = self.weak_index.guard();
self.weak_index
.get(key, &guard)
.is_some_and(|weak| weak.strong_count() > 0)
}
#[cfg(with_testing)]
pub fn cleanup_dead_entries(&self) {
let guard = self.weak_index.guard();
self.weak_index
.retain(|_, weak| weak.strong_count() > 0, &guard);
}
#[cfg(not(web))]
fn spawn_cleanup_task(
weak_index: Arc<papaya::HashMap<K, Weak<V>>>,
cleanup_interval: std::time::Duration,
) {
if tokio::runtime::Handle::try_current().is_err() {
return;
}
tokio::spawn(async move {
let mut interval = tokio::time::interval(cleanup_interval);
loop {
interval.tick().await;
let guard = weak_index.guard();
weak_index.retain(|_, weak| weak.strong_count() > 0, &guard);
}
});
}
fn dedup_insert(&self, key: &K, new_arc: crate::Arc<V>) -> crate::Arc<V> {
let guard = self.weak_index.guard();
let weak = Arc::downgrade(&new_arc.0);
let result = self.weak_index.compute(
key.clone(),
|entry| match entry {
Some((_k, existing_weak)) => match existing_weak.upgrade() {
Some(existing_arc) => Operation::Abort(existing_arc),
None => Operation::Insert(weak.clone()),
},
None => Operation::Insert(weak.clone()),
},
&guard,
);
let canonical_arc = match result {
Compute::Inserted(..) | Compute::Updated { .. } => new_arc,
Compute::Aborted(existing_arc) => crate::Arc(existing_arc),
_ => unreachable!(),
};
self.cache.insert(key.clone(), canonical_arc.clone());
canonical_arc
}
fn track_cache_usage(maybe_value: Option<crate::Arc<V>>) -> Option<crate::Arc<V>> {
#[cfg(with_metrics)]
{
let metric = if maybe_value.is_some() {
&metrics::CACHE_HIT_COUNT
} else {
&metrics::CACHE_MISS_COUNT
};
metric
.with_label_values(&[type_name::<K>(), type_name::<V>()])
.inc();
}
maybe_value
}
}
impl<V: Clone + Send + Sync + 'static> ValueCache<CryptoHash, V> {
pub fn insert_hashed<T>(&self, value: Cow<Hashed<T>>) -> crate::Arc<V>
where
T: Clone,
V: From<Hashed<T>>,
{
let hash = (*value).hash();
if let Some(arc) = self.cache.peek(&hash) {
return arc;
}
let guard = self.weak_index.guard();
if let Some(weak) = self.weak_index.get(&hash, &guard) {
if let Some(arc) = weak.upgrade() {
let arc = crate::Arc(arc);
self.cache.insert(hash, arc.clone());
return arc;
}
}
drop(guard);
self.dedup_insert(&hash, crate::Arc(Arc::new(value.into_owned().into())))
}
#[cfg(with_testing)]
pub fn insert_all_hashed<'a, T>(&self, values: impl IntoIterator<Item = Cow<'a, Hashed<T>>>)
where
T: Clone + 'a,
V: From<Hashed<T>>,
{
for value in values {
self.insert_hashed(value);
}
}
}
#[cfg(with_metrics)]
mod metrics {
use std::sync::LazyLock;
use linera_base::prometheus_util::register_int_counter_vec;
use prometheus::IntCounterVec;
pub static CACHE_HIT_COUNT: LazyLock<IntCounterVec> = LazyLock::new(|| {
register_int_counter_vec(
"value_cache_hit",
"Cache hits in `ValueCache`",
&["key_type", "value_type"],
)
});
pub static CACHE_MISS_COUNT: LazyLock<IntCounterVec> = LazyLock::new(|| {
register_int_counter_vec(
"value_cache_miss",
"Cache misses in `ValueCache`",
&["key_type", "value_type"],
)
});
}
#[cfg(test)]
mod tests {
use std::borrow::Cow;
use linera_base::{crypto::CryptoHash, hashed::Hashed};
use serde::{Deserialize, Serialize};
use super::{ValueCache, DEFAULT_CLEANUP_INTERVAL_SECS};
use crate::Arc as CacheArc;
const TEST_CACHE_SIZE: usize = 10;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
struct TestValue(u64);
impl linera_base::crypto::BcsHashable<'_> for TestValue {}
impl From<Hashed<TestValue>> for TestValue {
fn from(value: Hashed<TestValue>) -> Self {
value.into_inner()
}
}
fn create_test_value(n: u64) -> Hashed<TestValue> {
Hashed::new(TestValue(n))
}
fn create_test_values(iter: impl IntoIterator<Item = u64>) -> Vec<Hashed<TestValue>> {
iter.into_iter().map(create_test_value).collect()
}
fn new_hashed_cache(size: usize) -> ValueCache<CryptoHash, TestValue> {
ValueCache::new(size, DEFAULT_CLEANUP_INTERVAL_SECS)
}
fn new_string_cache(size: usize) -> ValueCache<u64, String> {
ValueCache::new(size, DEFAULT_CLEANUP_INTERVAL_SECS)
}
#[test]
fn test_retrieve_missing_value() {
let cache = new_hashed_cache(TEST_CACHE_SIZE);
let hash = CryptoHash::test_hash("Missing value");
assert!(cache.get(&hash).is_none());
assert!(!cache.contains(&hash));
}
#[test]
fn test_insert_and_get() {
let cache = new_hashed_cache(TEST_CACHE_SIZE);
let value = create_test_value(0);
let hash = value.hash();
cache.insert_hashed(Cow::Borrowed(&value));
assert!(cache.contains(&hash));
assert_eq!(cache.get(&hash).as_deref(), Some(value.inner()));
}
#[test]
fn test_insert_many_values() {
let cache = new_hashed_cache(TEST_CACHE_SIZE);
let values = create_test_values(0..TEST_CACHE_SIZE as u64);
for value in &values {
cache.insert_hashed(Cow::Borrowed(value));
}
for value in &values {
assert!(cache.contains(&value.hash()));
assert_eq!(cache.get(&value.hash()).as_deref(), Some(value.inner()));
}
let cache2 = new_hashed_cache(TEST_CACHE_SIZE);
cache2.insert_all_hashed(values.iter().map(Cow::Borrowed));
for value in &values {
assert_eq!(cache2.get(&value.hash()).as_deref(), Some(value.inner()));
}
}
#[test]
fn test_reinsertion_dedup() {
let cache = new_hashed_cache(TEST_CACHE_SIZE);
let values = create_test_values(0..TEST_CACHE_SIZE as u64);
let first_arcs: Vec<_> = values
.iter()
.map(|v| cache.insert_hashed(Cow::Borrowed(v)))
.collect();
for (value, first_arc) in values.iter().zip(&first_arcs) {
let second_arc = cache.insert_hashed(Cow::Borrowed(value));
assert!(CacheArc::ptr_eq(&second_arc, first_arc));
}
}
#[test]
fn test_eviction() {
let cache = new_hashed_cache(TEST_CACHE_SIZE);
let total = TEST_CACHE_SIZE * 3;
let values = create_test_values(0..total as u64);
for value in &values {
cache.insert_hashed(Cow::Borrowed(value));
}
let present_count = values.iter().filter(|v| cache.contains(&v.hash())).count();
assert!(
present_count <= TEST_CACHE_SIZE + 1,
"cache should not hold significantly more than its capacity, \
but has {present_count} entries for capacity {TEST_CACHE_SIZE}"
);
assert!(present_count > 0, "cache should still hold some entries");
}
#[test]
fn test_accessed_entry_survives_eviction() {
let cache = new_hashed_cache(TEST_CACHE_SIZE);
let promoted = create_test_value(0);
let promoted_hash = promoted.hash();
cache.insert_hashed(Cow::Borrowed(&promoted));
cache.get(&promoted_hash);
let extras = create_test_values(1..=TEST_CACHE_SIZE as u64 * 2);
for value in &extras {
cache.insert_hashed(Cow::Borrowed(value));
}
assert!(
cache.contains(&promoted_hash),
"recently accessed entry should survive eviction"
);
}
#[test]
fn test_promotion_of_reinsertion() {
let cache = new_hashed_cache(TEST_CACHE_SIZE);
let promoted = create_test_value(0);
let promoted_hash = promoted.hash();
let first = cache.insert_hashed(Cow::Borrowed(&promoted));
let second = cache.insert_hashed(Cow::Borrowed(&promoted));
assert!(CacheArc::ptr_eq(&first, &second));
let extras = create_test_values(1..=TEST_CACHE_SIZE as u64 * 2);
for value in &extras {
cache.insert_hashed(Cow::Borrowed(value));
}
assert!(
cache.contains(&promoted_hash),
"re-inserted entry should survive eviction"
);
}
#[test]
fn test_weak_index_dedup_after_eviction() {
let cache = new_string_cache(2);
let held = cache.insert(&1, "hello".to_string());
cache.insert(&2, "world".to_string());
cache.insert(&3, "foo".to_string());
cache.insert(&4, "bar".to_string());
let retrieved = cache
.get(&1)
.expect("held Arc should keep entry findable via weak index");
assert!(
CacheArc::ptr_eq(&retrieved, &held),
"must return same allocation, not a duplicate"
);
let reinserted = cache.insert(&1, "replacement".to_string());
assert!(CacheArc::ptr_eq(&reinserted, &held));
assert_eq!(&*reinserted, "hello");
}
#[test]
fn test_remove_preserves_weak_for_held_arcs() {
let cache = new_string_cache(TEST_CACHE_SIZE);
let held = cache.insert(&1, "hello".to_string());
cache.remove(&1);
let retrieved = cache.get(&1).expect("weak index should find held Arc");
assert!(CacheArc::ptr_eq(&retrieved, &held));
}
#[test]
fn test_remove_without_holder() {
let cache = new_string_cache(TEST_CACHE_SIZE);
cache.insert(&1, "hello".to_string());
cache.remove(&1);
assert!(!cache.contains(&1));
assert!(cache.get(&1).is_none());
}
#[test]
fn test_cleanup_dead_entries() {
let cache = new_string_cache(2);
cache.insert(&1, "alive".to_string());
let _held = cache.get(&1).expect("just inserted");
cache.insert(&2, "dead".to_string());
cache.insert(&3, "a".to_string());
cache.insert(&4, "b".to_string());
cache.insert(&5, "c".to_string());
cache.cleanup_dead_entries();
assert!(cache.contains(&1));
}
}