use std::collections::HashMap;
use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};
use crate::engine::PartitionTables;
#[derive(Debug, Clone)]
pub struct CacheConfig {
pub enabled: bool,
pub max_bytes: u64,
pub max_cached_source_events: usize,
}
impl CacheConfig {
#[must_use]
pub const fn disabled() -> Self {
Self {
enabled: false,
max_bytes: 0,
max_cached_source_events: 0,
}
}
#[must_use]
pub const fn new(replicas: Option<u32>) -> Self {
match replicas {
Some(1) => Self {
enabled: true,
max_bytes: 64 * 1024 * 1024,
max_cached_source_events: 1_000_000,
},
_ => Self::disabled(),
}
}
}
struct CacheEntry {
watermark: u64,
tables: PartitionTables,
last_used: u64,
}
pub(crate) enum Lookup {
Miss,
Tail {
base: PartitionTables,
from: u64,
},
Hit(PartitionTables),
}
pub(crate) struct DecodeCache {
config: CacheConfig,
entries: Mutex<HashMap<String, CacheEntry>>,
recency: AtomicU64,
}
impl DecodeCache {
pub(crate) fn new(config: CacheConfig) -> Self {
Self {
config,
entries: Mutex::new(HashMap::new()),
recency: AtomicU64::new(0),
}
}
pub(crate) const fn enabled(&self) -> bool {
self.config.enabled
}
pub(crate) const fn max_cached_source_events(&self) -> usize {
self.config.max_cached_source_events
}
pub(crate) fn lookup(&self, partition: &str, event_count: u64) -> Lookup {
if !self.config.enabled {
return Lookup::Miss;
}
let touch = self.touch();
let mut entries = self.entries.lock().expect("poison");
let Some(entry) = entries.get(partition) else {
drop(entries);
return Lookup::Miss;
};
let watermark = entry.watermark;
let tables = entry.tables.clone();
if let Some(entry) = entries.get_mut(partition) {
entry.last_used = touch;
}
match event_count.cmp(&watermark) {
std::cmp::Ordering::Equal => {
drop(entries);
crate::metrics::record_cache_hit();
Lookup::Hit(tables)
}
std::cmp::Ordering::Greater => {
drop(entries);
crate::metrics::record_cache_tail();
Lookup::Tail {
base: tables,
from: watermark,
}
}
std::cmp::Ordering::Less => {
entries.remove(partition);
drop(entries);
crate::metrics::record_cache_full_rebuild();
Lookup::Miss
}
}
}
pub(crate) fn store_full(&self, partition: &str, watermark: u64, tables: PartitionTables) {
self.insert(partition, watermark, tables);
}
pub(crate) fn store_tail(&self, partition: &str, watermark: u64, merged: PartitionTables) {
self.insert(partition, watermark, merged);
}
fn insert(&self, partition: &str, watermark: u64, tables: PartitionTables) {
if !self.config.enabled {
return;
}
let touch = self.touch();
{
let mut entries = self.entries.lock().expect("poison");
entries.insert(
partition.to_owned(),
CacheEntry {
watermark,
tables,
last_used: touch,
},
);
}
self.evict_over_budget();
}
#[allow(
clippy::significant_drop_tightening,
reason = "the lock must stay held across the whole re-check-then-evict loop (each \
iteration's total-bytes check depends on the previous iteration's removal) — \
metrics recording is deferred until after the block specifically so the lock \
is not held during it, but the loop itself cannot tighten further"
)]
fn evict_over_budget(&self) {
let evicted = {
let mut entries = self.entries.lock().expect("poison");
let mut evicted = 0_u32;
loop {
let total_bytes: u64 = entries
.values()
.map(|entry| entry.tables.memory_bytes() as u64)
.sum();
if total_bytes <= self.config.max_bytes || entries.len() <= 1 {
break;
}
let Some(victim) = entries
.iter()
.min_by_key(|(_, entry)| entry.last_used)
.map(|(partition, _)| partition.clone())
else {
break;
};
entries.remove(&victim);
evicted += 1;
}
evicted
};
for _ in 0..evicted {
crate::metrics::record_cache_eviction();
}
}
pub(crate) fn evict(&self, partition: &str) {
let mut entries = self.entries.lock().expect("poison");
if entries.remove(partition).is_some() {
crate::metrics::record_cache_eviction();
}
}
fn touch(&self) -> u64 {
self.recency.fetch_add(1, Ordering::Relaxed)
}
#[cfg(test)]
fn total_bytes(&self) -> u64 {
self.entries
.lock()
.expect("poison")
.values()
.map(|entry| entry.tables.memory_bytes() as u64)
.sum()
}
#[cfg(test)]
fn contains(&self, partition: &str) -> bool {
self.entries.lock().expect("poison").contains_key(partition)
}
}
#[cfg(test)]
mod tests {
use polyc_eventlog::Event;
use polyc_proto::kinds;
use uuid::Uuid;
use super::*;
use crate::engine::decode_partition_tables;
fn sample_tables(partition: &str, n: usize) -> PartitionTables {
let events: Vec<(u64, Event)> = (0..n)
.map(|i| {
let turn = Uuid::now_v7();
(
i as u64,
Event::new(kinds::tagged(kinds::TURN_START, &turn), Vec::new()),
)
})
.collect();
decode_partition_tables(
partition,
&events,
&[],
&crate::engine::fixture_handoff_trust(),
)
.expect("decode")
}
#[test]
fn disabled_cache_always_misses_and_never_stores() {
let cache = DecodeCache::new(CacheConfig::disabled());
assert!(matches!(cache.lookup("conv-a", 3), Lookup::Miss));
cache.store_full("conv-a", 3, sample_tables("conv-a", 3));
assert!(!cache.contains("conv-a"));
assert!(matches!(cache.lookup("conv-a", 3), Lookup::Miss));
}
#[test]
fn the_same_count_is_a_hit() {
let cache = DecodeCache::new(CacheConfig::new(Some(1)));
cache.store_full("conv-a", 3, sample_tables("conv-a", 3));
assert!(matches!(cache.lookup("conv-a", 3), Lookup::Hit(_)));
}
#[test]
fn a_higher_count_is_a_tail() {
let cache = DecodeCache::new(CacheConfig::new(Some(1)));
cache.store_full("conv-a", 3, sample_tables("conv-a", 3));
match cache.lookup("conv-a", 5) {
Lookup::Tail { from, .. } => assert_eq!(from, 3),
_ => panic!("expected a tail lookup"),
}
}
#[test]
fn a_lower_count_is_a_miss_and_drops_the_entry() {
let cache = DecodeCache::new(CacheConfig::new(Some(1)));
cache.store_full("conv-a", 3, sample_tables("conv-a", 3));
assert!(matches!(cache.lookup("conv-a", 2), Lookup::Miss));
assert!(!cache.contains("conv-a"));
}
#[test]
fn no_entry_is_a_miss() {
let cache = DecodeCache::new(CacheConfig::new(Some(1)));
assert!(matches!(cache.lookup("conv-never-seen", 0), Lookup::Miss));
}
#[test]
fn invalidation_evicts_the_named_partition() {
let cache = DecodeCache::new(CacheConfig::new(Some(1)));
cache.store_full("conv-a", 3, sample_tables("conv-a", 3));
assert!(cache.contains("conv-a"));
cache.evict("conv-a");
assert!(!cache.contains("conv-a"));
assert!(matches!(cache.lookup("conv-a", 3), Lookup::Miss));
}
#[test]
fn an_in_place_rewrite_is_refused_by_the_report_the_count_cannot_give() {
let cache = DecodeCache::new(CacheConfig::new(Some(1)));
cache.store_full("conv-a", 3, sample_tables("conv-a", 3));
assert!(
matches!(cache.lookup("conv-a", 3), Lookup::Hit(_)),
"an unchanged count is a hit, which is exactly why the report matters"
);
cache.evict("conv-a");
assert!(
matches!(cache.lookup("conv-a", 3), Lookup::Miss),
"the reported mutation is what refuses the pre-erasure bytes"
);
}
#[test]
fn eviction_keeps_total_bytes_under_the_cap_once_a_third_partition_is_stored() {
let one = sample_tables("conv-a", 50);
let per_entry_bytes = one.memory_bytes() as u64;
let max_bytes = per_entry_bytes * 5 / 2;
let cache = DecodeCache::new(CacheConfig {
enabled: true,
max_bytes,
max_cached_source_events: 1_000_000,
});
cache.store_full("conv-a", 50, sample_tables("conv-a", 50));
cache.store_full("conv-b", 50, sample_tables("conv-b", 50));
assert!(matches!(cache.lookup("conv-b", 50), Lookup::Hit(_)));
cache.store_full("conv-c", 50, sample_tables("conv-c", 50));
assert!(
cache.total_bytes() <= max_bytes,
"total cached bytes must stay under the configured cap after eviction"
);
assert!(!cache.contains("conv-a"), "the LRU entry must be evicted");
assert!(cache.contains("conv-b"), "the touched entry must survive");
assert!(
cache.contains("conv-c"),
"the just-inserted entry must survive"
);
}
#[test]
fn cache_config_new_enables_only_at_exactly_one_replica() {
assert!(CacheConfig::new(Some(1)).enabled);
assert!(!CacheConfig::new(Some(2)).enabled);
assert!(!CacheConfig::new(None).enabled);
assert!(!CacheConfig::new(Some(0)).enabled);
}
}