Coverage / crates/flow/src/microvm.rs 31.2%

pub fn environment(home: &Home, sandbox: &Sandbox, job: &Job) -> Result<Environment> {

    let cfg = home.config.microvm.as_ref().context("flow.toml has no [microvm] section")?;

        runtime: "microvm".into(),

        parts: vec![

            ("firecracker".into(), cfg.firecracker.blake3.clone()),

            ("kernel".into(), cfg.kernel.blake3.clone()),

            ("init".into(), cfg.init.blake3.clone()),

            ("memory_mib".into(), sandbox.memory_mib.to_string()),

            ("cpus".into(), sandbox.cpus.to_string()),

        .into_iter()

        .chain(job.mounts.iter().map(|m| (format!("mount {}", m.target), format!("{} {} MiB", m.source.as_str(), m.size_mib))))

        .collect(),

        rootfs: sandbox.rootfs.clone(),

        argv: job.exec.clone(),

        env: job.env.iter().map(|(k, v)| (k.clone(), v.clone())).collect(),

}

pub fn run(

    home: &Home,

    sandbox: &Sandbox,

    job: &Job,

    secrets: &[SecretInput],

    workspace: &Path,

    scratch: &Path,

    timeout: Duration,

    sink: &mut dyn Sink,

) -> Result<Outcome> {

    let cfg = home.config.microvm.as_ref().context("flow.toml has no [microvm] section")?;

    let firecracker = home.verified(&home.resolve(&cfg.firecracker.path), &cfg.firecracker.blake3)?;

    let kernel = home.verified(&home.resolve(&cfg.kernel.path), &cfg.kernel.blake3)?;

    let init = home.verified(&home.resolve(&cfg.init.path), &cfg.init.blake3)?;

    let rootfs = home.rootfs(&sandbox.rootfs);

    if !rootfs.is_file() {

        bail!("rootfs {} is not on this runner (import it with `flow rootfs import`)", sandbox.rootfs);

    }

    let rootfs = home.verified(&rootfs, &sandbox.rootfs)?;

    let environment = environment(home, sandbox, job)?;

    let percent = home.config.runner.as_ref().and_then(|r| r.cache_prune_percent).unwrap_or(flow_init::prune::DEFAULT_PERCENT);

    let mut spec = guest_spec(job, percent);

    spec["loopback"] = serde_json::Value::Bool(sandbox.network == flow_core::Network::Loopback);

    let egress = sandbox.network == flow_core::Network::Egress;

    let (egress_tx, egress_rx) = mpsc::channel::<String>();

    let mut _proxy = None;

    if egress {

        if !home.config.runner.as_ref().is_some_and(|r| r.egress) {

            bail!("this runner does not allow network egress (a runner allows it with `egress = true` under [runner] in flow.toml)");

        }

        _proxy = Some(

            crate::egress::start(&scratch.join(format!("v.sock_{}", crate::egress::VSOCK_PORT)), sandbox.egress_allow.clone(), egress_tx)

                .context("start the egress proxy")?,

        spec["egress"] = serde_json::Value::Bool(true);

        add_proxy_env(&mut spec);

        sink.record(Record {

            t: 0,

            stream: "flow".into(),

            text: format!("egress: through this runner, to {} only", sandbox.egress_allow.join(", ")),

        });

    }

    if !secrets.is_empty() {

        spec["secrets"] =

            serde_json::Value::Array(secrets.iter().map(|s| serde_json::json!({ "path": s.path, "value": s.value })).collect());

    }

    let mut disks = Vec::new();

    for (i, m) in job.mounts.iter().filter(|m| m.source == MountSource::Disk).enumerate() {

        let path = scratch.join(format!("mount-{i}.ext4"));

        if !path.exists() {

            crate::workspace::empty_ext4(&path, m.size_mib, "scratch")?;

        }

        disks.push(

            serde_json::json!({ "drive_id": format!("mount{i}"), "path_on_host": path, "is_root_device": false, "is_read_only": false }),

    let init_bytes = std::fs::read(&init)?;

    let spec_bytes = serde_json::to_vec(&spec)?;

    let initrd = crate::cpio::newc(&[

        crate::cpio::Entry { name: "init", mode: 0o755, data: &init_bytes },

        crate::cpio::Entry { name: "job.json", mode: 0o644, data: &spec_bytes },

    ]);

    let initrd_path = scratch.join("initrd.cpio");

    std::fs::write(&initrd_path, initrd)?;

    let drives: Vec<serde_json::Value> = [

        serde_json::json!({ "drive_id": "rootfs", "path_on_host": rootfs, "is_root_device": false, "is_read_only": true }),

        serde_json::json!({ "drive_id": "workspace", "path_on_host": workspace, "is_root_device": false, "is_read_only": false }),

    ]

    .into_iter()

    .chain(disks)

    .collect();

    let vm = serde_json::json!({

        "boot-source": {

            "kernel_image_path": kernel,

            "initrd_path": initrd_path,

            "boot_args": "console=ttyS0 reboot=k panic=-1 pci=off quiet loglevel=0 rdinit=/init",

        "drives": drives,

        "machine-config": { "vcpu_count": sandbox.cpus, "mem_size_mib": sandbox.memory_mib },

    let mut vm = vm;

    if egress {

        // Relative to the VMM's working directory (the scratch directory),

        // which keeps the socket path short.

        vm["vsock"] = serde_json::json!({ "guest_cid": 3, "uds_path": "v.sock" });

    }

    let vm_path = scratch.join("vm.json");

    std::fs::write(&vm_path, serde_json::to_vec_pretty(&vm)?)?;

    let started = Instant::now();

    let mut child = std::process::Command::new(&firecracker)

        .arg("--no-api")

        .arg("--config-file")

        .arg(&vm_path)

        .args(["--level", "Error"])

        .current_dir(scratch)

        .stdin(std::process::Stdio::null())

        .stdout(std::process::Stdio::piped())

        .stderr(std::process::Stdio::piped())

        .spawn()

        .with_context(|| format!("start {}", firecracker.display()))?;

    let (tx, rx) = mpsc::channel::<(bool, Vec<u8>)>();

    let out = child.stdout.take().unwrap();

    let err = child.stderr.take().unwrap();

    let tx2 = tx.clone();

    let readers = [std::thread::spawn(move || lines(out, true, tx)), std::thread::spawn(move || lines(err, false, tx2))];

    let mut exit_code = None;

    let mut timed_out = false;

    let mut cancelled = false;

    let deadline = started + timeout;

    let t = |s: &Instant| s.elapsed().as_millis() as u64;

    let mut peak_memory_mib = None;

    let mut pruned = Vec::new();

        if let Some(m) = crate::facts::peak_rss_mib(child.id()) {

            peak_memory_mib = Some(m);

        }

        while let Ok(line) = egress_rx.try_recv() {

            sink.record(Record { t: t(&started), stream: "flow".into(), text: line });

        }

        match rx.recv_timeout(Duration::from_millis(200)) {

            Ok((from_console, line)) => {

                let text = String::from_utf8_lossy(&line).trim_end_matches('\r').to_string();

                let rec = |stream: &str, text: &str| Record { t: t(&started), stream: stream.into(), text: text.into() };

                if !from_console {

                    sink.record(rec("console", &format!("firecracker: {text}")));

                } else if let Some(rest) = text.strip_prefix('\u{1e}') {

                    let (kind, body) = rest.split_at(rest.len().min(1));

                    let body = body.strip_prefix(' ').unwrap_or(body);

                    match kind {

                        "O" => sink.record(rec("out", body)),

                        "E" => sink.record(rec("err", body)),

                        "X" => exit_code = body.trim().parse().ok(),

                        "F" => sink.record(rec("flow", &format!("guest init failed: {body}"))),

                        "C" => match pruned_frame(body) {

                            Some(p) => pruned.push(p),

                            None => sink.record(rec("console", &text)),

                        _ => sink.record(rec("console", &text)),

                } else if !text.is_empty() {

                    sink.record(rec("console", &text));

                }

            Err(mpsc::RecvTimeoutError::Timeout) => {}

            Err(mpsc::RecvTimeoutError::Disconnected) => break,

        if !timed_out && !cancelled && sink.cancelled() {

            cancelled = true;

            let _ = child.kill();

            sink.record(Record { t: t(&started), stream: "flow".into(), text: "stopped: the run was cancelled".into() });

        }

        if Instant::now() >= deadline && !timed_out && !cancelled {

            timed_out = true;

            let _ = child.kill();

            sink.record(Record {

                t: t(&started),

                stream: "flow".into(),

                text: format!("stopped at the time limit of {} s", timeout.as_secs()),

            });

        }

    let status = child.wait()?;

    for r in readers {

        let _ = r.join();

    }

    if exit_code.is_none() && !timed_out && !cancelled {

        sink.record(Record {

            t: t(&started),

            stream: "flow".into(),

            text: format!("the microVM ended without the job's exit code ({status})"),

        });

    }

        exit_code: if timed_out || cancelled { None } else { exit_code },

        duration_ms: t(&started),

        environment,

        timed_out,

        cancelled,

        peak_memory_mib,

        pruned,

}

fn lines(r: impl Read, console: bool, tx: mpsc::Sender<(bool, Vec<u8>)>) {

    let mut r = BufReader::new(r);

    let mut buf = Vec::new();

        buf.clear();

        match r.read_until(b'\n', &mut buf) {

            Ok(0) | Err(_) => break,

                if buf.last() == Some(&b'\n') {

                    buf.pop();

                }

                if tx.send((console, buf.clone())).is_err() {

                    break;

                }

}