a3s_box_runtime/pool/

warm_pool.rs

//! WarmPool — Pre-warmed pool of ready-to-use MicroVMs.
//!
//! Maintains a set of pre-booted VMs in `Ready` state so that
//! `acquire()` can return a VM instantly without waiting for boot.

use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::{Duration, Instant};

use a3s_box_core::config::{BoxConfig, PoolConfig};
use a3s_box_core::error::{BoxError, Result};
use a3s_box_core::event::{BoxEvent, EventEmitter};
use tokio::sync::{watch, Mutex, OwnedSemaphorePermit, Semaphore};
use tokio::task::JoinHandle;

use crate::pool::scaler::PoolScaler;
use crate::vm::VmManager;

/// A pre-warmed VM waiting in the pool.
struct WarmVm {
    /// The ready VM manager instance.
    vm: VmManager,
    /// When this VM was added to the pool.
    created_at: Instant,
}

type BootVmFuture<'a> = Pin<Box<dyn Future<Output = Result<VmManager>> + Send + 'a>>;

/// Keeps the in-flight warm-pool boot gauge balanced even when a boot task is
/// cancelled or panics while its JoinSet is being drained.
struct BootMetricGuard {
    metrics: Option<crate::prom::RuntimeMetrics>,
}

impl BootMetricGuard {
    fn new(metrics: Option<crate::prom::RuntimeMetrics>) -> Self {
        if let Some(metrics) = &metrics {
            metrics.warm_pool_boots_inflight.inc();
        }
        Self { metrics }
    }
}

impl Drop for BootMetricGuard {
    fn drop(&mut self) {
        if let Some(metrics) = &self.metrics {
            metrics.warm_pool_boots_inflight.dec();
        }
    }
}

/// Acquire the per-pool and optional daemon-wide permits in one consistent
/// order. Keeping the order identical for eager and on-demand boots avoids a
/// cross-pool semaphore cycle while a daemon is filling multiple images.
async fn acquire_boot_permits(
    boot_limiter: Arc<Semaphore>,
    global_boot_limiter: Option<Arc<Semaphore>>,
) -> Result<(OwnedSemaphorePermit, Option<OwnedSemaphorePermit>)> {
    let pool_permit = boot_limiter
        .acquire_owned()
        .await
        .map_err(|_| BoxError::PoolError("Warm-pool boot limiter closed".to_string()))?;
    let global_permit = match global_boot_limiter {
        Some(limiter) => Some(limiter.acquire_owned().await.map_err(|_| {
            BoxError::PoolError("Warm-pool global boot limiter closed".to_string())
        })?),
        None => None,
    };
    Ok((pool_permit, global_permit))
}

/// Statistics about the warm pool.
#[derive(Debug, Clone)]
pub struct PoolStats {
    /// Number of idle VMs ready for acquisition.
    pub idle_count: usize,
    /// Total number of VMs created by this pool (including acquired ones).
    pub total_created: u64,
    /// Total number of VMs acquired from the pool.
    pub total_acquired: u64,
    /// Total number of VMs released back to the pool.
    pub total_released: u64,
    /// Total number of VMs evicted due to idle TTL.
    pub total_evicted: u64,
}

/// A pre-warmed pool of ready-to-use MicroVMs.
///
/// The pool maintains `min_idle` VMs in `Ready` state. When a VM is
/// acquired, the pool spawns a replacement in the background. Idle VMs
/// that exceed `idle_ttl_secs` are automatically evicted.
///
/// # Usage
///
/// ```ignore
/// let pool = WarmPool::start(pool_config, box_config, emitter).await?;
/// let vm = pool.acquire().await?;  // Instant if pool has capacity
/// // ... use vm ...
/// pool.release(vm).await?;         // Return to pool or destroy
/// pool.drain().await?;             // Graceful shutdown
/// ```
pub struct WarmPool {
    /// Pool configuration.
    config: PoolConfig,
    /// Base BoxConfig template for creating new VMs.
    box_config: BoxConfig,
    /// Idle VMs ready for acquisition.
    idle: Arc<Mutex<Vec<WarmVm>>>,
    /// Pool statistics.
    stats: Arc<Mutex<PoolStats>>,
    /// Event emitter for pool lifecycle events.
    event_emitter: EventEmitter,
    /// Background replenishment task handle.
    replenish_handle: Option<JoinHandle<()>>,
    /// Shutdown signal sender.
    shutdown_tx: watch::Sender<bool>,
    /// Shutdown signal receiver (cloned for background task).
    shutdown_rx: watch::Receiver<bool>,
    /// Autoscaler for dynamic min_idle adjustment (None if scaling disabled).
    scaler: Option<Arc<Mutex<PoolScaler>>>,
    /// Prometheus metrics (optional).
    metrics: Option<crate::prom::RuntimeMetrics>,
    /// Per-pool boot limiter shared by eager fill and on-demand misses.
    boot_limiter: Arc<Semaphore>,
    /// Optional daemon-wide boot limiter shared by all image pools.
    global_boot_limiter: Option<Arc<Semaphore>>,
    /// Snapshot-fork template state (built lazily on first fill when
    /// `config.snapshot_fork`): the file-backed RAM image + state file every other
    /// pool VM restores from. Caches an `Unavailable` verdict so a build failure
    /// (native VM snapshot unsupported on this build) is not re-attempted on every
    /// fill — the pool cold-boots instead.
    template: Arc<Mutex<TemplateState>>,
}

/// A built snapshot-fork template: the shared RAM image + state file that pool VMs
/// restore from (MAP_PRIVATE CoW of the RAM file).
#[derive(Clone)]
struct PoolTemplate {
    mem_file: String,
    state_file: String,
    rootfs_cache_key: Option<String>,
}

/// How many consecutive template-build failures are tolerated before the
/// verdict becomes permanently `Unavailable`. A transient failure (host
/// resource pressure, a source VM slow to bind its snapshot socket) presents
/// identically to "snapshot unsupported by this libkrun build" ("snapshot
/// socket never appeared"), so a bounded retry avoids permanently downgrading
/// the whole pool to cold-boot on a one-off hiccup, while still giving up on a
/// genuinely-unsupported host after a few attempts.
const MAX_TEMPLATE_BUILD_FAILURES: u32 = 3;

/// Cached state of the snapshot-fork template.
enum TemplateState {
    /// Not built yet — the first snapshot-fork fill attempts the build.
    Unbuilt,
    /// Built and ready; pool VMs restore from it.
    Ready(PoolTemplate),
    /// The last build failed but is still retryable; carries the consecutive
    /// failure count. A later fill retries until it reaches
    /// `MAX_TEMPLATE_BUILD_FAILURES`, then it becomes `Unavailable`.
    Failing(u32),
    /// The build failed permanently (native VM snapshot unavailable on this
    /// build/platform, or too many consecutive failures). Cached so it is not
    /// retried — `boot_or_restore` cold-boots instead.
    Unavailable,
}

#[derive(Clone, Copy)]
enum InitialFill {
    /// Boot the configured `min_idle` count before returning from `start`.
    Eager,
    /// Boot one ready VM before returning and let maintenance fill the rest.
    FirstReady,
}

impl WarmPool {
    /// Create and start the warm pool.
    ///
    /// Spawns `min_idle` VMs in the background and starts the
    /// replenishment/eviction loop.
    pub async fn start(
        config: PoolConfig,
        box_config: BoxConfig,
        event_emitter: EventEmitter,
    ) -> Result<Self> {
        Self::start_with_metrics(config, box_config, event_emitter, None).await
    }

