subms-rate-limiter 0.5.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)]
mod tests {
    use std::sync::Arc;

    use super::*;
    use crate::features::clock::TestClock;

    struct ArcClock(Arc<TestClock>);
    impl Clock for ArcClock {
        fn now_ns(&self) -> u64 {
            self.0.now_ns()
        }
    }

    fn build(
        parent_cap: u64,
        parent_rate: f64,
        children: usize,
        child_cap: u64,
        child_rate: f64,
    ) -> (HierarchicalLimiter, Arc<TestClock>) {
        let clk = Arc::new(TestClock::new());
        let c = clk.clone();
        let h = HierarchicalLimiter::with_clock_fn(
            parent_cap,
            parent_rate,
            children,
            child_cap,
            child_rate,
            || Box::new(ArcClock(c.clone())),
        );
        (h, clk)
    }

    #[test]
    fn parent_caps_total_across_children() {
        // Parent allows 5 total; each child can do 10. So even though
        // two children could grant up to 20 between them, the parent
        // caps the sum at 5.
        let (h, _clk) = build(5, 0.0, 2, 10, 0.0);
        let mut granted = 0;
        for _ in 0..10 {
            if h.try_acquire(0, 1) {
                granted += 1;
            }
        }
        for _ in 0..10 {
            if h.try_acquire(1, 1) {
                granted += 1;
            }
        }
        assert_eq!(granted, 5, "parent must cap total at 5");
    }

    #[test]
    fn child_caps_independent_when_parent_has_budget() {
        // Parent generous (1000); child capped at 3.
        let (h, _clk) = build(1000, 0.0, 1, 3, 0.0);
        for _ in 0..3 {
            assert!(h.try_acquire(0, 1));
        }
        assert!(!h.try_acquire(0, 1), "child capacity exhausted");
    }

    #[test]
    fn unknown_child_id_rejects() {
        let (h, _clk) = build(10, 0.0, 1, 5, 0.0);
        assert!(!h.try_acquire(99, 1), "out-of-range child id should reject");
    }

    #[test]
    fn refill_after_parent_exhaustion_unblocks_children() {
        // Parent 5 cap, 50 tokens/sec = 1 per 20 ms.
        let (h, clk) = build(5, 50.0, 2, 10, 0.0);
        // Drain parent via child 0.
        for _ in 0..5 {
            assert!(h.try_acquire(0, 1));
        }
        assert!(!h.try_acquire(1, 1), "parent exhausted");
        // 100 ms -> 5 tokens refilled.
        clk.advance_ms(100);
        let mut got = 0;
        for _ in 0..10 {
            if h.try_acquire(1, 1) {
                got += 1;
            }
        }
        assert_eq!(got, 5, "exactly 5 parent tokens refilled (then exhausted)");
    }

    #[test]
    fn batch_acquire_atomic_at_both_levels() {
        let (h, _clk) = build(10, 0.0, 1, 10, 0.0);
        assert!(h.try_acquire(0, 7));
        assert!(!h.try_acquire(0, 5), "would exceed parent capacity now");
        assert!(h.try_acquire(0, 3), "exactly 3 left should grant");
        assert!(!h.try_acquire(0, 1));
    }

    #[test]
    fn parent_and_child_accessors_expose_underlying_buckets() {
        let (h, _clk) = build(7, 0.0, 2, 3, 0.0);
        assert_eq!(h.parent().capacity(), 7);
        assert_eq!(h.child(0).map(|c| c.capacity()), Some(3));
        assert_eq!(h.child(99).map(|c| c.capacity()), None);
        assert_eq!(h.num_children(), 2);
    }
}