sausage/userland/brush/process.rs

//! Launches world components and pumps each child's bounded streams independently.
#![allow(clippy::all, missing_docs)]
use crate::{
    ExecutionResult,
    openfiles::OpenFile,
    processes::{ChildProcess, LaunchRequest},
};
use std::io::{Read, Write};
mod bindings {
    wit_bindgen::generate!({ path: "sausage-process", world: "program", generate_all });
}
use bindings::{
    sausage::process::host,
    wasi::io::streams::{InputStream, StreamError},
};
pub(crate) type ProcessId = i32;
fn drain(stream: &mut Option<InputStream>) -> std::io::Result<Vec<u8>> {
    let Some(reader) = stream.as_ref() else {
        return Ok(Vec::new());
    };
    match reader.read(8192) {
        Ok(bytes) => Ok(bytes),
        Err(StreamError::Closed) => {
            *stream = None;
            Ok(Vec::new())
        }
        Err(_) => Err(std::io::Error::other("child stream failed")),
    }
}
pub(crate) fn spawn(mut request: LaunchRequest) -> std::io::Result<ChildProcess> {
    let child = host::spawn(&host::Request {
        program: request.program.to_string_lossy().into_owned(),
        argv: std::iter::once(request.argv0).chain(request.args).collect(),
        env: request.env.into_iter().collect(),
        cwd: request.working_dir.to_string_lossy().into_owned(),
        group: match request.process_group {
            crate::processes::ProcessGroup::Join(group) => Some(group as u32),
            crate::processes::ProcessGroup::Inherit => None,
            _ => Some(0),
        },
        stdin: host::Stdio::Pipe,
        stdout: host::Stdio::Pipe,
        stderr: host::Stdio::Pipe,
    })
    .map_err(std::io::Error::other)?;
    let pid = child.pid() as i32;
    let group = child.group() as i32;
    let mut send = Some(child.stdin());
    let mut receive = Some(child.stdout());
    let mut errors = Some(child.stderr());
    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);
    let pump = tokio::spawn(async move {
        let mut pending = Vec::new();
        loop {
            if pending.is_empty() {
                if let Some(reader) = input.as_mut() {
                    let mut buf = [0; 8192];
                    match reader.read(&mut buf) {
                        Ok(0) => input = None,
                        Ok(n) => pending.extend_from_slice(&buf[..n]),
                        Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {}
                        Err(e) => return Err(e),
                    }
                }
            }
            if let Some(writer) = send.as_ref() {
                match writer.check_write() {
                    Ok(n) if n > 0 => {
                        let n = (n as usize).min(pending.len());
                        if n != 0 {
                            writer
                                .write(&pending[..n])
                                .map_err(|_| std::io::Error::other("child stdin failed"))?;
                            writer
                                .flush()
                                .map_err(|_| std::io::Error::other("child stdin flush failed"))?;
                            pending.drain(..n);
                        } else if input.is_none() {
                            // A completed flush must precede EOF, or the writer drops queued bytes.
                            send = None;
                        }
                    }
                    Ok(_) => {}
                    Err(_) => {
                        send = None;
                        input = None;
                        pending.clear();
                    }
                }
            }
            if output.is_none() {
                receive = None;
            }
            let out = drain(&mut receive)?;
            let err = drain(&mut errors)?;
            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);
                    }
                    output = None;
                }
            }
            if let Some(writer) = error.as_mut() {
                writer.write_all(&err)?;
            }
            if receive.is_none() && errors.is_none() {
                if let Some(code) = child.status() {
                    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)? })
            .native_child(Some(pid), Some(group)),
    )
}
pub(crate) fn exec(_request: LaunchRequest) -> std::io::Error {
    std::io::ErrorKind::Unsupported.into()
}