use crate::epoch::{Epoch, Snapshot};
use crate::page::CachedPage;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct CacheStats {
pub hits: u64,
pub misses: u64,
pub try_lock_misses: u64,
}
impl CacheStats {
pub fn hit_rate(&self) -> f64 {
let total = self.hits + self.misses;
if total == 0 {
0.0
} else {
self.hits as f64 / total as f64
}
}
}
pub struct PageCache {
map: HashMap<[u8; 32], Entry>,
ring: Vec<[u8; 32]>,
hand: usize,
capacity_bytes: u64,
used_bytes: u64,
dir: Option<PathBuf>,
persistent: bool,
hits: AtomicU64,
misses: AtomicU64,
}
struct Entry {
bytes: bytes::Bytes,
epoch: Epoch,
freq: u8,
}
const FREQ_MAX: u8 = 3;
impl PageCache {
pub fn new(capacity_bytes: u64) -> Self {
Self {
map: HashMap::new(),
ring: Vec::new(),
hand: 0,
capacity_bytes,
used_bytes: 0,
dir: None,
persistent: false,
hits: AtomicU64::new(0),
misses: AtomicU64::new(0),
}
}
pub fn stats(&self) -> CacheStats {
CacheStats {
hits: self.hits.load(Ordering::Relaxed),
misses: self.misses.load(Ordering::Relaxed),
try_lock_misses: 0,
}
}
pub fn reset_stats(&self) {
self.hits.store(0, Ordering::Relaxed);
self.misses.store(0, Ordering::Relaxed);
}
pub fn with_persistence(mut self, dir: PathBuf) -> Self {
self.persistent = true;
self.dir = Some(dir.clone());
self.load_from_disk();
self
}
pub fn insert(&mut self, page: CachedPage) {
let key = page.content_hash;
let size = page.bytes.len() as u64;
let epoch = page.committed_epoch;
let was_new = !self.map.contains_key(&key);
if let Some(old) = self.map.insert(
key,
Entry {
bytes: page.bytes.clone(),
epoch,
freq: 1,
},
) {
self.used_bytes = self.used_bytes.saturating_sub(old.bytes.len() as u64);
} else if was_new {
self.ring.push(key);
}
self.used_bytes += size;
self.evict_if_needed();
}
pub fn get(&mut self, content_hash: &[u8; 32], snapshot: Snapshot) -> Option<bytes::Bytes> {
let hit = self
.map
.get_mut(content_hash)
.filter(|e| e.epoch <= snapshot.epoch)
.map(|entry| {
if entry.freq < FREQ_MAX {
entry.freq += 1;
}
entry.bytes.clone()
});
let counter = if hit.is_some() {
&self.hits
} else {
&self.misses
};
counter.fetch_add(1, Ordering::Relaxed);
hit
}
pub fn try_get(&self, content_hash: &[u8; 32], snapshot: Snapshot) -> Option<bytes::Bytes> {
let hit = self.map.get(content_hash).and_then(|e| {
if e.epoch <= snapshot.epoch {
Some(e.bytes.clone())
} else {
None
}
});
let counter = if hit.is_some() {
&self.hits
} else {
&self.misses
};
counter.fetch_add(1, Ordering::Relaxed);
hit
}
pub fn used_bytes(&self) -> u64 {
self.used_bytes
}
pub fn len(&self) -> usize {
self.map.len()
}
pub fn is_empty(&self) -> bool {
self.map.is_empty()
}
pub fn flush_to_disk(&self) {
if !self.persistent {
return;
}
for (key, entry) in &self.map {
self.spill(key, entry.epoch, &entry.bytes);
}
}
fn spill(&self, key: &[u8; 32], epoch: Epoch, bytes: &[u8]) {
if !self.persistent {
return;
}
let Some(dir) = &self.dir else { return };
let _ = std::fs::create_dir_all(dir);
let hex = hex_key(key);
let path = dir.join(format!("{hex}.{}", epoch.0));
let tmp = dir.join(format!("{hex}.{}.tmp", epoch.0));
if std::fs::write(&tmp, bytes).is_ok() {
let _ = std::fs::rename(&tmp, &path);
}
}
fn load_from_disk(&mut self) {
let Some(dir) = &self.dir else { return };
let entries = match std::fs::read_dir(dir) {
Ok(e) => e,
Err(_) => return,
};
for entry in entries.flatten() {
let name = entry.file_name();
let Some(s) = name.to_str() else { continue };
if s.ends_with(".tmp") {
continue;
}
let (key, epoch) = match s.rsplit_once('.') {
Some((hex, suffix)) => {
let key = match decode_hex_key(hex) {
Some(k) => k,
None => continue,
};
let epoch = match suffix.parse::<u64>() {
Ok(e) => Epoch(e),
Err(_) => continue,
};
(key, epoch)
}
None => {
let key = match decode_hex_key(s) {
Some(k) => k,
None => continue,
};
(key, Epoch(u64::MAX))
}
};
if self.map.contains_key(&key) {
continue;
}
let Ok(bytes) = std::fs::read(entry.path()) else {
continue;
};
let size = bytes.len() as u64;
if self.used_bytes.saturating_add(size) > self.capacity_bytes {
break; }
self.map.insert(
key,
Entry {
bytes: bytes::Bytes::from(bytes),
epoch,
freq: 0,
},
);
self.ring.push(key);
self.used_bytes += size;
}
}
fn evict_if_needed(&mut self) {
if self.ring.is_empty() {
return;
}
while self.used_bytes > self.capacity_bytes {
let mut scanned = 0;
let len = self.ring.len();
loop {
let key = self.ring[self.hand];
let Some(entry) = self.map.get_mut(&key) else {
self.ring.swap_remove(self.hand);
if self.hand >= self.ring.len() && !self.ring.is_empty() {
self.hand %= self.ring.len();
}
break;
};
if entry.freq == 0 {
let removed = self.map.remove(&key).expect("entry present");
self.spill(&key, removed.epoch, &removed.bytes);
self.used_bytes = self.used_bytes.saturating_sub(removed.bytes.len() as u64);
self.ring.swap_remove(self.hand);
if !self.ring.is_empty() {
self.hand %= self.ring.len();
}
break;
} else {
entry.freq -= 1;
scanned += 1;
if scanned > len * 2 {
let removed = self.map.remove(&key).expect("entry present");
self.spill(&key, removed.epoch, &removed.bytes);
self.used_bytes =
self.used_bytes.saturating_sub(removed.bytes.len() as u64);
self.ring.swap_remove(self.hand);
if !self.ring.is_empty() {
self.hand %= self.ring.len();
}
break;
}
self.hand = (self.hand + 1) % len;
}
}
if self.ring.is_empty() {
break;
}
}
}
}
fn hex_key(key: &[u8; 32]) -> String {
let mut s = String::with_capacity(64);
for b in key {
s.push_str(&format!("{b:02x}"));
}
s
}
fn decode_hex_key(hex: &str) -> Option<[u8; 32]> {
if hex.len() != 64 {
return None;
}
let mut out = [0u8; 32];
for i in 0..32 {
out[i] = u8::from_str_radix(&hex[i * 2..i * 2 + 2], 16).ok()?;
}
Some(out)
}
pub struct DecodedPageCache {
map: HashMap<[u8; 32], std::sync::Arc<crate::columnar::NativeColumn>>,
ring: Vec<[u8; 32]>,
hand: usize,
capacity_bytes: u64,
used_bytes: u64,
hits: AtomicU64,
misses: AtomicU64,
}
impl DecodedPageCache {
pub fn new(capacity_bytes: u64) -> Self {
Self {
map: HashMap::new(),
ring: Vec::new(),
hand: 0,
capacity_bytes,
used_bytes: 0,
hits: AtomicU64::new(0),
misses: AtomicU64::new(0),
}
}
pub fn stats(&self) -> CacheStats {
CacheStats {
hits: self.hits.load(Ordering::Relaxed),
misses: self.misses.load(Ordering::Relaxed),
try_lock_misses: 0,
}
}
pub fn reset_stats(&self) {
self.hits.store(0, Ordering::Relaxed);
self.misses.store(0, Ordering::Relaxed);
}
pub fn try_get(&self, key: &[u8; 32]) -> Option<std::sync::Arc<crate::columnar::NativeColumn>> {
match self.map.get(key).cloned() {
Some(v) => {
self.hits.fetch_add(1, Ordering::Relaxed);
Some(v)
}
None => {
self.misses.fetch_add(1, Ordering::Relaxed);
None
}
}
}
pub fn insert(&mut self, key: [u8; 32], col: std::sync::Arc<crate::columnar::NativeColumn>) {
let size = col.approx_bytes();
let was_new = !self.map.contains_key(&key);
if let Some(old) = self.map.insert(key, col) {
self.used_bytes = self.used_bytes.saturating_sub(old.approx_bytes());
} else if was_new {
self.ring.push(key);
}
self.used_bytes += size;
self.evict_if_needed();
}
fn evict_if_needed(&mut self) {
while self.used_bytes > self.capacity_bytes && !self.ring.is_empty() {
let idx = self.hand.min(self.ring.len() - 1);
let key = self.ring.swap_remove(idx);
if let Some(col) = self.map.remove(&key) {
self.used_bytes = self.used_bytes.saturating_sub(col.approx_bytes());
}
if self.ring.is_empty() {
self.hand = 0;
} else {
self.hand %= self.ring.len();
}
}
}
pub fn len(&self) -> usize {
self.map.len()
}
pub fn is_empty(&self) -> bool {
self.map.is_empty()
}
pub fn used_bytes(&self) -> u64 {
self.used_bytes
}
}
pub struct Sharded<T> {
shards: Vec<parking_lot::Mutex<T>>,
try_lock_misses: AtomicU64,
}
pub const CACHE_SHARDS: usize = 16;
impl<T> Sharded<T> {
pub fn new(n_shards: usize, make: impl FnMut() -> T) -> Self {
let mut make = make;
let shards = (0..n_shards)
.map(|_| parking_lot::Mutex::new(make()))
.collect();
Self {
shards,
try_lock_misses: AtomicU64::new(0),
}
}
fn idx(&self, key: &[u8; 32]) -> usize {
let h = u32::from_le_bytes([key[0], key[1], key[2], key[3]]);
(h as usize) % self.shards.len()
}
pub fn try_lock(&self, key: &[u8; 32]) -> Option<parking_lot::MutexGuard<'_, T>> {
match self.shards[self.idx(key)].try_lock() {
Some(g) => Some(g),
None => {
self.try_lock_misses.fetch_add(1, Ordering::Relaxed);
None
}
}
}
pub fn lock(&self, key: &[u8; 32]) -> parking_lot::MutexGuard<'_, T> {
self.shards[self.idx(key)].lock()
}
pub fn try_lock_misses(&self) -> u64 {
self.try_lock_misses.load(Ordering::Relaxed)
}
}
impl Sharded<PageCache> {
pub fn stats(&self) -> CacheStats {
let mut total = CacheStats::default();
for s in &self.shards {
let st = s.lock().stats();
total.hits += st.hits;
total.misses += st.misses;
}
total.try_lock_misses = self.try_lock_misses();
total
}
pub fn reset_stats(&self) {
for s in &self.shards {
s.lock().reset_stats();
}
self.try_lock_misses.store(0, Ordering::Relaxed);
}
pub fn flush_to_disk(&self) {
for s in &self.shards {
s.lock().flush_to_disk();
}
}
pub fn len(&self) -> usize {
self.shards.iter().map(|s| s.lock().len()).sum()
}
#[allow(clippy::len_without_is_empty)]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
impl Sharded<DecodedPageCache> {
pub fn stats(&self) -> CacheStats {
let mut total = CacheStats::default();
for s in &self.shards {
let st = s.lock().stats();
total.hits += st.hits;
total.misses += st.misses;
}
total.try_lock_misses = self.try_lock_misses();
total
}
pub fn reset_stats(&self) {
for s in &self.shards {
s.lock().reset_stats();
}
self.try_lock_misses.store(0, Ordering::Relaxed);
}
pub fn len(&self) -> usize {
self.shards.iter().map(|s| s.lock().len()).sum()
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
fn page(hash: [u8; 32], epoch: u64, data: &[u8]) -> CachedPage {
CachedPage {
committed_epoch: Epoch(epoch),
content_hash: hash,
bytes: bytes::Bytes::copy_from_slice(data),
}
}
#[test]
fn mvcc_visibility() {
let mut cache = PageCache::new(1 << 20);
let hash = [1u8; 32];
cache.insert(page(hash, 3, b"v3"));
assert!(cache.get(&hash, Snapshot::at(Epoch(2))).is_none());
assert_eq!(
cache.get(&hash, Snapshot::at(Epoch(3))),
Some(bytes::Bytes::copy_from_slice(b"v3"))
);
}
#[test]
fn hit_miss_counters_track_lookups() {
let mut cache = PageCache::new(1 << 20);
let hash = [7u8; 32];
cache.insert(page(hash, 1, b"v"));
assert_eq!(
cache.stats(),
CacheStats {
hits: 0,
misses: 0,
try_lock_misses: 0
}
);
assert!(cache.get(&hash, Snapshot::at(Epoch(1))).is_some());
assert!(cache.try_get(&hash, Snapshot::at(Epoch(1))).is_some());
assert!(cache.get(&[9u8; 32], Snapshot::at(Epoch(1))).is_none());
assert!(cache.get(&hash, Snapshot::at(Epoch(0))).is_none());
let s = cache.stats();
assert_eq!(
s,
CacheStats {
hits: 2,
misses: 2,
try_lock_misses: 0
}
);
assert!((s.hit_rate() - 0.5).abs() < 1e-9);
cache.reset_stats();
assert_eq!(
cache.stats(),
CacheStats {
hits: 0,
misses: 0,
try_lock_misses: 0
}
);
assert_eq!(cache.stats().hit_rate(), 0.0);
}
#[test]
fn content_addressed_replacement_ages_out_old() {
let mut cache = PageCache::new(1 << 20);
let hash_a = [1u8; 32];
let hash_b = [2u8; 32];
cache.insert(page(hash_a, 1, b"old"));
cache.insert(page(hash_b, 1, b"new"));
assert_eq!(cache.len(), 2);
assert!(cache.get(&hash_a, Snapshot::at(Epoch(1))).is_some());
}
#[test]
fn capacity_eviction_removes_cold_first() {
let mut cache = PageCache::new(10);
cache.insert(page([1u8; 32], 1, b"0123456789")); assert_eq!(cache.used_bytes(), 10);
cache.insert(page([2u8; 32], 1, b"x")); assert_eq!(cache.used_bytes(), 1);
assert!(cache.get(&[1u8; 32], Snapshot::at(Epoch(1))).is_none());
assert!(cache.get(&[2u8; 32], Snapshot::at(Epoch(1))).is_some());
}
#[test]
fn frequency_keeps_hot_pages_alive() {
let mut cache = PageCache::new(2);
let a = [10u8; 32];
cache.insert(page(a, 1, b"A"));
for _ in 0..5 {
let _ = cache.get(&a, Snapshot::at(Epoch(1)));
}
cache.insert(page([20u8; 32], 1, b"B"));
cache.insert(page([30u8; 32], 1, b"C"));
assert!(
cache.get(&a, Snapshot::at(Epoch(1))).is_some(),
"hot page should survive eviction"
);
}
#[test]
fn persistent_cache_survives_restart() {
let dir = tempdir().unwrap();
let data = b"page-payload-1234";
let hash = [42u8; 32];
{
let mut cache = PageCache::new(1 << 20).with_persistence(dir.path().to_path_buf());
cache.insert(page(hash, 5, data));
cache.flush_to_disk();
}
let backing = dir.path().join(format!("{}.{}", hex_key(&hash), 5));
assert!(backing.exists(), "cache file should be spilled");
let mut cache = PageCache::new(1 << 20).with_persistence(dir.path().to_path_buf());
let got = cache.get(&hash, Snapshot::at(Epoch(u64::MAX)));
assert_eq!(got, Some(bytes::Bytes::copy_from_slice(data)));
}
#[test]
fn try_get_does_not_block() {
let mut cache = PageCache::new(1 << 20);
let hash = [7u8; 32];
cache.insert(page(hash, 1, b"hi"));
let got = cache.try_get(&hash, Snapshot::at(Epoch(1)));
assert!(got.is_some());
let miss = cache.try_get(&[8u8; 32], Snapshot::at(Epoch(1)));
assert!(miss.is_none());
}
#[test]
fn hot_eviction_storm_does_not_orphan_entries() {
let mut cache = PageCache::new(6); let mut next = 1u8;
for _ in 0..200 {
let key = [next; 32];
cache.insert(page(key, 1, b"aa")); for _ in 0..4 {
let _ = cache.get(&key, Snapshot::at(Epoch(1)));
}
next = next.wrapping_add(1);
}
assert!(
cache.used_bytes() <= cache.capacity_bytes + 2, "used_bytes {} leaked past capacity {}",
cache.used_bytes(),
cache.capacity_bytes
);
let mut ok = true;
for k in &cache.ring {
if !cache.map.contains_key(k) {
ok = false;
}
}
assert!(
ok,
"ring references a key absent from the map (orphan slot)"
);
assert_eq!(
cache.ring.len(),
cache.map.len(),
"ring/map size mismatch (orphaned entries or slots)"
);
}
#[test]
fn decoded_cache_hit_miss_counters() {
use crate::columnar::NativeColumn;
let mut cache = DecodedPageCache::new(1 << 20);
let key = [9u8; 32];
cache.insert(
key,
std::sync::Arc::new(NativeColumn::Int64 {
data: vec![1, 2, 3],
validity: vec![],
}),
);
assert!(cache.try_get(&key).is_some());
assert!(cache.try_get(&key).is_some());
assert!(cache.try_get(&[0u8; 32]).is_none());
let s = cache.stats();
assert_eq!(s.hits, 2);
assert_eq!(s.misses, 1);
assert!(s.hit_rate() > 0.66 && s.hit_rate() < 0.67);
cache.reset_stats();
assert_eq!(cache.stats(), CacheStats::default());
}
}