lgui_core/memory/

governor.rs

use std::{
    sync::{
        atomic::{AtomicU64, AtomicUsize, Ordering},
        Arc, Mutex,
    },
    time::{Duration, Instant},
};

use super::registry::RegistryState;
use super::{
    CacheDomain, CacheRegistration, CacheScope, CacheUsage, DomainInstanceId, DomainRegistration,
    DomainSnapshot, MemoryAction, MemoryEvent, MemoryOptions, MemorySnapshot, TrimReason,
    TrimRequest, TrimSnapshot,
};

#[cfg(feature = "persistent-cache")]
use super::PersistentCacheStore;

const FRAME_BUDGET_CHECK_INTERVAL: Duration = Duration::from_millis(250);

#[derive(Clone)]
pub struct MemoryGovernor {
    inner: Arc<MemoryGovernorInner>,
}

struct MemoryGovernorInner {
    options: Mutex<MemoryOptions>,
    registry: Arc<Mutex<RegistryState>>,
    epoch: AtomicU64,
    transient_reserved: Arc<AtomicUsize>,
    large_tasks_in_flight: Arc<AtomicUsize>,
    last_trim: Mutex<Option<TrimSnapshot>>,
    last_frame_budget_check: Mutex<Option<Instant>>,
    #[cfg(feature = "persistent-cache")]
    persistent: Option<Arc<dyn PersistentCacheStore>>,
}

impl MemoryGovernor {
    pub fn new(options: MemoryOptions) -> Self {
        Self::with_store(
            options,
            #[cfg(feature = "persistent-cache")]
            None,
        )
    }

    pub(crate) fn with_store(
        options: MemoryOptions,
        #[cfg(feature = "persistent-cache")] persistent: Option<Arc<dyn PersistentCacheStore>>,
    ) -> Self {
        options
            .validate()
            .expect("invalid application memory policy");
        Self {
            inner: Arc::new(MemoryGovernorInner {
                options: Mutex::new(options),
                registry: Arc::new(Mutex::new(RegistryState::default())),
                epoch: AtomicU64::new(1),
                transient_reserved: Arc::new(AtomicUsize::new(0)),
                large_tasks_in_flight: Arc::new(AtomicUsize::new(0)),
                last_trim: Mutex::new(None),
                last_frame_budget_check: Mutex::new(None),
                #[cfg(feature = "persistent-cache")]
                persistent,
            }),
        }
    }

    pub fn options(&self) -> MemoryOptions {
        *self.inner.options.lock().expect("memory options poisoned")
    }

    pub fn set_options(&self, options: MemoryOptions) {
        options
            .validate()
            .expect("invalid application memory policy");
        *self.inner.options.lock().expect("memory options poisoned") = options;
        self.bump_epoch();
        self.rebalance_budgets();
        self.enforce_budget(true);
        #[cfg(feature = "persistent-cache")]
        if let Some(store) = self.persistent_cache() {
            let _ = store.trim_to(options.budget.persistent_bytes);
        }
    }

    pub fn next_instance_id(&self) -> DomainInstanceId {
        let mut registry = self
            .inner
            .registry
            .lock()
            .expect("memory registry poisoned");
        let id = registry.next_instance;
        registry.next_instance = registry.next_instance.wrapping_add(1).max(1);
        DomainInstanceId(id)
    }

    pub fn register(&self, registration: DomainRegistration) -> CacheRegistration {
        let mut registry = self
            .inner
            .registry
            .lock()
            .expect("memory registry poisoned");
        let id = registry.next_registration;
        registry.next_registration = registry.next_registration.wrapping_add(1).max(1);
        registry.entries.insert(id, registration);
        drop(registry);
        self.rebalance_budgets();
        self.bump_epoch();
        let governor = self.clone();
        CacheRegistration::new(id, Arc::downgrade(&self.inner.registry), move || {
            governor.rebalance_budgets();
            governor.bump_epoch();
        })
    }

