supermachine 0.4.13

//! Typed VM configuration resources.
//!
//! This is the first reusable config boundary lifted out of the supermachine-worker
//! command-line harness. Keep this free of CLI parsing and host process state
//! so it can grow into the general Supermachine VM-library input surface.

use std::fmt;

pub const DEFAULT_CMDLINE: &str = "console=ttyAMA0 reboot=t panic=-1";
pub const DEFAULT_MEMORY_MIB: usize = 256;
pub const DEFAULT_VCPUS: u32 = 1;
pub const THROUGHPUT_PROFILE_VCPUS: u32 = 4;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum VmProfile {
    Latency,
    Throughput,
}

impl VmProfile {
    pub fn parse(s: &str) -> Option<Self> {
        match s {
            "latency" | "low-latency" => Some(Self::Latency),
            "throughput" => Some(Self::Throughput),
            _ => None,
        }
    }

    pub fn default_vcpus(self) -> u32 {
        match self {
            Self::Latency => DEFAULT_VCPUS,
            Self::Throughput => THROUGHPUT_PROFILE_VCPUS,
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct VmResources {
    pub kernel_path: Option<String>,
    pub initrd_path: Option<String>,
    pub cmdline: String,
    pub memory_mib: usize,
    /// Read-only block devices (squashfs layers, delta squashfs).
    /// Attached as virtio-blk before any volumes.
    pub block_devices: Vec<String>,
    /// Read-write block devices for `--volume` persistent volumes.
    /// Attached as virtio-blk after `block_devices`. Each is a host
    /// file the guest mounts as a writable filesystem.
    pub volumes: Vec<VolumeSpec>,
    pub vcpus: u32,
    pub restore_from: Option<String>,
    pub cow_restore: bool,
    pub snapshot: SnapshotResources,
    pub endpoints: EndpointResources,
    /// If `Some(N)`, the runner asks the guest's virtio-balloon
    /// driver to inflate by `N` 4 KiB pages after restore (or
    /// after kernel boot in the no-restore path). The host then
    /// reclaims those pages via `madvise(MADV_FREE)`, dropping
    /// per-worker idle RSS. Driven by metadata.json's
    /// `balloon_target_pages` field on the snapshot, plumbed
    /// through `--balloon-target-pages N` on the worker CLI.
    pub balloon_target_pages: Option<u32>,
}

/// A `--volume HOST_FILE:GUEST_PATH` entry. The host file is the
/// canonical store; on first use init-oci formats it ext4. Once
/// formatted, the same file persists across runs and across
/// snapshot restores. Snapshots don't capture the volume contents
/// (they live on the host); they only capture the mapping.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct VolumeSpec {
    /// Path to the host file backing this volume. Created (sparse,
    /// `size_bytes`) if missing.
    pub host_path: String,
    /// Mount point inside the guest (e.g. `/var/lib/postgres`).
    pub guest_path: String,
    /// Size of the volume in bytes. Default 1 GiB. Files are
    /// sparse — actual disk usage matches what the guest writes.
    pub size_bytes: u64,
}

impl VolumeSpec {
    pub const DEFAULT_SIZE_BYTES: u64 = 1024 * 1024 * 1024;

    pub fn new(host_path: impl Into<String>, guest_path: impl Into<String>) -> Self {
        Self {
            host_path: host_path.into(),
            guest_path: guest_path.into(),
            size_bytes: Self::DEFAULT_SIZE_BYTES,
        }
    }

    pub fn with_size_bytes(mut self, size_bytes: u64) -> Self {
        self.size_bytes = size_bytes;
        self
    }
}

#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct SnapshotResources {
    pub after_ms: Option<u64>,
    pub at_heartbeat: Option<u64>,
    pub on_listener: bool,
    pub quiesce_ms: u64,
    pub out_path: Option<String>,
}

#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct EndpointResources {
    pub vsock_mux: Option<String>,
    pub http_port: Option<String>,
    /// SCM_RIGHTS handoff acceptor path. The router connects here
    /// and passes accepted client TCP fds (plus a small prefix of
    /// already-consumed bytes) so the worker bridges directly to
    /// the guest without the router process being on the data path.
    pub vsock_mux_handoff: Option<String>,
    /// `<vsock_mux>-exec.sock` path. Each accepted unix-socket
    /// client gets bridged to a *native* AF_VSOCK connection to
    /// the guest's exec agent on `vsock_exec_guest_port` (default
    /// 1028 — see [`DEFAULT_EXEC_GUEST_PORT`]). Used by
    /// `supermachine exec` and `Vm::exec`.
    pub vsock_exec: Option<String>,
    /// Native AF_VSOCK port the guest agent listens on. None ⇒
    /// [`DEFAULT_EXEC_GUEST_PORT`].
    pub vsock_exec_guest_port: Option<u32>,
}

/// Default native AF_VSOCK port for the in-guest exec agent. Picked
/// to sit above the host-direction reserved range (env service is
/// `VSOCK_ENV_PORT=1026`, leaving room for future host-side
/// listeners) and below typical TSI vm_port allocations.
pub const DEFAULT_EXEC_GUEST_PORT: u32 = 1028;

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ResourceError {
    MissingKernel,
    ZeroMemory,
    ZeroVcpus,
    SnapshotTriggerWithoutOutput,
}

