use crate::error::LimiteronError;
use crate::limiters::traits::Limiter;
use async_trait::async_trait;
use std::sync::Arc;
use std::sync::atomic::{AtomicI64, AtomicU64, Ordering};
use std::time::Duration;
#[derive(Debug, Clone)]
pub struct AdaptiveConcurrencyConfig {
pub initial_max: u64,
pub min_max: u64,
pub max_limit: u64,
pub window_size: u64,
pub increase_step: u64,
pub decrease_factor: f64,
pub degrade_rate: f64,
pub slow_threshold: Duration,
}
impl Default for AdaptiveConcurrencyConfig {
fn default() -> Self {
Self {
initial_max: 16,
min_max: 4,
max_limit: 256,
window_size: 32,
increase_step: 1,
decrease_factor: 0.5,
degrade_rate: 0.5,
slow_threshold: Duration::from_millis(200),
}
}
}
#[derive(Debug, Default, Clone)]
struct WindowSignals {
samples: u64,
errors: u64,
slow: u64,
successes: u64,
}
pub struct AdaptiveConcurrencyLimiter {
config: AdaptiveConcurrencyConfig,
current_max: AtomicU64,
in_flight: Arc<AtomicI64>,
window: parking_lot::Mutex<WindowSignals>,
#[cfg(feature = "circuit-breaker")]
circuit_breaker: Option<Arc<crate::circuit::CircuitBreaker>>,
}
pub struct AdaptivePermit {
limiter: Arc<AdaptiveConcurrencyLimiter>,
units: u64,
}
impl Drop for AdaptivePermit {
fn drop(&mut self) {
self.limiter.release(self.units);
}
}
impl AdaptiveConcurrencyLimiter {
pub fn new(config: AdaptiveConcurrencyConfig) -> Self {
assert!(
config.decrease_factor > 0.0 && config.decrease_factor < 1.0,
"decrease_factor must be in (0, 1)"
);
assert!(
config.min_max <= config.initial_max && config.initial_max <= config.max_limit,
"require min_max <= initial_max <= max_limit"
);
let initial = config.initial_max;
Self {
config,
current_max: AtomicU64::new(initial),
in_flight: Arc::new(AtomicI64::new(0)),
window: parking_lot::Mutex::new(WindowSignals::default()),
#[cfg(feature = "circuit-breaker")]
circuit_breaker: None,
}
}
#[cfg(feature = "circuit-breaker")]
pub fn with_circuit_breaker(mut self, cb: Arc<crate::circuit::CircuitBreaker>) -> Self {
self.circuit_breaker = Some(cb);
self
}
pub fn current_max(&self) -> u64 {
self.current_max.load(Ordering::Relaxed)
}
pub fn in_flight(&self) -> i64 {
self.in_flight.load(Ordering::Relaxed)
}
pub fn record_success(&self, latency: Duration) {
let mut window = self.window.lock();
window.samples += 1;
window.successes += 1;
if latency > self.config.slow_threshold {
window.slow += 1;
}
self.maybe_adjust(&mut window);
}
pub fn record_error(&self) {
let mut window = self.window.lock();
window.samples += 1;
window.errors += 1;
self.maybe_adjust(&mut window);
}
fn maybe_adjust(&self, window: &mut WindowSignals) {
if window.samples < self.config.window_size {
return;
}
let degrade_rate = window.errors + window.slow;
let bad_rate = degrade_rate as f64 / window.samples as f64;
let current = self.current_max.load(Ordering::Relaxed);
let next = if bad_rate >= self.config.degrade_rate {
let decreased = (current as f64 * self.config.decrease_factor) as u64;
decreased.max(self.config.min_max)
} else if window.errors == 0 {
(current + self.config.increase_step).min(self.config.max_limit)
} else {
current
};
if next != current {
self.current_max.store(next, Ordering::Relaxed);
}
*window = WindowSignals::default();
}
pub fn try_acquire(self: &Arc<Self>, units: u64) -> Option<AdaptivePermit> {
loop {
let in_flight = self.in_flight.load(Ordering::Relaxed);
let max = self.current_max.load(Ordering::Relaxed) as i64;
if in_flight + units as i64 > max {
return None;
}
match self.in_flight.compare_exchange(
in_flight,
in_flight + units as i64,
Ordering::AcqRel,
Ordering::Relaxed,
) {
Ok(_) => {
return Some(AdaptivePermit {
limiter: Arc::clone(self),
units,
});
}
Err(_) => continue,
}
}
}
pub fn release(&self, units: u64) {
let _ = self
.in_flight
.fetch_update(Ordering::AcqRel, Ordering::Relaxed, |cur| {
Some(cur.saturating_sub(units as i64))
});
}
}
#[async_trait]
impl Limiter for AdaptiveConcurrencyLimiter {
async fn allow(&self, cost: u64) -> Result<bool, LimiteronError> {
if cost == 0 {
return Err(LimiteronError::ConfigError(
"Cost cannot be zero".to_string(),
));
}
#[cfg(feature = "circuit-breaker")]
if let Some(cb) = &self.circuit_breaker {
if cb.is_open().await {
self.record_error();
return Ok(false);
}
}
loop {
let in_flight = self.in_flight.load(Ordering::Relaxed);
let max = self.current_max.load(Ordering::Relaxed) as i64;
if in_flight + cost as i64 > max {
self.record_error();
return Ok(false);
}
match self.in_flight.compare_exchange(
in_flight,
in_flight + cost as i64,
Ordering::AcqRel,
Ordering::Relaxed,
) {
Ok(_) => {
self.record_success(Duration::ZERO);
if let Ok(handle) = tokio::runtime::Handle::try_current() {
let counter = Arc::clone(&self.in_flight);
handle.spawn(async move {
tokio::time::sleep(Duration::from_millis(50)).await;
counter.fetch_sub(cost as i64, Ordering::AcqRel);
});
} else {
self.release(cost);
}
return Ok(true);
}
Err(_) => continue,
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn config() -> AdaptiveConcurrencyConfig {
AdaptiveConcurrencyConfig {
initial_max: 8,
min_max: 2,
max_limit: 16,
window_size: 4,
increase_step: 1,
decrease_factor: 0.5,
degrade_rate: 0.5,
slow_threshold: Duration::from_millis(100),
}
}
#[test]
fn test_t610_initial_window() {
let limiter = AdaptiveConcurrencyLimiter::new(config());
assert_eq!(limiter.current_max(), 8);
assert_eq!(limiter.in_flight(), 0);
}
#[test]
fn test_t610_additive_increase_on_clean_window() {
let limiter = AdaptiveConcurrencyLimiter::new(config());
for _ in 0..4 {
limiter.record_success(Duration::from_millis(1));
}
assert_eq!(limiter.current_max(), 9, "无错窗口应加性增 1(8 → 9)");
}
#[test]
fn test_t610_multiplicative_decrease_on_errors() {
let limiter = AdaptiveConcurrencyLimiter::new(config());
limiter.record_success(Duration::from_millis(1));
limiter.record_success(Duration::from_millis(1));
limiter.record_error();
limiter.record_error();
assert_eq!(
limiter.current_max(),
4,
"错误率达阈值的窗口应乘性减半(8 → 4)"
);
}
#[test]
fn test_t610_slow_calls_degrade_window() {
let limiter = AdaptiveConcurrencyLimiter::new(config());
limiter.record_success(Duration::from_millis(1));
limiter.record_success(Duration::from_millis(1));
limiter.record_success(Duration::from_millis(500));
limiter.record_success(Duration::from_millis(500));
assert!(limiter.current_max() < 8, "慢调用率过半应收紧窗口(8 → 4)");
}
#[test]
fn test_t610_min_max_floor_respected() {
let limiter = AdaptiveConcurrencyLimiter::new(config());
for _ in 0..12 {
limiter.record_error();
}
assert_eq!(limiter.current_max(), 2, "乘性减不得低于 min_max 下限");
}
#[test]
fn test_t610_max_limit_cap_respected() {
let limiter = AdaptiveConcurrencyLimiter::new(config());
for _ in 0..400 {
limiter.record_success(Duration::from_millis(1));
}
assert_eq!(limiter.current_max(), 16, "加性增不得超过 max_limit 上限");
}
#[tokio::test]
async fn test_t610_reject_when_window_exhausted() {
let limiter = Arc::new(AdaptiveConcurrencyLimiter::new(config()));
let permit = limiter.try_acquire(8);
assert!(permit.is_some(), "空载应可占满窗口");
let second = limiter.try_acquire(1);
assert!(second.is_none(), "窗口占满后应拒绝");
drop(permit);
assert_eq!(limiter.in_flight(), 0, "许可 Drop 后应自动释放");
assert!(limiter.try_acquire(1).is_some(), "释放后应可再占用");
}
#[test]
fn test_t610_partial_units_reject_keeps_state() {
let limiter = Arc::new(AdaptiveConcurrencyLimiter::new(config()));
let permit = limiter.try_acquire(6);
assert!(permit.is_some());
assert!(limiter.try_acquire(3).is_none());
assert_eq!(limiter.in_flight(), 6);
assert!(limiter.try_acquire(2).is_some());
}
#[cfg(feature = "circuit-breaker")]
#[tokio::test]
async fn test_t610_circuit_open_rejects() {
use crate::CircuitState;
use crate::circuit::{CircuitBreaker, CircuitBreakerConfig};
let cb_config = CircuitBreakerConfig {
failure_threshold: 1,
..CircuitBreakerConfig::default()
};
let cb = Arc::new(CircuitBreaker::with_dependencies(cb_config));
let limiter =
Arc::new(AdaptiveConcurrencyLimiter::new(config()).with_circuit_breaker(cb.clone()));
let _ = cb
.execute(|| async { Err::<(), _>(LimiteronError::Other("boom".to_string())) })
.await;
assert_eq!(cb.get_state().await, CircuitState::Open, "熔断应已打开");
let allowed = <AdaptiveConcurrencyLimiter as Limiter>::allow(&*limiter, 1)
.await
.unwrap();
assert!(!allowed, "熔断打开期间自适应限流器应拒绝");
}
}