flux-platform 1.0.1

//! The shell runner: executes a command string in the project directory.
//!
//! Two modes:
//! * [`run`] hands the command Flux's own stdio (inherited) — used by the
//!   first-party tools (`flux fmt`, `flux lint`, deploy), which run one command
//!   at a time and want its raw output;
//! * [`run_streamed`] pipes stdout and stderr, delivers them to a [`LineSink`]
//!   one whole line at a time as they arrive, and enforces a wall-clock limit —
//!   used by the graph engine so a parallel run stays readable and a hung step
//!   cannot wedge the pipeline.

use std::io::{self, BufRead, BufReader, Read};
use std::path::Path;
use std::process::{Command, Stdio};
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

/// The result of running a single command (streamed).
#[derive(Debug, Clone)]
pub struct CommandResult {
    pub success: bool,
}

/// Which of the child's two streams a line came from.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Stream {
    Stdout,
    Stderr,
}

/// A consumer of the command's output, called once per complete line as it
/// arrives. Shared across the two reader threads, so it must be `Send + Sync`.
pub type LineSink = Arc<dyn Fn(Stream, &str) + Send + Sync>;

/// The result of running a single command with its output streamed.
#[derive(Debug, Clone)]
pub struct StreamedResult {
    pub success: bool,
    /// The command exceeded its limit and was killed.
    pub timed_out: bool,
    pub duration: Duration,
    /// Combined stdout + stderr, retained for failure diagnosis (Flux Assist).
    pub output: String,
}

/// Longest line handed to the sink in one piece. A command that emits binary
/// data (or a progress bar that never writes a newline) would otherwise grow
/// the line buffer without bound; splitting is nicer than dying.
const MAX_LINE: usize = 8 * 1024;

/// How long to keep draining after the process exits. Its pipes normally close
/// with it; they stay open only when the command left a grandchild holding the
/// other end, and waiting forever on that is how a build hangs after finishing.
const DRAIN_GRACE: Duration = Duration::from_secs(1);

/// Upper bound on the exit-poll interval. `std` has no `wait_with_timeout`, so
/// the deadline is enforced by polling `try_wait`; the backoff keeps a short
/// command responsive without spinning through a long one.
const MAX_POLL: Duration = Duration::from_millis(20);

/// Build a platform shell invocation for `cmd`.
#[cfg(windows)]
fn shell(cmd: &str) -> Command {
    let mut c = Command::new("cmd");
    c.arg("/C").arg(cmd);
    c
}

#[cfg(not(windows))]
fn shell(cmd: &str) -> Command {
    let mut c = Command::new("sh");
    c.arg("-c").arg(cmd);
    c
}

/// Run `cmd` in `dir`, letting it write straight to Flux's stdio. Returns
/// whether it succeeded.
pub fn run(cmd: &str, dir: &Path) -> io::Result<CommandResult> {
    let status = shell(cmd).current_dir(dir).status()?;
    Ok(CommandResult {
        success: status.success(),
    })
}

/// Run `cmd` in `dir` with extra environment variables, streaming each output
/// line to `sink` as it arrives and killing the command if it outlives
/// `timeout` (`None` runs it unbounded).
///
/// Killing reaches the shell Flux started. A command that spawns background
/// children of its own can leave those behind — Flux says the step timed out
/// rather than pretending it reaped a process tree it cannot see.
pub fn run_streamed(
    cmd: &str,
    dir: &Path,
    env: &[(String, String)],
    timeout: Option<Duration>,
    sink: LineSink,
) -> io::Result<StreamedResult> {
    let start = Instant::now();

    let mut command = shell(cmd);
    command
        .current_dir(dir)
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped());
    for (k, v) in env {
        command.env(k, v);
    }
    let mut child = command.spawn()?;

    let tap = Arc::new(Tap::new(sink));
    let (done_tx, done_rx) = channel::<()>();
    let mut readers = 0usize;
    if let Some(out) = child.stdout.take() {
        spawn_pump(out, Stream::Stdout, Arc::clone(&tap), done_tx.clone());
        readers += 1;
    }
    if let Some(err) = child.stderr.take() {
        spawn_pump(err, Stream::Stderr, Arc::clone(&tap), done_tx.clone());
        readers += 1;
    }
    drop(done_tx);

    let deadline = timeout.map(|t| start + t);
    let mut timed_out = false;
    let mut poll = Duration::from_millis(1);
    let status = loop {
        match child.try_wait()? {
            Some(status) => break Some(status),
            None => {
                if deadline.is_some_and(|d| Instant::now() >= d) {
                    // Kill, then reap: `kill` only asks, and an unreaped child
                    // would linger as a zombie for the life of the process.
                    let _ = child.kill();
                    let _ = child.wait();
                    timed_out = true;
                    break None;
                }
                std::thread::sleep(poll);
                poll = (poll * 2).min(MAX_POLL);
            }
        }
    };

    // Give the readers a moment to deliver whatever the command wrote last,
    // then stop consuming. Threads still blocked on an inherited pipe are left
    // to end on their own; `Tap::close` makes sure they can no longer print.
    drain(&done_rx, readers);
    let output = tap.close();

    Ok(StreamedResult {
        success: !timed_out && status.map(|s| s.success()).unwrap_or(false),
        timed_out,
        duration: start.elapsed(),
        output,
    })
}

