boxlite 0.10.1

//! Box initialization orchestration.
//!
//! ## Architecture
//!
//! Initialization is table-driven with different execution plans based on BoxStatus:
//!
//! ```text
//! Starting (new box):
//!   1. Filesystem           (create layout)
//!   2. BootAssets       ─┬─  (stage custom kernel/initramfs)
//!      ContainerRootfs  ├─  (pull image, create COW disk)
//!      GuestRootfs      ─┘  (prepare guest, create COW disk)
//!   3. VmmSpawn             (build config + spawn VM)
//!   4. GuestConnect         (wait for guest ready)
//!   5. GuestInit       ─┬─   (initialize container)
//!      PortPublish     ─┘    (publish host ports)
//!
//! Stopped (restart):
//!   1. Filesystem           (load existing layout)
//!   2. BootAssets       ─┬─  (reuse box-scoped boot assets)
//!      ContainerRootfs  ├─  (reuse existing COW disk - preserves user data)
//!      GuestRootfs      ─┘  (reuse existing COW disk)
//!   3. VmmSpawn             (build config + spawn NEW VM)
//!   4. GuestConnect         (wait for guest ready)
//!   5. GuestInit       ─┬─   (re-initialize container in new VM)
//!      PortPublish     ─┘    (republish host ports)
//!
//! Running (reattach):
//!   1. VmmAttach            (attach to running VM)
//!   2. GuestConnect         (reconnect to guest)
//!   3. PortPublish          (reconcile publication after a core crash)
//! ```
//!
//! `CleanupGuard` provides RAII cleanup on failure.

mod tasks;
mod types;

pub(crate) use crate::litebox::box_impl::LiveState;

use crate::litebox::BoxStatus;
use crate::litebox::config::BoxConfig;
use crate::metrics::BoxMetricsStorage;
use crate::pipeline::{
    BoxedTask, ExecutionPlan, PipelineBuilder, PipelineExecutor, PipelineMetrics, Stage,
};
use crate::runtime::rt_impl::SharedRuntimeImpl;
use crate::runtime::types::BoxState;
use boxlite_shared::errors::{BoxliteError, BoxliteResult};
use std::sync::Arc;
use tokio::sync::Mutex;

pub(crate) use tasks::PortPublishTask;
use tasks::{
    BootAssetsTask, ContainerRootfsTask, FilesystemTask, GuestConnectTask, GuestInitTask,
    GuestRootfsTask, InitCtx, VmmAttachTask, VmmSpawnTask,
};
use types::InitPipelineContext;

// ============================================================================
// EXECUTION PLAN
// ============================================================================

/// Get execution plan based on BoxStatus.
///
/// Returns Err for states that aren't valid init entry points
/// (e.g. Unknown from a corrupted DB row, or future variants). Callers
/// should treat this as a user-facing precondition error.
fn get_execution_plan(status: BoxStatus) -> BoxliteResult<ExecutionPlan<InitCtx>> {
    let stages: Vec<Stage<BoxedTask<InitCtx>>> = match status {
        BoxStatus::Configured => vec![
            // First start: Full pipeline
            // Phase 1: Setup filesystem layout first
            Stage::sequential(vec![Box::new(FilesystemTask)]),
            // Phase 2: Prepare independent boot and rootfs inputs after the
            // filesystem task has established their box-scoped destinations.
            Stage::parallel(vec![
                Box::new(BootAssetsTask),
                Box::new(ContainerRootfsTask),
                Box::new(GuestRootfsTask),
            ]),
            // Phase 3: Build config and spawn VM
            Stage::sequential(vec![Box::new(VmmSpawnTask)]),
            Stage::sequential(vec![Box::new(GuestConnectTask)]),
            Stage::parallel(vec![Box::new(GuestInitTask), Box::new(PortPublishTask)]),
        ],
        // Stopped and Failed both run the restart pipeline. A Failed box
        // has its rootfs preserved (per BoxStatus::Failed doc) and is
        // retryable per BoxStatus::can_start.
        BoxStatus::Stopped | BoxStatus::Failed => vec![
            // Restart: Same flow but rootfs tasks reuse existing COW disks
            // (preserves user modifications from previous run)
            Stage::sequential(vec![Box::new(FilesystemTask)]),
            Stage::parallel(vec![
                Box::new(BootAssetsTask),
                Box::new(ContainerRootfsTask),
                Box::new(GuestRootfsTask),
            ]),
            Stage::sequential(vec![Box::new(VmmSpawnTask)]),
            Stage::sequential(vec![Box::new(GuestConnectTask)]),
            // GuestInit must run - new VM process has fresh guest daemon
            Stage::parallel(vec![Box::new(GuestInitTask), Box::new(PortPublishTask)]),
        ],
        BoxStatus::Running => vec![
            // Reattach: vmm_attach gates on ProcessIdentity AND surfaces
            // any crash via exit_file, then connect to guest and reconcile
            // any publication interrupted by a prior core-process crash.
            Stage::sequential(vec![Box::new(VmmAttachTask)]),
            Stage::sequential(vec![Box::new(GuestConnectTask)]),
            Stage::sequential(vec![Box::new(PortPublishTask)]),
        ],
        other => {
            return Err(BoxliteError::InvalidState(format!(
                "Cannot initialize box in {other} state"
            )));
        }
    };

    Ok(ExecutionPlan::new(stages))
}

