From ac20e1432d162a259ba379dff04df04aa25253a8 Mon Sep 17 00:00:00 2001 From: jazelly Date: Thu, 3 Sep 2026 19:22:25 +0930 Subject: [PATCH] fix(search): bound grep and rg output memory --- src/cmds/system/search.rs | 45 ++++--- src/core/stream.rs | 233 +++++++++++++++++++++++++++++----- tests/search_faithful_test.rs | 56 ++++++++ 3 files changed, 290 insertions(+), 44 deletions(-) diff --git a/src/cmds/system/search.rs b/src/cmds/system/search.rs index b526c05e12..af98c03158 100644 --- a/src/cmds/system/search.rs +++ b/src/cmds/system/search.rs @@ -7,10 +7,10 @@ //! `file:line:content` shape. use crate::core::stream::{ - self, exec_capture, exec_capture_stdin, CaptureResult, FilterMode, StdinMode, StreamFilter, + self, exec_capture, CaptureResult, FilterMode, StdinMode, StreamFilter, }; use crate::core::tracking; -use crate::core::utils::{resolved_command, strip_ansi}; +use crate::core::utils::resolved_command; use crate::core::{args_utils, config}; use anyhow::{Context, Result}; use regex::Regex; @@ -301,9 +301,9 @@ fn engine_capture>( extra_args: &[T], patterns: &[String], paths: &[String], -) -> Result { +) -> Result<(CaptureResult, bool)> { let mut cmd = engine_command(engine, extra_args, patterns, paths, false); - exec_capture_stdin(&mut cmd).context("search failed") + stream::exec_capture_stdin_bounded(&mut cmd).context("search failed") } fn engine_command>( @@ -347,6 +347,13 @@ fn format_match_line(line: &str, show_file: bool, show_line: bool) -> Option String { + format!( + "[rtk] output capped; rerun the same command with rtk proxy {} ... for full output", + engine.label() + ) +} + /// Emits each piped match as it arrives. Buffered search waits for EOF, so /// `tail -f app.log | rtk grep ERROR` would otherwise show no matches. struct SearchStreamFilter { @@ -460,12 +467,9 @@ fn passthrough>( .context("search failed")? .exit_code } else { - let result = exec_capture_stdin(&mut cmd).context("search failed")?; - print!("{}", strip_ansi(&result.stdout)); - if !result.stderr.is_empty() { - eprint!("{}", result.stderr); - } - result.exit_code + stream::run_streaming(&mut cmd, StdinMode::Null, FilterMode::Passthrough) + .context("search failed")? + .exit_code }; timer.track_passthrough(real_cmd, &format!("rtk {} (passthrough)", real_cmd)); @@ -562,14 +566,13 @@ pub fn run( ); } - let result = engine_capture(engine, &extra_args, &patterns, &paths)?; - + let (result, capture_capped) = engine_capture(engine, &extra_args, &patterns, &paths)?; let exit_code = result.exit_code; let raw_output = result.stdout.clone(); // Unparseable shape re-runs verbatim below (with its own stderr), so handle it // before surfacing this run's stderr (#2333). - if unparsed_signal(&raw_output) > 0 { + if !capture_capped && unparsed_signal(&raw_output) > 0 { return passthrough(&timer, engine, &args, &real_cmd, false); } @@ -577,7 +580,10 @@ pub fn run( eprint!("{}", result.stderr); } - if result.stdout.trim().is_empty() { + if raw_output.trim().is_empty() { + if capture_capped { + println!("{}", recovery_hint(engine)); + } timer.track(&real_cmd, &rtk_label, &raw_output, ""); return Ok(exit_code); } @@ -709,7 +715,16 @@ pub fn run( body }; - let output = if capped && rtk_output.len() < plain.len() { + let output = if capture_capped { + let mut output = if capped && rtk_output.len() < plain.len() { + rtk_output + } else { + plain + }; + output.push_str(&recovery_hint(engine)); + output.push('\n'); + output + } else if capped && rtk_output.len() < plain.len() { rtk_output } else { plain diff --git a/src/core/stream.rs b/src/core/stream.rs index 43598d7b8b..c42919a9c8 100644 --- a/src/core/stream.rs +++ b/src/core/stream.rs @@ -23,20 +23,81 @@ use regex::Regex; /// [`decode_process_output`](super::utils::decode_process_output) exists to /// read, so the streamed path decodes them the same way the captured path /// does rather than going straight to U+FFFD. -fn read_lines_lossy(reader: impl Read) -> impl Iterator { - BufReader::new(reader).split(b'\n').filter_map(|res| { - let mut buf = match res { - Ok(buf) => buf, - Err(e) => { - eprintln!("rtk: stream read error: {}", e); - return None; +struct BoundedLines { + reader: BufReader, + line: Vec, + truncated: bool, + warned: bool, + finished: bool, +} + +impl BoundedLines { + fn new(reader: R) -> Self { + Self { + reader: BufReader::new(reader), + line: Vec::new(), + truncated: false, + warned: false, + finished: false, + } + } + + fn finish_line(&mut self) -> String { + if self.line.last() == Some(&b'\r') { + self.line.pop(); + } + let text = super::utils::decode_process_output(&self.line); + if self.truncated && !self.warned { + self.warned = true; + eprintln!("[rtk] warning: input line exceeds 10 MiB — line truncated"); + } + self.line.clear(); + text + } +} + +impl Iterator for BoundedLines { + type Item = String; + + fn next(&mut self) -> Option { + if self.finished { + return None; + } + self.line.clear(); + self.truncated = false; + + loop { + let buf = match self.reader.fill_buf() { + Ok(buf) => buf, + Err(e) => { + eprintln!("rtk: stream read error: {}", e); + self.finished = true; + return None; + } + }; + if buf.is_empty() { + self.finished = true; + return (!self.line.is_empty()).then(|| self.finish_line()); + } + + let newline = buf.iter().position(|byte| *byte == b'\n'); + let content_len = newline.unwrap_or(buf.len()); + let remaining = RAW_CAP.saturating_sub(self.line.len()); + let take = content_len.min(remaining); + let truncated = content_len > remaining; + let consumed = newline.map_or(buf.len(), |index| index + 1); + self.line.extend_from_slice(&buf[..take]); + self.reader.consume(consumed); + self.truncated |= truncated; + if newline.is_some() { + return Some(self.finish_line()); } - }; - if buf.last() == Some(&b'\r') { - buf.pop(); } - Some(super::utils::decode_process_output(&buf)) - }) + } +} + +fn read_lines_lossy(reader: impl Read) -> impl Iterator { + BoundedLines::new(reader) } pub trait StreamFilter { @@ -285,6 +346,47 @@ pub fn status_to_exit_code(status: std::process::ExitStatus) -> i32 { // ISSUE #897: ChildGuard RAII prevents zombie processes that caused kernel panic pub const RAW_CAP: usize = 10_485_760; // 10 MiB +fn read_bounded_output(mut reader: impl Read) -> std::io::Result<(String, bool)> { + let mut bytes = Vec::new(); + let mut chunk = [0u8; 8192]; + let mut truncated = false; + loop { + let count = reader.read(&mut chunk)?; + if count == 0 { + break; + } + if bytes.len() < RAW_CAP { + let take = (RAW_CAP - bytes.len()).min(count); + bytes.extend_from_slice(&chunk[..take]); + truncated |= take < count; + } else { + truncated = true; + } + } + Ok((super::utils::decode_process_output(&bytes), truncated)) +} + +fn retain_filtered(filtered: &mut String, capped: &mut bool, output: &str) { + if *capped { + return; + } + let remaining = RAW_CAP.saturating_sub(filtered.len()); + if output.len() <= remaining { + filtered.push_str(output); + return; + } + + let end = output + .char_indices() + .map(|(index, ch)| index + ch.len_utf8()) + .take_while(|end| *end <= remaining) + .last() + .unwrap_or(0); + filtered.push_str(&output[..end]); + *capped = true; + eprintln!("[rtk] warning: filtered output exceeds 10 MiB — retention truncated"); +} + pub fn run_streaming( cmd: &mut Command, stdin_mode: StdinMode, @@ -364,6 +466,7 @@ pub fn run_streaming( let mut raw_stdout = String::new(); let mut raw_stderr = String::new(); let mut filtered = String::new(); + let mut filtered_capped = false; let mut capped_out = false; let mut capped_err = false; let mut saved_filter: Option> = None; @@ -375,7 +478,7 @@ pub fn run_streaming( Stderr(String), } - let (tx, rx) = mpsc::channel(); + let (tx, rx) = mpsc::sync_channel(1); let tx_out = tx.clone(); let stdout_thread = std::thread::spawn(move || { for line in read_lines_lossy(stdout) { @@ -399,7 +502,9 @@ pub fn run_streaming( let stderr_handle = io::stderr(); let mut err_out = stderr_handle.lock(); - for msg in rx { + let receiver = rx; + let mut broken_pipe = false; + while let Ok(msg) = receiver.recv() { let (line, is_stderr) = match msg { StreamLine::Stderr(l) => (l, true), StreamLine::Stdout(l) => (l, false), @@ -424,29 +529,42 @@ pub fn run_streaming( } } filter_fd_is_stderr = is_stderr; - if let Some(output) = filter.feed_line(&line) { - filtered.push_str(&output); + let output = filter.feed_line(&line); + if let Some(output) = output { + retain_filtered(&mut filtered, &mut filtered_capped, &output); let dest: &mut dyn Write = if is_stderr { &mut err_out } else { &mut out }; match write!(dest, "{}", output) { - Err(e) if e.kind() == io::ErrorKind::BrokenPipe => break, + Err(e) if e.kind() == io::ErrorKind::BrokenPipe => { + child.0.kill().ok(); + broken_pipe = true; + drop(receiver); + break; + } Err(e) => return Err(e.into()), Ok(_) => {} } } } - let tail = filter.flush(); - filtered.push_str(&tail); - let flush_dest: &mut dyn Write = if filter_fd_is_stderr { - &mut err_out - } else { - &mut out - }; - match write!(flush_dest, "{}", tail) { - Err(e) if e.kind() == io::ErrorKind::BrokenPipe => {} - Err(e) => return Err(e.into()), - Ok(_) => {} + if !broken_pipe { + let tail = filter.flush(); + retain_filtered(&mut filtered, &mut filtered_capped, &tail); + let flush_dest: &mut dyn Write = if filter_fd_is_stderr { + &mut err_out + } else { + &mut out + }; + match write!(flush_dest, "{}", tail) { + Err(e) if e.kind() == io::ErrorKind::BrokenPipe => { + child.0.kill().ok(); + broken_pipe = true; + } + Err(e) => return Err(e.into()), + Ok(_) => {} + } + } + if !broken_pipe { + saved_filter = Some(filter); } - saved_filter = Some(filter); } stdout_thread.join().ok(); @@ -580,6 +698,47 @@ pub fn exec_capture_stdin(cmd: &mut Command) -> Result { capture(cmd) } +/// Like exec_capture_stdin, but drains both pipes while retaining at most +/// RAW_CAP bytes per stream. Returns the output and whether either stream was capped. +pub fn exec_capture_stdin_bounded(cmd: &mut Command) -> Result<(CaptureResult, bool)> { + let program = cmd.get_program().to_string_lossy().into_owned(); + cmd.stdin(Stdio::inherit()); + cmd.stdout(Stdio::piped()); + cmd.stderr(Stdio::piped()); + let mut child = cmd.spawn().context("Failed to execute command")?; + let stdout = child.stdout.take().context("No child stdout handle")?; + let stderr = child.stderr.take().context("No child stderr handle")?; + let stdout_thread = std::thread::spawn(move || read_bounded_output(stdout)); + let stderr_thread = std::thread::spawn(move || read_bounded_output(stderr)); + let stdout_result = stdout_thread + .join() + .unwrap_or_else(|_| Err(std::io::Error::other("stdout reader thread panicked"))); + if stdout_result.is_err() { + child.kill().ok(); + } + let stderr_result = stderr_thread + .join() + .unwrap_or_else(|_| Err(std::io::Error::other("stderr reader thread panicked"))); + if stderr_result.is_err() { + child.kill().ok(); + } + let status = child.wait().context("Failed to wait for child")?; + let (stdout, stdout_capped) = stdout_result?; + let (stderr, stderr_capped) = stderr_result?; + let truncated = stdout_capped || stderr_capped; + if truncated { + eprintln!("[rtk] warning: captured output exceeds 10 MiB — capture truncated"); + } + Ok(( + CaptureResult { + stdout, + stderr, + exit_code: super::utils::exit_code_from_status(&status, &program), + }, + truncated, + )) +} + /// Run `cmd` to completion, decode what it wrote, and report the exit code. /// /// A process killed by a signal has no exit code of its own, and returning @@ -910,6 +1069,22 @@ pub(crate) mod tests { ); } + #[test] + fn test_retain_filtered_cap_at_10mb() { + let mut filtered = String::new(); + let mut capped = false; + let output = "x".repeat(RAW_CAP + 1); + retain_filtered(&mut filtered, &mut capped, &output); + assert!( + filtered.len() <= RAW_CAP, + "retained filtered output should be capped at 10 MiB, got {} bytes", + filtered.len() + ); + assert!(capped); + retain_filtered(&mut filtered, &mut capped, "more"); + assert_eq!(filtered.len(), RAW_CAP); + } + #[test] fn test_child_guard_prevents_zombie() { let mut cmd = Command::new("true"); diff --git a/tests/search_faithful_test.rs b/tests/search_faithful_test.rs index 08417a42bc..55bd2ac333 100644 --- a/tests/search_faithful_test.rs +++ b/tests/search_faithful_test.rs @@ -80,6 +80,62 @@ fn rg_available() -> bool { .unwrap_or(false) } +#[test] +fn rg_file_search_signals_truncated_long_line() { + if !rg_available() { + return; + } + let input_dir = tempfile::tempdir().expect("input dir"); + let input = input_dir.path().join("input.txt"); + let content = format!("MATCH {}\n", "x".repeat(10_485_760 + 1024)); + std::fs::write(&input, content).expect("write input"); + + let out = rtk() + .args(["rg", "MATCH", input.to_str().unwrap()]) + .output() + .expect("run rtk rg"); + let stdout = String::from_utf8_lossy(&out.stdout); + assert!( + stdout.contains("rtk proxy"), + "capped long search output must include a recovery hint" + ); +} + +#[test] +fn grep_exits_promptly_after_stdout_broken_pipe() { + let mut source = Command::new("yes") + .arg("foo") + .stdout(Stdio::piped()) + .spawn() + .expect("spawn yes"); + let source_stdout = source.stdout.take().expect("yes stdout"); + let mut child = rtk() + .args(["grep", "foo"]) + .stdin(source_stdout) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .expect("spawn rtk grep"); + let stdout = child.stdout.take().expect("rtk stdout"); + let mut reader = BufReader::new(stdout); + let mut first = String::new(); + reader.read_line(&mut first).expect("read first match"); + drop(reader); + + let pid = child.id().to_string(); + let (tx, rx) = mpsc::channel(); + std::thread::spawn(move || tx.send(child.wait()).ok()); + let status = rx.recv_timeout(Duration::from_secs(2)); + if status.is_err() { + Command::new("kill").args(["-KILL", &pid]).status().ok(); + source.kill().ok(); + source.wait().ok(); + panic!("rtk grep did not exit promptly after stdout closed"); + } + source.kill().ok(); + source.wait().ok(); +} + fn write_temp(content: &str) -> (tempfile::TempDir, std::path::PathBuf) { let dir = tempfile::tempdir().expect("tempdir"); let path = dir.path().join("test.txt");