sausage/spikes/s2-shell/guest-process.rs

//! Launches external components through the S2 host import and pumps their streams.
#![allow(clippy::all, missing_docs)]
use crate::{
    ExecutionResult,
    openfiles::OpenFile,
    processes::{ChildProcess, LaunchRequest},
};
use std::io::{Read, Write};
mod bindings {
    wit_bindgen::generate!({ path: "s2-process.wit", world: "shell" });
}
pub(crate) type ProcessId = i32;
pub(crate) fn spawn(mut request: LaunchRequest) -> std::io::Result<ChildProcess> {
    let json = serde_json::json!({
        "argv": std::iter::once(request.argv0).chain(request.args).collect::<Vec<_>>(),
        "env": request.env,
        "cwd": request.working_dir,
    })
    .to_string();
    let pid = bindings::s2::process::host::spawn(&json);
    let mut input = request.open_files.remove_fd(0);
    if matches!(input, Some(OpenFile::Stdin(_))) {
        input = None;
    }
    let mut output = request.open_files.remove_fd(1);
    let mut error = request.open_files.remove_fd(2);
    // Brush waits for pipeline stages sequentially. Pump every child's IO independently.
    let pump = tokio::spawn(async move {
        let mut pending = Vec::new();
        let mut eof = input.is_none();
        let mut closed = output.is_none();
        loop {
            if pending.is_empty() && !eof {
                let mut buf = [0; 8192];
                match input
                    .as_mut()
                    .expect("input exists until EOF")
                    .read(&mut buf)
                {
                    Ok(0) => {
                        eof = true;
                        input = None;
                    }
                    Ok(n) => pending.extend_from_slice(&buf[..n]),
                    Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {}
                    Err(e) => return Err(e),
                }
            }
            let (consumed, done, code, out, err) =
                bindings::s2::process::host::exchange(pid, &pending, eof, closed);
            pending.drain(..consumed as usize);
            if let Some(writer) = output.as_mut() {
                if let Err(e) = writer.write_all(&out) {
                    if e.kind() != std::io::ErrorKind::BrokenPipe {
                        return Err(e);
                    }
                    closed = true;
                    output = None;
                }
            }
            if let Some(writer) = error.as_mut() {
                writer.write_all(&err)?;
            }
            if done {
                return Ok(ExecutionResult::new(code as u8));
            }
            tokio::task::yield_now().await;
        }
    });
    Ok(ChildProcess::from_future(async move {
        pump.await.map_err(std::io::Error::other)?
    }))
}
pub(crate) fn exec(_request: LaunchRequest) -> std::io::Error {
    std::io::ErrorKind::Unsupported.into()
}