    pub fn snapshot(&self) -> MemorySnapshot {
        let epoch = self.inner.epoch.load(Ordering::Acquire);
        let entries = self.registered_entries();
        let mut usage = CacheUsage::default();
        let mut domains = Vec::with_capacity(entries.len());
        for (id, registration) in entries {
            let domain_usage = registration.adapter.usage();
            usage.add_assign(domain_usage);
            domains.push(DomainSnapshot {
                registration_id: id,
                domain: registration.domain,
                instance: registration.instance,
                owner: registration.owner,
                usage: domain_usage,
            });
        }
        let options = self.options();
        let cache_soft = options.budget.cache_soft_bytes;
        MemorySnapshot {
            epoch,
            options,
            transient_reserved_bytes: self.inner.transient_reserved.load(Ordering::Acquire),
            large_tasks_in_flight: self.inner.large_tasks_in_flight.load(Ordering::Acquire),
            pinned_overflow_bytes: usage.pinned_bytes.saturating_sub(cache_soft),
            usage,
            domains,
            last_trim: *self
                .inner
                .last_trim
                .lock()
                .expect("memory trim state poisoned"),
        }
    }

    pub fn trim(&self, reason: TrimReason, scope: CacheScope, target_bytes: usize) -> usize {
        let epoch = self.bump_epoch();
        let started = Instant::now();
        let entries = if scope == CacheScope::Persistent {
            Vec::new()
        } else {
            self.registered_entries()
        };
        let assignments = entries
            .iter()
            .filter(|(_, registration)| {
                registration.domain != CacheDomain::Persistent
                    && (scope == CacheScope::AllRebuildable
                        || registration.domain != CacheDomain::HostScene)
            })
            .map(|(_, registration)| {
                (
                    registration,
                    self.assigned_budget(registration.domain, &entries),
                )
            })
            .collect::<Vec<_>>();
        let assigned_total = assignments
            .iter()
            .map(|(_, assigned)| *assigned as u128)
            .sum::<u128>();
        let mut released = 0usize;
        for (registration, assigned) in assignments {
            let domain_target = if target_bytes == 0 {
                0
            } else if target_bytes == usize::MAX {
                assigned
            } else if assigned_total == 0 {
                0
            } else {
                let proportional =
                    (assigned as u128).saturating_mul(target_bytes as u128) / assigned_total;
                usize::try_from(proportional).unwrap_or(usize::MAX)
            };
            let result = registration.adapter.trim(TrimRequest {
                reason,
                scope,
                target_bytes: domain_target,
            });
            released = released.saturating_add(result.released_bytes());
        }
        #[cfg(feature = "persistent-cache")]
        if scope == CacheScope::Persistent {
            if let Some(store) = self.inner.persistent.as_ref() {
                let target = u64::try_from(target_bytes).unwrap_or(u64::MAX);
                if let (Ok(before), Ok(after)) = (store.stats(), store.trim_to(target)) {
                    released = released.saturating_add(
                        usize::try_from(before.bytes.saturating_sub(after.bytes))
                            .unwrap_or(usize::MAX),
                    );
                }
            }
        }
        let duration_micros = started.elapsed().as_micros().min(u64::MAX as u128) as u64;
        *self
            .inner
            .last_trim
            .lock()
            .expect("memory trim state poisoned") = Some(TrimSnapshot {
            reason,
            requested_at_epoch: epoch,
            released_bytes: released,
            duration_micros,
        });
        released
    }

    pub fn invalidate_domain(&self, domain: CacheDomain) -> usize {
        let entries = self.registered_entries();
        let mut released = 0usize;
        for (_, registration) in entries {
            if registration.domain == domain {
                released = released.saturating_add(
                    registration
                        .adapter
                        .trim(TrimRequest {
                            reason: TrimReason::Explicit,
                            scope: CacheScope::Memory,
                            target_bytes: 0,
                        })
                        .released_bytes(),
                );
            }
        }
        self.bump_epoch();
        released
    }

