#[derive(Debug, Clone)]
pub struct CacheConfig {
pub l1_max_bytes: usize,
pub l2_max_bytes: usize,
pub eviction_policy: EvictionPolicy,
pub default_ttl: Option<Duration>,
pub prefetch_enabled: bool,
}
impl Default for CacheConfig {
fn default() -> Self {
Self {
l1_max_bytes: 64 * 1024 * 1024, l2_max_bytes: 1024 * 1024 * 1024, eviction_policy: EvictionPolicy::LRU,
default_ttl: None,
prefetch_enabled: true,
}
}
}
impl CacheConfig {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn embedded(l1_bytes: usize) -> Self {
Self {
l1_max_bytes: l1_bytes,
l2_max_bytes: 0, eviction_policy: EvictionPolicy::Fixed,
default_ttl: None,
prefetch_enabled: false,
}
}
#[must_use]
pub fn with_l1_size(mut self, bytes: usize) -> Self {
self.l1_max_bytes = bytes;
self
}
#[must_use]
pub fn with_l2_size(mut self, bytes: usize) -> Self {
self.l2_max_bytes = bytes;
self
}
#[must_use]
pub fn with_eviction_policy(mut self, policy: EvictionPolicy) -> Self {
self.eviction_policy = policy;
self
}
#[must_use]
pub fn with_ttl(mut self, ttl: Duration) -> Self {
self.default_ttl = Some(ttl);
self
}
#[must_use]
pub fn with_prefetch(mut self, enabled: bool) -> Self {
self.prefetch_enabled = enabled;
self
}
}
#[derive(Debug, Clone)]
pub struct ModelInfo {
pub name: String,
pub model_type: ModelType,
pub size_bytes: usize,
pub is_bundled: bool,
pub is_cached: bool,
pub cache_tier: Option<CacheTier>,
}
#[derive(Debug)]
pub struct ModelRegistry {
l1_cache: HashMap<String, CacheEntry>,
l2_cache: HashMap<String, CacheEntry>,
config: CacheConfig,
l1_current_bytes: usize,
l2_current_bytes: usize,
access_counter: u64,
created_at: Instant,
}
impl ModelRegistry {
#[must_use]
pub fn new(config: CacheConfig) -> Self {
Self {
l1_cache: HashMap::new(),
l2_cache: HashMap::new(),
config,
l1_current_bytes: 0,
l2_current_bytes: 0,
access_counter: 0,
created_at: Instant::now(),
}
}
pub fn insert_l1(&mut self, name: String, entry: CacheEntry) {
let size = entry.data.size();
while self.l1_current_bytes + size > self.config.l1_max_bytes && !self.l1_cache.is_empty() {
if let Some(evict_key) = self.find_eviction_candidate_l1() {
self.evict_l1(&evict_key);
} else {
break;
}
}
if let Some(old) = self.l1_cache.insert(name, entry) {
self.l1_current_bytes -= old.data.size();
}
self.l1_current_bytes += size;
}
pub fn insert_l2(&mut self, name: String, entry: CacheEntry) {
let size = entry.data.size();
while self.l2_current_bytes + size > self.config.l2_max_bytes && !self.l2_cache.is_empty() {
if let Some(evict_key) = self.find_eviction_candidate_l2() {
self.evict_l2(&evict_key);
} else {
break;
}
}
if let Some(old) = self.l2_cache.insert(name, entry) {
self.l2_current_bytes -= old.data.size();
}
self.l2_current_bytes += size;
}
pub fn get(&mut self, name: &str) -> Option<&CacheEntry> {
self.access_counter += 1;
let timestamp = self.access_counter;
if let Some(entry) = self.l1_cache.get_mut(name) {
entry.stats.record_hit(100, timestamp);
return self.l1_cache.get(name);
}
if let Some(entry) = self.l2_cache.get_mut(name) {
entry.stats.record_hit(1000, timestamp);
return self.l2_cache.get(name);
}
None
}
#[must_use]
pub fn contains(&self, name: &str) -> bool {
self.l1_cache.contains_key(name) || self.l2_cache.contains_key(name)
}
#[must_use]
pub fn get_tier(&self, name: &str) -> Option<CacheTier> {
if self.l1_cache.contains_key(name) {
Some(CacheTier::L1Hot)
} else if self.l2_cache.contains_key(name) {
Some(CacheTier::L2Warm)
} else {
None
}
}
pub fn remove(&mut self, name: &str) {
if let Some(entry) = self.l1_cache.remove(name) {
self.l1_current_bytes -= entry.data.size();
}
if let Some(entry) = self.l2_cache.remove(name) {
self.l2_current_bytes -= entry.data.size();
}
}
pub fn clear(&mut self) {
self.l1_cache.clear();
self.l2_cache.clear();
self.l1_current_bytes = 0;
self.l2_current_bytes = 0;
}
#[must_use]
pub fn stats(&self) -> CacheStats {
let l1_entries = self.l1_cache.len();
let l2_entries = self.l2_cache.len();
let (l1_hits, l1_misses) = self.l1_cache.values().fold((0, 0), |(h, m), e| {
(h + e.stats.hit_count, m + e.stats.miss_count)
});
let (l2_hits, l2_misses) = self.l2_cache.values().fold((0, 0), |(h, m), e| {
(h + e.stats.hit_count, m + e.stats.miss_count)
});
CacheStats {
l1_entries,
l1_bytes: self.l1_current_bytes,
l1_hits,
l1_misses,
l2_entries,
l2_bytes: self.l2_current_bytes,
l2_hits,
l2_misses,
uptime: self.created_at.elapsed(),
}
}
#[must_use]
pub fn list(&self) -> Vec<ModelInfo> {
let mut models = Vec::new();
for (name, entry) in &self.l1_cache {
models.push(ModelInfo {
name: name.clone(),
model_type: entry.model_type,
size_bytes: entry.data.size(),
is_bundled: false,
is_cached: true,
cache_tier: Some(CacheTier::L1Hot),
});
}
for (name, entry) in &self.l2_cache {
if !self.l1_cache.contains_key(name) {
models.push(ModelInfo {
name: name.clone(),
model_type: entry.model_type,
size_bytes: entry.data.size(),
is_bundled: false,
is_cached: true,
cache_tier: Some(CacheTier::L2Warm),
});
}
}
models
}
fn find_eviction_candidate_l1(&self) -> Option<String> {
match self.config.eviction_policy {
EvictionPolicy::Fixed => None,
EvictionPolicy::LRU | EvictionPolicy::Clock => self
.l1_cache
.iter()
.min_by_key(|(_, e)| e.stats.last_access)
.map(|(k, _)| k.clone()),
EvictionPolicy::LFU => self
.l1_cache
.iter()
.min_by_key(|(_, e)| e.stats.hit_count)
.map(|(k, _)| k.clone()),
EvictionPolicy::ARC => {
self.l1_cache
.iter()
.min_by_key(|(_, e)| {
e.stats.last_access.saturating_add(e.stats.hit_count * 100)
})
.map(|(k, _)| k.clone())
}
}
}
fn find_eviction_candidate_l2(&self) -> Option<String> {
match self.config.eviction_policy {
EvictionPolicy::Fixed => None,
_ => self
.l2_cache
.iter()
.min_by_key(|(_, e)| e.stats.last_access)
.map(|(k, _)| k.clone()),
}
}
fn evict_l1(&mut self, key: &str) {
if let Some(entry) = self.l1_cache.remove(key) {
self.l1_current_bytes -= entry.data.size();
}
}
fn evict_l2(&mut self, key: &str) {
if let Some(entry) = self.l2_cache.remove(key) {
self.l2_current_bytes -= entry.data.size();
}
}
}
#[derive(Debug, Clone)]
pub struct CacheStats {
pub l1_entries: usize,
pub l1_bytes: usize,
pub l1_hits: u64,
pub l1_misses: u64,
pub l2_entries: usize,
pub l2_bytes: usize,
pub l2_hits: u64,
pub l2_misses: u64,
pub uptime: Duration,
}
impl CacheStats {
#[must_use]
pub fn hit_rate(&self) -> f64 {
let total_hits = self.l1_hits + self.l2_hits;
let total_misses = self.l1_misses + self.l2_misses;
let total = total_hits + total_misses;
if total == 0 {
0.0
} else {
total_hits as f64 / total as f64
}
}
#[must_use]
pub fn total_bytes(&self) -> usize {
self.l1_bytes + self.l2_bytes
}
#[must_use]
pub fn total_entries(&self) -> usize {
self.l1_entries + self.l2_entries
}
}
#[cfg(test)]
mod tests;