use std::any::Any;
use std::collections::{BTreeMap, HashMap};
use std::sync::atomic::{AtomicI64, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::Duration;
use moka::policy::EvictionPolicy;
use moka::sync::Cache;
use crate::util::elapsed_time;
pub type CacheValue = Arc<dyn Any + Send + Sync>;
const DEFAULT_MAX_ITEMS: u64 = 2000;
const MIN_EXPIRY_MS: u64 = 1000;
const MAX_EXPIRY_MS: u64 = 100 * 365 * 24 * 60 * 60 * 1000;
const HOUSEKEEPING_INTERVAL: Duration = Duration::from_secs(600);
pub struct ManagedCache {
name: String,
expiry_ms: u64,
max_items: u64,
store: Cache<String, CacheValue>,
last_read: AtomicI64,
last_write: AtomicI64,
last_reset: AtomicI64,
}
fn registry() -> &'static Mutex<HashMap<String, Arc<ManagedCache>>> {
static REGISTRY: OnceLock<Mutex<HashMap<String, Arc<ManagedCache>>>> = OnceLock::new();
REGISTRY.get_or_init(|| Mutex::new(HashMap::new()))
}
fn now_ms() -> i64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_millis() as i64)
.unwrap_or_default()
}
impl ManagedCache {
pub fn create_cache(name: &str, expiry_ms: u64) -> Arc<ManagedCache> {
Self::create_cache_with_limit(name, expiry_ms, DEFAULT_MAX_ITEMS)
}
pub fn create_cache_with_limit(
name: &str,
expiry_ms: u64,
max_items: u64,
) -> Arc<ManagedCache> {
Self::create_clamped(
name,
expiry_ms.clamp(MIN_EXPIRY_MS, MAX_EXPIRY_MS),
max_items,
)
}
#[cfg(test)]
fn create_cache_unclamped(name: &str, expiry_ms: u64, max_items: u64) -> Arc<ManagedCache> {
Self::create_clamped(name, expiry_ms.min(MAX_EXPIRY_MS), max_items)
}
fn create_clamped(name: &str, expiry_ms: u64, max_items: u64) -> Arc<ManagedCache> {
let mut collection = registry().lock().expect("managed cache registry");
if let Some(existing) = collection.get(name) {
return existing.clone();
}
let store = Cache::builder()
.max_capacity(max_items)
.time_to_live(Duration::from_millis(expiry_ms))
.eviction_policy(EvictionPolicy::lru())
.build();
let cache = Arc::new(ManagedCache {
name: name.to_string(),
expiry_ms,
max_items,
store,
last_read: AtomicI64::new(0),
last_write: AtomicI64::new(0),
last_reset: AtomicI64::new(now_ms()),
});
collection.insert(name.to_string(), cache.clone());
log::info!(
"Created cache ({}), expiry {}, maxItems={}",
name,
elapsed_time(Duration::from_millis(expiry_ms)),
max_items
);
cache
}
pub fn get_instance(name: &str) -> Option<Arc<ManagedCache>> {
registry()
.lock()
.expect("managed cache registry")
.get(name)
.cloned()
}
pub fn get_cache_collection() -> BTreeMap<String, Arc<ManagedCache>> {
registry()
.lock()
.expect("managed cache registry")
.iter()
.map(|(k, v)| (k.clone(), v.clone()))
.collect()
}
pub fn put<V: Any + Send + Sync>(&self, key: &str, value: V) {
self.put_arc(key, Arc::new(value));
}
pub fn put_arc(&self, key: &str, value: CacheValue) {
if !key.is_empty() {
self.last_write.store(now_ms(), Ordering::Relaxed);
self.store.insert(key.to_string(), value);
}
}
pub fn get(&self, key: &str) -> Option<CacheValue> {
if key.is_empty() {
return None;
}
self.last_read.store(now_ms(), Ordering::Relaxed);
self.store.get(key)
}
pub fn get_as<T: Any + Send + Sync>(&self, key: &str) -> Option<Arc<T>> {
self.get(key).and_then(|value| value.downcast::<T>().ok())
}
pub fn exists(&self, key: &str) -> bool {
self.get(key).is_some()
}
pub fn remove(&self, key: &str) {
if !key.is_empty() {
self.last_write.store(now_ms(), Ordering::Relaxed);
self.store.invalidate(key);
}
}
pub fn clear(&self) {
self.last_reset.store(now_ms(), Ordering::Relaxed);
self.store.invalidate_all();
self.store.run_pending_tasks();
}
pub fn clean_up(&self) {
log::debug!("Cleaning up {}", self.name);
self.store.run_pending_tasks();
}
pub fn name(&self) -> &str {
&self.name
}
pub fn expiry_ms(&self) -> u64 {
self.expiry_ms
}
pub fn max_items(&self) -> u64 {
self.max_items
}
pub fn size(&self) -> u64 {
self.store.entry_count()
}
pub fn entries(&self) -> Vec<(String, CacheValue)> {
self.store.iter().map(|(k, v)| ((*k).clone(), v)).collect()
}
pub fn last_read(&self) -> i64 {
self.last_read.load(Ordering::Relaxed)
}
pub fn last_write(&self) -> i64 {
self.last_write.load(Ordering::Relaxed)
}
pub fn last_reset(&self) -> i64 {
self.last_reset.load(Ordering::Relaxed)
}
}
pub fn start_housekeeping() {
static STARTED: OnceLock<()> = OnceLock::new();
STARTED.get_or_init(|| {
log::info!("Housekeeper started");
tokio::spawn(async {
loop {
tokio::time::sleep(HOUSEKEEPING_INTERVAL).await;
housekeeping();
}
});
});
}
fn housekeeping() {
for cache in ManagedCache::get_cache_collection().into_values() {
cache.clean_up();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn round_trip_remove_clear_and_entries() {
let cache = ManagedCache::create_cache("unit.round.trip", 60_000);
cache.put("s", "text".to_string());
cache.put("n", 7_i64);
assert_eq!(cache.get_as::<String>("s").unwrap().as_str(), "text");
assert_eq!(*cache.get_as::<i64>("n").unwrap(), 7);
assert!(cache.get_as::<i64>("s").is_none());
let mut keys: Vec<String> = cache.entries().into_iter().map(|(k, _)| k).collect();
keys.sort();
assert_eq!(keys, ["n", "s"]);
cache.remove("s");
assert!(!cache.exists("s"));
cache.clear();
assert_eq!(cache.size(), 0);
assert!(cache.get("n").is_none());
}
#[test]
fn pre_wrapped_arc_is_the_documented_trap() {
let cache = ManagedCache::create_cache("unit.wrong.wrap", 60_000);
let wrapped: Arc<String> = Arc::new("hello".to_string());
cache.put("k", wrapped); assert!(cache.get_as::<String>("k").is_none());
assert!(cache.get_as::<Arc<String>>("k").is_some());
cache.put_arc("k2", Arc::new("hello".to_string()));
assert_eq!(cache.get_as::<String>("k2").unwrap().as_str(), "hello");
}
#[test]
fn empty_key_is_a_guarded_no_op() {
let cache = ManagedCache::create_cache("unit.empty.key", 60_000);
cache.put("", "x".to_string());
cache.clean_up();
assert_eq!(cache.size(), 0);
assert!(cache.get("").is_none());
assert!(!cache.exists(""));
cache.remove("");
assert_eq!(cache.last_read(), 0);
assert_eq!(cache.last_write(), 0);
}
#[test]
fn expiry_clamps_at_both_ends() {
let low = ManagedCache::create_cache("unit.clamp.low", 500);
assert_eq!(low.expiry_ms(), 1000);
let high = ManagedCache::create_cache("unit.clamp.high", u64::MAX);
assert_eq!(high.expiry_ms(), MAX_EXPIRY_MS);
}
#[test]
fn create_is_idempotent_first_params_win() {
let first = ManagedCache::create_cache_with_limit("unit.create.idempotent", 5_000, 10);
let second = ManagedCache::create_cache_with_limit("unit.create.idempotent", 9_000, 99);
assert!(Arc::ptr_eq(&first, &second));
assert_eq!(second.expiry_ms(), 5_000);
assert_eq!(second.max_items(), 10);
}
#[test]
fn registry_lookup_and_sorted_collection() {
ManagedCache::create_cache("unit.registry.zeta", 60_000);
ManagedCache::create_cache("unit.registry.alpha", 60_000);
assert!(ManagedCache::get_instance("unit.registry.alpha").is_some());
assert!(ManagedCache::get_instance("no.such.cache").is_none());
let all = ManagedCache::get_cache_collection();
let names: Vec<&str> = all
.keys()
.map(String::as_str)
.filter(|k| k.starts_with("unit.registry."))
.collect();
assert_eq!(names, ["unit.registry.alpha", "unit.registry.zeta"]);
assert_eq!(
all.get("unit.registry.alpha").unwrap().name(),
"unit.registry.alpha"
);
}
#[test]
fn ttl_expires_after_write_lazily() {
let cache = ManagedCache::create_cache_unclamped("unit.ttl.expiry", 200, 100);
cache.put("k", 42_i32);
assert!(cache.exists("k"));
std::thread::sleep(Duration::from_millis(600));
assert!(cache.get_as::<i32>("k").is_none());
assert!(!cache.exists("k"));
}
#[test]
fn ttl_resets_on_update_expire_after_write() {
let cache = ManagedCache::create_cache("unit.ttl.reset", 1200);
cache.put("k", "a".to_string());
std::thread::sleep(Duration::from_millis(800));
cache.put("k", "b".to_string());
std::thread::sleep(Duration::from_millis(800));
assert_eq!(cache.get_as::<String>("k").unwrap().as_str(), "b");
std::thread::sleep(Duration::from_millis(700));
assert!(cache.get("k").is_none());
}
#[test]
fn lru_eviction_is_deterministic() {
let cache = ManagedCache::create_cache_with_limit("unit.lru.eviction", 60_000, 3);
cache.put("a", 1_i32);
cache.put("b", 2_i32);
cache.put("c", 3_i32);
cache.clean_up();
assert!(cache.exists("a"));
assert!(cache.exists("b"));
cache.clean_up();
cache.put("d", 4_i32);
cache.clean_up();
assert!(!cache.exists("c"), "the LRU entry is the victim");
assert!(cache.exists("a"));
assert!(cache.exists("b"));
assert!(cache.exists("d"), "the newcomer is always admitted");
assert_eq!(cache.size(), 3);
}
#[test]
fn telemetry_stamps_follow_the_java_map() {
let cache = ManagedCache::create_cache("unit.telemetry.stamps", 60_000);
assert_eq!(cache.last_read(), 0);
assert_eq!(cache.last_write(), 0);
assert!(cache.last_reset() > 0);
assert!(cache.get("absent").is_none());
assert!(cache.last_read() > 0);
cache.put("k", "v".to_string());
let first_write = cache.last_write();
assert!(first_write > 0);
std::thread::sleep(Duration::from_millis(10));
cache.remove("no-such-key");
assert!(cache.last_write() > first_write);
let first_reset = cache.last_reset();
std::thread::sleep(Duration::from_millis(10));
cache.clear();
assert!(cache.last_reset() > first_reset);
}
#[test]
fn housekeeping_sweeps_idle_caches() {
let cache = ManagedCache::create_cache_unclamped("unit.housekeeping", 50, 100);
cache.put("k", 1_i32);
std::thread::sleep(Duration::from_millis(250));
housekeeping();
assert_eq!(cache.size(), 0);
}
#[test]
fn concurrent_put_get_smoke() {
let cache = ManagedCache::create_cache("unit.concurrent.smoke", 60_000);
let mut handles = Vec::new();
for t in 0..4 {
let cache = cache.clone();
handles.push(std::thread::spawn(move || {
for i in 0..250 {
let key = format!("k{t}-{i}");
cache.put(&key, i);
assert_eq!(*cache.get_as::<i32>(&key).unwrap(), i);
}
}));
}
for handle in handles {
handle.join().unwrap();
}
cache.clean_up();
assert_eq!(cache.size(), 1000);
}
}