fix: stabilize atomic publish e2e overlap detection

Author: Dicklesworthstone

tests/e2e_search_index.rs

@@ -985,6 +985,7 @@ fn force_rebuild_preserves_search_results_and_reader_surface_during_atomic_publi
 
     let stop = Arc::new(AtomicBool::new(false));
     let rebuild_running = Arc::new(AtomicBool::new(false));
+    let search_in_flight = Arc::new(AtomicBool::new(false));
     let reader_attempts_during_rebuild = Arc::new(AtomicUsize::new(0));
     let search_attempts_during_rebuild = Arc::new(AtomicUsize::new(0));
     let (ready_tx, ready_rx) = mpsc::channel();
@@ -1012,6 +1013,7 @@ fn force_rebuild_preserves_search_results_and_reader_surface_during_atomic_publi
     let search_ready_tx = ready_tx.clone();
     let search_stop = Arc::clone(&stop);
     let search_rebuild_running = Arc::clone(&rebuild_running);
+    let search_in_flight_thread = Arc::clone(&search_in_flight);
     let search_overlap = Arc::clone(&search_attempts_during_rebuild);
     let search_home = home.clone();
     let search_codex_home = codex_home.clone();
@@ -1020,11 +1022,9 @@ fn force_rebuild_preserves_search_results_and_reader_surface_during_atomic_publi
         let _ = search_ready_tx.send("search");
         let mut stats = SearchLoopStats::default();
         while !search_stop.load(Ordering::Relaxed) {
-            if search_rebuild_running.load(Ordering::Relaxed) {
-                search_overlap.fetch_add(1, Ordering::Relaxed);
-            }
-
             let search_started = Instant::now();
+            search_in_flight_thread.store(true, Ordering::Relaxed);
+            let started_during_rebuild = search_rebuild_running.load(Ordering::Relaxed);
             let output = cargo_bin_cmd!("cass")
                 .args([
                     "search",
@@ -1045,7 +1045,12 @@ fn force_rebuild_preserves_search_results_and_reader_surface_during_atomic_publi
                 .timeout(Duration::from_secs(20))
                 .output()
                 .expect("run concurrent cass search");
-            let elapsed_ms = search_started.elapsed().as_millis() as u64;
+            search_in_flight_thread.store(false, Ordering::Relaxed);
+            let search_finished = Instant::now();
+            if started_during_rebuild || search_rebuild_running.load(Ordering::Relaxed) {
+                search_overlap.fetch_add(1, Ordering::Relaxed);
+            }
+            let elapsed_ms = search_finished.duration_since(search_started).as_millis() as u64;
             stats.attempts += 1;
             stats.max_duration_ms = stats.max_duration_ms.max(elapsed_ms);
 
@@ -1086,6 +1091,9 @@ fn force_rebuild_preserves_search_results_and_reader_surface_during_atomic_publi
         Some("Run cass index --full --force-rebuild while a direct reader and cass search poll the same live index"),
     );
     rebuild_running.store(true, Ordering::Relaxed);
+    if search_in_flight.load(Ordering::Relaxed) {
+        search_attempts_during_rebuild.fetch_add(1, Ordering::Relaxed);
+    }
     let publish_pause_sentinel = home.join("atomic-publish-overlap-sentinel.json");
     let mut attempt = 0usize;
     let rebuild_output = loop {
@@ -1341,6 +1349,7 @@ fn force_rebuild_preserves_search_results_and_reader_surface_during_federated_at
 
     let stop = Arc::new(AtomicBool::new(false));
     let rebuild_running = Arc::new(AtomicBool::new(false));
+    let search_in_flight = Arc::new(AtomicBool::new(false));
     let reader_attempts_during_rebuild = Arc::new(AtomicUsize::new(0));
     let search_attempts_during_rebuild = Arc::new(AtomicUsize::new(0));
     let (ready_tx, ready_rx) = mpsc::channel();
@@ -1368,6 +1377,7 @@ fn force_rebuild_preserves_search_results_and_reader_surface_during_federated_at
     let search_ready_tx = ready_tx.clone();
     let search_stop = Arc::clone(&stop);
     let search_rebuild_running = Arc::clone(&rebuild_running);
+    let search_in_flight_thread = Arc::clone(&search_in_flight);
     let search_overlap = Arc::clone(&search_attempts_during_rebuild);
     let search_home = home.clone();
     let search_codex_home = codex_home.clone();
@@ -1376,11 +1386,9 @@ fn force_rebuild_preserves_search_results_and_reader_surface_during_federated_at
         let _ = search_ready_tx.send("search");
         let mut stats = SearchLoopStats::default();
         while !search_stop.load(Ordering::Relaxed) {
-            if search_rebuild_running.load(Ordering::Relaxed) {
-                search_overlap.fetch_add(1, Ordering::Relaxed);
-            }
-
             let search_started = Instant::now();
+            search_in_flight_thread.store(true, Ordering::Relaxed);
+            let started_during_rebuild = search_rebuild_running.load(Ordering::Relaxed);
             let output = cargo_bin_cmd!("cass")
                 .args([
                     "search",
@@ -1401,7 +1409,12 @@ fn force_rebuild_preserves_search_results_and_reader_surface_during_federated_at
                 .timeout(Duration::from_secs(20))
                 .output()
                 .expect("run concurrent federated cass search");
-            let elapsed_ms = search_started.elapsed().as_millis() as u64;
+            search_in_flight_thread.store(false, Ordering::Relaxed);
+            let search_finished = Instant::now();
+            if started_during_rebuild || search_rebuild_running.load(Ordering::Relaxed) {
+                search_overlap.fetch_add(1, Ordering::Relaxed);
+            }
+            let elapsed_ms = search_finished.duration_since(search_started).as_millis() as u64;
             stats.attempts += 1;
             stats.max_duration_ms = stats.max_duration_ms.max(elapsed_ms);
 
@@ -1442,6 +1455,9 @@ fn force_rebuild_preserves_search_results_and_reader_surface_during_federated_at
         Some("Run cass index --full --force-rebuild with forced multi-shard planning while a direct reader and cass search poll the same live index"),
     );
     rebuild_running.store(true, Ordering::Relaxed);
+    if search_in_flight.load(Ordering::Relaxed) {
+        search_attempts_during_rebuild.fetch_add(1, Ordering::Relaxed);
+    }
     let mut rebuild = cargo_bin_cmd!("cass");
     force_federated_publish_env(&mut rebuild);
     let publish_pause_sentinel = home.join("federated-atomic-publish-overlap-sentinel.json");