use crate::error::LimiteronError;
use crate::limiters::traits::{Limiter, RateLimitSnapshot};
use async_trait::async_trait;
use std::sync::atomic::{AtomicU64, Ordering};
#[derive(Debug, Clone)]
pub struct AdaptiveThresholdConfig {
pub base_limit: u64,
pub min_limit: u64,
pub max_limit: u64,
pub window_samples: usize,
pub min_samples_for_adjust: usize,
pub error_rate_increase_threshold: f64,
pub error_rate_decrease_threshold: f64,
pub latency_increase_ms: u64,
pub latency_decrease_ms: u64,
pub decrease_ratio: f64,
pub increase_step: u64,
pub cooldown_samples: usize,
}
impl AdaptiveThresholdConfig {
pub fn validate(&self) -> Result<(), LimiteronError> {
if self.min_limit == 0 {
return Err(LimiteronError::ConfigError(
"adaptive-threshold: min_limit must be > 0".to_string(),
));
}
if self.min_limit > self.base_limit || self.base_limit > self.max_limit {
return Err(LimiteronError::ConfigError(
"adaptive-threshold: require min_limit <= base_limit <= max_limit".to_string(),
));
}
if !(0.0 < self.decrease_ratio && self.decrease_ratio < 1.0) {
return Err(LimiteronError::ConfigError(
"adaptive-threshold: decrease_ratio must be in (0, 1)".to_string(),
));
}
if !(0.0..=1.0).contains(&self.error_rate_increase_threshold)
|| !(0.0..=1.0).contains(&self.error_rate_decrease_threshold)
{
return Err(LimiteronError::ConfigError(
"adaptive-threshold: error_rate_increase_threshold and \
error_rate_decrease_threshold must be in [0, 1]"
.to_string(),
));
}
if self.error_rate_decrease_threshold > self.error_rate_increase_threshold {
return Err(LimiteronError::ConfigError(
"adaptive-threshold: error_rate_decrease_threshold must be <= increase threshold"
.to_string(),
));
}
if self.latency_decrease_ms > self.latency_increase_ms {
return Err(LimiteronError::ConfigError(
"adaptive-threshold: latency_decrease_ms must be <= latency_increase_ms"
.to_string(),
));
}
if self.window_samples == 0 || self.cooldown_samples == 0 {
return Err(LimiteronError::ConfigError(
"adaptive-threshold: window_samples and cooldown_samples must be > 0".to_string(),
));
}
if self.min_samples_for_adjust == 0 || self.min_samples_for_adjust > self.window_samples {
return Err(LimiteronError::ConfigError(
"adaptive-threshold: min_samples_for_adjust must be in 1..=window_samples"
.to_string(),
));
}
Ok(())
}
}
#[derive(Debug)]
struct ObservationWindow {
latencies: Vec<u64>,
error_ring: Vec<bool>,
cursor: usize,
filled: usize,
errors: usize,
total_pushed: usize,
cooldown_left: usize,
}
impl ObservationWindow {
fn new(capacity: usize) -> Self {
Self {
latencies: vec![0; capacity],
error_ring: vec![false; capacity],
cursor: 0,
filled: 0,
errors: 0,
total_pushed: 0,
cooldown_left: 0,
}
}
fn capacity(&self) -> usize {
self.latencies.len()
}
fn push(&mut self, ok: bool, latency_ms: u64) -> bool {
if self.filled == self.latencies.len() {
if self.error_ring[self.cursor] {
self.errors = self.errors.saturating_sub(1);
}
} else {
self.filled += 1;
}
self.latencies[self.cursor] = latency_ms;
self.error_ring[self.cursor] = !ok;
if !ok {
self.errors += 1;
}
self.cursor = (self.cursor + 1) % self.latencies.len();
self.total_pushed += 1;
self.total_pushed % self.latencies.len() == 0
}
fn sample_count(&self) -> usize {
self.filled
}
fn error_rate(&self) -> f64 {
if self.filled == 0 {
return 0.0;
}
self.errors as f64 / self.filled as f64
}
fn latency_p95_ms(&self) -> u64 {
if self.filled == 0 {
return 0;
}
let mut latencies = self.latencies[..self.filled].to_vec();
latencies.sort_unstable();
let index = ((self.filled as f64 * 0.95).ceil() as usize)
.saturating_sub(1)
.min(self.filled - 1);
latencies[index]
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Feedback {
pub ok: bool,
pub latency_ms: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Adjustment {
None,
Cooling(usize),
Decreased(u64),
Increased(u64),
}
pub struct AdaptiveThresholdLimiter {
config: AdaptiveThresholdConfig,
dynamic_limit: AtomicU64,
consumed: AtomicU64,
observations: parking_lot::Mutex<ObservationWindow>,
last_adjustment: parking_lot::Mutex<Adjustment>,
}
impl AdaptiveThresholdLimiter {
#[must_use]
pub(crate) fn new(config: AdaptiveThresholdConfig) -> Self {
debug_assert!(
config.validate().is_ok(),
"adaptive-threshold: invalid config passed to new()"
);
let base = config.base_limit;
Self {
dynamic_limit: AtomicU64::new(base),
consumed: AtomicU64::new(0),
observations: parking_lot::Mutex::new(ObservationWindow::new(config.window_samples)),
last_adjustment: parking_lot::Mutex::new(Adjustment::None),
config,
}
}
pub fn with_validation(config: AdaptiveThresholdConfig) -> Result<Self, LimiteronError> {
config.validate()?;
Ok(Self::new(config))
}
#[must_use]
pub fn current_limit(&self) -> u64 {
self.dynamic_limit.load(Ordering::Relaxed)
}
#[must_use]
pub fn last_adjustment(&self) -> Adjustment {
*self.last_adjustment.lock()
}
#[must_use]
pub fn window_stats(&self) -> (usize, f64, u64) {
let obs = self.observations.lock();
(obs.sample_count(), obs.error_rate(), obs.latency_p95_ms())
}
pub fn report(&self, feedback: Feedback) {
let (window_rolled, decision) = {
let mut obs = self.observations.lock();
obs.cooldown_left = obs.cooldown_left.saturating_sub(1);
obs.push(feedback.ok, feedback.latency_ms);
let evaluate = obs.total_pushed % self.config.min_samples_for_adjust == 0
&& obs.sample_count() >= self.config.min_samples_for_adjust
&& obs.cooldown_left == 0;
if evaluate {
(
obs.total_pushed % obs.capacity() == 0,
Some((obs.error_rate(), obs.latency_p95_ms())),
)
} else {
(obs.total_pushed % obs.capacity() == 0, None)
}
};
if window_rolled {
self.consumed.store(0, Ordering::Relaxed);
}
if let Some((error_rate, p95)) = decision {
self.apply_adjustment(error_rate, p95);
}
}
fn apply_adjustment(&self, error_rate: f64, p95: u64) {
let should_decrease = error_rate >= self.config.error_rate_increase_threshold
|| p95 >= self.config.latency_increase_ms;
let should_increase = error_rate <= self.config.error_rate_decrease_threshold
&& p95 <= self.config.latency_decrease_ms;
if should_decrease {
let current = self.current_limit();
let next = ((current as f64 * self.config.decrease_ratio) as u64)
.max(self.config.min_limit)
.min(current);
if next < current {
self.dynamic_limit.store(next, Ordering::Relaxed);
self.observations.lock().cooldown_left = self.config.cooldown_samples;
*self.last_adjustment.lock() = Adjustment::Decreased(next);
}
} else if should_increase {
let current = self.current_limit();
let next = current
.saturating_add(self.config.increase_step)
.min(self.config.max_limit);
if next > current {
self.dynamic_limit.store(next, Ordering::Relaxed);
self.observations.lock().cooldown_left = self.config.cooldown_samples;
*self.last_adjustment.lock() = Adjustment::Increased(next);
}
}
}
#[must_use]
pub fn snapshot(&self) -> RateLimitSnapshot {
let limit = self.current_limit();
let consumed = self.consumed.load(Ordering::Relaxed);
RateLimitSnapshot {
limit,
remaining: limit.saturating_sub(consumed),
reset_secs: 0,
}
}
}
#[async_trait]
impl Limiter for AdaptiveThresholdLimiter {
async fn allow(&self, cost: u64) -> Result<bool, LimiteronError> {
let limit = self.current_limit();
let mut consumed = self.consumed.load(Ordering::Relaxed);
loop {
if consumed.saturating_add(cost) > limit {
return Ok(false);
}
match self.consumed.compare_exchange_weak(
consumed,
consumed + cost,
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => return Ok(true),
Err(actual) => consumed = actual,
}
}
}
async fn peek(&self, _cost: u64) -> Result<RateLimitSnapshot, LimiteronError> {
Ok(self.snapshot())
}
async fn remaining(&self) -> Result<RateLimitSnapshot, LimiteronError> {
Ok(self.snapshot())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn config() -> AdaptiveThresholdConfig {
AdaptiveThresholdConfig {
base_limit: 100,
min_limit: 20,
max_limit: 200,
window_samples: 10,
min_samples_for_adjust: 5,
error_rate_increase_threshold: 0.2,
error_rate_decrease_threshold: 0.01,
latency_increase_ms: 500,
latency_decrease_ms: 100,
decrease_ratio: 0.5,
increase_step: 10,
cooldown_samples: 10,
}
}
#[test]
fn config_validation_rejects_invariants() {
let mut c = config();
c.validate().unwrap();
c.min_limit = 0;
assert!(c.validate().is_err());
let mut c = config();
c.base_limit = 10;
c.min_limit = 20;
assert!(c.validate().is_err());
let mut c = config();
c.decrease_ratio = 1.0;
assert!(c.validate().is_err());
let mut c = config();
c.error_rate_decrease_threshold = 0.5;
assert!(c.validate().is_err());
}
#[test]
fn config_validation_bounds_min_samples_for_adjust() {
let mut c = config();
c.min_samples_for_adjust = 0;
let err = c.validate().unwrap_err();
assert!(
matches!(err, LimiteronError::ConfigError(ref m) if m.contains("min_samples_for_adjust")),
"the error message must contain the field name: {err}"
);
let mut c = config();
c.min_samples_for_adjust = c.window_samples + 1;
assert!(c.validate().is_err());
let mut c = config();
c.min_samples_for_adjust = 1;
c.validate().unwrap();
let mut c = config();
c.min_samples_for_adjust = c.window_samples;
c.validate().unwrap();
}
#[test]
fn config_validation_clamps_error_rate_thresholds_to_unit_interval() {
let mut c = config();
c.error_rate_increase_threshold = 1.5;
let err = c.validate().unwrap_err();
assert!(
matches!(err, LimiteronError::ConfigError(ref m) if m.contains("error_rate_increase_threshold")),
"the error message must contain the field name: {err}"
);
let mut c = config();
c.error_rate_decrease_threshold = -0.1;
assert!(c.validate().is_err());
let mut c = config();
c.error_rate_increase_threshold = 1.0;
c.error_rate_decrease_threshold = 0.0;
c.validate().unwrap();
}
fn small_config() -> AdaptiveThresholdConfig {
AdaptiveThresholdConfig {
base_limit: 3,
min_limit: 1,
max_limit: 3,
window_samples: 4,
min_samples_for_adjust: 2,
error_rate_increase_threshold: 0.5,
error_rate_decrease_threshold: 0.01,
latency_increase_ms: 500,
latency_decrease_ms: 100,
decrease_ratio: 0.5,
increase_step: 10,
cooldown_samples: 2,
}
}
#[test]
fn quota_period_resets_on_window_rollover() {
let limiter = AdaptiveThresholdLimiter::with_validation(small_config()).unwrap();
for _ in 0..3 {
assert!(tokio_test_block_on(limiter.allow(1)).unwrap());
}
assert!(
!tokio_test_block_on(limiter.allow(1)).unwrap(),
"must reject once the count is full"
);
for _ in 0..2 {
limiter.report(Feedback {
ok: true,
latency_ms: 1,
});
}
assert!(
!tokio_test_block_on(limiter.allow(1)).unwrap(),
"must not resume allowing before the window rolls"
);
for _ in 0..2 {
limiter.report(Feedback {
ok: true,
latency_ms: 1,
});
}
assert!(
tokio_test_block_on(limiter.allow(1)).unwrap(),
"should resume allowing after the window rolls"
);
for _ in 0..2 {
assert!(tokio_test_block_on(limiter.allow(1)).unwrap());
}
assert!(!tokio_test_block_on(limiter.allow(1)).unwrap());
for _ in 0..4 {
limiter.report(Feedback {
ok: true,
latency_ms: 1,
});
}
assert!(
tokio_test_block_on(limiter.allow(1)).unwrap(),
"the second window roll should resume allowing again"
);
}
#[test]
fn adjustment_decisions_fire_on_min_sample_boundaries() {
let limiter = AdaptiveThresholdLimiter::with_validation(small_config()).unwrap();
limiter.report(Feedback {
ok: false,
latency_ms: 1,
});
assert_eq!(
limiter.current_limit(),
3,
"must not adjust between decision points"
);
limiter.report(Feedback {
ok: false,
latency_ms: 1,
});
assert_eq!(limiter.current_limit(), 1);
assert_eq!(limiter.last_adjustment(), Adjustment::Decreased(1));
}
#[test]
fn allow_consumes_against_dynamic_limit() {
let limiter = AdaptiveThresholdLimiter::with_validation(config()).unwrap();
assert_eq!(limiter.current_limit(), 100);
for _ in 0..100 {
assert!(tokio_test_block_on(limiter.allow(1)).unwrap());
}
assert!(!tokio_test_block_on(limiter.allow(1)).unwrap());
assert!(!tokio_test_block_on(limiter.allow(u64::MAX)).unwrap());
}
fn tokio_test_block_on<F: std::future::Future>(fut: F) -> F::Output {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap()
.block_on(fut)
}
#[test]
fn high_error_rate_triggers_decrease_to_floor() {
let limiter = AdaptiveThresholdLimiter::with_validation(config()).unwrap();
for _ in 0..5 {
limiter.report(Feedback {
ok: false,
latency_ms: 1,
});
}
assert_eq!(limiter.current_limit(), 50);
assert_eq!(limiter.last_adjustment(), Adjustment::Decreased(50));
for _ in 0..9 {
limiter.report(Feedback {
ok: false,
latency_ms: 1,
});
}
assert_eq!(
limiter.current_limit(),
50,
"must not keep lowering during cooldown"
);
limiter.report(Feedback {
ok: false,
latency_ms: 1,
});
assert_eq!(limiter.current_limit(), 25);
for _ in 0..40 {
limiter.report(Feedback {
ok: false,
latency_ms: 1,
});
}
assert_eq!(limiter.current_limit(), 20, "must not go below min_limit");
}
#[test]
fn healthy_feedback_triggers_increase_to_ceiling() {
let limiter = AdaptiveThresholdLimiter::with_validation(config()).unwrap();
for _ in 0..20 {
for _ in 0..5 {
limiter.report(Feedback {
ok: true,
latency_ms: 10,
});
}
}
assert_eq!(
limiter.current_limit(),
200,
"additive increase must clamp at max_limit"
);
assert_eq!(limiter.last_adjustment(), Adjustment::Increased(200));
}
#[test]
fn cooldown_suppresses_immediate_reincrease() {
let limiter = AdaptiveThresholdLimiter::with_validation(config()).unwrap();
limiter.report(Feedback {
ok: false,
latency_ms: 1,
});
for _ in 0..4 {
limiter.report(Feedback {
ok: true,
latency_ms: 10,
});
}
assert_eq!(limiter.current_limit(), 50);
for _ in 0..5 {
limiter.report(Feedback {
ok: true,
latency_ms: 10,
});
}
assert_eq!(
limiter.current_limit(),
50,
"must not raise during cooldown"
);
assert!(
matches!(
limiter.last_adjustment(),
Adjustment::Decreased(50) | Adjustment::Cooling(_)
),
"the cooldown adjustment action must keep/show cooldown: {:?}",
limiter.last_adjustment()
);
for _ in 0..5 {
limiter.report(Feedback {
ok: true,
latency_ms: 10,
});
}
assert_eq!(
limiter.current_limit(),
60,
"additive increase should resume after cooldown ends"
);
}
#[test]
fn latency_p95_alone_triggers_decrease() {
let limiter = AdaptiveThresholdLimiter::with_validation(config()).unwrap();
for _ in 0..5 {
limiter.report(Feedback {
ok: true,
latency_ms: 800,
});
}
assert_eq!(
limiter.current_limit(),
50,
"delayed degradation must trigger tightening independently"
);
}
#[test]
fn cold_start_below_min_samples_skips_adjustment() {
let limiter = AdaptiveThresholdLimiter::with_validation(config()).unwrap();
for _ in 0..4 {
limiter.report(Feedback {
ok: false,
latency_ms: 900,
});
}
assert_eq!(
limiter.current_limit(),
100,
"must not jitter during cold start"
);
assert_eq!(limiter.last_adjustment(), Adjustment::None);
}
#[test]
fn window_stats_and_snapshot_reflect_state() {
let limiter = AdaptiveThresholdLimiter::with_validation(config()).unwrap();
limiter.report(Feedback {
ok: true,
latency_ms: 50,
});
limiter.report(Feedback {
ok: false,
latency_ms: 150,
});
let (count, error_rate, p95) = limiter.window_stats();
assert_eq!(count, 2);
assert!((error_rate - 0.5).abs() < f64::EPSILON);
assert_eq!(p95, 150);
let snap = limiter.snapshot();
assert_eq!(snap.limit, 100);
assert_eq!(snap.remaining, 100);
}
#[tokio::test]
async fn limiter_trait_check_maps_rejection_to_error() {
use crate::limiters::traits::Limiter;
let limiter = AdaptiveThresholdLimiter::with_validation(config()).unwrap();
limiter.allow(100).await.unwrap();
let result = limiter.check("any-key").await;
assert!(result.is_err());
}
}