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::ApproximateSlidingWindowCore;
use rate_guard_core::SimpleRateLimitError;
pub struct ApproximateSlidingWindow<P: Precision, T: TimeSource> {
core: ApproximateSlidingWindowCore,
time_source: T,
_precision: PhantomData<P>,
}
#[derive(Debug)]
pub struct ApproximateSlidingWindowBuilder {
capacity: Option<Uint>,
window_duration: Option<Duration>,
}
pub struct ApproximateSlidingWindowBuilderWithTime<T: TimeSource> {
capacity: Option<Uint>,
window_duration: Option<Duration>,
time_source: T,
}
pub struct ConfiguredApproximateSlidingWindowBuilder<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 ApproximateSlidingWindow<P, T>
where
T: std::fmt::Debug,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ApproximateSlidingWindow")
.field("time_source", &self.time_source)
.field("_precision", &std::any::type_name::<P>())
.finish_non_exhaustive()
}
}
impl<T: TimeSource> std::fmt::Debug for ApproximateSlidingWindowBuilderWithTime<T>
where
T: std::fmt::Debug,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ApproximateSlidingWindowBuilderWithTime")
.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 ConfiguredApproximateSlidingWindowBuilder<P, T>
where
T: std::fmt::Debug,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ConfiguredApproximateSlidingWindowBuilder")
.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 ApproximateSlidingWindow<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 ApproximateSlidingWindowBuilder {
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) -> ApproximateSlidingWindowBuilderWithTime<T> {
ApproximateSlidingWindowBuilderWithTime {
capacity: self.capacity,
window_duration: self.window_duration,
time_source,
}
}
}
impl Default for ApproximateSlidingWindowBuilder {
fn default() -> Self {
Self::builder()
}
}
impl<T: TimeSource> ApproximateSlidingWindowBuilderWithTime<T> {
pub fn with_precision<P: Precision>(self) -> ConfiguredApproximateSlidingWindowBuilder<P, T> {
ConfiguredApproximateSlidingWindowBuilder {
capacity: self.capacity,
window_duration: self.window_duration,
time_source: self.time_source,
_precision: PhantomData,
}
}
}
impl<P: Precision, T: TimeSource> ConfiguredApproximateSlidingWindowBuilder<P, T> {
pub fn build(self) -> BuildResult<ApproximateSlidingWindow<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_ticks = P::to_ticks(window_duration);
let core = ApproximateSlidingWindowCore::new(capacity, window_ticks);
Ok(ApproximateSlidingWindow {
core,
time_source: self.time_source,
_precision: PhantomData,
})
}
}
pub use ApproximateSlidingWindow as PrecisionApproximateSlidingWindow;