fn box_metrics_from_pipeline(pipeline_metrics: &PipelineMetrics) -> BoxMetricsStorage {
    let mut metrics = BoxMetricsStorage::new();

    if let Some(duration_ms) = pipeline_metrics.task_duration_ms("filesystem_setup") {
        metrics.set_stage_filesystem_setup(duration_ms);
    }
    if let Some(duration_ms) = pipeline_metrics.task_duration_ms("container_rootfs_prep") {
        metrics.set_stage_image_prepare(duration_ms);
    }
    if let Some(duration_ms) = pipeline_metrics.task_duration_ms("guest_rootfs_init") {
        metrics.set_stage_guest_rootfs(duration_ms);
    }
    if let Some(duration_ms) = pipeline_metrics.task_duration_ms("vmm_spawn") {
        metrics.set_stage_box_spawn(duration_ms);
    }
    if let Some(duration_ms) = pipeline_metrics.task_duration_ms("vmm_attach") {
        metrics.set_stage_box_spawn(duration_ms);
    }
    if let Some(_duration_ms) = pipeline_metrics.task_duration_ms("guest_connect") {
        // Track guest connection time
        // Could add a new metric field if needed
    }
    if let Some(duration_ms) = pipeline_metrics.task_duration_ms("guest_init") {
        metrics.set_stage_container_init(duration_ms);
    }

    metrics
}

/// Builds and initializes box components.
///
/// # Example
///
/// ```ignore
/// let inner = BoxBuilder::new(runtime, config, &state)
///     .build()
///     .await?;
/// ```
pub(crate) struct BoxBuilder {
    runtime: SharedRuntimeImpl,
    config: BoxConfig,
    state: BoxState,
}

fn validate_persisted_options(
    features: &crate::experimental::ExperimentalFeatures,
    options: &crate::BoxOptions,
) -> BoxliteResult<()> {
    features.require_for_options(options)?;
    options.sanitize_persisted()
}

impl BoxBuilder {
    /// Create a new builder from config and state.
    ///
    /// The state determines initialization mode:
    /// - `Starting`: normal init (pull image or use rootfs path)
    /// - `Stopped`: restart (reuse existing rootfs at box_home/rootfs)
    ///
    /// # Arguments
    ///
    /// * `runtime` - Runtime providing resources (layout, guest_rootfs, etc.)
    /// * `config` - Box configuration (immutable after creation)
    /// * `state` - Current box state (determines init mode)
    pub(crate) fn new(
        runtime: SharedRuntimeImpl,
        config: BoxConfig,
        state: BoxState,
    ) -> BoxliteResult<Self> {
        // External source paths were validated at create time. Persisted boxes
        // validate only their stored shape here; the boot-assets task reopens a
        // source only when no complete box-scoped generation exists.
        let options = &config.options;
        validate_persisted_options(&runtime.experimental_features, options)?;

        Ok(Self {
            runtime,
            config,
            state,
        })
    }

