use std::collections::VecDeque;
use std::time::Instant;
use a3s_box_core::config::ScalingPolicy;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ScaleDecision {
ScaleUp(usize),
ScaleDown(usize),
Hold,
}
struct PressureWindow {
events: VecDeque<(Instant, bool)>,
window: std::time::Duration,
}
impl PressureWindow {
fn new(window_secs: u64) -> Self {
Self {
events: VecDeque::new(),
window: std::time::Duration::from_secs(window_secs),
}
}
fn record(&mut self, hit: bool) {
self.events.push_back((Instant::now(), hit));
self.prune();
}
#[cfg(test)]
fn record_at(&mut self, at: Instant, hit: bool) {
self.events.push_back((at, hit));
self.prune();
}
fn prune(&mut self) {
let cutoff = Instant::now()
.checked_sub(self.window)
.unwrap_or_else(Instant::now);
while let Some(&(ts, _)) = self.events.front() {
if ts < cutoff {
self.events.pop_front();
} else {
break;
}
}
}
fn miss_rate(&mut self) -> Option<f64> {
self.prune();
if self.events.is_empty() {
return None;
}
let total = self.events.len() as f64;
let misses = self.events.iter().filter(|(_, hit)| !hit).count() as f64;
Some(misses / total)
}
fn len(&self) -> usize {
self.events.len()
}
}
pub struct PoolScaler {
policy: ScalingPolicy,
window: PressureWindow,
last_scale_at: Option<Instant>,
current_min_idle: usize,
max_min_idle: usize,
gateway_pressure: f64,
gateway_pressure_at: Option<Instant>,
}
impl PoolScaler {
pub fn new(policy: ScalingPolicy, initial_min_idle: usize, max_size: usize) -> Self {
let max_min_idle = if policy.max_min_idle > 0 {
policy.max_min_idle.min(max_size)
} else {
max_size
};
Self {
window: PressureWindow::new(policy.window_secs),
policy,
last_scale_at: None,
current_min_idle: initial_min_idle,
max_min_idle,
gateway_pressure: 0.0,
gateway_pressure_at: None,
}
}
pub fn record_acquire(&mut self, hit: bool) {
self.window.record(hit);
}
pub fn current_min_idle(&self) -> usize {
self.current_min_idle
}
pub fn update_gateway_pressure(&mut self, pressure: f64) {
self.gateway_pressure = pressure.clamp(0.0, 1.0);
self.gateway_pressure_at = Some(Instant::now());
}
pub fn gateway_pressure(&self) -> f64 {
self.gateway_pressure
}
fn is_gateway_pressure_fresh(&self) -> bool {
match self.gateway_pressure_at {
Some(at) => at.elapsed().as_secs() < self.policy.window_secs * 2,
None => false,
}
}
fn effective_miss_rate(&mut self) -> Option<f64> {
let local = self.window.miss_rate()?;
if self.is_gateway_pressure_fresh() && self.gateway_pressure > 0.0 {
Some(local * 0.6 + self.gateway_pressure * 0.4)
} else {
Some(local)
}
}
pub fn evaluate(&mut self) -> ScaleDecision {
if let Some(last) = self.last_scale_at {
if last.elapsed().as_secs() < self.policy.cooldown_secs {
return ScaleDecision::Hold;
}
}
if self.window.len() < 3 {
return ScaleDecision::Hold;
}
let miss_rate = match self.effective_miss_rate() {
Some(rate) => rate,
None => return ScaleDecision::Hold,
};
if miss_rate > self.policy.scale_up_threshold {
let new_min = (self.current_min_idle + 1).min(self.max_min_idle);
if new_min > self.current_min_idle {
self.current_min_idle = new_min;
self.last_scale_at = Some(Instant::now());
ScaleDecision::ScaleUp(1)
} else {
ScaleDecision::Hold }
} else if miss_rate < self.policy.scale_down_threshold {
let new_min = self.current_min_idle.saturating_sub(1).max(1);
if new_min < self.current_min_idle {
self.current_min_idle = new_min;
self.last_scale_at = Some(Instant::now());
ScaleDecision::ScaleDown(1)
} else {
ScaleDecision::Hold }
} else {
ScaleDecision::Hold
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn default_policy() -> ScalingPolicy {
ScalingPolicy {
enabled: true,
scale_up_threshold: 0.3,
scale_down_threshold: 0.05,
max_min_idle: 0, cooldown_secs: 60,
window_secs: 120,
}
}
#[test]
fn test_window_empty_miss_rate_is_none() {
let mut w = PressureWindow::new(120);
assert_eq!(w.miss_rate(), None);
assert_eq!(w.len(), 0);
}
#[test]
fn test_window_all_hits() {
let mut w = PressureWindow::new(120);
for _ in 0..10 {
w.record(true);
}
assert_eq!(w.miss_rate(), Some(0.0));
}
#[test]
fn test_window_all_misses() {
let mut w = PressureWindow::new(120);
for _ in 0..10 {
w.record(false);
}
assert_eq!(w.miss_rate(), Some(1.0));
}
#[test]
fn test_window_mixed_events() {
let mut w = PressureWindow::new(120);
for _ in 0..3 {
w.record(false);
}
for _ in 0..7 {
w.record(true);
}
let rate = w.miss_rate().unwrap();
assert!((rate - 0.3).abs() < 0.001);
}
#[test]
fn test_window_prunes_old_events() {
let mut w = PressureWindow::new(1);
let old = Instant::now() - std::time::Duration::from_secs(5);
w.record_at(old, false);
w.record_at(old, false);
w.record(true);
assert_eq!(w.len(), 1);
assert_eq!(w.miss_rate(), Some(0.0));
}
#[test]
fn test_window_record_at() {
let mut w = PressureWindow::new(300);
let now = Instant::now();
w.record_at(now, true);
w.record_at(now, false);
assert_eq!(w.len(), 2);
assert_eq!(w.miss_rate(), Some(0.5));
}
#[test]
fn test_scale_decision_equality() {
assert_eq!(ScaleDecision::Hold, ScaleDecision::Hold);
assert_eq!(ScaleDecision::ScaleUp(1), ScaleDecision::ScaleUp(1));
assert_eq!(ScaleDecision::ScaleDown(1), ScaleDecision::ScaleDown(1));
assert_ne!(ScaleDecision::ScaleUp(1), ScaleDecision::ScaleDown(1));
}
#[test]
fn test_scale_decision_debug() {
let d = ScaleDecision::ScaleUp(2);
assert!(format!("{:?}", d).contains("ScaleUp"));
}
#[test]
fn test_scaler_initial_min_idle() {
let scaler = PoolScaler::new(default_policy(), 2, 10);
assert_eq!(scaler.current_min_idle(), 2);
}
#[test]
fn test_scaler_max_min_idle_defaults_to_max_size() {
let scaler = PoolScaler::new(default_policy(), 2, 10);
assert_eq!(scaler.max_min_idle, 10);
}
#[test]
fn test_scaler_max_min_idle_capped_by_max_size() {
let mut policy = default_policy();
policy.max_min_idle = 20; let scaler = PoolScaler::new(policy, 2, 10);
assert_eq!(scaler.max_min_idle, 10); }
#[test]
fn test_scaler_max_min_idle_explicit() {
let mut policy = default_policy();
policy.max_min_idle = 5;
let scaler = PoolScaler::new(policy, 2, 10);
assert_eq!(scaler.max_min_idle, 5);
}
#[test]
fn test_scaler_hold_with_few_events() {
let mut scaler = PoolScaler::new(default_policy(), 2, 10);
scaler.record_acquire(false);
scaler.record_acquire(false);
assert_eq!(scaler.evaluate(), ScaleDecision::Hold);
}
#[test]
fn test_scaler_scale_up_on_high_miss_rate() {
let mut policy = default_policy();
policy.cooldown_secs = 0; let mut scaler = PoolScaler::new(policy, 2, 10);
for _ in 0..4 {
scaler.record_acquire(false);
}
scaler.record_acquire(true);
assert_eq!(scaler.evaluate(), ScaleDecision::ScaleUp(1));
assert_eq!(scaler.current_min_idle(), 3);
}
#[test]
fn test_scaler_scale_down_on_low_miss_rate() {
let mut policy = default_policy();
policy.cooldown_secs = 0;
let mut scaler = PoolScaler::new(policy, 3, 10);
for _ in 0..20 {
scaler.record_acquire(true);
}
assert_eq!(scaler.evaluate(), ScaleDecision::ScaleDown(1));
assert_eq!(scaler.current_min_idle(), 2);
}
#[test]
fn test_scaler_hold_in_normal_range() {
let mut policy = default_policy();
policy.cooldown_secs = 0;
let mut scaler = PoolScaler::new(policy, 3, 10);
for _ in 0..10 {
scaler.record_acquire(true);
}
scaler.record_acquire(false);
assert_eq!(scaler.evaluate(), ScaleDecision::Hold);
assert_eq!(scaler.current_min_idle(), 3);
}
#[test]
fn test_scaler_wont_exceed_max_min_idle() {
let mut policy = default_policy();
policy.cooldown_secs = 0;
policy.max_min_idle = 3;
let mut scaler = PoolScaler::new(policy, 3, 10);
for _ in 0..5 {
scaler.record_acquire(false);
}
assert_eq!(scaler.evaluate(), ScaleDecision::Hold);
assert_eq!(scaler.current_min_idle(), 3);
}
#[test]
fn test_scaler_wont_go_below_one() {
let mut policy = default_policy();
policy.cooldown_secs = 0;
let mut scaler = PoolScaler::new(policy, 1, 10);
for _ in 0..10 {
scaler.record_acquire(true);
}
assert_eq!(scaler.evaluate(), ScaleDecision::Hold);
assert_eq!(scaler.current_min_idle(), 1);
}
#[test]
fn test_scaler_cooldown_prevents_rapid_scaling() {
let policy = default_policy(); let mut scaler = PoolScaler::new(policy, 2, 10);
for _ in 0..5 {
scaler.record_acquire(false);
}
assert_eq!(scaler.evaluate(), ScaleDecision::ScaleUp(1));
for _ in 0..5 {
scaler.record_acquire(false);
}
assert_eq!(scaler.evaluate(), ScaleDecision::Hold);
}
#[test]
fn test_scaler_successive_scale_ups() {
let mut policy = default_policy();
policy.cooldown_secs = 0;
let mut scaler = PoolScaler::new(policy, 1, 5);
for _ in 0..3 {
for _ in 0..5 {
scaler.record_acquire(false);
}
let decision = scaler.evaluate();
assert_eq!(decision, ScaleDecision::ScaleUp(1));
}
assert_eq!(scaler.current_min_idle(), 4);
}
#[test]
fn test_scaler_no_events_holds() {
let mut policy = default_policy();
policy.cooldown_secs = 0;
let mut scaler = PoolScaler::new(policy, 2, 10);
assert_eq!(scaler.evaluate(), ScaleDecision::Hold);
}
#[test]
fn test_scaling_policy_default() {
let policy = ScalingPolicy::default();
assert!(!policy.enabled);
assert!((policy.scale_up_threshold - 0.3).abs() < 0.001);
assert!((policy.scale_down_threshold - 0.05).abs() < 0.001);
assert_eq!(policy.max_min_idle, 0);
assert_eq!(policy.cooldown_secs, 60);
assert_eq!(policy.window_secs, 120);
}
#[test]
fn test_scaling_policy_serde_roundtrip() {
let policy = ScalingPolicy {
enabled: true,
scale_up_threshold: 0.4,
scale_down_threshold: 0.1,
max_min_idle: 8,
cooldown_secs: 30,
window_secs: 60,
};
let json = serde_json::to_string(&policy).unwrap();
let parsed: ScalingPolicy = serde_json::from_str(&json).unwrap();
assert!(parsed.enabled);
assert!((parsed.scale_up_threshold - 0.4).abs() < 0.001);
assert!((parsed.scale_down_threshold - 0.1).abs() < 0.001);
assert_eq!(parsed.max_min_idle, 8);
assert_eq!(parsed.cooldown_secs, 30);
assert_eq!(parsed.window_secs, 60);
}
#[test]
fn test_scaling_policy_deserialize_with_defaults() {
let json = r#"{"enabled": true}"#;
let policy: ScalingPolicy = serde_json::from_str(json).unwrap();
assert!(policy.enabled);
assert!((policy.scale_up_threshold - 0.3).abs() < 0.001);
assert_eq!(policy.cooldown_secs, 60);
assert_eq!(policy.window_secs, 120);
}
#[test]
fn test_gateway_pressure_default_zero() {
let scaler = PoolScaler::new(default_policy(), 2, 10);
assert_eq!(scaler.gateway_pressure(), 0.0);
}
#[test]
fn test_update_gateway_pressure() {
let mut scaler = PoolScaler::new(default_policy(), 2, 10);
scaler.update_gateway_pressure(0.75);
assert!((scaler.gateway_pressure() - 0.75).abs() < 0.001);
}
#[test]
fn test_gateway_pressure_clamped() {
let mut scaler = PoolScaler::new(default_policy(), 2, 10);
scaler.update_gateway_pressure(1.5);
assert!((scaler.gateway_pressure() - 1.0).abs() < 0.001);
scaler.update_gateway_pressure(-0.5);
assert!((scaler.gateway_pressure() - 0.0).abs() < 0.001);
}
#[test]
fn test_gateway_pressure_boosts_scale_up() {
let mut policy = default_policy();
policy.cooldown_secs = 0;
policy.scale_up_threshold = 0.3;
let mut scaler = PoolScaler::new(policy, 2, 10);
scaler.record_acquire(false);
for _ in 0..9 {
scaler.record_acquire(true);
}
assert_eq!(scaler.evaluate(), ScaleDecision::Hold);
scaler.update_gateway_pressure(0.8);
assert_eq!(scaler.evaluate(), ScaleDecision::ScaleUp(1));
}
#[test]
fn test_gateway_pressure_zero_no_effect() {
let mut policy = default_policy();
policy.cooldown_secs = 0;
let mut scaler = PoolScaler::new(policy, 3, 10);
for _ in 0..10 {
scaler.record_acquire(true);
}
scaler.update_gateway_pressure(0.0);
assert_eq!(scaler.evaluate(), ScaleDecision::ScaleDown(1));
}
#[test]
fn test_gateway_pressure_stale_ignored() {
let mut policy = default_policy();
policy.cooldown_secs = 0;
policy.window_secs = 1; let mut scaler = PoolScaler::new(policy, 2, 10);
scaler.gateway_pressure = 0.9;
scaler.gateway_pressure_at = Some(Instant::now() - std::time::Duration::from_secs(10));
scaler.record_acquire(false);
for _ in 0..9 {
scaler.record_acquire(true);
}
assert_eq!(scaler.evaluate(), ScaleDecision::Hold);
}
#[test]
fn test_effective_miss_rate_blending() {
let mut policy = default_policy();
policy.cooldown_secs = 0;
let mut scaler = PoolScaler::new(policy, 2, 10);
for _ in 0..5 {
scaler.record_acquire(false);
}
for _ in 0..5 {
scaler.record_acquire(true);
}
let rate = scaler.effective_miss_rate().unwrap();
assert!((rate - 0.5).abs() < 0.001);
scaler.update_gateway_pressure(0.2);
let rate = scaler.effective_miss_rate().unwrap();
assert!((rate - 0.38).abs() < 0.001);
}
#[test]
fn test_gateway_pressure_prevents_scale_down() {
let mut policy = default_policy();
policy.cooldown_secs = 0;
policy.scale_down_threshold = 0.05;
let mut scaler = PoolScaler::new(policy, 3, 10);
for _ in 0..10 {
scaler.record_acquire(true);
}
scaler.update_gateway_pressure(0.3);
assert_eq!(scaler.evaluate(), ScaleDecision::Hold);
}
}