hyperhive/hive-jobq/src/resources.rs

//! Named counting-semaphore resources — the [`Dep::Resource`] side of the model.
//!
//! A [`ResourceTable`] tracks a set of named counters. Each name has a
//! *capacity* (how many units exist) and a *held* count (how many are currently
//! in use); a name that was never configured is assumed to have capacity 1 and
//! is created lazily on first use — mara's "unknown resources assumed
//! available-once". Typical names: `build-slot` (capacity = number of build
//! slots), `agent/<name>` (capacity 1 — the per-agent lifecycle lock).
//!
//! The one operation that matters is [`ResourceTable::try_acquire_all`]: it
//! takes *all* of a node's resource requests and either grants every one or
//! grants none, touching nothing on failure. Because a node acquires all its
//! resources atomically at start (never holds one while waiting for another),
//! there is no hold-and-wait and therefore no deadlock and no cycle detection.
//!
//! This table is runtime scheduler state, not persisted: on restart the held
//! counts are rederived from the graph's currently-running nodes, which are the
//! source of truth. Guard objects (RAII release, recursive re-entrancy within a
//! group) wrap this table in a follow-up; here it is the raw counter machinery.
//!
//! [`Dep::Resource`]: crate::Dep::Resource

use std::collections::HashMap;
use std::hash::Hash;

/// A set of named counting semaphores.
///
/// Configure known capacities with [`ResourceTable::set_capacity`]; any name
/// left unconfigured has the default capacity (1). Acquire and release move
/// units atomically via [`ResourceTable::try_acquire_all`] /
/// [`ResourceTable::release_all`].
#[derive(Debug, Clone)]
pub struct ResourceTable<R> {
    /// Configured capacities, keyed by resource. Missing ⇒ `default_capacity`.
    capacities: HashMap<R, u32>,
    /// Units currently held, keyed by resource. Missing ⇒ 0.
    held: HashMap<R, u32>,
    /// Capacity assumed for a resource with no configured entry.
    default_capacity: u32,
}

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

impl<R: Clone + Eq + Hash> ResourceTable<R> {
    /// An empty table whose unconfigured names default to capacity 1.
    #[must_use]
    pub fn new() -> Self {
        Self {
            capacities: HashMap::new(),
            held: HashMap::new(),
            default_capacity: 1,
        }
    }

    /// Configure the capacity of a named resource (e.g. `build-slot` = N).
    ///
    /// Overwrites any previous capacity for that name. Lowering capacity below
    /// the currently-held count is allowed — the table simply reports zero
    /// available until enough is released; it never rejects a config change.
    pub fn set_capacity(&mut self, name: R, capacity: u32) {
        self.capacities.insert(name, capacity);
    }

    /// The capacity of a name — its configured value, or the default (1) if it
    /// was never configured.
    #[must_use]
    pub fn capacity(&self, name: &R) -> u32 {
        self.capacities
            .get(name)
            .copied()
            .unwrap_or(self.default_capacity)
    }

    /// Units of a name currently held (0 if none).
    #[must_use]
    pub fn held(&self, name: &R) -> u32 {
        self.held.get(name).copied().unwrap_or(0)
    }

    /// Units of a name available to acquire right now (`capacity - held`).
    #[must_use]
    pub fn available(&self, name: &R) -> u32 {
        self.capacity(name).saturating_sub(self.held(name))
    }

    /// Atomically acquire every requested `(name, count)` or none of them.
    ///
    /// Duplicate names in `reqs` are summed, so `[(a, 1), (a, 1)]` asks for two
    /// units of `a`. Returns `true` and records the units as held only if every
    /// aggregated request fits in what is available *now*; otherwise returns
    /// `false` and leaves the table completely untouched. A request for more
    /// units than a name's capacity can therefore never succeed — the scheduler
    /// should reject such a node at insert time so it does not wait forever.
    pub(crate) fn try_acquire_all(&mut self, reqs: &[(R, u32)]) -> bool {
        let wanted = aggregate(reqs);
        // All-or-nothing: bail before mutating if any request cannot be met.
        for (name, &count) in &wanted {
            if count > self.available(name) {
                return false;
            }
        }
        for (name, count) in wanted {
            *self.held.entry(name.clone()).or_insert(0) += count;
        }
        true
    }

