From 9d5508980ffc2249581641f2a44b1d1252992a42 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 1 Apr 2026 20:46:21 +0200 Subject: [PATCH] perf(pi-natives): optimized pipe I/O to use async instead of blocking workers - Replaced blocking Tokio worker with non-blocking async I/O for pipe reading to prevent stalling the runtime when child processes keep pipes open. - Refactored background job termination to use process group and process tree APIs instead of deprecated signal parsing, improving reliability across Unix and Windows. - Added non-blocking pipe registration and fcntl-based flag management for Unix platforms to enable efficient async I/O without blocking workers. - Added test case verifying read_output cancellation before pipe EOF to ensure proper cleanup behavior. --- crates/brush-core-vendored/src/commands.rs | 65 ++++++++- crates/pi-natives/src/shell.rs | 161 +++++++++++++++++---- packages/natives/test/native.test.ts | 24 +++ 3 files changed, 214 insertions(+), 36 deletions(-) diff --git a/crates/brush-core-vendored/src/commands.rs b/crates/brush-core-vendored/src/commands.rs index 54d01fd1d..3a4a6e00a 100644 --- a/crates/brush-core-vendored/src/commands.rs +++ b/crates/brush-core-vendored/src/commands.rs @@ -560,13 +560,14 @@ pub(crate) async fn invoke_command_in_subshell_and_get_output( rt.block_on(run_substitution_command(subshell, params, s)) }); - // Read subshell output on a blocking thread to avoid stalling the - // async runtime when the pipe stays open (e.g. a hung child process). - let output_join_handle = tokio::task::spawn_blocking(move || io::read_to_string(reader)); + // Read subshell output without pinning a Tokio blocking worker for the + // lifetime of the pipe. Detached descendants can keep the pipe open long + // after the direct child exits. + let output_future = read_pipe_to_string(reader); // Wait for both the output reader and the command to complete. - let (output_result, cmd_result) = tokio::join!(output_join_handle, cmd_join_handle); - let output_str = output_result.map_err(io::Error::other)??; + let (output_result, cmd_result) = tokio::join!(output_future, cmd_join_handle); + let output_str = output_result?; let cmd_result = cmd_result.map_err(io::Error::other)??; // Store the status. @@ -602,6 +603,60 @@ async fn run_substitution_command( .await } +#[cfg(unix)] +async fn read_pipe_to_string(reader: io::PipeReader) -> io::Result { + let reader = register_nonblocking_pipe(reader)?; + let mut output = Vec::new(); + let mut buf = [0u8; 8192]; + + loop { + let mut readiness = reader.readable().await?; + let n = match readiness.try_io(|inner| read_nonblocking(inner.get_ref(), &mut buf)) { + Ok(Ok(0)) => break, + Ok(Ok(n)) => n, + Ok(Err(err)) if err.kind() == io::ErrorKind::Interrupted => continue, + Ok(Err(err)) => return Err(err), + Err(_would_block) => continue, + }; + output.extend_from_slice(&buf[..n]); + } + + String::from_utf8(output).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err)) +} + +#[cfg(not(unix))] +async fn read_pipe_to_string(reader: io::PipeReader) -> io::Result { + tokio::task::spawn_blocking(move || io::read_to_string(reader)) + .await + .map_err(io::Error::other)? +} + +#[cfg(unix)] +fn register_nonblocking_pipe( + reader: io::PipeReader, +) -> io::Result> { + set_nonblocking(&reader)?; + tokio::io::unix::AsyncFd::new(reader) +} + +#[cfg(unix)] +fn set_nonblocking(file: &T) -> io::Result<()> { + let flags = nix::fcntl::fcntl(file, nix::fcntl::FcntlArg::F_GETFL).map_err(io::Error::from)?; + let flags = nix::fcntl::OFlag::from_bits_truncate(flags); + if flags.contains(nix::fcntl::OFlag::O_NONBLOCK) { + return Ok(()); + } + + nix::fcntl::fcntl(file, nix::fcntl::FcntlArg::F_SETFL(flags | nix::fcntl::OFlag::O_NONBLOCK)) + .map_err(io::Error::from)?; + Ok(()) +} + +#[cfg(unix)] +fn read_nonblocking(file: &T, buf: &mut [u8]) -> io::Result { + nix::unistd::read(file, buf).map_err(io::Error::from) +} + // Detects a subshell command that consists solely of a single input redirection // (e.g., "< file"), returning the IoRedirect when present. fn try_unwrap_bare_input_redir_program(program: &ast::Program) -> Option<&ast::IoRedirect> { diff --git a/crates/pi-natives/src/shell.rs b/crates/pi-natives/src/shell.rs index 3c0ee4991..d300bbc1f 100644 --- a/crates/pi-natives/src/shell.rs +++ b/crates/pi-natives/src/shell.rs @@ -32,7 +32,6 @@ use brush_core::{ ProcessGroupPolicy, Shell as BrushShell, ShellValue, ShellVariable, builtins, env::EnvironmentScope, openfiles::{self, OpenFile, OpenFiles}, - sys, traps, }; use clap::Parser; use napi::{ @@ -45,6 +44,7 @@ use napi::{ }, }; use napi_derive::napi; +#[cfg(not(unix))] use tokio::io::AsyncReadExt as _; use tokio_util::sync::CancellationToken; #[cfg(windows)] @@ -52,6 +52,9 @@ use windows::configure_windows_path; use crate::task; +const TERM_SIGNAL: i32 = 15; +const KILL_SIGNAL: i32 = 9; + struct ShellSessionCore { shell: BrushShell, } @@ -662,26 +665,38 @@ fn terminate_background_jobs(shell: &BrushShell) { if shell.jobs.jobs.is_empty() { return; } - let Ok(signal) = "TERM".parse::() else { - return; - }; let mut pgids = Vec::new(); + let mut pids = Vec::new(); for job in &shell.jobs.jobs { - if let Some(pid) = job.process_group_id().or_else(|| job.representative_pid()) { - let _ = sys::signal::kill_process(pid, signal); - pgids.push(pid); + if let Some(pgid) = job.process_group_id() + && !pgids.contains(&pgid) + { + pgids.push(pgid); + } + if let Some(pid) = job.representative_pid() + && !pids.contains(&pid) + { + pids.push(pid); } } - if pgids.is_empty() { + if pgids.is_empty() && pids.is_empty() { return; } + + for &pgid in &pgids { + let _ = crate::ps::kill_process_group(pgid, TERM_SIGNAL); + } + for &pid in &pids { + let _ = crate::ps::kill_tree(pid, TERM_SIGNAL); + } + tokio::spawn(async move { time::sleep(Duration::from_millis(500)).await; - let Ok(signal) = "KILL".parse::() else { - return; - }; for pid in pgids { - let _ = sys::signal::kill_process(pid, signal); + let _ = crate::ps::kill_process_group(pid, KILL_SIGNAL); + } + for pid in pids { + let _ = crate::ps::kill_tree(pid, KILL_SIGNAL); } }); } @@ -691,26 +706,26 @@ fn terminate_background_jobs(shell: &BrushShell) { if shell.jobs.jobs.is_empty() { return; } - let Ok(signal) = "TERM".parse::() else { - return; - }; let mut pids = Vec::new(); for job in &shell.jobs.jobs { - if let Some(pid) = job.process_group_id().or_else(|| job.representative_pid()) { - let _ = sys::signal::kill_process(pid, signal); + if let Some(pid) = job.representative_pid() + && !pids.contains(&pid) + { pids.push(pid); } } if pids.is_empty() { return; } + + for &pid in &pids { + let _ = crate::ps::kill_tree(pid, TERM_SIGNAL); + } + tokio::spawn(async move { time::sleep(Duration::from_millis(500)).await; - let Ok(signal) = "KILL".parse::() else { - return; - }; for pid in pids { - let _ = sys::signal::kill_process(pid, signal); + let _ = crate::ps::kill_tree(pid, KILL_SIGNAL); } }); } @@ -794,20 +809,47 @@ async fn read_output( let mut buf = vec![0u8; BUF + 4]; // +4 for max UTF-8 char let mut it = 0; + #[cfg(unix)] + let reader = match register_nonblocking_pipe(reader) { + Ok(reader) => reader, + Err(_) => return, + }; + #[cfg(not(unix))] let reader = tokio::fs::File::from_std(reader); + #[cfg(not(unix))] tokio::pin!(reader); loop { - let read_future = reader.read(&mut buf[it..BUF]); - tokio::pin!(read_future); - let n = match tokio::select! { - res = &mut read_future => res, - () = cancel_token.cancelled() => break, - } { - Ok(0) => break, // EOF - Ok(n) => n, - Err(e) if e.kind() == io::ErrorKind::Interrupted => continue, - Err(_) => break, + #[cfg(unix)] + let n = { + let mut readiness = match tokio::select! { + ready = reader.readable() => ready, + () = cancel_token.cancelled() => break, + } { + Ok(readiness) => readiness, + Err(_) => break, + }; + match readiness.try_io(|inner| read_nonblocking(inner.get_ref(), &mut buf[it..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[it..BUF]); + tokio::pin!(read_future); + match tokio::select! { + res = &mut read_future => res, + () = cancel_token.cancelled() => break, + } { + Ok(0) => break, // EOF + Ok(n) => n, + Err(e) if e.kind() == io::ErrorKind::Interrupted => continue, + Err(_) => break, + } }; if n > 0 { let _ = activity.try_send(()); @@ -866,6 +908,46 @@ async fn read_output( } } +#[cfg(unix)] +fn register_nonblocking_pipe(reader: fs::File) -> io::Result> { + set_nonblocking(&reader)?; + tokio::io::unix::AsyncFd::new(reader) +} + +#[cfg(unix)] +fn set_nonblocking(file: &T) -> io::Result<()> { + let fd = file.as_raw_fd(); + // SAFETY: `fd` is owned by `file` and remains valid for the duration of + // these `fcntl` calls. + let flags = unsafe { libc::fcntl(fd, libc::F_GETFL) }; + if flags < 0 { + return Err(io::Error::last_os_error()); + } + if flags & libc::O_NONBLOCK != 0 { + return Ok(()); + } + + // SAFETY: `fd` remains valid here and we are only toggling `O_NONBLOCK`. + let result = unsafe { libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) }; + if result < 0 { + Err(io::Error::last_os_error()) + } else { + Ok(()) + } +} + +#[cfg(unix)] +fn read_nonblocking(file: &T, buf: &mut [u8]) -> io::Result { + // SAFETY: `buf` is writable for `buf.len()` bytes, and the raw fd obtained + // from `file` stays valid for the duration of the syscall. + let read = unsafe { libc::read(file.as_raw_fd(), buf.as_mut_ptr().cast(), buf.len()) }; + if read < 0 { + Err(io::Error::last_os_error()) + } else { + Ok(read as usize) + } +} + fn emit_chunk(text: &str, callback: Option<&ThreadsafeFunction>) { if let Some(callback) = callback { callback.call(Ok(text.to_string()), ThreadsafeFunctionCallMode::NonBlocking); @@ -1071,4 +1153,21 @@ mod tests { .expect("cancel token should be signalled"); assert!(matches!(reason, task::AbortReason::Signal)); } + + #[cfg(unix)] + #[tokio::test] + async fn read_output_stops_when_cancelled_before_pipe_eof() { + let (reader, _writer) = pipe_to_files("test").expect("test pipe should be created"); + let cancel = CancellationToken::new(); + let (activity_tx, _activity_rx) = mpsc::channel(1); + let handle = tokio::spawn(read_output(reader, None, cancel.clone(), activity_tx)); + + time::sleep(Duration::from_millis(10)).await; + cancel.cancel(); + + time::timeout(Duration::from_millis(100), handle) + .await + .expect("reader task should stop after cancellation") + .expect("reader task should not panic"); + } } diff --git a/packages/natives/test/native.test.ts b/packages/natives/test/native.test.ts index 94aaf1297..e9305183d 100644 --- a/packages/natives/test/native.test.ts +++ b/packages/natives/test/native.test.ts @@ -3,6 +3,7 @@ import * as fs from "node:fs/promises"; import * as os from "node:os"; import * as path from "node:path"; import { + executeShell, FileType, fuzzyFind, type GlobMatch, @@ -385,6 +386,29 @@ describe("pi-natives", () => { } }); }); + + describe("shell", () => { + it("should time out background workloads without leaving delayed writers behind", async () => { + if (process.platform === "win32") { + return; + } + + const markerPath = path.join(testDir, "shell-timeout-marker.txt"); + const markerEscaped = markerPath.replace(/'/g, "'\\''"); + await fs.rm(markerPath, { force: true }); + + const result = await executeShell({ + command: `{ sleep 2; echo done > '${markerEscaped}'; } & sleep 10`, + cwd: testDir, + timeoutMs: 100, + }); + + expect(result.timedOut).toBe(true); + + await Bun.sleep(3000); + expect(await Bun.file(markerPath).exists()).toBe(false); + }); + }); describe("htmlToMarkdown", () => { it("should convert basic HTML to markdown", async () => { const html = "

Hello World

This is a paragraph.

";