    /// Create and start the warm pool with an optional shared metrics sink.
    ///
    /// The sink is installed before the initial fill so the first pre-warmed
    /// VMs contribute boot, cache, and pool metrics. [`Self::start`] remains
    /// the compatibility entry point for callers that do not need metrics.
    pub async fn start_with_metrics(
        config: PoolConfig,
        box_config: BoxConfig,
        event_emitter: EventEmitter,
        metrics: Option<crate::prom::RuntimeMetrics>,
    ) -> Result<Self> {
        Self::start_with_metrics_and_boot_limiter(config, box_config, event_emitter, metrics, None)
            .await
    }

    /// Create and start the warm pool with optional metrics and a shared
    /// daemon-wide boot limiter.
    ///
    /// `max_concurrent_boots` still limits each pool independently. When a
    /// shared limiter is supplied, it additionally caps aggregate boots across
    /// all pools that use it (for example, a multi-image pool daemon).
    pub async fn start_with_metrics_and_boot_limiter(
        config: PoolConfig,
        box_config: BoxConfig,
        event_emitter: EventEmitter,
        metrics: Option<crate::prom::RuntimeMetrics>,
        global_boot_limiter: Option<Arc<Semaphore>>,
    ) -> Result<Self> {
        Self::start_with_metrics_and_boot_limiter_with_fill(
            config,
            box_config,
            event_emitter,
            metrics,
            global_boot_limiter,
            InitialFill::Eager,
        )
        .await
    }

    /// Create a warm pool with one ready VM on the critical path.
    ///
    /// This is intended for lazy, multi-image daemons: the first request for a
    /// new pool can run as soon as one VM is ready while the maintenance loop
    /// fills the remaining `min_idle` capacity in the background. The regular
    /// [`Self::start_with_metrics_and_boot_limiter`] API keeps its eager-fill
    /// behavior for explicit pool startup.
    pub async fn start_with_metrics_and_boot_limiter_first_ready(
        config: PoolConfig,
        box_config: BoxConfig,
        event_emitter: EventEmitter,
        metrics: Option<crate::prom::RuntimeMetrics>,
        global_boot_limiter: Option<Arc<Semaphore>>,
    ) -> Result<Self> {
        Self::start_with_metrics_and_boot_limiter_with_fill(
            config,
            box_config,
            event_emitter,
            metrics,
            global_boot_limiter,
            InitialFill::FirstReady,
        )
        .await
    }

    async fn start_with_metrics_and_boot_limiter_with_fill(
        config: PoolConfig,
        box_config: BoxConfig,
        event_emitter: EventEmitter,
        metrics: Option<crate::prom::RuntimeMetrics>,
        global_boot_limiter: Option<Arc<Semaphore>>,
        initial_fill: InitialFill,
    ) -> Result<Self> {
        if config.max_size == 0 {
            return Err(BoxError::PoolError(
                "Pool max_size must be greater than 0".to_string(),
            ));
        }
        if config.min_idle > config.max_size {
            return Err(BoxError::PoolError(format!(
                "Pool min_idle ({}) cannot exceed max_size ({})",
                config.min_idle, config.max_size
            )));
        }
        if config.max_concurrent_boots == 0 {
            return Err(BoxError::PoolError(
                "Pool max_concurrent_boots must be greater than 0".to_string(),
            ));
        }

        let idle = Arc::new(Mutex::new(Vec::with_capacity(config.max_size)));
        let stats = Arc::new(Mutex::new(PoolStats {
            idle_count: 0,
            total_created: 0,
            total_acquired: 0,
            total_released: 0,
            total_evicted: 0,
        }));
        let (shutdown_tx, shutdown_rx) = watch::channel(false);

        let scaler = if config.scaling.enabled {
            Some(Arc::new(Mutex::new(PoolScaler::new(
                config.scaling.clone(),
                config.min_idle,
                config.max_size,
            ))))
        } else {
            None
        };

        let boot_limiter = Arc::new(Semaphore::new(config.max_concurrent_boots));
        let mut pool = Self {
            config,
            box_config,
            idle,
            stats,
            event_emitter,
            replenish_handle: None,
            shutdown_tx,
            shutdown_rx,
            scaler,
            metrics,
            boot_limiter,
            global_boot_limiter,
            template: Arc::new(Mutex::new(TemplateState::Unbuilt)),
        };

        if let Some(metrics) = &pool.metrics {
            metrics.warm_pool_capacity.set(pool.config.max_size as i64);
        }

        // Initial fill. Lazy pools only wait for the first ready VM; the
        // maintenance loop immediately schedules the remaining capacity.
        let initial_target = match initial_fill {
            InitialFill::Eager => pool.config.min_idle,
            InitialFill::FirstReady => pool.config.min_idle.min(1),
        };
        let initial_fill_started = Instant::now();
        pool.fill_to_target(initial_target).await;
        if let Some(metrics) = &pool.metrics {
            metrics
                .warm_pool_initial_fill_duration
                .observe(initial_fill_started.elapsed().as_secs_f64());
        }

        // Start background maintenance loop
        let handle = pool.spawn_maintenance_loop();
        pool.replenish_handle = Some(handle);

        tracing::info!(
            min_idle = pool.config.min_idle,
            max_size = pool.config.max_size,
            idle_ttl_secs = pool.config.idle_ttl_secs,
            "Warm pool started"
        );

        Ok(pool)
    }

    /// Attach Prometheus metrics to this pool.
    pub fn set_metrics(&mut self, metrics: crate::prom::RuntimeMetrics) {
        metrics.warm_pool_capacity.set(self.config.max_size as i64);
        metrics.warm_pool_size.set(
            self.idle
                .try_lock()
                .map(|idle| idle.len() as i64)
                .unwrap_or_default(),
        );
        self.metrics = Some(metrics);
    }

    fn sync_idle_metric(metrics: Option<&crate::prom::RuntimeMetrics>, idle_count: usize) {
        if let Some(metrics) = metrics {
            metrics.warm_pool_size.set(idle_count as i64);
        }
    }

    /// Acquire a ready VM from the pool.
    ///
    /// If an idle VM is available, returns it immediately.
    /// Otherwise, boots a new VM on demand (slower path).
    pub async fn acquire(&self) -> Result<VmManager> {
        // Try to pop an idle VM
        {
            let mut idle = self.idle.lock().await;
            if let Some(warm_vm) = idle.pop() {
                let mut stats = self.stats.lock().await;
                stats.total_acquired += 1;
                stats.idle_count = idle.len();

                // Record hit for autoscaler
                if let Some(ref scaler) = self.scaler {
                    scaler.lock().await.record_acquire(true);
                }

                if let Some(ref m) = self.metrics {
                    m.warm_pool_hits.inc();
                    m.warm_pool_size.set(idle.len() as i64);
                }

                self.event_emitter.emit(BoxEvent::with_string(
                    "pool.vm.acquired",
                    format!("Acquired VM {} from pool", warm_vm.vm.box_id()),
                ));

                tracing::debug!(
                    box_id = %warm_vm.vm.box_id(),
                    idle_remaining = idle.len(),
                    "Acquired VM from warm pool"
                );

                return Ok(warm_vm.vm);
            }
        }

        // No idle VM available — boot one on demand (miss)
        tracing::info!("No idle VM in pool, booting on demand");

        // Record miss for autoscaler
        if let Some(ref scaler) = self.scaler {
            scaler.lock().await.record_acquire(false);
        }

        if let Some(ref m) = self.metrics {
            m.warm_pool_misses.inc();
        }

        let vm = self.boot_new_vm().await?;

        let mut stats = self.stats.lock().await;
        stats.total_acquired += 1;

        Ok(vm)
    }

