gpu_handle_types/

budget.rs

// SPDX-License-Identifier: MIT OR Apache-2.0

//! Memory budget — three-tier pressure response (callbacks → evict →
//! retry). Keeps downstream pools within caller-declared caps.
//!
//! ## Why the single `dyn` exception
//!
//! [`EvictablePool`] is the one `dyn` trait-object in this crate. The
//! dispatch fires only when the budget is under pressure (soft/hard cap
//! breach or `try_reserve` needs to reclaim) — never on the per-frame
//! hot path — so the principle "no dyn on the hot path" still holds.

use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::{Arc, Weak};

use parking_lot::Mutex;
use smallvec::SmallVec;

use crate::{BackendKind, Error};

/// Per-backend budget caps. `None` disables that side (e.g. soft_cap
/// only, hard_cap only, or unbounded).
#[derive(Debug, Default, Copy, Clone)]
pub struct BudgetCaps {
    pub soft_cap_bytes: Option<u64>,
    pub hard_cap_bytes: Option<u64>,
}

/// Named-field set of per-backend caps. Deliberately a struct with
/// cfg-gated fields rather than a `[BudgetCaps; N]` — no
/// `BackendKind::COUNT` to mislead under `--no-default-features`, and
/// partial configs compose via `..Default::default()`.
#[derive(Debug, Default, Clone)]
pub struct BudgetCapsSet {
    pub cpu: BudgetCaps,
    #[cfg(feature = "wgpu")]
    pub wgpu: BudgetCaps,
    #[cfg(feature = "opencl")]
    pub opencl: BudgetCaps,
    #[cfg(feature = "cuda")]
    pub cuda: BudgetCaps,
}

impl BudgetCapsSet {
    fn for_backend(&self, backend: BackendKind) -> BudgetCaps {
        match backend {
            BackendKind::Cpu => self.cpu,
            #[cfg(feature = "wgpu")]
            BackendKind::Wgpu => self.wgpu,
            #[cfg(feature = "opencl")]
            BackendKind::OpenCl => self.opencl,
            #[cfg(feature = "cuda")]
            BackendKind::Cuda => self.cuda,
            // Native-API variants (Vulkan / D3D12 / D3D11 / Metal /
            // OpenGL) account against the wgpu budget — every native
            // submit ultimately backs a wgpu texture or aliases
            // its memory.
            #[cfg(feature = "wgpu")]
            BackendKind::Vulkan
            | BackendKind::D3D12
            | BackendKind::D3D11
            | BackendKind::Metal
            | BackendKind::OpenGL => self.wgpu,
            #[cfg(not(feature = "wgpu"))]
            BackendKind::Vulkan
            | BackendKind::D3D12
            | BackendKind::D3D11
            | BackendKind::Metal
            | BackendKind::OpenGL => self.cpu,
            // Android routing variants account against the
            // wgpu budget: the producer is a wgpu texture and the
            // external handle aliases its memory (AHB export is
            // refcounted and the underlying VkDeviceMemory belongs to the
            // wgpu device).
            #[cfg(all(target_os = "android", feature = "wgpu"))]
            BackendKind::AHardwareBuffer
            | BackendKind::AndroidPresentation
            | BackendKind::AndroidSurfaceControl
            | BackendKind::AImageWriter
            | BackendKind::MediaCodec
            | BackendKind::ForeignGl
            | BackendKind::ForeignVulkan => self.wgpu,
            // Same cfg with wgpu off — Android build without the
            // wgpu feature falls back to the CPU bucket for
            // accounting. In practice an Android consumer enables
            // wgpu, so this branch only serves builds of this crate
            // without it.
            #[cfg(all(target_os = "android", not(feature = "wgpu")))]
            BackendKind::AHardwareBuffer
            | BackendKind::AndroidPresentation
            | BackendKind::AndroidSurfaceControl
            | BackendKind::AImageWriter
            | BackendKind::MediaCodec
            | BackendKind::ForeignGl
            | BackendKind::ForeignVulkan => self.cpu,
            // WebCodecs shared allocation backs its OffscreenCanvas swap
            // targets with wgpu textures, so it accounts against the wgpu
            // budget (cpu bucket when wgpu is compiled out).
            #[cfg(all(feature = "web-codecs", feature = "wgpu"))]
            BackendKind::WebCodecs => self.wgpu,
            #[cfg(all(feature = "web-codecs", not(feature = "wgpu")))]
            BackendKind::WebCodecs => self.cpu,
        }
    }
}