    /// Build and initialize LiveState.
    ///
    /// Executes all initialization stages with automatic cleanup on failure.
    /// Returns (LiveState, CleanupGuard) - caller must disarm guard after all
    /// operations succeed (including DB persist).
    pub(crate) async fn build(self) -> BoxliteResult<(LiveState, types::CleanupGuard)> {
        use std::time::Instant;

        let total_start = Instant::now();

        let BoxBuilder {
            runtime,
            config,
            state,
        } = self;

        let status = state.status;
        let reuse_rootfs = status == BoxStatus::Stopped;
        let skip_guest_wait = status == BoxStatus::Running;

        let ctx = InitPipelineContext::new(config, runtime.clone(), reuse_rootfs, skip_guest_wait);
        let ctx = Arc::new(Mutex::new(ctx));
        let ctx_for_cleanup = Arc::clone(&ctx);

        // Note: Guard stays armed until caller disarms it after DB persist succeeds.
        // On any error path inside `inner`, we capture the error onto the guard
        // so `CleanupGuard::drop` can persist it as `BoxState::Failed` with
        // `error_reason`. This matches Daytona/Kata/containerd/Docker semantics
        // of "preserve the record on init failure".
        let inner = async move {
            let plan = get_execution_plan(status)?;
            let pipeline = PipelineBuilder::from_plan(plan);
            let pipeline_metrics = PipelineExecutor::execute(pipeline, Arc::clone(&ctx)).await?;

            let mut ctx = ctx.lock().await;
            let total_create_duration_ms = total_start.elapsed().as_millis();
            let handler = ctx
                .guard
                .take_handler()
                .ok_or_else(|| BoxliteError::Internal("handler was not set".into()))?;

            let mut metrics = box_metrics_from_pipeline(&pipeline_metrics);
            metrics.set_total_create_duration(total_create_duration_ms);

            metrics.log_init_stages();

            // Note: Guard is NOT disarmed here. Caller is responsible for disarming
            // after all operations succeed (including DB persist).

            // Get guest_session from GuestConnectTask
            let guest_session = ctx.guest_session.take().ok_or_else(|| {
                BoxliteError::Internal("guest_connect task must run first".into())
            })?;

            // Get disks from context (for Running, create disk reference directly)
            let (container_disk, guest_disk) = if status == BoxStatus::Running {
                // Reattach: create disk reference to existing qcow2
                use crate::disk::DiskFormat;
                use crate::disk::constants::filenames;
                let disk = crate::disk::Disk::new(
                    ctx.config.box_home.join(filenames::CONTAINER_DISK),
                    DiskFormat::Qcow2,
                    true,
                );
                (disk, None)
            } else {
                // Starting/Stopped: get disks from rootfs tasks
                let container_disk = ctx
                    .container_disk
                    .take()
                    .ok_or_else(|| BoxliteError::Internal("rootfs task must run first".into()))?;
                (container_disk, ctx.guest_disk.take())
            };

            #[cfg(target_os = "linux")]
            let bind_mount = ctx.bind_mount.take();

            // Take the guard out of context, replacing with a disarmed placeholder.
            // The caller is responsible for disarming the returned guard after all
            // operations succeed (including DB persist).
            let mut placeholder =
                types::CleanupGuard::new(ctx.runtime.clone(), ctx.config.id.clone());
            placeholder.disarm();
            let guard = std::mem::replace(&mut ctx.guard, placeholder);

            // Own the box's one network backend beside guest_session — created by
            // vmm_spawn (which also used it to produce the wire spec) or, on the
            // reattach path, by vmm_attach; both thread it here for runtime control.
            let network = ctx.network_backend.take();
            let published_ports = ctx.published_ports.take();

            // Build LiveState
            let live_state = LiveState::new(
                handler,
                guest_session,
                network,
                published_ports,
                metrics,
                container_disk,
                guest_disk,
                #[cfg(target_os = "linux")]
                bind_mount,
            );

            Ok::<(LiveState, types::CleanupGuard), BoxliteError>((live_state, guard))
        };

        match inner.await {
            Ok(parts) => Ok(parts),
            Err(e) => {
                // Capture the cause onto the still-armed guard inside ctx so
                // its Drop can mark the box `Failed` with this reason.
                ctx_for_cleanup.lock().await.guard.set_last_error(&e);
                Err(e)
            }
        }
    }
}

