use std::{
sync::{
Arc, Weak,
atomic::{AtomicBool, AtomicU64, Ordering},
},
thread,
time::Duration,
};
use crate::{
BucketSize, HardLimitFactor, RateLimiterBuilder, SuppressionFactorCachePeriod, TrypemaError,
WindowSize,
builder::{CleanupConfig, ProviderConfig},
};
use super::{AbsoluteLocalRateLimiter, SuppressedLocalRateLimiter};
#[derive(Clone, Copy, Debug, Default)]
pub struct LocalRateLimiterBuilder {
provider: ProviderConfig,
cleanup: CleanupConfig,
}
impl RateLimiterBuilder for LocalRateLimiterBuilder {
type Provider = LocalRateLimiterProvider;
fn window_size(mut self, value: WindowSize) -> Self {
self.provider.window_size = value;
self
}
fn bucket_size(mut self, value: BucketSize) -> Self {
self.provider.bucket_size = value;
self
}
fn hard_limit_factor(mut self, value: HardLimitFactor) -> Self {
self.provider.hard_limit_factor = value;
self
}
fn suppression_factor_cache_period(mut self, value: SuppressionFactorCachePeriod) -> Self {
self.provider.suppression_factor_cache_period = value;
self
}
fn stale_after(mut self, value: Duration) -> Self {
self.cleanup.stale_after = value;
self
}
fn cleanup_interval(mut self, value: Duration) -> Self {
self.cleanup.interval = value;
self
}
fn cleanup_enabled(mut self, enabled: bool) -> Self {
self.cleanup.enabled = enabled;
self
}
fn build(self) -> Result<Arc<Self::Provider>, TrypemaError> {
let config = self.provider.validate()?;
let cleanup = self.cleanup.validate()?;
let provider = Arc::new(LocalRateLimiterProvider::new(config, cleanup));
if cleanup.enabled {
provider.start_cleanup_loop();
}
Ok(provider)
}
}
#[derive(Debug)]
pub struct LocalRateLimiterProvider {
absolute: AbsoluteLocalRateLimiter,
suppressed: SuppressedLocalRateLimiter,
cleanup: CleanupConfig,
is_cleanup_loop_running: AtomicBool,
cleanup_generation: AtomicU64,
}
impl LocalRateLimiterProvider {
pub fn builder() -> LocalRateLimiterBuilder {
LocalRateLimiterBuilder::default()
}
pub(crate) fn new(config: ProviderConfig, cleanup: CleanupConfig) -> Self {
Self {
absolute: AbsoluteLocalRateLimiter::new(config),
suppressed: SuppressedLocalRateLimiter::new(config),
cleanup,
is_cleanup_loop_running: AtomicBool::new(false),
cleanup_generation: AtomicU64::new(0),
}
}
pub fn start_cleanup_loop(self: &Arc<Self>) {
if self.is_cleanup_loop_running.swap(true, Ordering::AcqRel) {
return;
}
let generation = self
.cleanup_generation
.fetch_add(1, Ordering::AcqRel)
.wrapping_add(1);
let provider = Arc::downgrade(self);
thread::spawn(move || Self::run_cleanup_loop(provider, generation));
}
fn run_cleanup_loop(weak_provider: Weak<Self>, generation: u64) {
loop {
let Some(provider) = weak_provider.upgrade() else {
break;
};
let interval = provider.cleanup.interval;
drop(provider);
thread::sleep(interval);
let Some(provider) = weak_provider.upgrade() else {
break;
};
if !provider.is_cleanup_loop_running.load(Ordering::Acquire)
|| provider.cleanup_generation.load(Ordering::Acquire) != generation
{
break;
}
provider.cleanup(provider.cleanup.stale_after_ms());
}
}
pub fn stop_cleanup_loop(&self) {
self.cleanup_generation.fetch_add(1, Ordering::AcqRel);
self.is_cleanup_loop_running.store(false, Ordering::Release);
}
pub fn absolute(&self) -> &AbsoluteLocalRateLimiter {
&self.absolute
}
pub fn suppressed(&self) -> &SuppressedLocalRateLimiter {
&self.suppressed
}
pub(crate) fn cleanup(&self, stale_after_ms: u64) {
self.absolute.cleanup(stale_after_ms);
self.suppressed.cleanup(stale_after_ms);
}
}