    /// Release every requested `(name, count)`, returning the units to the pool.
    ///
    /// Duplicate names are summed. Releasing more than is held saturates at zero
    /// rather than underflowing, so a double release is harmless.
    pub(crate) fn release_all(&mut self, reqs: &[(R, u32)]) {
        for (name, count) in aggregate(reqs) {
            if let Some(h) = self.held.get_mut(name) {
                *h = h.saturating_sub(count);
            }
        }
    }
}

/// Sum a request list into per-name totals so duplicate names are one entry.
fn aggregate<R: Eq + Hash>(reqs: &[(R, u32)]) -> HashMap<&R, u32> {
    let mut wanted: HashMap<&R, u32> = HashMap::new();
    for (name, count) in reqs {
        *wanted.entry(name).or_insert(0) += *count;
    }
    wanted
}

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

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

    #[test]
    fn unconfigured_name_defaults_to_capacity_one() {
        let t = ResourceTable::new();
        let a = res("agent/foo");
        assert_eq!(t.capacity(&a), 1);
        assert_eq!(t.available(&a), 1);
        assert_eq!(t.held(&a), 0);
    }

    #[test]
    fn set_capacity_overrides_default() {
        let mut t = ResourceTable::new();
        let slot = res("build-slot");
        t.set_capacity(slot.clone(), 3);
        assert_eq!(t.capacity(&slot), 3);
        assert_eq!(t.available(&slot), 3);
    }

    #[test]
    fn acquire_all_or_nothing_leaves_table_untouched_on_failure() {
        let mut t = ResourceTable::new();
        let slot = res("build-slot");
        let agent = res("agent/foo");
        t.set_capacity(slot.clone(), 2);
        // First node grabs both slots and the agent lock — fits exactly.
        assert!(t.try_acquire_all(&[(slot.clone(), 2), (agent.clone(), 1)]));
        assert_eq!(t.held(&slot), 2);
        assert_eq!(t.held(&agent), 1);
        // Second node wants a slot (none free) + the agent lock (held): must
        // fail AND must not have grabbed the agent lock on the way.
        assert!(!t.try_acquire_all(&[(slot.clone(), 1), (agent.clone(), 1)]));
        assert_eq!(t.held(&slot), 2);
        assert_eq!(t.held(&agent), 1);
    }

    #[test]
    fn release_returns_units_to_the_pool() {
        let mut t = ResourceTable::new();
        let slot = res("build-slot");
        t.set_capacity(slot.clone(), 2);
        assert!(t.try_acquire_all(&[(slot.clone(), 2)]));
        assert_eq!(t.available(&slot), 0);
        t.release_all(&[(slot.clone(), 2)]);
        assert_eq!(t.available(&slot), 2);
        // Over-release saturates, never underflows.
        t.release_all(&[(slot.clone(), 5)]);
        assert_eq!(t.held(&slot), 0);
    }

    #[test]
    fn request_over_capacity_never_acquires() {
        let mut t = ResourceTable::new();
        let slot = res("build-slot");
        t.set_capacity(slot.clone(), 2);
        assert!(!t.try_acquire_all(&[(slot.clone(), 3)]));
        assert_eq!(t.held(&slot), 0);
    }

    #[test]
    fn duplicate_names_in_one_request_are_summed() {
        let mut t = ResourceTable::new();
        let a = res("thing");
        // Default capacity 1: two units in one request cannot both fit.
        assert!(!t.try_acquire_all(&[(a.clone(), 1), (a.clone(), 1)]));
        assert_eq!(t.held(&a), 0);
        // Raise capacity to 2 and the same aggregated request fits.
        t.set_capacity(a.clone(), 2);
        assert!(t.try_acquire_all(&[(a.clone(), 1), (a.clone(), 1)]));
        assert_eq!(t.held(&a), 2);
    }
}