use super::config::CacheConfig;
use super::policy::{
CacheAdmission, CachePolicy, CachePolicyConfig, CachePolicyKind, CachePolicyMetrics,
build_cache_policy,
};
use crate::cache::{CacheKey, GraphNodeSummary};
use dashmap::DashMap;
use lru::LruCache;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
#[derive(Debug, Clone)]
struct CacheEntry {
summaries: Arc<[GraphNodeSummary]>,
size_bytes: usize,
}
impl CacheEntry {
fn new(summaries: Vec<GraphNodeSummary>) -> Self {
let size_bytes = if summaries.is_empty() {
0
} else {
let sample_size = postcard::to_allocvec(&summaries[0])
.map(|bytes| bytes.len())
.unwrap_or(256); sample_size * summaries.len()
};
Self {
summaries: Arc::from(summaries.into_boxed_slice()),
size_bytes,
}
}
}
pub struct CacheStorage {
entries: DashMap<CacheKey, CacheEntry>,
lru: Mutex<LruCache<CacheKey, ()>>,
max_bytes: u64,
total_bytes: AtomicU64,
hits: AtomicUsize,
misses: AtomicUsize,
evictions: AtomicUsize,
policy: Arc<dyn CachePolicy<CacheKey>>,
}
impl CacheStorage {
#[must_use]
pub fn new(max_bytes: u64) -> Self {
Self::with_policy(&CachePolicyConfig::new(
CachePolicyKind::Lru,
max_bytes,
CacheConfig::DEFAULT_POLICY_WINDOW_RATIO,
))
}
#[must_use]
pub fn with_policy(config: &CachePolicyConfig) -> Self {
let policy = build_cache_policy::<CacheKey>(config);
Self {
entries: DashMap::new(),
lru: Mutex::new(LruCache::unbounded()),
max_bytes: config.max_bytes,
total_bytes: AtomicU64::new(0),
hits: AtomicUsize::new(0),
misses: AtomicUsize::new(0),
evictions: AtomicUsize::new(0),
policy,
}
}
fn handle_policy_evictions(&self) {
for eviction in self.policy.drain_evictions() {
let key = eviction.key;
if let Some((_, removed)) = self.entries.remove(&key) {
self.total_bytes
.fetch_sub(removed.size_bytes as u64, Ordering::Relaxed);
self.evictions.fetch_add(1, Ordering::Relaxed);
}
if let Ok(mut lru) = self.lru.lock() {
let _ = lru.pop(&key);
}
}
}
pub fn get(&self, key: &CacheKey) -> Option<Arc<[GraphNodeSummary]>> {
if let Some(entry) = self.entries.get(key) {
let _ = self.policy.record_hit(key);
if let Ok(mut lru) = self.lru.lock() {
lru.get_or_insert(key.clone(), || ());
}
self.hits.fetch_add(1, Ordering::Relaxed);
Some(Arc::clone(&entry.summaries))
} else {
self.misses.fetch_add(1, Ordering::Relaxed);
None
}
}
pub fn insert(&self, key: CacheKey, summaries: Vec<GraphNodeSummary>) {
let entry = CacheEntry::new(summaries);
let entry_size = entry.size_bytes as u64;
let key_for_lru = key.clone();
if let Some((_, old_entry)) = self.entries.remove(&key) {
self.total_bytes
.fetch_sub(old_entry.size_bytes as u64, Ordering::Relaxed);
self.policy.invalidate(&key);
}
if matches!(
self.policy.admit(&key, entry_size),
CacheAdmission::Rejected
) {
log::debug!(
"cache policy {:?} rejected entry {:?} ({} bytes)",
self.policy.kind(),
&key,
entry_size
);
return;
}
self.entries.insert(key, entry);
self.total_bytes.fetch_add(entry_size, Ordering::Relaxed);
if let Ok(mut lru) = self.lru.lock() {
lru.put(key_for_lru, ());
}
self.handle_policy_evictions();
if self.total_bytes.load(Ordering::Relaxed) > self.max_bytes {
self.evict_lru();
}
}
fn evict_lru(&self) {
let Ok(mut lru) = self.lru.lock() else {
log::warn!("Failed to acquire LRU lock for eviction");
return;
};
let mut current_size = self.total_bytes.load(Ordering::Relaxed);
while current_size > self.max_bytes {
let Some((key, ())) = lru.pop_lru() else {
break; };
if let Some((_, removed)) = self.entries.remove(&key) {
current_size = current_size.saturating_sub(removed.size_bytes as u64);
self.total_bytes
.fetch_sub(removed.size_bytes as u64, Ordering::Relaxed);
self.evictions.fetch_add(1, Ordering::Relaxed);
self.policy.invalidate(&key);
log::debug!(
"Evicted cache entry: {} ({} bytes)",
key,
removed.size_bytes
);
}
}
}
pub fn clear(&self) {
self.entries.clear();
self.total_bytes.store(0, Ordering::Relaxed);
self.hits.store(0, Ordering::Relaxed);
self.misses.store(0, Ordering::Relaxed);
self.evictions.store(0, Ordering::Relaxed);
self.policy.reset();
if let Ok(mut lru) = self.lru.lock() {
lru.clear();
}
log::debug!("Cache cleared");
}
pub fn stats(&self) -> CacheStats {
CacheStats {
entry_count: self.entries.len(),
total_bytes: self.total_bytes.load(Ordering::Relaxed),
max_bytes: self.max_bytes,
hits: self.hits.load(Ordering::Relaxed),
misses: self.misses.load(Ordering::Relaxed),
evictions: self.evictions.load(Ordering::Relaxed),
policy: self.policy.stats(),
}
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct CacheStats {
pub entry_count: usize,
pub total_bytes: u64,
pub max_bytes: u64,
pub hits: usize,
pub misses: usize,
pub evictions: usize,
pub policy: CachePolicyMetrics,
}
impl CacheStats {
fn usize_to_f64(value: usize) -> f64 {
#[allow(clippy::cast_precision_loss)]
{
value as f64
}
}
fn u64_to_f64(value: u64) -> f64 {
#[allow(clippy::cast_precision_loss)]
{
value as f64
}
}
#[must_use]
pub fn hit_rate(&self) -> f64 {
let total = self.hits + self.misses;
if total == 0 {
0.0
} else {
Self::usize_to_f64(self.hits) / Self::usize_to_f64(total)
}
}
#[must_use]
pub fn utilization(&self) -> f64 {
if self.max_bytes == 0 {
0.0
} else {
Self::u64_to_f64(self.total_bytes) / Self::u64_to_f64(self.max_bytes)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::cache::policy::{CachePolicyConfig, CachePolicyKind};
use crate::graph::unified::node::NodeKind;
use crate::hash::Blake3Hash;
use approx::assert_abs_diff_eq;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::thread;
fn make_test_key(name: &str, lang: &str) -> CacheKey {
let hash = Blake3Hash::from_bytes([name.as_bytes()[0]; 32]);
CacheKey::from_raw_path(PathBuf::from(name), lang, hash)
}
fn make_test_summary(name: &str) -> GraphNodeSummary {
GraphNodeSummary::new(
Arc::from(name),
NodeKind::Function,
Arc::from(Path::new("test.rs")),
1,
0,
1,
10,
)
}
#[test]
fn test_storage_new() {
let storage = CacheStorage::new(1024);
let stats = storage.stats();
assert_eq!(stats.entry_count, 0);
assert_eq!(stats.total_bytes, 0);
assert_eq!(stats.max_bytes, 1024);
assert_eq!(stats.hits, 0);
assert_eq!(stats.misses, 0);
}
#[test]
fn test_storage_insert_and_get() {
let storage = CacheStorage::new(10 * 1024);
let key = make_test_key("file.rs", "rust");
let summaries = vec![make_test_summary("test_fn")];
storage.insert(key.clone(), summaries.clone());
let retrieved = storage.get(&key).unwrap();
assert_eq!(retrieved.len(), 1);
assert_eq!(retrieved[0].name.as_ref(), "test_fn");
let stats = storage.stats();
assert_eq!(stats.hits, 1);
assert_eq!(stats.misses, 0);
assert_eq!(stats.entry_count, 1);
}
#[test]
fn test_storage_miss() {
let storage = CacheStorage::new(10 * 1024);
let key = make_test_key("file.rs", "rust");
assert!(storage.get(&key).is_none());
let stats = storage.stats();
assert_eq!(stats.hits, 0);
assert_eq!(stats.misses, 1);
}
#[test]
fn test_storage_update() {
let storage = CacheStorage::new(10 * 1024);
let key = make_test_key("file.rs", "rust");
let summaries1 = vec![make_test_summary("fn1")];
let summaries2 = vec![make_test_summary("fn2"), make_test_summary("fn3")];
storage.insert(key.clone(), summaries1);
storage.insert(key.clone(), summaries2.clone());
let retrieved = storage.get(&key).unwrap();
assert_eq!(retrieved.len(), 2);
assert_eq!(retrieved[0].name.as_ref(), "fn2");
}
#[test]
fn test_storage_clear() {
let storage = CacheStorage::new(10 * 1024);
let key = make_test_key("file.rs", "rust");
storage.insert(key.clone(), vec![make_test_summary("test")]);
assert!(storage.get(&key).is_some());
storage.clear();
assert!(storage.get(&key).is_none());
let stats = storage.stats();
assert_eq!(stats.entry_count, 0);
assert_eq!(stats.total_bytes, 0);
}
#[test]
fn test_storage_eviction() {
let storage = CacheStorage::new(100);
for i in 0..10 {
let key = make_test_key(&format!("file{i}.rs"), "rust");
let summaries = vec![make_test_summary(&format!("fn{i}"))];
storage.insert(key, summaries);
}
let stats = storage.stats();
assert!(stats.evictions > 0, "Expected evictions, got 0");
assert!(
stats.entry_count < 10,
"Expected < 10 entries due to eviction"
);
assert!(
stats.total_bytes <= 100,
"Cache size should be under cap, got {}",
stats.total_bytes
);
}
#[test]
fn test_storage_lru_order() {
let storage = CacheStorage::new(80);
let key1 = make_test_key("file1.rs", "rust");
let key2 = make_test_key("file2.rs", "rust");
let key3 = make_test_key("file3.rs", "rust");
storage.insert(key1.clone(), vec![make_test_summary("fn1")]);
storage.insert(key2.clone(), vec![make_test_summary("fn2")]);
storage.insert(key3.clone(), vec![make_test_summary("fn3")]);
storage.get(&key1);
for i in 4..20 {
let key = make_test_key(&format!("file{i}.rs"), "rust");
storage.insert(key, vec![make_test_summary(&format!("fn{i}"))]);
}
let stats = storage.stats();
assert!(
stats.evictions > 0,
"Expected evictions with small cache, got 0"
);
let key1_present = storage.get(&key1).is_some();
let key2_present = storage.get(&key2).is_some();
assert!(
key1_present || !key2_present,
"LRU violation: older key2 survived but recently accessed key1 didn't"
);
}
#[test]
fn test_concurrent_insert_and_get() {
let storage = Arc::new(CacheStorage::new(10 * 1024));
let mut handles = vec![];
for i in 0..10 {
let storage = Arc::clone(&storage);
let handle = thread::spawn(move || {
let key = make_test_key(&format!("file{i}.rs"), "rust");
let summaries = vec![make_test_summary(&format!("fn{i}"))];
storage.insert(key.clone(), summaries.clone());
let retrieved = storage.get(&key).expect("Should retrieve inserted value");
assert_eq!(retrieved.len(), 1);
assert_eq!(retrieved[0].name.as_ref(), &format!("fn{i}"));
});
handles.push(handle);
}
for handle in handles {
handle.join().unwrap();
}
let stats = storage.stats();
assert_eq!(stats.entry_count, 10);
assert_eq!(stats.hits, 10); }
#[test]
fn test_concurrent_eviction() {
let storage = Arc::new(CacheStorage::new(100));
let mut handles = vec![];
for i in 0..20 {
let storage = Arc::clone(&storage);
let handle = thread::spawn(move || {
let key = make_test_key(&format!("file{i}.rs"), "rust");
let summaries = vec![make_test_summary(&format!("fn{i}"))];
storage.insert(key, summaries);
});
handles.push(handle);
}
for handle in handles {
handle.join().unwrap();
}
let stats = storage.stats();
assert!(
stats.evictions > 0,
"Expected evictions with small cache and concurrent inserts"
);
assert!(
stats.total_bytes <= 100,
"Cache size should be under cap, got {}",
stats.total_bytes
);
assert!(stats.entry_count > 0, "Should have entries after eviction");
assert!(
stats.entry_count < 20,
"Should have evicted some of 20 entries"
);
}
#[test]
fn test_cache_stats_hit_rate() {
let stats = CacheStats {
entry_count: 10,
total_bytes: 1000,
max_bytes: 2000,
hits: 75,
misses: 25,
evictions: 0,
policy: CachePolicyMetrics::default(),
};
assert_abs_diff_eq!(stats.hit_rate(), 0.75, epsilon = 1e-10);
}
#[test]
fn test_cache_stats_utilization() {
let stats = CacheStats {
entry_count: 10,
total_bytes: 1000,
max_bytes: 2000,
hits: 0,
misses: 0,
evictions: 0,
policy: CachePolicyMetrics::default(),
};
assert_abs_diff_eq!(stats.utilization(), 0.5, epsilon = 1e-10);
}
#[test]
fn test_cache_stats_empty() {
let stats = CacheStats {
entry_count: 0,
total_bytes: 0,
max_bytes: 1000,
hits: 0,
misses: 0,
evictions: 0,
policy: CachePolicyMetrics::default(),
};
assert_abs_diff_eq!(stats.hit_rate(), 0.0, epsilon = 1e-10);
assert_abs_diff_eq!(stats.utilization(), 0.0, epsilon = 1e-10);
}
#[test]
fn test_tiny_lfu_rejects_cold_workload() {
let storage =
CacheStorage::with_policy(&CachePolicyConfig::new(CachePolicyKind::TinyLfu, 1024, 0.2));
let hot_key = make_test_key("hot.rs", "rust");
storage.insert(hot_key.clone(), vec![make_test_summary("hot_fn")]);
for _ in 0..8 {
assert!(storage.get(&hot_key).is_some());
}
for i in 0..50 {
let key = make_test_key(&format!("cold{i}.rs"), "rust");
storage.insert(key.clone(), vec![make_test_summary(&format!("cold{i}"))]);
let _ = storage.get(&key);
}
let stats = storage.stats();
assert!(
stats.policy.lfu_rejects > 0,
"expected TinyLFU policy to reject some cold inserts"
);
}
}