quota 0.3.3

Fastest Lane-parallel Rate-limiter for Rust
Documentation
use super::{QuotaPolicy, QuotaPool, QuotaSlot};
use crate::refill_rate::RefillRate;
use std::sync::{Arc, Barrier};
use std::time::Duration;

#[test]
fn pool_consumes_initial_tokens_for_one_key() {
    let policy = QuotaPolicy::new()
        .set_capacity(10.0)
        .set_refill_rate(RefillRate::per_sec(3));

    let pool = QuotaPool::new(policy, 10);
    let mut results = vec![];

    for _ in 0..100 {
        results.push(pool.consume("testing", 1));
    }

    assert_eq!(results.iter().filter(|r| r.is_ok()).count(), 10);
    assert_eq!(results.iter().filter(|r| r.is_err()).count(), 90);
}

#[test]
fn pool_limits_one_key_across_threads() {
    let policy = QuotaPolicy::new().set_capacity(10.0);
    let pool = Arc::new(QuotaPool::new(policy, 10));
    let barrier = Arc::new(Barrier::new(100));
    let mut handles = Vec::new();

    for _ in 0..100 {
        let pool = Arc::clone(&pool);
        let barrier = Arc::clone(&barrier);
        handles.push(std::thread::spawn(move || {
            barrier.wait();
            pool.consume("testing", 1).is_ok()
        }));
    }

    let accepted = handles
        .into_iter()
        .map(|handle| handle.join().expect("worker thread panicked"))
        .filter(|accepted| *accepted)
        .count();

    assert_eq!(accepted, 10);
    assert_eq!(pool.len(), 1);
}

#[test]
fn pool_consumes_owned_key_and_keeps_entry() {
    let policy = QuotaPolicy::new().set_capacity(1.0);
    let pool = QuotaPool::new(policy, 1);

    assert!(pool.consume_owned(String::from("testing"), 1).is_ok());
    assert!(pool.consume("testing", 1).is_err());
    assert!(pool.contains_key("testing"));
}

#[test]
fn pool_removes_key_and_recreates_fresh_quota() {
    let policy = QuotaPolicy::new().set_capacity(1.0);
    let pool = QuotaPool::new(policy, 1);

    assert!(pool.consume("testing", 1).is_ok());
    assert!(pool.consume("testing", 1).is_err());
    assert!(pool.remove("testing"));
    assert!(!pool.contains_key("testing"));
    assert_eq!(pool.len(), 0);

    assert!(pool.consume("testing", 1).is_ok());
}

#[test]
fn pool_cache_does_not_use_removed_slot_after_reuse() {
    let policy = QuotaPolicy::new().set_capacity(1.0);
    let pool = QuotaPool::with_capacity(policy, 1, 2);

    assert!(pool.consume("a", 1).is_ok());
    assert!(pool.remove("a"));

    pool.insert_key("b");

    assert!(pool.consume("a", 1).is_ok());
    assert!(pool.consume("b", 1).is_ok());
}

#[test]
fn pool_tracks_spread_keys_independently() {
    let policy = QuotaPolicy::new().set_capacity(1.0);
    let keys: Vec<_> = (0..64).map(|id| format!("key-{id}")).collect();
    let pool = Arc::new(QuotaPool::with_capacity(policy, 1, keys.len()));
    pool.insert_keys(&keys);

    assert_eq!(pool.len(), keys.len());

    let mut handles = Vec::new();

    for key in keys {
        let pool = Arc::clone(&pool);
        handles.push(std::thread::spawn(move || {
            pool.consume(key.as_str(), 1).is_ok()
        }));
    }

    let accepted = handles
        .into_iter()
        .map(|handle| handle.join().expect("worker thread panicked"))
        .filter(|accepted| *accepted)
        .count();

    assert_eq!(accepted, 64);
    assert_eq!(pool.len(), 64);
}

#[test]
fn slot_skips_refill_before_interval() {
    let policy = QuotaPolicy::new()
        .set_capacity(10.0)
        .set_refill_rate(RefillRate::per_sec(10))
        .set_refill_interval(Duration::from_secs(1));
    let slot = QuotaSlot::new(&policy, 0);

    let result = slot.consume(&policy, 500_000_000, 1);

    assert!(result.is_err());
    assert_eq!(result.unwrap_err().current().available(), 0.0);
}

#[test]
fn slot_refills_after_interval() {
    let policy = QuotaPolicy::new()
        .set_capacity(72.0)
        .set_refill_rate(RefillRate::per_sec(10))
        .set_refill_interval(Duration::from_secs(1));
    let slot = QuotaSlot::new(&policy, 0);

    assert!(slot.consume(&policy, 500_000_000, 1).is_err());
    let result = slot
        .consume(&policy, 1_000_000_000, 1)
        .expect("interval refill should add tokens");

    assert_eq!(result.available(), 9.0);
}

#[test]
fn slot_keeps_initial_tokens_when_interval_skips_refill() {
    let policy = QuotaPolicy::new()
        .set_capacity(1.0)
        .set_refill_rate(RefillRate::per_sec(10))
        .set_refill_interval(Duration::from_secs(1));
    let slot = QuotaSlot::new(&policy, 10);

    let result = slot
        .consume(&policy, 500_000_000, 5)
        .expect("skipped refill should not clamp existing quota");

    assert_eq!(result.available(), 5.0);
}