use super::CypherQuery;
use std::collections::{HashMap, VecDeque};
use std::sync::{Arc, OnceLock, RwLock};
pub(crate) const CACHE_CAPACITY: usize = 512;
type PlanKey = (u64, u64, bool, u64);
#[derive(Clone)]
pub struct CachedPlan {
pub plan: Arc<CypherQuery>,
pub warnings: Arc<[String]>,
}
struct PlanCache {
map: HashMap<PlanKey, CachedPlan>,
order: VecDeque<PlanKey>,
}
impl PlanCache {
fn new() -> Self {
Self {
map: HashMap::with_capacity(CACHE_CAPACITY),
order: VecDeque::with_capacity(CACHE_CAPACITY),
}
}
}
static CACHE: OnceLock<RwLock<PlanCache>> = OnceLock::new();
fn cache() -> &'static RwLock<PlanCache> {
CACHE.get_or_init(|| RwLock::new(PlanCache::new()))
}
fn hash_query(query: &str) -> u64 {
use std::hash::{Hash, Hasher};
let mut hasher = std::collections::hash_map::DefaultHasher::new();
query.hash(&mut hasher);
hasher.finish()
}
pub fn get(graph_id: u64, version: u64, lazy: bool, query: &str) -> Option<CachedPlan> {
let key = (graph_id, version, lazy, hash_query(query));
let guard = cache().read().expect("plan_cache RwLock poisoned");
let hit = guard.map.get(&key).cloned();
#[cfg(test)]
instrumentation::record_lookup(hit.is_some());
hit
}
pub fn insert(
graph_id: u64,
version: u64,
lazy: bool,
query: &str,
plan: Arc<CypherQuery>,
warnings: Arc<[String]>,
) {
let key = (graph_id, version, lazy, hash_query(query));
let mut guard = cache().write().expect("plan_cache RwLock poisoned");
if guard.map.contains_key(&key) {
return; }
if guard.map.len() >= CACHE_CAPACITY {
if let Some(oldest) = guard.order.pop_front() {
guard.map.remove(&oldest);
#[cfg(test)]
instrumentation::record_eviction();
}
}
guard.order.push_back(key);
guard.map.insert(key, CachedPlan { plan, warnings });
#[cfg(test)]
instrumentation::record_insertion();
}
#[cfg(test)]
pub mod instrumentation {
use std::cell::Cell;
#[derive(Default, Clone, Copy, Debug, PartialEq, Eq)]
pub struct CacheStats {
pub lookups: u64,
pub hits: u64,
pub insertions: u64,
pub evictions: u64,
}
const EMPTY: CacheStats = CacheStats {
lookups: 0,
hits: 0,
insertions: 0,
evictions: 0,
};
impl CacheStats {
fn add(self, other: CacheStats) -> CacheStats {
CacheStats {
lookups: self.lookups + other.lookups,
hits: self.hits + other.hits,
insertions: self.insertions + other.insertions,
evictions: self.evictions + other.evictions,
}
}
}
#[derive(Default, Clone, Copy, Debug, PartialEq, Eq)]
pub struct CallerStats {
pub read: CacheStats,
pub mutation: CacheStats,
pub unclassified: CacheStats,
}
thread_local! {
static PENDING: Cell<CacheStats> = const { Cell::new(EMPTY) };
static TOTALS: Cell<CallerStats> = const {
Cell::new(CallerStats { read: EMPTY, mutation: EMPTY, unclassified: EMPTY })
};
}
fn bump(f: impl FnOnce(&mut CacheStats)) {
PENDING.with(|pending| {
let mut stats = pending.get();
f(&mut stats);
pending.set(stats);
});
}
pub(super) fn record_lookup(hit: bool) {
bump(|stats| {
stats.lookups += 1;
stats.hits += u64::from(hit);
});
}
pub(super) fn record_insertion() {
bump(|stats| stats.insertions += 1);
}
pub(super) fn record_eviction() {
bump(|stats| stats.evictions += 1);
}
fn take_pending() -> CacheStats {
PENDING.with(|pending| pending.replace(EMPTY))
}
pub fn begin_prepare() {
let leftover = take_pending();
if leftover != EMPTY {
TOTALS.with(|totals| {
let mut all = totals.get();
all.unclassified = all.unclassified.add(leftover);
totals.set(all);
});
}
}
pub fn classify_pending(is_mutation: bool) {
let pending = take_pending();
TOTALS.with(|totals| {
let mut all = totals.get();
if is_mutation {
all.mutation = all.mutation.add(pending);
} else {
all.read = all.read.add(pending);
}
totals.set(all);
});
}
pub fn reset() {
take_pending();
TOTALS.with(|totals| totals.set(CallerStats::default()));
}
pub fn totals() -> CallerStats {
TOTALS.with(|totals| totals.get())
}
}
#[cfg(test)]
pub static TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
#[cfg(test)]
pub fn clear_for_tests() {
let mut guard = cache().write().expect("plan_cache RwLock poisoned");
guard.map.clear();
guard.order.clear();
}
#[cfg(test)]
pub fn entry_count_for_tests() -> usize {
cache()
.read()
.expect("plan_cache RwLock poisoned")
.map
.len()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::graph::languages::cypher::parser::parse_cypher;
fn plan(q: &str) -> Arc<CypherQuery> {
Arc::new(parse_cypher(q).expect("parse"))
}
fn no_warnings() -> Arc<[String]> {
Vec::new().into()
}
#[test]
fn miss_then_hit_same_key() {
let _g = TEST_LOCK.lock().unwrap_or_else(|p| p.into_inner());
clear_for_tests();
let q = "MATCH (n:T) RETURN n";
assert!(get(1, 0, false, q).is_none(), "cold miss");
insert(
1,
0,
false,
q,
plan(q),
vec!["typo'd label".to_string()].into(),
);
let hit = get(1, 0, false, q).expect("warm hit");
assert_eq!(
&*hit.warnings,
["typo'd label".to_string()],
"a hit carries the warnings its miss computed"
);
}
#[test]
fn version_graph_id_and_lazy_partition_the_key() {
let _g = TEST_LOCK.lock().unwrap_or_else(|p| p.into_inner());
clear_for_tests();
let q = "MATCH (n:T) RETURN n";
insert(7, 3, false, q, plan(q), no_warnings());
assert!(get(7, 4, false, q).is_none(), "version change invalidates");
assert!(
get(8, 3, false, q).is_none(),
"different graph never collides"
);
assert!(get(7, 3, true, q).is_none(), "lazy mode is part of the key");
assert!(get(7, 3, false, q).is_some(), "exact key hits");
}
#[test]
fn evicts_at_capacity() {
let _g = TEST_LOCK.lock().unwrap_or_else(|p| p.into_inner());
clear_for_tests();
for i in 0..(CACHE_CAPACITY as u64 + 5) {
insert(
1,
i,
false,
"MATCH (n:T) RETURN n",
plan("MATCH (n:T) RETURN n"),
no_warnings(),
);
}
assert_eq!(entry_count_for_tests(), CACHE_CAPACITY);
}
}