    pub fn notify(&self, event: MemoryEvent) {
        match self.options().event_action(event) {
            MemoryAction::None => {}
            MemoryAction::EnforceBudget => {
                if event == MemoryEvent::FrameCommitted {
                    if self.begin_frame_budget_check() {
                        self.finish_frame_budget_check();
                    }
                } else {
                    self.enforce_budget(true);
                }
            }
            MemoryAction::Trim {
                scope,
                target_bytes,
            } => {
                self.trim(event.trim_reason(), scope, target_bytes);
            }
        }
    }

    pub fn try_reserve(&self, bytes: usize) -> Option<MemoryReservation> {
        let hard = self.options().budget.transient_hard_bytes;
        let reserved = &self.inner.transient_reserved;
        if hard == usize::MAX {
            return Some(MemoryReservation {
                bytes,
                accounted_bytes: 0,
                reserved: Arc::clone(reserved),
            });
        }
        let mut current = reserved.load(Ordering::Acquire);
        loop {
            let next = current.checked_add(bytes)?;
            if next > hard {
                self.trim(TrimReason::HardBudget, CacheScope::Memory, 0);
                return None;
            }
            match reserved.compare_exchange_weak(current, next, Ordering::AcqRel, Ordering::Acquire)
            {
                Ok(_) => {
                    return Some(MemoryReservation {
                        bytes,
                        accounted_bytes: bytes,
                        reserved: Arc::clone(reserved),
                    })
                }
                Err(actual) => current = actual,
            }
        }
    }

    pub fn try_reserve_task(&self, bytes: usize) -> Option<MemoryTaskReservation> {
        let limit = self.options().budget.max_parallel_large_tasks;
        let in_flight = &self.inner.large_tasks_in_flight;
        let mut current = in_flight.load(Ordering::Acquire);
        loop {
            if current >= limit {
                return None;
            }
            match in_flight.compare_exchange_weak(
                current,
                current + 1,
                Ordering::AcqRel,
                Ordering::Acquire,
            ) {
                Ok(_) => break,
                Err(actual) => current = actual,
            }
        }
        let Some(reservation) = self.try_reserve(bytes) else {
            in_flight.fetch_sub(1, Ordering::AcqRel);
            return None;
        };
        Some(MemoryTaskReservation {
            _reservation: reservation,
            in_flight: Arc::clone(in_flight),
        })
    }

    #[cfg(feature = "persistent-cache")]
    pub fn persistent_cache(&self) -> Option<Arc<dyn PersistentCacheStore>> {
        self.options()
            .persistent_cache_enabled
            .then(|| self.inner.persistent.as_ref().map(Arc::clone))
            .flatten()
    }

    #[cfg(feature = "persistent-cache")]
    pub fn persistent_cache_stats(
        &self,
    ) -> Result<super::PersistentCacheStats, super::CacheStoreError> {
        self.inner.persistent.as_ref().map_or_else(
            || Ok(super::PersistentCacheStats::default()),
            |store| store.stats(),
        )
    }

    #[cfg(feature = "persistent-cache")]
    pub fn clear_persistent_cache(
        &self,
        namespace: Option<&str>,
    ) -> Result<super::PersistentCacheStats, super::CacheStoreError> {
        if let Some(store) = self.inner.persistent.as_ref() {
            store.clear(namespace)?;
            self.bump_epoch();
            return store.stats();
        }
        Ok(super::PersistentCacheStats::default())
    }

