blinkvm-sdk 0.3.9

//! Shared BoxLite runtime for long-lived server processes.

use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;

use anyhow::{Context, Result, bail};
use boxlite::{
    BoxCommand, BoxliteRuntime, LiteBox, SnapshotInfo,
    runtime::options::{
        BoxArchive, BoxOptions, ExportOptions, NetworkSpec, RootfsSpec,
        SnapshotOptions, VolumeSpec,
    },
};
use tokio::sync::Semaphore;
use tracing::{info, warn};

use crate::boxlite_options::load_boxlite_options;
use crate::exec::exec_agent_script;
use crate::network::resolve_network_spec;
use crate::runner::run_agent_script;
use crate::AgentResult;
use blink_shared::AGENT_MEMORY_DIR;

const DEFAULT_VM_CONCURRENCY: usize = 3;
const DEFAULT_CONTROL_TIMEOUT_SECS: u64 = 120;
const DEFAULT_RUN_TIMEOUT_SECS: u64 = 300;

#[derive(Clone, Debug, serde::Deserialize)]
pub struct SessionVolume {
    pub host_path: String,
    pub guest_path: String,
    #[serde(default)]
    pub read_only: bool,
}

impl From<SessionVolume> for VolumeSpec {
    fn from(volume: SessionVolume) -> Self {
        Self {
            host_path: volume.host_path,
            guest_path: volume.guest_path,
            read_only: volume.read_only,
        }
    }
}

/// VM resource limits passed through to BoxLite `BoxOptions`.
///
/// Unset fields use BoxLite defaults (4 CPUs, 4096 MiB RAM, 10 GB disk).
#[derive(Clone, Debug, Default, serde::Deserialize, serde::Serialize)]
pub struct SandboxResources {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cpus: Option<u8>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub memory_mib: Option<u32>,
    /// Sparse container rootfs virtual size in GB (at least the base image size).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub disk_size_gb: Option<u64>,
}

impl SandboxResources {
    pub fn validate(&self) -> Result<()> {
        if let Some(0) = self.cpus {
            bail!("cpus must be >= 1");
        }
        if let Some(0) = self.memory_mib {
            bail!("memory_mib must be >= 1");
        }
        if let Some(0) = self.disk_size_gb {
            bail!("disk_size_gb must be >= 1");
        }
        Ok(())
    }
}

/// Options for creating or reusing a named session.
#[derive(Clone, Debug, Default)]
pub struct OpenSessionOptions {
    pub volumes: Vec<SessionVolume>,
    pub network: Option<boxlite::runtime::options::NetworkConfig>,
    pub resources: SandboxResources,
}

#[derive(Clone)]
pub struct BlinkContext {
    runtime: BoxliteRuntime,
    export_dir: PathBuf,
    /// Limits concurrent VM operations so worker threads are not exhausted,
    /// keeping HTTP handlers (including /health) responsive under load.
    vm_concurrency: Arc<Semaphore>,
    /// Timeout for lightweight VM control ops (open/stop/remove/spawn).
    control_timeout: Duration,
    /// Timeout for potentially long ops (run/checkpoint/restore/export/import).
    run_timeout: Duration,
}

#[derive(Clone, Debug, serde::Serialize)]
pub struct SessionInfo {
    pub name: Option<String>,
    pub box_id: String,
    pub status: String,
    pub running: bool,
}

impl BlinkContext {
    pub fn new() -> Result<Self> {
        let home = dirs::home_dir().context("could not resolve home directory")?;
        let export_dir = home.join(".blink").join("exports");
        std::fs::create_dir_all(&export_dir).context("failed to create export directory")?;
        Ok(Self {
            runtime: BoxliteRuntime::new(load_boxlite_options()?)
                .context("failed to initialize BoxLite runtime")?,
            export_dir,
            vm_concurrency: Arc::new(Semaphore::new(vm_concurrency_limit())),
            control_timeout: duration_from_env("BLINK_VM_TIMEOUT_SECS", DEFAULT_CONTROL_TIMEOUT_SECS),
            run_timeout: duration_from_env("BLINK_RUN_TIMEOUT_SECS", DEFAULT_RUN_TIMEOUT_SECS),
        })
    }

