use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
#[derive(Debug, Clone)]
pub struct PoolScalerConfig {
pub scale_up_threshold: f64,
pub scale_down_threshold: f64,
pub check_interval: Duration,
pub max_connections: usize,
pub min_connections: usize,
}
impl Default for PoolScalerConfig {
fn default() -> Self {
Self {
scale_up_threshold: 0.3,
scale_down_threshold: 0.7,
check_interval: Duration::from_secs(30),
max_connections: 100,
min_connections: 5,
}
}
}
#[derive(Debug, Clone)]
pub struct PoolMetrics {
pub current_connections: usize,
pub idle_connections: usize,
pub timeout_count: u64,
pub total_acquire: u64,
}
impl PoolMetrics {
pub fn timeout_rate(&self) -> f64 {
if self.total_acquire == 0 {
0.0
} else {
self.timeout_count as f64 / self.total_acquire as f64
}
}
pub fn idle_rate(&self) -> f64 {
if self.current_connections == 0 {
0.0
} else {
self.idle_connections as f64 / self.current_connections as f64
}
}
}
pub struct PoolScaler {
config: PoolScalerConfig,
target_connections: Arc<AtomicUsize>,
running: Arc<AtomicBool>,
scale_up_count: Arc<AtomicUsize>,
scale_down_count: Arc<AtomicUsize>,
}
impl PoolScaler {
pub fn new(config: PoolScalerConfig) -> Self {
let initial = config.min_connections;
Self {
config,
target_connections: Arc::new(AtomicUsize::new(initial)),
running: Arc::new(AtomicBool::new(false)),
scale_up_count: Arc::new(AtomicUsize::new(0)),
scale_down_count: Arc::new(AtomicUsize::new(0)),
}
}
pub fn target_connections(&self) -> usize {
self.target_connections.load(Ordering::Acquire)
}
pub fn is_running(&self) -> bool {
self.running.load(Ordering::Relaxed)
}
pub fn scale_up_count(&self) -> usize {
self.scale_up_count.load(Ordering::Relaxed)
}
pub fn scale_down_count(&self) -> usize {
self.scale_down_count.load(Ordering::Relaxed)
}
pub fn scale_up(&self, metrics: &PoolMetrics) {
if metrics.timeout_rate() > self.config.scale_up_threshold {
let current = self.target_connections.load(Ordering::Relaxed);
let new_target = (current + (current / 4)).min(self.config.max_connections);
self.target_connections.store(new_target, Ordering::Release);
self.scale_up_count.fetch_add(1, Ordering::Relaxed);
}
}
pub fn scale_down(&self, metrics: &PoolMetrics) {
if metrics.idle_rate() > self.config.scale_down_threshold {
let current = self.target_connections.load(Ordering::Relaxed);
let new_target = (current - (current / 4)).max(self.config.min_connections);
self.target_connections.store(new_target, Ordering::Release);
self.scale_down_count.fetch_add(1, Ordering::Relaxed);
}
}
pub fn adjust(&self, metrics: &PoolMetrics) {
if metrics.timeout_rate() > self.config.scale_up_threshold {
self.scale_up(metrics);
} else if metrics.idle_rate() > self.config.scale_down_threshold {
self.scale_down(metrics);
}
}
}
impl std::fmt::Debug for PoolScaler {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let target = self.target_connections.load(Ordering::Acquire);
write!(
f,
"PoolScaler {{ target: {target}, running: {} }}",
self.running.load(Ordering::Relaxed)
)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_pool_metrics_timeout_rate() {
let metrics = PoolMetrics {
current_connections: 10,
idle_connections: 3,
timeout_count: 5,
total_acquire: 100,
};
assert_eq!(metrics.timeout_rate(), 0.05);
}
#[test]
fn test_pool_metrics_timeout_rate_zero() {
let metrics = PoolMetrics {
current_connections: 10,
idle_connections: 3,
timeout_count: 0,
total_acquire: 0,
};
assert_eq!(metrics.timeout_rate(), 0.0);
}
#[test]
fn test_pool_metrics_idle_rate() {
let metrics = PoolMetrics {
current_connections: 10,
idle_connections: 7,
timeout_count: 0,
total_acquire: 100,
};
assert_eq!(metrics.idle_rate(), 0.7);
}
#[test]
fn test_config_default() {
let config = PoolScalerConfig::default();
assert_eq!(config.scale_up_threshold, 0.3);
assert_eq!(config.scale_down_threshold, 0.7);
assert_eq!(config.check_interval, Duration::from_secs(30));
}
#[test]
fn test_scaler_new() {
let config = PoolScalerConfig::default();
let scaler = PoolScaler::new(config);
assert_eq!(scaler.target_connections(), 5);
assert!(!scaler.is_running());
}
#[test]
fn test_scale_up() {
let config = PoolScalerConfig::default();
let scaler = PoolScaler::new(config);
scaler.target_connections.store(10, Ordering::Relaxed);
let metrics = PoolMetrics {
current_connections: 10,
idle_connections: 2,
timeout_count: 50,
total_acquire: 100,
};
scaler.scale_up(&metrics);
assert!(scaler.target_connections() > 10);
assert_eq!(scaler.scale_up_count(), 1);
}
#[test]
fn test_scale_down() {
let config = PoolScalerConfig::default();
let scaler = PoolScaler::new(config);
scaler.target_connections.store(20, Ordering::Relaxed);
let metrics = PoolMetrics {
current_connections: 20,
idle_connections: 18,
timeout_count: 0,
total_acquire: 100,
};
scaler.scale_down(&metrics);
assert!(scaler.target_connections() < 20);
assert_eq!(scaler.scale_down_count(), 1);
}
#[test]
fn test_scale_up_max_cap() {
let config = PoolScalerConfig::default();
let scaler = PoolScaler::new(config);
scaler.target_connections.store(95, Ordering::Relaxed);
let metrics = PoolMetrics {
current_connections: 95,
idle_connections: 0,
timeout_count: 50,
total_acquire: 100,
};
scaler.scale_up(&metrics);
assert_eq!(scaler.target_connections(), 100);
}
#[test]
fn test_scale_down_min_cap() {
let config = PoolScalerConfig::default();
let scaler = PoolScaler::new(config);
scaler.target_connections.store(6, Ordering::Relaxed);
let metrics = PoolMetrics {
current_connections: 6,
idle_connections: 5,
timeout_count: 0,
total_acquire: 100,
};
scaler.scale_down(&metrics);
assert_eq!(scaler.target_connections(), 5);
}
#[test]
fn test_adjust_scale_up() {
let config = PoolScalerConfig::default();
let scaler = PoolScaler::new(config);
scaler.target_connections.store(10, Ordering::Relaxed);
let metrics = PoolMetrics {
current_connections: 10,
idle_connections: 0,
timeout_count: 50,
total_acquire: 100,
};
scaler.adjust(&metrics);
assert_eq!(scaler.scale_up_count(), 1);
assert_eq!(scaler.scale_down_count(), 0);
}
#[test]
fn test_adjust_scale_down() {
let config = PoolScalerConfig::default();
let scaler = PoolScaler::new(config);
scaler.target_connections.store(20, Ordering::Relaxed);
let metrics = PoolMetrics {
current_connections: 20,
idle_connections: 18,
timeout_count: 0,
total_acquire: 100,
};
scaler.adjust(&metrics);
assert_eq!(scaler.scale_up_count(), 0);
assert_eq!(scaler.scale_down_count(), 1);
}
}