    /// Release a VM back to the pool.
    ///
    /// If the pool is at capacity, the VM is destroyed instead.
    pub async fn release(&self, vm: VmManager) -> Result<()> {
        let mut idle = self.idle.lock().await;

        // Don't return a VM to a pool that is shutting down: drain_idle has (or
        // soon will have) cleared `idle` and won't run again, so a push here leaks
        // the VM (no Drop reaper). Checked under the idle lock so it is atomic with
        // a concurrent drain_idle. Destroy the VM instead.
        if *self.shutdown_rx.borrow() {
            drop(idle);
            let mut vm = vm;
            vm.destroy().await?;
            return Ok(());
        }

        if idle.len() >= self.config.max_size {
            // Pool is full — destroy the VM
            drop(idle); // Release lock before async destroy
            let mut vm = vm;
            vm.destroy().await?;

            tracing::debug!(
                box_id = %vm.box_id(),
                "Pool full, destroyed released VM"
            );
            return Ok(());
        }

        let box_id = vm.box_id().to_string();
        idle.push(WarmVm {
            vm,
            created_at: Instant::now(),
        });

        let mut stats = self.stats.lock().await;
        stats.total_released += 1;
        stats.idle_count = idle.len();

        if let Some(ref m) = self.metrics {
            m.warm_pool_size.set(idle.len() as i64);
        }

        self.event_emitter.emit(BoxEvent::with_string(
            "pool.vm.released",
            format!("Released VM {} back to pool", box_id),
        ));

        tracing::debug!(
            box_id = %box_id,
            idle_count = idle.len(),
            "Released VM back to warm pool"
        );

        Ok(())
    }

    /// Get current pool statistics.
    pub async fn stats(&self) -> PoolStats {
        self.stats.lock().await.clone()
    }

    /// Get the number of idle VMs currently in the pool.
    pub async fn idle_count(&self) -> usize {
        self.idle.lock().await.len()
    }

    /// Signal the pool to shutdown. This signals the background task to stop
    /// replenishing and sets the shutdown flag. VMs will continue to exist
    /// until the pool is drained or dropped.
    pub fn signal_shutdown(&self) {
        let _ = self.shutdown_tx.send(true);
        tracing::info!("Warm pool shutdown signaled");
    }

    /// Gracefully drain all VMs and stop the pool.
    pub async fn drain(&mut self) -> Result<()> {
        // Signal shutdown to background task
        let _ = self.shutdown_tx.send(true);

        // Wait for background task to finish
        if let Some(handle) = self.replenish_handle.take() {
            let _ = handle.await;
        }

        // Detach idle VMs before destroying them. VM teardown is asynchronous
        // and must not hold the pool lock, otherwise acquire/release and the
        // maintenance loop can be blocked for the entire drain duration.
        let idle_vms = {
            let mut idle = self.idle.lock().await;
            let idle_vms = idle.drain(..).collect::<Vec<_>>();
            Self::sync_idle_metric(self.metrics.as_ref(), idle.len());
            idle_vms
        };
        let count = idle_vms.len();

        for warm_vm in idle_vms {
            let mut vm = warm_vm.vm;
            if let Err(e) = vm.destroy().await {
                tracing::warn!(
                    box_id = %vm.box_id(),
                    error = %e,
                    "Failed to destroy pooled VM during drain"
                );
            }
        }

        let mut stats = self.stats.lock().await;
        stats.idle_count = 0;

        self.event_emitter.emit(BoxEvent::empty("pool.drained"));

        tracing::info!(destroyed = count, "Warm pool drained");

        Ok(())
    }

    /// Destroy all idle VMs without consuming the pool (`&self`), so it can be
    /// shut down from behind an `Arc` (e.g. a daemon serving concurrent requests).
    /// Pair with [`Self::signal_shutdown`] first to stop the background replenisher;
    /// its task then exits on its own (it watches the shutdown channel).
    pub async fn drain_idle(&self) -> Result<()> {
        // Detach idle VMs before destroying them. Keeping `idle` locked while
        // awaiting VM teardown blocks concurrent acquire/release operations and
        // widens the shutdown race window for a replenishment batch.
        let idle_vms = {
            let mut idle = self.idle.lock().await;
            let idle_vms = idle.drain(..).collect::<Vec<_>>();
            Self::sync_idle_metric(self.metrics.as_ref(), idle.len());
            idle_vms
        };
        let count = idle_vms.len();
        for warm_vm in idle_vms {
            let mut vm = warm_vm.vm;
            if let Err(e) = vm.destroy().await {
                tracing::warn!(
                    box_id = %vm.box_id(),
                    error = %e,
                    "Failed to destroy pooled VM during drain_idle"
                );
            }
        }
        self.stats.lock().await.idle_count = 0;
        tracing::info!(destroyed = count, "Warm pool idle VMs drained");
        Ok(())
    }

    /// Remove and destroy specific idle VMs by their box IDs.
    ///
    /// Used when a fill partially fails and needs to roll back
    /// successfully added VMs.
    async fn remove_idle_vms(&self, box_ids: &[String]) {
        // First pass: collect indices of VMs to remove
        let indices_to_remove: Vec<usize> = {
            let idle = self.idle.lock().await;
            idle.iter()
                .enumerate()
                .filter(|(_, wm)| box_ids.iter().any(|id| id == wm.vm.box_id()))
                .map(|(i, _)| i)
                .collect()
        };

        if indices_to_remove.is_empty() {
            return;
        }

        // Second pass: remove and collect VMs to destroy
        // We do this in reverse order to avoid index shifting issues
        let mut to_destroy: Vec<WarmVm> = Vec::new();
        {
            let mut idle = self.idle.lock().await;
            for idx in indices_to_remove.into_iter().rev() {
                if idx < idle.len() {
                    let warm_vm = idle.remove(idx);
                    to_destroy.push(warm_vm);
                }
            }
        }

        // Update stats before destroying (approximate, since VMs still exist in to_destroy)
        {
            let idle_count = self.idle.lock().await.len();
            if let Ok(mut stats) = self.stats.try_lock() {
                stats.idle_count = idle_count;
            }
            Self::sync_idle_metric(self.metrics.as_ref(), idle_count);
        }

        // Destroy collected VMs (outside of pool lock)
        for warm_vm in to_destroy {
            let box_id = warm_vm.vm.box_id().to_string();
            let mut vm = warm_vm.vm;
            if let Err(e) = vm.destroy().await {
                tracing::warn!(
                    box_id = %box_id,
                    error = %e,
                    "Failed to destroy VM during pool fill rollback"
                );
            } else {
                tracing::debug!(box_id = %box_id, "Destroyed VM during pool fill rollback");
            }
        }
    }

    /// Boot a new VM using the pool's template config.
    async fn boot_new_vm(&self) -> Result<VmManager> {
        let _boot_permits =
            acquire_boot_permits(self.boot_limiter.clone(), self.global_boot_limiter.clone())
                .await?;
        let _boot_guard = BootMetricGuard::new(self.metrics.clone());
        let result = Self::boot_or_restore(
            self.config.snapshot_fork,
            &self.box_config,
            &self.event_emitter,
            &self.template,
        )
        .await;
        if result.is_err() {
            if let Some(metrics) = &self.metrics {
                metrics.warm_pool_boot_failures_total.inc();
            }
        }
        let vm = result?;

        let mut stats = self.stats.lock().await;
        stats.total_created += 1;

        self.event_emitter.emit(BoxEvent::with_string(
            "pool.vm.created",
            format!("Booted new VM {}", vm.box_id()),
        ));

        Ok(vm)
    }

