use parking_lot::RwLock;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::{Duration, Instant};
#[derive(Debug, Clone)]
pub struct QueryCacheConfig {
pub max_entries: usize,
pub ttl_secs: u64,
}
impl Default for QueryCacheConfig {
fn default() -> Self {
Self {
max_entries: 1_000,
ttl_secs: 60, }
}
}
pub struct CachedQueryResult {
pub result: Arc<Vec<serde_json::Value>>,
pub cached_at: Instant,
collections: Vec<String>,
}
#[derive(Default)]
struct Inner {
entries: HashMap<String, CachedQueryResult>,
by_collection: HashMap<String, HashSet<String>>,
}
impl Inner {
fn remove_entry(&mut self, key: &str) {
if let Some(entry) = self.entries.remove(key) {
for coll in &entry.collections {
let emptied = match self.by_collection.get_mut(coll) {
Some(set) => {
set.remove(key);
set.is_empty()
}
None => false,
};
if emptied {
self.by_collection.remove(coll);
}
}
}
}
}
pub struct QueryCache {
inner: RwLock<Inner>,
max_entries: usize,
ttl: Duration,
}
impl QueryCache {
pub fn new(max_entries: usize, ttl_secs: u64) -> Self {
Self {
inner: RwLock::new(Inner::default()),
max_entries,
ttl: Duration::from_secs(ttl_secs),
}
}
pub fn with_config(config: &QueryCacheConfig) -> Self {
Self::new(config.max_entries, config.ttl_secs)
}
pub fn get(&self, query_hash: &str) -> Option<Arc<Vec<serde_json::Value>>> {
let inner = self.inner.read();
if let Some(cached) = inner.entries.get(query_hash) {
if cached.cached_at.elapsed() < self.ttl {
return Some(cached.result.clone());
}
}
None
}
pub fn put(&self, query_hash: String, result: Vec<serde_json::Value>) {
let collections = extract_collections_from_key(&query_hash);
let mut inner = self.inner.write();
if inner.entries.len() >= self.max_entries {
let keys_to_remove: Vec<String> = inner
.entries
.keys()
.take(self.max_entries / 2)
.cloned()
.collect();
for key in &keys_to_remove {
inner.remove_entry(key);
}
}
for coll in &collections {
inner
.by_collection
.entry(coll.clone())
.or_default()
.insert(query_hash.clone());
}
inner.entries.insert(
query_hash,
CachedQueryResult {
result: Arc::new(result),
cached_at: Instant::now(),
collections,
},
);
}
pub fn invalidate_all(&self) {
let mut inner = self.inner.write();
inner.entries.clear();
inner.by_collection.clear();
}
pub fn invalidate_collection(&self, collection_name: &str) {
let mut inner = self.inner.write();
let Some(keys_to_remove) = inner.by_collection.remove(collection_name) else {
return;
};
for key in &keys_to_remove {
inner.remove_entry(key);
}
}
pub fn stats(&self) -> QueryCacheStats {
QueryCacheStats {
entries: self.inner.read().entries.len(),
max_entries: self.max_entries,
ttl_secs: self.ttl.as_secs(),
}
}
#[cfg(test)]
fn index_len(&self) -> usize {
self.inner.read().by_collection.len()
}
}
impl Default for QueryCache {
fn default() -> Self {
Self::new(1_000, 60)
}
}
#[derive(Debug, Clone)]
pub struct QueryCacheStats {
pub entries: usize,
pub max_entries: usize,
pub ttl_secs: u64,
}
static QUERY_CACHE: std::sync::OnceLock<QueryCache> = std::sync::OnceLock::new();
pub fn init_query_cache(config: &QueryCacheConfig) {
let _ = QUERY_CACHE.set(QueryCache::with_config(config));
}
pub fn get_query_cache() -> &'static QueryCache {
QUERY_CACHE.get_or_init(QueryCache::default)
}
fn extract_collections_from_key(key: &str) -> Vec<String> {
let Some(slash_pos) = key.find('/') else {
return vec![];
};
let after_slash = &key[slash_pos + 1..];
let Some(colon_pos) = after_slash.rfind(':') else {
return vec![];
};
let colls_str = &after_slash[..colon_pos];
if colls_str.is_empty() {
return vec![];
}
colls_str.split(',').map(|s| s.to_string()).collect()
}
pub fn hash_query(
db_name: &str,
query: &str,
bind_vars: &std::collections::HashMap<String, serde_json::Value>,
) -> String {
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
let mut hasher = DefaultHasher::new();
db_name.hash(&mut hasher);
query.hash(&mut hasher);
let mut sorted_vars: Vec<_> = bind_vars.iter().collect();
sorted_vars.sort_by(|a, b| a.0.cmp(b.0));
for (k, v) in sorted_vars {
k.hash(&mut hasher);
v.hash(&mut hasher);
}
let mut collections = Vec::new();
let upper = query.to_uppercase();
for (i, _) in upper.match_indices(" IN ") {
let after = &query[i + 4..];
if let Some(name) = after.split_whitespace().next() {
if !name.starts_with('@') && !name.starts_with('(') {
collections.push(name.to_string());
}
}
}
collections.sort();
collections.dedup();
if collections.is_empty() {
format!("{}/:{:x}", db_name, hasher.finish())
} else {
format!(
"{}/{}:{:x}",
db_name,
collections.join(","),
hasher.finish()
)
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn test_put_and_get() {
let cache = QueryCache::new(10, 60);
cache.put("db/coll:abc".to_string(), vec![json!({"a": 1})]);
let got = cache.get("db/coll:abc");
assert!(got.is_some());
assert_eq!(got.unwrap().len(), 1);
}
#[test]
fn test_get_missing() {
let cache = QueryCache::new(10, 60);
assert!(cache.get("db/coll:missing").is_none());
}
#[test]
fn test_invalidate_collection() {
let cache = QueryCache::new(10, 60);
cache.put("db/users:1".to_string(), vec![json!({"a": 1})]);
cache.put("db/orders:2".to_string(), vec![json!({"b": 2})]);
cache.put("db/users,orders:3".to_string(), vec![json!({"c": 3})]);
cache.invalidate_collection("users");
assert!(cache.get("db/users:1").is_none());
assert!(cache.get("db/orders:2").is_some());
assert!(cache.get("db/users,orders:3").is_none());
}
#[test]
fn test_invalidate_all() {
let cache = QueryCache::new(10, 60);
cache.put("db/users:1".to_string(), vec![json!({"a": 1})]);
cache.put("db/orders:2".to_string(), vec![json!({"b": 2})]);
cache.invalidate_all();
assert!(cache.get("db/users:1").is_none());
assert!(cache.get("db/orders:2").is_none());
}
#[test]
fn test_extract_collections() {
let mut got = extract_collections_from_key("db/users,orders:abc");
got.sort();
assert_eq!(got, vec!["orders".to_string(), "users".to_string()]);
assert_eq!(
extract_collections_from_key("db/:abc"),
Vec::<String>::new()
);
assert_eq!(
extract_collections_from_key("no_slash"),
Vec::<String>::new()
);
}
#[test]
fn test_eviction() {
let cache = QueryCache::new(2, 60);
cache.put("db/a:1".to_string(), vec![json!({"a": 1})]);
cache.put("db/b:2".to_string(), vec![json!({"b": 2})]);
cache.put("db/c:3".to_string(), vec![json!({"c": 3})]);
let stats = cache.stats();
assert!(stats.entries <= 2);
assert!(cache.index_len() <= stats.entries);
}
#[test]
fn test_invalidate_prunes_index() {
let cache = QueryCache::new(10, 60);
cache.put("db/users:1".to_string(), vec![json!({"a": 1})]);
cache.put("db/users,orders:2".to_string(), vec![json!({"b": 2})]);
assert_eq!(cache.index_len(), 2);
cache.invalidate_collection("users");
assert_eq!(cache.index_len(), 0);
assert_eq!(cache.stats().entries, 0);
}
#[test]
fn test_hash_query_format() {
let key = hash_query(
"mydb",
"FOR doc IN users RETURN doc",
&std::collections::HashMap::new(),
);
assert!(key.starts_with("mydb/users:"));
}
}