use std::collections::{HashMap, VecDeque};
use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{mpsc, Arc, Mutex, OnceLock};
use std::time::{Duration, Instant};
pub type CachedBackend<T> = Arc<Mutex<T>>;
struct FinalHandleReap {
released: Box<dyn Fn() -> bool + Send>,
cleanup: Box<dyn Fn() + Send>,
}
struct ReaperBackoff {
initial: Duration,
current: Duration,
maximum: Duration,
}
impl ReaperBackoff {
fn new(initial: Duration, maximum: Duration) -> Self {
Self {
initial,
current: initial,
maximum,
}
}
fn current(&self) -> Duration {
self.current
}
fn progress(&mut self) {
self.current = self.initial;
}
fn no_progress(&mut self) {
self.current = self.current.saturating_mul(2).min(self.maximum);
}
}
fn final_handle_reaper() -> &'static mpsc::Sender<FinalHandleReap> {
static SENDER: OnceLock<mpsc::Sender<FinalHandleReap>> = OnceLock::new();
SENDER.get_or_init(|| {
let (sender, receiver) = mpsc::channel::<FinalHandleReap>();
let _ = std::thread::Builder::new()
.name("car-backend-handle-reaper".into())
.spawn(move || {
let mut pending = Vec::<FinalHandleReap>::new();
let mut backoff =
ReaperBackoff::new(Duration::from_millis(10), Duration::from_millis(250));
loop {
if pending.is_empty() {
match receiver.recv() {
Ok(item) => pending.push(item),
Err(_) => break,
}
} else {
match receiver.recv_timeout(backoff.current()) {
Ok(item) => pending.push(item),
Err(mpsc::RecvTimeoutError::Timeout) => {}
Err(mpsc::RecvTimeoutError::Disconnected) => {}
}
}
while let Ok(item) = receiver.try_recv() {
pending.push(item);
}
let before = pending.len();
let mut index = 0;
while index < pending.len() {
if (pending[index].released)() {
let item = pending.swap_remove(index);
(item.cleanup)();
} else {
index += 1;
}
}
if pending.len() < before {
backoff.progress();
} else {
backoff.no_progress();
}
}
});
sender
})
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum BackendRetention {
Resident,
Transient,
}
pub struct SharedModelBudget {
budget_bytes: AtomicU64,
total_bytes: AtomicU64,
}
impl SharedModelBudget {
pub fn new(budget_bytes: u64) -> Arc<Self> {
Arc::new(Self {
budget_bytes: AtomicU64::new(budget_bytes),
total_bytes: AtomicU64::new(0),
})
}
pub fn from_env_or(default_mb: u64) -> Arc<Self> {
let mb = std::env::var("CAR_INFERENCE_MODEL_CACHE_MB")
.ok()
.and_then(|v| v.parse::<u64>().ok())
.unwrap_or(default_mb);
Self::new(mb.saturating_mul(1024 * 1024))
}
fn add(&self, n: u64) {
self.total_bytes.fetch_add(n, Ordering::Relaxed);
}
fn sub(&self, n: u64) {
let _ = self
.total_bytes
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |cur| {
Some(cur.saturating_sub(n))
});
}
fn total(&self) -> u64 {
self.total_bytes.load(Ordering::Relaxed)
}
fn over_budget(&self) -> bool {
let budget = self.budget_bytes();
budget != 0 && self.total() > budget
}
pub fn is_disabled(&self) -> bool {
self.budget_bytes() == 0
}
pub fn set_budget_bytes(&self, budget_bytes: u64) {
self.budget_bytes.store(budget_bytes, Ordering::Release);
}
pub fn budget_bytes(&self) -> u64 {
self.budget_bytes.load(Ordering::Acquire)
}
}
pub fn default_model_cache_mb() -> u64 {
crate::hardware::HardwareInfo::detect().max_model_mb
}
pub fn idle_ttl_from_env() -> Option<Duration> {
let idle_secs = std::env::var("CAR_INFERENCE_MODEL_IDLE_SECS")
.ok()
.and_then(|v| v.parse::<u64>().ok())
.unwrap_or(300);
(idle_secs > 0).then(|| Duration::from_secs(idle_secs))
}
struct Entry<T> {
backend: CachedBackend<T>,
allocation_id: String,
size_bytes: u64,
last_used: Instant,
invalidated: bool,
}
struct Inner<T> {
map: HashMap<String, Entry<T>>,
lru: VecDeque<String>,
}
pub struct BackendCache<T: Send + 'static> {
inner: Mutex<Inner<T>>,
budget: Arc<SharedModelBudget>,
idle_ttl: Option<Duration>,
load_locks: Mutex<HashMap<String, Arc<Mutex<()>>>>,
resident_accounting: Option<Arc<crate::resource_policy::LocalAdmissionCoordinator>>,
allocation_scope: u64,
}
fn next_cache_allocation_scope() -> u64 {
static NEXT: AtomicU64 = AtomicU64::new(1);
NEXT.fetch_add(1, Ordering::Relaxed)
}
impl<T: Send + 'static> Drop for BackendCache<T> {
fn drop(&mut self) {
let inner = self
.inner
.get_mut()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let idle = inner
.map
.iter()
.filter(|(_, entry)| Arc::strong_count(&entry.backend) == 1)
.map(|(_, entry)| (entry.allocation_id.clone(), entry.size_bytes))
.collect::<Vec<_>>();
let active = inner
.map
.iter()
.filter(|(_, entry)| Arc::strong_count(&entry.backend) > 1)
.map(|(_, entry)| {
(
entry.allocation_id.clone(),
entry.size_bytes,
Arc::downgrade(&entry.backend),
)
})
.collect::<Vec<_>>();
let bytes = idle
.iter()
.map(|(_, size_bytes)| *size_bytes)
.fold(0_u64, u64::saturating_add);
inner.lru.clear();
self.budget.sub(bytes);
if let Some(accounting) = &self.resident_accounting {
for (key, _) in idle {
accounting.mark_evicted(&key);
}
}
for (key, size_bytes, backend) in active {
let accounting = self.resident_accounting.clone();
let budget = self.budget.clone();
let item = FinalHandleReap {
released: Box::new(move || backend.strong_count() == 0),
cleanup: Box::new(move || {
budget.sub(size_bytes);
if let Some(accounting) = &accounting {
accounting.mark_evicted(&key);
}
}),
};
if let Err(error) = final_handle_reaper().send(item) {
std::mem::forget(error.0);
}
}
inner.map.clear();
}
}
impl<T: Send + 'static> BackendCache<T> {
pub fn new(budget_bytes: u64) -> Self {
Self::from_shared(SharedModelBudget::new(budget_bytes), None)
}
pub fn with_idle_ttl(budget_bytes: u64, idle_ttl: Option<Duration>) -> Self {
Self::from_shared(SharedModelBudget::new(budget_bytes), idle_ttl)
}
pub fn from_shared(budget: Arc<SharedModelBudget>, idle_ttl: Option<Duration>) -> Self {
Self::from_shared_with_admission(budget, idle_ttl, None)
}
pub fn from_shared_with_admission(
budget: Arc<SharedModelBudget>,
idle_ttl: Option<Duration>,
resident_accounting: Option<Arc<crate::resource_policy::LocalAdmissionCoordinator>>,
) -> Self {
Self {
inner: Mutex::new(Inner {
map: HashMap::new(),
lru: VecDeque::new(),
}),
budget,
idle_ttl,
load_locks: Mutex::new(HashMap::new()),
resident_accounting,
allocation_scope: next_cache_allocation_scope(),
}
}
pub fn from_env() -> Self {
Self::from_shared(
SharedModelBudget::from_env_or(default_model_cache_mb()),
idle_ttl_from_env(),
)
}
pub fn is_disabled(&self) -> bool {
self.budget.is_disabled()
}
pub fn get_or_load<E>(
&self,
key: &str,
size_bytes: u64,
loader: impl FnOnce() -> Result<T, E>,
) -> Result<CachedBackend<T>, E> {
self.get_or_load_with_publish(key, size_bytes, loader, |allocation_id| {
if let Some(accounting) = &self.resident_accounting {
accounting.mark_resident_allocation(
key,
allocation_id,
size_bytes.div_ceil(1024 * 1024),
);
}
Ok(())
})
.map(|(handle, _retained)| handle)
}
pub fn get_or_load_admitted(
&self,
key: &str,
size_bytes: u64,
reservation: &mut crate::resource_policy::LocalLoadReservation,
loader: impl FnOnce() -> Result<T, crate::InferenceError>,
) -> Result<(CachedBackend<T>, BackendRetention), crate::InferenceError> {
let allocation_id = format!("cache:{}:{key}", self.allocation_scope);
reservation.bind_allocation_id(&allocation_id);
reservation
.reconcile_measured_weights(size_bytes)
.map_err(crate::InferenceError::from)?;
let (handle, retained) =
self.get_or_load_with_publish(key, size_bytes, loader, |allocation_id| {
reservation.publish_resident_weights_as(allocation_id, size_bytes);
Ok(())
})?;
if retained {
reservation.publish_resident_weights(size_bytes);
}
Ok((
handle,
if retained {
BackendRetention::Resident
} else {
BackendRetention::Transient
},
))
}
fn get_or_load_with_publish<E>(
&self,
key: &str,
size_bytes: u64,
loader: impl FnOnce() -> Result<T, E>,
on_publish: impl FnOnce(&str) -> Result<(), E>,
) -> Result<(CachedBackend<T>, bool), E> {
{
let mut guard = self.inner.lock().expect("backend cache poisoned");
let stale_is_idle = guard
.map
.get(key)
.is_some_and(|entry| entry.invalidated && Arc::strong_count(&entry.backend) == 1);
if stale_is_idle {
let stale = guard.map.remove(key).expect("entry checked above");
guard.lru.retain(|candidate| candidate != key);
self.budget.sub(stale.size_bytes);
if let Some(accounting) = &self.resident_accounting {
accounting.mark_evicted(&stale.allocation_id);
}
}
if let Some(entry) = guard.map.get_mut(key) {
let handle = Arc::clone(&entry.backend);
entry.last_used = Instant::now();
guard.lru.retain(|k| k != key);
guard.lru.push_back(key.to_string());
return Ok((handle, true));
}
}
let load_gate = {
let mut gates = self.load_locks.lock().expect("backend load gates poisoned");
Arc::clone(
gates
.entry(key.to_string())
.or_insert_with(|| Arc::new(Mutex::new(()))),
)
};
let _singleflight = load_gate.lock().expect("backend load gate poisoned");
{
let mut guard = self.inner.lock().expect("backend cache poisoned");
if let Some(entry) = guard.map.get_mut(key) {
let handle = Arc::clone(&entry.backend);
entry.last_used = Instant::now();
guard.lru.retain(|candidate| candidate != key);
guard.lru.push_back(key.to_string());
drop(_singleflight);
self.release_load_gate(key, &load_gate);
return Ok((handle, true));
}
}
let backend = match loader() {
Ok(backend) => backend,
Err(error) => {
drop(_singleflight);
self.release_load_gate(key, &load_gate);
return Err(error);
}
};
let handle = Arc::new(Mutex::new(backend));
if self.budget.is_disabled() {
drop(_singleflight);
self.release_load_gate(key, &load_gate);
return Ok((handle, false));
}
let mut guard = self.inner.lock().expect("backend cache poisoned");
if let Some(existing) = guard.map.get(key) {
let existing = Arc::clone(&existing.backend);
drop(guard);
drop(_singleflight);
self.release_load_gate(key, &load_gate);
return Ok((existing, true));
}
let allocation_id = format!("cache:{}:{key}", self.allocation_scope);
if let Err(error) = on_publish(&allocation_id) {
drop(guard);
drop(_singleflight);
self.release_load_gate(key, &load_gate);
return Err(error);
}
self.budget.add(size_bytes);
guard.map.insert(
key.to_string(),
Entry {
backend: Arc::clone(&handle),
allocation_id,
size_bytes,
last_used: Instant::now(),
invalidated: false,
},
);
guard.lru.push_back(key.to_string());
let mut examined = guard.lru.len();
while self.budget.over_budget() && examined > 0 {
examined -= 1;
let Some(victim_key) = guard.lru.pop_front() else {
break;
};
if victim_key == key {
guard.lru.push_front(victim_key);
break;
}
let victim_is_idle = guard
.map
.get(&victim_key)
.is_some_and(|entry| Arc::strong_count(&entry.backend) == 1);
if !victim_is_idle {
guard.lru.push_back(victim_key);
continue;
}
if let Some(victim) = guard.map.remove(&victim_key) {
self.budget.sub(victim.size_bytes);
if let Some(accounting) = &self.resident_accounting {
accounting.mark_evicted(&victim.allocation_id);
}
drop(victim);
}
}
drop(guard);
drop(_singleflight);
self.release_load_gate(key, &load_gate);
Ok((handle, true))
}
fn release_load_gate(&self, key: &str, gate: &Arc<Mutex<()>>) {
let mut gates = self.load_locks.lock().expect("backend load gates poisoned");
if gates
.get(key)
.is_some_and(|current| Arc::ptr_eq(current, gate) && Arc::strong_count(current) == 2)
{
gates.remove(key);
}
}
pub fn invalidate(&self, key: &str) {
let mut guard = self.inner.lock().expect("backend cache poisoned");
let idle = guard
.map
.get(key)
.is_some_and(|entry| Arc::strong_count(&entry.backend) == 1);
if idle {
let Some(entry) = guard.map.remove(key) else {
return;
};
self.budget.sub(entry.size_bytes);
if let Some(accounting) = &self.resident_accounting {
accounting.mark_evicted(&entry.allocation_id);
}
guard.lru.retain(|k| k != key);
} else if let Some(entry) = guard.map.get_mut(key) {
entry.invalidated = true;
}
}
pub fn contains(&self, key: &str) -> bool {
self.inner
.lock()
.expect("backend cache poisoned")
.map
.contains_key(key)
}
pub fn evict_if_idle(&self, key: &str) -> bool {
let mut guard = self.inner.lock().expect("backend cache poisoned");
let Some(entry) = guard.map.get(key) else {
return true;
};
if Arc::strong_count(&entry.backend) != 1 {
return false;
}
let entry = guard.map.remove(key).expect("entry checked above");
guard.lru.retain(|candidate| candidate != key);
self.budget.sub(entry.size_bytes);
if let Some(accounting) = &self.resident_accounting {
accounting.mark_evicted(&entry.allocation_id);
}
true
}
pub fn enforce_budget(&self) -> (usize, u64) {
let mut guard = self.inner.lock().expect("backend cache poisoned");
let mut evicted = 0usize;
let mut bytes = 0u64;
let mut examined = guard.lru.len();
while (self.budget.over_budget() || self.budget.is_disabled()) && examined > 0 {
examined -= 1;
let Some(key) = guard.lru.pop_front() else {
break;
};
let is_idle = guard
.map
.get(&key)
.is_some_and(|entry| Arc::strong_count(&entry.backend) == 1);
if !is_idle {
guard.lru.push_back(key);
continue;
}
if let Some(entry) = guard.map.remove(&key) {
self.budget.sub(entry.size_bytes);
if let Some(accounting) = &self.resident_accounting {
accounting.mark_evicted(&entry.allocation_id);
}
evicted += 1;
bytes = bytes.saturating_add(entry.size_bytes);
}
}
(evicted, bytes)
}
pub fn evict_idle(&self) -> (usize, u64) {
let Some(ttl) = self.idle_ttl else {
return (0, 0);
};
let now = Instant::now();
let mut guard = self.inner.lock().expect("backend cache poisoned");
let stale: Vec<String> = guard
.map
.iter()
.filter(|(_, e)| now.duration_since(e.last_used) >= ttl)
.map(|(k, _)| k.clone())
.collect();
let mut entries = 0usize;
let mut bytes = 0u64;
for key in stale {
let is_idle = guard
.map
.get(&key)
.is_some_and(|entry| Arc::strong_count(&entry.backend) == 1);
if !is_idle {
continue;
}
if let Some(victim) = guard.map.remove(&key) {
self.budget.sub(victim.size_bytes);
if let Some(accounting) = &self.resident_accounting {
accounting.mark_evicted(&victim.allocation_id);
}
guard.lru.retain(|k| k != &key);
entries += 1;
bytes = bytes.saturating_add(victim.size_bytes);
drop(victim);
}
}
(entries, bytes)
}
pub fn clear(&self) {
let mut guard = self.inner.lock().expect("backend cache poisoned");
let idle_keys = guard
.map
.iter()
.filter(|(_, entry)| Arc::strong_count(&entry.backend) == 1)
.map(|(key, _)| key.clone())
.collect::<Vec<_>>();
let mut freed = 0u64;
for key in idle_keys {
if let Some(entry) = guard.map.remove(&key) {
freed = freed.saturating_add(entry.size_bytes);
if let Some(accounting) = &self.resident_accounting {
accounting.mark_evicted(&entry.allocation_id);
}
}
guard.lru.retain(|candidate| candidate != &key);
}
self.budget.sub(freed);
}
pub fn stats(&self) -> (usize, u64, u64) {
let guard = self.inner.lock().expect("backend cache poisoned");
(
guard.map.len(),
self.budget.total(),
self.budget.budget_bytes(),
)
}
}
pub fn estimate_model_size(model_dir: &Path) -> u64 {
fn visit(dir: &Path, total: &mut u64) {
if let Ok(meta) = dir.metadata() {
if meta.is_file() {
*total = total.saturating_add(meta.len());
return;
}
}
let Ok(entries) = std::fs::read_dir(dir) else {
return;
};
for entry in entries.flatten() {
let path = entry.path();
let Ok(link_meta) = path.symlink_metadata() else {
continue;
};
if link_meta.file_type().is_symlink() {
if let Ok(target_meta) = path.metadata() {
if target_meta.is_file() {
*total = total.saturating_add(target_meta.len());
}
}
continue;
}
if link_meta.is_dir() {
visit(&path, total);
continue;
}
if link_meta.is_file() {
*total = total.saturating_add(link_meta.len());
}
}
}
let mut total = 0u64;
visit(model_dir, &mut total);
total
}
#[cfg(test)]
mod tests {
use super::*;
struct FixedProbe;
impl crate::resource_policy::LiveMemoryProbe for FixedProbe {
fn available_memory_mb(
&self,
) -> Result<Option<u64>, crate::resource_policy::ResourcePolicyError> {
Ok(Some(24_000))
}
}
#[test]
fn final_handle_reaper_backoff_is_bounded_and_resets_after_progress() {
let mut backoff = ReaperBackoff::new(Duration::from_millis(10), Duration::from_millis(250));
assert_eq!(backoff.current(), Duration::from_millis(10));
for _ in 0..10 {
backoff.no_progress();
}
assert_eq!(backoff.current(), Duration::from_millis(250));
backoff.progress();
assert_eq!(backoff.current(), Duration::from_millis(10));
}
#[test]
fn cache_hit_returns_same_handle() {
let cache: BackendCache<u32> = BackendCache::new(1024);
let a = cache.get_or_load::<()>("a", 100, || Ok(42)).unwrap();
let b = cache
.get_or_load::<()>("a", 100, || panic!("should not reload"))
.unwrap();
assert!(Arc::ptr_eq(&a, &b));
}
#[test]
fn evicts_lru_when_over_budget() {
let cache: BackendCache<u32> = BackendCache::new(250);
let _a = cache.get_or_load::<()>("a", 100, || Ok(1)).unwrap();
let b = cache.get_or_load::<()>("b", 100, || Ok(2)).unwrap();
let _a_again = cache
.get_or_load::<()>("a", 100, || panic!("cached"))
.unwrap();
drop(b);
let _c = cache.get_or_load::<()>("c", 100, || Ok(3)).unwrap();
let (n, bytes, budget) = cache.stats();
assert_eq!(n, 2, "a + c should remain, b evicted");
assert_eq!(bytes, 200);
assert_eq!(budget, 250);
}
#[test]
fn zero_budget_disables_cache_but_returns_handle() {
let cache: BackendCache<u32> = BackendCache::new(0);
let mut load_count = 0u32;
let a = cache
.get_or_load::<()>("a", 100, || {
load_count += 1;
Ok(1)
})
.unwrap();
assert_eq!(*a.lock().unwrap(), 1);
let b = cache
.get_or_load::<()>("a", 100, || {
load_count += 1;
Ok(1)
})
.unwrap();
assert_eq!(*b.lock().unwrap(), 1);
assert_eq!(load_count, 2, "disabled cache reloads every call");
assert!(!Arc::ptr_eq(&a, &b));
}
#[test]
fn evict_idle_drops_stale_entries_below_capacity() {
let cache: BackendCache<u32> =
BackendCache::with_idle_ttl(1_000_000, Some(Duration::from_millis(20)));
let _a = cache.get_or_load::<()>("a", 100, || Ok(1)).unwrap();
let b = cache.get_or_load::<()>("b", 100, || Ok(2)).unwrap();
assert_eq!(cache.stats().0, 2);
assert_eq!(cache.evict_idle(), (0, 0));
std::thread::sleep(Duration::from_millis(40));
drop(b);
let _a_again = cache
.get_or_load::<()>("a", 100, || panic!("cached"))
.unwrap();
let (entries, bytes) = cache.evict_idle();
assert_eq!((entries, bytes), (1, 100), "only b should be swept");
let (n, total, _) = cache.stats();
assert_eq!(n, 1, "a remains");
assert_eq!(total, 100);
}
#[test]
fn evict_idle_noop_when_disabled() {
let cache: BackendCache<u32> = BackendCache::new(1024); let _a = cache.get_or_load::<()>("a", 100, || Ok(1)).unwrap();
assert_eq!(cache.evict_idle(), (0, 0));
assert_eq!(cache.stats().0, 1, "disabled idle eviction keeps the entry");
}
#[test]
fn invalidate_removes_key() {
let cache: BackendCache<u32> = BackendCache::new(1024);
let a = cache.get_or_load::<()>("a", 100, || Ok(1)).unwrap();
assert_eq!(cache.stats().0, 1);
drop(a);
cache.invalidate("a");
assert_eq!(cache.stats().0, 0);
}
#[test]
fn active_invalidation_is_deferred_without_losing_resident_accounting() {
let cache: BackendCache<u32> = BackendCache::new(1024);
let active = cache.get_or_load::<()>("a", 100, || Ok(1)).unwrap();
cache.invalidate("a");
assert_eq!(cache.stats().1, 100, "active weights remain accounted");
drop(active);
let replacement = cache.get_or_load::<()>("a", 100, || Ok(2)).unwrap();
assert_eq!(*replacement.lock().unwrap(), 2);
assert_eq!(cache.stats().1, 100, "replacement is not double-counted");
}
#[test]
fn installed_size_measurement_covers_single_files_and_non_safetensor_weights() {
let dir = tempfile::tempdir().unwrap();
std::fs::write(dir.path().join("model.gguf"), vec![0_u8; 17]).unwrap();
std::fs::write(dir.path().join("encoder.onnx"), vec![0_u8; 23]).unwrap();
let single = dir.path().join("whisper.bin");
std::fs::write(&single, vec![0_u8; 31]).unwrap();
assert_eq!(estimate_model_size(dir.path()), 71);
assert_eq!(estimate_model_size(&single), 31);
}
#[cfg(unix)]
#[test]
fn installed_size_measurement_follows_hf_leaf_symlinks_but_not_directory_symlinks() {
use std::os::unix::fs::symlink;
let dir = tempfile::tempdir().unwrap();
let blobs = dir.path().join("blobs");
let snapshot = dir.path().join("snapshots/revision");
std::fs::create_dir_all(&blobs).unwrap();
std::fs::create_dir_all(&snapshot).unwrap();
std::fs::write(blobs.join("weights"), vec![0_u8; 47]).unwrap();
symlink("../../blobs/weights", snapshot.join("model.safetensors")).unwrap();
symlink("../../blobs", snapshot.join("directory-link")).unwrap();
symlink("../../blobs/missing", snapshot.join("dangling")).unwrap();
assert_eq!(estimate_model_size(&snapshot), 47);
}
#[test]
fn separate_cache_generations_keep_exact_resident_owners_until_final_handle_drop() {
let coordinator = Arc::new(crate::resource_policy::LocalAdmissionCoordinator::new(
crate::ResourcePolicy::custom_gb(8.0).unwrap(),
crate::hardware::HardwareInfo::detect(),
));
let first: BackendCache<u32> = BackendCache::from_shared_with_admission(
SharedModelBudget::new(1024),
None,
Some(coordinator.clone()),
);
let second: BackendCache<u32> = BackendCache::from_shared_with_admission(
SharedModelBudget::new(1024),
None,
Some(coordinator.clone()),
);
let mut first_reservation = coordinator
.reserve_measured_host("same-model", 100, 0)
.unwrap();
let (first_handle, _) = first
.get_or_load_admitted("same-model", 100, &mut first_reservation, || {
Ok::<_, crate::InferenceError>(1)
})
.unwrap();
drop(first_reservation);
let mut second_reservation = coordinator
.reserve_measured_host("same-model", 100, 0)
.unwrap();
let (_second_handle, _) = second
.get_or_load_admitted("same-model", 100, &mut second_reservation, || {
Ok::<_, crate::InferenceError>(2)
})
.unwrap();
drop(second_reservation);
assert_eq!(coordinator.resident_allocation_ids("same-model").len(), 2);
drop(first);
assert_eq!(coordinator.resident_allocation_ids("same-model").len(), 2);
drop(first_handle);
for _ in 0..100 {
if coordinator.resident_allocation_ids("same-model").len() == 1 {
break;
}
std::thread::sleep(Duration::from_millis(5));
}
assert_eq!(coordinator.resident_allocation_ids("same-model").len(), 1);
}
#[test]
fn shared_budget_spans_caches_and_each_self_trims() {
let budget = SharedModelBudget::new(250);
let a: BackendCache<u32> = BackendCache::from_shared(budget.clone(), None);
let b: BackendCache<u32> = BackendCache::from_shared(budget.clone(), None);
let a1 = a.get_or_load::<()>("a1", 100, || Ok(1)).unwrap();
let _ = b.get_or_load::<()>("b1", 100, || Ok(2)).unwrap();
assert_eq!(budget.total(), 200);
drop(a1);
let _ = a.get_or_load::<()>("a2", 100, || Ok(3)).unwrap(); assert_eq!(budget.total(), 200, "A self-trimmed to the shared budget");
assert_eq!(a.stats().0, 1, "A kept a2, evicted a1");
assert_eq!(b.stats().0, 1, "B's b1 not evicted by A");
}
#[test]
fn default_budget_is_ram_derived_not_flat() {
let mb = default_model_cache_mb();
assert!(mb > 0, "RAM-derived default should be positive");
}
#[test]
fn backend_cache_policy_decrease_evicts_idle_entries() {
let budget = SharedModelBudget::new(1_000);
let cache: BackendCache<u32> = BackendCache::from_shared(budget.clone(), None);
let handle = cache.get_or_load::<()>("model", 600, || Ok(1)).unwrap();
drop(handle);
budget.set_budget_bytes(500);
assert_eq!(cache.enforce_budget(), (1, 600));
assert_eq!(cache.stats(), (0, 0, 500));
}
#[test]
fn backend_cache_policy_decrease_does_not_kill_active_handle() {
let budget = SharedModelBudget::new(1_000);
let cache: BackendCache<u32> = BackendCache::from_shared(budget.clone(), None);
let handle = cache.get_or_load::<()>("model", 600, || Ok(1)).unwrap();
budget.set_budget_bytes(0);
assert_eq!(cache.enforce_budget(), (0, 0));
assert!(cache.contains("model"));
drop(handle);
assert_eq!(cache.enforce_budget(), (1, 600));
}
#[test]
fn backend_cache_same_key_load_is_singleflight() {
let cache = Arc::new(BackendCache::new(1_000));
let loads = Arc::new(AtomicU64::new(0));
let mut threads = Vec::new();
for _ in 0..8 {
let cache = Arc::clone(&cache);
let loads = Arc::clone(&loads);
threads.push(std::thread::spawn(move || {
cache
.get_or_load::<()>("model", 100, || {
loads.fetch_add(1, Ordering::SeqCst);
std::thread::sleep(Duration::from_millis(5));
Ok(42_u32)
})
.unwrap()
}));
}
let handles = threads
.into_iter()
.map(|thread| thread.join().unwrap())
.collect::<Vec<_>>();
assert_eq!(loads.load(Ordering::SeqCst), 1);
assert!(handles
.windows(2)
.all(|pair| Arc::ptr_eq(&pair[0], &pair[1])));
}
#[test]
fn backend_cache_eviction_updates_admission_residency() {
let hw = crate::hardware::HardwareInfo {
os: "test".into(),
arch: "test".into(),
cpu_cores: 8,
total_ram_mb: 32 * 1024,
gpu_backend: crate::hardware::GpuBackend::Metal,
gpu_memory_mb: None,
gpu_devices: Vec::new(),
recommended_model: "fixture".into(),
recommended_context: 4_096,
max_model_mb: 32 * 1024,
};
let admission = Arc::new(crate::resource_policy::LocalAdmissionCoordinator::new(
crate::resource_policy::ResourcePolicy::everyday(),
hw,
));
let budget = SharedModelBudget::new(2 * 1024 * 1024);
let cache: BackendCache<u32> =
BackendCache::from_shared_with_admission(budget.clone(), None, Some(admission.clone()));
let handle = cache
.get_or_load::<()>("model", 1024 * 1024, || Ok(1))
.unwrap();
assert_eq!(admission.resident_model_mb(), 1);
drop(handle);
budget.set_budget_bytes(0);
assert_eq!(cache.enforce_budget().0, 1);
assert_eq!(admission.resident_model_mb(), 0);
}
#[test]
fn backend_cache_drop_never_peer_discounts_an_outstanding_handle() {
let admission = Arc::new(crate::resource_policy::LocalAdmissionCoordinator::new(
crate::resource_policy::ResourcePolicy::custom_gb(8.0).unwrap(),
crate::hardware::HardwareInfo {
total_ram_mb: 32 * 1024,
..crate::hardware::HardwareInfo::detect()
},
));
let budget = SharedModelBudget::new(8 * 1024 * 1024);
let cache: BackendCache<u32> =
BackendCache::from_shared_with_admission(budget.clone(), None, Some(admission.clone()));
let active = cache
.get_or_load::<()>("model", 1024 * 1024, || Ok(1))
.unwrap();
assert_eq!(admission.resident_model_mb(), 1);
drop(cache);
assert_eq!(
admission.resident_model_mb(),
1,
"dropping a cache must not call an outstanding allocation safe"
);
assert_eq!(budget.total(), 1024 * 1024);
drop(active);
for _ in 0..100 {
if admission.resident_model_mb() == 0 && budget.total() == 0 {
break;
}
std::thread::sleep(Duration::from_millis(5));
}
assert_eq!(
admission.resident_model_mb(),
0,
"the final external handle drop must release resident accounting"
);
assert_eq!(budget.total(), 0);
}
#[test]
fn backend_cache_targeted_eviction_refuses_active_then_releases_idle() {
let cache: BackendCache<u32> = BackendCache::new(1_000);
let handle = cache.get_or_load::<()>("model", 100, || Ok(1)).unwrap();
assert!(!cache.evict_if_idle("model"));
assert!(cache.contains("model"));
drop(handle);
assert!(cache.evict_if_idle("model"));
assert!(!cache.contains("model"));
}
#[test]
fn admitted_cache_publication_atomically_transfers_residency() {
let hw = crate::hardware::HardwareInfo {
os: "test".into(),
arch: "test".into(),
cpu_cores: 8,
total_ram_mb: 32 * 1024,
gpu_backend: crate::hardware::GpuBackend::Metal,
gpu_memory_mb: None,
gpu_devices: Vec::new(),
recommended_model: "fixture".into(),
recommended_context: 4_096,
max_model_mb: 32 * 1024,
};
let admission = Arc::new(
crate::resource_policy::LocalAdmissionCoordinator::with_probe(
crate::resource_policy::ResourcePolicy::local_focused(),
hw,
Arc::new(FixedProbe),
),
);
let model = crate::registry::builtin_catalog()
.into_iter()
.find(|model| model.id == "mlx/qwen3-4b:4bit")
.unwrap();
let budget = SharedModelBudget::new(8 * 1024 * 1024 * 1024);
let cache: BackendCache<u32> =
BackendCache::from_shared_with_admission(budget, None, Some(admission.clone()));
let mut reservation = admission.reserve(&model, 2_048).unwrap();
let (handle, retention) = cache
.get_or_load_admitted(&model.id, 3 * 1024 * 1024 * 1024, &mut reservation, || {
Ok::<_, crate::InferenceError>(42)
})
.unwrap();
assert_eq!(*handle.lock().unwrap(), 42);
assert_eq!(retention, BackendRetention::Resident);
assert_eq!(admission.resident_model_mb(), 3 * 1024);
assert_eq!(
admission.preflight(&model, 2_048).active_reservations_mb,
reservation.reserved_incremental_mb()
);
}
#[test]
fn admitted_zero_cache_success_is_explicitly_transient_and_not_resident() {
let admission = Arc::new(
crate::resource_policy::LocalAdmissionCoordinator::with_probe(
crate::resource_policy::ResourcePolicy::custom_gb(8.0).unwrap(),
crate::hardware::HardwareInfo {
total_ram_mb: 32 * 1024,
..crate::hardware::HardwareInfo::detect()
},
Arc::new(FixedProbe),
),
);
let model = crate::registry::builtin_catalog()
.into_iter()
.find(|model| model.id == "mlx/qwen3-4b:4bit")
.unwrap();
let cache: BackendCache<u32> = BackendCache::from_shared_with_admission(
SharedModelBudget::new(0),
None,
Some(admission.clone()),
);
let mut reservation = admission.reserve(&model, 64).unwrap();
let (handle, retention) = cache
.get_or_load_admitted(&model.id, 1024 * 1024, &mut reservation, || {
Ok::<_, crate::InferenceError>(42)
})
.unwrap();
assert_eq!(*handle.lock().unwrap(), 42);
assert_eq!(retention, BackendRetention::Transient);
drop(handle);
drop(reservation);
assert!(!admission.is_resident(&model.id));
}
}