/// Wait for up to `readers` completion signals, at most [`DRAIN_GRACE`] each.
fn drain(done_rx: &Receiver<()>, readers: usize) {
    for _ in 0..readers {
        if done_rx.recv_timeout(DRAIN_GRACE).is_err() {
            break;
        }
    }
}

/// The shared destination for one command's output: it accumulates the text for
/// failure diagnosis and forwards each line to the sink, under one lock so that
/// closing it is a hard stop rather than a race with an in-flight line.
struct Tap {
    sink: LineSink,
    /// `None` once closed — late lines from a detached reader are dropped.
    buffer: Mutex<Option<String>>,
}

impl Tap {
    fn new(sink: LineSink) -> Self {
        Tap {
            sink,
            buffer: Mutex::new(Some(String::new())),
        }
    }

    fn push(&self, stream: Stream, line: &str) {
        let mut guard = self.buffer.lock().unwrap_or_else(|e| e.into_inner());
        let Some(buffer) = guard.as_mut() else {
            return;
        };
        buffer.push_str(line);
        buffer.push('\n');
        (self.sink)(stream, line);
    }

    /// Stop accepting lines and take the accumulated output.
    fn close(&self) -> String {
        let mut guard = self.buffer.lock().unwrap_or_else(|e| e.into_inner());
        guard.take().unwrap_or_default()
    }
}

/// Read `reader` line by line into `tap` until EOF, then signal `done`.
///
/// Deliberately detached rather than scoped: after a timeout the pipe can stay
/// open (a surviving grandchild holds it), and joining that thread would hang
/// the very build the timeout exists to unblock.
fn spawn_pump<R: Read + Send + 'static>(
    reader: R,
    stream: Stream,
    tap: Arc<Tap>,
    done: Sender<()>,
) {
    std::thread::spawn(move || {
        let mut reader = BufReader::new(reader);
        let mut raw: Vec<u8> = Vec::new();
        loop {
            match read_line_bounded(&mut reader, &mut raw) {
                Ok(0) | Err(_) => break,
                Ok(_) => {
                    let line = String::from_utf8_lossy(trim_eol(&raw));
                    tap.push(stream, &line);
                }
            }
        }
        let _ = done.send(());
    });
}

fn trim_eol(bytes: &[u8]) -> &[u8] {
    let mut end = bytes.len();
    if end > 0 && bytes[end - 1] == b'\n' {
        end -= 1;
    }
    if end > 0 && bytes[end - 1] == b'\r' {
        end -= 1;
    }
    &bytes[..end]
}

