gpu-handle-types 0.2.0

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

// `BudgetCapsSet`'s `cuda` / `opencl` / `wgpu` fields are feature-gated, so
// the `..Default::default()` tails below are only redundant on a build with
// none of those features.
#![allow(clippy::needless_update)]

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

use gpu_handle_types::{BackendKind, BudgetCaps, BudgetCapsSet, BudgetPressure, Error, MemoryBudget, PressureCallback};

#[test]
fn reserve_and_drop_updates_usage() {
    let budget = MemoryBudget::new(BudgetCapsSet {
        cpu: BudgetCaps { soft_cap_bytes: None, hard_cap_bytes: Some(1024) },
        ..Default::default()
    });
    {
        let _r = budget.try_reserve(BackendKind::Cpu, 512).unwrap();
        assert_eq!(budget.current_usage(BackendKind::Cpu), 512);
    }
    assert_eq!(budget.current_usage(BackendKind::Cpu), 0);
}

#[test]
fn reserve_over_hard_cap_fires_callbacks_and_fails() {
    let budget = MemoryBudget::new(BudgetCapsSet {
        cpu: BudgetCaps { soft_cap_bytes: None, hard_cap_bytes: Some(100) },
        ..Default::default()
    });
    let fired = Arc::new(AtomicU64::new(0));
    let f = fired.clone();
    let cb: PressureCallback = Arc::new(move |_| {
        f.fetch_add(1, Ordering::Relaxed);
    });
    budget.register_pressure_callback(cb);

    let err = budget.try_reserve(BackendKind::Cpu, 101).unwrap_err();
    assert!(matches!(err, Error::OutOfGpuMemory { .. }));
    assert_eq!(fired.load(Ordering::Relaxed), 1);
}

#[test]
fn soft_cap_warns_but_succeeds() {
    let budget = MemoryBudget::new(BudgetCapsSet {
        cpu: BudgetCaps { soft_cap_bytes: Some(50), hard_cap_bytes: Some(1000) },
        ..Default::default()
    });
    let fired = Arc::new(AtomicU64::new(0));
    let f = fired.clone();
    let cb: PressureCallback = Arc::new(move |_| {
        f.fetch_add(1, Ordering::Relaxed);
    });
    budget.register_pressure_callback(cb);

    // Soft breach: the pressure callbacks fire (soft is warning-only),
    // and the commit still succeeds because only the hard cap blocks.
    let _r = budget.try_reserve(BackendKind::Cpu, 80).unwrap();
    assert_eq!(budget.pressure(BackendKind::Cpu, 0), BudgetPressure::Soft);
}

#[test]
fn panicking_callback_is_caught() {
    let budget = MemoryBudget::new(BudgetCapsSet {
        cpu: BudgetCaps { soft_cap_bytes: None, hard_cap_bytes: Some(10) },
        ..Default::default()
    });
    let cb: PressureCallback = Arc::new(|_| panic!("intentional"));
    budget.register_pressure_callback(cb);
    // Must return an error, not propagate the panic.
    assert!(budget.try_reserve(BackendKind::Cpu, 100).is_err());
}

#[test]
fn uncapped_backend_is_allocation_free_on_reserve() {
    // No caps → fast path (no mutex, no vec).
    let budget = MemoryBudget::new(BudgetCapsSet::default());
    let r = budget.try_reserve(BackendKind::Cpu, 1 << 30).unwrap();
    assert_eq!(r.bytes(), 1 << 30);
}