/// Per-backend usage counters. Parallel to [`BudgetCapsSet`].
#[derive(Default)]
struct BudgetCountersSet {
    cpu: AtomicU64,
    #[cfg(feature = "wgpu")]
    wgpu: AtomicU64,
    #[cfg(feature = "opencl")]
    opencl: AtomicU64,
    #[cfg(feature = "cuda")]
    cuda: AtomicU64,
}

impl BudgetCountersSet {
    fn get(&self, backend: BackendKind) -> &AtomicU64 {
        match backend {
            BackendKind::Cpu => &self.cpu,
            #[cfg(feature = "wgpu")]
            BackendKind::Wgpu => &self.wgpu,
            #[cfg(feature = "opencl")]
            BackendKind::OpenCl => &self.opencl,
            #[cfg(feature = "cuda")]
            BackendKind::Cuda => &self.cuda,
            // Native-API variants — see `BudgetCapsSet::for_backend`.
            #[cfg(feature = "wgpu")]
            BackendKind::Vulkan
            | BackendKind::D3D12
            | BackendKind::D3D11
            | BackendKind::Metal
            | BackendKind::OpenGL => &self.wgpu,
            #[cfg(not(feature = "wgpu"))]
            BackendKind::Vulkan
            | BackendKind::D3D12
            | BackendKind::D3D11
            | BackendKind::Metal
            | BackendKind::OpenGL => &self.cpu,
            // See `BudgetCapsSet::for_backend` for the rationale: all
            // Android routing variants account against wgpu.
            #[cfg(all(target_os = "android", feature = "wgpu"))]
            BackendKind::AHardwareBuffer
            | BackendKind::AndroidPresentation
            | BackendKind::AndroidSurfaceControl
            | BackendKind::AImageWriter
            | BackendKind::MediaCodec
            | BackendKind::ForeignGl
            | BackendKind::ForeignVulkan => &self.wgpu,
            #[cfg(all(target_os = "android", not(feature = "wgpu")))]
            BackendKind::AHardwareBuffer
            | BackendKind::AndroidPresentation
            | BackendKind::AndroidSurfaceControl
            | BackendKind::AImageWriter
            | BackendKind::MediaCodec
            | BackendKind::ForeignGl
            | BackendKind::ForeignVulkan => &self.cpu,
            // See `BudgetCapsSet::for_backend` — WebCodecs accounts
            // against wgpu (cpu bucket when wgpu is compiled out).
            #[cfg(all(feature = "web-codecs", feature = "wgpu"))]
            BackendKind::WebCodecs => &self.wgpu,
            #[cfg(all(feature = "web-codecs", not(feature = "wgpu")))]
            BackendKind::WebCodecs => &self.cpu,
        }
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum BudgetPressure {
    Ok,
    Soft,
    Hard,
}

#[derive(Debug, Copy, Clone)]
pub struct BudgetPressureEvent {
    pub backend: BackendKind,
    pub level: BudgetPressure,
    pub current_usage_bytes: u64,
    pub required_bytes: u64,
    pub cap_bytes: u64,
}

// `Send + Sync` off wasm; dropped on wasm (single-threaded, the callback
// may close over thread-affine state). An explicit `+ Send + Sync` on a
// `dyn Fn` cannot be spelled with a non-auto marker trait, so the alias is
// cfg-split directly.
#[cfg(not(target_family = "wasm"))]
pub type PressureCallback = Arc<dyn Fn(&BudgetPressureEvent) + Send + Sync>;
#[cfg(target_family = "wasm")]
pub type PressureCallback = Arc<dyn Fn(&BudgetPressureEvent)>;

/// Pool trait — drop unused resources on request. This is the single
/// `dyn` exception in this crate; see the module doc-comment. The
/// `MaybeSendSync` supertrait is `Send + Sync` off wasm and empty on wasm.
pub trait EvictablePool: crate::MaybeSendSync {
    /// Reclaim unused resources; returns bytes freed.
    fn evict_unused(&self) -> u64;
    /// Current in-use byte count.
    fn bytes_in_use(&self) -> u64;
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct PoolHandle(pub u32);

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct PressureCallbackHandle(pub u32);

type PoolEntry = (PoolHandle, BackendKind, Weak<dyn EvictablePool>);
type CallbackEntry = (PressureCallbackHandle, PressureCallback);

pub struct MemoryBudget {
    caps: BudgetCapsSet,
    usage: BudgetCountersSet,
    pools: Mutex<SmallVec<[PoolEntry; 4]>>,
    callbacks: Mutex<SmallVec<[CallbackEntry; 4]>>,
    next_pool_handle: AtomicU32,
    next_callback_handle: AtomicU32,
}

impl core::fmt::Debug for MemoryBudget {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("MemoryBudget").field("caps", &self.caps).finish_non_exhaustive()
    }
}

/// RAII token. Dropping the token releases the reservation back to the
/// pool.
pub struct BudgetReservation {
    budget: Arc<MemoryBudget>,
    backend: BackendKind,
    bytes: u64,
}

impl BudgetReservation {
    pub fn backend(&self) -> BackendKind {
        self.backend
    }
    pub fn bytes(&self) -> u64 {
        self.bytes
    }
}

impl Drop for BudgetReservation {
    fn drop(&mut self) {
        self.budget.usage.get(self.backend).fetch_sub(self.bytes, Ordering::Release);
    }
}

impl core::fmt::Debug for BudgetReservation {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("BudgetReservation").field("backend", &self.backend).field("bytes", &self.bytes).finish()
    }
}

impl MemoryBudget {
    // On wasm the pool and callback entries are not `Send + Sync` (see
    // `PressureCallback`), and neither is the `Arc` — by design, the web
    // target is single-threaded.
    #[cfg_attr(target_family = "wasm", allow(clippy::arc_with_non_send_sync))]
    pub fn new(caps: BudgetCapsSet) -> Arc<Self> {
        Arc::new(Self {
            caps,
            usage: BudgetCountersSet::default(),
            pools: Mutex::new(SmallVec::new()),
            callbacks: Mutex::new(SmallVec::new()),
            next_pool_handle: AtomicU32::new(1),
            next_callback_handle: AtomicU32::new(1),
        })
    }

