use core::{convert::TryInto, fmt, num::NonZeroU32, time::Duration};
use crate::nanos::Nanos;
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub struct Quota {
pub(crate) max_burst: NonZeroU32,
pub(crate) replenish_1_per: Duration,
}
impl Quota {
pub fn per_second<B>(max_burst: B) -> Self
where
B: TryInto<NonZeroU32>,
B::Error: fmt::Debug,
{
let max_burst = max_burst.try_into().unwrap();
let replenish_interval_ns = Duration::from_secs(1).as_nanos() / (max_burst.get() as u128);
Self::new(max_burst, replenish_interval_ns)
}
pub fn per_minute<B>(max_burst: B) -> Self
where
B: TryInto<NonZeroU32>,
B::Error: fmt::Debug,
{
let max_burst = max_burst.try_into().unwrap();
let replenish_interval_ns = Duration::from_secs(60).as_nanos() / (max_burst.get() as u128);
Self::new(max_burst, replenish_interval_ns)
}
pub fn per_hour<B>(max_burst: B) -> Self
where
B: TryInto<NonZeroU32>,
B::Error: fmt::Debug,
{
let max_burst = max_burst.try_into().unwrap();
let replenish_interval_ns = Duration::from_secs(60 * 60).as_nanos() / (max_burst.get() as u128);
Self::new(max_burst, replenish_interval_ns)
}
pub fn with_period(replenish_1_per: Duration) -> Option<Quota> {
if replenish_1_per.as_nanos() == 0 {
None
} else {
Some(Quota {
max_burst: NonZeroU32::new(1).unwrap(),
replenish_1_per,
})
}
}
pub fn allow_burst<B>(mut self, max_burst: B) -> Self
where
B: TryInto<NonZeroU32>,
B::Error: fmt::Debug,
{
self.max_burst = max_burst.try_into().unwrap();
self
}
}
impl Quota {
fn new(max_burst: NonZeroU32, dur_ns: u128) -> Self {
Self {
max_burst,
replenish_1_per: Duration::from_nanos(dur_ns as u64),
}
}
pub(crate) const fn burst_size(&self) -> NonZeroU32 {
self.max_burst
}
#[cfg(test)]
const fn replenish_interval(&self) -> Duration {
self.replenish_1_per
}
}
impl Quota {
pub(crate) fn from_gcra_parameters(t: Nanos, tau: Nanos) -> Quota {
let max_burst = NonZeroU32::new((tau.as_u64() / t.as_u64()) as u32).unwrap();
let replenish_1_per = t.into();
Quota {
max_burst,
replenish_1_per,
}
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn time_multiples() {
let hourly = Quota::per_hour(1);
let minutely = Quota::per_minute(1);
let secondly = Quota::per_second(1);
assert_eq!(hourly.replenish_interval() / 60, minutely.replenish_interval());
assert_eq!(minutely.replenish_interval() / 60, secondly.replenish_interval());
}
#[test]
fn period_error_cases() {
assert!(Quota::with_period(Duration::from_secs(0)).is_none());
}
}