impl VmResources {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn for_kernel(kernel_path: impl Into<String>, initrd_path: impl Into<String>) -> Self {
        Self::new()
            .with_kernel_path(kernel_path)
            .with_initramfs(initrd_path)
    }

    pub fn from_snapshot(path: impl Into<String>) -> Self {
        Self::new().with_restore(path)
    }

    pub fn with_kernel_path(mut self, path: impl Into<String>) -> Self {
        self.kernel_path = Some(path.into());
        self
    }

    pub fn with_initramfs(mut self, path: impl Into<String>) -> Self {
        self.initrd_path = Some(path.into());
        self
    }

    pub fn with_cmdline(mut self, cmdline: impl Into<String>) -> Self {
        self.cmdline = cmdline.into();
        self
    }

    pub fn with_memory_mib(mut self, memory_mib: usize) -> Self {
        self.memory_mib = memory_mib;
        self
    }

    pub fn with_profile(mut self, profile: VmProfile) -> Self {
        self.apply_profile_defaults(profile);
        self
    }

    pub fn with_vcpus(mut self, vcpus: u32) -> Self {
        self.vcpus = vcpus;
        self
    }

    pub fn with_block_device(mut self, path: impl Into<String>) -> Self {
        self.block_devices.push(path.into());
        self
    }

    /// Attach a writable volume. See [`VolumeSpec`].
    pub fn with_volume(mut self, volume: VolumeSpec) -> Self {
        self.volumes.push(volume);
        self
    }

    pub fn with_restore(mut self, path: impl Into<String>) -> Self {
        self.restore_from = Some(path.into());
        self
    }

    pub fn with_cow_restore(mut self, enabled: bool) -> Self {
        self.cow_restore = enabled;
        self
    }

    pub fn with_snapshot_after_ms(mut self, after_ms: u64, out_path: impl Into<String>) -> Self {
        self.snapshot.after_ms = Some(after_ms);
        self.snapshot.out_path = Some(out_path.into());
        self
    }

    pub fn with_snapshot_at_heartbeat(
        mut self,
        at_heartbeat: u64,
        out_path: impl Into<String>,
    ) -> Self {
        self.snapshot.at_heartbeat = Some(at_heartbeat);
        self.snapshot.out_path = Some(out_path.into());
        self
    }

    pub fn with_snapshot_on_listener(mut self, out_path: impl Into<String>) -> Self {
        self.snapshot.on_listener = true;
        self.snapshot.out_path = Some(out_path.into());
        self
    }

    pub fn with_quiesce_ms(mut self, quiesce_ms: u64) -> Self {
        self.snapshot.quiesce_ms = quiesce_ms;
        self
    }

    pub fn with_vsock_mux(mut self, path: impl Into<String>) -> Self {
        self.endpoints.vsock_mux = Some(path.into());
        self
    }

    pub fn with_http_port(mut self, port: impl Into<String>) -> Self {
        self.endpoints.http_port = Some(port.into());
        self
    }

    pub fn with_vsock_mux_handoff(mut self, path: impl Into<String>) -> Self {
        self.endpoints.vsock_mux_handoff = Some(path.into());
        self
    }

    /// Set the path of the `<vsock_mux>-exec.sock` frontend. See
    /// [`EndpointResources::vsock_exec`].
    pub fn with_vsock_exec(mut self, path: impl Into<String>) -> Self {
        self.endpoints.vsock_exec = Some(path.into());
        self
    }

    /// Override the default guest-side AF_VSOCK port the exec
    /// agent listens on. Match this to whatever your guest agent
    /// binds. Default is [`DEFAULT_EXEC_GUEST_PORT`].
    pub fn with_vsock_exec_guest_port(mut self, port: u32) -> Self {
        self.endpoints.vsock_exec_guest_port = Some(port);
        self
    }

    pub fn memory_bytes(&self) -> usize {
        self.memory_mib * 1024 * 1024
    }

    pub fn is_restore(&self) -> bool {
        self.restore_from.is_some()
    }

    pub fn apply_profile_defaults(&mut self, profile: VmProfile) {
        self.vcpus = profile.default_vcpus();
    }

    pub fn validate_for_run(&self) -> Result<(), ResourceError> {
        if self.memory_mib == 0 {
            return Err(ResourceError::ZeroMemory);
        }
        if self.vcpus == 0 {
            return Err(ResourceError::ZeroVcpus);
        }
        if self.kernel_path.is_none() && self.restore_from.is_none() {
            return Err(ResourceError::MissingKernel);
        }
        let wants_snapshot = self.snapshot.after_ms.is_some()
            || self.snapshot.at_heartbeat.is_some()
            || self.snapshot.on_listener;
        if wants_snapshot && self.snapshot.out_path.is_none() {
            return Err(ResourceError::SnapshotTriggerWithoutOutput);
        }
        Ok(())
    }
}

