fix the search cutoff and add a test

milli/src/lib.rs

@@ -105,10 +105,15 @@ pub const MAX_WORD_LENGTH: usize = MAX_LMDB_KEY_LENGTH / 2;
 
 pub const MAX_POSITION_PER_ATTRIBUTE: u32 = u16::MAX as u32 + 1;
 
-#[derive(Clone, Copy)]
+#[derive(Clone)]
 pub struct TimeBudget {
     started_at: std::time::Instant,
     budget: std::time::Duration,
+
+    /// When testing the time budget, ensuring we did more than iteration of the bucket sort can be useful.
+    /// But to avoid being flaky, the only option is to add the ability to stop after a specific number of calls instead of a `Duration`.
+    #[cfg(test)]
+    stop_after: Option<(std::sync::Arc<std::sync::atomic::AtomicUsize>, usize)>,
 }
 
 impl fmt::Debug for TimeBudget {
@@ -129,18 +134,40 @@ impl Default for TimeBudget {
 
 impl TimeBudget {
     pub fn new(budget: std::time::Duration) -> Self {
-        Self { started_at: std::time::Instant::now(), budget }
+        Self {
+            started_at: std::time::Instant::now(),
+            budget,
+
+            #[cfg(test)]
+            stop_after: None,
+        }
     }
 
     pub fn max() -> Self {
         Self::new(std::time::Duration::from_secs(u64::MAX))
     }
 
-    pub fn exceeded(&self) -> bool {
-        self.must_stop()
+    #[cfg(test)]
+    pub fn with_stop_after(mut self, stop_after: usize) -> Self {
+        use std::sync::atomic::AtomicUsize;
+        use std::sync::Arc;
+
+        self.stop_after = Some((Arc::new(AtomicUsize::new(0)), stop_after));
+        self
     }
 
-    pub fn must_stop(&self) -> bool {
+    pub fn exceeded(&self) -> bool {
+        #[cfg(test)]
+        if let Some((current, stop_after)) = &self.stop_after {
+            let current = current.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+            if current >= *stop_after {
+                return true;
+            } else {
+                // if a number has been specified then we ignore entirely the time budget
+                return false;
+            }
+        }
+
         self.started_at.elapsed() > self.budget
     }
 }