use std::{
sync::{
Arc,
atomic::{AtomicU64, Ordering},
},
time::Duration,
};
use dashmap::{DashMap, DashSet};
use fraiseql_db::ViewName;
use moka::sync::Cache as MokaCache;
use serde_json::Value;
use crate::{error::Result, security::SecurityContext};
#[derive(Debug, Clone, Copy)]
pub struct ResponseCacheConfig {
pub enabled: bool,
pub max_entries: usize,
pub ttl_seconds: u64,
}
impl Default for ResponseCacheConfig {
fn default() -> Self {
Self {
enabled: false, max_entries: 10_000,
ttl_seconds: 300,
}
}
}
struct ResponseEntry {
response: Arc<Value>,
accessed_views: Box<[ViewName]>,
}
pub struct ResponseCache {
store: MokaCache<(u64, u64), Arc<ResponseEntry>>,
view_index: Arc<DashMap<ViewName, DashSet<(u64, u64)>>>,
enabled: bool,
hits: AtomicU64,
misses: AtomicU64,
}
impl ResponseCache {
#[must_use]
pub fn new(config: ResponseCacheConfig) -> Self {
let view_index: Arc<DashMap<ViewName, DashSet<(u64, u64)>>> = Arc::new(DashMap::new());
let vi = Arc::clone(&view_index);
let mut builder = MokaCache::builder()
.max_capacity(config.max_entries as u64)
.eviction_listener(move |key: Arc<(u64, u64)>, value: Arc<ResponseEntry>, _cause| {
for view in &value.accessed_views {
if let Some(keys) = vi.get(view) {
keys.remove(&*key);
}
}
});
if config.ttl_seconds > 0 {
builder = builder.time_to_live(Duration::from_secs(config.ttl_seconds));
}
let store = builder.build();
Self {
store,
view_index,
enabled: config.enabled,
hits: AtomicU64::new(0),
misses: AtomicU64::new(0),
}
}
#[must_use]
pub const fn is_enabled(&self) -> bool {
self.enabled
}
pub fn get(&self, query_key: u64, security_hash: u64) -> Result<Option<Arc<Value>>> {
if !self.enabled {
return Ok(None);
}
let key = (query_key, security_hash);
if let Some(entry) = self.store.get(&key) {
self.hits.fetch_add(1, Ordering::Relaxed);
Ok(Some(Arc::clone(&entry.response)))
} else {
self.misses.fetch_add(1, Ordering::Relaxed);
Ok(None)
}
}
pub fn put(
&self,
query_key: u64,
security_hash: u64,
response: Arc<Value>,
accessed_views: Vec<String>,
) -> Result<()> {
if !self.enabled {
return Ok(());
}
let key = (query_key, security_hash);
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(key);
}
let entry = Arc::new(ResponseEntry {
response,
accessed_views,
});
self.store.insert(key, entry);
Ok(())
}
pub fn invalidate_views(&self, views: &[ViewName]) -> Result<u64> {
let mut total = 0_u64;
for view in views {
if let Some(keys) = self.view_index.get(view.as_str()) {
let to_remove: Vec<(u64, u64)> = keys.iter().map(|k| *k).collect();
drop(keys);
for key in to_remove {
self.store.invalidate(&key);
total += 1;
}
}
}
Ok(total)
}
#[must_use]
pub fn metrics(&self) -> (u64, u64) {
(self.hits.load(Ordering::Relaxed), self.misses.load(Ordering::Relaxed))
}
#[cfg(test)]
pub(crate) fn run_pending_tasks(&self) {
self.store.run_pending_tasks();
}
}
#[must_use]
pub fn hash_security_context(ctx: Option<&SecurityContext>) -> u64 {
use std::hash::{Hash, Hasher};
let Some(ctx) = ctx else {
return 0;
};
let mut hasher = ahash::AHasher::default();
ctx.user_id.hash(&mut hasher);
let mut sorted_roles = ctx.roles.clone();
sorted_roles.sort();
for role in &sorted_roles {
role.hash(&mut hasher);
}
ctx.tenant_id.hash(&mut hasher);
let mut sorted_scopes = ctx.scopes.clone();
sorted_scopes.sort();
for scope in &sorted_scopes {
scope.hash(&mut hasher);
}
if !ctx.attributes.is_empty() {
let mut attr_keys: Vec<&String> = ctx.attributes.keys().collect();
attr_keys.sort();
for key in attr_keys {
key.hash(&mut hasher);
serde_json::to_string(&ctx.attributes[key])
.unwrap_or_default()
.hash(&mut hasher);
}
}
hasher.finish()
}