/// Read one line (including its newline) into `out`, or up to [`MAX_LINE`]
/// bytes of one that never ends. Returns the number of bytes read; `0` is EOF.
fn read_line_bounded(reader: &mut impl BufRead, out: &mut Vec<u8>) -> io::Result<usize> {
    out.clear();
    loop {
        let available = match reader.fill_buf() {
            Ok(bytes) => bytes,
            Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
            Err(e) => return Err(e),
        };
        if available.is_empty() {
            return Ok(out.len()); // EOF, possibly with a trailing partial line
        }
        match available.iter().position(|&b| b == b'\n') {
            Some(i) => {
                out.extend_from_slice(&available[..=i]);
                reader.consume(i + 1);
                return Ok(out.len());
            }
            None => {
                let n = available.len();
                out.extend_from_slice(available);
                reader.consume(n);
                if out.len() >= MAX_LINE {
                    return Ok(out.len());
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering};

    /// Every line a test collector saw, in arrival order.
    type Collected = Arc<Mutex<Vec<(Stream, String)>>>;

    /// Collect every streamed line as `(stream, text)`.
    fn collector() -> (LineSink, Collected) {
        let lines = Arc::new(Mutex::new(Vec::new()));
        let seen = Arc::clone(&lines);
        let sink: LineSink = Arc::new(move |stream, line: &str| {
            seen.lock().unwrap().push((stream, line.to_string()));
        });
        (sink, lines)
    }

    /// A command that blocks for far longer than any test wants to wait.
    fn long_command() -> &'static str {
        if cfg!(windows) {
            "ping -n 60 127.0.0.1 >nul"
        } else {
            "sleep 60"
        }
    }

    #[test]
    fn streams_stdout_and_stderr_separately() {
        let (sink, lines) = collector();
        let res =
            run_streamed("echo one && echo two 1>&2", Path::new("."), &[], None, sink).unwrap();
        assert!(res.success, "{res:?}");
        assert!(!res.timed_out);

        let lines = lines.lock().unwrap();
        let out: Vec<&(Stream, String)> =
            lines.iter().filter(|(s, _)| *s == Stream::Stdout).collect();
        let err: Vec<&(Stream, String)> =
            lines.iter().filter(|(s, _)| *s == Stream::Stderr).collect();
        assert!(
            out.iter().any(|(_, l)| l.trim() == "one"),
            "stdout not streamed: {lines:?}"
        );
        assert!(
            err.iter().any(|(_, l)| l.trim() == "two"),
            "stderr not streamed: {lines:?}"
        );
        // Both are also retained for failure diagnosis.
        assert!(res.output.contains("one") && res.output.contains("two"));
    }

    /// Lines must reach the sink while the command is still running, otherwise
    /// "streaming" is just a buffer with extra steps.
    #[test]
    fn lines_arrive_before_the_command_exits() {
        let seen = Arc::new(AtomicUsize::new(0));
        let counter = Arc::clone(&seen);
        let sink: LineSink = Arc::new(move |_, _: &str| {
            counter.fetch_add(1, Ordering::SeqCst);
        });
        let slow = if cfg!(windows) {
            "echo first && ping -n 3 127.0.0.1 >nul && echo second"
        } else {
            "echo first && sleep 2 && echo second"
        };

        let watcher = Arc::clone(&seen);
        let probe = std::thread::spawn(move || {
            std::thread::sleep(Duration::from_millis(800));
            watcher.load(Ordering::SeqCst)
        });

        let res = run_streamed(slow, Path::new("."), &[], None, sink).unwrap();
        assert!(res.success, "{res:?}");
        assert_eq!(
            probe.join().unwrap(),
            1,
            "the first line should have been delivered while the command was still running"
        );
        assert_eq!(seen.load(Ordering::SeqCst), 2);
    }

    #[test]
    fn timeout_kills_a_hanging_command() {
        let (sink, _) = collector();
        let start = Instant::now();
        let res = run_streamed(
            long_command(),
            Path::new("."),
            &[],
            Some(Duration::from_millis(300)),
            sink,
        )
        .unwrap();

        assert!(res.timed_out, "should report a timeout: {res:?}");
        assert!(!res.success, "a killed command has not succeeded");
        assert!(
            start.elapsed() < Duration::from_secs(20),
            "the command should have been killed, not waited out: {:?}",
            start.elapsed()
        );
    }

    #[test]
    fn a_command_finishing_inside_its_limit_is_untouched() {
        let (sink, _) = collector();
        let res = run_streamed(
            "echo quick",
            Path::new("."),
            &[],
            Some(Duration::from_secs(30)),
            sink,
        )
        .unwrap();
        assert!(res.success && !res.timed_out, "{res:?}");
        assert!(res.output.contains("quick"));
    }

    #[test]
    fn injected_environment_reaches_the_command() {
        let (sink, _) = collector();
        let echo = if cfg!(windows) {
            "echo token=%TOKEN%"
        } else {
            "echo token=$TOKEN"
        };
        let res = run_streamed(
            echo,
            Path::new("."),
            &[("TOKEN".to_string(), "abc123".to_string())],
            None,
            sink,
        )
        .unwrap();
        assert!(res.output.contains("token=abc123"), "{}", res.output);
    }

    #[test]
    fn a_failing_command_reports_failure_without_a_timeout() {
        let (sink, _) = collector();
        let res = run_streamed("exit 3", Path::new("."), &[], None, sink).unwrap();
        assert!(!res.success);
        assert!(!res.timed_out, "a plain failure is not a timeout");
    }

    #[test]
    fn a_line_that_never_ends_is_split_rather_than_buffered_forever() {
        let mut reader = BufReader::new(io::Cursor::new(vec![b'x'; MAX_LINE * 2 + 5]));
        let mut raw = Vec::new();

        assert_eq!(read_line_bounded(&mut reader, &mut raw).unwrap(), MAX_LINE);
        assert_eq!(read_line_bounded(&mut reader, &mut raw).unwrap(), MAX_LINE);
        assert_eq!(read_line_bounded(&mut reader, &mut raw).unwrap(), 5);
        assert_eq!(read_line_bounded(&mut reader, &mut raw).unwrap(), 0);
    }

    #[test]
    fn crlf_and_lf_line_endings_both_trim() {
        assert_eq!(trim_eol(b"hello\n"), b"hello");
        assert_eq!(trim_eol(b"hello\r\n"), b"hello");
        assert_eq!(trim_eol(b"hello"), b"hello");
        assert_eq!(trim_eol(b"\n"), b"");
    }
}