diff --git a/codex-rs/file-search/src/lib.rs b/codex-rs/file-search/src/lib.rs index faf96acd..8f7bce3e 100644 --- a/codex-rs/file-search/src/lib.rs +++ b/codex-rs/file-search/src/lib.rs @@ -11,6 +11,8 @@ use std::cmp::Reverse; use std::collections::BinaryHeap; use std::num::NonZero; use std::path::Path; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use tokio::process::Command; @@ -71,10 +73,18 @@ pub async fn run_main( } }; + let cancel_flag = Arc::new(AtomicBool::new(false)); let FileSearchResults { total_match_count, matches, - } = run(&pattern_text, limit, &search_directory, exclude, threads).await?; + } = run( + &pattern_text, + limit, + &search_directory, + exclude, + threads, + cancel_flag, + )?; let match_count = matches.len(); let matches_truncated = total_match_count > match_count; @@ -88,12 +98,15 @@ pub async fn run_main( Ok(()) } -pub async fn run( +/// The worker threads will periodically check `cancel_flag` to see if they +/// should stop processing files. +pub fn run( pattern_text: &str, limit: NonZero, search_directory: &Path, exclude: Vec, threads: NonZero, + cancel_flag: Arc, ) -> anyhow::Result { let pattern = create_pattern(pattern_text); // Create one BestMatchesList per worker thread so that each worker can @@ -136,11 +149,25 @@ pub async fn run( let index = index_counter.fetch_add(1, Ordering::Relaxed); let best_list_ptr = best_matchers_per_worker[index].get(); let best_list = unsafe { &mut *best_list_ptr }; + + // Each worker keeps a local counter so we only read the atomic flag + // every N entries which is cheaper than checking on every file. + const CHECK_INTERVAL: usize = 1024; + let mut processed = 0; + + let cancel = cancel_flag.clone(); + Box::new(move |entry| { if let Some(path) = get_file_path(&entry, search_directory) { best_list.insert(path); } - ignore::WalkState::Continue + + processed += 1; + if processed % CHECK_INTERVAL == 0 && cancel.load(Ordering::Relaxed) { + ignore::WalkState::Quit + } else { + ignore::WalkState::Continue + } }) }); @@ -162,6 +189,14 @@ pub async fn run( } } + // If the cancel flag is set, we return early with an empty result. + if cancel_flag.load(Ordering::Relaxed) { + return Ok(FileSearchResults { + matches: Vec::new(), + total_match_count: 0, + }); + } + // Merge results across best_matchers_per_worker. let mut global_heap: BinaryHeap> = BinaryHeap::new(); let mut total_match_count = 0;