    /// Fill one slot: restore from the snapshot-fork template when enabled, else cold
    /// boot. Static so both `boot_new_vm` and the background replenish task use it.
    fn boot_or_restore<'a>(
        snapshot_fork: bool,
        box_config: &'a BoxConfig,
        event_emitter: &'a EventEmitter,
        template: &'a Arc<Mutex<TemplateState>>,
    ) -> BootVmFuture<'a> {
        Box::pin(async move {
            if snapshot_fork && crate::vm::native_snapshot_fork_supported() {
                // Try the snapshot-fork template. If it can't be built (native VM
                // snapshot unavailable — the verdict is cached so this is attempted at
                // most once), fall back to a normal cold boot so the warm pool still
                // fills rather than failing outright.
                match Self::ensure_template(box_config, event_emitter, template).await {
                    Ok(tpl) => {
                        let mut cfg = box_config.clone();
                        cfg.snapshot_mem_file = Some(tpl.mem_file.clone());
                        cfg.restore_from = Some(tpl.state_file.clone());
                        cfg.snapshot_sock = None;
                        let mut vm = VmManager::new(cfg, event_emitter.clone());
                        vm.restore_rootfs_cache_key = tpl.rootfs_cache_key.clone();
                        let restored = async {
                            vm.boot().await?;
                            vm.wait_for_exec_available(std::time::Duration::from_secs(120))
                                .await
                        }
                        .await;
                        match restored {
                            Ok(()) => return Ok(vm),
                            Err(error) => {
                                let _ = vm.destroy_with_timeout(2000).await;
                                tracing::warn!(
                                    %error,
                                    "snapshot-fork restore failed; cold-booting this pool VM"
                                );
                            }
                        }
                    }
                    Err(error) => {
                        tracing::debug!(%error, "snapshot-fork unavailable; cold-booting this pool VM");
                    }
                }
            } else if snapshot_fork {
                tracing::debug!(
                    "snapshot-fork is unavailable on this build; cold-booting without snapshot side effects"
                );
            }
            let mut vm = VmManager::new(box_config.clone(), event_emitter.clone());
            vm.boot().await?;
            vm.wait_for_exec_available(std::time::Duration::from_secs(120))
                .await?;
            Ok(vm)
        })
    }

    /// Boot a bounded batch of VMs and return results in completion order.
    ///
    /// VM boot is expensive in both host CPU and memory. A JoinSet without a
    /// concurrency limit turns a large min_idle or autoscaler step into a
    /// resource burst, so only `max_concurrent_boots` tasks are in flight at
    /// once. Each task owns cloned inputs, allowing the set to remain
    /// `'static` while the caller retains its pool state.
    async fn boot_batch(
        snapshot_fork: bool,
        box_config: &BoxConfig,
        event_emitter: &EventEmitter,
        template: &Arc<Mutex<TemplateState>>,
        needed: usize,
        max_concurrent_boots: usize,
        metrics: Option<crate::prom::RuntimeMetrics>,
        boot_limiter: Arc<Semaphore>,
        global_boot_limiter: Option<Arc<Semaphore>>,
    ) -> Vec<Result<VmManager>> {
        if needed == 0 {
            return Vec::new();
        }

        let limit = bounded_boot_limit(needed, max_concurrent_boots);
        let mut set = tokio::task::JoinSet::new();
        let mut launched = 0usize;
        let mut results = Vec::with_capacity(needed);

        while launched < needed || !set.is_empty() {
            while launched < needed && set.len() < limit {
                let config = box_config.clone();
                let emitter = event_emitter.clone();
                let shared_template = Arc::clone(template);
                let boot_metrics = metrics.clone();
                let pool_boot_limiter = boot_limiter.clone();
                let daemon_boot_limiter = global_boot_limiter.clone();
                set.spawn(async move {
                    let _boot_permits =
                        acquire_boot_permits(pool_boot_limiter, daemon_boot_limiter).await?;
                    let _boot_guard = BootMetricGuard::new(boot_metrics);
                    WarmPool::boot_or_restore(snapshot_fork, &config, &emitter, &shared_template)
                        .await
                });
                launched += 1;
            }

            if let Some(result) = set.join_next().await {
                let result = match result {
                    Ok(result) => result,
                    Err(error) => Err(BoxError::PoolError(format!(
                        "Warm-pool boot task failed: {error}"
                    ))),
                };
                if result.is_err() {
                    if let Some(metrics) = &metrics {
                        metrics.warm_pool_boot_failures_total.inc();
                    }
                }
                results.push(result);
            }
        }

        results
    }

    /// Get the snapshot-fork template, building it once lazily. Concurrent callers
    /// wait on the lock and reuse the first result — a built template OR a cached
    /// `Unavailable` verdict, so a failed build (native VM snapshot unsupported on
    /// this build) is attempted at most once rather than re-tried (and re-timed-out)
    /// on every pool fill. Returns `Err` when unavailable so `boot_or_restore` cold
    /// boots instead.
    async fn ensure_template(
        box_config: &BoxConfig,
        event_emitter: &EventEmitter,
        template: &Arc<Mutex<TemplateState>>,
    ) -> Result<PoolTemplate> {
        if !crate::vm::native_snapshot_fork_supported() {
            return Err(BoxError::PoolError(
                "snapshot-fork requires the Linux x86_64 KVM build".to_string(),
            ));
        }
        let mut guard = template.lock().await;
        let prior_failures = match &*guard {
            TemplateState::Ready(t) => return Ok(t.clone()),
            TemplateState::Unavailable => {
                return Err(BoxError::PoolError(
                    "snapshot-fork template unavailable (native VM snapshot unsupported)"
                        .to_string(),
                ));
            }
            // Unbuilt or a still-retryable prior failure: (re)attempt the build.
            TemplateState::Failing(n) => *n,
            TemplateState::Unbuilt => 0,
        };

        match Self::build_template(box_config, event_emitter).await {
            Ok(tpl) => {
                *guard = TemplateState::Ready(tpl.clone());
                event_emitter.emit(BoxEvent::with_string(
                    "pool.template.built",
                    format!(
                        "Snapshot-fork template built for image {}",
                        box_config.image
                    ),
                ));
                Ok(tpl)
            }
            Err(error) => {
                // Bounded retry: a transient failure presents identically to
                // "snapshot unsupported", so only give up permanently after a few
                // consecutive failures rather than downgrading the pool to
                // cold-boot forever on a one-off hiccup.
                let failures = prior_failures + 1;
                if failures >= MAX_TEMPLATE_BUILD_FAILURES {
                    tracing::warn!(
                        %error, failures,
                        "snapshot-fork template build failed repeatedly; marking \
                         unavailable — the warm pool will cold-boot"
                    );
                    *guard = TemplateState::Unavailable;
                } else {
                    tracing::warn!(
                        %error, failures,
                        "snapshot-fork template build failed; will retry on a later fill"
                    );
                    *guard = TemplateState::Failing(failures);
                }
                Err(error)
            }
        }
    }

    /// Cold-boot one source VM with file-backed RAM + a trigger socket, snapshot it,
    /// and tear it down — leaving the RAM image + state file as the template.
    async fn build_template(
        box_config: &BoxConfig,
        event_emitter: &EventEmitter,
    ) -> Result<PoolTemplate> {
        let dir = a3s_box_core::dirs_home().join("pool").join(format!(
            "tpl-{:016x}",
            crate::vm::fnv1a_hash(&box_config.image)
        ));
        std::fs::create_dir_all(&dir).map_err(BoxError::IoError)?;

        // Cross-process lock on the per-image template dir. The dir is keyed only
        // by the image hash, so two processes building the same image's template
        // would write the same template.ram/template.state concurrently and
        // corrupt them. Held (via a Send File handle) across the boot+snapshot
        // awaits below; acquired off-runtime so a contended flock doesn't block a
        // worker thread.
        let lock_target = dir.clone();
        let _lock =
            tokio::task::spawn_blocking(move || crate::file_lock::FileLock::acquire(&lock_target))
                .await
                .map_err(|e| BoxError::PoolError(format!("Template lock task failed: {e}")))?
                .map_err(|e| BoxError::PoolError(format!("Failed to lock template dir: {e}")))?;

        let mem_file = dir.join("template.ram");
        let sock = dir.join("template.sock");
        let state_file = dir.join("template.state");
        let _ = std::fs::remove_file(&sock);

        // Cold-boot the source as a snapshot TEMPLATE (file-backed RAM + trigger sock).
        let mut cfg = box_config.clone();
        cfg.snapshot_mem_file = Some(mem_file.to_string_lossy().into_owned());
        cfg.snapshot_sock = Some(sock.to_string_lossy().into_owned());
        cfg.restore_from = None;
        let mut src = VmManager::new(cfg, event_emitter.clone());
        src.boot().await?;
        let rootfs_cache_key = match src.current_rootfs_cache_key() {
            Ok(key) => key,
            Err(error) => {
                let _ = src.destroy_with_timeout(2000).await;
                return Err(error);
            }
        };

        // Trigger the snapshot over libkrun's socket, then tear down the source (it is
        // left paused by the snapshot; the RAM + state files are the template).
        //
        // Destroy the source UNCONDITIONALLY: `trigger_snapshot` fails on any
        // libkrun without snapshot support (the common case), and `?`-ing out
        // here would leak the fully-booted source VM (shim process, overlay
        // mount, box dir, sockets) — neither VmManager nor ShimHandler reaps on
        // drop. Capture the result, tear down, then propagate.
        let snapshot = Self::trigger_snapshot(&sock, &state_file).await;
        let _ = src.destroy_with_timeout(2000).await;
        snapshot?;

        Ok(PoolTemplate {
            mem_file: mem_file.to_string_lossy().into_owned(),
            state_file: state_file.to_string_lossy().into_owned(),
            rootfs_cache_key,
        })
    }

    /// Send a `snapshot <state>` request to libkrun's per-template trigger socket and
    /// wait for the `ok` reply (the socket appears once the template's vCPUs run).
    ///
    /// Snapshot-fork is a Linux/KVM (Unix) feature; on non-Unix hosts the trigger
    /// socket does not exist, so this is unavailable (see the `not(unix)` stub).
    #[cfg(unix)]
    async fn trigger_snapshot(sock: &std::path::Path, state_file: &std::path::Path) -> Result<()> {
        use tokio::io::{AsyncReadExt, AsyncWriteExt};
        // The socket is bound by libkrun after the guest starts; poll briefly.
        let mut stream = None;
        for _ in 0..200 {
            match tokio::net::UnixStream::connect(sock).await {
                Ok(s) => {
                    stream = Some(s);
                    break;
                }
                Err(_) => tokio::time::sleep(std::time::Duration::from_millis(25)).await,
            }
        }
        let mut stream = stream.ok_or_else(|| {
            BoxError::PoolError(format!("snapshot socket {} never appeared", sock.display()))
        })?;
        let cmd = format!("snapshot {}\n", state_file.display());
        stream
            .write_all(cmd.as_bytes())
            .await
            .map_err(BoxError::IoError)?;
        let mut buf = [0u8; 64];
        let n = stream.read(&mut buf).await.map_err(BoxError::IoError)?;
        let reply = String::from_utf8_lossy(&buf[..n]);
        if reply.trim() == "ok" {
            Ok(())
        } else {
            Err(BoxError::PoolError(format!(
                "snapshot trigger failed: {}",
                reply.trim()
            )))
        }
    }

    /// Non-Unix stub: snapshot-fork relies on libkrun's Unix trigger socket and KVM
    /// state save/restore, neither of which exist on Windows. `--snapshot-fork` is
    /// Linux/KVM-only, so this path is never reached there in practice.
    #[cfg(not(unix))]
    async fn trigger_snapshot(
        _sock: &std::path::Path,
        _state_file: &std::path::Path,
    ) -> Result<()> {
        Err(BoxError::PoolError(
            "snapshot-fork is only supported on Linux/KVM hosts".to_string(),
        ))
    }

    /// Fill the pool to a specific idle target.
    async fn fill_to_target(&self, target: usize) {
        let current = self.idle.lock().await.len();
        let needed = target.saturating_sub(current);

        if needed == 0 {
            return;
        }

        tracing::debug!(current, needed, target, "Replenishing warm pool");

        // Track VMs added in this fill attempt so we can clean up on failure.
        let mut added_ids: Vec<String> = Vec::new();
        let mut failed = false;
        let results = Self::boot_batch(
            self.config.snapshot_fork,
            &self.box_config,
            &self.event_emitter,
            &self.template,
            needed,
            self.config.max_concurrent_boots,
            self.metrics.clone(),
            self.boot_limiter.clone(),
            self.global_boot_limiter.clone(),
        )
        .await;

        for result in results {
            match result {
                Ok(vm) => {
                    let box_id = vm.box_id().to_string();
                    let mut idle = self.idle.lock().await;
                    idle.push(WarmVm {
                        vm,
                        created_at: Instant::now(),
                    });
                    Self::sync_idle_metric(self.metrics.as_ref(), idle.len());
                    let mut stats = self.stats.lock().await;
                    stats.total_created += 1;
                    stats.idle_count = idle.len();
                    added_ids.push(box_id.clone());

                    self.event_emitter.emit(BoxEvent::with_string(
                        "pool.vm.created",
                        format!("Booted new VM {box_id}"),
                    ));
                    tracing::debug!(box_id = %box_id, "Added VM to warm pool");
                }
                Err(error) => {
                    failed = true;
                    tracing::warn!(error = %error, "Failed to boot VM for warm pool");
                }
            }
        }

        if failed && !added_ids.is_empty() {
            tracing::info!(
                count = added_ids.len(),
                "Cleaning up VMs added before pool fill failed"
            );
            self.remove_idle_vms(&added_ids).await;
        }

        self.event_emitter.emit(BoxEvent::empty("pool.replenish"));
    }

    /// Spawn the background maintenance loop.
    ///
    /// Periodically checks for:
    /// 1. Autoscaler evaluation → adjust min_idle dynamically
    /// 2. Pool below min_idle → replenish
    /// 3. Idle VMs past TTL → evict
    fn spawn_maintenance_loop(&self) -> JoinHandle<()> {
        let idle = Arc::clone(&self.idle);
        let stats = Arc::clone(&self.stats);
        let config = self.config.clone();
        let box_config = self.box_config.clone();
        let event_emitter = self.event_emitter.clone();
        let mut shutdown_rx = self.shutdown_rx.clone();
        let scaler = self.scaler.clone();
        let template = Arc::clone(&self.template);
        let metrics = self.metrics.clone();
        let boot_limiter = self.boot_limiter.clone();
        let global_boot_limiter = self.global_boot_limiter.clone();

        tokio::spawn(async move {
            let check_interval = std::time::Duration::from_secs(
                // Check every 1/5 of TTL, minimum 5 seconds
                if config.idle_ttl_secs > 0 {
                    (config.idle_ttl_secs / 5).max(5)
                } else {
                    30
                },
            );
            let mut maintenance = tokio::time::interval(check_interval);
            // The first tick is immediate, which lets a first-ready lazy pool
            // continue filling without waiting for the full maintenance period.
            maintenance.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);

            // Dynamic min_idle starts from config, adjusted by scaler
            let mut effective_min_idle = config.min_idle;
            // A provider failure should not cause the maintenance loop to
            // repeatedly launch expensive doomed boots. Back off retries while
            // keeping the normal maintenance/eviction cadence unchanged.
            let mut replenish_failures = 0u32;
            let mut next_replenish_at = Instant::now();

            loop {
                tokio::select! {
                    result = shutdown_rx.changed() => {
                        if result.is_ok() && *shutdown_rx.borrow() {
                            tracing::debug!("Pool maintenance loop shutting down");
                            break;
                        }
                    }
                    _ = maintenance.tick() => {
                        // Evict expired VMs
                        if config.idle_ttl_secs > 0 {
                            Self::evict_expired_static(
                                &idle,
                                &stats,
                                &event_emitter,
                                metrics.as_ref(),
                                config.idle_ttl_secs,
                            ).await;
                        }

                        // Evaluate autoscaler
                        if let Some(ref scaler) = scaler {
                            let mut s = scaler.lock().await;
                            let decision = s.evaluate();
                            let new_min = s.current_min_idle();
                            if new_min != effective_min_idle {
                                tracing::info!(
                                    old_min_idle = effective_min_idle,
                                    new_min_idle = new_min,
                                    ?decision,
                                    "Autoscaler adjusted min_idle"
                                );
                                event_emitter.emit(BoxEvent::with_string(
                                    "pool.autoscale",
                                    format!(
                                        "min_idle adjusted {} → {} ({:?})",
                                        effective_min_idle, new_min, decision
                                    ),
                                ));
                                effective_min_idle = new_min;
                            }
                        }

                        // Replenish if below effective min_idle
                        let current = idle.lock().await.len();
                        if current < effective_min_idle && Instant::now() >= next_replenish_at {
                            let needed = effective_min_idle - current;
                            tracing::debug!(current, needed, min_idle = effective_min_idle, "Replenishing warm pool");

                            // Overlap readiness waits while keeping the number of
                            // expensive VM boots bounded by configuration.
                            let results = Self::boot_batch(
                                config.snapshot_fork,
                                &box_config,
                                &event_emitter,
                                &template,
                                needed,
                                config.max_concurrent_boots,
                                metrics.clone(),
                                boot_limiter.clone(),
                                global_boot_limiter.clone(),
                            )
                            .await;
                            let mut batch_failed = false;
                            for result in results {
                                match result {
                                    Ok(mut vm) => {
                                        let box_id = vm.box_id().to_string();
                                        // If shutdown landed while this batch was
                                        // booting, drain_idle has already cleared
                                        // `idle` and will not run again, so a VM
                                        // pushed now leaks (no Drop reaper). Destroy
                                        // it instead. Acquire the idle lock FIRST and
                                        // re-check shutdown UNDER it: drain_idle drains
                                        // while holding this same lock (always after
                                        // signal_shutdown), so the check-and-push is
                                        // atomic against it — closing the TOCTOU window
                                        // that an unlocked `borrow()` check left open.
                                        let mut pool = idle.lock().await;
                                        if *shutdown_rx.borrow() {
                                            drop(pool);
                                            tracing::debug!(
                                                box_id = %box_id,
                                                "Pool shutting down mid-replenish; destroying freshly-booted VM"
                                            );
                                            let _ = vm.destroy_with_timeout(2000).await;
                                            continue;
                                        }
                                        pool.push(WarmVm {
                                            vm,
                                            created_at: Instant::now(),
                                        });
                                        Self::sync_idle_metric(metrics.as_ref(), pool.len());
                                        let mut s = stats.lock().await;
                                        s.total_created += 1;
                                        s.idle_count = pool.len();
                                        drop(s);
                                        drop(pool);

                                        event_emitter.emit(BoxEvent::with_string(
                                            "pool.vm.created",
                                            format!("Replenished VM {}", box_id),
                                        ));
                                    }
                                    Err(error) => {
                                        batch_failed = true;
                                        tracing::warn!(error = %error, "Failed to replenish warm pool");
                                    }
                                }
                            }

                            if batch_failed {
                                replenish_failures = replenish_failures.saturating_add(1);
                                let delay = replenish_backoff_delay(
                                    replenish_failures,
                                    check_interval,
                                );
                                next_replenish_at = Instant::now() + delay;
                                tracing::warn!(
                                    failures = replenish_failures,
                                    retry_in_secs = delay.as_secs(),
                                    "Backing off warm-pool replenishment after boot failure"
                                );
                            } else {
                                replenish_failures = 0;
                                next_replenish_at = Instant::now();
                            }

                            event_emitter.emit(BoxEvent::empty("pool.replenish"));
                        }
                    }
                }
            }
        })
    }

    /// Static version of evict_expired for use in the spawned task.
    async fn evict_expired_static(
        idle: &Arc<Mutex<Vec<WarmVm>>>,
        stats: &Arc<Mutex<PoolStats>>,
        event_emitter: &EventEmitter,
        metrics: Option<&crate::prom::RuntimeMetrics>,
        idle_ttl_secs: u64,
    ) {
        let ttl = std::time::Duration::from_secs(idle_ttl_secs);

        let mut pool = idle.lock().await;
        let mut kept = Vec::new();
        let mut expired = Vec::new();

        for warm_vm in pool.drain(..) {
            if warm_vm.created_at.elapsed() > ttl {
                expired.push(warm_vm);
            } else {
                kept.push(warm_vm);
            }
        }
        *pool = kept;
        let after_count = pool.len();
        drop(pool);

        let evicted_count = expired.len();
        Self::sync_idle_metric(metrics, after_count);
        for warm_vm in expired {
            let mut vm = warm_vm.vm;
            let _ = vm.destroy().await;
        }

        if evicted_count > 0 {
            let mut s = stats.lock().await;
            s.total_evicted += evicted_count as u64;
            s.idle_count = after_count;

            event_emitter.emit(BoxEvent::with_string(
                "pool.vm.evicted",
                format!("Evicted {} expired VMs", evicted_count),
            ));
        }
    }
}