impl Default for VmResources {
    fn default() -> Self {
        Self {
            kernel_path: None,
            initrd_path: None,
            cmdline: DEFAULT_CMDLINE.to_string(),
            memory_mib: DEFAULT_MEMORY_MIB,
            block_devices: Vec::new(),
            volumes: Vec::new(),
            vcpus: DEFAULT_VCPUS,
            restore_from: None,
            cow_restore: false,
            snapshot: SnapshotResources::default(),
            endpoints: EndpointResources::default(),
            balloon_target_pages: None,
        }
    }
}

impl fmt::Display for ResourceError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ResourceError::MissingKernel => {
                write!(f, "kernel path or restore snapshot is required")
            }
            ResourceError::ZeroMemory => write!(f, "memory must be greater than zero"),
            ResourceError::ZeroVcpus => write!(f, "vCPU count must be greater than zero"),
            ResourceError::SnapshotTriggerWithoutOutput => {
                write!(f, "snapshot trigger requires snapshot output path")
            }
        }
    }
}

impl std::error::Error for ResourceError {}

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

    #[test]
    fn defaults_match_supermachine_cli() {
        let resources = VmResources::default();
        assert_eq!(resources.cmdline, DEFAULT_CMDLINE);
        assert_eq!(resources.memory_mib, 256);
        assert_eq!(resources.vcpus, 1);
        assert_eq!(resources.memory_bytes(), 256 * 1024 * 1024);
    }

    #[test]
    fn profile_defaults_are_stable() {
        assert_eq!(VmProfile::parse("latency"), Some(VmProfile::Latency));
        assert_eq!(VmProfile::parse("low-latency"), Some(VmProfile::Latency));
        assert_eq!(VmProfile::parse("throughput"), Some(VmProfile::Throughput));
        assert_eq!(VmProfile::parse("unknown"), None);
        assert_eq!(VmProfile::Latency.default_vcpus(), 1);
        assert_eq!(VmProfile::Throughput.default_vcpus(), 4);
    }

    #[test]
    fn applies_profile_defaults_to_resources() {
        let mut resources = VmResources::default();
        resources.apply_profile_defaults(VmProfile::Throughput);
        assert_eq!(resources.vcpus, 4);
        resources.apply_profile_defaults(VmProfile::Latency);
        assert_eq!(resources.vcpus, 1);
    }

    #[test]
    fn convenience_constructors_cover_kernel_and_restore() {
        let kernel = VmResources::for_kernel("kernel", "initrd");
        assert_eq!(kernel.kernel_path.as_deref(), Some("kernel"));
        assert_eq!(kernel.initrd_path.as_deref(), Some("initrd"));
        assert!(!kernel.is_restore());

        let restore = VmResources::from_snapshot("snap.sm");
        assert_eq!(restore.restore_from.as_deref(), Some("snap.sm"));
        assert!(restore.is_restore());
    }

    #[test]
    fn builder_style_methods_cover_common_library_config() {
        let resources = VmResources::new()
            .with_kernel_path("kernel")
            .with_initramfs("initrd")
            .with_cmdline("console=ttyS0")
            .with_memory_mib(512)
            .with_profile(VmProfile::Throughput)
            .with_vcpus(2)
            .with_block_device("rootfs.squashfs")
            .with_snapshot_at_heartbeat(1, "snap.sm")
            .with_quiesce_ms(7)
            .with_vsock_mux("/tmp/vsock.sock")
            .with_http_port("8080");

        assert_eq!(resources.kernel_path.as_deref(), Some("kernel"));
        assert_eq!(resources.initrd_path.as_deref(), Some("initrd"));
        assert_eq!(resources.cmdline, "console=ttyS0");
        assert_eq!(resources.memory_mib, 512);
        assert_eq!(resources.vcpus, 2);
        assert_eq!(resources.block_devices, vec!["rootfs.squashfs"]);
        assert_eq!(resources.snapshot.at_heartbeat, Some(1));
        assert_eq!(resources.snapshot.quiesce_ms, 7);
        assert_eq!(resources.snapshot.out_path.as_deref(), Some("snap.sm"));
        assert_eq!(
            resources.endpoints.vsock_mux.as_deref(),
            Some("/tmp/vsock.sock")
        );
        assert_eq!(resources.endpoints.http_port.as_deref(), Some("8080"));
    }

    #[test]
    fn builder_style_restore_config_is_valid_without_kernel() {
        let resources = VmResources::new()
            .with_restore("snap.sm")
            .with_cow_restore(true);

        assert!(resources.is_restore());
        assert!(resources.cow_restore);
        assert_eq!(resources.validate_for_run(), Ok(()));
    }

    #[test]
    fn validates_kernel_or_restore() {
        let mut resources = VmResources::default();
        assert_eq!(
            resources.validate_for_run(),
            Err(ResourceError::MissingKernel)
        );

        resources.kernel_path = Some("vmlinux".to_string());
        assert_eq!(resources.validate_for_run(), Ok(()));

        resources.kernel_path = None;
        resources.restore_from = Some("snap.sm".to_string());
        assert_eq!(resources.validate_for_run(), Ok(()));
    }
}