    /// Acquire a VM concurrency permit. Prevents all tokio worker threads from
    /// being consumed by simultaneous VM operations, which would deadlock the
    /// runtime and make even /health unresponsive.
    async fn acquire_permit(&self) -> Result<tokio::sync::SemaphorePermit<'_>> {
        self.vm_concurrency
            .acquire()
            .await
            .map_err(|_| anyhow::anyhow!("VM concurrency semaphore closed"))
    }

    pub fn export_dir(&self) -> &Path {
        &self.export_dir
    }

    fn session_options(
        image: &str,
        volumes: &[VolumeSpec],
        network: NetworkSpec,
        resources: &SandboxResources,
    ) -> BoxOptions {
        let working_dir = volumes
            .first()
            .map(|volume| volume.guest_path.clone());
        BoxOptions {
            rootfs: RootfsSpec::Image(image.to_string()),
            network,
            auto_remove: false,
            detach: true,
            volumes: volumes.to_vec(),
            working_dir,
            cpus: resources.cpus,
            memory_mib: resources.memory_mib,
            disk_size_gb: resources.disk_size_gb,
            ..Default::default()
        }
    }

    async fn get_box(&self, name: &str) -> Result<LiteBox> {
        self.runtime
            .get(name)
            .await
            .context("failed to lookup session")?
            .with_context(|| format!("session '{name}' not found"))
    }

    pub async fn list_sessions(&self) -> Result<Vec<SessionInfo>> {
        let timeout = self.control_timeout;
        let infos = tokio::time::timeout(timeout, self.runtime.list_info())
            .await
            .map_err(|_| {
                anyhow::anyhow!("list sessions timed out after {}s", timeout.as_secs())
            })?
            .context("list sessions")?;
        Ok(infos
            .into_iter()
            .map(|b| SessionInfo {
                name: b.name,
                box_id: b.id.to_string(),
                status: format!("{:?}", b.status),
                running: b.pid.is_some(),
            })
            .collect())
    }

    pub async fn open_session(
        &self,
        name: &str,
        image: &str,
        warm: bool,
        options: OpenSessionOptions,
    ) -> Result<(String, bool)> {
        let _permit = self.acquire_permit().await?;
        let timeout = self.control_timeout;
        tokio::time::timeout(timeout, async move {
            options.resources.validate()?;
            info!(
                name,
                image,
                warm,
                volume_count = options.volumes.len(),
                cpus = ?options.resources.cpus,
                memory_mib = ?options.resources.memory_mib,
                disk_size_gb = ?options.resources.disk_size_gb,
                "opening session"
            );
            let network_spec = resolve_network_spec(options.network)?;
            let volume_specs: Vec<VolumeSpec> =
                options.volumes.into_iter().map(VolumeSpec::from).collect();
            let (litebox, created) = self
                .runtime
                .get_or_create(
                    Self::session_options(image, &volume_specs, network_spec, &options.resources),
                    Some(name.to_string()),
                )
                .await
                .context("open session")?;
            ensure_memory_dir(&litebox).await?;
            Ok((litebox.id().as_str().to_string(), created))
        })
        .await
        .map_err(|_| anyhow::anyhow!("open session timed out after {}s", timeout.as_secs()))?
    }

    pub async fn run_agent_ephemeral(
        &self,
        script_path: &Path,
        image: Option<&str>,
        resources: SandboxResources,
    ) -> Result<AgentResult> {
        let _permit = self.acquire_permit().await?;
        let timeout = self.run_timeout;
        let image = image.unwrap_or(blink_shared::DEFAULT_ROOTFS_IMAGE);
        tokio::time::timeout(timeout, run_agent_script(script_path, image, resources))
            .await
            .map_err(|_| {
                anyhow::anyhow!("ephemeral run timed out after {}s", timeout.as_secs())
            })?
    }

    pub async fn run_in_session(&self, name: &str, script_path: &Path) -> Result<AgentResult> {
        let _permit = self.acquire_permit().await?;
        let timeout = self.run_timeout;
        tokio::time::timeout(timeout, async {
            let litebox = self.get_box(name).await?;
            exec_agent_script(&litebox, script_path).await
        })
        .await
        .map_err(|_| anyhow::anyhow!("run in session timed out after {}s", timeout.as_secs()))?
    }

    pub async fn spawn_in_session(&self, name: &str, spec: crate::SpawnSpec) -> Result<crate::Execution> {
        let _permit = self.acquire_permit().await?;
        let timeout = self.control_timeout;
        tokio::time::timeout(timeout, async move {
            let litebox = self.get_box(name).await?;
            crate::spawn_exec(&litebox, spec).await
        })
        .await
        .map_err(|_| anyhow::anyhow!("spawn in session timed out after {}s", timeout.as_secs()))?
    }

    pub async fn checkpoint_session(&self, name: &str, snapshot: &str) -> Result<SnapshotInfo> {
        let _permit = self.acquire_permit().await?;
        let timeout = self.run_timeout;
        tokio::time::timeout(timeout, async {
            let litebox = self.get_box(name).await?;
            litebox
                .snapshots()
                .create(SnapshotOptions::default(), snapshot)
                .await
                .context("checkpoint")
        })
        .await
        .map_err(|_| anyhow::anyhow!("checkpoint session timed out after {}s", timeout.as_secs()))?
    }

    pub async fn restore_session(&self, name: &str, snapshot: &str) -> Result<()> {
        let _permit = self.acquire_permit().await?;
        let timeout = self.run_timeout;
        tokio::time::timeout(timeout, async {
            let litebox = self.get_box(name).await?;
            if litebox.info().status.is_active() {
                litebox.stop().await.context("stop before restore")?;
            }
            litebox.snapshots().restore(snapshot).await.context("restore")
        })
        .await
        .map_err(|_| anyhow::anyhow!("restore session timed out after {}s", timeout.as_secs()))?
    }

    pub async fn list_checkpoints(&self, name: &str) -> Result<Vec<SnapshotInfo>> {
        let timeout = self.control_timeout;
        tokio::time::timeout(timeout, async {
            let litebox = self.get_box(name).await?;
            litebox
                .snapshots()
                .list()
                .await
                .context("list checkpoints")
        })
        .await
        .map_err(|_| {
            anyhow::anyhow!("list checkpoints timed out after {}s", timeout.as_secs())
        })?
    }

    pub async fn stop_session(&self, name: &str) -> Result<()> {
        let _permit = self.acquire_permit().await?;
        let timeout = self.control_timeout;
        tokio::time::timeout(timeout, async {
            let litebox = self.get_box(name).await?;
            if litebox.info().status.is_active() {
                litebox.stop().await.context("stop session")?;
            }
            Ok(())
        })
        .await
        .map_err(|_| anyhow::anyhow!("stop session timed out after {}s", timeout.as_secs()))?
    }

    pub async fn remove_session(&self, name: &str, force: bool) -> Result<()> {
        let _permit = self.acquire_permit().await?;
        let timeout = self.control_timeout;
        tokio::time::timeout(timeout, async {
            // Stop the box first to ensure proper child process reaping
            // (boxlite-shim + libkrun VM). Without this, removing a running
            // box leaves zombie subprocesses that accumulate over time and
            // can deadlock the server.
            match self.get_box(name).await {
                Ok(litebox) => {
                    if litebox.info().status.is_active() {
                        match litebox.stop().await {
                            Ok(()) => {}
                            Err(err) if !force => {
                                return Err(err).context("stop session before remove");
                            }
                            Err(err) => {
                                warn!(
                                    name,
                                    error = %err,
                                    "stop before remove failed; continuing with force"
                                );
                            }
                        }
                    }
                }
                Err(_) => {
                    // Box not found or lookup failed; proceed with
                    // runtime.remove() which can clean up partial/stale state.
                }
            }
            self.runtime.remove(name, force).await.context("remove session")
        })
        .await
        .map_err(|_| anyhow::anyhow!("remove session timed out after {}s", timeout.as_secs()))?
    }

    pub async fn export_session(&self, name: &str) -> Result<PathBuf> {
        let _permit = self.acquire_permit().await?;
        let timeout = self.run_timeout;
        tokio::time::timeout(timeout, async {
            let litebox = self.get_box(name).await?;
            let stamp = chrono::Utc::now().format("%Y%m%d-%H%M%S");
            let dest = self
                .export_dir
                .join(format!("{name}-{stamp}.boxlite"));
            let archive = litebox
                .export(ExportOptions::default(), &dest)
                .await
                .context("export session")?;
            Ok(archive.path().to_path_buf())
        })
        .await
        .map_err(|_| anyhow::anyhow!("export session timed out after {}s", timeout.as_secs()))?
    }

    pub async fn import_session(&self, archive_path: &Path, name: Option<&str>) -> Result<String> {
        if !archive_path.exists() {
            bail!("archive not found: {}", archive_path.display());
        }
        let _permit = self.acquire_permit().await?;
        let timeout = self.run_timeout;
        let archive = BoxArchive::new(archive_path);
        let name_owned = name.map(String::from);
        tokio::time::timeout(timeout, async move {
            let litebox = self
                .runtime
                .import_box(archive, name_owned)
                .await
                .context("import session")?;
            ensure_memory_dir(&litebox).await?;
            Ok(litebox.id().as_str().to_string())
        })
        .await
        .map_err(|_| anyhow::anyhow!("import session timed out after {}s", timeout.as_secs()))?
    }
}

