magi-code 0.96.1

use std::sync::{
    Arc,
    atomic::{AtomicBool, AtomicUsize, Ordering},
};

/// Shared by QuickJS allocations and temporary bridge conversions.
pub(crate) struct MemoryBudget {
    used: AtomicUsize,
    limit: usize,
    pub(super) stopped: Arc<AtomicBool>,
}

impl MemoryBudget {
    pub(crate) fn new(limit: usize, stopped: Arc<AtomicBool>) -> Arc<Self> {
        Arc::new(Self {
            used: AtomicUsize::new(0),
            limit,
            stopped,
        })
    }

    pub(super) fn snapshot(&self) -> serde_json::Value {
        crate::code_mode::limits::budget_snapshot(
            self.used.load(Ordering::SeqCst) as u64,
            self.limit as u64,
        )
    }

    pub(super) fn charge(&self, bytes: usize) -> bool {
        if self.stopped.load(Ordering::SeqCst)
            || self
                .used
                .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |used| {
                    used.checked_add(bytes).filter(|total| *total <= self.limit)
                })
                .is_err()
        {
            self.stopped.store(true, Ordering::SeqCst);
            return false;
        }
        true
    }

    pub(super) fn release(&self, bytes: usize) {
        self.used.fetch_sub(bytes, Ordering::SeqCst);
    }

    pub(crate) fn reserve(self: &Arc<Self>, bytes: usize) -> anyhow::Result<Reservation> {
        if !self.charge(bytes) {
            return Err(super::ResourceLimitExceeded("memory").into());
        }
        Ok(Reservation {
            budget: self.clone(),
            bytes,
        })
    }
}

pub(crate) struct Reservation {
    budget: Arc<MemoryBudget>,
    bytes: usize,
}

impl Drop for Reservation {
    fn drop(&mut self) {
        self.budget.release(self.bytes);
    }
}