feat(sort): per-phase --sort-threads / --merge-threads control for sort and runall

crates/fgumi-pipeline-io/src/sort/spill_decompress.rs

@@ -4,6 +4,7 @@
 
 use std::io;
 use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
 
 use fgumi_sort::{PHASE2_DECOMP_CAP, SortMergeSlot, SpillBlockDecompressor};
 use parking_lot::Mutex;
@@ -21,6 +22,42 @@ use fgumi_pipeline_core::{
 /// Max raw blocks read+decompressed by a single `try_fill_some_slot` call.
 const MAX_BATCH_PER_CALL: usize = 4;
 
+/// CAS-acquire one decompress permit. `max == None` ⇒ unbounded (always succeeds).
+///
+/// Returns the [`DecompressPermit`] on success so ownership of the slot is
+/// encoded in the type: the count is released exactly once, when the returned
+/// permit drops, and acquisition cannot be separated from release. `None` ⇒ the
+/// cap is full and no slot was taken.
+#[must_use]
+fn try_acquire(active: &Arc<AtomicUsize>, max: Option<usize>) -> Option<DecompressPermit> {
+    let Some(max) = max else {
+        active.fetch_add(1, Ordering::AcqRel);
+        return Some(DecompressPermit { active: Arc::clone(active) });
+    };
+    let max = max.max(1);
+    let mut cur = active.load(Ordering::Acquire);
+    loop {
+        if cur >= max {
+            return None;
+        }
+        match active.compare_exchange_weak(cur, cur + 1, Ordering::AcqRel, Ordering::Acquire) {
+            Ok(_) => return Some(DecompressPermit { active: Arc::clone(active) }),
+            Err(observed) => cur = observed,
+        }
+    }
+}
+
+/// RAII release of one decompress permit. Holds an OWNED `Arc` clone so it does
+/// not borrow `self` while `try_fill_some_slot(&mut self)` runs.
+struct DecompressPermit {
+    active: Arc<AtomicUsize>,
+}
+impl Drop for DecompressPermit {
+    fn drop(&mut self) {
+        self.active.fetch_sub(1, Ordering::AcqRel);
+    }
+}
+
 struct RegisteredSpill {
     slot: Arc<SortMergeSlot>,
 }
@@ -33,6 +70,11 @@ pub struct SortSpillDecompress {
     block_dec: SpillBlockDecompressor,
     held: HeldSlot<Unpushed<SortPhase2Event>>,
     output_byte_limit: u64,
+    /// Shared admission counter: how many worker clones are currently inside the
+    /// decompress branch. Shared across all clones so the cap is global.
+    active: Arc<AtomicUsize>,
+    /// Maximum concurrent decompress workers. `None` ⇒ unbounded.
+    max_decompress: Option<usize>,
 }
 
 impl SortSpillDecompress {
@@ -49,9 +91,19 @@ impl SortSpillDecompress {
             block_dec: SpillBlockDecompressor::new(),
             held: HeldSlot::new(),
             output_byte_limit,
+            active: Arc::new(AtomicUsize::new(0)),
+            max_decompress: None,
         }
     }
 
+    /// Cap the number of concurrent spill-decompression workers (Phase-2 CPU
+    /// control via `--merge-threads`). `None` (default) leaves it unbounded.
+    #[must_use]
+    pub fn with_max_concurrency(mut self, max: Option<usize>) -> Self {
+        self.max_decompress = max;
+        self
+    }
+
     fn flush_held(&mut self, ctx: &mut StepCtx<'_, Self>) -> bool {
         let Some(unpushed) = self.held.take() else {
             return true;
@@ -147,6 +199,9 @@ impl Clone for SortSpillDecompress {
             block_dec: SpillBlockDecompressor::new(),
             held: HeldSlot::new(),
             output_byte_limit: self.output_byte_limit,
+            // Share the SAME counter so the cap is global across all worker clones.
+            active: Arc::clone(&self.active),
+            max_decompress: self.max_decompress,
         }
     }
 }
@@ -189,9 +244,12 @@ impl Step for SortSpillDecompress {
             return Ok(StepOutcome::Progress);
         }
 
-        // 3. Greedy slot-fill step.
-        if self.try_fill_some_slot()? {
-            return Ok(StepOutcome::Progress);
+        // 3. Greedy slot-fill, admission-controlled by --merge-threads.
+        if let Some(_permit) = try_acquire(&self.active, self.max_decompress) {
+            if self.try_fill_some_slot()? {
+                return Ok(StepOutcome::Progress);
+            }
+            // permit drops here (also on the `?` early return), releasing the count
         }
 
         // 4. No fill work.

crates/fgumi-pipeline-io/src/sort/spill_decompress/tests.rs

@@ -1 +1,60 @@
-// No unit tests for spill_decompress directly; integration tests live in sort/tests.rs.
+use super::*;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+#[test]
+fn admission_counter_caps_concurrency() {
+    let active = Arc::new(AtomicUsize::new(0));
+    let max = Some(2usize);
+    // Acquire up to the cap; hold the permits so the count accumulates.
+    let p1 = try_acquire(&active, max);
+    assert!(p1.is_some());
+    let p2 = try_acquire(&active, max);
+    assert!(p2.is_some());
+    assert!(try_acquire(&active, max).is_none()); // at cap
+    drop(p1); // releasing one permit frees a slot
+    assert!(try_acquire(&active, max).is_some()); // freed one slot
+    drop(p2);
+}
+
+#[test]
+fn admission_counter_unbounded_when_none() {
+    let active = Arc::new(AtomicUsize::new(0));
+    for _ in 0..1000 {
+        assert!(try_acquire(&active, None).is_some());
+    }
+}
+
+/// Under real thread contention, the shared counter never exceeds the cap and
+/// every acquired permit is released via `DecompressPermit::drop` (the counter
+/// returns to zero). Mirrors how `new_worker_copy` clones share one `active`.
+#[test]
+fn admission_counter_concurrent_never_exceeds_cap() {
+    use std::thread;
+
+    let active = Arc::new(AtomicUsize::new(0));
+    let cap = 3usize;
+    let handles: Vec<_> = (0..16)
+        .map(|_| {
+            let active = Arc::clone(&active);
+            thread::spawn(move || {
+                for _ in 0..5000 {
+                    // The returned permit IS the ownership token; its Drop at the
+                    // end of the block exercises the decrement.
+                    if let Some(_permit) = try_acquire(&active, Some(cap)) {
+                        // Occupancy observed while holding a permit can never
+                        // exceed the cap: `try_acquire` only increments past a
+                        // CAS that checks `cur < cap`, and the count only drops
+                        // otherwise.
+                        let occupancy = active.load(Ordering::Acquire);
+                        assert!(occupancy <= cap, "occupancy {occupancy} exceeded cap {cap}");
+                    }
+                }
+            })
+        })
+        .collect();
+    for h in handles {
+        h.join().expect("worker thread panicked");
+    }
+    assert_eq!(active.load(Ordering::Acquire), 0, "every permit must be released on drop");
+}

crates/fgumi-sort-cli/src/chains.rs

@@ -135,8 +135,14 @@ pub struct SortStepCaptures {
     pub input_path: PathBuf,
     /// Output BAM path.
     pub output_path: PathBuf,
-    /// Number of threads for the sort engine's worker pool.
+    /// Number of threads for the sort engine's worker pool (the overall /
+    /// Phase-2 "writing" count).
     pub num_sorter_threads: usize,
+    /// Worker count for Phase 1 (accumulate/sort/spill); defaults to
+    /// `num_sorter_threads`. Lower to cede CPU to an upstream producer.
+    pub num_phase1_threads: usize,
+    /// Worker count for Phase 2 (merge/write); defaults to `num_sorter_threads`.
+    pub num_phase2_threads: usize,
     /// Resolved memory budget, already computed by the caller via
     /// [`resolve_memory_budget`](fgumi_cli_common::resolve_memory_budget).
     /// Passed in (rather than recomputed here) so
@@ -172,6 +178,8 @@ pub fn build_sort_step(cap: SortStepCaptures) -> Result<SortBamFile> {
     let mut sorter = RawExternalSorter::new(sort_order)
         .memory_limit(effective_memory)
         .threads(cap.num_sorter_threads)
+        .sort_threads(cap.num_phase1_threads)
+        .merge_threads(cap.num_phase2_threads)
         .output_compression(cap.output_compression)
         .temp_compression(cap.sort.temp_compression)
         .spill_codec(cap.sort.temp_codec)
@@ -210,6 +218,8 @@ pub fn log_sort_start(
     input_path: &std::path::Path,
     output_path: &std::path::Path,
     num_sorter_threads: usize,
+    num_phase1_threads: usize,
+    num_phase2_threads: usize,
     effective_memory: usize,
 ) {
     let cell_tag = parse_cell_tag(sort.order).unwrap_or(None);
@@ -233,7 +243,13 @@ pub fn log_sort_start(
             info!("Max memory: {} (fixed)", ByteSize(effective_memory as u64));
         }
     }
-    info!("Threads: {num_sorter_threads}");
+    if num_phase1_threads == num_sorter_threads && num_phase2_threads == num_sorter_threads {
+        info!("Threads: {num_sorter_threads}");
+    } else {
+        info!(
+            "Threads: {num_sorter_threads} (sort phase {num_phase1_threads}, merge phase {num_phase2_threads})"
+        );
+    }
     info!("Temp compression level: {}", sort.temp_compression);
     let env_value = std::env::var(TMP_DIRS_ENV).ok();
     let resolved_tmp_dirs = resolve_tmp_dirs(&sort.tmp_dirs, env_value.as_deref());

crates/fgumi-sort-cli/src/sort.rs

@@ -325,6 +325,25 @@ pub struct SortOptions {
     #[arg(long = "key-types", value_parser = parse_key_types)]
     pub key_types: Option<KeyTypesSpec>,
 
+    /// Worker threads for the accumulation/sort/spill phase (Phase 1).
+    ///
+    /// Defaults to `--threads`. Lower this to cede CPU to an upstream producer
+    /// (e.g. `bwa mem | fgumi sort`, or runall's align→sort chain); use
+    /// `--merge-threads` for the merge/write phase. Output is byte-identical
+    /// regardless of this value.
+    #[arg(long = "sort-threads")]
+    pub sort_threads: Option<usize>,
+
+    /// Worker threads for the merge/write phase (Phase 2).
+    ///
+    /// Defaults to `--threads`. In standalone `fgumi sort` (and runall's
+    /// sole-stage sort) this caps the merge/write workers of the sort engine's
+    /// pool. In a multi-stage runall chain it caps the number of concurrent
+    /// spill-decompression workers (the merge step itself is serial). Output is
+    /// byte-identical regardless of this value.
+    #[arg(long = "merge-threads")]
+    pub merge_threads: Option<usize>,
+
     /// Sort order (chain-builder slot).
     ///
     /// Carried here so the chain builder can read it from the bag without
@@ -344,6 +363,8 @@ impl Default for SortOptions {
             temp_compression: 1,
             temp_codec: SpillCodec::default(),
             key_types: None,
+            sort_threads: None,
+            merge_threads: None,
             order: SortOrderArg::TemplateCoordinate,
         }
     }
@@ -505,6 +526,12 @@ impl Sort {
         let input_path = self.input.clone();
         let output_path = output.to_path_buf();
         let num_sorter_threads = self.threads.max(1);
+        // Per-phase worker counts; each defaults to --threads. Memory budget
+        // intentionally stays tied to --threads (these knobs are CPU-only). Read
+        // from `sort_opts` (the resolved options used by the rest of the method)
+        // so the phase counts cannot diverge from the logged / built config.
+        let num_phase1_threads = sort_opts.sort_threads.unwrap_or(num_sorter_threads).max(1);
+        let num_phase2_threads = sort_opts.merge_threads.unwrap_or(num_sorter_threads).max(1);
 
         let effective_memory = resolve_memory_budget(
             sort_opts.max_memory,
@@ -514,7 +541,15 @@ impl Sort {
         )?;
 
         let timer = OperationTimer::new("Sorting BAM");
-        log_sort_start(&sort_opts, &input_path, &output_path, num_sorter_threads, effective_memory);
+        log_sort_start(
+            &sort_opts,
+            &input_path,
+            &output_path,
+            num_sorter_threads,
+            num_phase1_threads,
+            num_phase2_threads,
+            effective_memory,
+        );
 
         let stats_slot = Arc::new(Mutex::new(None::<fgumi_sort::SortStats>));
 
@@ -523,6 +558,8 @@ impl Sort {
             input_path,
             output_path: output_path.clone(),
             num_sorter_threads,
+            num_phase1_threads,
+            num_phase2_threads,
             effective_memory,
             output_compression: self.compression.compression_level,
             command_line: command_line.to_string(),
@@ -849,6 +886,44 @@ mod tests {
         assert!(output.exists(), "sorted output not written");
     }
 
+    #[test]
+    fn test_phase_thread_flags_parse() {
+        use clap::Parser;
+        // Standalone: flags now live on SortOptions, still spelled --sort-threads / --merge-threads.
+        let cmd = Sort::parse_from([
+            "sort",
+            "-i",
+            "in.bam",
+            "-o",
+            "out.bam",
+            "--sort-threads",
+            "2",
+            "--merge-threads",
+            "5",
+        ]);
+        assert_eq!(cmd.options.sort_threads, Some(2));
+        assert_eq!(cmd.options.merge_threads, Some(5));
+
+        // Defaults: unset => None (engine falls back to --threads).
+        let cmd = Sort::parse_from(["sort", "-i", "in.bam", "-o", "out.bam"]);
+        assert_eq!(cmd.options.sort_threads, None);
+        assert_eq!(cmd.options.merge_threads, None);
+    }
+
+    #[test]
+    fn test_runall_prefixed_phase_thread_flags_parse() {
+        use clap::Parser;
+        #[derive(clap::Parser)]
+        struct Probe {
+            #[command(flatten)]
+            sort: MultiSortOptions,
+        }
+        let p = Probe::parse_from(["x", "--sort::sort-threads", "3", "--sort::merge-threads", "7"]);
+        let opts = p.sort.validate().expect("validate");
+        assert_eq!(opts.sort_threads, Some(3));
+        assert_eq!(opts.merge_threads, Some(7));
+    }
+
     /// Write a minimal valid coordinate-sorted BAM (one mapped record) to
     /// `path`, so a pre-placed output file looks indexable to the BAI hook.
     #[allow(clippy::cast_possible_truncation)]