RustProxy/src/cache/lifecycle.rs

impl Cache {
    pub fn new(revision: u64) -> Arc<Self> {
        Arc::new(Self {
            store: Arc::new(Mutex::new(Store {
                revision,
                ..Default::default()
            })),
            revision: AtomicU64::new(revision),
            domain_stats: DomainStatsStore::default(),
            cold: Arc::new(Mutex::new(ColdStore::default())),
            tier_settings: Mutex::new(Settings::default()),
            scope_revisions: Mutex::new(HashMap::new()),
            data_root: None,
            cold_tx: None,
            cold_delete_fallbacks: Arc::new(Semaphore::new(COLD_DELETE_FALLBACKS)),
            cleanup_slots: Arc::new(Semaphore::new(COLD_CLEANUP_SLOTS)),
            fills: Arc::new(Semaphore::new(FILL_SLOTS)),
            fill_bytes_active: AtomicU64::new(0),
            transient_bytes_active: Arc::new(AtomicU64::new(0)),
            cold_queue_items: AtomicU64::new(0),
            cold_queue_bytes: AtomicU64::new(0),
            cold_write_bytes_active: AtomicU64::new(0),
            cold_read_bytes_active: AtomicU64::new(0),
            revalidating: Mutex::new(HashSet::new()),
            fill_flights: Mutex::new(HashMap::new()),
            uncacheable_until: Mutex::new(HashMap::new()),
            fill_flight_sequence: AtomicU64::new(0),
            collapse_waiters: AtomicU64::new(0),
            bans: Mutex::new(Vec::new()),
            ban_rule_count: AtomicU64::new(0),
            ban_sequence: AtomicU64::new(0),
            proactive_scheduled_keys: Mutex::new(HashSet::new()),
            proactive_inflight: AtomicU64::new(0),
            cleanup_jobs: Mutex::new(HashMap::new()),
            cleanup_sequence: AtomicU64::new(0),
            counters: Counters::default(),
        })
    }
    pub fn start(
        revision: u64,
        data_root: Option<PathBuf>,
        settings: Settings,
        shutdown: CancellationToken,
    ) -> Arc<Self> {
        Self::start_with_scope_revisions(revision, data_root, settings, shutdown, HashMap::new())
    }

    pub fn start_with_scope_revisions(
        revision: u64,
        data_root: Option<PathBuf>,
        settings: Settings,
        shutdown: CancellationToken,
        scope_revisions: HashMap<String, u64>,
    ) -> Arc<Self> {
        let (tx, rx) = mpsc::channel(COLD_QUEUE);
        let cache = Arc::new(Self {
            store: Arc::new(Mutex::new(Store {
                revision,
                ..Default::default()
            })),
            revision: AtomicU64::new(revision),
            domain_stats: DomainStatsStore::default(),
            cold: Arc::new(Mutex::new(ColdStore::default())),
            tier_settings: Mutex::new(settings),
            scope_revisions: Mutex::new(scope_revisions),
            data_root,
            cold_tx: Some(tx),
            cold_delete_fallbacks: Arc::new(Semaphore::new(COLD_DELETE_FALLBACKS)),
            cleanup_slots: Arc::new(Semaphore::new(COLD_CLEANUP_SLOTS)),
            fills: Arc::new(Semaphore::new(FILL_SLOTS)),
            fill_bytes_active: AtomicU64::new(0),
            transient_bytes_active: Arc::new(AtomicU64::new(0)),
            cold_queue_items: AtomicU64::new(0),
            cold_queue_bytes: AtomicU64::new(0),
            cold_write_bytes_active: AtomicU64::new(0),
            cold_read_bytes_active: AtomicU64::new(0),
            revalidating: Mutex::new(HashSet::new()),
            fill_flights: Mutex::new(HashMap::new()),
            uncacheable_until: Mutex::new(HashMap::new()),
            fill_flight_sequence: AtomicU64::new(0),
            collapse_waiters: AtomicU64::new(0),
            bans: Mutex::new(Vec::new()),
            ban_rule_count: AtomicU64::new(0),
            ban_sequence: AtomicU64::new(0),
            proactive_scheduled_keys: Mutex::new(HashSet::new()),
            proactive_inflight: AtomicU64::new(0),
            cleanup_jobs: Mutex::new(HashMap::new()),
            cleanup_sequence: AtomicU64::new(0),
            counters: Counters::default(),
        });
        cache.load_cold_startup();
        let weak = Arc::downgrade(&cache);
        let worker_shutdown = shutdown.clone();
        tokio::spawn(async move { Self::cold_worker(weak, rx, worker_shutdown).await });
        let weak = Arc::downgrade(&cache);
        let maintenance_shutdown = shutdown.clone();
        tokio::spawn(async move {
            let mut tick = tokio::time::interval(Duration::from_secs(5));
            tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
            loop {
                tokio::select! {
                    _ = maintenance_shutdown.cancelled() => break,
                    _ = tick.tick() => {
                        let Some(cache) = weak.upgrade() else { break; };
                        cache.trim_hot_idle();
                        cache.demote_idle();
                        cache.sweep_cold().await;
                        cache.gc_banned(256);
                    }
                }
            }
        });
        let weak = Arc::downgrade(&cache);
        tokio::spawn(async move {
            loop {
                let Some(cache) = weak.upgrade() else {
                    break;
                };
                let interval = cache
                    .tier_settings
                    .lock()
                    .unwrap_or_else(|e| e.into_inner())
                    .cold_gc_interval_seconds;
                drop(cache);
                tokio::select! {
                    _ = shutdown.cancelled() => break,
                    _ = tokio::time::sleep(Duration::from_secs(interval)) => {
                        let Some(cache) = weak.upgrade() else { break; };
                        cache.cold_gc_pass().await;
                    }
                }
            }
        });
        cache
    }
    fn fill_headroom_bytes(max_bytes: usize) -> usize {
        let proportional = (max_bytes / 64).clamp(4 * 1024, 1024 * 1024);
        proportional.min((max_bytes / 8).max(1))
    }

