use std::{
collections::HashSet,
sync::{
Arc,
atomic::{AtomicBool, 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,
fence::{InvalidationFence, Reach},
tviews_feed::Collection,
};
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 collected_views: Box<[ViewName]>,
pub collections: Box<[(Collection, String)]>,
pub by_id: Option<(String, String)>,
pub epoch: u64,
}
type IndexRef = (u64, u64);
type CollectionIndex = DashMap<String, DashMap<(String, String), DashSet<IndexRef>>>;
type ViewIndex = DashMap<ViewName, DashSet<IndexRef>>;
type RowIndex = DashMap<String, DashMap<String, DashSet<IndexRef>>>;
#[derive(Clone, Copy)]
enum Sweep {
Full,
Write,
Feed,
}
#[derive(Debug, Clone, Default)]
pub struct ComposedTags {
pub collected_views: Vec<String>,
pub collections: Vec<(Collection, String)>,
pub entities: Vec<(String, String)>,
}
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<IndexRef>>>,
entity_index: Arc<DashMap<String, DashMap<String, DashSet<IndexRef>>>>,
collected_view_index: Arc<DashMap<ViewName, DashSet<IndexRef>>>,
collection_index: Arc<CollectionIndex>,
by_id_view_index: Arc<ViewIndex>,
by_id_index: Arc<RowIndex>,
next_epoch: AtomicU64,
fence: InvalidationFence,
suspended: AtomicBool,
}
#[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
}
struct ReverseIndexes {
views: Arc<ViewIndex>,
entities: Arc<RowIndex>,
collected_views: Arc<ViewIndex>,
collections: Arc<CollectionIndex>,
by_id_views: Arc<ViewIndex>,
by_id: Arc<RowIndex>,
}
fn build_store(
config: &CacheConfig,
memory_bytes: Arc<AtomicUsize>,
indexes: ReverseIndexes,
) -> MokaCache<u64, Arc<CachedResult>> {
let max_cap = config.max_entries as u64;
let mb = memory_bytes;
let vi = indexes.views;
let ei = indexes.entities;
let cvi = indexes.collected_views;
let ci = indexes.collections;
let bvi = indexes.by_id_views;
let bi = indexes.by_id;
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);
let reference: IndexRef = (*key, value.epoch);
let views = if value.by_id.is_some() { &bvi } else { &vi };
for view in &value.accessed_views {
if let Some(keys) = views.get(view) {
keys.remove(&reference);
}
}
if let Some((type_name, id)) = &value.by_id {
if let Some(by_type) = bi.get(type_name) {
if let Some(keys) = by_type.get(id) {
keys.remove(&reference);
}
}
}
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(&reference);
}
}
}
for view in &value.collected_views {
if let Some(keys) = cvi.get(view) {
keys.remove(&reference);
}
}
for (collection, parent) in &*value.collections {
if let Some(by_type) = ci.get(&collection.child_type) {
if let Some(keys) = by_type.get(&(collection.column.clone(), parent.clone())) {
keys.remove(&reference);
}
}
}
})
.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<IndexRef>>> = Arc::new(DashMap::new());
let entity_index: Arc<DashMap<String, DashMap<String, DashSet<IndexRef>>>> =
Arc::new(DashMap::new());
let collected_view_index: Arc<DashMap<ViewName, DashSet<IndexRef>>> =
Arc::new(DashMap::new());
let collection_index: Arc<CollectionIndex> = Arc::new(DashMap::new());
let by_id_view_index: Arc<ViewIndex> = Arc::new(DashMap::new());
let by_id_index: Arc<RowIndex> = Arc::new(DashMap::new());
let store = build_store(
&config,
Arc::clone(&memory_bytes),
ReverseIndexes {
views: Arc::clone(&view_index),
entities: Arc::clone(&entity_index),
collected_views: Arc::clone(&collected_view_index),
collections: Arc::clone(&collection_index),
by_id_views: Arc::clone(&by_id_view_index),
by_id: Arc::clone(&by_id_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,
collected_view_index,
collection_index,
by_id_view_index,
by_id_index,
next_epoch: AtomicU64::new(0),
fence: InvalidationFence::new(),
suspended: AtomicBool::new(false),
}
}
#[must_use]
pub fn invalidation_generation(&self) -> u64 {
self.fence.snapshot()
}
#[must_use]
pub const fn is_enabled(&self) -> bool {
self.config.enabled
}
#[must_use]
pub const fn config(&self) -> &CacheConfig {
&self.config
}
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>,
fence: Option<u64>,
) -> Result<()> {
self.put_arc_with(
cache_key,
result,
accessed_views,
ttl_override,
entity_type,
fence,
|| {},
|| {},
)
}
#[allow(clippy::too_many_arguments)] pub(super) fn put_arc_with(
&self,
cache_key: u64,
result: Arc<Vec<JsonbValue>>,
accessed_views: Vec<String>,
ttl_override: Option<u64>,
entity_type: Option<&str>,
fence: Option<u64>,
before_insert: impl FnOnce(),
after_insert: impl FnOnce(),
) -> Result<()> {
self.put_tagged(
cache_key,
result,
accessed_views,
ttl_override,
entity_type,
ComposedTags::default(),
None,
fence,
before_insert,
after_insert,
)
}
pub fn put_by_id(
&self,
cache_key: u64,
result: Arc<Vec<JsonbValue>>,
accessed_views: Vec<String>,
ttl_override: Option<u64>,
by_id: (String, String),
fence: Option<u64>,
) -> Result<()> {
self.put_by_id_with(
cache_key,
result,
accessed_views,
ttl_override,
by_id,
fence,
|| {},
|| {},
)
}
#[allow(clippy::too_many_arguments)] pub(super) fn put_by_id_with(
&self,
cache_key: u64,
result: Arc<Vec<JsonbValue>>,
accessed_views: Vec<String>,
ttl_override: Option<u64>,
by_id: (String, String),
fence: Option<u64>,
before_insert: impl FnOnce(),
after_insert: impl FnOnce(),
) -> Result<()> {
self.put_tagged(
cache_key,
result,
accessed_views,
ttl_override,
None,
ComposedTags::default(),
Some(by_id),
fence,
before_insert,
after_insert,
)
}
#[allow(clippy::too_many_arguments)] fn put_tagged(
&self,
cache_key: u64,
result: Arc<Vec<JsonbValue>>,
accessed_views: Vec<String>,
ttl_override: Option<u64>,
entity_type: Option<&str>,
tags: ComposedTags,
by_id: Option<(String, String)>,
fence: Option<u64>,
before_insert: impl FnOnce(),
after_insert: impl FnOnce(),
) -> Result<()> {
if !self.config.enabled || self.suspended.load(Ordering::SeqCst) {
return Ok(());
}
let entity_type = entity_type.filter(|_| by_id.is_none());
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()))
})
.chain(tags.entities.iter().cloned())
.collect::<Vec<_>>()
.into_boxed_slice()
} else {
tags.entities.clone().into_boxed_slice()
};
if fence.is_some_and(|snapshot| {
self.fence.moved_since(
snapshot,
&Reach {
views: by_id
.is_none()
.then(|| accessed_views.iter().map(String::as_str))
.into_iter()
.flatten()
.collect(),
collected_views: tags.collected_views.iter().map(String::as_str).collect(),
entities: entity_refs
.iter()
.map(|(t, id)| (t.as_str(), id.as_str()))
.collect(),
collections: tags
.collections
.iter()
.map(|(c, p)| (c.child_type.as_str(), c.column.as_str(), p.as_str()))
.collect(),
by_id: by_id.as_ref().map(|(t, id)| {
(
(t.as_str(), id.as_str()),
accessed_views.iter().map(String::as_str).collect(),
)
}),
},
)
}) {
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 accessed_views: Box<[ViewName]> =
accessed_views.into_iter().map(ViewName::from).collect();
let collected_views: Box<[ViewName]> =
tags.collected_views.into_iter().map(ViewName::from).collect();
let collections: Box<[(Collection, String)]> = tags.collections.into_boxed_slice();
let epoch = self.next_epoch.fetch_add(1, Ordering::Relaxed);
let reference: IndexRef = (cache_key, epoch);
let views = if by_id.is_some() {
&self.by_id_view_index
} else {
&self.view_index
};
for view in &accessed_views {
views.entry(view.clone()).or_default().insert(reference);
}
if let Some((type_name, id)) = &by_id {
self.by_id_index
.entry(type_name.clone())
.or_default()
.entry(id.clone())
.or_default()
.insert(reference);
}
for (et, id) in &*entity_refs {
self.entity_index
.entry(et.clone())
.or_default()
.entry(id.clone())
.or_default()
.insert(reference);
}
for view in &collected_views {
self.collected_view_index.entry(view.clone()).or_default().insert(reference);
}
for (collection, parent) in &*collections {
self.collection_index
.entry(collection.child_type.clone())
.or_default()
.entry((collection.column.clone(), parent.clone()))
.or_default()
.insert(reference);
}
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,
collected_views,
collections,
by_id,
epoch,
};
before_insert();
self.memory_bytes.fetch_add(entry_overhead(), Ordering::Relaxed);
let entry = Arc::new(cached);
self.store.insert(cache_key, Arc::clone(&entry));
after_insert();
if fence.is_some_and(|snapshot| {
let views = || entry.accessed_views.iter().map(ViewName::as_str);
self.fence.moved_since_inserted(
snapshot,
&Reach {
views: if entry.by_id.is_some() {
Vec::new()
} else {
views().collect()
},
collected_views: entry.collected_views.iter().map(ViewName::as_str).collect(),
entities: entry
.entity_refs
.iter()
.map(|(t, id)| (t.as_str(), id.as_str()))
.collect(),
collections: entry
.collections
.iter()
.map(|(c, p)| (c.child_type.as_str(), c.column.as_str(), p.as_str()))
.collect(),
by_id: entry
.by_id
.as_ref()
.map(|(t, id)| ((t.as_str(), id.as_str()), views().collect())),
},
)
}) {
self.store.invalidate(&cache_key);
return Ok(());
}
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>,
fence: Option<u64>,
) -> Result<()> {
self.put_arc(cache_key, Arc::new(result), accessed_views, ttl_override, entity_type, fence)
}
#[allow(clippy::too_many_arguments)] pub fn put_composed(
&self,
cache_key: u64,
result: Arc<Vec<JsonbValue>>,
accessed_views: Vec<String>,
ttl_override: Option<u64>,
entity_type: Option<&str>,
tags: ComposedTags,
fence: Option<u64>,
) -> Result<()> {
self.put_tagged(
cache_key,
result,
accessed_views,
ttl_override,
entity_type,
tags,
None,
fence,
|| {},
|| {},
)
}
pub fn invalidate_by_id(&self, type_name: &str, id: &str) -> Result<u64> {
if !self.config.enabled {
return Ok(0);
}
self.fence.stamp_by_id(type_name, id);
let keys: HashSet<u64> = self
.by_id_index
.get(type_name)
.and_then(|by_type| by_type.get(id).map(|refs| refs.iter().map(|r| r.0).collect()))
.unwrap_or_default();
Ok(self.evict(keys))
}
pub fn invalidate_by_id_type(&self, type_name: &str) -> Result<u64> {
if !self.config.enabled {
return Ok(0);
}
self.fence.stamp_by_id_type(type_name);
let keys: HashSet<u64> = self
.by_id_index
.get(type_name)
.map(|by_type| {
by_type
.iter()
.flat_map(|refs| refs.iter().map(|r| r.0).collect::<Vec<_>>())
.collect()
})
.unwrap_or_default();
Ok(self.evict(keys))
}
pub fn invalidate_collection(&self, collection: &Collection, parent_id: &str) -> Result<u64> {
if !self.config.enabled {
return Ok(0);
}
self.fence
.stamp_collection(&collection.child_type, &collection.column, parent_id);
let keys: HashSet<u64> = self
.collection_index
.get(&collection.child_type)
.and_then(|by_type| {
by_type
.get(&(collection.column.clone(), parent_id.to_string()))
.map(|refs| refs.iter().map(|r| r.0).collect())
})
.unwrap_or_default();
Ok(self.evict(keys))
}
pub fn invalidate_collections_of(&self, child_type: &str) -> Result<u64> {
if !self.config.enabled {
return Ok(0);
}
self.fence.stamp_collection_type(child_type);
let keys: HashSet<u64> = self
.collection_index
.get(child_type)
.map(|by_type| {
by_type
.iter()
.flat_map(|refs| refs.iter().map(|r| r.0).collect::<Vec<_>>())
.collect()
})
.unwrap_or_default();
Ok(self.evict(keys))
}
fn evict(&self, keys: HashSet<u64>) -> u64 {
#[allow(clippy::cast_possible_truncation)]
let count = keys.len() as u64;
for key in keys {
self.store.invalidate(&key);
}
self.invalidations.fetch_add(count, Ordering::Relaxed);
count
}
pub fn invalidate_views_except_collections(&self, views: &[ViewName]) -> Result<u64> {
self.sweep_views(views, Sweep::Feed)
}
pub fn invalidate_views_except_by_id(&self, views: &[ViewName]) -> Result<u64> {
self.sweep_views(views, Sweep::Write)
}
pub fn invalidate_views(&self, views: &[ViewName]) -> Result<u64> {
self.sweep_views(views, Sweep::Full)
}
fn sweep_views(&self, views: &[ViewName], sweep: Sweep) -> Result<u64> {
if !self.config.enabled {
return Ok(0);
}
let names = views.iter().map(ViewName::as_str);
match sweep {
Sweep::Full => self.fence.stamp_views(names),
Sweep::Write => self.fence.stamp_write_views(names),
Sweep::Feed => self.fence.stamp_feed_views(names),
}
let indexes: &[&ViewIndex] = match sweep {
Sweep::Full => &[
&self.view_index,
&self.collected_view_index,
&self.by_id_view_index,
],
Sweep::Write => &[&self.view_index, &self.collected_view_index],
Sweep::Feed => &[&self.view_index],
};
let mut keys_to_invalidate: HashSet<u64> = HashSet::new();
for view in views {
for index in indexes {
if let Some(keys) = index.get(view.as_str()) {
for reference in keys.iter() {
keys_to_invalidate.insert(reference.0);
}
}
}
}
Ok(self.evict(keys_to_invalidate))
}
pub fn invalidate_by_entity(&self, entity_type: &str, entity_id: &str) -> Result<u64> {
if !self.config.enabled {
return Ok(0);
}
let by_id = self.invalidate_by_id(entity_type, entity_id)?;
self.fence.stamp_entity(entity_type, entity_id);
if !self.entity_index.contains_key(entity_type) {
return Ok(by_id);
}
let keys_to_invalidate: HashSet<u64> = self
.entity_index
.get(entity_type)
.and_then(|by_type| {
by_type.get(entity_id).map(|keys| keys.iter().map(|r| r.0).collect())
})
.unwrap_or_default();
Ok(by_id + self.evict(keys_to_invalidate))
}
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.fence.stamp_everything();
self.store.invalidate_all();
self.view_index.clear();
self.entity_index.clear();
self.collected_view_index.clear();
self.collection_index.clear();
self.by_id_view_index.clear();
self.by_id_index.clear();
self.memory_bytes.store(0, Ordering::Relaxed);
Ok(())
}
pub fn suspend(&self) {
self.suspended.store(true, Ordering::SeqCst);
let _ = self.clear();
}
pub fn resume(&self) {
self.suspended.store(false, Ordering::SeqCst);
let _ = self.clear();
}
pub 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
}
}
#[cfg(test)]
#[path = "result_collections_tests.rs"]
mod result_collections_tests;
#[cfg(test)]
#[path = "result_by_id_tests.rs"]
mod result_by_id_tests;