fix(grep): enforce_time_budget option for zero-match searches (#826)

README.md

@@ -278,6 +278,7 @@ local r = require('fff').content_search('TODO', {
   page_size             = 50,
   file_offset           = 0,
   time_budget_ms        = 0,
+  enforce_time_budget   = false,    -- also bound zero-match searches
   trim_whitespace       = false,
   cwd                   = nil,      -- switch indexed root if different
   wait_for_index_ms     = nil,      -- override the default scan wait timeout
@@ -431,6 +432,7 @@ require('fff').setup({
     max_matches_per_file = 100,
     smart_case = true,
     time_budget_ms = 150,
+    enforce_time_budget = false, -- apply time_budget_ms even before anything matched (off = zero-match queries scan everything)
     modes = { 'plain', 'regex', 'fuzzy' },
     trim_whitespace = false,
     enable_filename_constraint = false, -- treat filename-like tokens (e.g. `score.rs`) in a grep query as a file-path filter scoping the search; off = searched as literal text

crates/fff-c/include/fff.h

@@ -636,6 +636,27 @@ struct FffResult *fff_live_grep(void *fff_handle,
                                 uint32_t after_context,
                                 bool classify_definitions);
 
+/**
+ * [`fff_live_grep`] plus `enforce_time_budget`: when true the budget also bounds
+ * zero-match searches and `next_file_offset` resumes at the first unsearched file.
+ *
+ * ## Safety
+ * Same as [`fff_live_grep`].
+ */
+struct FffResult *fff_live_grep_ex(void *fff_handle,
+                                   const char *query,
+                                   uint8_t mode,
+                                   uint64_t max_file_size,
+                                   uint32_t max_matches_per_file,
+                                   bool smart_case,
+                                   uint32_t file_offset,
+                                   uint32_t page_limit,
+                                   uint64_t time_budget_ms,
+                                   bool enforce_time_budget,
+                                   uint32_t before_context,
+                                   uint32_t after_context,
+                                   bool classify_definitions);
+
 /**
  * Multi-pattern OR search (SIMD Aho-Corasick): lines matching ANY pattern.
  *
@@ -660,6 +681,26 @@ struct FffResult *fff_multi_grep(void *fff_handle,
                                  uint32_t after_context,
                                  bool classify_definitions);
 
+/**
+ * [`fff_multi_grep`] plus `enforce_time_budget`, as in [`fff_live_grep_ex`].
+ *
+ * ## Safety
+ * Same as [`fff_multi_grep`].
+ */
+struct FffResult *fff_multi_grep_ex(void *fff_handle,
+                                    const char *patterns_joined,
+                                    const char *constraints,
+                                    uint64_t max_file_size,
+                                    uint32_t max_matches_per_file,
+                                    bool smart_case,
+                                    uint32_t file_offset,
+                                    uint32_t page_limit,
+                                    uint64_t time_budget_ms,
+                                    bool enforce_time_budget,
+                                    uint32_t before_context,
+                                    uint32_t after_context,
+                                    bool classify_definitions);
+
 /**
  * Trigger a rescan of the file index.
  *
@@ -1321,7 +1362,7 @@ struct FffResult *fff_watch_args(void *fff_handle,
 struct FffResult *fff_unwatch(void *fff_handle, uint64_t watch_id);
 
 /**
- * Number of events in a batch; 0 if `batch` is null.
+ * Number of events in a batch, 0 if `batch` is null.
  *
  * ## Safety
  * `batch` must be a valid `FffWatchEventBatch` pointer or null.

crates/fff-c/src/lib.rs

@@ -626,6 +626,46 @@ pub unsafe extern "C" fn fff_live_grep(
     before_context: u32,
     after_context: u32,
     classify_definitions: bool,
+) -> *mut FffResult {
+    unsafe {
+        fff_live_grep_ex(
+            fff_handle,
+            query,
+            mode,
+            max_file_size,
+            max_matches_per_file,
+            smart_case,
+            file_offset,
+            page_limit,
+            time_budget_ms,
+            false,
+            before_context,
+            after_context,
+            classify_definitions,
+        )
+    }
+}
+
+/// [`fff_live_grep`] plus `enforce_time_budget`: when true the budget also bounds
+/// zero-match searches and `next_file_offset` resumes at the first unsearched file.
+///
+/// ## Safety
+/// Same as [`fff_live_grep`].
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn fff_live_grep_ex(
+    fff_handle: *mut c_void,
+    query: *const c_char,
+    mode: u8,
+    max_file_size: u64,
+    max_matches_per_file: u32,
+    smart_case: bool,
+    file_offset: u32,
+    page_limit: u32,
+    time_budget_ms: u64,
+    enforce_time_budget: bool,
+    before_context: u32,
+    after_context: u32,
+    classify_definitions: bool,
 ) -> *mut FffResult {
     let inst = match unsafe { instance_ref(fff_handle) } {
         Ok(i) => i,
@@ -664,6 +704,7 @@ pub unsafe extern "C" fn fff_live_grep(
         page_limit: default_u32(page_limit, 50) as usize,
         mode: grep_mode_from_u8(mode),
         time_budget_ms,
+        enforce_time_budget,
         before_context: before_context as usize,
         after_context: after_context as usize,
         classify_definitions,
@@ -699,6 +740,45 @@ pub unsafe extern "C" fn fff_multi_grep(
     before_context: u32,
     after_context: u32,
     classify_definitions: bool,
+) -> *mut FffResult {
+    unsafe {
+        fff_multi_grep_ex(
+            fff_handle,
+            patterns_joined,
+            constraints,
+            max_file_size,
+            max_matches_per_file,
+            smart_case,
+            file_offset,
+            page_limit,
+            time_budget_ms,
+            false,
+            before_context,
+            after_context,
+            classify_definitions,
+        )
+    }
+}
+
+/// [`fff_multi_grep`] plus `enforce_time_budget`, as in [`fff_live_grep_ex`].
+///
+/// ## Safety
+/// Same as [`fff_multi_grep`].
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn fff_multi_grep_ex(
+    fff_handle: *mut c_void,
+    patterns_joined: *const c_char,
+    constraints: *const c_char,
+    max_file_size: u64,
+    max_matches_per_file: u32,
+    smart_case: bool,
+    file_offset: u32,
+    page_limit: u32,
+    time_budget_ms: u64,
+    enforce_time_budget: bool,
+    before_context: u32,
+    after_context: u32,
+    classify_definitions: bool,
 ) -> *mut FffResult {
     let inst = match unsafe { instance_ref(fff_handle) } {
         Ok(i) => i,
@@ -746,6 +826,7 @@ pub unsafe extern "C" fn fff_multi_grep(
         page_limit: default_u32(page_limit, 50) as usize,
         mode: fff::GrepMode::PlainText, // ignored by multi_grep_search
         time_budget_ms,
+        enforce_time_budget,
         before_context: before_context as usize,
         after_context: after_context as usize,
         classify_definitions,

crates/fff-core/src/grep/fuzzy_grep.rs

@@ -3,7 +3,7 @@ use crate::types::{ContentCacheBudget, FileItem, MmapSlot};
 use fff_grep::lines::LineStep;
 use rayon::prelude::*;
 use std::path::Path;
-use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 
 use super::sink::{
     char_indices_to_byte_offsets, classify_definition, strip_line_terminators,
@@ -107,7 +107,8 @@ pub(super) fn fuzzy_grep_search<'a>(
         None
     };
     let search_start = std::time::Instant::now();
-    let budget_exceeded = AtomicBool::new(false);
+    // Lowest index an abort skipped, i.e. the resume point; usize::MAX means no abort.
+    let first_skipped = AtomicUsize::new(usize::MAX);
     let max_matches_per_file = options.max_matches_per_file;
 
     // for fuzzy match we need a bit smarter chunking as the amount of work we have to perform is
@@ -153,15 +154,14 @@ pub(super) fn fuzzy_grep_search<'a>(
                     )
                 },
                 |(matcher, buf, mmap_slot), (local_idx, file)| {
-                    if abort_signal.load(Ordering::Relaxed) {
-                        budget_exceeded.store(true, Ordering::Relaxed);
-                        return None;
-                    }
-
-                    if let Some(budget) = time_budget
-                        && search_start.elapsed() > budget
+                    // File 0 is never skipped so a page always consumes at least one
+                    // file (cursor 0 means "done").
+                    let idx = chunk_offset + local_idx;
+                    if idx > 0
+                        && (abort_signal.load(Ordering::Relaxed)
+                            || time_budget.is_some_and(|b| search_start.elapsed() > b))
                     {
-                        budget_exceeded.store(true, Ordering::Relaxed);
+                        first_skipped.fetch_min(idx, Ordering::Relaxed);
                         return None;
                     }
 
@@ -331,7 +331,7 @@ pub(super) fn fuzzy_grep_search<'a>(
                         return None;
                     }
 
-                    Some((chunk_offset + local_idx, *file, file_matches))
+                    Some((idx, *file, file_matches))
                 },
             )
             .flatten()
@@ -342,17 +342,18 @@ pub(super) fn fuzzy_grep_search<'a>(
             per_file_results.push(result);
         }
 
-        if running_matches >= page_limit || budget_exceeded.load(Ordering::Relaxed) {
+        if running_matches >= page_limit || first_skipped.load(Ordering::Relaxed) != usize::MAX {
             break;
         }
     }
 
+    let abort_resume = Some(first_skipped.load(Ordering::Relaxed)).filter(|&i| i != usize::MAX);
     GrepResult::collect(
         per_file_results,
         files_to_search.len(),
         options,
         total_files,
         filtered_file_count,
-        budget_exceeded.load(Ordering::Relaxed),
+        abort_resume,
     )
 }

crates/fff-core/src/grep/grep.rs

@@ -16,7 +16,7 @@ use fff_query_parser::{FFFQuery, GrepConfig, QueryParser};
 use rayon::prelude::*;
 use smallvec::SmallVec;
 use std::path::Path;
-use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 use tracing::Level;
 
 #[allow(clippy::large_enum_variant)]
@@ -561,12 +561,15 @@ where
 
     let search_start = std::time::Instant::now();
     let page_limit = options.page_limit;
-    let budget_exceeded = AtomicBool::new(false);
+    // Lowest index an abort skipped; everything below it was searched, so it is
+    // the resume point for the next page. usize::MAX means no abort happened.
+    let first_skipped = AtomicUsize::new(usize::MAX);
 
     let mut result_files: Vec<&'a FileItem> = Vec::new();
     let mut all_matches: Vec<GrepMatch> = Vec::new();
     let mut files_consumed: usize = 0;
     let mut page_filled = false;
+    let mut aborted = false;
 
     // Each chunk is a rayon barrier. A flat small chunk over 500k files = ~7800
     // barriers; x2 growth makes it logarithmic. But a too-aggressive growth
@@ -593,30 +596,29 @@ where
         chunk_size = (chunk_size * growth).min(max_chunk);
         let chunk_offset = files_consumed;
 
+        // Unless enforced, the budget stays dormant until something matched.
+        let budget = time_budget.filter(|_| options.enforce_time_budget || all_matches.len() > 1);
+
         let chunk_results: Vec<(usize, &'a FileItem, Vec<GrepMatch>)> = chunk
             .par_iter()
             .enumerate()
             .map_init(
                 // tested it out a few times, this is just fine for rayon worker in this specific
                 // case it doesn't reallocate this many times and it is actually faster than using
                 // scoped threads with a predefined local scratch buffers because of spawn cost
-                || (Vec::with_capacity(64 * 1024), MmapSlot::default()),
-                |(buf, mmap_slot), (local_idx, file)| {
-                    // perform all the atomic machinery on every 8th
-                    if local_idx % 8 == 0 {
-                        let mut need_abort = ctx.abort_signal.load(Ordering::Relaxed);
-                        if !need_abort
-                            && let Some(budget) = time_budget
-                            && all_matches.len() > 1
-                            && search_start.elapsed() > budget
-                        {
-                            need_abort = true;
-                        }
-
-                        if need_abort {
-                            budget_exceeded.store(true, Ordering::Relaxed);
-                            return None;
-                        }
+                || (Vec::with_capacity(64 * 1024), MmapSlot::default(), false),
+                |(buf, mmap_slot, aborted), (local_idx, file)| {
+                    // check the clock on every 8th file only; `aborted` latches for the
+                    // rest of this worker's range. File 0 is never skipped so a page
+                    // always consumes at least one file (cursor 0 means "done").
+                    let idx = chunk_offset + local_idx;
+                    if !*aborted && local_idx % 8 == 0 && idx > 0 {
+                        *aborted = ctx.abort_signal.load(Ordering::Relaxed)
+                            || budget.is_some_and(|b| search_start.elapsed() > b);
+                    }
+                    if *aborted {
+                        first_skipped.fetch_min(idx, Ordering::Relaxed);
+                        return None;
                     }
 
                     let content = file.get_content_for_search(
@@ -642,17 +644,23 @@ where
                         return None;
                     }
 
-                    Some((chunk_offset + local_idx, *file, file_matches))
+                    Some((idx, *file, file_matches))
                 },
             )
             .flatten()
             .collect();
 
-        // Every file in the chunk was visited by rayon (matched or not).
-        files_consumed = chunk_offset + chunk.len();
+        // Every file in the chunk was visited unless an abort cut it short.
+        let resume_at = first_skipped.load(Ordering::Relaxed);
+        aborted = resume_at != usize::MAX;
+        files_consumed = resume_at.min(chunk_offset + chunk.len());
 
-        // Flatten this chunk's results into the accumulator.
         for (batch_idx, file, file_matches) in chunk_results {
+            // Matches past the abort point are dropped; the next page re-searches them.
+            if batch_idx >= resume_at {
+                continue;
+            }
+
             let file_result_idx = result_files.len();
             result_files.push(file);
 
@@ -673,18 +681,17 @@ where
             }
         }
 
-        if page_filled || budget_exceeded.load(Ordering::Relaxed) {
+        if page_filled || aborted {
             break;
         }
     }
 
-    // If no file had any match, we searched the entire slice.
-    if result_files.is_empty() {
+    // No match and no abort means the whole slice was searched.
+    if result_files.is_empty() && !aborted {
         files_consumed = files_to_search.len();
     }
 
-    let has_more = budget_exceeded.load(Ordering::Relaxed)
-        || (page_filled && files_consumed < files_to_search.len());
+    let has_more = (aborted || page_filled) && files_consumed < files_to_search.len();
 
     let next_file_offset = if has_more {
         options.file_offset + files_consumed