    fn try_reserve_transient(&self, bytes: usize, settings: &Settings) -> Option<TransientLease> {
        let bytes = u64::try_from(bytes).ok()?;
        let limit = u64::try_from(settings.transient_max_bytes).unwrap_or(u64::MAX);
        loop {
            let current = self.transient_bytes_active.load(Ordering::Relaxed);
            let next = current.checked_add(bytes)?;
            if next > limit {
                self.counters
                    .memory_pressure_rejections
                    .fetch_add(1, Ordering::Relaxed);
                return None;
            }
            if self
                .transient_bytes_active
                .compare_exchange_weak(current, next, Ordering::AcqRel, Ordering::Relaxed)
                .is_ok()
            {
                return Some(TransientLease {
                    _inner: Arc::new(TransientLeaseInner {
                        counter: self.transient_bytes_active.clone(),
                        bytes,
                    }),
                });
            }
        }
    }

    fn reserve_fill_allocation(&self, bytes: usize, settings: &Settings) -> Option<TransientLease> {
        if bytes == 0 {
            return self.try_reserve_transient(0, settings);
        }
        let lease = self.try_reserve_transient(bytes, settings)?;
        let mut store = self.store.lock().unwrap_or_else(|e| e.into_inner());
        let fill_now =
            usize::try_from(self.fill_bytes_active.load(Ordering::Relaxed)).unwrap_or(usize::MAX);
        let allowed = settings
            .max_bytes
            .saturating_add(Self::fill_headroom_bytes(settings.max_bytes));
        let revision = store.revision;
        let generation = store.generation;
        let mut removed = 0usize;
        let mut released_hot = false;
        while store.bytes.saturating_add(fill_now).saturating_add(bytes) > allowed {
            if removed >= settings.hot_eviction_max_objects {
                self.counters
                    .hot_eviction_limited
                    .fetch_add(1, Ordering::Relaxed);
                break;
            }
            let Some((_, key)) = store.lru.first_key_value().map(|(seq, key)| (*seq, *key)) else {
                break;
            };
            let Some(old) = store.remove(&key) else {
                continue;
            };
            removed = removed.saturating_add(1);
            released_hot = true;
            let scope_generation = store
                .scope_generations
                .get(&old.scope)
                .copied()
                .unwrap_or(0);
            let key_generation = store
                .key_generations
                .get(&old.base_key)
                .copied()
                .unwrap_or(0);
            if !(settings.cold_enabled
                && settings.cold_on_lru
                && self.enqueue_demotion(
                    key,
                    old,
                    revision,
                    generation,
                    scope_generation,
                    key_generation,
                ))
            {
                self.counters.evictions.fetch_add(1, Ordering::Relaxed);
            }
        }
        let fits = store.bytes.saturating_add(fill_now).saturating_add(bytes) <= allowed;
        if fits {
            self.fill_bytes_active
                .fetch_add(bytes as u64, Ordering::Relaxed);
        }
        drop(store);
        if released_hot {
            release_allocator_memory();
        }
        if fits {
            Some(lease)
        } else {
            self.counters
                .memory_pressure_rejections
                .fetch_add(1, Ordering::Relaxed);
            drop(lease);
            None
        }
    }