/// Return the number of boot tasks that may be in flight for one batch.
///
/// Configuration validation rejects zero for normal pool construction. The
/// defensive `max(1)` keeps this scheduler total for internal callers and
/// prevents a zero limit from deadlocking the JoinSet loop.
fn bounded_boot_limit(needed: usize, max_concurrent_boots: usize) -> usize {
    needed.min(max_concurrent_boots.max(1))
}

/// Calculate the retry delay after a failed background replenishment batch.
///
/// The regular maintenance tick remains responsible for eviction and scaling,
/// while only replenishment is delayed. Capping the delay keeps a transient
/// provider outage recoverable without allowing a permanently unavailable
/// backend to churn the host indefinitely.
fn replenish_backoff_delay(failures: u32, check_interval: Duration) -> Duration {
    let exponent = failures.saturating_sub(1).min(8);
    let multiplier = 1u64 << exponent;
    let delay_secs = check_interval
        .as_secs()
        .max(1)
        .saturating_mul(multiplier)
        .min(300);
    Duration::from_secs(delay_secs)
}

#[cfg(test)]
mod tests {
    use super::*;
    use a3s_box_core::config::PoolConfig;

    fn test_pool_config(min_idle: usize, max_size: usize) -> PoolConfig {
        PoolConfig {
            enabled: true,
            min_idle,
            max_size,
            idle_ttl_secs: 300,
            ..Default::default()
        }
    }

