1use std::collections::{HashMap, VecDeque};
45use std::path::Path;
46use std::sync::atomic::{AtomicU64, Ordering};
47use std::sync::{mpsc, Arc, Mutex, OnceLock};
48use std::time::{Duration, Instant};
49
50pub type CachedBackend<T> = Arc<Mutex<T>>;
54
55struct FinalHandleReap {
56 released: Box<dyn Fn() -> bool + Send>,
57 cleanup: Box<dyn Fn() + Send>,
58}
59
60struct ReaperBackoff {
61 initial: Duration,
62 current: Duration,
63 maximum: Duration,
64}
65
66impl ReaperBackoff {
67 fn new(initial: Duration, maximum: Duration) -> Self {
68 Self {
69 initial,
70 current: initial,
71 maximum,
72 }
73 }
74
75 fn current(&self) -> Duration {
76 self.current
77 }
78
79 fn progress(&mut self) {
80 self.current = self.initial;
81 }
82
83 fn no_progress(&mut self) {
84 self.current = self.current.saturating_mul(2).min(self.maximum);
85 }
86}
87
88fn final_handle_reaper() -> &'static mpsc::Sender<FinalHandleReap> {
89 static SENDER: OnceLock<mpsc::Sender<FinalHandleReap>> = OnceLock::new();
90 SENDER.get_or_init(|| {
91 let (sender, receiver) = mpsc::channel::<FinalHandleReap>();
92 let _ = std::thread::Builder::new()
93 .name("car-backend-handle-reaper".into())
94 .spawn(move || {
95 let mut pending = Vec::<FinalHandleReap>::new();
96 let mut backoff =
97 ReaperBackoff::new(Duration::from_millis(10), Duration::from_millis(250));
98 loop {
99 if pending.is_empty() {
100 match receiver.recv() {
101 Ok(item) => pending.push(item),
102 Err(_) => break,
103 }
104 } else {
105 match receiver.recv_timeout(backoff.current()) {
106 Ok(item) => pending.push(item),
107 Err(mpsc::RecvTimeoutError::Timeout) => {}
108 Err(mpsc::RecvTimeoutError::Disconnected) => {}
109 }
110 }
111 while let Ok(item) = receiver.try_recv() {
112 pending.push(item);
113 }
114 let before = pending.len();
115 let mut index = 0;
116 while index < pending.len() {
117 if (pending[index].released)() {
118 let item = pending.swap_remove(index);
119 (item.cleanup)();
120 } else {
121 index += 1;
122 }
123 }
124 if pending.len() < before {
125 backoff.progress();
126 } else {
127 backoff.no_progress();
128 }
129 }
130 });
131 sender
132 })
133}
134
135#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
139#[serde(rename_all = "snake_case")]
140pub enum BackendRetention {
141 Resident,
142 Transient,
143}
144
145pub struct SharedModelBudget {
158 budget_bytes: AtomicU64,
159 total_bytes: AtomicU64,
160}
161
162impl SharedModelBudget {
163 pub fn new(budget_bytes: u64) -> Arc<Self> {
165 Arc::new(Self {
166 budget_bytes: AtomicU64::new(budget_bytes),
167 total_bytes: AtomicU64::new(0),
168 })
169 }
170
171 pub fn from_env_or(default_mb: u64) -> Arc<Self> {
174 let mb = std::env::var("CAR_INFERENCE_MODEL_CACHE_MB")
175 .ok()
176 .and_then(|v| v.parse::<u64>().ok())
177 .unwrap_or(default_mb);
178 Self::new(mb.saturating_mul(1024 * 1024))
179 }
180
181 fn add(&self, n: u64) {
182 self.total_bytes.fetch_add(n, Ordering::Relaxed);
183 }
184 fn sub(&self, n: u64) {
185 let _ = self
188 .total_bytes
189 .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |cur| {
190 Some(cur.saturating_sub(n))
191 });
192 }
193 fn total(&self) -> u64 {
194 self.total_bytes.load(Ordering::Relaxed)
195 }
196 fn over_budget(&self) -> bool {
198 let budget = self.budget_bytes();
199 budget != 0 && self.total() > budget
200 }
201 pub fn is_disabled(&self) -> bool {
203 self.budget_bytes() == 0
204 }
205
206 pub fn set_budget_bytes(&self, budget_bytes: u64) {
209 self.budget_bytes.store(budget_bytes, Ordering::Release);
210 }
211
212 pub fn budget_bytes(&self) -> u64 {
213 self.budget_bytes.load(Ordering::Acquire)
214 }
215}
216
217pub fn default_model_cache_mb() -> u64 {
221 crate::hardware::HardwareInfo::detect().max_model_mb
222}
223
224pub fn idle_ttl_from_env() -> Option<Duration> {
226 let idle_secs = std::env::var("CAR_INFERENCE_MODEL_IDLE_SECS")
227 .ok()
228 .and_then(|v| v.parse::<u64>().ok())
229 .unwrap_or(300);
230 (idle_secs > 0).then(|| Duration::from_secs(idle_secs))
231}
232
233struct Entry<T> {
234 backend: CachedBackend<T>,
235 allocation_id: String,
236 size_bytes: u64,
237 last_used: Instant,
240 invalidated: bool,
243}
244
245struct Inner<T> {
246 map: HashMap<String, Entry<T>>,
248 lru: VecDeque<String>,
250}
251
252pub struct BackendCache<T: Send + 'static> {
258 inner: Mutex<Inner<T>>,
259 budget: Arc<SharedModelBudget>,
262 idle_ttl: Option<Duration>,
265 load_locks: Mutex<HashMap<String, Arc<Mutex<()>>>>,
269 resident_accounting: Option<Arc<crate::resource_policy::LocalAdmissionCoordinator>>,
270 allocation_scope: u64,
271}
272
273fn next_cache_allocation_scope() -> u64 {
274 static NEXT: AtomicU64 = AtomicU64::new(1);
275 NEXT.fetch_add(1, Ordering::Relaxed)
276}
277
278impl<T: Send + 'static> Drop for BackendCache<T> {
279 fn drop(&mut self) {
280 let inner = self
281 .inner
282 .get_mut()
283 .unwrap_or_else(std::sync::PoisonError::into_inner);
284 let idle = inner
289 .map
290 .iter()
291 .filter(|(_, entry)| Arc::strong_count(&entry.backend) == 1)
292 .map(|(_, entry)| (entry.allocation_id.clone(), entry.size_bytes))
293 .collect::<Vec<_>>();
294 let active = inner
295 .map
296 .iter()
297 .filter(|(_, entry)| Arc::strong_count(&entry.backend) > 1)
298 .map(|(_, entry)| {
299 (
300 entry.allocation_id.clone(),
301 entry.size_bytes,
302 Arc::downgrade(&entry.backend),
303 )
304 })
305 .collect::<Vec<_>>();
306 let bytes = idle
307 .iter()
308 .map(|(_, size_bytes)| *size_bytes)
309 .fold(0_u64, u64::saturating_add);
310 inner.lru.clear();
311 self.budget.sub(bytes);
312 if let Some(accounting) = &self.resident_accounting {
313 for (key, _) in idle {
314 accounting.mark_evicted(&key);
315 }
316 }
317
318 for (key, size_bytes, backend) in active {
322 let accounting = self.resident_accounting.clone();
323 let budget = self.budget.clone();
324 let item = FinalHandleReap {
325 released: Box::new(move || backend.strong_count() == 0),
326 cleanup: Box::new(move || {
327 budget.sub(size_bytes);
328 if let Some(accounting) = &accounting {
329 accounting.mark_evicted(&key);
330 }
331 }),
332 };
333 if let Err(error) = final_handle_reaper().send(item) {
334 std::mem::forget(error.0);
336 }
337 }
338 inner.map.clear();
339 }
340}
341
342impl<T: Send + 'static> BackendCache<T> {
343 pub fn new(budget_bytes: u64) -> Self {
346 Self::from_shared(SharedModelBudget::new(budget_bytes), None)
347 }
348
349 pub fn with_idle_ttl(budget_bytes: u64, idle_ttl: Option<Duration>) -> Self {
351 Self::from_shared(SharedModelBudget::new(budget_bytes), idle_ttl)
352 }
353
354 pub fn from_shared(budget: Arc<SharedModelBudget>, idle_ttl: Option<Duration>) -> Self {
357 Self::from_shared_with_admission(budget, idle_ttl, None)
358 }
359
360 pub fn from_shared_with_admission(
361 budget: Arc<SharedModelBudget>,
362 idle_ttl: Option<Duration>,
363 resident_accounting: Option<Arc<crate::resource_policy::LocalAdmissionCoordinator>>,
364 ) -> Self {
365 Self {
366 inner: Mutex::new(Inner {
367 map: HashMap::new(),
368 lru: VecDeque::new(),
369 }),
370 budget,
371 idle_ttl,
372 load_locks: Mutex::new(HashMap::new()),
373 resident_accounting,
374 allocation_scope: next_cache_allocation_scope(),
375 }
376 }
377
378 pub fn from_env() -> Self {
388 Self::from_shared(
389 SharedModelBudget::from_env_or(default_model_cache_mb()),
390 idle_ttl_from_env(),
391 )
392 }
393
394 pub fn is_disabled(&self) -> bool {
396 self.budget.is_disabled()
397 }
398
399 pub fn get_or_load<E>(
409 &self,
410 key: &str,
411 size_bytes: u64,
412 loader: impl FnOnce() -> Result<T, E>,
413 ) -> Result<CachedBackend<T>, E> {
414 self.get_or_load_with_publish(key, size_bytes, loader, |allocation_id| {
415 if let Some(accounting) = &self.resident_accounting {
416 accounting.mark_resident_allocation(
417 key,
418 allocation_id,
419 size_bytes.div_ceil(1024 * 1024),
420 );
421 }
422 Ok(())
423 })
424 .map(|(handle, _retained)| handle)
425 }
426
427 pub fn get_or_load_admitted(
431 &self,
432 key: &str,
433 size_bytes: u64,
434 reservation: &mut crate::resource_policy::LocalLoadReservation,
435 loader: impl FnOnce() -> Result<T, crate::InferenceError>,
436 ) -> Result<(CachedBackend<T>, BackendRetention), crate::InferenceError> {
437 let allocation_id = format!("cache:{}:{key}", self.allocation_scope);
438 reservation.bind_allocation_id(&allocation_id);
439 reservation
440 .reconcile_measured_weights(size_bytes)
441 .map_err(crate::InferenceError::from)?;
442 let (handle, retained) =
443 self.get_or_load_with_publish(key, size_bytes, loader, |allocation_id| {
444 reservation.publish_resident_weights_as(allocation_id, size_bytes);
445 Ok(())
446 })?;
447 if retained {
448 reservation.publish_resident_weights(size_bytes);
449 }
450 Ok((
451 handle,
452 if retained {
453 BackendRetention::Resident
454 } else {
455 BackendRetention::Transient
456 },
457 ))
458 }
459
460 fn get_or_load_with_publish<E>(
461 &self,
462 key: &str,
463 size_bytes: u64,
464 loader: impl FnOnce() -> Result<T, E>,
465 on_publish: impl FnOnce(&str) -> Result<(), E>,
466 ) -> Result<(CachedBackend<T>, bool), E> {
467 {
469 let mut guard = self.inner.lock().expect("backend cache poisoned");
470 let stale_is_idle = guard
471 .map
472 .get(key)
473 .is_some_and(|entry| entry.invalidated && Arc::strong_count(&entry.backend) == 1);
474 if stale_is_idle {
475 let stale = guard.map.remove(key).expect("entry checked above");
476 guard.lru.retain(|candidate| candidate != key);
477 self.budget.sub(stale.size_bytes);
478 if let Some(accounting) = &self.resident_accounting {
479 accounting.mark_evicted(&stale.allocation_id);
480 }
481 }
482 if let Some(entry) = guard.map.get_mut(key) {
483 let handle = Arc::clone(&entry.backend);
484 entry.last_used = Instant::now();
487 guard.lru.retain(|k| k != key);
488 guard.lru.push_back(key.to_string());
489 return Ok((handle, true));
490 }
491 }
492
493 let load_gate = {
494 let mut gates = self.load_locks.lock().expect("backend load gates poisoned");
495 Arc::clone(
496 gates
497 .entry(key.to_string())
498 .or_insert_with(|| Arc::new(Mutex::new(()))),
499 )
500 };
501 let _singleflight = load_gate.lock().expect("backend load gate poisoned");
502
503 {
505 let mut guard = self.inner.lock().expect("backend cache poisoned");
506 if let Some(entry) = guard.map.get_mut(key) {
507 let handle = Arc::clone(&entry.backend);
508 entry.last_used = Instant::now();
509 guard.lru.retain(|candidate| candidate != key);
510 guard.lru.push_back(key.to_string());
511 drop(_singleflight);
512 self.release_load_gate(key, &load_gate);
513 return Ok((handle, true));
514 }
515 }
516
517 let backend = match loader() {
519 Ok(backend) => backend,
520 Err(error) => {
521 drop(_singleflight);
522 self.release_load_gate(key, &load_gate);
523 return Err(error);
524 }
525 };
526 let handle = Arc::new(Mutex::new(backend));
527
528 if self.budget.is_disabled() {
529 drop(_singleflight);
531 self.release_load_gate(key, &load_gate);
532 return Ok((handle, false));
533 }
534
535 let mut guard = self.inner.lock().expect("backend cache poisoned");
536
537 if let Some(existing) = guard.map.get(key) {
539 let existing = Arc::clone(&existing.backend);
540 drop(guard);
541 drop(_singleflight);
542 self.release_load_gate(key, &load_gate);
543 return Ok((existing, true));
544 }
545
546 let allocation_id = format!("cache:{}:{key}", self.allocation_scope);
547 if let Err(error) = on_publish(&allocation_id) {
548 drop(guard);
549 drop(_singleflight);
550 self.release_load_gate(key, &load_gate);
551 return Err(error);
552 }
553 self.budget.add(size_bytes);
554 guard.map.insert(
555 key.to_string(),
556 Entry {
557 backend: Arc::clone(&handle),
558 allocation_id,
559 size_bytes,
560 last_used: Instant::now(),
561 invalidated: false,
562 },
563 );
564 guard.lru.push_back(key.to_string());
565
566 let mut examined = guard.lru.len();
570 while self.budget.over_budget() && examined > 0 {
571 examined -= 1;
572 let Some(victim_key) = guard.lru.pop_front() else {
573 break;
574 };
575 if victim_key == key {
576 guard.lru.push_front(victim_key);
581 break;
582 }
583 let victim_is_idle = guard
584 .map
585 .get(&victim_key)
586 .is_some_and(|entry| Arc::strong_count(&entry.backend) == 1);
587 if !victim_is_idle {
588 guard.lru.push_back(victim_key);
589 continue;
590 }
591 if let Some(victim) = guard.map.remove(&victim_key) {
592 self.budget.sub(victim.size_bytes);
593 if let Some(accounting) = &self.resident_accounting {
594 accounting.mark_evicted(&victim.allocation_id);
595 }
596 drop(victim);
600 }
601 }
602
603 drop(guard);
604 drop(_singleflight);
605 self.release_load_gate(key, &load_gate);
606 Ok((handle, true))
607 }
608
609 fn release_load_gate(&self, key: &str, gate: &Arc<Mutex<()>>) {
610 let mut gates = self.load_locks.lock().expect("backend load gates poisoned");
611 if gates
612 .get(key)
613 .is_some_and(|current| Arc::ptr_eq(current, gate) && Arc::strong_count(current) == 2)
614 {
615 gates.remove(key);
616 }
617 }
618
619 pub fn invalidate(&self, key: &str) {
623 let mut guard = self.inner.lock().expect("backend cache poisoned");
624 let idle = guard
625 .map
626 .get(key)
627 .is_some_and(|entry| Arc::strong_count(&entry.backend) == 1);
628 if idle {
629 let Some(entry) = guard.map.remove(key) else {
630 return;
631 };
632 self.budget.sub(entry.size_bytes);
633 if let Some(accounting) = &self.resident_accounting {
634 accounting.mark_evicted(&entry.allocation_id);
635 }
636 guard.lru.retain(|k| k != key);
637 } else if let Some(entry) = guard.map.get_mut(key) {
638 entry.invalidated = true;
639 }
640 }
641
642 pub fn contains(&self, key: &str) -> bool {
643 self.inner
644 .lock()
645 .expect("backend cache poisoned")
646 .map
647 .contains_key(key)
648 }
649
650 pub fn evict_if_idle(&self, key: &str) -> bool {
654 let mut guard = self.inner.lock().expect("backend cache poisoned");
655 let Some(entry) = guard.map.get(key) else {
656 return true;
657 };
658 if Arc::strong_count(&entry.backend) != 1 {
659 return false;
660 }
661 let entry = guard.map.remove(key).expect("entry checked above");
662 guard.lru.retain(|candidate| candidate != key);
663 self.budget.sub(entry.size_bytes);
664 if let Some(accounting) = &self.resident_accounting {
665 accounting.mark_evicted(&entry.allocation_id);
666 }
667 true
668 }
669
670 pub fn enforce_budget(&self) -> (usize, u64) {
674 let mut guard = self.inner.lock().expect("backend cache poisoned");
675 let mut evicted = 0usize;
676 let mut bytes = 0u64;
677 let mut examined = guard.lru.len();
678 while (self.budget.over_budget() || self.budget.is_disabled()) && examined > 0 {
679 examined -= 1;
680 let Some(key) = guard.lru.pop_front() else {
681 break;
682 };
683 let is_idle = guard
684 .map
685 .get(&key)
686 .is_some_and(|entry| Arc::strong_count(&entry.backend) == 1);
687 if !is_idle {
688 guard.lru.push_back(key);
689 continue;
690 }
691 if let Some(entry) = guard.map.remove(&key) {
692 self.budget.sub(entry.size_bytes);
693 if let Some(accounting) = &self.resident_accounting {
694 accounting.mark_evicted(&entry.allocation_id);
695 }
696 evicted += 1;
697 bytes = bytes.saturating_add(entry.size_bytes);
698 }
699 }
700 (evicted, bytes)
701 }
702
703 pub fn evict_idle(&self) -> (usize, u64) {
713 let Some(ttl) = self.idle_ttl else {
714 return (0, 0);
715 };
716 let now = Instant::now();
717 let mut guard = self.inner.lock().expect("backend cache poisoned");
718 let stale: Vec<String> = guard
719 .map
720 .iter()
721 .filter(|(_, e)| now.duration_since(e.last_used) >= ttl)
722 .map(|(k, _)| k.clone())
723 .collect();
724 let mut entries = 0usize;
725 let mut bytes = 0u64;
726 for key in stale {
727 let is_idle = guard
728 .map
729 .get(&key)
730 .is_some_and(|entry| Arc::strong_count(&entry.backend) == 1);
731 if !is_idle {
732 continue;
733 }
734 if let Some(victim) = guard.map.remove(&key) {
735 self.budget.sub(victim.size_bytes);
736 if let Some(accounting) = &self.resident_accounting {
737 accounting.mark_evicted(&victim.allocation_id);
738 }
739 guard.lru.retain(|k| k != &key);
740 entries += 1;
741 bytes = bytes.saturating_add(victim.size_bytes);
742 drop(victim);
743 }
744 }
745 (entries, bytes)
746 }
747
748 pub fn clear(&self) {
750 let mut guard = self.inner.lock().expect("backend cache poisoned");
751 let idle_keys = guard
754 .map
755 .iter()
756 .filter(|(_, entry)| Arc::strong_count(&entry.backend) == 1)
757 .map(|(key, _)| key.clone())
758 .collect::<Vec<_>>();
759 let mut freed = 0u64;
760 for key in idle_keys {
761 if let Some(entry) = guard.map.remove(&key) {
762 freed = freed.saturating_add(entry.size_bytes);
763 if let Some(accounting) = &self.resident_accounting {
764 accounting.mark_evicted(&entry.allocation_id);
765 }
766 }
767 guard.lru.retain(|candidate| candidate != &key);
768 }
769 self.budget.sub(freed);
770 }
771
772 pub fn stats(&self) -> (usize, u64, u64) {
776 let guard = self.inner.lock().expect("backend cache poisoned");
777 (
778 guard.map.len(),
779 self.budget.total(),
780 self.budget.budget_bytes(),
781 )
782 }
783}
784
785pub fn estimate_model_size(model_dir: &Path) -> u64 {
791 fn visit(dir: &Path, total: &mut u64) {
792 if let Ok(meta) = dir.metadata() {
796 if meta.is_file() {
797 *total = total.saturating_add(meta.len());
798 return;
799 }
800 }
801 let Ok(entries) = std::fs::read_dir(dir) else {
802 return;
803 };
804 for entry in entries.flatten() {
805 let path = entry.path();
806 let Ok(link_meta) = path.symlink_metadata() else {
807 continue;
808 };
809 if link_meta.file_type().is_symlink() {
810 if let Ok(target_meta) = path.metadata() {
811 if target_meta.is_file() {
812 *total = total.saturating_add(target_meta.len());
813 }
814 }
815 continue;
816 }
817 if link_meta.is_dir() {
818 visit(&path, total);
819 continue;
820 }
821 if link_meta.is_file() {
822 *total = total.saturating_add(link_meta.len());
823 }
824 }
825 }
826 let mut total = 0u64;
827 visit(model_dir, &mut total);
828 total
829}
830
831#[cfg(test)]
832mod tests {
833 use super::*;
834
835 struct FixedProbe;
836
837 impl crate::resource_policy::LiveMemoryProbe for FixedProbe {
838 fn available_memory_mb(
839 &self,
840 ) -> Result<Option<u64>, crate::resource_policy::ResourcePolicyError> {
841 Ok(Some(24_000))
842 }
843 }
844
845 #[test]
846 fn final_handle_reaper_backoff_is_bounded_and_resets_after_progress() {
847 let mut backoff = ReaperBackoff::new(Duration::from_millis(10), Duration::from_millis(250));
848 assert_eq!(backoff.current(), Duration::from_millis(10));
849 for _ in 0..10 {
850 backoff.no_progress();
851 }
852 assert_eq!(backoff.current(), Duration::from_millis(250));
853 backoff.progress();
854 assert_eq!(backoff.current(), Duration::from_millis(10));
855 }
856
857 #[test]
858 fn cache_hit_returns_same_handle() {
859 let cache: BackendCache<u32> = BackendCache::new(1024);
860 let a = cache.get_or_load::<()>("a", 100, || Ok(42)).unwrap();
861 let b = cache
862 .get_or_load::<()>("a", 100, || panic!("should not reload"))
863 .unwrap();
864 assert!(Arc::ptr_eq(&a, &b));
865 }
866
867 #[test]
868 fn evicts_lru_when_over_budget() {
869 let cache: BackendCache<u32> = BackendCache::new(250);
870 let _a = cache.get_or_load::<()>("a", 100, || Ok(1)).unwrap();
871 let b = cache.get_or_load::<()>("b", 100, || Ok(2)).unwrap();
872 let _a_again = cache
874 .get_or_load::<()>("a", 100, || panic!("cached"))
875 .unwrap();
876 drop(b);
877 let _c = cache.get_or_load::<()>("c", 100, || Ok(3)).unwrap();
879 let (n, bytes, budget) = cache.stats();
880 assert_eq!(n, 2, "a + c should remain, b evicted");
881 assert_eq!(bytes, 200);
882 assert_eq!(budget, 250);
883 }
884
885 #[test]
886 fn zero_budget_disables_cache_but_returns_handle() {
887 let cache: BackendCache<u32> = BackendCache::new(0);
888 let mut load_count = 0u32;
889 let a = cache
890 .get_or_load::<()>("a", 100, || {
891 load_count += 1;
892 Ok(1)
893 })
894 .unwrap();
895 assert_eq!(*a.lock().unwrap(), 1);
896 let b = cache
897 .get_or_load::<()>("a", 100, || {
898 load_count += 1;
899 Ok(1)
900 })
901 .unwrap();
902 assert_eq!(*b.lock().unwrap(), 1);
903 assert_eq!(load_count, 2, "disabled cache reloads every call");
904 assert!(!Arc::ptr_eq(&a, &b));
905 }
906
907 #[test]
908 fn evict_idle_drops_stale_entries_below_capacity() {
909 let cache: BackendCache<u32> =
912 BackendCache::with_idle_ttl(1_000_000, Some(Duration::from_millis(20)));
913 let _a = cache.get_or_load::<()>("a", 100, || Ok(1)).unwrap();
914 let b = cache.get_or_load::<()>("b", 100, || Ok(2)).unwrap();
915 assert_eq!(cache.stats().0, 2);
916 assert_eq!(cache.evict_idle(), (0, 0));
918 std::thread::sleep(Duration::from_millis(40));
919 drop(b);
920 let _a_again = cache
922 .get_or_load::<()>("a", 100, || panic!("cached"))
923 .unwrap();
924 let (entries, bytes) = cache.evict_idle();
925 assert_eq!((entries, bytes), (1, 100), "only b should be swept");
926 let (n, total, _) = cache.stats();
927 assert_eq!(n, 1, "a remains");
928 assert_eq!(total, 100);
929 }
930
931 #[test]
932 fn evict_idle_noop_when_disabled() {
933 let cache: BackendCache<u32> = BackendCache::new(1024); let _a = cache.get_or_load::<()>("a", 100, || Ok(1)).unwrap();
935 assert_eq!(cache.evict_idle(), (0, 0));
936 assert_eq!(cache.stats().0, 1, "disabled idle eviction keeps the entry");
937 }
938
939 #[test]
940 fn invalidate_removes_key() {
941 let cache: BackendCache<u32> = BackendCache::new(1024);
942 let a = cache.get_or_load::<()>("a", 100, || Ok(1)).unwrap();
943 assert_eq!(cache.stats().0, 1);
944 drop(a);
945 cache.invalidate("a");
946 assert_eq!(cache.stats().0, 0);
947 }
948
949 #[test]
950 fn active_invalidation_is_deferred_without_losing_resident_accounting() {
951 let cache: BackendCache<u32> = BackendCache::new(1024);
952 let active = cache.get_or_load::<()>("a", 100, || Ok(1)).unwrap();
953
954 cache.invalidate("a");
955 assert_eq!(cache.stats().1, 100, "active weights remain accounted");
956 drop(active);
957
958 let replacement = cache.get_or_load::<()>("a", 100, || Ok(2)).unwrap();
959 assert_eq!(*replacement.lock().unwrap(), 2);
960 assert_eq!(cache.stats().1, 100, "replacement is not double-counted");
961 }
962
963 #[test]
964 fn installed_size_measurement_covers_single_files_and_non_safetensor_weights() {
965 let dir = tempfile::tempdir().unwrap();
966 std::fs::write(dir.path().join("model.gguf"), vec![0_u8; 17]).unwrap();
967 std::fs::write(dir.path().join("encoder.onnx"), vec![0_u8; 23]).unwrap();
968 let single = dir.path().join("whisper.bin");
969 std::fs::write(&single, vec![0_u8; 31]).unwrap();
970
971 assert_eq!(estimate_model_size(dir.path()), 71);
972 assert_eq!(estimate_model_size(&single), 31);
973 }
974
975 #[cfg(unix)]
976 #[test]
977 fn installed_size_measurement_follows_hf_leaf_symlinks_but_not_directory_symlinks() {
978 use std::os::unix::fs::symlink;
979
980 let dir = tempfile::tempdir().unwrap();
981 let blobs = dir.path().join("blobs");
982 let snapshot = dir.path().join("snapshots/revision");
983 std::fs::create_dir_all(&blobs).unwrap();
984 std::fs::create_dir_all(&snapshot).unwrap();
985 std::fs::write(blobs.join("weights"), vec![0_u8; 47]).unwrap();
986 symlink("../../blobs/weights", snapshot.join("model.safetensors")).unwrap();
987 symlink("../../blobs", snapshot.join("directory-link")).unwrap();
988 symlink("../../blobs/missing", snapshot.join("dangling")).unwrap();
989
990 assert_eq!(estimate_model_size(&snapshot), 47);
991 }
992
993 #[test]
994 fn separate_cache_generations_keep_exact_resident_owners_until_final_handle_drop() {
995 let coordinator = Arc::new(crate::resource_policy::LocalAdmissionCoordinator::new(
996 crate::ResourcePolicy::custom_gb(8.0).unwrap(),
997 crate::hardware::HardwareInfo::detect(),
998 ));
999 let first: BackendCache<u32> = BackendCache::from_shared_with_admission(
1000 SharedModelBudget::new(1024),
1001 None,
1002 Some(coordinator.clone()),
1003 );
1004 let second: BackendCache<u32> = BackendCache::from_shared_with_admission(
1005 SharedModelBudget::new(1024),
1006 None,
1007 Some(coordinator.clone()),
1008 );
1009 let mut first_reservation = coordinator
1010 .reserve_measured_host("same-model", 100, 0)
1011 .unwrap();
1012 let (first_handle, _) = first
1013 .get_or_load_admitted("same-model", 100, &mut first_reservation, || {
1014 Ok::<_, crate::InferenceError>(1)
1015 })
1016 .unwrap();
1017 drop(first_reservation);
1018 let mut second_reservation = coordinator
1019 .reserve_measured_host("same-model", 100, 0)
1020 .unwrap();
1021 let (_second_handle, _) = second
1022 .get_or_load_admitted("same-model", 100, &mut second_reservation, || {
1023 Ok::<_, crate::InferenceError>(2)
1024 })
1025 .unwrap();
1026 drop(second_reservation);
1027
1028 assert_eq!(coordinator.resident_allocation_ids("same-model").len(), 2);
1029 drop(first);
1030 assert_eq!(coordinator.resident_allocation_ids("same-model").len(), 2);
1031 drop(first_handle);
1032 for _ in 0..100 {
1033 if coordinator.resident_allocation_ids("same-model").len() == 1 {
1034 break;
1035 }
1036 std::thread::sleep(Duration::from_millis(5));
1037 }
1038 assert_eq!(coordinator.resident_allocation_ids("same-model").len(), 1);
1039 }
1040
1041 #[test]
1042 fn shared_budget_spans_caches_and_each_self_trims() {
1043 let budget = SharedModelBudget::new(250);
1045 let a: BackendCache<u32> = BackendCache::from_shared(budget.clone(), None);
1046 let b: BackendCache<u32> = BackendCache::from_shared(budget.clone(), None);
1047
1048 let a1 = a.get_or_load::<()>("a1", 100, || Ok(1)).unwrap();
1049 let _ = b.get_or_load::<()>("b1", 100, || Ok(2)).unwrap();
1050 assert_eq!(budget.total(), 200);
1052
1053 drop(a1);
1056 let _ = a.get_or_load::<()>("a2", 100, || Ok(3)).unwrap(); assert_eq!(budget.total(), 200, "A self-trimmed to the shared budget");
1058 assert_eq!(a.stats().0, 1, "A kept a2, evicted a1");
1059 assert_eq!(b.stats().0, 1, "B's b1 not evicted by A");
1060 }
1061
1062 #[test]
1063 fn default_budget_is_ram_derived_not_flat() {
1064 let mb = default_model_cache_mb();
1067 assert!(mb > 0, "RAM-derived default should be positive");
1068 }
1069
1070 #[test]
1071 fn backend_cache_policy_decrease_evicts_idle_entries() {
1072 let budget = SharedModelBudget::new(1_000);
1073 let cache: BackendCache<u32> = BackendCache::from_shared(budget.clone(), None);
1074 let handle = cache.get_or_load::<()>("model", 600, || Ok(1)).unwrap();
1075 drop(handle);
1076
1077 budget.set_budget_bytes(500);
1078 assert_eq!(cache.enforce_budget(), (1, 600));
1079 assert_eq!(cache.stats(), (0, 0, 500));
1080 }
1081
1082 #[test]
1083 fn backend_cache_policy_decrease_does_not_kill_active_handle() {
1084 let budget = SharedModelBudget::new(1_000);
1085 let cache: BackendCache<u32> = BackendCache::from_shared(budget.clone(), None);
1086 let handle = cache.get_or_load::<()>("model", 600, || Ok(1)).unwrap();
1087
1088 budget.set_budget_bytes(0);
1089 assert_eq!(cache.enforce_budget(), (0, 0));
1090 assert!(cache.contains("model"));
1091 drop(handle);
1092 assert_eq!(cache.enforce_budget(), (1, 600));
1093 }
1094
1095 #[test]
1096 fn backend_cache_same_key_load_is_singleflight() {
1097 let cache = Arc::new(BackendCache::new(1_000));
1098 let loads = Arc::new(AtomicU64::new(0));
1099 let mut threads = Vec::new();
1100 for _ in 0..8 {
1101 let cache = Arc::clone(&cache);
1102 let loads = Arc::clone(&loads);
1103 threads.push(std::thread::spawn(move || {
1104 cache
1105 .get_or_load::<()>("model", 100, || {
1106 loads.fetch_add(1, Ordering::SeqCst);
1107 std::thread::sleep(Duration::from_millis(5));
1108 Ok(42_u32)
1109 })
1110 .unwrap()
1111 }));
1112 }
1113 let handles = threads
1114 .into_iter()
1115 .map(|thread| thread.join().unwrap())
1116 .collect::<Vec<_>>();
1117 assert_eq!(loads.load(Ordering::SeqCst), 1);
1118 assert!(handles
1119 .windows(2)
1120 .all(|pair| Arc::ptr_eq(&pair[0], &pair[1])));
1121 }
1122
1123 #[test]
1124 fn backend_cache_eviction_updates_admission_residency() {
1125 let hw = crate::hardware::HardwareInfo {
1126 os: "test".into(),
1127 arch: "test".into(),
1128 cpu_cores: 8,
1129 total_ram_mb: 32 * 1024,
1130 gpu_backend: crate::hardware::GpuBackend::Metal,
1131 gpu_memory_mb: None,
1132 gpu_devices: Vec::new(),
1133 recommended_model: "fixture".into(),
1134 recommended_context: 4_096,
1135 max_model_mb: 32 * 1024,
1136 };
1137 let admission = Arc::new(crate::resource_policy::LocalAdmissionCoordinator::new(
1138 crate::resource_policy::ResourcePolicy::everyday(),
1139 hw,
1140 ));
1141 let budget = SharedModelBudget::new(2 * 1024 * 1024);
1142 let cache: BackendCache<u32> =
1143 BackendCache::from_shared_with_admission(budget.clone(), None, Some(admission.clone()));
1144 let handle = cache
1145 .get_or_load::<()>("model", 1024 * 1024, || Ok(1))
1146 .unwrap();
1147 assert_eq!(admission.resident_model_mb(), 1);
1148 drop(handle);
1149
1150 budget.set_budget_bytes(0);
1151 assert_eq!(cache.enforce_budget().0, 1);
1152 assert_eq!(admission.resident_model_mb(), 0);
1153 }
1154
1155 #[test]
1156 fn backend_cache_drop_never_peer_discounts_an_outstanding_handle() {
1157 let admission = Arc::new(crate::resource_policy::LocalAdmissionCoordinator::new(
1158 crate::resource_policy::ResourcePolicy::custom_gb(8.0).unwrap(),
1159 crate::hardware::HardwareInfo {
1160 total_ram_mb: 32 * 1024,
1161 ..crate::hardware::HardwareInfo::detect()
1162 },
1163 ));
1164 let budget = SharedModelBudget::new(8 * 1024 * 1024);
1165 let cache: BackendCache<u32> =
1166 BackendCache::from_shared_with_admission(budget.clone(), None, Some(admission.clone()));
1167 let active = cache
1168 .get_or_load::<()>("model", 1024 * 1024, || Ok(1))
1169 .unwrap();
1170 assert_eq!(admission.resident_model_mb(), 1);
1171
1172 drop(cache);
1173 assert_eq!(
1174 admission.resident_model_mb(),
1175 1,
1176 "dropping a cache must not call an outstanding allocation safe"
1177 );
1178 assert_eq!(budget.total(), 1024 * 1024);
1179 drop(active);
1180 for _ in 0..100 {
1181 if admission.resident_model_mb() == 0 && budget.total() == 0 {
1182 break;
1183 }
1184 std::thread::sleep(Duration::from_millis(5));
1185 }
1186 assert_eq!(
1187 admission.resident_model_mb(),
1188 0,
1189 "the final external handle drop must release resident accounting"
1190 );
1191 assert_eq!(budget.total(), 0);
1192 }
1193
1194 #[test]
1195 fn backend_cache_targeted_eviction_refuses_active_then_releases_idle() {
1196 let cache: BackendCache<u32> = BackendCache::new(1_000);
1197 let handle = cache.get_or_load::<()>("model", 100, || Ok(1)).unwrap();
1198 assert!(!cache.evict_if_idle("model"));
1199 assert!(cache.contains("model"));
1200 drop(handle);
1201 assert!(cache.evict_if_idle("model"));
1202 assert!(!cache.contains("model"));
1203 }
1204
1205 #[test]
1206 fn admitted_cache_publication_atomically_transfers_residency() {
1207 let hw = crate::hardware::HardwareInfo {
1208 os: "test".into(),
1209 arch: "test".into(),
1210 cpu_cores: 8,
1211 total_ram_mb: 32 * 1024,
1212 gpu_backend: crate::hardware::GpuBackend::Metal,
1213 gpu_memory_mb: None,
1214 gpu_devices: Vec::new(),
1215 recommended_model: "fixture".into(),
1216 recommended_context: 4_096,
1217 max_model_mb: 32 * 1024,
1218 };
1219 let admission = Arc::new(
1220 crate::resource_policy::LocalAdmissionCoordinator::with_probe(
1221 crate::resource_policy::ResourcePolicy::local_focused(),
1222 hw,
1223 Arc::new(FixedProbe),
1224 ),
1225 );
1226 let model = crate::registry::builtin_catalog()
1227 .into_iter()
1228 .find(|model| model.id == "mlx/qwen3-4b:4bit")
1229 .unwrap();
1230 let budget = SharedModelBudget::new(8 * 1024 * 1024 * 1024);
1231 let cache: BackendCache<u32> =
1232 BackendCache::from_shared_with_admission(budget, None, Some(admission.clone()));
1233 let mut reservation = admission.reserve(&model, 2_048).unwrap();
1234
1235 let (handle, retention) = cache
1236 .get_or_load_admitted(&model.id, 3 * 1024 * 1024 * 1024, &mut reservation, || {
1237 Ok::<_, crate::InferenceError>(42)
1238 })
1239 .unwrap();
1240
1241 assert_eq!(*handle.lock().unwrap(), 42);
1242 assert_eq!(retention, BackendRetention::Resident);
1243 assert_eq!(admission.resident_model_mb(), 3 * 1024);
1244 assert_eq!(
1245 admission.preflight(&model, 2_048).active_reservations_mb,
1246 reservation.reserved_incremental_mb()
1247 );
1248 }
1249
1250 #[test]
1251 fn admitted_zero_cache_success_is_explicitly_transient_and_not_resident() {
1252 let admission = Arc::new(
1253 crate::resource_policy::LocalAdmissionCoordinator::with_probe(
1254 crate::resource_policy::ResourcePolicy::custom_gb(8.0).unwrap(),
1255 crate::hardware::HardwareInfo {
1256 total_ram_mb: 32 * 1024,
1257 ..crate::hardware::HardwareInfo::detect()
1258 },
1259 Arc::new(FixedProbe),
1260 ),
1261 );
1262 let model = crate::registry::builtin_catalog()
1263 .into_iter()
1264 .find(|model| model.id == "mlx/qwen3-4b:4bit")
1265 .unwrap();
1266 let cache: BackendCache<u32> = BackendCache::from_shared_with_admission(
1267 SharedModelBudget::new(0),
1268 None,
1269 Some(admission.clone()),
1270 );
1271 let mut reservation = admission.reserve(&model, 64).unwrap();
1272 let (handle, retention) = cache
1273 .get_or_load_admitted(&model.id, 1024 * 1024, &mut reservation, || {
1274 Ok::<_, crate::InferenceError>(42)
1275 })
1276 .unwrap();
1277
1278 assert_eq!(*handle.lock().unwrap(), 42);
1279 assert_eq!(retention, BackendRetention::Transient);
1280 drop(handle);
1281 drop(reservation);
1282 assert!(!admission.is_resident(&model.id));
1283 }
1284}