use crate::error::{AgentError, Result};
use crate::metrics_providers::{LockedServiceManagerContainerProvider, RuntimeStatsProvider};
use crate::runtime::{ContainerId, ContainerResourceUpdate, Runtime};
use crate::service::ServiceManager;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use tracing::{debug, error, info, warn};
use zlayer_scheduler::metrics::{
CgroupsMetricsSource, ContainerStatsProvider, MetricsCollector, MetricsContainerId,
MetricsSource, RawContainerStats,
};
use zlayer_scheduler::Autoscaler;
use zlayer_spec::{ScaleSpec, ServiceSpec, VerticalMode, VerticalScaleSpec};
pub const DEFAULT_AUTOSCALE_INTERVAL: Duration = Duration::from_secs(10);
const IDLE_CPU_RATE_USEC_PER_SEC: f64 = 5_000.0;
const VERTICAL_DEADBAND: f64 = 0.10;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct VpaRecommendation {
pub cpu_millis: u32,
pub memory_mib: u32,
}
#[derive(Debug, Default)]
pub struct VpaEngine {
history: HashMap<String, ContainerUsageHistory>,
}
#[derive(Debug, Default)]
struct ContainerUsageHistory {
cpu_millis: std::collections::VecDeque<f64>,
memory_mib: std::collections::VecDeque<f64>,
last_cpu: Option<(u64, Instant)>,
}
impl ContainerUsageHistory {
const CAPACITY: usize = 32;
fn push_cpu(&mut self, millis: f64) {
if self.cpu_millis.len() == Self::CAPACITY {
self.cpu_millis.pop_front();
}
self.cpu_millis.push_back(millis);
}
fn push_memory(&mut self, mib: f64) {
if self.memory_mib.len() == Self::CAPACITY {
self.memory_mib.pop_front();
}
self.memory_mib.push_back(mib);
}
fn percentile(samples: &std::collections::VecDeque<f64>, pct: u8) -> Option<f64> {
if samples.is_empty() {
return None;
}
let mut sorted: Vec<f64> = samples.iter().copied().collect();
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
let pct = f64::from(pct.min(100)) / 100.0;
#[allow(
clippy::cast_possible_truncation,
clippy::cast_sign_loss,
clippy::cast_precision_loss
)]
let idx = ((pct * sorted.len() as f64).ceil() as usize)
.saturating_sub(1)
.min(sorted.len() - 1);
Some(sorted[idx])
}
}
impl VpaEngine {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn observe(&mut self, container: &str, stats: &RawContainerStats) -> f64 {
let now = Instant::now();
let hist = self.history.entry(container.to_string()).or_default();
let cpu_millis = if let Some((prev_usec, prev_at)) = hist.last_cpu {
let elapsed = now.duration_since(prev_at).as_secs_f64();
let delta_usec = stats.cpu_usage_usec.saturating_sub(prev_usec);
if elapsed > 0.0 {
#[allow(clippy::cast_precision_loss)]
let rate = delta_usec as f64 / elapsed / 1000.0;
hist.push_cpu(rate);
rate
} else {
0.0
}
} else {
0.0
};
hist.last_cpu = Some((stats.cpu_usage_usec, now));
#[allow(clippy::cast_precision_loss)]
let mem_mib = stats.memory_bytes as f64 / (1024.0 * 1024.0);
hist.push_memory(mem_mib);
cpu_millis
}
#[must_use]
pub fn recommend(
&self,
container: &str,
spec: &VerticalScaleSpec,
) -> Option<VpaRecommendation> {
let hist = self.history.get(container)?;
let cpu_pct = ContainerUsageHistory::percentile(&hist.cpu_millis, spec.percentile)?;
let mem_pct = ContainerUsageHistory::percentile(&hist.memory_mib, spec.percentile)?;
#[allow(
clippy::cast_possible_truncation,
clippy::cast_sign_loss,
clippy::cast_precision_loss
)]
let cpu_millis = {
let mut v = cpu_pct.ceil().max(0.0) as u32;
if let Some(min) = spec.min_cpu_millis {
v = v.max(min);
}
if let Some(max) = spec.max_cpu_millis {
v = v.min(max);
}
v.max(1)
};
#[allow(
clippy::cast_possible_truncation,
clippy::cast_sign_loss,
clippy::cast_precision_loss
)]
let memory_mib = {
let mut v = mem_pct.ceil().max(0.0) as u32;
if let Some(min) = spec.min_memory_mib {
v = v.max(min);
}
if let Some(max) = spec.max_memory_mib {
v = v.min(max);
}
v.max(1)
};
Some(VpaRecommendation {
cpu_millis,
memory_mib,
})
}
pub fn forget(&mut self, container: &str) {
self.history.remove(container);
}
}
fn outside_deadband(prev: Option<VpaRecommendation>, next: VpaRecommendation) -> bool {
let Some(prev) = prev else { return true };
let exceeds = |old: u32, new: u32| {
if old == 0 {
return new != 0;
}
let delta = (f64::from(new) - f64::from(old)).abs();
delta / f64::from(old) > VERTICAL_DEADBAND
};
exceeds(prev.cpu_millis, next.cpu_millis) || exceeds(prev.memory_mib, next.memory_mib)
}
fn resource_update_for(rec: VpaRecommendation) -> ContainerResourceUpdate {
ContainerResourceUpdate {
cpu_period: Some(100_000),
cpu_quota: Some(100_000 * i64::from(rec.cpu_millis) / 1000),
memory: Some(i64::from(rec.memory_mib) * 1024 * 1024),
..Default::default()
}
}
pub struct AutoscaleController {
service_manager: Arc<RwLock<ServiceManager>>,
metrics: Arc<MetricsCollector>,
autoscaler: Arc<RwLock<Autoscaler>>,
service_specs: Arc<RwLock<HashMap<String, ScaleSpec>>>,
last_scale_times: Arc<RwLock<HashMap<String, Instant>>>,
interval: Duration,
shutdown: Arc<tokio::sync::Notify>,
runtime: Arc<dyn Runtime + Send + Sync>,
stats_provider: Arc<RuntimeStatsProvider>,
last_active: Arc<RwLock<HashMap<String, Instant>>>,
idle_window: Arc<RwLock<HashMap<String, Duration>>>,
min_replicas: Arc<RwLock<HashMap<String, u32>>>,
vertical_specs: Arc<RwLock<HashMap<String, VerticalScaleSpec>>>,
service_templates: Arc<RwLock<HashMap<String, ServiceSpec>>>,
vpa: Arc<RwLock<VpaState>>,
}
#[derive(Default)]
struct VpaState {
engine: VpaEngine,
last_applied: HashMap<String, VpaRecommendation>,
}
impl AutoscaleController {
#[must_use]
pub fn new(
service_manager: Arc<RwLock<ServiceManager>>,
runtime: Arc<dyn Runtime + Send + Sync>,
interval: Duration,
) -> Self {
let mut metrics = MetricsCollector::new();
let runtime_for_controller = runtime.clone();
let stats_provider = Arc::new(RuntimeStatsProvider::new(runtime));
let service_provider = Arc::new(LockedServiceManagerContainerProvider::new(
service_manager.clone(),
));
let source: Arc<dyn MetricsSource> = Arc::new(CgroupsMetricsSource::new(
service_provider,
stats_provider.clone(),
));
metrics.add_source(source);
Self {
service_manager,
metrics: Arc::new(metrics),
autoscaler: Arc::new(RwLock::new(Autoscaler::new())),
service_specs: Arc::new(RwLock::new(HashMap::new())),
last_scale_times: Arc::new(RwLock::new(HashMap::new())),
interval,
shutdown: Arc::new(tokio::sync::Notify::new()),
runtime: runtime_for_controller,
stats_provider,
last_active: Arc::new(RwLock::new(HashMap::new())),
idle_window: Arc::new(RwLock::new(HashMap::new())),
min_replicas: Arc::new(RwLock::new(HashMap::new())),
vertical_specs: Arc::new(RwLock::new(HashMap::new())),
service_templates: Arc::new(RwLock::new(HashMap::new())),
vpa: Arc::new(RwLock::new(VpaState::default())),
}
}
#[must_use]
pub fn with_custom_metrics(
service_manager: Arc<RwLock<ServiceManager>>,
runtime: Arc<dyn Runtime + Send + Sync>,
metrics: MetricsCollector,
interval: Duration,
) -> Self {
let runtime_for_controller = runtime.clone();
let stats_provider = Arc::new(RuntimeStatsProvider::new(runtime));
Self {
service_manager,
metrics: Arc::new(metrics),
autoscaler: Arc::new(RwLock::new(Autoscaler::new())),
service_specs: Arc::new(RwLock::new(HashMap::new())),
last_scale_times: Arc::new(RwLock::new(HashMap::new())),
interval,
shutdown: Arc::new(tokio::sync::Notify::new()),
runtime: runtime_for_controller,
stats_provider,
last_active: Arc::new(RwLock::new(HashMap::new())),
idle_window: Arc::new(RwLock::new(HashMap::new())),
min_replicas: Arc::new(RwLock::new(HashMap::new())),
vertical_specs: Arc::new(RwLock::new(HashMap::new())),
service_templates: Arc::new(RwLock::new(HashMap::new())),
vpa: Arc::new(RwLock::new(VpaState::default())),
}
}
#[must_use]
pub fn with_extra_metrics_source(mut self, source: Arc<dyn MetricsSource>) -> Self {
if let Some(collector) = Arc::get_mut(&mut self.metrics) {
collector.add_source(source);
} else {
warn!(
"with_extra_metrics_source called after the metrics collector was shared; \
source ignored"
);
}
self
}
pub async fn register_service(&self, name: &str, spec: &ScaleSpec, initial_replicas: u32) {
let ScaleSpec::Adaptive {
min,
idle_window,
vertical,
..
} = spec
else {
debug!(
service = name,
"Skipping registration for non-adaptive service"
);
return;
};
{
let mut autoscaler = self.autoscaler.write().await;
autoscaler.register_service(name, spec.clone(), initial_replicas);
}
{
let mut specs = self.service_specs.write().await;
specs.insert(name.to_string(), spec.clone());
}
{
let mut mins = self.min_replicas.write().await;
mins.insert(name.to_string(), *min);
}
if let Some(window) = idle_window {
self.idle_window
.write()
.await
.insert(name.to_string(), *window);
} else {
self.idle_window.write().await.remove(name);
}
self.last_active
.write()
.await
.insert(name.to_string(), Instant::now());
if let Some(v) = vertical {
if matches!(v.mode, VerticalMode::Recommend | VerticalMode::Auto) {
self.vertical_specs
.write()
.await
.insert(name.to_string(), v.clone());
} else {
self.vertical_specs.write().await.remove(name);
}
} else {
self.vertical_specs.write().await.remove(name);
}
info!(
service = name,
initial_replicas,
idle_window_secs = idle_window.as_ref().map(Duration::as_secs),
min = *min,
vertical = vertical.is_some(),
"Registered service for autoscaling"
);
}
pub async fn set_service_template(&self, name: &str, spec: ServiceSpec) {
self.service_templates
.write()
.await
.insert(name.to_string(), spec);
}
pub fn mark_active(&self, service: &str) {
let last_active = self.last_active.clone();
let service = service.to_string();
if let Ok(mut guard) = last_active.try_write() {
guard.insert(service, Instant::now());
return;
}
tokio::spawn(async move {
last_active.write().await.insert(service, Instant::now());
});
}
pub async fn mark_active_async(&self, service: &str) {
self.last_active
.write()
.await
.insert(service.to_string(), Instant::now());
}
pub async fn unregister_service(&self, name: &str) {
{
let mut autoscaler = self.autoscaler.write().await;
autoscaler.unregister_service(name);
}
self.service_specs.write().await.remove(name);
self.last_scale_times.write().await.remove(name);
self.last_active.write().await.remove(name);
self.idle_window.write().await.remove(name);
self.min_replicas.write().await.remove(name);
self.vertical_specs.write().await.remove(name);
self.service_templates.write().await.remove(name);
self.vpa.write().await.last_applied.remove(name);
info!(service = name, "Unregistered service from autoscaling");
}
pub async fn is_registered(&self, name: &str) -> bool {
let specs = self.service_specs.read().await;
specs.contains_key(name)
}
async fn should_scale(&self, service_name: &str) -> bool {
let cooldown = {
let specs = self.service_specs.read().await;
match specs.get(service_name) {
Some(ScaleSpec::Adaptive { cooldown, .. }) => {
cooldown.unwrap_or(zlayer_scheduler::DEFAULT_COOLDOWN)
}
_ => return false, }
};
let last_scale_times = self.last_scale_times.read().await;
if let Some(last_time) = last_scale_times.get(service_name) {
if last_time.elapsed() < cooldown {
let remaining = cooldown
.checked_sub(last_time.elapsed())
.unwrap_or_default();
debug!(
service = service_name,
remaining_secs = remaining.as_secs(),
"Service in cooldown"
);
return false;
}
}
true
}
async fn record_scale_action(&self, service_name: &str) {
let mut times = self.last_scale_times.write().await;
times.insert(service_name.to_string(), Instant::now());
}
#[allow(clippy::cast_possible_truncation)]
pub async fn run_loop(&self) -> Result<()> {
let mut ticker = tokio::time::interval(self.interval);
info!(
interval_ms = self.interval.as_millis() as u64,
"Starting autoscale controller loop"
);
loop {
tokio::select! {
_ = ticker.tick() => {
Box::pin(self.evaluate_all_services()).await;
}
() = self.shutdown.notified() => {
info!("Autoscale controller shutting down");
break;
}
}
}
Ok(())
}
async fn discover_services(&self) {
let live = self.service_manager.read().await.service_specs().await;
let mut seen_adaptive: Vec<String> = Vec::new();
for (name, spec) in &live {
if !matches!(spec.scale, ScaleSpec::Adaptive { .. }) {
continue;
}
seen_adaptive.push(name.clone());
let needs_register = {
let specs = self.service_specs.read().await;
specs.get(name) != Some(&spec.scale)
};
if needs_register {
let initial = u32::try_from(
self.service_manager
.read()
.await
.service_replica_count(name)
.await
.unwrap_or(0),
)
.unwrap_or(0);
self.register_service(name, &spec.scale, initial).await;
}
self.set_service_template(name, spec.clone()).await;
}
let registered: Vec<String> = {
let specs = self.service_specs.read().await;
specs.keys().cloned().collect()
};
for name in registered {
if !seen_adaptive.contains(&name) {
self.unregister_service(&name).await;
}
}
}
async fn evaluate_all_services(&self) {
self.discover_services().await;
let service_names: Vec<String> = {
let specs = self.service_specs.read().await;
specs.keys().cloned().collect()
};
for service_name in service_names {
if let Err(e) = Box::pin(self.evaluate_vertical(&service_name)).await {
warn!(
service = %service_name,
error = %e,
"Failed vertical (VPA) evaluation"
);
}
match self.evaluate_idle(&service_name).await {
Ok(true) => continue, Ok(false) => {}
Err(e) => warn!(
service = %service_name,
error = %e,
"Failed scale-to-zero evaluation"
),
}
if let Err(e) = self.evaluate_and_scale(&service_name).await {
warn!(
service = %service_name,
error = %e,
"Failed to evaluate/scale service"
);
}
}
}
async fn evaluate_idle(&self, service_name: &str) -> Result<bool> {
let containers = self
.service_manager
.read()
.await
.get_service_containers(service_name)
.await;
let mut busiest_cpu_millis = 0.0_f64;
{
let mut vpa = self.vpa.write().await;
for id in &containers {
let metrics_id = MetricsContainerId {
service: id.service.clone(),
replica: id.replica,
};
if let Ok(stats) = self.stats_provider.get_stats(&metrics_id).await {
let rate = vpa.engine.observe(&id.to_string(), &stats);
if rate > busiest_cpu_millis {
busiest_cpu_millis = rate;
}
}
}
}
if busiest_cpu_millis * 1000.0 >= IDLE_CPU_RATE_USEC_PER_SEC {
self.last_active
.write()
.await
.insert(service_name.to_string(), Instant::now());
}
let window = {
let windows = self.idle_window.read().await;
match windows.get(service_name) {
Some(w) => *w,
None => return Ok(false),
}
};
let min = self
.min_replicas
.read()
.await
.get(service_name)
.copied()
.unwrap_or(1);
if min != 0 {
return Ok(false);
}
let current = self
.service_manager
.read()
.await
.service_replica_count(service_name)
.await
.unwrap_or(0);
if current == 0 {
return Ok(false);
}
let idle_for = {
let last_active = self.last_active.read().await;
last_active
.get(service_name)
.map_or(Duration::ZERO, Instant::elapsed)
};
if idle_for <= window {
return Ok(false);
}
if !self.should_scale(service_name).await {
return Ok(false);
}
info!(
service = service_name,
idle_secs = idle_for.as_secs(),
window_secs = window.as_secs(),
"Scaling service to zero (idle past window)"
);
self.service_manager
.read()
.await
.scale_service(service_name, 0)
.await?;
self.record_scale_action(service_name).await;
{
let mut autoscaler = self.autoscaler.write().await;
if let Err(e) = autoscaler.record_scale_action(service_name, 0) {
warn!(
service = service_name,
error = %e,
"Failed to record scale-to-zero in autoscaler"
);
}
}
{
let mut vpa = self.vpa.write().await;
for id in &containers {
vpa.engine.forget(&id.to_string());
}
}
Ok(true)
}
async fn evaluate_vertical(&self, service_name: &str) -> Result<()> {
let spec = {
let specs = self.vertical_specs.read().await;
match specs.get(service_name) {
Some(s) => s.clone(),
None => return Ok(()), }
};
let containers = self
.service_manager
.read()
.await
.get_service_containers(service_name)
.await;
if containers.is_empty() {
return Ok(());
}
let mut chosen: Option<VpaRecommendation> = None;
{
let mut vpa = self.vpa.write().await;
for id in &containers {
let metrics_id = MetricsContainerId {
service: id.service.clone(),
replica: id.replica,
};
match self.stats_provider.get_stats(&metrics_id).await {
Ok(stats) => {
vpa.engine.observe(&id.to_string(), &stats);
if let Some(rec) = vpa.engine.recommend(&id.to_string(), &spec) {
chosen = Some(match chosen {
Some(c) => VpaRecommendation {
cpu_millis: c.cpu_millis.max(rec.cpu_millis),
memory_mib: c.memory_mib.max(rec.memory_mib),
},
None => rec,
});
}
}
Err(e) => debug!(
service = service_name,
container = %id,
error = %e,
"vertical: no stats for replica; skipping"
),
}
}
}
let Some(rec) = chosen else {
return Ok(());
};
match spec.mode {
VerticalMode::Off => Ok(()),
VerticalMode::Recommend => {
info!(
service = service_name,
cpu_millis = rec.cpu_millis,
memory_mib = rec.memory_mib,
"vertical recommendation (recommend mode; not applied)"
);
Ok(())
}
VerticalMode::Auto => {
Box::pin(self.apply_vertical(service_name, rec, &containers)).await
}
}
}
async fn apply_vertical(
&self,
service_name: &str,
rec: VpaRecommendation,
containers: &[ContainerId],
) -> Result<()> {
let prev = self
.vpa
.read()
.await
.last_applied
.get(service_name)
.copied();
if !outside_deadband(prev, rec) {
debug!(
service = service_name,
cpu_millis = rec.cpu_millis,
memory_mib = rec.memory_mib,
"vertical recommendation within deadband; skipping"
);
return Ok(());
}
let update = resource_update_for(rec);
let mut needs_rolling_restart = false;
for id in containers {
match self.runtime.update_container_resources(id, &update).await {
Ok(outcome) => {
if !outcome.warnings.is_empty() {
debug!(
service = service_name,
container = %id,
warnings = ?outcome.warnings,
"vertical apply produced warnings"
);
}
}
Err(AgentError::Unsupported(reason)) => {
debug!(
service = service_name,
container = %id,
reason = %reason,
"runtime cannot live-update resources; will roll the service"
);
needs_rolling_restart = true;
break;
}
Err(AgentError::NotFound { .. }) => {
debug!(
service = service_name,
container = %id,
"vertical apply: container vanished; skipping"
);
}
Err(e) => return Err(e),
}
}
if needs_rolling_restart {
Box::pin(self.rolling_restart_with_resources(service_name, rec)).await?;
}
info!(
service = service_name,
cpu_millis = rec.cpu_millis,
memory_mib = rec.memory_mib,
rolled = needs_rolling_restart,
"applied vertical recommendation"
);
self.vpa
.write()
.await
.last_applied
.insert(service_name.to_string(), rec);
Ok(())
}
async fn rolling_restart_with_resources(
&self,
service_name: &str,
rec: VpaRecommendation,
) -> Result<()> {
let template = self
.service_templates
.read()
.await
.get(service_name)
.cloned();
let containers = self
.service_manager
.read()
.await
.get_service_containers(service_name)
.await;
let Some(mut spec) = template else {
warn!(
service = service_name,
"rolling restart without a base ServiceSpec template; bouncing the service \
(call set_service_template to enable one-at-a-time recreation)"
);
let count = u32::try_from(containers.len()).unwrap_or(u32::MAX);
{
let sm = self.service_manager.read().await;
sm.scale_service(service_name, 0).await?;
sm.scale_service(service_name, count).await?;
}
self.record_scale_action(service_name).await;
return Ok(());
};
let mut resources = spec.resources.clone();
resources.cpu = Some(f64::from(rec.cpu_millis) / 1000.0);
resources.memory = Some(format!("{}Mi", rec.memory_mib));
spec.resources = resources;
for id in &containers {
info!(
service = service_name,
container = %id,
cpu_millis = rec.cpu_millis,
memory_mib = rec.memory_mib,
"rolling restart: recreating replica with new resources"
);
if let Err(e) = self
.runtime
.stop_container(id, Duration::from_secs(10))
.await
{
debug!(service = service_name, container = %id, error = %e, "rolling restart: stop failed (continuing)");
}
if let Err(e) = self.runtime.remove_container(id).await {
debug!(service = service_name, container = %id, error = %e, "rolling restart: remove failed (continuing)");
}
self.vpa.write().await.engine.forget(&id.to_string());
if let Err(e) = self.runtime.create_container(id, &spec).await {
error!(service = service_name, container = %id, error = %e, "rolling restart: recreate failed");
return Err(e);
}
if let Err(e) = self.runtime.start_container(id).await {
error!(service = service_name, container = %id, error = %e, "rolling restart: start failed");
return Err(e);
}
}
self.record_scale_action(service_name).await;
Ok(())
}
async fn evaluate_and_scale(&self, service_name: &str) -> Result<()> {
if !self.should_scale(service_name).await {
return Ok(());
}
let aggregated = match self.metrics.collect(service_name).await {
Ok(m) => m,
Err(e) => {
debug!(
service = service_name,
error = %e,
"No metrics available for service"
);
return Ok(());
}
};
let decision = {
let mut autoscaler = self.autoscaler.write().await;
match autoscaler.evaluate(service_name, &aggregated) {
Ok(d) => d,
Err(e) => {
debug!(
service = service_name,
error = %e,
"Failed to evaluate scaling"
);
return Ok(());
}
}
};
debug!(
service = service_name,
?decision,
cpu = aggregated.avg_cpu_percent,
memory = aggregated.avg_memory_percent,
instances = aggregated.instance_count,
"Autoscale evaluation"
);
if let Some(target) = decision.target_replicas() {
info!(
service = service_name,
target_replicas = target,
decision = ?decision,
"Executing autoscale"
);
if let Err(e) = self
.service_manager
.read()
.await
.scale_service(service_name, target)
.await
{
error!(
service = service_name,
target = target,
error = %e,
"Failed to scale service"
);
return Err(e);
}
if target > 0 {
self.last_active
.write()
.await
.insert(service_name.to_string(), Instant::now());
}
self.record_scale_action(service_name).await;
{
let mut autoscaler = self.autoscaler.write().await;
if let Err(e) = autoscaler.record_scale_action(service_name, target) {
warn!(
service = service_name,
error = %e,
"Failed to record scale action in autoscaler"
);
}
}
}
Ok(())
}
pub fn shutdown(&self) {
self.shutdown.notify_one();
}
#[must_use]
pub fn interval(&self) -> Duration {
self.interval
}
pub async fn registered_service_count(&self) -> usize {
let specs = self.service_specs.read().await;
specs.len()
}
}
#[must_use]
#[allow(clippy::implicit_hasher)]
pub fn has_adaptive_scaling(services: &HashMap<String, zlayer_spec::ServiceSpec>) -> bool {
services
.values()
.any(|s| matches!(s.scale, ScaleSpec::Adaptive { .. }))
}
#[cfg(test)]
#[allow(deprecated)]
mod tests {
use super::*;
use crate::runtime::MockRuntime;
use zlayer_scheduler::metrics::{MockMetricsSource, ServiceMetrics};
use zlayer_spec::ScaleTargets;
fn mock_spec() -> zlayer_spec::ServiceSpec {
serde_yaml::from_str::<zlayer_spec::DeploymentSpec>(
r"
version: v1
deployment: test
services:
test:
rtype: service
image:
name: test:latest
endpoints:
- name: http
protocol: http
port: 8080
scale:
mode: fixed
replicas: 1
",
)
.unwrap()
.services
.remove("test")
.unwrap()
}
fn adaptive_spec(
min: u32,
max: u32,
cpu_target: Option<u8>,
memory_target: Option<u8>,
) -> ScaleSpec {
ScaleSpec::Adaptive {
min,
max,
cooldown: Some(Duration::from_secs(0)), targets: ScaleTargets {
cpu: cpu_target,
memory: memory_target,
rps: None,
custom: Vec::new(),
external: Vec::new(),
},
behavior: None,
triggers: Vec::new(),
idle_window: None,
vertical: None,
predictive: None,
}
}
fn raw_stats(cpu_usec: u64, mem_bytes: u64) -> RawContainerStats {
RawContainerStats {
cpu_usage_usec: cpu_usec,
memory_bytes: mem_bytes,
memory_limit: 512 * 1024 * 1024,
timestamp: Instant::now(),
}
}
#[test]
fn test_vpa_percentile_nearest_rank() {
let mut samples = std::collections::VecDeque::new();
for v in [10.0, 20.0, 30.0, 40.0, 50.0] {
samples.push_back(v);
}
assert_eq!(ContainerUsageHistory::percentile(&samples, 100), Some(50.0));
assert_eq!(ContainerUsageHistory::percentile(&samples, 0), Some(10.0));
assert_eq!(ContainerUsageHistory::percentile(&samples, 90), Some(50.0));
}
#[test]
fn test_vpa_recommend_clamps_to_bounds() {
let mut engine = VpaEngine::new();
let id = "svc-rep-1";
engine.observe(id, &raw_stats(1_000_000, 300 * 1024 * 1024));
std::thread::sleep(Duration::from_millis(5));
engine.observe(id, &raw_stats(1_000_000, 300 * 1024 * 1024));
let spec = VerticalScaleSpec {
mode: VerticalMode::Auto,
min_cpu_millis: Some(500),
max_cpu_millis: Some(2000),
min_memory_mib: Some(128),
max_memory_mib: Some(256),
percentile: 90,
};
let rec = engine.recommend(id, &spec).expect("recommendation");
assert_eq!(rec.cpu_millis, 500);
assert_eq!(rec.memory_mib, 256);
}
#[test]
fn test_deadband() {
let base = VpaRecommendation {
cpu_millis: 1000,
memory_mib: 512,
};
assert!(outside_deadband(None, base));
assert!(!outside_deadband(
Some(base),
VpaRecommendation {
cpu_millis: 1050,
memory_mib: 512
}
));
assert!(outside_deadband(
Some(base),
VpaRecommendation {
cpu_millis: 1200,
memory_mib: 512
}
));
}
#[test]
fn test_resource_update_for() {
let rec = VpaRecommendation {
cpu_millis: 1500,
memory_mib: 256,
};
let update = resource_update_for(rec);
assert_eq!(update.cpu_period, Some(100_000));
assert_eq!(update.cpu_quota, Some(150_000));
assert_eq!(update.memory, Some(256 * 1024 * 1024));
}
fn locked(runtime: &Arc<dyn Runtime + Send + Sync>) -> Arc<RwLock<ServiceManager>> {
Arc::new(RwLock::new(ServiceManager::new(runtime.clone())))
}
#[tokio::test]
async fn test_autoscale_controller_creation() {
let runtime: Arc<dyn Runtime + Send + Sync> = Arc::new(MockRuntime::new());
let manager = locked(&runtime);
let controller = AutoscaleController::new(manager, runtime, Duration::from_secs(10));
assert_eq!(controller.interval(), Duration::from_secs(10));
assert_eq!(controller.registered_service_count().await, 0);
}
#[tokio::test]
async fn test_register_service() {
let runtime: Arc<dyn Runtime + Send + Sync> = Arc::new(MockRuntime::new());
let manager = locked(&runtime);
let controller = AutoscaleController::new(manager, runtime, Duration::from_secs(10));
let spec = adaptive_spec(1, 10, Some(70), None);
controller.register_service("api", &spec, 2).await;
assert!(controller.is_registered("api").await);
assert_eq!(controller.registered_service_count().await, 1);
}
#[tokio::test]
async fn test_register_fixed_service_ignored() {
let runtime: Arc<dyn Runtime + Send + Sync> = Arc::new(MockRuntime::new());
let manager = locked(&runtime);
let controller = AutoscaleController::new(manager, runtime, Duration::from_secs(10));
let spec = ScaleSpec::Fixed { replicas: 3 };
controller.register_service("api", &spec, 3).await;
assert!(!controller.is_registered("api").await);
assert_eq!(controller.registered_service_count().await, 0);
}
#[tokio::test]
async fn test_unregister_service() {
let runtime: Arc<dyn Runtime + Send + Sync> = Arc::new(MockRuntime::new());
let manager = locked(&runtime);
let controller = AutoscaleController::new(manager, runtime, Duration::from_secs(10));
let spec = adaptive_spec(1, 10, Some(70), None);
controller.register_service("api", &spec, 2).await;
assert!(controller.is_registered("api").await);
controller.unregister_service("api").await;
assert!(!controller.is_registered("api").await);
assert_eq!(controller.registered_service_count().await, 0);
}
#[tokio::test]
async fn test_has_adaptive_scaling() {
let mut services = HashMap::new();
let mut fixed_spec = mock_spec();
fixed_spec.scale = ScaleSpec::Fixed { replicas: 3 };
services.insert("web".to_string(), fixed_spec);
assert!(!has_adaptive_scaling(&services));
let mut adaptive = mock_spec();
adaptive.scale = adaptive_spec(1, 10, Some(70), None);
services.insert("api".to_string(), adaptive);
assert!(has_adaptive_scaling(&services));
}
#[tokio::test]
async fn test_autoscale_controller_with_mock_metrics() {
let runtime: Arc<dyn Runtime + Send + Sync> = Arc::new(MockRuntime::new());
let manager = locked(&runtime);
let mock = Arc::new(MockMetricsSource::new());
mock.set_metrics(
"api",
vec![
ServiceMetrics {
cpu_percent: 85.0,
memory_bytes: 100 * 1024 * 1024,
memory_limit: 512 * 1024 * 1024,
rps: None,
timestamp: Some(Instant::now()),
..Default::default()
},
ServiceMetrics {
cpu_percent: 90.0,
memory_bytes: 150 * 1024 * 1024,
memory_limit: 512 * 1024 * 1024,
rps: None,
timestamp: Some(Instant::now()),
..Default::default()
},
],
)
.await;
let mut metrics = MetricsCollector::new();
metrics.add_source(mock);
let controller = AutoscaleController::with_custom_metrics(
manager.clone(),
runtime,
metrics,
Duration::from_secs(10),
);
Box::pin(
manager
.read()
.await
.upsert_service("api".to_string(), mock_spec()),
)
.await
.unwrap();
manager.read().await.scale_service("api", 2).await.unwrap();
let spec = adaptive_spec(1, 10, Some(70), None);
controller.register_service("api", &spec, 2).await;
controller.evaluate_and_scale("api").await.unwrap();
let count = manager
.read()
.await
.service_replica_count("api")
.await
.unwrap();
assert_eq!(count, 3);
}
#[tokio::test]
async fn test_autoscale_controller_cooldown() {
let runtime: Arc<dyn Runtime + Send + Sync> = Arc::new(MockRuntime::new());
let manager = locked(&runtime);
let controller = AutoscaleController::new(manager, runtime, Duration::from_secs(10));
let spec = ScaleSpec::Adaptive {
min: 1,
max: 10,
cooldown: Some(Duration::from_secs(60)), targets: ScaleTargets {
cpu: Some(70),
memory: None,
rps: None,
custom: Vec::new(),
external: Vec::new(),
},
behavior: None,
triggers: Vec::new(),
idle_window: None,
vertical: None,
predictive: None,
};
controller.register_service("api", &spec, 2).await;
assert!(controller.should_scale("api").await);
controller.record_scale_action("api").await;
assert!(!controller.should_scale("api").await);
}
#[tokio::test]
async fn test_scale_to_zero_after_idle_window() {
let runtime: Arc<dyn Runtime + Send + Sync> = Arc::new(MockRuntime::new());
let manager = locked(&runtime);
let controller =
AutoscaleController::new(manager.clone(), runtime, Duration::from_secs(10));
Box::pin(
manager
.read()
.await
.upsert_service("api".to_string(), mock_spec()),
)
.await
.unwrap();
manager.read().await.scale_service("api", 2).await.unwrap();
let spec = ScaleSpec::Adaptive {
min: 0,
max: 10,
cooldown: Some(Duration::from_secs(0)),
targets: ScaleTargets::default(),
behavior: None,
triggers: Vec::new(),
idle_window: Some(Duration::from_millis(10)),
vertical: None,
predictive: None,
};
controller.register_service("api", &spec, 2).await;
controller.last_active.write().await.insert(
"api".to_string(),
Instant::now().checked_sub(Duration::from_secs(60)).unwrap(),
);
let reaped = controller.evaluate_idle("api").await.unwrap();
assert!(reaped, "service should have been reaped to zero");
assert_eq!(
manager
.read()
.await
.service_replica_count("api")
.await
.unwrap(),
0
);
}
#[tokio::test]
async fn test_mark_active_resets_idle_clock() {
let runtime: Arc<dyn Runtime + Send + Sync> = Arc::new(MockRuntime::new());
let manager = locked(&runtime);
let controller =
AutoscaleController::new(manager.clone(), runtime, Duration::from_secs(10));
Box::pin(
manager
.read()
.await
.upsert_service("api".to_string(), mock_spec()),
)
.await
.unwrap();
manager.read().await.scale_service("api", 1).await.unwrap();
let spec = ScaleSpec::Adaptive {
min: 0,
max: 10,
cooldown: Some(Duration::from_secs(0)),
targets: ScaleTargets::default(),
behavior: None,
triggers: Vec::new(),
idle_window: Some(Duration::from_secs(300)),
vertical: None,
predictive: None,
};
controller.register_service("api", &spec, 1).await;
controller.last_active.write().await.insert(
"api".to_string(),
Instant::now()
.checked_sub(Duration::from_secs(600))
.unwrap(),
);
controller.mark_active_async("api").await;
let reaped = controller.evaluate_idle("api").await.unwrap();
assert!(!reaped, "marked-active service must not be reaped");
assert_eq!(
manager
.read()
.await
.service_replica_count("api")
.await
.unwrap(),
1
);
}
#[tokio::test]
async fn test_no_scale_to_zero_when_min_nonzero() {
let runtime: Arc<dyn Runtime + Send + Sync> = Arc::new(MockRuntime::new());
let manager = locked(&runtime);
let controller =
AutoscaleController::new(manager.clone(), runtime, Duration::from_secs(10));
Box::pin(
manager
.read()
.await
.upsert_service("api".to_string(), mock_spec()),
)
.await
.unwrap();
manager.read().await.scale_service("api", 2).await.unwrap();
let spec = ScaleSpec::Adaptive {
min: 1,
max: 10,
cooldown: Some(Duration::from_secs(0)),
targets: ScaleTargets::default(),
behavior: None,
triggers: Vec::new(),
idle_window: Some(Duration::from_millis(1)),
vertical: None,
predictive: None,
};
controller.register_service("api", &spec, 2).await;
controller.last_active.write().await.insert(
"api".to_string(),
Instant::now().checked_sub(Duration::from_secs(60)).unwrap(),
);
let reaped = controller.evaluate_idle("api").await.unwrap();
assert!(!reaped, "min>0 must never scale to zero");
assert_eq!(
manager
.read()
.await
.service_replica_count("api")
.await
.unwrap(),
2
);
}
#[tokio::test]
async fn test_autoscale_controller_shutdown() {
let runtime: Arc<dyn Runtime + Send + Sync> = Arc::new(MockRuntime::new());
let manager = locked(&runtime);
let controller = Arc::new(AutoscaleController::new(
manager,
runtime,
Duration::from_millis(100), ));
let controller_clone = controller.clone();
let handle = tokio::spawn(async move { Box::pin(controller_clone.run_loop()).await });
tokio::time::sleep(Duration::from_millis(50)).await;
controller.shutdown();
let result = handle.await.unwrap();
assert!(result.is_ok());
}
}