From b7a32013050ab68d845ea1c0d47775fa11496e44 Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Thu, 25 Jun 2026 14:38:49 -0700 Subject: [PATCH 1/4] feat(sort): per-phase worker-thread control in RawExternalSorter Add --sort-threads (Phase 1 = accumulate/sort/spill) and --merge-threads (Phase 2 = merge/write) worker control to the sort engine. SortWorkerPool gains an active-worker cap (active_worker_limit + set_active_workers); RawExternalSorter gains sort_threads/merge_threads builders and phase1_threads()/phase2_threads(). The monolithic sort() sizes the pool to max(phase1, phase2) and switches the active cap at the ingest->write boundary; the streaming into_*_stream constructors size their Phase-1-only pool to phase1_threads(). Capped workers still drain their per-worker held items (only new-work acquisition is blocked) so lowering the cap can never strand output. Output is byte-identical (a pure scheduling knob); memory stays tied to --threads. --- crates/fgumi-sort/src/external.rs | 334 +++++++++++++++++++++++++-- crates/fgumi-sort/src/worker_pool.rs | 128 ++++++++++ 2 files changed, 440 insertions(+), 22 deletions(-) diff --git a/crates/fgumi-sort/src/external.rs b/crates/fgumi-sort/src/external.rs index 49f1003a6..c21dc94a7 100644 --- a/crates/fgumi-sort/src/external.rs +++ b/crates/fgumi-sort/src/external.rs @@ -1322,8 +1322,17 @@ pub struct RawExternalSorter { /// distributed across them in free-space-aware round-robin order via /// [`TmpDirAllocator`]. temp_dirs: Vec, - /// Number of threads for parallel operations. + /// Number of threads for parallel operations — the base/fallback count used + /// by both phases unless overridden by `sort_threads` / `merge_threads`. threads: usize, + /// Optional override for the Phase-1 (accumulate/sort/spill) worker count. + /// `None` ⇒ Phase 1 uses `threads`. Lower it to cede cores to an upstream + /// producer (e.g. `bwa mem` in a pipeline) during accumulation. + sort_threads: Option, + /// Optional override for the Phase-2 (merge/write) worker count. + /// `None` ⇒ Phase 2 uses `threads`. The merge/write phase scales with cores, + /// so this is typically left at (or raised toward) the full core count. + merge_threads: Option, /// Compression level for output. output_compression: u32, /// Compression level for temporary chunk files (0 = uncompressed). @@ -1420,8 +1429,9 @@ impl RawExternalSorter { // detects each spill chunk's codec from its file magic (see // `slots_for_chunk_files`) and dispatches accordingly — so the streaming // path no longer has to pin BGZF. + // Phase 1 (streaming) is capped by --sort-threads; this pool is Phase-1-only (dropped at finalize). let pool = Arc::new(SortWorkerPool::new( - self.threads.max(1), + self.phase1_threads(), self.temp_compression, self.output_compression, self.spill_codec, @@ -1491,8 +1501,9 @@ impl RawExternalSorter { // Spill codec from config (see `into_coordinate_stream`): the streaming // decoder is codec-aware, so honor `self.spill_codec` (zstd default). + // Phase 1 (streaming) is capped by --sort-threads; this pool is Phase-1-only (dropped at finalize). let pool = Arc::new(SortWorkerPool::new( - self.threads.max(1), + self.phase1_threads(), self.temp_compression, self.output_compression, self.spill_codec, @@ -1580,6 +1591,8 @@ impl RawExternalSorter { memory_limit: 512 * 1024 * 1024, // 512 MB default temp_dirs: Vec::new(), threads: 1, + sort_threads: None, + merge_threads: None, output_compression: 6, temp_compression: 1, // Default: fast compression spill_codec: crate::codec::SpillCodec::default(), @@ -1626,6 +1639,43 @@ impl RawExternalSorter { self } + /// Set the Phase-1 (accumulate/sort/spill) worker count. Defaults to + /// [`threads`](Self::threads) when unset. Lower it to cede cores to an + /// upstream producer in a pipeline. + /// + /// Honored by [`sort`](Self::sort) (via the worker-pool active cap) and by + /// the streaming `into_*_stream` entry points (which size their Phase-1-only + /// pool via `phase1_threads()`). + #[must_use] + pub fn sort_threads(mut self, n: usize) -> Self { + self.sort_threads = Some(n); + self + } + + /// Set the Phase-2 (merge/write) worker count. Defaults to + /// [`threads`](Self::threads) when unset. + /// + /// Honored by [`sort`](Self::sort) (via the worker-pool active cap). The + /// streaming `into_*_stream` entry points do not read this value — their + /// Phase 2 (decompress/merge) runs on the framework scheduler, so the chain + /// builder caps streaming Phase-2 concurrency at the pipeline layer + /// (`SortSpillDecompress::with_max_concurrency`) instead. + #[must_use] + pub fn merge_threads(mut self, n: usize) -> Self { + self.merge_threads = Some(n); + self + } + + /// Effective Phase-1 worker count: `sort_threads` if set, else `threads`. + fn phase1_threads(&self) -> usize { + self.sort_threads.unwrap_or(self.threads).max(1) + } + + /// Effective Phase-2 worker count: `merge_threads` if set, else `threads`. + fn phase2_threads(&self) -> usize { + self.merge_threads.unwrap_or(self.threads).max(1) + } + /// Set the output compression level. #[must_use] pub fn output_compression(mut self, level: u32) -> Self { @@ -1742,7 +1792,7 @@ impl RawExternalSorter { /// idle at the moment rayon fans out. fn build_sort_rayon_pool(&self) -> Result { rayon::ThreadPoolBuilder::new() - .num_threads(self.threads.max(1)) + .num_threads(self.phase1_threads()) .thread_name(|i| format!("fgumi-sort-rayon-{i}")) .build() .map_err(|e| anyhow::anyhow!("failed to build rayon sort pool: {e}")) @@ -1808,8 +1858,11 @@ impl RawExternalSorter { // Create merged output file let merged_path = namer.next_merged_path()?; - // Open readers with semaphore to cap concurrent I/O. - let sem = make_reader_semaphore(self.threads); + // Open readers with semaphore to cap concurrent I/O. Consolidation runs + // during Phase 1 (from `drain_pending_spill`, before the Phase-2 worker + // switch), so bound the readers by the Phase-1 thread count to match the + // rest of the Phase-1 sizing rather than the base thread count. + let sem = make_reader_semaphore(self.phase1_threads()); let mut readers: Vec> = files_to_merge .iter() .map(|p| GenericKeyedChunkReader::::open(p, Some(Arc::clone(&sem)))) @@ -1912,13 +1965,20 @@ impl RawExternalSorter { debug!("Memory limit: {} MB", self.memory_limit / (1024 * 1024)); debug!("Threads: {}", self.threads); - // Shared worker pool for parallel BGZF compress/decompress across all phases + // Shared worker pool for parallel BGZF compress/decompress across all + // phases. Sized to the larger of the two phase counts; the active-worker + // cap (set below, then switched to the Phase-2 count at the merge/write + // boundary — which may raise or lower it) governs how many run in each + // phase. let pool = Arc::new(SortWorkerPool::new( - self.threads.max(1), + self.phase1_threads().max(self.phase2_threads()), self.temp_compression, self.output_compression, self.spill_codec, )); + // Phase 1 (accumulate/sort/spill) runs on the Phase-1 count; the + // merge/write phase switches the cap to the Phase-2 count. + pool.set_active_workers(self.phase1_threads()); // Open input BAM and create record source // N+2 model: workers do ReadInputBlocks + DecompressInput, @@ -2218,6 +2278,11 @@ impl RawExternalSorter { )?; probe.phase1_end(buffer.memory_usage() as u64); + // Ingest (accumulate/sort/spill) is done. Switch the worker cap to the + // Phase-2 count for the writing phase (merge or in-memory write), which + // is the part that scales and runs after any upstream producer. + pool.set_active_workers(self.phase2_threads()); + if chunk_files.is_empty() { // All records fit in memory - no merge needed debug!("All records fit in memory, performing in-memory sort"); @@ -2243,8 +2308,10 @@ impl RawExternalSorter { // chunk becomes its own in-memory merge source. let memory_chunks: Vec> = if buffer.is_empty() { Vec::new() - } else if self.threads > 1 { - timer.time_sort(|| rayon_pool.install(|| buffer.par_sort_into_chunks(self.threads))) + } else if self.phase1_threads() > 1 { + timer.time_sort(|| { + rayon_pool.install(|| buffer.par_sort_into_chunks(self.phase1_threads())) + }) } else { timer.time_sort(|| { rayon_pool.install(|| buffer.par_sort()); @@ -2430,6 +2497,9 @@ impl RawExternalSorter { )?; probe.phase1_end(memory_used as u64); + // Ingest done; switch the worker cap to the Phase-2 count for the write phase. + pool.set_active_workers(self.phase2_threads()); + if chunk_files.is_empty() { // All records fit in memory debug!("All records fit in memory, performing in-memory sort"); @@ -2463,10 +2533,10 @@ impl RawExternalSorter { // `InMemoryChunk`s sharing an `Arc`. let keyed_chunks: Vec> = if entries.is_empty() { Vec::new() - } else if self.threads > 1 { + } else if self.phase1_threads() > 1 { timer.time_sort(|| { use rayon::prelude::*; - let chunk_size = entries.len().div_ceil(self.threads.max(1)); + let chunk_size = entries.len().div_ceil(self.phase1_threads()); rayon_pool.install(|| { entries.par_chunks_mut(chunk_size).for_each(|chunk| { chunk.sort_unstable_by(|a, b| a.0.cmp(&b.0)); @@ -2789,6 +2859,9 @@ impl RawExternalSorter { )?; probe.phase1_end(buffer.memory_usage() as u64); + // Ingest done; switch the worker cap to the Phase-2 count for the write phase. + pool.set_active_workers(self.phase2_threads()); + if chunk_files.is_empty() { // All records fit in memory debug!("All records fit in memory, performing in-memory sort"); @@ -2813,8 +2886,10 @@ impl RawExternalSorter { // intermediate merge back into a single sorted buffer) let memory_chunks: Vec> = if buffer.is_empty() { Vec::new() - } else if self.threads > 1 { - timer.time_sort(|| rayon_pool.install(|| buffer.par_sort_into_chunks(self.threads))) + } else if self.phase1_threads() > 1 { + timer.time_sort(|| { + rayon_pool.install(|| buffer.par_sort_into_chunks(self.phase1_threads())) + }) } else { timer.time_sort(|| { rayon_pool.install(|| buffer.par_sort()); @@ -3825,10 +3900,10 @@ impl CoordinateSortStream { .collect::>() }); vec![chunk] - } else if self.sorter.threads > 1 { + } else if self.sorter.phase1_threads() > 1 { let buffer = &mut self.buffer; let rayon_pool = &self.rayon_pool; - let threads = self.sorter.threads; + let threads = self.sorter.phase1_threads(); inmem_chunks_to_keyed( self.timer .time_sort(|| rayon_pool.install(|| buffer.par_sort_into_chunks(threads))), @@ -3955,8 +4030,9 @@ impl QuerynameSortStreamGeneric { // Spill codec from config (see `into_coordinate_stream`): the streaming // decoder is codec-aware, so honor `sorter.spill_codec` (zstd default). + // Phase 1 (streaming) is capped by --sort-threads; this pool is Phase-1-only (dropped at finalize). let pool = Arc::new(SortWorkerPool::new( - sorter.threads.max(1), + sorter.phase1_threads(), sorter.temp_compression, sorter.output_compression, sorter.spill_codec, @@ -4146,7 +4222,7 @@ impl QuerynameSortStreamGeneric { } else if self.chunk_files.is_empty() { let entries = &mut self.entries; let rayon_pool = &self.rayon_pool; - if self.sorter.threads > 1 { + if self.sorter.phase1_threads() > 1 { self.timer.time_sort(|| { use rayon::prelude::*; rayon_pool.install(|| entries.par_sort_unstable_by(|a, b| a.0.cmp(&b.0))); @@ -4157,10 +4233,10 @@ impl QuerynameSortStreamGeneric { }); } vec![std::mem::take(&mut self.entries)] - } else if self.sorter.threads > 1 { + } else if self.sorter.phase1_threads() > 1 { let entries = &mut self.entries; let rayon_pool = &self.rayon_pool; - let threads = self.sorter.threads; + let threads = self.sorter.phase1_threads(); self.timer.time_sort(|| { use rayon::prelude::*; let chunk_size = entries.len().div_ceil(threads.max(1)); @@ -4414,10 +4490,10 @@ impl TemplateCoordinateSortStream { .collect::>() }); vec![chunk] - } else if self.sorter.threads > 1 { + } else if self.sorter.phase1_threads() > 1 { let buffer = &mut self.buffer; let rayon_pool = &self.rayon_pool; - let threads = self.sorter.threads; + let threads = self.sorter.phase1_threads(); inmem_chunks_to_keyed( self.timer .time_sort(|| rayon_pool.install(|| buffer.par_sort_into_chunks(threads))), @@ -5178,6 +5254,220 @@ mod tests { assert_eq!(sorter.max_temp_files, 128); } + #[test] + fn test_raw_sorter_phase_threads_builders() { + let new = || RawExternalSorter::new(SortOrder::Coordinate).threads(8); + // Default: both phases fall back to threads. + let base = new(); + assert_eq!((base.sort_threads, base.merge_threads), (None, None)); + assert_eq!((base.phase1_threads(), base.phase2_threads()), (8, 8)); + // sort_threads overrides only Phase 1. + let sort_only = new().sort_threads(2); + assert_eq!((sort_only.phase1_threads(), sort_only.phase2_threads()), (2, 8)); + // merge_threads overrides only Phase 2. + let merge_only = new().merge_threads(3); + assert_eq!((merge_only.phase1_threads(), merge_only.phase2_threads()), (8, 3)); + // Both override independently. + let both = new().sort_threads(2).merge_threads(16); + assert_eq!((both.phase1_threads(), both.phase2_threads()), (2, 16)); + // Clamp to >= 1. + let zero = new().sort_threads(0).merge_threads(0); + assert_eq!((zero.phase1_threads(), zero.phase2_threads()), (1, 1)); + } + + /// Read every record's raw bytes from a BAM, in file order. + fn records_bytes(path: &Path) -> Vec> { + let (mut reader, _h) = create_raw_bam_reader(path, 1).expect("open bam"); + let mut out = Vec::new(); + let mut rec = fgumi_raw_bam::RawRecord::new(); + while reader.read_record(&mut rec).expect("read record") != 0 { + out.push(rec.view().as_bytes().to_vec()); + } + out + } + + /// Splitting Phase-1 (`sort_threads`) and Phase-2 (`merge_threads`) from the + /// base `threads` is a pure scheduling knob: output must be byte-identical to + /// the unsplit sort, with spills forced. Matrixed over every order (`#[values]`) + /// and every thread-split variant (`#[case]`) so a failure names the exact order + /// and split. The cases mirror the prior inline matrix: `sort_low` (Phase-1 on + /// 1), `merge_low` (Phase-2 on 1), `both` (Phase-1 on 1 / Phase-2 on 2), and + /// `sort_high` (Phase-1 above base). + #[rstest::rstest] + #[case::sort_low(Some(1), None)] + #[case::merge_low(None, Some(1))] + #[case::both(Some(1), Some(2))] + #[case::sort_high(Some(8), Some(2))] + fn phase_thread_split_is_byte_identical( + #[values( + SortOrder::Coordinate, + SortOrder::Queryname(QuerynameComparator::Natural), + SortOrder::Queryname(QuerynameComparator::Lexicographic), + SortOrder::TemplateCoordinate + )] + order: SortOrder, + #[case] sort_t: Option, + #[case] merge_t: Option, + ) { + use fgumi_sam::SamBuilder; + let mut builder = SamBuilder::new(); + for i in 0..2000 { + let _ = builder + .add_pair() + .name(&format!("read{i}")) + .start1(i * 50 + 1) + .start2(i * 50 + 101) + .build(); + } + let dir = tempfile::tempdir().expect("temp dir"); + let input = dir.path().join("in.bam"); + builder.write_bam(&input).expect("write bam"); + + let mk = |out: &Path, sort_t: Option, merge_t: Option| { + let mut s = RawExternalSorter::new(order) + .threads(4) + .memory_limit(32 * 1024) // force spills + .spill_codec(crate::codec::SpillCodec::Bgzf) + .temp_compression(0) + .output_compression(0); + if let Some(j) = sort_t { + s = s.sort_threads(j); + } + if let Some(k) = merge_t { + s = s.merge_threads(k); + } + s.sort(&input, out).expect("sort"); + }; + // Base: both phases on the unsplit `threads(4)`. + let base = dir.path().join("base.bam"); + mk(&base, None, None); + let golden = records_bytes(&base); + // The split variant must reproduce the base output exactly. + let out = dir.path().join("variant.bam"); + mk(&out, sort_t, merge_t); + assert_eq!( + records_bytes(&out), + golden, + "order {order:?} split sort_t={sort_t:?} merge_t={merge_t:?}: \ + thread split must not change output" + ); + } + + /// Regression guard for the residual Phase-1 sort honoring `sort_threads` + /// rather than the base `threads`. With `threads(1).sort_threads(4)` the + /// residual sort/chunk decision must take the *parallel* `phase1_threads()` + /// branch (previously it gated on `threads`, so it ran serially and ignored + /// `sort_threads`). Spills are forced so the keyed-chunk residual path + /// (`chunk_files` non-empty) is exercised; output must be byte-identical to + /// the serial `threads(1)` sort for every order. + #[rstest::rstest] + #[case::coordinate(SortOrder::Coordinate)] + #[case::queryname_natural(SortOrder::Queryname(QuerynameComparator::Natural))] + #[case::queryname_lex(SortOrder::Queryname(QuerynameComparator::Lexicographic))] + #[case::template_coordinate(SortOrder::TemplateCoordinate)] + fn residual_phase1_honors_sort_threads_with_single_base_thread(#[case] order: SortOrder) { + use fgumi_sam::SamBuilder; + let mut builder = SamBuilder::new(); + for i in 0..2000 { + let _ = builder + .add_pair() + .name(&format!("read{i}")) + .start1(i * 50 + 1) + .start2(i * 50 + 101) + .build(); + } + let dir = tempfile::tempdir().expect("temp dir"); + let input = dir.path().join("in.bam"); + builder.write_bam(&input).expect("write bam"); + + let mk = |out: &Path, sort_t: Option| { + let mut s = RawExternalSorter::new(order) + .threads(1) + .memory_limit(32 * 1024) // force spills -> exercise residual chunk path + .spill_codec(crate::codec::SpillCodec::Bgzf) + .temp_compression(0) + .output_compression(0); + if let Some(j) = sort_t { + s = s.sort_threads(j); + } + s.sort(&input, out).expect("sort"); + }; + let serial = dir.path().join("serial.bam"); + mk(&serial, None); // threads(1): residual sort runs serially + let golden = records_bytes(&serial); + + let parallel = dir.path().join("parallel.bam"); + mk(¶llel, Some(4)); // sort_threads(4): residual sort now parallel + assert_eq!( + records_bytes(¶llel), + golden, + "order {order:?}: threads(1).sort_threads(4) must match the serial sort output" + ); + } + + /// Build a streaming sorter for `order` with `threads(4)` and the optional + /// `sort_threads` cap, returning its Phase-1 `SortWorkerPool` worker count. + /// Centralizes the per-order stream-constructor dispatch so the pool-sizing + /// assertion can be parameterized. + fn streaming_phase1_pool_workers(order: SortOrder, sort_threads: Option) -> usize { + use noodles::sam::Header; + + let header = Header::default(); + let mut sorter = RawExternalSorter::new(order).threads(4); + if let Some(j) = sort_threads { + sorter = sorter.sort_threads(j); + } + match order { + SortOrder::Coordinate => sorter + .into_coordinate_stream(&header) + .expect("coordinate stream") + .pool + .num_workers(), + SortOrder::TemplateCoordinate => sorter + .into_template_coordinate_stream(&header) + .expect("template-coordinate stream") + .pool + .num_workers(), + SortOrder::Queryname(_) => { + match sorter.into_queryname_stream(&header).expect("queryname stream") { + QuerynameSortStream::Lex(s) => s.pool.num_workers(), + QuerynameSortStream::Natural(s) => s.pool.num_workers(), + } + } + } + } + + /// The streaming Phase-1 `SortWorkerPool` must be sized to `phase1_threads()`, + /// not `threads`. When `sort_threads` is set, the pool must reflect the cap; + /// when it is unset, the pool has the full `threads` count. One case per + /// (order, split-state) so a failure names the exact combination. + #[rstest::rstest] + #[case::coordinate_split(SortOrder::Coordinate, Some(1), 1)] + #[case::coordinate_nosplit(SortOrder::Coordinate, None, 4)] + #[case::template_coordinate_split(SortOrder::TemplateCoordinate, Some(1), 1)] + #[case::template_coordinate_nosplit(SortOrder::TemplateCoordinate, None, 4)] + #[case::queryname_split( + SortOrder::Queryname(crate::keys::QuerynameComparator::Lexicographic), + Some(1), + 1 + )] + #[case::queryname_nosplit( + SortOrder::Queryname(crate::keys::QuerynameComparator::Lexicographic), + None, + 4 + )] + fn streaming_pool_is_capped_to_phase1_threads( + #[case] order: SortOrder, + #[case] sort_threads: Option, + #[case] expected_workers: usize, + ) { + assert_eq!( + streaming_phase1_pool_workers(order, sort_threads), + expected_workers, + "order {order:?} sort_threads {sort_threads:?}: stream pool must size to phase1_threads" + ); + } + #[test] fn test_raw_sorter_memory_limit() { let sorter = RawExternalSorter::new(SortOrder::Coordinate).memory_limit(256 * 1024 * 1024); diff --git a/crates/fgumi-sort/src/worker_pool.rs b/crates/fgumi-sort/src/worker_pool.rs index 2868ab214..98e5ee50f 100644 --- a/crates/fgumi-sort/src/worker_pool.rs +++ b/crates/fgumi-sort/src/worker_pool.rs @@ -742,6 +742,16 @@ pub(crate) struct SharedPipelineState { /// Current phase: 0=shutdown, 1=Phase1, 2=Phase2, 255=Legacy. pub(crate) phase: AtomicU8, + /// Number of workers permitted to be active in the current phase. Workers + /// with `worker_id >= active_worker_limit` idle (backoff) instead of taking + /// work. Lets the sort driver run Phase 1 (accumulate/sort/spill) on fewer + /// threads than Phase 2 (merge/write) — e.g. to cede cores to an upstream + /// producer in a pipeline. Defaults to `num_workers` (all active), so the + /// behavior is unchanged unless the driver calls `set_active_workers`. + /// Capped workers re-check this within `MAX_BACKOFF_US`, so raising the + /// limit reactivates them without an explicit wake. + pub(crate) active_worker_limit: AtomicUsize, + // --- Phase 1 queues --- /// Input BAM file (Mutex because only one worker reads at a time). pub(crate) input_file: std::sync::Mutex>>, @@ -831,6 +841,7 @@ impl SharedPipelineState { Self { phase: AtomicU8::new(phase::LEGACY), + active_worker_limit: AtomicUsize::new(num_workers), input_file: std::sync::Mutex::new(None), input_eof: AtomicBool::new(false), @@ -1158,6 +1169,25 @@ impl SortWorkerPool { break; } + // 1b. Honor the per-phase active-worker cap: workers above the limit + // take no NEW work, but must still drain any items they already + // hold — held items are per-worker, so a capped worker that + // froze with held work would strand that output (no other worker + // can advance it). So advance held items first, then idle without + // acquiring fresh work. Capped workers re-check the limit within + // MAX_BACKOFF_US, so a later raise reactivates them. Default + // limit == num_workers ⇒ no worker is ever capped (no-op). + if worker.worker_id >= shared.active_worker_limit.load(Ordering::Acquire) { + if Self::try_advance_all_held(shared, worker) { + worker.backoff_us = MIN_BACKOFF_US; + } else { + sleep_with_jitter(worker.backoff_us, worker.worker_id, worker.idle_iter); + worker.idle_iter = worker.idle_iter.wrapping_add(1); + worker.backoff_us = (worker.backoff_us * 2).min(MAX_BACKOFF_US); + } + continue; + } + // 2. Check phase completion — wait for next phase, only exit on SHUTDOWN. // Workers must survive across Phase 1 → Phase 2 transitions. if Self::is_phase_complete(shared, current_phase) && !worker.has_any_held_items() { @@ -1951,6 +1981,16 @@ impl SortWorkerPool { self.shared.phase.store(new_phase, Ordering::Release); } + /// Cap the number of workers active in the current phase to `n` (clamped to + /// `[1, num_workers]`). Workers with `worker_id >= n` idle until the limit is + /// raised. Used to run Phase 1 on fewer threads than Phase 2; raising the + /// limit reactivates idled workers within `MAX_BACKOFF_US` (no explicit wake + /// needed). Defaults to `num_workers` (all active) when never called. + pub fn set_active_workers(&self, n: usize) { + let n = n.clamp(1, self.num_workers); + self.shared.active_worker_limit.store(n, Ordering::Release); + } + /// Set the input file for Phase 1 reading. /// /// Must be called before `set_phase(PHASE1)`. @@ -2369,6 +2409,94 @@ mod tests { pool.shutdown(); } + #[test] + fn set_active_workers_clamps() { + let pool = SortWorkerPool::new(4, 1, 6, crate::codec::SpillCodec::Bgzf); + // Clamp to [1, num_workers]. + pool.set_active_workers(0); + assert_eq!(pool.shared.active_worker_limit.load(Ordering::Acquire), 1); + pool.set_active_workers(100); + assert_eq!(pool.shared.active_worker_limit.load(Ordering::Acquire), 4); + pool.set_active_workers(2); + assert_eq!(pool.shared.active_worker_limit.load(Ordering::Acquire), 2); + pool.shutdown(); + } + + #[test] + fn active_worker_limit_caps_then_reactivates() { + // Reactivation must be observable on the SAME pool: a fresh pool can never + // prove that raising the cap re-enables workers idled by an earlier cap. + // With a cap of 2 over 6 workers, only workers 0..2 may run compress steps; + // workers >= 2 stay idle. Raising the cap to 6 and running a second batch + // must bring those workers online — and because per-worker counts are + // cumulative and they did nothing under the cap, any work they show after + // the second batch was done exclusively in that batch. + let cs = SortStep::CompressSpill as usize; + + // Run one batch on `pool` (moved in, returned out so the next batch can + // reuse it) and report the cumulative per-worker CompressSpill counts. + let run_batch = |pool: SortWorkerPool, num_jobs: usize| -> (SortWorkerPool, Vec) { + let (result_tx, result_rx) = pool.compress_result_channel(); + let submit_tx = result_tx.clone(); + let submit_handle = std::thread::spawn(move || { + for i in 0..num_jobs { + pool.submit_compress(CompressJob { + data: vec![b'X'; 4096], + serial: i as u64, + result_tx: submit_tx.clone(), + codec: SpillCodec::Bgzf, + }); + } + drop(submit_tx); + pool + }); + drop(result_tx); + let mut received = 0; + while result_rx.recv().is_ok() { + received += 1; + } + assert_eq!(received, num_jobs); + let pool = submit_handle.join().expect("submit thread should not panic"); + let counts: Vec = (0..6) + .map(|t| pool.pipeline_stats.per_thread_step_counts[t][cs].load(Ordering::Relaxed)) + .collect(); + (pool, counts) + }; + + let pool = SortWorkerPool::new(6, 1, 6, crate::codec::SpillCodec::Bgzf); + + // Batch 1: capped at 2 — only workers 0..2 may run. + pool.set_active_workers(2); + let (pool, after_capped) = run_batch(pool, 300); + assert!( + after_capped.iter().filter(|&&c| c > 0).count() <= 2, + "cap=2 must bound active workers; per-worker counts {after_capped:?}" + ); + assert!( + after_capped[2..].iter().all(|&c| c == 0), + "workers >= cap must do no work; per-worker counts {after_capped:?}" + ); + assert_eq!(after_capped.iter().sum::(), 300, "all batch-1 jobs accounted for"); + + // Batch 2: raise the cap on the SAME pool and run again. + pool.set_active_workers(6); + let (pool, after_full) = run_batch(pool, 300); + // Workers >= the old cap did nothing in batch 1, so any cumulative work + // they now show was processed solely in batch 2 — i.e. reactivation worked. + let reactivated_work: u64 = after_full[2..].iter().sum(); + assert!( + reactivated_work > 0, + "raising the cap on the same pool must reactivate workers >= old cap; \ + after_capped {after_capped:?}, after_full {after_full:?}" + ); + assert_eq!( + after_full.iter().sum::(), + 600, + "all jobs across both batches accounted for; per-worker counts {after_full:?}" + ); + pool.shutdown(); + } + #[test] fn test_pool_stats() { let pool = SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf); From 3c74508a616f4e53bfc4f290d8617764014d7db4 Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Thu, 25 Jun 2026 14:38:49 -0700 Subject: [PATCH 2/4] feat(sort): cap streaming Phase-2 decompress concurrency via admission counter In a multi-stage runall chain, Phase 2 (decompress/merge) runs as framework pipeline steps rather than on the sort worker pool, and the merge is serial, so the only parallel CPU consumer is SortSpillDecompress. Add a global atomic admission counter (with_max_concurrency) shared across all worker clones that caps concurrent decompress workers. Only the decompress branch is gated, so event forwarding can never stall the merge. Byte-identical: decompression order is deterministic regardless of concurrency. --- .../src/sort/spill_decompress.rs | 64 ++++++++++++++++++- .../src/sort/spill_decompress/tests.rs | 61 +++++++++++++++++- 2 files changed, 121 insertions(+), 4 deletions(-) diff --git a/crates/fgumi-pipeline-io/src/sort/spill_decompress.rs b/crates/fgumi-pipeline-io/src/sort/spill_decompress.rs index 4a16d848b..87326b57f 100644 --- a/crates/fgumi-pipeline-io/src/sort/spill_decompress.rs +++ b/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, max: Option) -> Option { + 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, +} +impl Drop for DecompressPermit { + fn drop(&mut self) { + self.active.fetch_sub(1, Ordering::AcqRel); + } +} + struct RegisteredSpill { slot: Arc, } @@ -33,6 +70,11 @@ pub struct SortSpillDecompress { block_dec: SpillBlockDecompressor, held: HeldSlot>, 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, + /// Maximum concurrent decompress workers. `None` ⇒ unbounded. + max_decompress: Option, } 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) -> 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. diff --git a/crates/fgumi-pipeline-io/src/sort/spill_decompress/tests.rs b/crates/fgumi-pipeline-io/src/sort/spill_decompress/tests.rs index a22bfec24..a4f557e40 100644 --- a/crates/fgumi-pipeline-io/src/sort/spill_decompress/tests.rs +++ b/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"); +} From 7470adb2590fb94080c85e1221f429743fd5ac77 Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Thu, 25 Jun 2026 14:38:50 -0700 Subject: [PATCH 3/4] feat(sort-cli): --sort-threads / --merge-threads per-phase flags on SortOptions Expose the per-phase thread knobs on SortOptions so standalone fgumi sort accepts --sort-threads/--merge-threads and runall forwards them as --sort::sort-threads / --sort::merge-threads via the #[multi_options] macro. Each defaults to --threads. Wire them through SortStepCaptures into the RawExternalSorter builders and log the split when set. --- crates/fgumi-sort-cli/src/chains.rs | 20 +++++++- crates/fgumi-sort-cli/src/sort.rs | 77 ++++++++++++++++++++++++++++- 2 files changed, 94 insertions(+), 3 deletions(-) diff --git a/crates/fgumi-sort-cli/src/chains.rs b/crates/fgumi-sort-cli/src/chains.rs index 8f86b3d9e..784fb3bc8 100644 --- a/crates/fgumi-sort-cli/src/chains.rs +++ b/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 { 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()); diff --git a/crates/fgumi-sort-cli/src/sort.rs b/crates/fgumi-sort-cli/src/sort.rs index dce26b822..0fb9a0289 100644 --- a/crates/fgumi-sort-cli/src/sort.rs +++ b/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, + /// 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, + + /// 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, + /// 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::)); @@ -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)] From 7d219fcd7af83aec4938321d6a31f955dd7f2a44 Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Thu, 25 Jun 2026 14:38:50 -0700 Subject: [PATCH 4/4] feat(runall): honor per-phase sort threads in sole-stage and streaming chains Resolve per-phase worker counts from the forwarded SortOptions (defaulting to --threads) in both runall sort paths. The sole-stage (monolithic) path passes them through SortStepCaptures; the streaming chain sizes the Phase-1 pool with sort_threads and caps Phase-2 decompress concurrency via SortSpillDecompress::with_max_concurrency (merge_threads). --- src/lib/pipeline/chains/builder.rs | 75 +++++++++++++++++++++++++++++- 1 file changed, 73 insertions(+), 2 deletions(-) diff --git a/src/lib/pipeline/chains/builder.rs b/src/lib/pipeline/chains/builder.rs index f4e9bb7de..62024cff9 100644 --- a/src/lib/pipeline/chains/builder.rs +++ b/src/lib/pipeline/chains/builder.rs @@ -2084,6 +2084,11 @@ impl<'a> ChainBuilder<'a> { let num_sorter_threads = self.spec.threading.num_threads(); + // Per-phase worker counts default to --threads; runall forwards + // --sort::sort-threads / --sort::merge-threads via SortOptions. + let num_phase1_threads = resolve_phase_threads(sort.sort_threads, num_sorter_threads); + let num_phase2_threads = resolve_phase_threads(sort.merge_threads, num_sorter_threads); + let effective_memory = resolve_memory_budget( sort.max_memory, sort.memory_reserve, @@ -2093,7 +2098,15 @@ impl<'a> ChainBuilder<'a> { let timer = crate::logging::OperationTimer::new("Sorting BAM"); - log_sort_start(&sort, &input_path, &output_path, num_sorter_threads, effective_memory); + log_sort_start( + &sort, + &input_path, + &output_path, + num_sorter_threads, + num_phase1_threads, + num_phase2_threads, + effective_memory, + ); let stats_slot = Arc::new(parking_lot::Mutex::new(None::)); @@ -2102,6 +2115,8 @@ impl<'a> ChainBuilder<'a> { input_path: input_path.clone(), output_path: output_path.clone(), num_sorter_threads, + num_phase1_threads, + num_phase2_threads, effective_memory, output_compression: self.spec.compression.compression_level, command_line: self.spec.command_line.clone(), @@ -2195,6 +2210,18 @@ impl<'a> ChainBuilder<'a> { ); } let num_threads = self.spec.threading.num_threads(); + // Per-phase worker counts default to --threads; runall forwards + // --sort::sort-threads / --sort::merge-threads via SortOptions. + // Phase 1 caps the streaming sort/spill worker pool; Phase 2 caps + // concurrent spill-decompression (the merge step itself is serial). + let num_phase1_threads = resolve_phase_threads(sort.sort_threads, num_threads); + let num_phase2_threads = resolve_phase_threads(sort.merge_threads, num_threads); + if sort.sort_threads.is_some() || sort.merge_threads.is_some() { + log::debug!( + "streaming sort: per-phase thread split (phase1={num_phase1_threads}, \ + phase2={num_phase2_threads}); phase2 caps decompress concurrency, merge is serial" + ); + } // Reuse the same helper as the standalone path so `--sort::memory-per-thread` // is honored (the previous inline `× num_threads` ignored it). let total_memory = resolve_memory_budget( @@ -2207,9 +2234,16 @@ impl<'a> ChainBuilder<'a> { let sort_order = SortOrder::TemplateCoordinate; // NB: the spill codec is pinned to BGZF inside `build_stream` // (`SortSpillDecompress` is BGZF-block-only); see the comment there. + // NB: `sort_threads` sizes the streaming Phase-1 pool; `merge_threads` + // is inert on the streaming sorter (its Phase 2 runs on the framework + // scheduler, not this pool) — streaming Phase-2 concurrency is capped + // on `SortSpillDecompress` below. It is set here only for config + // symmetry with the standalone `sort()` path. let mut sorter = RawExternalSorter::new(sort_order) .memory_limit(total_memory) .threads(num_threads.max(1)) + .sort_threads(num_phase1_threads) + .merge_threads(num_phase2_threads) .output_compression(1) .temp_compression(sort.temp_compression) .cell_tag(SamTag::CB); @@ -2225,7 +2259,8 @@ impl<'a> ChainBuilder<'a> { SortAndSpill::from_sorter(sorter, &self.header, self.tuning.per_step_byte_limit) .map_err(|e| anyhow!("SortAndSpill::from_sorter: {e}"))? .with_affinity(affinity); - let decompress = SortSpillDecompress::new(self.tuning.per_step_byte_limit); + let decompress = SortSpillDecompress::new(self.tuning.per_step_byte_limit) + .with_max_concurrency(Some(num_phase2_threads)); let merge = SortMerge::new(sort_order, self.tuning.per_step_byte_limit).with_affinity(affinity); @@ -3835,3 +3870,39 @@ impl<'a> ChainBuilder<'a> { Ok(()) } } + +/// Resolve a sort phase's worker-thread count. +/// +/// A per-phase override (`--sort::sort-threads` / `--sort::merge-threads`, or +/// the standalone `--sort-threads` / `--merge-threads`) is used as-is when +/// present; otherwise the phase falls back to the chain's base sorter-thread +/// count (`--threads`). The result is clamped to at least 1 so a `0` override +/// never disables a phase. This is the single source of truth for both the +/// sole-stage (`SortStepCaptures`) and streaming (`RawExternalSorter`) sort +/// paths in `add_sort`. +fn resolve_phase_threads(override_threads: Option, num_sorter_threads: usize) -> usize { + override_threads.unwrap_or(num_sorter_threads).max(1) +} + +#[cfg(test)] +mod tests { + use super::resolve_phase_threads; + + /// Pin the per-phase thread resolution contract shared by the standalone and + /// streaming sort branches in `add_sort`: each phase falls back to the base + /// sorter-thread count when no explicit override is supplied, and is clamped + /// to at least 1. Exercises the production `resolve_phase_threads` helper + /// directly so it cannot drift from the code `add_sort` actually runs. + #[test] + fn sole_stage_sort_resolves_per_phase_threads() { + let num_sorter_threads = 8usize; + + // Explicit overrides are used as-is (clamped ≥ 1). + assert_eq!(resolve_phase_threads(Some(1), num_sorter_threads), 1); + assert_eq!(resolve_phase_threads(Some(2), num_sorter_threads), 2); + // Zero is clamped to 1. + assert_eq!(resolve_phase_threads(Some(0), num_sorter_threads), 1); + // None falls back to num_sorter_threads. + assert_eq!(resolve_phase_threads(None, num_sorter_threads), 8); + } +}