use std::time::Duration;
use std::marker::PhantomData;
use crate::limits::RateLimit;
use crate::{Uint, SimpleRateLimitResult, RateLimitResult, BuildResult};
use crate::precision::Precision;
use crate::time_source::TimeSource;
use crate::error::{RateLimitError, BuildError};
use rate_guard_core::cores::FixedWindowCounterCore;
use rate_guard_core::SimpleRateLimitError;
pub struct FixedWindowCounter<P: Precision, T: TimeSource> {
core: FixedWindowCounterCore,
time_source: T,
_precision: PhantomData<P>,
}
#[derive(Debug)]
pub struct FixedWindowCounterBuilder {
capacity: Option<Uint>,
window_duration: Option<Duration>,
}
pub struct FixedWindowCounterBuilderWithTime<T: TimeSource> {
capacity: Option<Uint>,
window_duration: Option<Duration>,
time_source: T,
}
pub struct ConfiguredFixedWindowCounterBuilder<P: Precision, T: TimeSource> {
capacity: Option<Uint>,
window_duration: Option<Duration>,
time_source: T,
_precision: PhantomData<P>,
}
impl<P: Precision, T: TimeSource> std::fmt::Debug for FixedWindowCounter<P, T>
where
T: std::fmt::Debug,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("FixedWindowCounter")
.field("time_source", &self.time_source)
.field("_precision", &std::any::type_name::<P>())
.finish_non_exhaustive()
}
}
impl<T: TimeSource> std::fmt::Debug for FixedWindowCounterBuilderWithTime<T>
where
T: std::fmt::Debug,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("FixedWindowCounterBuilderWithTime")
.field("capacity", &self.capacity)
.field("window_duration", &self.window_duration)
.field("time_source", &self.time_source)
.finish()
}
}
impl<P: Precision, T: TimeSource> std::fmt::Debug for ConfiguredFixedWindowCounterBuilder<P, T>
where
T: std::fmt::Debug,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ConfiguredFixedWindowCounterBuilder")
.field("capacity", &self.capacity)
.field("window_duration", &self.window_duration)
.field("time_source", &self.time_source)
.field("_precision", &std::any::type_name::<P>())
.finish()
}
}
impl<P: Precision, T: TimeSource> RateLimit for FixedWindowCounter<P, T> {
#[inline(always)]
fn try_acquire(&self, tokens: Uint) -> SimpleRateLimitResult {
let elapsed = self.time_source.now();
let current_tick = P::to_ticks(elapsed);
self.core.try_acquire_at(current_tick, tokens)
}
#[inline(always)]
fn try_acquire_verbose(&self, tokens: Uint) -> RateLimitResult {
let elapsed = self.time_source.now();
let current_tick = P::to_ticks(elapsed);
self.core.try_acquire_verbose_at(current_tick, tokens)
.map_err(|e| RateLimitError::from_core_error(e, |ticks| P::from_ticks(ticks)))
}
#[inline(always)]
fn capacity_remaining(&self) -> Result<Uint, SimpleRateLimitError> {
let elapsed = self.time_source.now();
let current_tick = P::to_ticks(elapsed);
self.core.capacity_remaining(current_tick)
}
}
impl FixedWindowCounterBuilder {
pub fn builder() -> Self {
Self {
capacity: None,
window_duration: None,
}
}
pub fn capacity(mut self, capacity: Uint) -> Self {
self.capacity = Some(capacity);
self
}
pub fn window_duration(mut self, duration: Duration) -> Self {
self.window_duration = Some(duration);
self
}
pub fn with_time<T: TimeSource>(self, time_source: T) -> FixedWindowCounterBuilderWithTime<T> {
FixedWindowCounterBuilderWithTime {
capacity: self.capacity,
window_duration: self.window_duration,
time_source,
}
}
}
impl Default for FixedWindowCounterBuilder {
fn default() -> Self {
Self::builder()
}
}
impl<T: TimeSource> FixedWindowCounterBuilderWithTime<T> {
pub fn with_precision<P: Precision>(self) -> ConfiguredFixedWindowCounterBuilder<P, T> {
ConfiguredFixedWindowCounterBuilder {
capacity: self.capacity,
window_duration: self.window_duration,
time_source: self.time_source,
_precision: PhantomData,
}
}
}
impl<P: Precision, T: TimeSource> ConfiguredFixedWindowCounterBuilder<P, T> {
pub fn build(self) -> BuildResult<FixedWindowCounter<P, T>> {
let capacity = self.capacity.ok_or(BuildError::MissingArgument("capacity"))?;
let window_duration = self.window_duration.ok_or(BuildError::MissingArgument("window_duration"))?;
if capacity == 0 {
return Err(BuildError::InvalidArgument {
field: "capacity",
reason: "must be greater than 0"
});
}
if window_duration == Duration::ZERO {
return Err(BuildError::InvalidArgument {
field: "window_duration",
reason: "must be greater than zero"
});
}
let window_size_ticks = P::to_ticks(window_duration);
let core = FixedWindowCounterCore::new(capacity, window_size_ticks);
Ok(FixedWindowCounter {
core,
time_source: self.time_source,
_precision: PhantomData,
})
}
}
pub use FixedWindowCounter as PrecisionFixedWindowCounter;