async fn ensure_memory_dir(litebox: &LiteBox) -> Result<()> {
    let cmd = format!("mkdir -p {AGENT_MEMORY_DIR}");
    for attempt in 0..3 {
        let execution = match litebox
            .exec(BoxCommand::new("sh").arg("-c").arg(&cmd))
            .await
        {
            Ok(e) => e,
            Err(e) => {
                if attempt < 2 {
                    tokio::time::sleep(std::time::Duration::from_millis(500)).await;
                    continue;
                }
                return Err(e).context("mkdir memory dir");
            }
        };
        let status = execution.wait().await.context("mkdir wait")?;
        if status.exit_code == 0 {
            return Ok(());
        }
        if attempt < 2 {
            tokio::time::sleep(std::time::Duration::from_millis(500)).await;
        }
    }
    warn!(exit_code = "unknown", "mkdir {AGENT_MEMORY_DIR} failed after 3 attempts, agent may need to create it");
    Ok(())
}

fn vm_concurrency_limit() -> usize {
    std::env::var("BLINK_VM_CONCURRENCY")
        .ok()
        .and_then(|v| v.trim().parse::<usize>().ok())
        .filter(|v| *v > 0)
        .unwrap_or(DEFAULT_VM_CONCURRENCY)
}

fn duration_from_env(var: &str, default_secs: u64) -> Duration {
    std::env::var(var)
        .ok()
        .and_then(|v| v.trim().parse::<u64>().ok())
        .filter(|v| *v > 0)
        .map(Duration::from_secs)
        .unwrap_or_else(|| Duration::from_secs(default_secs))
}