use std::{
collections::HashSet,
sync::{
Arc,
atomic::{AtomicU64, AtomicUsize, Ordering},
},
time::Duration,
};
use dashmap::{DashMap, DashSet};
use fraiseql_db::ViewName;
use moka::sync::Cache as MokaCache;
use serde::{Deserialize, Serialize};
use super::config::CacheConfig;
use crate::{db::types::JsonbValue, error::Result};
#[derive(Debug, Clone)]
pub struct CachedResult {
pub result: Arc<Vec<JsonbValue>>,
pub accessed_views: Box<[ViewName]>,
pub cached_at: u64,
pub ttl_seconds: u64,
pub entity_refs: Box<[(String, String)]>,
pub is_list_query: bool,
}
struct CacheEntryExpiry;
impl moka::Expiry<u64, Arc<CachedResult>> for CacheEntryExpiry {
fn expire_after_create(
&self,
_key: &u64,
value: &Arc<CachedResult>,
_created_at: std::time::Instant,
) -> Option<Duration> {
if value.ttl_seconds == 0 {
None
} else {
Some(Duration::from_secs(value.ttl_seconds))
}
}
}
pub struct QueryResultCache {
store: MokaCache<u64, Arc<CachedResult>>,
config: CacheConfig,
hits: AtomicU64,
misses: AtomicU64,
total_cached: AtomicU64,
invalidations: AtomicU64,
memory_bytes: Arc<AtomicUsize>,
view_index: Arc<DashMap<ViewName, DashSet<u64>>>,
entity_index: Arc<DashMap<String, DashMap<String, DashSet<u64>>>>,
list_index: Arc<DashMap<ViewName, DashSet<u64>>>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CacheMetrics {
pub hits: u64,
pub misses: u64,
pub total_cached: u64,
pub invalidations: u64,
pub size: usize,
pub memory_bytes: usize,
}
const fn entry_overhead() -> usize {
std::mem::size_of::<CachedResult>() + std::mem::size_of::<u64>() * 2
}
fn build_store(
config: &CacheConfig,
memory_bytes: Arc<AtomicUsize>,
view_index: Arc<DashMap<ViewName, DashSet<u64>>>,
entity_index: Arc<DashMap<String, DashMap<String, DashSet<u64>>>>,
list_index: Arc<DashMap<ViewName, DashSet<u64>>>,
) -> MokaCache<u64, Arc<CachedResult>> {
let max_cap = config.max_entries as u64;
let mb = memory_bytes;
let vi = view_index;
let ei = entity_index;
let li = list_index;
MokaCache::builder()
.max_capacity(max_cap)
.expire_after(CacheEntryExpiry)
.eviction_listener(move |key: Arc<u64>, value: Arc<CachedResult>, _cause| {
mb.fetch_sub(entry_overhead(), Ordering::Relaxed);
for view in &value.accessed_views {
if let Some(keys) = vi.get(view) {
keys.remove(&*key);
}
}
if value.is_list_query {
for view in &value.accessed_views {
if let Some(keys) = li.get(view) {
keys.remove(&*key);
}
}
}
for (et, id) in &*value.entity_refs {
if let Some(by_type) = ei.get(et) {
if let Some(keys) = by_type.get(id) {
keys.remove(&*key);
}
}
}
})
.build()
}
impl QueryResultCache {
#[must_use]
pub fn new(config: CacheConfig) -> Self {
assert!(config.max_entries > 0, "max_entries must be > 0");
let memory_bytes = Arc::new(AtomicUsize::new(0));
let view_index: Arc<DashMap<ViewName, DashSet<u64>>> = Arc::new(DashMap::new());
let entity_index: Arc<DashMap<String, DashMap<String, DashSet<u64>>>> =
Arc::new(DashMap::new());
let list_index: Arc<DashMap<ViewName, DashSet<u64>>> = Arc::new(DashMap::new());
let store = build_store(
&config,
Arc::clone(&memory_bytes),
Arc::clone(&view_index),
Arc::clone(&entity_index),
Arc::clone(&list_index),
);
Self {
store,
config,
hits: AtomicU64::new(0),
misses: AtomicU64::new(0),
total_cached: AtomicU64::new(0),
invalidations: AtomicU64::new(0),
memory_bytes,
view_index,
entity_index,
list_index,
}
}
#[must_use]
pub const fn is_enabled(&self) -> bool {
self.config.enabled
}
pub fn get(&self, cache_key: u64) -> Result<Option<Arc<Vec<JsonbValue>>>> {
if !self.config.enabled {
return Ok(None);
}
if let Some(cached) = self.store.get(&cache_key) {
self.hits.fetch_add(1, Ordering::Relaxed);
Ok(Some(Arc::clone(&cached.result)))
} else {
self.misses.fetch_add(1, Ordering::Relaxed);
Ok(None)
}
}
pub fn put_arc(
&self,
cache_key: u64,
result: Arc<Vec<JsonbValue>>,
accessed_views: Vec<String>,
ttl_override: Option<u64>,
entity_type: Option<&str>,
) -> Result<()> {
if !self.config.enabled {
return Ok(());
}
let ttl_seconds = ttl_override.unwrap_or(self.config.ttl_seconds);
if !self.config.cache_list_queries && result.len() > 1 {
return Ok(());
}
if let Some(max_entry) = self.config.max_entry_bytes {
let estimated = serde_json::to_vec(&*result).map_or(0, |v| v.len());
if estimated > max_entry {
return Ok(()); }
}
if let Some(max_total) = self.config.max_total_bytes {
if self.memory_bytes.load(Ordering::Relaxed) >= max_total {
return Ok(()); }
}
let is_list_query = result.len() > 1;
let entity_refs: Box<[(String, String)]> = if let Some(et) = entity_type {
result
.iter()
.filter_map(|row| {
row.as_value()
.as_object()?
.get("id")?
.as_str()
.map(|id| (et.to_string(), id.to_string()))
})
.collect::<Vec<_>>()
.into_boxed_slice()
} else {
Box::default()
};
let accessed_views: Box<[ViewName]> =
accessed_views.into_iter().map(ViewName::from).collect();
for view in &accessed_views {
self.view_index.entry(view.clone()).or_default().insert(cache_key);
}
if is_list_query {
for view in &accessed_views {
self.list_index.entry(view.clone()).or_default().insert(cache_key);
}
}
for (et, id) in &*entity_refs {
self.entity_index
.entry(et.clone())
.or_default()
.entry(id.clone())
.or_default()
.insert(cache_key);
}
let cached = CachedResult {
result,
accessed_views,
cached_at: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_or(0, |d| d.as_secs()),
ttl_seconds,
entity_refs,
is_list_query,
};
self.memory_bytes.fetch_add(entry_overhead(), Ordering::Relaxed);
self.store.insert(cache_key, Arc::new(cached));
self.total_cached.fetch_add(1, Ordering::Relaxed);
Ok(())
}
pub fn put(
&self,
cache_key: u64,
result: Vec<JsonbValue>,
accessed_views: Vec<String>,
ttl_override: Option<u64>,
entity_type: Option<&str>,
) -> Result<()> {
self.put_arc(cache_key, Arc::new(result), accessed_views, ttl_override, entity_type)
}
pub fn invalidate_views(&self, views: &[ViewName]) -> Result<u64> {
if !self.config.enabled {
return Ok(0);
}
let mut keys_to_invalidate: HashSet<u64> = HashSet::new();
for view in views {
if let Some(keys) = self.view_index.get(view.as_str()) {
for key in keys.iter() {
keys_to_invalidate.insert(*key);
}
}
}
#[allow(clippy::cast_possible_truncation)]
let count = keys_to_invalidate.len() as u64;
for key in keys_to_invalidate {
self.store.invalidate(&key);
}
self.invalidations.fetch_add(count, Ordering::Relaxed);
Ok(count)
}
pub fn invalidate_list_queries(&self, views: &[ViewName]) -> Result<u64> {
if !self.config.enabled {
return Ok(0);
}
let mut keys_to_invalidate: HashSet<u64> = HashSet::new();
for view in views {
if let Some(keys) = self.list_index.get(view.as_str()) {
for k in keys.iter() {
keys_to_invalidate.insert(*k);
}
}
}
#[allow(clippy::cast_possible_truncation)]
let count = keys_to_invalidate.len() as u64;
for key in keys_to_invalidate {
self.store.invalidate(&key);
}
self.invalidations.fetch_add(count, Ordering::Relaxed);
Ok(count)
}
pub fn invalidate_by_entity(&self, entity_type: &str, entity_id: &str) -> Result<u64> {
if !self.config.enabled {
return Ok(0);
}
if !self.entity_index.contains_key(entity_type) {
return Ok(0);
}
let keys_to_invalidate: Vec<u64> = self
.entity_index
.get(entity_type)
.and_then(|by_type| {
by_type.get(entity_id).map(|keys| keys.iter().map(|k| *k).collect())
})
.unwrap_or_default();
#[allow(clippy::cast_possible_truncation)]
let count = keys_to_invalidate.len() as u64;
for key in keys_to_invalidate {
self.store.invalidate(&key);
}
self.invalidations.fetch_add(count, Ordering::Relaxed);
Ok(count)
}
pub fn metrics(&self) -> Result<CacheMetrics> {
Ok(CacheMetrics {
hits: self.hits.load(Ordering::Relaxed),
misses: self.misses.load(Ordering::Relaxed),
total_cached: self.total_cached.load(Ordering::Relaxed),
invalidations: self.invalidations.load(Ordering::Relaxed),
#[allow(clippy::cast_possible_truncation)]
size: self.store.entry_count() as usize,
memory_bytes: self.memory_bytes.load(Ordering::Relaxed),
})
}
pub fn clear(&self) -> Result<()> {
self.store.invalidate_all();
self.view_index.clear();
self.entity_index.clear();
self.list_index.clear();
self.memory_bytes.store(0, Ordering::Relaxed);
Ok(())
}
#[cfg(test)]
pub(crate) fn run_pending_tasks(&self) {
self.store.run_pending_tasks();
}
}
impl CacheMetrics {
#[must_use]
pub fn hit_rate(&self) -> f64 {
let total = self.hits + self.misses;
if total == 0 {
return 0.0;
}
#[allow(clippy::cast_precision_loss)]
{
self.hits as f64 / total as f64
}
}
#[must_use]
pub fn is_healthy(&self) -> bool {
self.hit_rate() > 0.6
}
}