From 72d0e2b4765707776f1a1deb56a23bdc21c4826d Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Sat, 25 Jul 2026 10:56:15 -0700 Subject: [PATCH] refactor(sort): pick the compressor from the job, not the pipeline phase MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit One `compress_queue` carries both spill and output blocks, and `CompressJob` recorded only the *codec* (BGZF vs zstd), never which of the two it was. So a worker popping a job re-derived the compression level from `shared.phase`, a mutable atomic that advances independently of what is already sitting in the queue. Every block that outlived a phase transition was compressed at the wrong level, silently. That single coupling produced the same bug twice, in opposite directions: - Output blocks popped outside Phase 2 went through the spill compressor at `temp_compression`, so `--compression-level` was silently discarded. Fixed twice at the symptom layer — by entering Phase 2 even for an in-memory sort (#633), then by keeping the pool in Phase 2 across the writer's `finish()` (#653). - Phase 1 spill blocks still queued when `begin_phase2` fired went through the output compressor at `output_compression`. This one was never a live bug: it was held off only by `drain_pending_spill` happening to sit immediately before `enter_output_phase`, and was written down as a caller obligation on `begin_phase2` rather than enforced. Its cost is wasted CPU and disk, not wrong bytes, since any BGZF level decodes. The identity was already known for certain at submit time, by the one caller who cannot be wrong about it: there is a single production submit site, `StagingBuffer::flush`, and its two constructors are statically distinct — the output writer and the spill writer. So record it rather than infer it. `CompressJob` gains a `target: CompressTarget { Spill, Output }` beside the existing `codec`; `StagingBuffer` carries it and stamps every job it submits; `try_compress` selects `compressor` or `output_compressor` from `job.target`. Nothing about the choice reads shared mutable state, so a block is compressed at the level its writer asked for however long it waited in the queue and whatever phase the pool is in when a worker gets to it. This is the second form of that selection. `5df7e176`, which introduced the worker pool, read it from `shared.phase` at pop time and then narrowed it to the dispatched `SortStep` to close the window where `set_phase` fired between the pop and the choice. But a step is itself chosen from the phase, so blocks that outlived a transition were still mis-compressed — the window was narrowed, not closed. With the level travelling on the job, `SortStep::CompressSpill` and `CompressOutput` did the same thing, so they collapse into one `Compress`. There was only ever one queue; two steps implied two, and the per-step stats buckets they fed ("CmpSpl" / "CmpOut") could not be trusted to mean what they said. `SortStep::COUNT` drops from 5 to 4. `SortStep` is `pub` but `worker_pool` is `pub(crate)`, so this is not an API change. Also drops `begin_phase2`'s caller obligation, which no longer exists, and pins zstd's spill-only invariant now that it is expressible: there is one zstd compressor per worker, fixed at `temp_compression`, because the output BAM is always BGZF. That one asserts unconditionally rather than in debug only — an `Output` job reaching it would be the same silent wrong-level output this mechanism exists to prevent, and the check is one compare inside the zstd arm, which the BGZF path never evaluates. The ordering #653 introduced — release the drained sources, finalize, then leave Phase 2 — is kept, but for the reason that survives: not holding the merge's file descriptors and 2 MiB-per-chunk reorder buffers across a slow `finish`. Its compression rationale is gone, and its docs and tests say so. Reverting that ordering now leaves `test_temp_compression_does_not_reach_the_ output_bam` passing, which is the check that this commit fixes the cause rather than a third symptom. `test_compress_target_decides_level_regardless_of_phase` sweeps the phase across `LEGACY`/`PHASE1`/`PHASE2` and asserts a `Spill` job stays stored at `temp_compression = 0` while an `Output` job compresses at `output_compression = 9`. All three cases fail if the compressor is selected from the phase, including the mirror direction that had no coverage before. No measurable throughput cost. The job grows by one byte in a struct that already owns a `Vec`, and the per-block work is one `match` on a `Copy` enum replacing one `match` on the step; scheduling is untouched, since the two collapsed steps shared an eligibility predicate and neither was ever exclusive (only `ReadInputBlocks` is). Measured on a spilling coordinate sort of a 3,689,310-record BAM (idt-cfdna, 162 MB, 4 spill chunks, 8 threads, `-m 256m`, output level 6, temp level 1), 10 interleaved reps per binary, baseline being #653's merge commit (`ef0f6f8b`, identical content in these files): CPU (user+sys) baseline 11.497s ± 0.312 candidate 11.326s ± 0.462 -1.5% wall baseline 5.663s ± 0.941 candidate 5.910s ± 1.044 +4.4% Welch t on CPU time is -0.97, so the difference is not significant; wall clock on this host is too noisy (σ ≈ 1 s) to resolve anything under ~15%, which is why CPU time is the reported metric. All 20 runs exited 0 and wrote exactly 3,689,310 records. Byte-for-byte, the two binaries produce the same output: the sorted BAMs differ only in the `@PG` `CL:` field (it records the binary's path), and the record streams have identical checksums. --- crates/fgumi-sort/src/bgzf_io.rs | 52 +++- crates/fgumi-sort/src/external.rs | 107 +++---- crates/fgumi-sort/src/pooled_bam_writer.rs | 12 +- crates/fgumi-sort/src/pooled_chunk_writer.rs | 12 +- crates/fgumi-sort/src/worker_pool.rs | 276 +++++++++++++----- .../integration/test_output_compression.rs | 47 +-- 6 files changed, 358 insertions(+), 148 deletions(-) diff --git a/crates/fgumi-sort/src/bgzf_io.rs b/crates/fgumi-sort/src/bgzf_io.rs index d71d9e4a2..4a8cd9c8b 100644 --- a/crates/fgumi-sort/src/bgzf_io.rs +++ b/crates/fgumi-sort/src/bgzf_io.rs @@ -5,7 +5,9 @@ //! accumulating data into ~64KB blocks before submitting compression jobs. use crate::codec::SpillCodec; -use crate::worker_pool::{BufferPool, CompressJob, CompressResult, PermitPool, SortWorkerPool}; +use crate::worker_pool::{ + BufferPool, CompressJob, CompressResult, CompressTarget, PermitPool, SortWorkerPool, +}; use anyhow::Result; use crossbeam_channel::{Receiver, Sender}; use fgumi_bgzf::{BGZF_EOF, BGZF_MAX_BLOCK_SIZE}; @@ -39,6 +41,10 @@ pub(crate) struct StagingBuffer { result_tx: Sender, permit_pool: Arc, codec: SpillCodec, + /// Stamped onto every job this buffer submits, so the worker that pops it + /// compresses at the level this writer asked for regardless of the pool's + /// phase at that moment. + target: CompressTarget, } impl StagingBuffer { @@ -49,6 +55,7 @@ impl StagingBuffer { result_tx: Sender, permit_pool: Arc, codec: SpillCodec, + target: CompressTarget, ) -> Self { Self { pool, @@ -57,6 +64,7 @@ impl StagingBuffer { result_tx, permit_pool, codec, + target, } } @@ -123,6 +131,7 @@ impl StagingBuffer { serial, result_tx: self.result_tx.clone(), codec: self.codec, + target: self.target, }); Ok(()) } @@ -331,7 +340,13 @@ mod tests { io_writer_loop(writer, result_rx, buffer_pool, pp, codec, None) }); - let mut staging = StagingBuffer::new(Arc::clone(&pool), result_tx, permit_pool, codec); + let mut staging = StagingBuffer::new( + Arc::clone(&pool), + result_tx, + permit_pool, + codec, + CompressTarget::Spill, + ); staging.write_chunked(data).unwrap(); staging.flush().unwrap(); drop(staging); // closes result_tx senders → io_writer_loop exits @@ -352,7 +367,13 @@ mod tests { let (result_tx, _result_rx) = pool.compress_result_channel(); let permit_pool = make_permit_pool(&pool); - let mut staging = StagingBuffer::new(Arc::clone(&pool), result_tx, permit_pool, codec); + let mut staging = StagingBuffer::new( + Arc::clone(&pool), + result_tx, + permit_pool, + codec, + CompressTarget::Spill, + ); // Flush with empty buffer: should not submit a compress job staging.flush().unwrap(); @@ -373,7 +394,13 @@ mod tests { let pool = Arc::new(SortWorkerPool::new(1, 1, 6, codec)); let (result_tx, _result_rx) = pool.compress_result_channel(); let permit_pool = make_permit_pool(&pool); - let mut staging = StagingBuffer::new(Arc::clone(&pool), result_tx, permit_pool, codec); + let mut staging = StagingBuffer::new( + Arc::clone(&pool), + result_tx, + permit_pool, + codec, + CompressTarget::Spill, + ); assert!(!staging.is_full(), "empty buffer should not be full"); staging.buf().extend(vec![0u8; BGZF_MAX_BLOCK_SIZE]); @@ -404,7 +431,13 @@ mod tests { io_writer_loop(writer, result_rx, buffer_pool, pp, codec, None) }); - let mut staging = StagingBuffer::new(Arc::clone(&pool), result_tx, permit_pool, codec); + let mut staging = StagingBuffer::new( + Arc::clone(&pool), + result_tx, + permit_pool, + codec, + CompressTarget::Spill, + ); staging.write_chunked(&large).unwrap(); staging.flush().unwrap(); drop(staging); @@ -452,9 +485,16 @@ mod tests { serial: 1, result_tx: result_tx.clone(), codec, + target: CompressTarget::Spill, }); permit_pool.acquire().unwrap(); - pool.submit_compress(CompressJob { data: data1, serial: 0, result_tx, codec }); + pool.submit_compress(CompressJob { + data: data1, + serial: 0, + result_tx, + codec, + target: CompressTarget::Spill, + }); // Wait for both compress results to be received by io_writer_loop io_handle.join().unwrap().unwrap(); diff --git a/crates/fgumi-sort/src/external.rs b/crates/fgumi-sort/src/external.rs index 50ae96cc5..44aab6362 100644 --- a/crates/fgumi-sort/src/external.rs +++ b/crates/fgumi-sort/src/external.rs @@ -1679,32 +1679,25 @@ impl Phase2Guard<'_, K> { /// Release the drained merge sources, finalize the output, then leave Phase 2. /// - /// **The output must be finalized inside Phase 2.** Phase 2 is what puts - /// [`SortStep::CompressOutput`](crate::worker_pool::SortStep::CompressOutput) - /// on a worker's priority list, and that step is what selects - /// `output_compressor`. In `LEGACY` the only compress step workers schedule - /// is `CompressSpill`, so every output block still queued — plus the - /// writer's final flush — would be compressed at `temp_compression`, and the - /// tail of the BAM would silently ignore the caller's `output_compression`. - /// This method exists so that ordering cannot be got wrong at a call site: - /// `finish` runs between the source release and the return to `LEGACY`. + /// Releasing the sources first is the point: it keeps a slow `finish` from + /// holding the merge's file descriptors and per-file reorder buffers (a 2 MiB + /// `BufReader` per spill file) alive while it drains. Leaving Phase 2 last + /// then costs nothing, and keeps the pool's phase an accurate description of + /// what the pool is doing for as long as it is doing it. /// - /// Releasing the sources first is what keeps a slow `finish` from holding - /// the merge's file descriptors and per-file reorder buffers (a 2 MiB - /// `BufReader` per spill file) alive while it drains. - /// - /// The underlying coupling — "which compressor a block gets" being a - /// function of the pool's *global phase* rather than of the block's own - /// origin — is what makes this ordering load-bearing at all. Tagging each - /// `CompressJob` with its target at submit time would retire this contract - /// (and the mirror-image one on `set_phase`); that is a separate change. + /// Note what this ordering is *not* responsible for. The output's compression + /// level does not depend on it: each block carries its own + /// [`CompressTarget`](crate::worker_pool::CompressTarget), stamped by the + /// writer that produced it, so trailing output blocks are compressed at + /// `output_compression` whatever phase the pool is in when a worker pops + /// them. That was not always true — the compressor used to be selected from + /// the pool's phase at pop time, which is why finalizing after the return to + /// `LEGACY` silently wrote the tail of the BAM at `temp_compression`. /// /// Takes `self` by value: finalizing the output is the guard's terminal - /// operation, so a second call — which would run `finish` with the pool - /// already back in `LEGACY`, silently compressing the output at - /// `temp_compression` and reporting success — is a compile error rather than - /// a runtime check. The remaining `Drop` at the end of this method is a - /// no-op, since `deactivate` has already cleared `active`. + /// operation, so a second call is a compile error rather than a runtime + /// check. The `Drop` at the end of this method is a no-op, since + /// `deactivate` has already cleared `active`. fn finish_output(mut self, finish: impl FnOnce() -> Result) -> Result { self.release_sources(); let finished = finish(); @@ -1837,14 +1830,18 @@ impl RawExternalSorter { /// /// Every sort path calls this exactly once, immediately after ingest /// completes, to hand the pool over to Phase 2 and lift the active-worker cap - /// to `merge_threads`. Centralized here so the invariant cannot drift between - /// sort orders: it is precisely the transition whose omission left the - /// in-memory output writing at the wrong compression level. + /// to `merge_threads`. Centralized here so the transition cannot drift + /// between sort orders. /// /// Workers idled by the Phase-1 cap re-check within `MAX_BACKOFF_US`, so no /// explicit wake is needed; capped workers always drained their held items, so - /// raising the cap can never be required for correctness, only for width. The - /// phase half is what routes output blocks to `output_compressor` — see + /// raising the cap can never be required for correctness, only for width. + /// + /// Omitting this call once left in-memory sorts writing their output at + /// `temp_compression`, because the phase used to decide which compressor a + /// block went through. It no longer does — blocks carry their own + /// [`CompressTarget`](crate::worker_pool::CompressTarget) — so what the + /// phase buys now is Phase 2 scheduling and the wider worker cap. See /// [`SortWorkerPool::begin_phase2`]. fn enter_output_phase(&self, pool: &SortWorkerPool) { pool.begin_phase2(self.phase2_threads()); @@ -3491,16 +3488,17 @@ impl RawExternalSorter { /// single-sourced and the two paths cannot drift. Returns the sources plus /// an RAII [`Phase2Guard`] (borrowing `pool`); callers finish through /// [`Phase2Guard::finish_output`], and `Drop` resets Phase 2 on any path - /// that does not get that far. Phase 2 is armed whether or not there are disk - /// chunks: it is what makes workers reach for `output_compressor` - /// (`SortStep::CompressOutput`) rather than the Phase 1 spill compressor, so - /// leaving an all-in-memory merge in Phase 1 would silently write the output - /// BAM at `temp_compression` and discard the caller's `output_compression`. - /// With no spill files `try_phase2_file_work` sees an empty file vector and - /// returns immediately, so the only Phase 2 work available is output - /// compression. The consumer is created only for disk sources — it drains the - /// pool's per-file reorder buffers, and there is nothing to drain when - /// everything is already in memory. + /// that does not get that far. + /// + /// Phase 2 is armed whether or not there are disk chunks: it is what lets + /// workers schedule `SortStep::Phase2FileWork` and reach the merge's + /// per-file reorder buffers. With no spill files `try_phase2_file_work` sees + /// an empty file vector and returns immediately, so the only work left is + /// compression — which is correct in any phase, since each block carries its + /// own [`CompressTarget`](crate::worker_pool::CompressTarget). The consumer + /// is created only for disk sources — it drains the pool's per-file reorder + /// buffers, and there is nothing to drain when everything is already in + /// memory. fn setup_phase2_merge<'a, K: RawSortKey + Default + 'static>( chunk_files: &[PathBuf], memory_chunks: MemorySources, @@ -3518,8 +3516,10 @@ impl RawExternalSorter { let sources = Self::build_chunk_sources::(chunk_files, memory_chunks, pool)?; debug!("Merging from {} sources...", sources.len()); - // Activate Phase 2 for the whole merge, spill files or not, so the output - // BAM is compressed at `output_compression` rather than `temp_compression`. + // Activate Phase 2 for the whole merge, spill files or not, so workers can + // schedule `SortStep::Phase2FileWork` and reach the per-file reorder + // buffers. Compression level is no longer a reason to arm it — each block + // carries its own `CompressTarget` and is compressed correctly in any phase. // The consumer is created only for disk sources: it holds an Arc snapshot of // the pool's per-file Phase 2 state — workers populate per-file reorder // buffers, the consumer pops from them. There is nothing to consume when @@ -4827,19 +4827,22 @@ mod tests { } } - /// The output must be finalized while the pool is still in Phase 2. + /// The merge sources are released *before* the output is finalized, and the + /// pool leaves Phase 2 only *after*. /// - /// This is the invariant [`Phase2Guard::finish_output`] exists to enforce, - /// and the one the bug violated: the trailing output blocks are compressed - /// while the writer is being finished, and `CompressOutput` — the step that - /// selects `output_compressor` — is only on a worker's priority list in - /// Phase 2. Finalizing in `LEGACY` sends those blocks through the spill - /// compressor at `temp_compression`. + /// [`Phase2Guard::finish_output`] exists for the first half: a `finish` that + /// blocks draining the writer must not still be pinning the merge's file + /// descriptors and 2 MiB-per-chunk reorder buffers. So the observation is + /// made from *inside* the finalize closure, where the real writer's + /// `finish()` runs — not merely before and after it. /// - /// So the assertion is made from *inside* the finalize closure, where the - /// real writer's `finish()` would run — not merely before and after it. + /// The phase is checked in the same place, but note it is no longer + /// load-bearing for compression: blocks carry their own `CompressTarget`, so + /// the tail is written at `output_compression` in any phase. What it pins is + /// that the pool's phase still describes what the pool is doing while it is + /// doing it. #[test] - fn test_phase2_guard_finishes_output_while_still_in_phase2() { + fn test_phase2_guard_releases_sources_before_finalizing_output() { use crate::keys::RawCoordinateKey; use crate::worker_pool::phase; @@ -4862,8 +4865,8 @@ mod tests { assert_eq!( observed, (phase::PHASE2, 0), - "the output must be finalized in Phase 2 (so its blocks are compressed at \ - output_compression) with the drained spill files already unpublished" + "the output must be finalized with the drained spill files already unpublished, and \ + before the pool leaves Phase 2" ); // `finish_output` consumed the guard, so its `Drop` has already run — a // second deactivate must not have disturbed the phase it just set. diff --git a/crates/fgumi-sort/src/pooled_bam_writer.rs b/crates/fgumi-sort/src/pooled_bam_writer.rs index 293a6d23a..3c8554731 100644 --- a/crates/fgumi-sort/src/pooled_bam_writer.rs +++ b/crates/fgumi-sort/src/pooled_bam_writer.rs @@ -22,7 +22,7 @@ use crate::bgzf_io::{BlockOffset, StagingBuffer, io_writer_loop}; use crate::codec::SpillCodec; -use crate::worker_pool::{CompressResult, PermitPool, SortWorkerPool}; +use crate::worker_pool::{CompressResult, CompressTarget, PermitPool, SortWorkerPool}; use anyhow::Result; use crossbeam_channel::{Receiver, bounded, unbounded}; use fgumi_bam_io::BaiBuilder; @@ -124,7 +124,15 @@ impl PooledBamWriter { io_writer_loop(writer, result_rx, buffer_pool, pp, SpillCodec::Bgzf, block_offset_tx) }); - let mut staging = StagingBuffer::new(pool, result_tx, permit_pool, SpillCodec::Bgzf); + // `CompressTarget::Output`: every block this writer submits is the sort's + // output BAM and must be compressed at `output_compression`. + let mut staging = StagingBuffer::new( + pool, + result_tx, + permit_pool, + SpillCodec::Bgzf, + CompressTarget::Output, + ); // Write BAM header into a temporary buffer then flush in BGZF-sized chunks. // Headers can exceed BGZF_MAX_BLOCK_SIZE; write_chunked handles the splitting. diff --git a/crates/fgumi-sort/src/pooled_chunk_writer.rs b/crates/fgumi-sort/src/pooled_chunk_writer.rs index e57555e0a..3100aec9f 100644 --- a/crates/fgumi-sort/src/pooled_chunk_writer.rs +++ b/crates/fgumi-sort/src/pooled_chunk_writer.rs @@ -25,7 +25,7 @@ use crate::bgzf_io::{StagingBuffer, io_writer_loop}; use crate::codec::{SpillCodec, ZSPILL_MAGIC}; use crate::keys::RawSortKey; -use crate::worker_pool::{CompressResult, PermitPool, SortWorkerPool}; +use crate::worker_pool::{CompressResult, CompressTarget, PermitPool, SortWorkerPool}; use anyhow::Result; use crossbeam_channel::bounded; use fgumi_bgzf::BGZF_MAX_BLOCK_SIZE; @@ -78,7 +78,15 @@ impl PooledChunkWriter { thread::spawn(move || io_writer_loop(writer, result_rx, buffer_pool, pp, codec, None)); Ok(Self { - staging: Some(StagingBuffer::new(pool, result_tx, permit_pool, codec)), + // `CompressTarget::Spill`: every block this writer submits is a + // Phase 1 spill chunk and must be compressed at `temp_compression`. + staging: Some(StagingBuffer::new( + pool, + result_tx, + permit_pool, + codec, + CompressTarget::Spill, + )), key_buf: Vec::new(), io_handle: Some(io_handle), _phantom: PhantomData, diff --git a/crates/fgumi-sort/src/worker_pool.rs b/crates/fgumi-sort/src/worker_pool.rs index 06ecefc3b..c0be993b8 100644 --- a/crates/fgumi-sort/src/worker_pool.rs +++ b/crates/fgumi-sort/src/worker_pool.rs @@ -14,10 +14,14 @@ //! # Design //! //! - **Phase-aware scheduling**: Workers check the current phase to pick eligible steps. -//! Phase 1: `DecompressInput` > `ReadInputBlocks` > `CompressSpill`. +//! Phase 1: `DecompressInput` > `ReadInputBlocks` > `Compress`. //! Phase 2: `Phase2FileWork` (read+decompress the next block of any file that has -//! room in its reorder buffer) > `CompressOutput`. -//! - **Per-worker state**: Each worker owns an `InlineBgzfCompressor`, a +//! room in its reorder buffer) > `Compress`. +//! The phase orders the steps; it never decides which *compressor* a block goes +//! through — that is carried per-job as a `CompressTarget`. +//! - **Per-worker state**: Each worker owns two `InlineBgzfCompressor`s (one at +//! `temp_compression` for spill blocks, one at `output_compression` for output +//! blocks, selected by the job's `CompressTarget`), a //! `libdeflater::Decompressor`, a `zstd::bulk::Compressor`, and a //! `zstd::bulk::Decompressor` (the latter pair are used only when the //! selected spill codec is `Zstd`). Per-worker contexts avoid cross-thread @@ -133,11 +137,33 @@ const ZSTD_MAX_CLEVEL: u32 = 22; // Job and Result Types // ============================================================================ +/// Which of the pool's two compressors a queued block must go through, and so +/// which compression level it is written at. +/// +/// One `compress_queue` carries both spill and output blocks, so a worker +/// popping a job has to know which kind it has. That is recorded here, at submit +/// time, by the only code that knows it for certain — the staging buffer of the +/// writer that produced the block. It is deliberately *not* inferred from the +/// pool's current phase: the phase is shared mutable state that advances +/// independently of what is already sitting in the queue, so inferring from it +/// silently mis-compresses every block that outlives the transition (an output +/// block popped in Phase 1 or `LEGACY` at `temp_compression`, a Phase 1 spill +/// block popped after `begin_phase2` at `output_compression`). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CompressTarget { + /// A Phase 1 spill chunk — compress at `temp_compression`. + Spill, + /// The sort's output BAM — compress at `output_compression`. + Output, +} + /// A compression job: compress uncompressed data into one compressed block. /// /// The codec determines the output framing: /// - `SpillCodec::Bgzf` produces one BGZF block (header + deflate + footer). /// - `SpillCodec::Zstd` produces `[u32 LE frame-len][zstd frame]`. +/// +/// The target determines the compression *level* — see [`CompressTarget`]. pub struct CompressJob { /// Uncompressed data to compress. pub data: Vec, @@ -147,6 +173,8 @@ pub struct CompressJob { pub result_tx: Sender, /// Codec to use when compressing. pub codec: SpillCodec, + /// Which compressor (and therefore which level) this block belongs to. + pub target: CompressTarget, } /// Result of a compression job. @@ -188,17 +216,19 @@ pub enum SortStep { ReadInputBlocks = 0, /// Decompress input BGZF blocks during Phase 1. DecompressInput = 1, - /// Compress data for spill files during Phase 1. - CompressSpill = 2, + /// Compress one queued block — spill or output. + /// + /// There is a single `compress_queue`, so this is a single step. Which + /// compressor the popped block goes through is carried on the job as a + /// [`CompressTarget`], never derived from the step or the phase. + Compress = 2, /// Read+decompress one unit of work for some Phase 2 spill file (work-stealing). Phase2FileWork = 3, - /// Compress data for merge output during Phase 2. - CompressOutput = 4, } impl SortStep { /// Number of distinct sort steps. - pub const COUNT: usize = 5; + pub const COUNT: usize = 4; /// Short label for display. #[must_use] @@ -206,9 +236,8 @@ impl SortStep { match self { Self::ReadInputBlocks => "RdInp", Self::DecompressInput => "DecInp", - Self::CompressSpill => "CmpSpl", + Self::Compress => "Cmprs", Self::Phase2FileWork => "P2File", - Self::CompressOutput => "CmpOut", } } } @@ -280,9 +309,8 @@ impl SortPipelineStats { let all_steps = [ SortStep::ReadInputBlocks, SortStep::DecompressInput, - SortStep::CompressSpill, + SortStep::Compress, SortStep::Phase2FileWork, - SortStep::CompressOutput, ]; for &step in &all_steps { @@ -662,10 +690,10 @@ pub mod phase { /// # Phase-Aware Scheduling /// /// Workers check the current phase to determine eligible steps: -/// - **Phase 1**: `DecompressInput` > `ReadInputBlocks` > `CompressSpill` +/// - **Phase 1**: `DecompressInput` > `ReadInputBlocks` > `Compress` /// - **Phase 2**: `Phase2FileWork` (per-file work stealing: read next raw block /// batch from any file with FIFO room, OR decompress the next queued raw -/// block for any file whose reorder buffer has capacity) > `CompressOutput` +/// block for any file whose reorder buffer has capacity) > `Compress` pub struct SortWorkerPool { // Shared pipeline state (visible to workers and main thread) shared: Arc, @@ -918,10 +946,10 @@ fn get_sort_priorities(bp: &SortBackpressureState) -> &'static [SortStep] { } else if bp.compress_has_items && !bp.decompressed_input_low { // Spill compression is the bottleneck (13.7s at t4). Drain compress // while decompressed blocks are plentiful for the main thread. - &[SortStep::CompressSpill, SortStep::DecompressInput, SortStep::ReadInputBlocks] + &[SortStep::Compress, SortStep::DecompressInput, SortStep::ReadInputBlocks] } else { // Default/starving: feed the main thread first, compress if available - &[SortStep::DecompressInput, SortStep::ReadInputBlocks, SortStep::CompressSpill] + &[SortStep::DecompressInput, SortStep::ReadInputBlocks, SortStep::Compress] } } phase::PHASE2 => { @@ -931,13 +959,14 @@ fn get_sort_priorities(bp: &SortBackpressureState) -> &'static [SortStep] { // continue until all per-file reorder buffers empty. if bp.compress_has_items { // Drain output compression while we can; it's the writer-side bottleneck. - &[SortStep::CompressOutput, SortStep::Phase2FileWork] + &[SortStep::Compress, SortStep::Phase2FileWork] } else { - &[SortStep::Phase2FileWork, SortStep::CompressOutput] + &[SortStep::Phase2FileWork, SortStep::Compress] } } - // Legacy/transition: compress only (drain any remaining jobs) - _ => &[SortStep::CompressSpill], + // Legacy/transition: compress only (drain any remaining jobs, of either + // kind — each carries its own `CompressTarget`). + _ => &[SortStep::Compress], } } @@ -1372,7 +1401,7 @@ impl SortWorkerPool { || (shared.input_eof.load(Ordering::Acquire) && !shared.decompressed_input_done.load(Ordering::Acquire))) } - SortStep::CompressSpill | SortStep::CompressOutput => !shared.compress_queue.is_empty(), + SortStep::Compress => !shared.compress_queue.is_empty(), SortStep::Phase2FileWork => current_phase == phase::PHASE2, } } @@ -1386,18 +1415,7 @@ impl SortWorkerPool { match step { SortStep::ReadInputBlocks => Self::try_read_input_blocks(shared, worker), SortStep::DecompressInput => Self::try_decompress_input(shared, worker), - // Bind the compressor to the dispatched step, not shared.phase, to avoid - // the race where a worker pops a spill job then set_phase(PHASE2) fires - // before the compressor is chosen, causing spill data to be compressed - // at the output level. - SortStep::CompressSpill => { - Self::try_compress(shared, &mut worker.compressor, &mut worker.zstd_compressor) - } - SortStep::CompressOutput => Self::try_compress( - shared, - &mut worker.output_compressor, - &mut worker.zstd_compressor, - ), + SortStep::Compress => Self::try_compress(shared, worker), SortStep::Phase2FileWork => Self::try_phase2_file_work(shared, worker), } } @@ -1851,18 +1869,28 @@ impl SortWorkerPool { /// Try to pick up a compress job from the `ArrayQueue` (non-blocking). /// - /// The compressor is passed in by the caller (dispatched from `execute_step`) - /// so the choice is bound to the scheduled `SortStep`, not `shared.phase`. - /// This avoids the race where a worker pops a Phase-1 spill job and then - /// `set_phase(PHASE2)` fires before the compressor is selected. - fn try_compress( - shared: &SharedPipelineState, - bgzf_compressor: &mut InlineBgzfCompressor, - zstd_compressor: &mut ZstdCompressor<'static>, - ) -> StepResult { + /// The compressor is selected from the job's own [`CompressTarget`], which + /// the producing staging buffer stamped on it at submit time. Nothing about + /// this choice depends on shared mutable state, so a block is compressed at + /// the level its writer asked for no matter how long it waited in the queue + /// or which phase the pool is in when a worker gets to it. + /// + /// This is the third and final form of that selection. It was originally + /// read from `shared.phase` at pop time; that was narrowed to the dispatched + /// `SortStep` to close the window where `set_phase` fired between the pop + /// and the choice — but a step is still chosen from the phase, so blocks + /// that outlived a phase transition were still mis-compressed. + fn try_compress(shared: &SharedPipelineState, worker: &mut SortWorkerState) -> StepResult { let Some(job) = shared.compress_queue.pop() else { return StepResult::InputEmpty; }; + // Destructured so the BGZF compressor and the zstd one are borrowed as + // disjoint fields rather than through `worker` twice. + let SortWorkerState { compressor, output_compressor, zstd_compressor, .. } = worker; + let bgzf_compressor = match job.target { + CompressTarget::Spill => compressor, + CompressTarget::Output => output_compressor, + }; Self::handle_compress_job(shared, job, bgzf_compressor, zstd_compressor); StepResult::Success } @@ -1953,11 +1981,12 @@ impl SortWorkerPool { /// The current pipeline phase (see [`phase`]), the read counterpart to /// [`set_phase`](Self::set_phase). /// - /// The phase selects which compress step workers schedule — and therefore - /// which compressor the output blocks go through — so it is the observable - /// the Phase 2 teardown tests assert on. Workers read the phase constantly, - /// but no caller *outside* the pool needs it, so this accessor exists for - /// those tests. + /// The phase orders the steps workers schedule, and nothing else — notably + /// *not* which compressor a block goes through, which `try_compress` takes + /// from the job's own [`CompressTarget`]. It is the observable the Phase 2 + /// teardown tests assert on. Workers read the phase constantly, but no + /// caller *outside* the pool needs it, so this accessor exists for those + /// tests. #[cfg(test)] pub(crate) fn current_phase(&self) -> u8 { self.shared.phase.load(Ordering::Acquire) @@ -1976,17 +2005,19 @@ impl SortWorkerPool { /// Hand the pool over to Phase 2 once ingest is done, widening it to /// `active_workers`. /// - /// Both halves belong to the same transition and must happen together. - /// Widening alone was not enough: the phase is what makes workers schedule - /// [`SortStep::CompressOutput`], and that step is what selects - /// `output_compressor` over the Phase 1 spill compressor. A sort that stayed in - /// Phase 1 through its output write therefore compressed the output BAM at - /// `temp_compression`, silently discarding the caller's `output_compression`. + /// Both halves belong to the same transition and must happen together: the + /// phase is what makes workers schedule [`SortStep::Phase2FileWork`], and + /// widening is what gives them the threads to do it on. /// /// Call this once ingest has finished — after that point every remaining byte - /// the pool touches is merge input or output. Callers must not leave spill - /// compression outstanding: Phase 1 spill jobs still queued when the phase flips - /// would be picked up by `CompressOutput` and written at the output level. + /// the pool touches is merge input or output. + /// + /// The phase deliberately says nothing about *compression levels*. Each + /// queued block carries its own [`CompressTarget`], so spill jobs still + /// outstanding when the phase flips are compressed at `temp_compression` + /// regardless, and output blocks are compressed at `output_compression` even + /// if the pool has already left Phase 2. Callers used to owe this method a + /// drained spill queue for exactly that reason; they no longer do. pub fn begin_phase2(&self, active_workers: usize) { self.set_active_workers(active_workers); self.set_phase(phase::PHASE2); @@ -2167,6 +2198,20 @@ impl SortWorkerPool { out } SpillCodec::Zstd => { + // Unlike BGZF, there is one zstd compressor per worker, fixed at + // `temp_compression` — zstd is a spill-only codec (the output BAM + // is always BGZF, so `PooledBamWriter` always submits + // `SpillCodec::Bgzf`). An Output job reaching here would be + // compressed at the temp level, which is the silent wrong-level + // output this whole selection mechanism exists to prevent — so + // assert unconditionally rather than only in debug builds. The + // check is one compare on a `Copy` enum per spill block, inside + // the zstd arm, so the BGZF path never evaluates it. + assert_eq!( + job.target, + CompressTarget::Spill, + "zstd is spill-only; its per-worker compressor is fixed at temp_compression" + ); // One self-contained zstd frame per job, length-prefixed so the // reader can split frames without scanning the stream. let frame = @@ -2473,6 +2518,98 @@ mod tests { assert_eq!(stats.compress_jobs_submitted.load(Ordering::Relaxed), 42); } + /// A job's [`CompressTarget`] decides its compression level, in every phase. + /// + /// This is the invariant that replaced deriving the compressor from the + /// pool's phase at pop time. That derivation produced the same bug twice: + /// output blocks popped outside Phase 2 were written at `temp_compression`, + /// and Phase 1 spill blocks still queued when `begin_phase2` fired were + /// written at `output_compression`. Both are impossible if the level travels + /// with the block, so the phase is swept across all three values here and + /// must not change either result. + /// + /// `temp_compression = 0` and `output_compression = 9` make the two levels + /// separable by size on compressible data: a `Spill` job must stay at the + /// (larger) stored size and an `Output` job must come back compressed. + #[rstest] + #[case::legacy(phase::LEGACY)] + #[case::phase1(phase::PHASE1)] + #[case::phase2(phase::PHASE2)] + fn test_compress_target_decides_level_regardless_of_phase(#[case] phase_under_test: u8) { + // Compressible: 8 KiB of one byte shrinks to almost nothing at level 9. + let data = vec![b'A'; 8192]; + let pool = SortWorkerPool::new(2, 0, 9, crate::codec::SpillCodec::Bgzf); + pool.set_phase(phase_under_test); + + let compressed_len = |target: CompressTarget| { + let (result_tx, result_rx) = pool.compress_result_channel(); + pool.submit_compress(CompressJob { + data: data.clone(), + serial: 0, + result_tx, + codec: SpillCodec::Bgzf, + target, + }); + result_rx.recv().expect("compress result").compressed.len() + }; + + let spill_len = compressed_len(CompressTarget::Spill); + let output_len = compressed_len(CompressTarget::Output); + + assert!( + spill_len > data.len(), + "a Spill job must be compressed at temp_compression=0 (stored), so it must not shrink \ + below the {} input bytes; got {spill_len} in phase {phase_under_test}", + data.len() + ); + assert!( + output_len < data.len() / 4, + "an Output job must be compressed at output_compression=9, so 8 KiB of one repeated \ + byte must shrink far below a quarter of its size; got {output_len} in phase \ + {phase_under_test}" + ); + + pool.shutdown(); + } + + /// zstd is spill-only, and `handle_compress_job` asserts that unconditionally. + /// + /// The assert is the backstop for the one target/codec combination the + /// per-job `CompressTarget` cannot express correctly: there is a single zstd + /// compressor per worker, fixed at `temp_compression`, so an `Output` job + /// arriving on the zstd arm would be written at the *spill* level — exactly + /// the silent wrong-level output this mechanism exists to prevent. Without a + /// test, a future output-zstd feature (or a mis-stamped call site) would only + /// trip it in production. + /// + /// `handle_compress_job` is called directly, on the test thread, rather than + /// through `submit_compress`: a panic on a pool worker unwinds only that + /// worker — its guard publishes `worker_panicked` and the unwind stops at the + /// thread boundary — so `#[should_panic]` would never observe it, and the + /// test would need a sleep just to reach the panic. Calling the handler + /// inline puts the assert on the test thread and makes it deterministic. + #[test] + #[should_panic(expected = "zstd is spill-only")] + fn test_zstd_output_job_trips_the_spill_only_assert() { + let shared = SharedPipelineState::new(1, std::thread::current()); + let mut bgzf_compressor = InlineBgzfCompressor::new(0); + let mut zstd_compressor = ZstdCompressor::new(1).expect("zstd compressor init"); + let (result_tx, _result_rx) = bounded(1); + + SortWorkerPool::handle_compress_job( + &shared, + CompressJob { + data: vec![b'A'; 64], + serial: 0, + result_tx, + codec: SpillCodec::Zstd, + target: CompressTarget::Output, + }, + &mut bgzf_compressor, + &mut zstd_compressor, + ); + } + #[test] fn test_pool_compress_roundtrip() { let pool = SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf); @@ -2480,7 +2617,13 @@ mod tests { // Submit a compress job let data = vec![b'A'; 1000]; - pool.submit_compress(CompressJob { data, serial: 0, result_tx, codec: SpillCodec::Bgzf }); + pool.submit_compress(CompressJob { + data, + serial: 0, + result_tx, + codec: SpillCodec::Bgzf, + target: CompressTarget::Spill, + }); // Wait for result let result = result_rx.recv().expect("should receive compress result"); @@ -2513,6 +2656,7 @@ mod tests { serial: i as u64, result_tx: submit_tx.clone(), codec: SpillCodec::Bgzf, + target: CompressTarget::Spill, }); } drop(submit_tx); @@ -2552,10 +2696,10 @@ mod tests { // 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; + let cs = SortStep::Compress 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. + // reuse it) and report the cumulative per-worker Compress counts. // `cumulative_jobs` is the total submitted across all batches so far — // the counters are cumulative, so it is what they must settle at. let run_batch = |pool: SortWorkerPool, @@ -2571,6 +2715,7 @@ mod tests { serial: i as u64, result_tx: submit_tx.clone(), codec: SpillCodec::Bgzf, + target: CompressTarget::Spill, }); } drop(submit_tx); @@ -2659,6 +2804,7 @@ mod tests { serial: 0, result_tx: c_tx, codec: SpillCodec::Bgzf, + target: CompressTarget::Spill, }); let _ = c_rx.recv(); @@ -2674,13 +2820,13 @@ mod tests { stats.record_step(0, SortStep::ReadInputBlocks, 1_000_000); stats.record_step(0, SortStep::ReadInputBlocks, 500_000); stats.record_step(1, SortStep::DecompressInput, 2_000_000); - stats.record_step(0, SortStep::CompressSpill, 300_000); + stats.record_step(0, SortStep::Compress, 300_000); stats.record_idle(0, 100_000); stats.record_idle(1, 200_000); let read_idx = SortStep::ReadInputBlocks as usize; let decomp_idx = SortStep::DecompressInput as usize; - let compress_idx = SortStep::CompressSpill as usize; + let compress_idx = SortStep::Compress as usize; assert_eq!(stats.step_count[read_idx].load(Ordering::Relaxed), 2); assert_eq!(stats.step_ns[read_idx].load(Ordering::Relaxed), 1_500_000); @@ -2696,7 +2842,7 @@ mod tests { fn test_pipeline_stats_log_summary_does_not_panic() { let stats = SortPipelineStats::new(4); stats.record_step(0, SortStep::ReadInputBlocks, 1_000_000_000); - stats.record_step(1, SortStep::CompressSpill, 500_000_000); + stats.record_step(1, SortStep::Compress, 500_000_000); stats.record_idle(0, 10_000_000); // Verify log_summary doesn't panic (output goes to log, not captured in tests) stats.log_summary(); @@ -2748,7 +2894,7 @@ mod tests { phase: phase::PHASE1, }; let priorities = get_sort_priorities(&bp); - assert_eq!(priorities[0], SortStep::CompressSpill); + assert_eq!(priorities[0], SortStep::Compress); } #[test] @@ -2786,7 +2932,7 @@ mod tests { phase: phase::PHASE2, }; let priorities = get_sort_priorities(&bp); - assert_eq!(priorities[0], SortStep::CompressOutput); + assert_eq!(priorities[0], SortStep::Compress); } #[test] @@ -2815,7 +2961,7 @@ mod tests { }; let priorities = get_sort_priorities(&bp); assert_eq!(priorities.len(), 1); - assert_eq!(priorities[0], SortStep::CompressSpill); + assert_eq!(priorities[0], SortStep::Compress); } #[test] diff --git a/crates/fgumi-sort/tests/integration/test_output_compression.rs b/crates/fgumi-sort/tests/integration/test_output_compression.rs index 0f501514d..cb9e29ecf 100644 --- a/crates/fgumi-sort/tests/integration/test_output_compression.rs +++ b/crates/fgumi-sort/tests/integration/test_output_compression.rs @@ -1,23 +1,27 @@ //! Integration test: `output_compression` must reach the BAM the sorter writes, //! on both the spilling and the fully-in-memory merge paths. //! -//! Output BGZF blocks are compressed by the worker pool, and which compressor a -//! worker reaches for is chosen by the pool's phase. Phase 2 used to be entered -//! only when there were spill files to merge, so an in-memory sort left the pool -//! in Phase 1 — where the only compress step is the *spill* one, using -//! `temp_compression`. The user's `--compression-level` was silently replaced by -//! the temp level. +//! Output BGZF blocks are compressed by the worker pool. Which compressor a +//! worker reaches for *used to* be chosen by the pool's phase at the moment the +//! block was popped, and that produced the same bug twice, from opposite sides: //! -//! The spilling merge had the mirror-image defect at the other end: it returned -//! the pool to `LEGACY` *before* finishing the writer, so the output blocks -//! still queued at teardown — plus the writer's final flush — went through the -//! spill compressor and the tail of the BAM came out at `temp_compression`. -//! Both defects are the same coupling seen from opposite sides: which -//! compressor a block gets is decided by the pool's phase when a worker pops -//! it, not by where the block came from. +//! - Phase 2 was entered only when there were spill files to merge, so an +//! in-memory sort left the pool in Phase 1 — where the only compress step was +//! the *spill* one, at `temp_compression`. +//! - A spilling merge returned the pool to `LEGACY` *before* finishing the +//! writer, so the output blocks still queued at teardown — plus the writer's +//! final flush — went through the spill compressor, and the tail of the BAM +//! came out at `temp_compression`. +//! +//! Either way the user's `--compression-level` was silently replaced by the temp +//! level. Each block now carries a `CompressTarget` stamped on it by the writer +//! that produced it, so the level no longer depends on the pool's phase, on +//! scheduling, or on how long a block sat in the queue. //! //! Hence the two tests below: the first asserts `output_compression` reaches -//! the file, the second that `temp_compression` never does. +//! the file, the second that `temp_compression` never does. They are the +//! regression tests for that whole class, so they are deliberately end-to-end — +//! they assert on the bytes of the written BAM, not on any pool internals. use fgumi_raw_bam::{IndexedRawBamReader, RawRecord}; use fgumi_sam::SamBuilder; @@ -249,8 +253,9 @@ fn sort_at_level( /// /// Both merge paths are covered: a memory limit large enough to hold everything /// (no spill, `chunks_written == 0`) and one small enough to force spill files. -/// Only the in-memory case was broken — it never entered Phase 2, so the pool -/// compressed output blocks with the Phase 1 spill compressor. +/// Only the in-memory case was broken — it never entered Phase 2, and back then +/// the phase was what chose the compressor, so the pool used the Phase 1 spill +/// one. /// /// Every sort order that routes its in-memory output through `PooledBamWriter` /// is exercised, because the `begin_phase2` transition is duplicated per order @@ -336,17 +341,17 @@ fn test_output_compression_level_reaches_the_output_bam( /// `temp_compression` must not reach the output BAM. /// /// The mirror image of the check above, at the other end of the sort: see this -/// file's module header for the defect. The output must be finalized while the -/// pool is still in Phase 2, or the trailing blocks go through the *temp* -/// compressor — so raising `temp_compression` alone would change the output. +/// file's module header for the defect. If any output block can reach the spill +/// compressor, raising `temp_compression` alone would change the output. /// /// So: sort twice at `output_compression = 0`, changing only the temp level. /// The written BAM must be exactly the same size both times. Level 0 is used /// for the output because it makes any block that slipped through the temp /// compressor strictly smaller, and therefore visible in the total. /// -/// Only the spilling path is exercised — the in-memory path never builds a -/// `Phase2Guard`, so it has no teardown to get wrong. +/// Only the spilling path is exercised, because it is the only one that produces +/// spill blocks at all — an in-memory sort has no `temp_compression` for the +/// output to be confused with. /// /// Both spilling finalizations are covered: `merge_chunks_generic` /// (`writer.finish()`) and, via the `coordinate_indexed` case,