    fn test_event_emitter() -> EventEmitter {
        EventEmitter::new(100)
    }

    #[test]
    fn boot_or_restore_future_stays_heap_indirected() {
        let config = BoxConfig::default();
        let emitter = test_event_emitter();
        let template = Arc::new(Mutex::new(TemplateState::Unbuilt));
        let future = WarmPool::boot_or_restore(false, &config, &emitter, &template);

        assert!(
            std::mem::size_of_val(&future) <= 2 * std::mem::size_of::<usize>(),
            "boot_or_restore future must remain pointer-sized so pool misses fit on Tokio worker stacks; got {} bytes",
            std::mem::size_of_val(&future)
        );
    }

    #[cfg(not(all(target_os = "linux", target_arch = "x86_64")))]
    #[tokio::test]
    async fn unsupported_snapshot_fork_is_rejected_before_template_construction() {
        let template = Arc::new(Mutex::new(TemplateState::Unbuilt));
        let result =
            WarmPool::ensure_template(&BoxConfig::default(), &test_event_emitter(), &template)
                .await;
        let error = match result {
            Ok(_) => panic!("unsupported host unexpectedly built a snapshot template"),
            Err(error) => error.to_string(),
        };

        assert!(error.contains("Linux x86_64 KVM"), "{error}");
        assert!(matches!(&*template.lock().await, TemplateState::Unbuilt));
    }

