a3s_box_runtime/scale/

manager.rs

//! ScaleManager implementation.

use std::collections::{HashMap, VecDeque};

use a3s_box_core::scale::{
    InstanceEvent, InstanceHealth, InstanceInfo, InstanceState, ScaleDirection, ScaleObservation,
    ScaleOperationConflict, ScaleOperationRequest, ScaleOperationResponse, ScaleRequest,
    ScaleResponse, SCALE_OPERATION_SCHEMA_VERSION,
};
use chrono::Utc;
use serde::{Deserialize, Serialize};

use super::{ServiceHealth, ServiceInstances, TrackedInstance};

/// Tracks all instances managed by this Box host.
pub struct ScaleManager {
    /// Maximum total instances across all services
    max_instances: u32,
    /// Per-service instance tracking
    services: HashMap<String, ServiceInstances>,
    /// Event log (bounded ring buffer)
    events: Vec<InstanceEvent>,
    /// Maximum events to retain
    max_events: usize,
    /// Monotonic desired-state revision per service.
    revisions: HashMap<String, u64>,
    /// Bounded replay receipts keyed by operation identity.
    operation_receipts: HashMap<String, ScaleOperationReceipt>,
    operation_order: VecDeque<String>,
    max_operation_receipts: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub(super) struct ScaleOperationReceipt {
    request: ScaleOperationRequest,
    response: ScaleOperationResponse,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub(super) struct ScaleAuthorityState {
    pub(super) schema_version: u32,
    pub(super) targets: HashMap<String, u32>,
    pub(super) revisions: HashMap<String, u64>,
    pub(super) receipts: Vec<ScaleOperationReceipt>,
}

impl ScaleManager {
    /// Create a new scale manager with the given capacity.
    pub fn new(max_instances: u32) -> Self {
        Self {
            max_instances,
            services: HashMap::new(),
            events: Vec::new(),
            max_events: 1000,
            revisions: HashMap::new(),
            operation_receipts: HashMap::new(),
            operation_order: VecDeque::new(),
            max_operation_receipts: 10_000,
        }
    }

    /// Return the authoritative desired replica count and CAS revision.
    pub fn scale_observation(&self, service: &str) -> ScaleObservation {
        ScaleObservation {
            replicas: self
                .services
                .get(service)
                .map_or(0, |state| state.target_replicas),
            revision: Some(
                self.revisions
                    .get(service)
                    .copied()
                    .unwrap_or(0)
                    .to_string(),
            ),
            ready_replicas: 0,
            endpoints: Vec::new(),
        }
    }

    /// Apply one versioned scale mutation exactly once.
    pub fn apply_operation(
        &mut self,
        request: &ScaleOperationRequest,
    ) -> Result<ScaleOperationResponse, ScaleOperationConflict> {
        let observation = self.scale_observation(&request.service);

        if let Some(receipt) = self.operation_receipts.get(&request.operation_id) {
            return if receipt.request == *request {
                Ok(receipt.response.clone())
            } else {
                Err(scale_conflict(
                    "operation_conflict",
                    "operation_id was already used with a different scale request",
                    observation,
                ))
            };
        }

        if request.schema_version != SCALE_OPERATION_SCHEMA_VERSION {
            return Err(scale_conflict(
                "unsupported_schema",
                format!(
                    "scale schema {} is unsupported; expected {}",
                    request.schema_version, SCALE_OPERATION_SCHEMA_VERSION
                ),
                observation,
            ));
        }
        if request.operation_id.trim().is_empty() || request.service.trim().is_empty() {
            return Err(scale_conflict(
                "invalid_request",
                "operation_id and service must be non-empty",
                observation,
            ));
        }
        if request.expected_revision.as_deref() != observation.revision.as_deref()
            || request.current_replicas != observation.replicas
        {
            return Err(scale_conflict(
                "stale_revision",
                "scale request was not derived from the current desired state",
                observation,
            ));
        }
        let direction_matches = match request.direction {
            ScaleDirection::Up => request.desired_replicas > request.current_replicas,
            ScaleDirection::Down => request.desired_replicas < request.current_replicas,
        };
        if !direction_matches {
            return Err(scale_conflict(
                "invalid_direction",
                "scale direction does not match the requested replica transition",
                observation,
            ));
        }

        let total_other: u32 = self
            .services
            .iter()
            .filter(|(service, _)| service.as_str() != request.service)
            .map(|(_, state)| state.target_replicas)
            .sum();
        let available = self.max_instances.saturating_sub(total_other);
        if request.desired_replicas > available {
            return Err(scale_conflict(
                "capacity_exceeded",
                format!(
                    "requested {} replicas but only {} fit within host capacity",
                    request.desired_replicas, available
                ),
                observation,
            ));
        }

        let state = self
            .services
            .entry(request.service.clone())
            .or_insert_with(|| ServiceInstances {
                target_replicas: 0,
                instances: Vec::new(),
            });
        state.target_replicas = request.desired_replicas;
        let revision = self.revisions.entry(request.service.clone()).or_insert(0);
        *revision = revision.saturating_add(1);
        let response = ScaleOperationResponse {
            accepted: true,
            actual_replicas: request.desired_replicas,
            revision: Some(revision.to_string()),
            message: format!(
                "Box accepted service '{}' at {} desired replicas",
                request.service, request.desired_replicas
            ),
        };
        self.remember_operation(request.clone(), response.clone());
        Ok(response)
    }

    pub(super) fn finalize_operation_response(
        &mut self,
        request: &ScaleOperationRequest,
        response: ScaleOperationResponse,
    ) -> Result<(), String> {
        let receipt = self
            .operation_receipts
            .get_mut(&request.operation_id)
            .ok_or_else(|| {
                format!(
                    "scale operation {} has no durable acceptance receipt",
                    request.operation_id
                )
            })?;
        if receipt.request != *request {
            return Err(format!(
                "scale operation {} was accepted with different intent",
                request.operation_id
            ));
        }
        if response.revision != receipt.response.revision {
            return Err(format!(
                "scale operation {} completion changed its authority revision",
                request.operation_id
            ));
        }
        receipt.response = response;
        Ok(())
    }

    fn remember_operation(
        &mut self,
        request: ScaleOperationRequest,
        response: ScaleOperationResponse,
    ) {
        while self.operation_order.len() >= self.max_operation_receipts {
            if let Some(expired) = self.operation_order.pop_front() {
                self.operation_receipts.remove(&expired);
            }
        }
        self.operation_order.push_back(request.operation_id.clone());
        self.operation_receipts.insert(
            request.operation_id.clone(),
            ScaleOperationReceipt { request, response },
        );
    }

    pub(super) fn authority_state(&self) -> ScaleAuthorityState {
        let receipts = self
            .operation_order
            .iter()
            .filter_map(|operation| self.operation_receipts.get(operation).cloned())
            .collect();
        ScaleAuthorityState {
            schema_version: SCALE_OPERATION_SCHEMA_VERSION,
            targets: self
                .services
                .iter()
                .map(|(service, state)| (service.clone(), state.target_replicas))
                .collect(),
            revisions: self.revisions.clone(),
            receipts,
        }
    }

    pub(super) fn restore_authority_state(
        &mut self,
        state: ScaleAuthorityState,
    ) -> Result<(), String> {
        if state.schema_version != SCALE_OPERATION_SCHEMA_VERSION {
            return Err(format!(
                "unsupported scale authority schema {}",
                state.schema_version
            ));
        }
        if state
            .targets
            .keys()
            .any(|service| service.trim().is_empty())
            || state
                .revisions
                .keys()
                .any(|service| service.trim().is_empty())
        {
            return Err("scale authority contains an empty service identity".to_string());
        }
        if state
            .targets
            .keys()
            .any(|service| !state.revisions.contains_key(service))
        {
            return Err("scale authority target is missing its revision".to_string());
        }

        let mut receipts = HashMap::new();
        let mut order = VecDeque::new();
        for receipt in state.receipts {
            let operation_id = receipt.request.operation_id.clone();
            if operation_id.trim().is_empty() || receipts.contains_key(&operation_id) {
                return Err("scale authority contains an invalid operation receipt".to_string());
            }
            order.push_back(operation_id.clone());
            receipts.insert(operation_id, receipt);
        }
        while order.len() > self.max_operation_receipts {
            if let Some(expired) = order.pop_front() {
                receipts.remove(&expired);
            }
        }

        self.services = state
            .targets
            .into_iter()
            .map(|(service, target_replicas)| {
                (
                    service,
                    ServiceInstances {
                        target_replicas,
                        instances: Vec::new(),
                    },
                )
            })
            .collect();
        self.revisions = state.revisions;
        self.operation_receipts = receipts;
        self.operation_order = order;
        Ok(())
    }

    /// Process a scale request and return the response.
    ///
    /// This determines how many instances to create or destroy
    /// but does not actually start/stop VMs — the caller is responsible
    /// for that based on the response.
    pub fn process_request(&mut self, request: &ScaleRequest) -> ScaleResponse {
        let service = &request.service;
        let desired = request.replicas;

        // Compute total instances in other services first (before mutable borrow)
        let total_other: u32 = self
            .services
            .iter()
            .filter(|(k, _)| k.as_str() != service)
            .map(|(_, v)| v.instances.len() as u32)
            .sum();

        // Get or create service entry
        let svc = self
            .services
            .entry(service.clone())
            .or_insert_with(|| ServiceInstances {
                target_replicas: 0,
                instances: Vec::new(),
            });

        let current = svc.instances.len() as u32;

        // Check capacity
        let available = self.max_instances.saturating_sub(total_other);
        let target = desired.min(available);

        svc.target_replicas = target;

        let accepted = target == desired;
        let error = if !accepted {
            Some(format!(
                "Capped to {} instances (max {} total, {} used by other services)",
                target, self.max_instances, total_other
            ))
        } else {
            None
        };

        let instances: Vec<InstanceInfo> = svc
            .instances
            .iter()
            .map(|inst| InstanceInfo {
                id: inst.id.clone(),
                state: inst.state,
                service: service.clone(),
                created_at: inst.created_at,
                ready_at: inst.ready_at,
                endpoint: inst.endpoint.clone(),
                health: inst.health.clone(),
            })
            .collect();

        ScaleResponse {
            request_id: request.request_id.clone(),
            accepted,
            current_replicas: current,
            target_replicas: target,
            instances,
            error,
        }
    }

    /// Register a new instance for a service.
    pub fn register_instance(&mut self, service: &str, instance_id: &str, endpoint: Option<&str>) {
        let svc = self
            .services
            .entry(service.to_string())
            .or_insert_with(|| ServiceInstances {
                target_replicas: 0,
                instances: Vec::new(),
            });

        // Don't register duplicates
        if svc.instances.iter().any(|i| i.id == instance_id) {
            return;
        }

        svc.instances.push(TrackedInstance {
            id: instance_id.to_string(),
            state: InstanceState::Creating,
            created_at: Utc::now(),
            ready_at: None,
            endpoint: endpoint.map(|s| s.to_string()),
            health: InstanceHealth::default(),
        });
    }

    /// Update an instance's state and emit a transition event.
    pub fn update_state(
        &mut self,
        service: &str,
        instance_id: &str,
        new_state: InstanceState,
    ) -> Option<InstanceEvent> {
        let svc = self.services.get_mut(service)?;
        let inst = svc.instances.iter_mut().find(|i| i.id == instance_id)?;

        let old_state = inst.state;
        if old_state == new_state {
            return None;
        }

        inst.state = new_state;
        if new_state == InstanceState::Ready && inst.ready_at.is_none() {
            inst.ready_at = Some(Utc::now());
        }

        let event = InstanceEvent::transition(instance_id, service, old_state, new_state);
        self.push_event(event.clone());
        Some(event)
    }

    /// Update an instance's health metrics.
    pub fn update_health(&mut self, service: &str, instance_id: &str, health: InstanceHealth) {
        if let Some(svc) = self.services.get_mut(service) {
            if let Some(inst) = svc.instances.iter_mut().find(|i| i.id == instance_id) {
                inst.health = health;
            }
        }
    }

    /// Update an instance's endpoint.
    pub fn update_endpoint(&mut self, service: &str, instance_id: &str, endpoint: &str) {
        if let Some(svc) = self.services.get_mut(service) {
            if let Some(inst) = svc.instances.iter_mut().find(|i| i.id == instance_id) {
                inst.endpoint = Some(endpoint.to_string());
            }
        }
    }

    /// Remove an instance from tracking.
    pub fn deregister_instance(&mut self, service: &str, instance_id: &str) -> bool {
        if let Some(svc) = self.services.get_mut(service) {
            let before = svc.instances.len();
            svc.instances.retain(|i| i.id != instance_id);
            return svc.instances.len() < before;
        }
        false
    }

    /// Get instances that need to be created (target > current running).
    pub fn instances_to_create(&self, service: &str) -> u32 {
        if let Some(svc) = self.services.get(service) {
            let active = svc
                .instances
                .iter()
                .filter(|i| !matches!(i.state, InstanceState::Stopped | InstanceState::Failed))
                .count() as u32;
            svc.target_replicas.saturating_sub(active)
        } else {
            0
        }
    }

    /// Get instances that should be stopped (current > target).
    /// Returns instance IDs to stop, preferring non-busy instances.
    pub fn instances_to_stop(&self, service: &str) -> Vec<String> {
        if let Some(svc) = self.services.get(service) {
            let active: Vec<&TrackedInstance> = svc
                .instances
                .iter()
                .filter(|i| {
                    !matches!(
                        i.state,
                        InstanceState::Stopped
                            | InstanceState::Failed
                            | InstanceState::Stopping
                            | InstanceState::Draining
                    )
                })
                .collect();

            let excess = (active.len() as u32).saturating_sub(svc.target_replicas);
            if excess == 0 {
                return Vec::new();
            }

            // Prefer stopping idle (Ready) instances over Busy ones
            let mut candidates: Vec<&TrackedInstance> = active;
            candidates.sort_by_key(|i| match i.state {
                InstanceState::Ready => 0,    // Stop idle first
                InstanceState::Creating => 1, // Then creating
                InstanceState::Booting => 2,  // Then booting
                InstanceState::Busy => 3,     // Busy last
                _ => 4,
            });

            candidates
                .iter()
                .take(excess as usize)
                .map(|i| i.id.clone())
                .collect()
        } else {
            Vec::new()
        }
    }

    /// Get all ready instances for a service (for traffic routing).
    pub fn ready_instances(&self, service: &str) -> Vec<InstanceInfo> {
        if let Some(svc) = self.services.get(service) {
            svc.instances
                .iter()
                .filter(|i| i.state == InstanceState::Ready)
                .map(|i| InstanceInfo {
                    id: i.id.clone(),
                    state: i.state,
                    service: service.to_string(),
                    created_at: i.created_at,
                    ready_at: i.ready_at,
                    endpoint: i.endpoint.clone(),
                    health: i.health.clone(),
                })
                .collect()
        } else {
            Vec::new()
        }
    }

    /// Get the total number of instances across all services.
    pub fn total_instances(&self) -> u32 {
        self.services
            .values()
            .map(|s| s.instances.len() as u32)
            .sum()
    }

    /// Get the number of instances for a specific service.
    pub fn service_instance_count(&self, service: &str) -> u32 {
        self.services
            .get(service)
            .map(|s| s.instances.len() as u32)
            .unwrap_or(0)
    }

    /// List all tracked services.
    pub fn services(&self) -> Vec<String> {
        self.services.keys().cloned().collect()
    }

    /// Get recent events.
    pub fn recent_events(&self, limit: usize) -> &[InstanceEvent] {
        let start = self.events.len().saturating_sub(limit);
        &self.events[start..]
    }

    fn push_event(&mut self, event: InstanceEvent) {
        self.events.push(event);
        if self.events.len() > self.max_events {
            self.events.drain(..self.events.len() - self.max_events);
        }
    }

    /// Aggregate health metrics for a service (for autoscaler decisions).
    pub fn service_health(&self, service: &str) -> ServiceHealth {
        let svc = match self.services.get(service) {
            Some(s) => s,
            None => return ServiceHealth::default(),
        };

        let active: Vec<&TrackedInstance> = svc
            .instances
            .iter()
            .filter(|i| matches!(i.state, InstanceState::Ready | InstanceState::Busy))
            .collect();

        if active.is_empty() {
            return ServiceHealth {
                active_instances: 0,
                ready_instances: 0,
                busy_instances: 0,
                ..Default::default()
            };
        }

        let ready_count = active
            .iter()
            .filter(|i| i.state == InstanceState::Ready)
            .count() as u32;
        let busy_count = active
            .iter()
            .filter(|i| i.state == InstanceState::Busy)
            .count() as u32;

        let mut total_cpu = 0.0f64;
        let mut total_mem = 0u64;
        let mut total_inflight = 0u32;
        let mut cpu_count = 0u32;
        let mut unhealthy = 0u32;

        for inst in &active {
            if let Some(cpu) = inst.health.cpu_percent {
                total_cpu += cpu as f64;
                cpu_count += 1;
            }
            if let Some(mem) = inst.health.memory_bytes {
                total_mem += mem;
            }
            total_inflight += inst.health.inflight_requests;
            if !inst.health.healthy {
                unhealthy += 1;
            }
        }

        ServiceHealth {
            active_instances: active.len() as u32,
            ready_instances: ready_count,
            busy_instances: busy_count,
            avg_cpu_percent: if cpu_count > 0 {
                Some((total_cpu / cpu_count as f64) as f32)
            } else {
                None
            },
            total_memory_bytes: total_mem,
            total_inflight_requests: total_inflight,
            unhealthy_instances: unhealthy,
        }
    }

    /// Initiate graceful drain for an instance.
    ///
    /// Transitions the instance to `Draining` state. The caller should:
    /// 1. Stop routing new requests to this instance
    /// 2. Wait for in-flight requests to complete (or timeout)
    /// 3. Call `complete_drain()` to transition to `Stopping`
    pub fn start_drain(&mut self, service: &str, instance_id: &str) -> Option<InstanceEvent> {
        let svc = self.services.get_mut(service)?;
        let inst = svc.instances.iter_mut().find(|i| i.id == instance_id)?;

        // Can only drain from Ready or Busy
        if !matches!(inst.state, InstanceState::Ready | InstanceState::Busy) {
            return None;
        }

        let old_state = inst.state;
        inst.state = InstanceState::Draining;

        let event =
            InstanceEvent::transition(instance_id, service, old_state, InstanceState::Draining)
                .with_message("Graceful drain initiated");
        self.push_event(event.clone());
        Some(event)
    }

    /// Complete a drain and transition to Stopping.
    ///
    /// Called after in-flight requests have completed or the drain timeout expired.
    pub fn complete_drain(&mut self, service: &str, instance_id: &str) -> Option<InstanceEvent> {
        let svc = self.services.get_mut(service)?;
        let inst = svc.instances.iter_mut().find(|i| i.id == instance_id)?;

        if inst.state != InstanceState::Draining {
            return None;
        }

        inst.state = InstanceState::Stopping;

        let event = InstanceEvent::transition(
            instance_id,
            service,
            InstanceState::Draining,
            InstanceState::Stopping,
        )
        .with_message("Drain complete, stopping instance");
        self.push_event(event.clone());
        Some(event)
    }

    /// Check if a draining instance has no in-flight requests.
    pub fn is_drain_complete(&self, service: &str, instance_id: &str) -> bool {
        if let Some(svc) = self.services.get(service) {
            if let Some(inst) = svc.instances.iter().find(|i| i.id == instance_id) {
                return inst.state == InstanceState::Draining && inst.health.inflight_requests == 0;
            }
        }
        false
    }

    /// Get all instances currently draining.
    pub fn draining_instances(&self, service: &str) -> Vec<String> {
        if let Some(svc) = self.services.get(service) {
            svc.instances
                .iter()
                .filter(|i| i.state == InstanceState::Draining)
                .map(|i| i.id.clone())
                .collect()
        } else {
            Vec::new()
        }
    }
}

fn scale_conflict(
    code: impl Into<String>,
    message: impl Into<String>,
    observation: ScaleObservation,
) -> ScaleOperationConflict {
    ScaleOperationConflict {
        code: code.into(),
        message: message.into(),
        observation,
    }
}