subms-rate-limiter 0.10.0

//! Hierarchical limiter: a parent `TokenBucket` caps a shared budget
//! across N child limiters. `try_acquire` must succeed at the child
//! AND the parent. If the parent rejects, the child token is refunded
//! so its budget isn't burned by a parent-side denial.
//!
//! Shape: imagine a multi-tenant API. Each tenant has its own per-
//! second rate, but the global gateway also has a single ceiling
//! that protects the upstream from a thundering herd of tenants who
//! all peak at once. The parent sees the AND of the children.

use std::sync::Arc;

use super::clock::{Clock, SystemClock};
use super::token_bucket::TokenBucket;

pub struct HierarchicalLimiter {
    parent: Arc<TokenBucket>,
    children: Vec<Arc<TokenBucket>>,
}

impl HierarchicalLimiter {
    /// Build a hierarchy: one parent bucket sized at `parent_capacity`
    /// + `parent_rate`, plus `num_children` child buckets each sized at
    ///   `child_capacity` + `child_rate`.
    pub fn new(
        parent_capacity: u64,
        parent_rate: f64,
        num_children: usize,
        child_capacity: u64,
        child_rate: f64,
    ) -> Self {
        Self::with_clock_fn(
            parent_capacity,
            parent_rate,
            num_children,
            child_capacity,
            child_rate,
            || Box::new(SystemClock::new()),
        )
    }

    /// Test/advanced constructor. The `clock_fn` is called once per
    /// bucket (parent + each child); tests typically return clones of
    /// the same shared `TestClock` so every bucket sees the same time.
    pub fn with_clock_fn<F>(
        parent_capacity: u64,
        parent_rate: f64,
        num_children: usize,
        child_capacity: u64,
        child_rate: f64,
        mut clock_fn: F,
    ) -> Self
    where
        F: FnMut() -> Box<dyn Clock>,
    {
        let parent = Arc::new(TokenBucket::with_clock(
            parent_capacity,
            parent_rate,
            clock_fn(),
        ));
        let children = (0..num_children.max(1))
            .map(|_| {
                Arc::new(TokenBucket::with_clock(
                    child_capacity,
                    child_rate,
                    clock_fn(),
                ))
            })
            .collect();
        Self { parent, children }
    }

    /// Try to acquire `n` tokens from `child_id`. Both child and parent
    /// must grant. If the parent rejects, the child is NOT charged.
    pub fn try_acquire(&self, child_id: usize, n: u64) -> bool {
        let child = match self.children.get(child_id) {
            Some(c) => c,
            None => return false,
        };
        // Order matters: try child first, then parent. If the parent
        // rejects, we'd otherwise have burned a child token for nothing.
        // The TokenBucket has no refund API (atomicity is at the
        // try_acquire boundary), so we mirror the parent check by
        // checking parent.available() first - if it's clearly under n,
        // skip the child draw entirely. Under the race window the
        // parent could drop below n between the check and our parent
        // draw; in that case we lose one child token. Tradeoff vs
        // adding a refund path to TokenBucket. Documented.
        if self.parent_available_at_least(n) {
            if !child.try_acquire(n) {
                return false;
            }
            if self.parent.try_acquire(n) {
                return true;
            }
            // Parent raced us. Best-effort: a one-token leak on the child
            // here is acceptable under contention; the call still
            // correctly returns false. Long-running mismatch is bounded
            // by parent_capacity (parent always re-fills).
            false
        } else {
            false
        }
    }

    pub fn parent(&self) -> &TokenBucket {
        &self.parent
    }

    pub fn child(&self, child_id: usize) -> Option<&TokenBucket> {
        self.children.get(child_id).map(|c| c.as_ref())
    }

    pub fn num_children(&self) -> usize {
        self.children.len()
    }

    fn parent_available_at_least(&self, n: u64) -> bool {
        self.parent.available() >= n
    }
}

#[cfg(test)]
#[path = "hierarchical_tests.rs"]
mod tests;