    // --- PoolConfig validation tests ---

    #[tokio::test]
    async fn test_pool_rejects_zero_max_size() {
        let config = test_pool_config(0, 0);
        let result = WarmPool::start(config, BoxConfig::default(), test_event_emitter()).await;
        match result {
            Err(e) => assert!(e.to_string().contains("max_size must be greater than 0")),
            Ok(_) => panic!("Expected error for zero max_size"),
        }
    }

    #[tokio::test]
    async fn test_pool_rejects_min_idle_exceeds_max() {
        let config = test_pool_config(10, 5);
        let result = WarmPool::start(config, BoxConfig::default(), test_event_emitter()).await;
        match result {
            Err(e) => assert!(e.to_string().contains("cannot exceed max_size")),
            Ok(_) => panic!("Expected error for min_idle > max_size"),
        }
    }

    #[tokio::test]
    async fn test_pool_rejects_zero_max_concurrent_boots() {
        let mut config = test_pool_config(0, 1);
        config.max_concurrent_boots = 0;
        let result = WarmPool::start(config, BoxConfig::default(), test_event_emitter()).await;
        match result {
            Err(error) => assert!(error.to_string().contains("max_concurrent_boots")),
            Ok(_) => panic!("Expected error for zero max_concurrent_boots"),
        }
    }

    #[tokio::test]
    async fn acquire_boot_permits_releases_both_scopes() {
        let pool_limiter = Arc::new(Semaphore::new(1));
        let global_limiter = Arc::new(Semaphore::new(1));
        let (pool_permit, global_permit) =
            acquire_boot_permits(pool_limiter.clone(), Some(global_limiter.clone()))
                .await
                .expect("both boot limiters should grant a permit");
        assert_eq!(pool_limiter.available_permits(), 0);
        assert_eq!(global_limiter.available_permits(), 0);
        drop((pool_permit, global_permit));
        assert_eq!(pool_limiter.available_permits(), 1);
        assert_eq!(global_limiter.available_permits(), 1);
    }

    #[test]
    fn boot_batch_limit_is_bounded_and_never_deadlocks() {
        assert_eq!(bounded_boot_limit(0, 2), 0);
        assert_eq!(bounded_boot_limit(8, 2), 2);
        assert_eq!(bounded_boot_limit(2, 8), 2);
        assert_eq!(bounded_boot_limit(8, 0), 1);
    }

    #[test]
    fn replenish_backoff_is_exponential_and_capped() {
        let base = Duration::from_secs(5);
        assert_eq!(replenish_backoff_delay(0, base), Duration::from_secs(5));
        assert_eq!(replenish_backoff_delay(1, base), Duration::from_secs(5));
        assert_eq!(replenish_backoff_delay(2, base), Duration::from_secs(10));
        assert_eq!(replenish_backoff_delay(7, base), Duration::from_secs(300));
        assert_eq!(replenish_backoff_delay(20, base), Duration::from_secs(300));
    }

    #[test]
    fn boot_metric_guard_balances_inflight_gauge() {
        let metrics = crate::prom::RuntimeMetrics::new();
        {
            let _guard = BootMetricGuard::new(Some(metrics.clone()));
            assert_eq!(metrics.warm_pool_boots_inflight.get(), 1);
        }
        assert_eq!(metrics.warm_pool_boots_inflight.get(), 0);
    }

    // --- PoolStats tests ---

    #[test]
    fn test_pool_stats_default() {
        let stats = PoolStats {
            idle_count: 0,
            total_created: 0,
            total_acquired: 0,
            total_released: 0,
            total_evicted: 0,
        };
        assert_eq!(stats.idle_count, 0);
        assert_eq!(stats.total_created, 0);
    }

    #[test]
    fn test_pool_stats_clone() {
        let stats = PoolStats {
            idle_count: 3,
            total_created: 10,
            total_acquired: 7,
            total_released: 5,
            total_evicted: 2,
        };
        let cloned = stats.clone();
        assert_eq!(cloned.idle_count, 3);
        assert_eq!(cloned.total_created, 10);
        assert_eq!(cloned.total_acquired, 7);
        assert_eq!(cloned.total_released, 5);
        assert_eq!(cloned.total_evicted, 2);
    }

    #[test]
    fn test_pool_stats_debug() {
        let stats = PoolStats {
            idle_count: 1,
            total_created: 2,
            total_acquired: 3,
            total_released: 4,
            total_evicted: 5,
        };
        let debug = format!("{:?}", stats);
        assert!(debug.contains("idle_count"));
        assert!(debug.contains("total_created"));
    }

    // --- PoolConfig serialization tests ---

    #[test]
    fn test_pool_config_roundtrip() {
        let config = PoolConfig {
            enabled: true,
            min_idle: 3,
            max_size: 10,
            idle_ttl_secs: 600,
            ..Default::default()
        };

        let json = serde_json::to_string(&config).unwrap();
        let parsed: PoolConfig = serde_json::from_str(&json).unwrap();

        assert!(parsed.enabled);
        assert_eq!(parsed.min_idle, 3);
        assert_eq!(parsed.max_size, 10);
        assert_eq!(parsed.idle_ttl_secs, 600);
    }

    #[test]
    fn test_pool_config_default_values() {
        let config = PoolConfig::default();
        assert!(!config.enabled);
        assert_eq!(config.min_idle, 1);
        assert_eq!(config.max_size, 5);
        assert_eq!(config.idle_ttl_secs, 300);
    }

    #[test]
    fn test_pool_config_deserialization_with_defaults() {
        let json = r#"{"enabled": true}"#;
        let config: PoolConfig = serde_json::from_str(json).unwrap();
        assert!(config.enabled);
        assert_eq!(config.min_idle, 1);
        assert_eq!(config.max_size, 5);
        assert_eq!(config.idle_ttl_secs, 300);
    }

    // --- PoolConfig validation edge cases ---

    #[tokio::test]
    async fn test_pool_accepts_min_idle_equals_max() {
        let config = test_pool_config(3, 3);
        // This should be accepted (min_idle == max_size is valid)
        // It will fail at boot (no shim), but config validation should pass
        let result = WarmPool::start(config, BoxConfig::default(), test_event_emitter()).await;
        // The error should be about VM boot, not config validation
        match result {
            Err(e) => assert!(!e.to_string().contains("cannot exceed max_size")),
            Ok(mut pool) => {
                let _ = pool.drain().await;
            }
        }
    }