#[cfg(test)]
mod plan_tests {
    use super::*;
    use crate::experimental::ExperimentalFeatures;
    use crate::experimental::custom_kernel::{KernelFormat, KernelOptions};
    use crate::pipeline::ExecutionMode;

    fn plan_shape(status: BoxStatus) -> Vec<(ExecutionMode, Vec<String>)> {
        get_execution_plan(status)
            .unwrap()
            .stages()
            .into_iter()
            .map(|stage| {
                (
                    stage.execution,
                    stage
                        .tasks
                        .into_iter()
                        .map(|task| task.name().to_string())
                        .collect(),
                )
            })
            .collect()
    }

    fn boot_pipeline_shape() -> Vec<(ExecutionMode, Vec<String>)> {
        vec![
            (
                ExecutionMode::Sequential,
                vec!["filesystem_setup".to_string()],
            ),
            (
                ExecutionMode::Parallel,
                vec![
                    "boot_assets_prepare".to_string(),
                    "container_rootfs_prep".to_string(),
                    "guest_rootfs_init".to_string(),
                ],
            ),
            (ExecutionMode::Sequential, vec!["vmm_spawn".to_string()]),
            (ExecutionMode::Sequential, vec!["guest_connect".to_string()]),
            // Publication needs a live gvproxy, which the shim binds before
            // it boots the VM; it has no reason to gate container init.
            (
                ExecutionMode::Parallel,
                vec!["guest_init".to_string(), "port_publish".to_string()],
            ),
        ]
    }

    #[test]
    fn configured_box_prepares_boot_assets_before_vmm_spawn() {
        assert_eq!(plan_shape(BoxStatus::Configured), boot_pipeline_shape());
    }

    #[test]
    fn restartable_boxes_prepare_boot_assets_before_vmm_spawn() {
        let expected = boot_pipeline_shape();
        assert_eq!(plan_shape(BoxStatus::Stopped), expected);
        assert_eq!(plan_shape(BoxStatus::Failed), boot_pipeline_shape());
    }

    #[test]
    fn running_box_reattach_does_not_prepare_boot_assets() {
        assert_eq!(
            plan_shape(BoxStatus::Running),
            vec![
                (ExecutionMode::Sequential, vec!["vmm_attach".to_string()]),
                (ExecutionMode::Sequential, vec!["guest_connect".to_string()]),
                // Reattach has no guest_init to overlap with.
                (ExecutionMode::Sequential, vec!["port_publish".to_string()]),
            ]
        );
    }

    #[test]
    #[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
    fn persisted_custom_kernel_uses_injected_feature_state() {
        #[cfg(target_arch = "x86_64")]
        let format = KernelFormat::Elf;
        #[cfg(target_arch = "aarch64")]
        let format = KernelFormat::PeGz;

        let mut options = crate::BoxOptions::default();
        options.advanced.kernel =
            Some(KernelOptions::new("/source/no-longer-needed").with_format(format));

        let error = validate_persisted_options(&ExperimentalFeatures::default(), &options)
            .expect_err("persisted custom kernel must be disabled by default");
        assert!(
            error
                .to_string()
                .contains("ExperimentalFeature::CustomKernel")
        );

        let enabled = ExperimentalFeatures::parse("custom-kernel").unwrap();
        validate_persisted_options(&enabled, &options).unwrap();
    }

    #[test]
    fn persisted_nested_virtualization_uses_injected_feature_state() {
        let mut advanced = crate::runtime::advanced_options::AdvancedBoxOptions::default();
        advanced.nested_virtualization = true;
        let options = crate::BoxOptions {
            advanced,
            ..Default::default()
        };

        let error = validate_persisted_options(&ExperimentalFeatures::default(), &options)
            .expect_err("persisted nested virtualization must be disabled by default");
        assert!(
            error
                .to_string()
                .contains("ExperimentalFeature::NestedVirtualization")
        );

        let enabled = ExperimentalFeatures::parse("nested-virtualization").unwrap();
        validate_persisted_options(&enabled, &options).unwrap();
    }
}