feat(pi-shell): add execute_shell_streams with per-stream raw byte sinks
This commit is contained in:
can1357
Generated
+1
@@ -2505,6 +2505,7 @@ dependencies = [
|
||||
"brush-builtins",
|
||||
"brush-core",
|
||||
"brush-parser 0.3.0",
|
||||
"bytes",
|
||||
"clap",
|
||||
"libc",
|
||||
"os_pipe",
|
||||
|
||||
@@ -189,6 +189,7 @@ toml = "1.1"
|
||||
# Error Handling
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
anyhow = "1.0"
|
||||
bytes = "1"
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# CLI & Configuration
|
||||
|
||||
@@ -11,6 +11,7 @@ workspace = true
|
||||
|
||||
[dependencies]
|
||||
anyhow.workspace = true
|
||||
bytes.workspace = true
|
||||
brush-builtins.workspace = true
|
||||
brush-core.workspace = true
|
||||
brush-parser.workspace = true
|
||||
|
||||
@@ -8,5 +8,5 @@ pub mod windows;
|
||||
pub use brush_core::commands::{ChildSessionAction, child_session_action};
|
||||
pub use shell::{
|
||||
MinimizerResult, Shell, ShellExecuteOptions, ShellExecuteResult, ShellOptions, ShellRunOptions,
|
||||
ShellRunResult, execute_shell,
|
||||
ShellRunResult, StreamSinks, execute_shell, execute_shell_streams,
|
||||
};
|
||||
|
||||
@@ -12,6 +12,7 @@ use std::{
|
||||
};
|
||||
|
||||
use anyhow::{Error, Result};
|
||||
use bytes::Bytes;
|
||||
use brush_builtins::{BuiltinSet, default_builtins};
|
||||
use brush_core::{
|
||||
ExecutionContext, ExecutionControlFlow, ExecutionExitCode, ExecutionResult, ProcessGroupPolicy,
|
||||
@@ -196,6 +197,40 @@ pub async fn execute_shell(
|
||||
run_shell_oneshot(config, run_config, on_chunk, cancel_token).await
|
||||
}
|
||||
|
||||
/// Optional per-stream raw byte sinks for [`execute_shell_streams`].
|
||||
///
|
||||
/// When a sink is `Some`, that stream's pipe is drained directly into the
|
||||
/// channel with no UTF-8 decoding and no merging. When `None`, the
|
||||
/// corresponding pipe is still drained (to avoid blocking the child) but
|
||||
/// its bytes are dropped.
|
||||
#[derive(Default)]
|
||||
pub struct StreamSinks {
|
||||
pub stdout: Option<mpsc::UnboundedSender<Bytes>>,
|
||||
pub stderr: Option<mpsc::UnboundedSender<Bytes>>,
|
||||
}
|
||||
|
||||
/// One-shot execution that delivers stdout and stderr as raw byte chunks on
|
||||
/// separate channels. The minimizer is intentionally disabled — its
|
||||
/// `MinimizerResult.text` contract presumes a single merged transcript.
|
||||
pub async fn execute_shell_streams(
|
||||
options: ShellExecuteOptions,
|
||||
streams: StreamSinks,
|
||||
cancel_token: CancelToken,
|
||||
) -> Result<ShellExecuteResult> {
|
||||
let config = ShellConfig {
|
||||
session_env: options.session_env,
|
||||
snapshot_path: options.snapshot_path,
|
||||
minimizer: None,
|
||||
};
|
||||
let run_config = ShellRunConfig {
|
||||
command: options.command,
|
||||
cwd: options.cwd,
|
||||
env: options.env,
|
||||
minimizer: None,
|
||||
};
|
||||
run_shell_oneshot_streams(config, run_config, streams, cancel_token).await
|
||||
}
|
||||
|
||||
async fn run_shell_session(
|
||||
session: Arc<TokioMutex<Option<ShellSessionCore>>>,
|
||||
abort_state: ShellAbortState,
|
||||
@@ -309,6 +344,51 @@ async fn run_shell_oneshot(
|
||||
})
|
||||
}
|
||||
|
||||
async fn run_shell_oneshot_streams(
|
||||
config: ShellConfig,
|
||||
run_config: ShellRunConfig,
|
||||
streams: StreamSinks,
|
||||
ct: CancelToken,
|
||||
) -> Result<ShellExecuteResult> {
|
||||
let tokio_cancel = CancellationToken::new();
|
||||
|
||||
let mut task = tokio::spawn({
|
||||
let tokio_cancel = tokio_cancel.clone();
|
||||
async move {
|
||||
let mut session = create_session(&config).await?;
|
||||
run_shell_command_streams(&mut session, &run_config, streams, tokio_cancel).await
|
||||
}
|
||||
});
|
||||
|
||||
let run_result = tokio::select! {
|
||||
result = &mut task => result,
|
||||
reason = ct.wait() => {
|
||||
tokio_cancel.cancel();
|
||||
let graceful = time::timeout(Duration::from_secs(2), &mut task).await;
|
||||
if graceful.is_err() {
|
||||
task.abort();
|
||||
let _ = task.await;
|
||||
}
|
||||
return Ok(ShellExecuteResult {
|
||||
exit_code: None,
|
||||
cancelled: matches!(reason, AbortReason::Signal),
|
||||
timed_out: matches!(reason, AbortReason::Timeout),
|
||||
minimized: None,
|
||||
});
|
||||
},
|
||||
};
|
||||
|
||||
let res = run_result
|
||||
.unwrap_or_else(|err| Err(Error::msg(format!("Shell execution task failed: {err}"))));
|
||||
let exec = res?;
|
||||
Ok(ShellExecuteResult {
|
||||
exit_code: Some(exit_code(&exec)),
|
||||
cancelled: false,
|
||||
timed_out: false,
|
||||
minimized: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn null_file() -> Result<OpenFile> {
|
||||
openfiles::null().map_err(|err| Error::msg(format!("Failed to create null file: {err}")))
|
||||
}
|
||||
@@ -729,6 +809,259 @@ async fn run_shell_command(
|
||||
Ok((result, minimized_out))
|
||||
}
|
||||
|
||||
async fn run_shell_command_streams(
|
||||
session: &mut ShellSessionCore,
|
||||
options: &ShellRunConfig,
|
||||
streams: StreamSinks,
|
||||
cancel_token: CancellationToken,
|
||||
) -> Result<ExecutionResult> {
|
||||
if let Some(cwd) = options.cwd.as_deref() {
|
||||
session
|
||||
.shell
|
||||
.set_working_dir(cwd)
|
||||
.map_err(|err| Error::msg(format!("Failed to set cwd: {err}")))?;
|
||||
}
|
||||
|
||||
let mut env_scope_pushed = false;
|
||||
if let Some(env) = options.env.as_ref() {
|
||||
session
|
||||
.shell
|
||||
.env_mut()
|
||||
.push_scope(EnvironmentScope::Command);
|
||||
env_scope_pushed = true;
|
||||
for (key, value) in env {
|
||||
let normalized_key = normalize_env_key(key);
|
||||
if should_skip_env_var(normalized_key) {
|
||||
continue;
|
||||
}
|
||||
let mut var = ShellVariable::new(ShellValue::String(value.clone()));
|
||||
var.export();
|
||||
if let Err(err) =
|
||||
session
|
||||
.shell
|
||||
.env_mut()
|
||||
.add(normalized_key, var, EnvironmentScope::Command)
|
||||
{
|
||||
let _ = session.shell.env_mut().pop_scope(EnvironmentScope::Command);
|
||||
return Err(Error::msg(format!("Failed to set env: {err}")));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let (stdout_reader, stdout_writer) = pipe_to_files("stdout")?;
|
||||
let (stderr_reader, stderr_writer) = pipe_to_files("stderr")?;
|
||||
|
||||
let stdout_file = OpenFile::from(stdout_writer);
|
||||
let stderr_file = OpenFile::from(stderr_writer);
|
||||
|
||||
let mut params = session.shell.default_exec_params();
|
||||
params.set_fd(OpenFiles::STDIN_FD, null_file()?);
|
||||
params.set_fd(OpenFiles::STDOUT_FD, stdout_file);
|
||||
params.set_fd(OpenFiles::STDERR_FD, stderr_file);
|
||||
params.process_group_policy = ProcessGroupPolicy::NewProcessGroup;
|
||||
params.set_cancel_token(cancel_token.clone());
|
||||
let baseline_descendants = process::current_descendant_pids();
|
||||
let reader_cancel = CancellationToken::new();
|
||||
let (activity_tx, mut activity_rx) = mpsc::channel::<()>(1);
|
||||
|
||||
let StreamSinks { stdout: stdout_sink, stderr: stderr_sink } = streams;
|
||||
let mut stdout_handle = tokio::spawn(Box::pin(read_output_bytes(
|
||||
stdout_reader,
|
||||
stdout_sink,
|
||||
reader_cancel.clone(),
|
||||
activity_tx.clone(),
|
||||
)));
|
||||
let mut stderr_handle = tokio::spawn(Box::pin(read_output_bytes(
|
||||
stderr_reader,
|
||||
stderr_sink,
|
||||
reader_cancel.clone(),
|
||||
activity_tx,
|
||||
)));
|
||||
|
||||
let cancel_bridge = tokio::spawn({
|
||||
let cancel_token = cancel_token.clone();
|
||||
let reader_cancel = reader_cancel.clone();
|
||||
async move {
|
||||
cancel_token.cancelled().await;
|
||||
reader_cancel.cancel();
|
||||
}
|
||||
});
|
||||
let process_cancel_bridge = tokio::spawn({
|
||||
let cancel_token = cancel_token.clone();
|
||||
let baseline_descendants = baseline_descendants.clone();
|
||||
async move {
|
||||
cancel_token.cancelled().await;
|
||||
const WAVES: u32 = 3;
|
||||
for wave in 0..WAVES {
|
||||
let mut targets = process::TerminationTargets::new();
|
||||
process::add_new_descendants(&mut targets, &baseline_descendants);
|
||||
if targets.is_empty() {
|
||||
return;
|
||||
}
|
||||
let signal = if wave == 0 {
|
||||
process::TERM_SIGNAL
|
||||
} else {
|
||||
process::KILL_SIGNAL
|
||||
};
|
||||
targets.signal(signal);
|
||||
if wave + 1 < WAVES {
|
||||
let pause = if wave == 0 {
|
||||
Duration::from_millis(75)
|
||||
} else {
|
||||
Duration::from_millis(150)
|
||||
};
|
||||
time::sleep(pause).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
let source_info = SourceInfo::from("pi-shell:streams");
|
||||
let result = session
|
||||
.shell
|
||||
.run_string(options.command.clone(), &source_info, ¶ms)
|
||||
.await;
|
||||
|
||||
if cancel_token.is_cancelled() {
|
||||
terminate_background_jobs(&session.shell);
|
||||
}
|
||||
|
||||
if env_scope_pushed {
|
||||
session
|
||||
.shell
|
||||
.env_mut()
|
||||
.pop_scope(EnvironmentScope::Command)
|
||||
.map_err(|err| Error::msg(format!("Failed to pop env scope: {err}")))?;
|
||||
}
|
||||
|
||||
drop(params);
|
||||
|
||||
const POST_EXIT_IDLE: Duration = Duration::from_millis(250);
|
||||
const POST_EXIT_MAX: Duration = Duration::from_secs(2);
|
||||
const READER_SHUTDOWN_TIMEOUT: Duration = Duration::from_millis(250);
|
||||
|
||||
let mut stdout_finished = false;
|
||||
let mut stderr_finished = false;
|
||||
let mut idle_timer = Box::pin(time::sleep(POST_EXIT_IDLE));
|
||||
let mut max_timer = Box::pin(time::sleep(POST_EXIT_MAX));
|
||||
|
||||
loop {
|
||||
if stdout_finished && stderr_finished {
|
||||
break;
|
||||
}
|
||||
tokio::select! {
|
||||
res = &mut stdout_handle, if !stdout_finished => {
|
||||
let _ = res;
|
||||
stdout_finished = true;
|
||||
}
|
||||
res = &mut stderr_handle, if !stderr_finished => {
|
||||
let _ = res;
|
||||
stderr_finished = true;
|
||||
}
|
||||
msg = activity_rx.recv() => {
|
||||
if msg.is_none() {
|
||||
break;
|
||||
}
|
||||
idle_timer.as_mut().reset(time::Instant::now() + POST_EXIT_IDLE);
|
||||
}
|
||||
() = &mut idle_timer => break,
|
||||
() = &mut max_timer => break,
|
||||
}
|
||||
}
|
||||
|
||||
if !stdout_finished || !stderr_finished {
|
||||
reader_cancel.cancel();
|
||||
}
|
||||
if !stdout_finished {
|
||||
if time::timeout(READER_SHUTDOWN_TIMEOUT, &mut stdout_handle)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
stdout_handle.abort();
|
||||
let _ = stdout_handle.await;
|
||||
}
|
||||
}
|
||||
if !stderr_finished {
|
||||
if time::timeout(READER_SHUTDOWN_TIMEOUT, &mut stderr_handle)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
stderr_handle.abort();
|
||||
let _ = stderr_handle.await;
|
||||
}
|
||||
}
|
||||
cancel_bridge.abort();
|
||||
let _ = cancel_bridge.await;
|
||||
if cancel_token.is_cancelled() {
|
||||
// Let the kill-wave bridge finish all three signal passes so stragglers
|
||||
// have a chance to receive SIGKILL.
|
||||
let _ = process_cancel_bridge.await;
|
||||
} else {
|
||||
process_cancel_bridge.abort();
|
||||
let _ = process_cancel_bridge.await;
|
||||
}
|
||||
|
||||
let result = result.map_err(|err| Error::msg(format!("Shell execution failed: {err}")))?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn read_output_bytes(
|
||||
reader: fs::File,
|
||||
sink: Option<mpsc::UnboundedSender<Bytes>>,
|
||||
cancel_token: CancellationToken,
|
||||
activity: mpsc::Sender<()>,
|
||||
) {
|
||||
const BUF: usize = 65536;
|
||||
|
||||
#[cfg(unix)]
|
||||
let Ok(reader) = register_nonblocking_pipe(reader) else {
|
||||
return;
|
||||
};
|
||||
#[cfg(not(unix))]
|
||||
let mut reader = tokio::fs::File::from_std(reader);
|
||||
|
||||
loop {
|
||||
let mut buf = vec![0u8; BUF];
|
||||
#[cfg(unix)]
|
||||
let n = {
|
||||
let Ok(mut readiness) = (tokio::select! {
|
||||
ready = reader.readable() => ready,
|
||||
() = cancel_token.cancelled() => break,
|
||||
}) else {
|
||||
break;
|
||||
};
|
||||
match readiness.try_io(|inner| read_nonblocking(inner.get_ref(), &mut buf)) {
|
||||
Ok(Ok(0)) => break,
|
||||
Ok(Ok(n)) => n,
|
||||
Ok(Err(e)) if e.kind() == io::ErrorKind::Interrupted => continue,
|
||||
Ok(Err(_)) => break,
|
||||
Err(_would_block) => continue,
|
||||
}
|
||||
};
|
||||
#[cfg(not(unix))]
|
||||
let n = {
|
||||
let read_future = reader.read(&mut buf);
|
||||
tokio::pin!(read_future);
|
||||
match tokio::select! {
|
||||
res = &mut read_future => res,
|
||||
() = cancel_token.cancelled() => break,
|
||||
} {
|
||||
Ok(0) => break,
|
||||
Ok(n) => n,
|
||||
Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
|
||||
Err(_) => break,
|
||||
}
|
||||
};
|
||||
let _ = activity.try_send(());
|
||||
buf.truncate(n);
|
||||
if let Some(sink) = sink.as_ref()
|
||||
&& sink.send(Bytes::from(buf)).is_err()
|
||||
{
|
||||
// Receiver dropped — stop forwarding and let the pipe close.
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn terminate_background_jobs(shell: &BrushShell) {
|
||||
let mut targets = process::TerminationTargets::new();
|
||||
for job in &shell.jobs().jobs {
|
||||
@@ -1505,4 +1838,46 @@ mod tests {
|
||||
.expect("reader task should stop after cancellation")
|
||||
.expect("reader task should not panic");
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn execute_shell_streams_separates_stdout_and_stderr() {
|
||||
let (stdout_tx, mut stdout_rx) = mpsc::unbounded_channel::<Bytes>();
|
||||
let (stderr_tx, mut stderr_rx) = mpsc::unbounded_channel::<Bytes>();
|
||||
let options = ShellExecuteOptions {
|
||||
command: "echo out; echo err 1>&2".to_string(),
|
||||
..Default::default()
|
||||
};
|
||||
let streams = StreamSinks { stdout: Some(stdout_tx), stderr: Some(stderr_tx) };
|
||||
let result = execute_shell_streams(options, streams, CancelToken::default())
|
||||
.await
|
||||
.expect("execute should succeed");
|
||||
assert_eq!(result.exit_code, Some(0));
|
||||
assert!(!result.cancelled);
|
||||
|
||||
let mut stdout = Vec::new();
|
||||
while let Some(chunk) = stdout_rx.recv().await {
|
||||
stdout.extend_from_slice(&chunk);
|
||||
}
|
||||
let mut stderr = Vec::new();
|
||||
while let Some(chunk) = stderr_rx.recv().await {
|
||||
stderr.extend_from_slice(&chunk);
|
||||
}
|
||||
assert_eq!(stdout, b"out\n");
|
||||
assert_eq!(stderr, b"err\n");
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn execute_shell_streams_works_when_sinks_are_none() {
|
||||
// Both sinks `None` — pipes must still drain so the child can exit.
|
||||
let options = ShellExecuteOptions {
|
||||
command: "yes done | head -n 100 1>&2; echo final".to_string(),
|
||||
..Default::default()
|
||||
};
|
||||
let result = execute_shell_streams(options, StreamSinks::default(), CancelToken::default())
|
||||
.await
|
||||
.expect("execute should succeed");
|
||||
assert_eq!(result.exit_code, Some(0));
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user