hyperhive/hive-jobq/src/guard.rs

//! RAII guard objects over [`ResourceTable`] — owning resource grants.
//!
//! A running node acquires its resources through [`SharedResources::acquire`],
//! which hands back a [`ResourceGuard`] owning those units. Dropping the guard
//! releases exactly what it acquired, so a node's resources are freed when its
//! grant goes out of scope — there is no explicit release call to forget.
//!
//! Re-entrancy (a sub-node reusing a resource its ancestor group already holds)
//! is not expressed here: the scheduler tracks it with a single borrow slot per
//! `(holder, resource)` and never re-acquires, so these guards are always owning.
//!
//! Single-owner by design: the scheduler drives one settle loop, so the shared
//! table is `Rc<RefCell<…>>` (single-threaded interior mutability), not
//! `Arc<Mutex<…>>` — there is no cross-thread contention to guard against.

use std::cell::RefCell;
use std::rc::Rc;

use std::hash::Hash;

use crate::resources::ResourceTable;

/// A [`ResourceTable`] shared between the scheduler and the live guards that
/// release back into it on drop. Cheap to clone — an `Rc` refcount bump.
#[derive(Debug, Clone)]
pub struct SharedResources<R>(Rc<RefCell<ResourceTable<R>>>);

impl<R: Clone + Eq + Hash> Default for SharedResources<R> {
    fn default() -> Self {
        Self::new(ResourceTable::new())
    }
}

impl<R: Clone + Eq + Hash> SharedResources<R> {
    /// Wrap an existing table so guards can release into it.
    #[must_use]
    pub fn new(table: ResourceTable<R>) -> Self {
        Self(Rc::new(RefCell::new(table)))
    }

    /// Atomically acquire every requested `(name, count)` or none of them.
    ///
    /// Returns an owning [`ResourceGuard`] (releases on drop) when the whole
    /// request fits in what is available right now; returns `None` and leaves
    /// the table completely untouched otherwise. Duplicate names are summed and
    /// an over-capacity request can never succeed — same all-or-nothing
    /// semantics as [`ResourceTable::try_acquire_all`].
    #[must_use]
    pub fn acquire(&self, reqs: Vec<(R, u32)>) -> Option<ResourceGuard<R>> {
        if self.0.borrow_mut().try_acquire_all(&reqs) {
            Some(ResourceGuard {
                table: self.clone(),
                reqs,
            })
        } else {
            None
        }
    }

    /// Observe the underlying table — test-only (the scheduler's tests assert
    /// on resource availability). Gated `#[cfg(test)]` so it is compiled out of
    /// the shipped crate: no consumer can reach the raw table through it.
    #[cfg(test)]
    pub fn with<T>(&self, f: impl FnOnce(&ResourceTable<R>) -> T) -> T {
        f(&self.0.borrow())
    }
}

/// An RAII grant of resources: dropping it releases exactly the units it
/// acquired back into the shared table.
#[derive(Debug)]
pub struct ResourceGuard<R: Clone + Eq + Hash> {
    table: SharedResources<R>,
    reqs: Vec<(R, u32)>,
}

impl<R: Clone + Eq + Hash> ResourceGuard<R> {
    /// The `(name, count)` units this guard releases on drop.
    #[must_use]
    pub fn held(&self) -> &[(R, u32)] {
        &self.reqs
    }
}

impl<R: Clone + Eq + Hash> Drop for ResourceGuard<R> {
    fn drop(&mut self) {
        self.table.0.borrow_mut().release_all(&self.reqs);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn res(name: &str) -> String {
        name.to_owned()
    }

    fn shared_with(slots: u32) -> SharedResources<String> {
        let mut t = ResourceTable::new();
        t.set_capacity(res("build-slot"), slots);
        SharedResources::new(t)
    }

    #[test]
    fn owning_guard_releases_on_drop() {
        let sr = shared_with(2);
        let slot = res("build-slot");
        {
            let g = sr.acquire(vec![(slot.clone(), 2)]).expect("fits");
            assert_eq!(g.held(), &[(slot.clone(), 2)]);
            // Both units held → any further acquire fails.
            assert!(sr.acquire(vec![(slot.clone(), 1)]).is_none());
        } // guard dropped here → its units are released
        // Full capacity is available again.
        assert!(sr.acquire(vec![(slot.clone(), 2)]).is_some());
    }

    #[test]
    fn acquire_returns_none_and_leaves_table_untouched_when_it_does_not_fit() {
        let sr = shared_with(1);
        let slot = res("build-slot");
        let held = sr.acquire(vec![(slot.clone(), 1)]).expect("first fits");
        assert!(sr.acquire(vec![(slot.clone(), 1)]).is_none());
        // The failed acquire took nothing extra: dropping the one real grant
        // frees exactly one unit, so a single-unit acquire then fits.
        drop(held);
        assert!(sr.acquire(vec![(slot.clone(), 1)]).is_some());
    }
}