    pub fn caps(&self, backend: BackendKind) -> BudgetCaps {
        self.caps.for_backend(backend)
    }

    pub fn current_usage(&self, backend: BackendKind) -> u64 {
        self.usage.get(backend).load(Ordering::Acquire)
    }

    pub fn available(&self, backend: BackendKind) -> u64 {
        let cap = self.caps.for_backend(backend).hard_cap_bytes.unwrap_or(u64::MAX);
        cap.saturating_sub(self.current_usage(backend))
    }

    pub fn pressure(&self, backend: BackendKind, required: u64) -> BudgetPressure {
        let caps = self.caps.for_backend(backend);
        let usage = self.current_usage(backend);
        let projected = usage.saturating_add(required);
        if let Some(hard) = caps.hard_cap_bytes
            && projected > hard
        {
            return BudgetPressure::Hard;
        }
        if let Some(soft) = caps.soft_cap_bytes
            && projected > soft
        {
            return BudgetPressure::Soft;
        }
        BudgetPressure::Ok
    }

    pub fn try_reserve(self: &Arc<Self>, backend: BackendKind, bytes: u64) -> Result<BudgetReservation, Error> {
        // Fast-path: no caps configured → pure atomic add, no locks, no
        // allocations. This is the single-stream configuration and must
        // stay allocation-free.
        let caps = self.caps.for_backend(backend);
        if caps.soft_cap_bytes.is_none() && caps.hard_cap_bytes.is_none() {
            self.usage.get(backend).fetch_add(bytes, Ordering::Release);
            return Ok(BudgetReservation { budget: self.clone(), backend, bytes });
        }

        // Three-tier response: callbacks → evict →
        // commit. Pressure is sampled *before* committing so soft-cap
        // callbacks fire on the boundary crossing even when the
        // reservation itself fits under the hard cap (soft is
        // warning-only, but the warning must still fire on a soft-cap
        // or hard-cap boundary crossing).
        let level = self.pressure(backend, bytes);
        if level != BudgetPressure::Ok {
            let cap = caps.hard_cap_bytes.or(caps.soft_cap_bytes).unwrap_or(u64::MAX);
            let event = BudgetPressureEvent {
                backend,
                level,
                current_usage_bytes: self.current_usage(backend),
                required_bytes: bytes,
                cap_bytes: cap,
            };

            // Tier 1: pressure callbacks. Misbehaving callbacks must not
            // deadlock the reservation path.
            {
                let cbs = self.callbacks.lock().clone();
                for (_, cb) in cbs.iter() {
                    let cb = cb.clone();
                    let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| cb(&event)));
                }
            }

