use core::time::Duration;
use crate::error::RateLimiterError;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Quota {
limit: u32,
period: Duration,
burst: u32,
}
impl Quota {
#[must_use]
pub const fn per_second(limit: u32) -> Self {
Self {
limit,
period: Duration::from_secs(1),
burst: limit,
}
}
#[must_use]
pub const fn per_minute(limit: u32) -> Self {
Self {
limit,
period: Duration::from_secs(60),
burst: limit,
}
}
pub const fn rate(limit: u32, period: Duration) -> Result<Self, RateLimiterError> {
if limit == 0 {
return Err(RateLimiterError::ZeroQuota);
}
if period.is_zero() {
return Err(RateLimiterError::ZeroPeriod);
}
Ok(Self {
limit,
period,
burst: limit,
})
}
#[must_use]
pub const fn limit(&self) -> u32 {
self.limit
}
#[must_use]
pub const fn period(&self) -> Duration {
self.period
}
#[must_use]
pub const fn burst(&self) -> u32 {
self.burst
}
#[must_use]
pub const fn with_burst(mut self, burst: u32) -> Self {
self.burst = burst;
self
}
pub(crate) const fn from_parts(limit: u32, period: Duration, burst: u32) -> Self {
Self {
limit,
period,
burst,
}
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used)]
use super::Quota;
use crate::error::RateLimiterError;
use core::time::Duration;
#[test]
fn test_per_second_sets_one_second_period() {
let quota = Quota::per_second(10);
assert_eq!(quota.limit(), 10);
assert_eq!(quota.period(), Duration::from_secs(1));
}
#[test]
fn test_per_minute_sets_sixty_second_period() {
assert_eq!(Quota::per_minute(10).period(), Duration::from_secs(60));
}
#[test]
fn test_burst_defaults_to_limit_and_overrides() {
assert_eq!(Quota::per_second(10).burst(), 10);
assert_eq!(Quota::per_minute(10).burst(), 10);
let q = Quota::rate(10, Duration::from_secs(1)).unwrap();
assert_eq!(q.burst(), 10);
assert_eq!(q.with_burst(25).burst(), 25);
assert_eq!(q.with_burst(25).limit(), 10);
}
#[test]
fn test_rate_accepts_valid_values() {
let quota = Quota::rate(5, Duration::from_millis(100)).unwrap();
assert_eq!(quota.limit(), 5);
assert_eq!(quota.period(), Duration::from_millis(100));
}
#[test]
fn test_rate_rejects_zero_limit() {
assert_eq!(
Quota::rate(0, Duration::from_secs(1)),
Err(RateLimiterError::ZeroQuota)
);
}
#[test]
fn test_rate_rejects_zero_period() {
assert_eq!(
Quota::rate(10, Duration::ZERO),
Err(RateLimiterError::ZeroPeriod)
);
}
#[test]
fn test_per_second_zero_limit_is_allowed() {
assert_eq!(Quota::per_second(0).limit(), 0);
}
}