feat: run shell commands asynchronously from command mode

src/shell.rs

@@ -0,0 +1,240 @@
+#[cfg(unix)]
+use std::{
+    ffi::OsString,
+    io::{self, Read},
+    os::{
+        fd::OwnedFd,
+        unix::{net::UnixStream, process::CommandExt},
+    },
+    process::{Child, Command, Stdio},
+    sync::mpsc::{self, RecvTimeoutError},
+    thread,
+    time::Duration,
+};
+use std::{
+    path::PathBuf,
+    sync::mpsc::{Receiver, Sender, TryRecvError},
+    thread::JoinHandle,
+};
+
+#[cfg(unix)]
+const OUTPUT_LIMIT: usize = 4 * 1024 * 1024;
+#[cfg(unix)]
+const TRUNCATION_NOTE: &[u8] = b"\n[vamp: output truncated after 4 MiB]\n";
+
+#[derive(Debug)]
+pub(crate) struct Output {
+    pub(crate) code: Option<i32>,
+    pub(crate) stdout: Vec<u8>,
+    pub(crate) stderr: Vec<u8>,
+}
+
+pub(crate) struct Worker {
+    cancel: Sender<()>,
+    results: Receiver<Result<Output, String>>,
+    thread: Option<JoinHandle<()>>,
+    finished: bool,
+}
+
+impl Worker {
+    pub(crate) fn start(command: String, root: PathBuf) -> Result<Self, String> {
+        #[cfg(unix)]
+        {
+            let shell = std::env::var_os("SHELL")
+                .filter(|shell| !shell.is_empty())
+                .unwrap_or_else(|| "/bin/sh".into());
+            Self::start_with_shell(command, root, shell)
+        }
+        #[cfg(not(unix))]
+        {
+            let _ = (command, root);
+            Err("Shell commands require Unix process-group cancellation on this platform".into())
+        }
+    }
+
+    #[cfg(unix)]
+    fn start_with_shell(command: String, root: PathBuf, shell: OsString) -> Result<Self, String> {
+        let (cancel, cancelled) = mpsc::channel();
+        let (send, results) = mpsc::channel();
+        let thread = thread::Builder::new()
+            .name("vamp-shell".into())
+            .spawn(move || {
+                let _ = send.send(run(command, root, shell, cancelled));
+            })
+            .map_err(|error| format!("Cannot start shell worker: {error}"))?;
+        Ok(Self {
+            cancel,
+            results,
+            thread: Some(thread),
+            finished: false,
+        })
+    }
+
+    pub(crate) fn poll(&mut self) -> Option<Result<Output, String>> {
+        if self.finished {
+            return None;
+        }
+        let result = match self.results.try_recv() {
+            Ok(output) => output,
+            Err(TryRecvError::Empty) => return None,
+            Err(TryRecvError::Disconnected) => Err("Shell worker stopped unexpectedly".into()),
+        };
+        self.finished = true;
+        Some(result)
+    }
+}
+
+impl Drop for Worker {
+    fn drop(&mut self) {
+        let _ = self.cancel.send(());
+        if let Some(thread) = self.thread.take() {
+            let _ = thread.join();
+        }
+    }
+}
+
+#[cfg(unix)]
+struct Process(Child);
+
+#[cfg(unix)]
+impl Drop for Process {
+    fn drop(&mut self) {
+        unsafe extern "C" {
+            fn kill(pid: i32, signal: i32) -> i32;
+        }
+        // The shell and its descendants share a private process group; killing
+        // only the leader would leave children holding the output descriptors.
+        unsafe { kill(-(self.0.id() as i32), 9) };
+        let _ = self.0.kill();
+        let _ = self.0.wait();
+    }
+}
+
+#[cfg(unix)]
+struct Capture {
+    reader: UnixStream,
+    bytes: Vec<u8>,
+    truncated: bool,
+    eof: bool,
+}
+
+#[cfg(unix)]
+impl Capture {
+    fn new() -> Result<(Self, Stdio), String> {
+        let (reader, writer) = UnixStream::pair()
+            .map_err(|error| format!("Cannot create shell output stream: {error}"))?;
+        reader
+            .set_nonblocking(true)
+            .map_err(|error| format!("Cannot configure shell output stream: {error}"))?;
+        Ok((
+            Self {
+                reader,
+                bytes: Vec::new(),
+                truncated: false,
+                eof: false,
+            },
+            Stdio::from(OwnedFd::from(writer)),
+        ))
+    }
+
+    fn read(&mut self) -> io::Result<bool> {
+        if self.eof {
+            return Ok(false);
+        }
+        let mut chunk = [0; 16 * 1024];
+        match self.reader.read(&mut chunk) {
+            Ok(0) => {
+                self.eof = true;
+                Ok(false)
+            }
+            Ok(count) => {
+                let keep = count.min(OUTPUT_LIMIT.saturating_sub(self.bytes.len()));
+                self.bytes.extend_from_slice(&chunk[..keep]);
+                self.truncated |= keep != count;
+                Ok(true)
+            }
+            Err(error) if error.kind() == io::ErrorKind::WouldBlock => Ok(false),
+            Err(error) if error.kind() == io::ErrorKind::Interrupted => Ok(true),
+            Err(error) => Err(error),
+        }
+    }
+
+    fn finish(mut self) -> Vec<u8> {
+        if self.truncated {
+            self.bytes.extend_from_slice(TRUNCATION_NOTE);
+        }
+        self.bytes
+    }
+}
+
+#[cfg(unix)]
+fn run(
+    command: String,
+    root: PathBuf,
+    shell: OsString,
+    cancelled: Receiver<()>,
+) -> Result<Output, String> {
+    if cancelled.try_recv().is_ok() {
+        return Err("Shell command cancelled".into());
+    }
+    let (mut stdout, stdout_writer) = Capture::new()?;
+    let (mut stderr, stderr_writer) = Capture::new()?;
+    let child = Command::new(shell)
+        .arg("-c")
+        .arg(command)
+        .current_dir(root)
+        .stdin(Stdio::null())
+        .stdout(stdout_writer)
+        .stderr(stderr_writer)
+        .process_group(0)
+        .spawn()
+        .map_err(|error| format!("Cannot start shell command: {error}"))?;
+    let mut process = Some(Process(child));
+    let mut status = None;
+    loop {
+        if cancelled.try_recv().is_ok() {
+            return Err("Shell command cancelled".into());
+        }
+        // Read each stream once per turn so a busy writer cannot starve the
+        // other stream or cancellation. Overflow is still drained, not retained.
+        let out = stdout
+            .read()
+            .map_err(|error| format!("Cannot read shell stdout: {error}"))?;
+        let err = stderr
+            .read()
+            .map_err(|error| format!("Cannot read shell stderr: {error}"))?;
+        if let Some(child) = process.as_mut() {
+            status = child
+                .0
+                .try_wait()
+                .map_err(|error| format!("Cannot wait for shell command: {error}"))?;
+            if status.is_some() {
+                // Background jobs must not outlive the command. Killing the
+                // remaining group also closes inherited output descriptors.
+                drop(process.take());
+            }
+        }
+        if let Some(status) = status
+            && stdout.eof
+            && stderr.eof
+        {
+            return Ok(Output {
+                code: status.code(),
+                stdout: stdout.finish(),
+                stderr: stderr.finish(),
+            });
+        }
+        if !out && !err {
+            match cancelled.recv_timeout(Duration::from_millis(10)) {
+                Ok(()) | Err(RecvTimeoutError::Disconnected) => {
+                    return Err("Shell command cancelled".into());
+                }
+                Err(RecvTimeoutError::Timeout) => {}
+            }
+        }
+    }
+}
+
+#[cfg(all(test, unix))]
+#[path = "../tests/unit/shell.rs"]
+mod tests;