            // Tier 2: evict each registered pool on the target backend.
            {
                let pools = self.pools.lock().clone();
                for (_, bk, weak) in pools.iter() {
                    if *bk != backend {
                        continue;
                    }
                    if let Some(pool) = weak.upgrade() {
                        let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| pool.evict_unused()));
                    }
                }
            }
        }

        // Tier 3: commit. Soft pressure passes through (soft is
        // warning-only); hard pressure passes only if eviction freed
        // enough room.
        if self.try_commit(backend, bytes) {
            return Ok(BudgetReservation { budget: self.clone(), backend, bytes });
        }

        // Over hard cap → hard failure.
        let available = self.available(backend);
        Err(Error::OutOfGpuMemory { required_bytes: bytes, available_bytes: available, backend })
    }

    /// Attempt an atomic commit under the hard cap. Returns `true` on
    /// success. Soft-cap breach does not block commit — it only triggers
    /// callbacks above.
    fn try_commit(&self, backend: BackendKind, bytes: u64) -> bool {
        let counter = self.usage.get(backend);
        let hard = self.caps.for_backend(backend).hard_cap_bytes;
        loop {
            let cur = counter.load(Ordering::Acquire);
            let next = cur.saturating_add(bytes);
            if let Some(h) = hard
                && next > h
            {
                return false;
            }
            if counter.compare_exchange(cur, next, Ordering::AcqRel, Ordering::Acquire).is_ok() {
                return true;
            }
        }
    }

    pub fn register_pool(&self, backend: BackendKind, pool: &Arc<dyn EvictablePool>) -> PoolHandle {
        let h = PoolHandle(self.next_pool_handle.fetch_add(1, Ordering::Relaxed));
        self.pools.lock().push((h, backend, Arc::downgrade(pool)));
        h
    }

    pub fn unregister_pool(&self, handle: PoolHandle) {
        self.pools.lock().retain(|(h, _, _)| *h != handle);
    }

    pub fn register_pressure_callback(&self, cb: PressureCallback) -> PressureCallbackHandle {
        let h = PressureCallbackHandle(self.next_callback_handle.fetch_add(1, Ordering::Relaxed));
        self.callbacks.lock().push((h, cb));
        h
    }

    pub fn unregister_pressure_callback(&self, handle: PressureCallbackHandle) {
        self.callbacks.lock().retain(|(h, _)| *h != handle);
    }
}