refactor(code-index): route chunk fan-outs through one yielding helper

crates/tracedecay-code-index/src/chunks.rs

@@ -13,7 +13,6 @@ use std::{
     sync::{Arc, OnceLock},
 };
 
-use rayon::prelude::*;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 use tracedecay_code_extraction::{ExtractedCloneBodyV1, ExtractionArtifactV1};
@@ -213,6 +212,33 @@ unsafe impl ExtractionAdmittedChunkV1 for ExtractionAdmittedCodeSearchChunkV1 {
 /// for the coarser per-file fan-out above this layer.
 const PARALLEL_CHUNK_THRESHOLD: usize = 16;
 
+/// The only route from this module to the rayon pool.
+///
+/// Every chunk fan-out runs on a worker that already holds one background-CPU
+/// unit (a `collect_bounded_ordered` leaf), and its stolen halves admit
+/// themselves one unit at a time. Holding the parent's unit across the join
+/// while a full-width request sits at the FIFO head wedges the process, so
+/// the yield is welded to the fan-out here rather than left at each call
+/// site, where a merge resolution once kept the `par_iter` and dropped the
+/// yield. `rayon` is imported nowhere else in `chunks`, so a bare `par_iter`
+/// outside this module does not compile.
+mod fan_out {
+    use rayon::prelude::*;
+
+    /// Map `items` across the pool and return the outputs in input order,
+    /// with the caller's admitted units yielded for the duration of the join.
+    pub(super) fn map_yielding<I, R>(items: I, map: impl Fn(I::Item) -> R + Send + Sync) -> Vec<R>
+    where
+        I: IntoParallelIterator,
+        I::Iter: IndexedParallelIterator,
+        R: Send,
+    {
+        crate::parallelism::with_yielded_background_cpu_permits(|| {
+            items.into_par_iter().map(map).collect()
+        })
+    }
+}
+
 /// Run `operation` over every chunk for its failure only, fanning out across
 /// the pool once the batch is large enough. The lowest-index failure is
 /// returned, matching the sequential sweep's short-circuit outcome.
@@ -229,23 +255,12 @@ where
     if chunks.len() < PARALLEL_CHUNK_THRESHOLD {
         return chunks.iter().try_for_each(&operation);
     }
-    // Leaves are admitted one unit at a time on whichever worker runs them,
-    // so the caller's own unit must not be held across the join.
-    let failure = crate::parallelism::with_yielded_background_cpu_permits(|| {
-        chunks
-            .par_iter()
-            .enumerate()
-            .filter_map(|(index, chunk)| {
-                admit(&mut || operation(chunk))
-                    .err()
-                    .map(|error| (index, error))
-            })
-            .min_by_key(|(index, _)| *index)
-    });
-    match failure {
-        Some((_, error)) => Err(error),
-        None => Ok(()),
-    }
+    // Outputs come back in input order, so the first failure is the
+    // lowest-index one.
+    fan_out::map_yielding(chunks, |chunk| admit(&mut || operation(chunk)).err())
+        .into_iter()
+        .find_map(|failure| failure)
+        .map_or(Ok(()), Err)
 }
 
 impl ExactExtractionAuthorityV1 {
@@ -365,13 +380,11 @@ impl ExactExtractionAuthorityV1 {
         if chunks.len() < PARALLEL_CHUNK_THRESHOLD {
             return chunks.into_iter().map(|chunk| self.admit(chunk)).collect();
         }
-        let admitted = crate::parallelism::with_yielded_background_cpu_permits(|| {
-            chunks
-                .into_par_iter()
-                .map(|chunk| crate::parallelism::with_background_cpu_permit(|| self.admit(chunk)))
-                .collect::<Vec<_>>()
-        });
-        admitted.into_iter().collect()
+        fan_out::map_yielding(chunks, |chunk| {
+            crate::parallelism::with_background_cpu_permit(|| self.admit(chunk))
+        })
+        .into_iter()
+        .collect()
     }
 
     /// Rebind an exact authority only after every prior parser-backed chunk
@@ -2929,6 +2942,61 @@ mod tests {
         assert_eq!(authority.waiting_work_units(), 0);
     }
 
+    /// Structural pin for the nested-admission fix: the yield lives inside
+    /// `fan_out`, so it survives only while `fan_out` stays the sole pool
+    /// entry point in this file. A merge resolution that reintroduces a bare
+    /// `par_iter` elsewhere fails here even before it fails to compile.
+    #[test]
+    fn every_pool_fan_out_in_chunks_goes_through_the_yielding_helper() {
+        let source = include_str!("chunks.rs");
+        let module_start = source
+            .find("\nmod fan_out {\n")
+            .expect("chunks.rs declares `mod fan_out`");
+        let module_end = module_start
+            + source[module_start..]
+                .find("\n}\n")
+                .expect("`mod fan_out` closes at column zero");
+        let module = &source[module_start..=module_end];
+        assert!(
+            module.contains("with_yielded_background_cpu_permits("),
+            "`fan_out` must yield the caller's admitted units around the join"
+        );
+
+        // Built with `concat!` so this test's own source is not an occurrence.
+        let pool_tokens = [
+            concat!("ray", "on"),
+            concat!("par_", "iter("),
+            concat!("par_", "chunks"),
+            concat!("par_", "bridge"),
+        ];
+        let code_lines = |text: &str| {
+            text.lines()
+                .map(str::trim_start)
+                .filter(|line| !line.starts_with("//"))
+                .map(str::to_owned)
+                .collect::<Vec<_>>()
+        };
+        let outside = code_lines(&source[..module_start])
+            .into_iter()
+            .chain(code_lines(&source[module_end + 1..]))
+            .collect::<Vec<_>>();
+        let inside = code_lines(module);
+        for token in pool_tokens {
+            let stray = outside
+                .iter()
+                .filter(|line| line.contains(token))
+                .collect::<Vec<_>>();
+            assert!(
+                stray.is_empty(),
+                "`{token}` reaches the pool outside `fan_out`, bypassing the yield: {stray:?}"
+            );
+        }
+        assert!(
+            inside.iter().any(|line| line.contains(pool_tokens[0])),
+            "`fan_out` is expected to be where the pool crate is imported"
+        );
+    }
+
     const RUST_SOURCE: &str = "//! Module documentation.\n\nuse std::collections::HashMap;\n\n/// Doc comment.\npub fn alpha(x: u32) -> u32 {\n    x + 1\n}\n\npub struct Holder {\n    map: HashMap<u32, u32>,\n}\n\nimpl Holder {\n    pub fn get(&self, key: u32) -> Option<u32> {\n        self.map.get(&key).copied()\n    }\n}\n\n// A trailing free-floating comment.\n";
 
     fn chunker() -> DeterministicCodeChunker {