    fn enforce_budget(&self, include_soft_limit: bool) {
        let options = self.options();
        let usage = self.budget_usage();
        let evictable = usage.managed_bytes.saturating_sub(usage.protected_bytes);
        let reason = if evictable > options.budget.cache_hard_bytes {
            Some(TrimReason::HardBudget)
        } else if include_soft_limit && evictable > options.budget.cache_soft_bytes {
            Some(TrimReason::SoftBudget)
        } else {
            None
        };
        if let Some(reason) = reason {
            self.trim(reason, CacheScope::Memory, options.budget.cache_soft_bytes);
        }
    }

    pub(crate) fn begin_frame_budget_check(&self) -> bool {
        if self.options().event_action(MemoryEvent::FrameCommitted) != MemoryAction::EnforceBudget {
            return false;
        }
        let now = Instant::now();
        let mut last = self
            .inner
            .last_frame_budget_check
            .lock()
            .expect("memory budget check state poisoned");
        if last.is_some_and(|previous| {
            now.saturating_duration_since(previous) < FRAME_BUDGET_CHECK_INTERVAL
        }) {
            return false;
        }
        *last = Some(now);
        true
    }

    pub(crate) fn finish_frame_budget_check(&self) {
        self.enforce_budget(false);
    }

    fn budget_usage(&self) -> BudgetUsage {
        let probes = self.registered_usage_probes();
        let mut result = BudgetUsage::default();
        for (domain, adapter) in probes {
            let usage = adapter.usage();
            let managed = usage.managed_bytes();
            result.managed_bytes = result.managed_bytes.saturating_add(managed);
            let protected = if domain == CacheDomain::HostScene {
                managed
            } else {
                usage.pinned_bytes.min(managed)
            };
            result.protected_bytes = result.protected_bytes.saturating_add(protected);
        }
        result
    }

    fn registered_usage_probes(&self) -> Vec<(CacheDomain, super::CacheAdapter)> {
        self.inner
            .registry
            .lock()
            .expect("memory registry poisoned")
            .entries
            .values()
            .map(|registration| (registration.domain, registration.adapter.clone()))
            .collect()
    }

    fn registered_entries(&self) -> Vec<(u64, DomainRegistration)> {
        self.inner
            .registry
            .lock()
            .expect("memory registry poisoned")
            .entries
            .iter()
            .map(|(id, registration)| (*id, registration.clone()))
            .collect()
    }

    fn rebalance_budgets(&self) {
        let entries = self.registered_entries();
        for (_, registration) in &entries {
            registration
                .adapter
                .set_budget(self.assigned_budget(registration.domain, &entries));
        }
    }

    fn assigned_budget(&self, domain: CacheDomain, entries: &[(u64, DomainRegistration)]) -> usize {
        let options = self.options();
        if domain == CacheDomain::Persistent {
            return options.budget.persistent_bytes.min(usize::MAX as u64) as usize;
        }
        let instances = entries
            .iter()
            .filter(|(_, registration)| registration.domain == domain)
            .count()
            .max(1);
        options.domain_budget(domain).saturating_div(instances)
    }

    fn bump_epoch(&self) -> u64 {
        self.inner.epoch.fetch_add(1, Ordering::AcqRel) + 1
    }
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
struct BudgetUsage {
    managed_bytes: usize,
    protected_bytes: usize,
}

pub struct MemoryReservation {
    bytes: usize,
    accounted_bytes: usize,
    reserved: Arc<AtomicUsize>,
}

pub struct MemoryTaskReservation {
    _reservation: MemoryReservation,
    in_flight: Arc<AtomicUsize>,
}

impl MemoryTaskReservation {
    pub const fn bytes(&self) -> usize {
        self._reservation.bytes()
    }
}

impl Drop for MemoryTaskReservation {
    fn drop(&mut self) {
        self.in_flight.fetch_sub(1, Ordering::AcqRel);
    }
}

impl MemoryReservation {
    pub const fn bytes(&self) -> usize {
        self.bytes
    }
}

impl Drop for MemoryReservation {
    fn drop(&mut self) {
        self.reserved
            .fetch_sub(self.accounted_bytes, Ordering::AcqRel);
    }
}