    fn trim_hot_idle(&self) {
        let settings = self
            .tier_settings
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .clone();
        if !settings.hot_auto_trim_enabled {
            return;
        }
        let target = settings
            .max_bytes
            .saturating_mul(settings.hot_auto_trim_target_percent as usize)
            / 100;
        let idle_cutoff = Duration::from_secs(settings.hot_auto_trim_after_seconds);
        let mut store = self.store.lock().unwrap_or_else(|e| e.into_inner());
        if store.bytes <= target {
            return;
        }
        let revision = store.revision;
        let generation = store.generation;
        let mut removed = 0usize;
        while store.bytes > target && removed < settings.hot_eviction_max_objects {
            let Some((_, key)) = store.lru.first_key_value().map(|(seq, key)| (*seq, *key)) else {
                break;
            };
            let idle = store
                .entries
                .get(&key)
                .is_some_and(|entry| entry.last_access.elapsed() >= idle_cutoff);
            if !idle {
                break;
            }
            let Some(entry) = store.remove(&key) else {
                continue;
            };
            removed += 1;
            let scope_generation = store
                .scope_generations
                .get(&entry.scope)
                .copied()
                .unwrap_or(0);
            let key_generation = store
                .key_generations
                .get(&entry.base_key)
                .copied()
                .unwrap_or(0);
            if !(settings.cold_enabled
                && settings.cold_on_lru
                && self.enqueue_demotion(
                    key,
                    entry,
                    revision,
                    generation,
                    scope_generation,
                    key_generation,
                ))
            {
                self.counters.evictions.fetch_add(1, Ordering::Relaxed);
            }
        }
        drop(store);
        if removed > 0 {
            release_allocator_memory();
        }
    }

    pub fn set_control_revision(&self, revision: u64) {
        self.store
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .control_revision = revision;
    }

    pub fn replace_scope_revisions(&self, scopes: HashMap<String, u64>) {
        *self
            .scope_revisions
            .lock()
            .unwrap_or_else(|e| e.into_inner()) = scopes;
    }

    pub fn configure(&self, settings: Settings) {
        *self.tier_settings.lock().unwrap_or_else(|e| e.into_inner()) = settings.clone();
        let mut store = self.store.lock().unwrap_or_else(|e| e.into_inner());
        let revision = store.revision;
        let generation = store.generation;
        let mut released_hot = false;
        while store.entries.len() > settings.max_entries || store.bytes > settings.max_bytes {
            let Some((_, key)) = store.lru.first_key_value().map(|(seq, key)| (*seq, *key)) else {
                break;
            };
            let Some(old) = store.remove(&key) else {
                continue;
            };
            released_hot = true;
            let scope_generation = store
                .scope_generations
                .get(&old.scope)
                .copied()
                .unwrap_or(0);
            let key_generation = store
                .key_generations
                .get(&old.base_key)
                .copied()
                .unwrap_or(0);
            if !(settings.cold_enabled
                && settings.cold_on_lru
                && self.enqueue_demotion(
                    key,
                    old,
                    revision,
                    generation,
                    scope_generation,
                    key_generation,
                ))
            {
                self.counters.evictions.fetch_add(1, Ordering::Relaxed);
            }
        }
        drop(store);
        if released_hot {
            release_allocator_memory();
        }
    }
    pub fn revision(&self) -> u64 {
        self.revision.load(Ordering::Relaxed)
    }
}