    #[tokio::test]
    async fn test_pool_accepts_min_idle_zero() {
        let config = test_pool_config(0, 5);
        // min_idle=0 means no pre-warming, should be valid
        let result = WarmPool::start(config, BoxConfig::default(), test_event_emitter()).await;
        match result {
            Ok(mut pool) => {
                // Pool should start with 0 idle VMs
                assert_eq!(pool.idle_count().await, 0);
                let stats = pool.stats().await;
                assert_eq!(stats.idle_count, 0);
                assert_eq!(stats.total_created, 0);
                let _ = pool.drain().await;
            }
            Err(e) => {
                // If it fails, it should NOT be a config validation error
                assert!(!e.to_string().contains("max_size"));
                assert!(!e.to_string().contains("min_idle"));
            }
        }
    }

    // --- WarmPool internal state tests (using min_idle=0 to avoid boot) ---

    #[tokio::test]
    async fn test_pool_stats_initial() {
        let config = test_pool_config(0, 5);
        let result = WarmPool::start(config, BoxConfig::default(), test_event_emitter()).await;
        if let Ok(mut pool) = result {
            let stats = pool.stats().await;
            assert_eq!(stats.idle_count, 0);
            assert_eq!(stats.total_created, 0);
            assert_eq!(stats.total_acquired, 0);
            assert_eq!(stats.total_released, 0);
            assert_eq!(stats.total_evicted, 0);
            let _ = pool.drain().await;
        }
    }

    #[tokio::test]
    async fn test_pool_idle_count_initial() {
        let config = test_pool_config(0, 5);
        let result = WarmPool::start(config, BoxConfig::default(), test_event_emitter()).await;
        if let Ok(mut pool) = result {
            assert_eq!(pool.idle_count().await, 0);
            let _ = pool.drain().await;
        }
    }

    #[tokio::test]
    async fn test_pool_drain_empty_pool() {
        let config = test_pool_config(0, 5);
        let result = WarmPool::start(config, BoxConfig::default(), test_event_emitter()).await;
        if let Ok(mut pool) = result {
            // Draining an empty pool should succeed without error
            let drain_result = pool.drain().await;
            assert!(drain_result.is_ok());

            let stats = pool.stats().await;
            assert_eq!(stats.idle_count, 0);
        }
    }

    #[tokio::test]
    async fn test_pool_drain_emits_event() {
        let emitter = test_event_emitter();
        let mut receiver = emitter.subscribe();
        let config = test_pool_config(0, 5);

        let result = WarmPool::start(config, BoxConfig::default(), emitter).await;
        if let Ok(mut pool) = result {
            pool.drain().await.unwrap();

            // Check that pool.drained event was emitted
            let mut found_drain_event = false;
            // Drain all events from the receiver
            while let Ok(event) = receiver.try_recv() {
                if event.key == "pool.drained" {
                    found_drain_event = true;
                }
            }
            assert!(found_drain_event, "Expected pool.drained event");
        }
    }

    #[tokio::test]
    async fn test_pool_acquire_from_empty_pool_fails_without_shim() {
        let config = test_pool_config(0, 5);
        let result = WarmPool::start(config, BoxConfig::default(), test_event_emitter()).await;
        if let Ok(pool) = result {
            // Acquire from empty pool should try to boot a VM, which will fail
            // because there's no shim binary available in test environment
            let acquire_result = pool.acquire().await;
            assert!(acquire_result.is_err());
        }
    }

    // --- Maintenance loop check interval calculation ---

    #[test]
    #[allow(clippy::unnecessary_min_or_max)]
    fn test_maintenance_check_interval_with_ttl() {
        // TTL = 300s → check every 60s (300/5)
        let interval = if 300_u64 > 0 {
            (300_u64 / 5).max(5)
        } else {
            30
        };
        assert_eq!(interval, 60);
    }

    #[test]
    #[allow(clippy::unnecessary_min_or_max)]
    fn test_maintenance_check_interval_short_ttl() {
        // TTL = 10s → check every 5s (min 5)
        let interval = if 10_u64 > 0 { (10_u64 / 5).max(5) } else { 30 };
        assert_eq!(interval, 5);
    }

    #[test]
    #[allow(clippy::unnecessary_min_or_max)]
    fn test_maintenance_check_interval_very_short_ttl() {
        // TTL = 1s → check every 5s (min 5)
        let interval = if 1_u64 > 0 { (1_u64 / 5).max(5) } else { 30 };
        assert_eq!(interval, 5);
    }

    #[test]
    #[allow(
        clippy::absurd_extreme_comparisons,
        clippy::erasing_op,
        clippy::unnecessary_min_or_max,
        unused_comparisons
    )]
    fn test_maintenance_check_interval_no_ttl() {
        // TTL = 0 → check every 30s
        let interval = if 0_u64 > 0 { (0_u64 / 5).max(5) } else { 30 };
        assert_eq!(interval, 30);
    }

    // --- WarmVm struct tests ---

    #[test]
    fn test_warm_vm_created_at_is_recent() {
        let before = Instant::now();
        let created_at = Instant::now();
        let after = Instant::now();

        assert!(created_at >= before);
        assert!(created_at <= after);
    }

    // --- PoolStats field coverage ---

    #[test]
    fn test_pool_stats_all_fields() {
        let stats = PoolStats {
            idle_count: 10,
            total_created: 100,
            total_acquired: 80,
            total_released: 70,
            total_evicted: 15,
        };

        assert_eq!(stats.idle_count, 10);
        assert_eq!(stats.total_created, 100);
        assert_eq!(stats.total_acquired, 80);
        assert_eq!(stats.total_released, 70);
        assert_eq!(stats.total_evicted, 15);

        // Verify debug output contains all fields
        let debug = format!("{:?}", stats);
        assert!(debug.contains("10"));
        assert!(debug.contains("100"));
        assert!(debug.contains("80"));
        assert!(debug.contains("70"));
        assert!(debug.contains("15"));
    }

    // Note: Full integration tests for acquire/release/drain with actual VMs
    // require a working VM runtime (shim binary + libkrun). These are tested
    // in integration tests with the full box environment. The unit tests here
    // validate configuration, statistics, error handling, and pool lifecycle
    // with min_idle=0 (no VM boot required).

    #[tokio::test]
    async fn test_pool_set_metrics_attaches() {
        let config = test_pool_config(0, 5);
        let result = WarmPool::start(config, BoxConfig::default(), test_event_emitter()).await;
        match result {
            Ok(mut pool) => {
                let metrics = crate::prom::RuntimeMetrics::new();
                pool.set_metrics(metrics.clone());
                assert!(pool.metrics.is_some());
                // Metrics start at zero
                assert_eq!(metrics.warm_pool_hits.get(), 0);
                assert_eq!(metrics.warm_pool_misses.get(), 0);
                assert_eq!(metrics.warm_pool_size.get(), 0);
                let _ = pool.drain().await;
            }
            Err(_) => {
                // Boot failure is acceptable in unit test environment
            }
        }
    }

    #[tokio::test]
    async fn test_pool_start_with_metrics_installs_sink_before_fill() {
        let config = test_pool_config(0, 5);
        let metrics = crate::prom::RuntimeMetrics::new();
        let result = WarmPool::start_with_metrics(
            config,
            BoxConfig::default(),
            test_event_emitter(),
            Some(metrics.clone()),
        )
        .await;

        match result {
            Ok(mut pool) => {
                assert!(pool.metrics.is_some());
                assert_eq!(metrics.warm_pool_capacity.get(), 5);
                assert_eq!(
                    metrics.warm_pool_initial_fill_duration.get_sample_count(),
                    1
                );
                let _ = pool.drain().await;
            }
            Err(_) => {
                // Boot failure is acceptable in unit test environments without
                // a usable VM provider; min_idle=0 normally avoids this path.
            }
        }
    }
}