diff --git a/.cargo/config.toml b/.cargo/config.toml index e1b57b931..2c455a011 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -34,7 +34,7 @@ ci-test-samtools = "nextest run --workspace --features compare,simulate,profile- # Run the concurrency stress tests (behind the `stress-tests` feature). Timing- # sensitive, so run on a nightly schedule (see .github/workflows/stress.yml) rather # than per-PR, where a flake would block unrelated changes. -ci-test-stress = "nextest run --workspace --features compare,simulate,profile-adjacency,stress-tests --locked" +ci-test-stress = "nextest run --workspace --features compare,simulate,profile-adjacency,stress-tests,fgumi-pipeline-io/stress-tests --locked" # Run tests with test-utils feature enabled (allows binary tests to use library test utilities) t = "test --features test-utils" # Generate and serve documentation locally (runs xtask then mdbook serve) diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index a47a7888d..c74a12bbf 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -109,7 +109,11 @@ jobs: with: tool: nextest - name: Generate coverage - run: cargo llvm-cov nextest --workspace --features compare,simulate,profile-adjacency --no-tests=pass --lcov --output-path lcov.info + # Includes `fgumi-pipeline-io/stress-tests`: the soak/matrix/proptest + # suites are gated off the fast `test` job, but they cover ~200 lines that + # nothing else reaches. Measuring without them under-reports patch + # coverage for code that IS tested, just not on the PR-latency path. + run: cargo llvm-cov nextest --workspace --features compare,simulate,profile-adjacency,fgumi-pipeline-io/stress-tests --no-tests=pass --lcov --output-path lcov.info - name: Upload to Codecov uses: codecov/codecov-action@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v6.0.2 with: diff --git a/Cargo.lock b/Cargo.lock index f1c3582d8..d7ee5dc34 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -871,6 +871,26 @@ dependencies = [ "trybuild", ] +[[package]] +name = "fgumi-pipeline-io" +version = "0.5.0" +dependencies = [ + "anyhow", + "criterion", + "fgumi-bam-io", + "fgumi-bgzf", + "fgumi-pipeline-core", + "fgumi-raw-bam", + "fgumi-sort", + "log", + "noodles", + "parking_lot", + "proptest", + "rayon", + "rstest", + "tempfile", +] + [[package]] name = "fgumi-raw-bam" version = "0.5.0" diff --git a/Cargo.toml b/Cargo.toml index be1d7b01c..c84c3f0ee 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [workspace] -members = [".", "crates/fgumi-raw-bam", "crates/fgumi-dna", "crates/fgumi-bgzf", "crates/fgumi-fmt", "crates/fgumi-metrics", "crates/fgumi-sam", "crates/fgumi-simd-fastq", "crates/fgumi-tag", "crates/fgumi-umi", "crates/fgumi-consensus", "crates/fgumi-bam-io", "crates/fgumi-sort", "crates/fgumi-cli-common", "crates/fgumi-cli-macros", "crates/fgumi-pipeline-core", "crates/xtask"] +members = [".", "crates/fgumi-raw-bam", "crates/fgumi-dna", "crates/fgumi-bgzf", "crates/fgumi-fmt", "crates/fgumi-metrics", "crates/fgumi-sam", "crates/fgumi-simd-fastq", "crates/fgumi-tag", "crates/fgumi-umi", "crates/fgumi-consensus", "crates/fgumi-bam-io", "crates/fgumi-sort", "crates/fgumi-cli-common", "crates/fgumi-cli-macros", "crates/fgumi-pipeline-core", "crates/fgumi-pipeline-io", "crates/xtask"] resolver = "2" [workspace.package] @@ -29,6 +29,7 @@ fgumi-consensus = { version = "0.5.0", path = "crates/fgumi-consensus", default- fgumi-dna = { version = "0.5.0", path = "crates/fgumi-dna" } fgumi-metrics = { version = "0.5.0", path = "crates/fgumi-metrics" } fgumi-pipeline-core = { version = "0.5.0", path = "crates/fgumi-pipeline-core" } +fgumi-pipeline-io = { version = "0.5.0", path = "crates/fgumi-pipeline-io" } fgumi-raw-bam = { version = "0.5.0", path = "crates/fgumi-raw-bam" } fgumi-sam = { version = "0.5.0", path = "crates/fgumi-sam" } fgumi-simd-fastq = { version = "0.5.0", path = "crates/fgumi-simd-fastq" } diff --git a/crates/fgumi-bam-io/src/prefetch_reader.rs b/crates/fgumi-bam-io/src/prefetch_reader.rs index 129a33288..a2e57b141 100644 --- a/crates/fgumi-bam-io/src/prefetch_reader.rs +++ b/crates/fgumi-bam-io/src/prefetch_reader.rs @@ -205,7 +205,6 @@ impl PrefetchReader { /// Total bytes served to callers of [`Read::read`] so far. #[must_use] - #[allow(dead_code)] pub fn bytes_consumed(&self) -> u64 { self.bytes_consumed } @@ -214,7 +213,6 @@ impl PrefetchReader { /// to deliver the next chunk. Useful as a prototype-phase signal for /// whether [`DEFAULT_PREFETCH_DEPTH`] is large enough. #[must_use] - #[allow(dead_code)] pub fn consumer_stalls(&self) -> u64 { self.consumer_stalls } diff --git a/crates/fgumi-pipeline-io/Cargo.toml b/crates/fgumi-pipeline-io/Cargo.toml new file mode 100644 index 000000000..8ca7db535 --- /dev/null +++ b/crates/fgumi-pipeline-io/Cargo.toml @@ -0,0 +1,43 @@ +[package] +name = "fgumi-pipeline-io" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +description = "BAM pipeline I/O types and steps for fgumi" +repository.workspace = true +license.workspace = true + +[dependencies] +fgumi-bam-io = { workspace = true } +fgumi-bgzf = { workspace = true } +fgumi-pipeline-core = { workspace = true } +fgumi-raw-bam = { workspace = true, features = ["noodles"] } +fgumi-sort = { workspace = true } +anyhow = "1.0.102" +log = "0" +noodles = { version = "0.111.0", features = ["bam", "sam", "bgzf"] } +parking_lot = "0.12" +rayon = "1.10" +tempfile = "3.4" + +[features] +test-utils = [] +# Multi-minute soak / matrix / proptest suites. Off by default so PR CI stays +# fast; run on the nightly schedule via `cargo ci-test-stress`, matching the +# root crate's `stress-tests` convention. +stress-tests = [] + +[dev-dependencies] +fgumi-bam-io = { workspace = true } +fgumi-raw-bam = { workspace = true, features = ["noodles", "test-utils"] } +fgumi-sort = { workspace = true, features = ["test-utils"] } +tempfile = "3.4" +noodles = { version = "0.111.0", features = ["bam", "sam"] } +fgumi-bgzf = { workspace = true } +rstest = "0.26" +proptest = "1.10" +criterion = { version = "0.8", features = ["html_reports"] } + +[[bench]] +name = "serial_ingest" +harness = false diff --git a/crates/fgumi-pipeline-io/benches/serial_ingest.rs b/crates/fgumi-pipeline-io/benches/serial_ingest.rs new file mode 100644 index 000000000..32b90ffe4 --- /dev/null +++ b/crates/fgumi-pipeline-io/benches/serial_ingest.rs @@ -0,0 +1,119 @@ +#![deny(unsafe_code)] + +//! Go/no-go microbenchmark for the parallel-inflate redesign (increment 3b.0). +//! +//! The redesign's whole wall-clock premise is that the *serial* Phase-1 ingest +//! step gets cheaper per record by NOT copying each record's bytes: today's +//! `SortBlockBuffer` does `BoundaryState::scan` → `CoordinateChunkSorter::push`, +//! and `push` memcpys every record body into the sorter's own arena (~the +//! 121 ns/rec the profile flagged). The redesign instead builds a lightweight +//! `(key, offset, len)` ref over bytes already in a shared arena — no copy. +//! +//! This bench measures EXACTLY that per-record delta with today's public API — +//! no new pipeline steps, no arena, no visibility changes: +//! * `scan_copy_push` — scan + `CoordinateChunkSorter::push` (copies body) +//! * `scan_refbuild` — scan + `extract_coordinate_key_inline` + `RecordRef::new` (no copy) +//! +//! It is a NECESSARY-condition gate, not the full wall proof: it shows whether +//! removing the copy cuts serial per-record cost. If `scan_refbuild` is not +//! materially faster than `scan_copy_push`, the redesign cannot beat legacy on +//! wall and we stop before building the pipeline. If it is faster, the full wall +//! win still has to be confirmed by the end-to-end gp3 measurement (the overlap +//! the freeze-per-run model trades away is not captured here). +//! +//! cargo bench -p fgumi-pipeline-io --bench serial_ingest + +use criterion::{BatchSize, Criterion, Throughput, criterion_group, criterion_main}; +use fgumi_pipeline_io::boundaries::BoundaryState; +use fgumi_raw_bam::testutil::make_bam_bytes; +use fgumi_sort::{RawExternalSorter, RecordRef, SortOrder, extract_coordinate_key_inline}; +use noodles::sam::Header; + +/// Build a decompressed-BAM-style buffer of `n` framed records: each record is +/// `[block_size: u32 LE][body]`, exactly what `BoundaryState::scan` walks. Bodies +/// are ~100 bp aligned reads at scattered positions on tid 0 (realistic per-record +/// size; the sort would do real reordering work). +fn synth_framed(n: usize) -> Vec { + let mut buf = Vec::new(); + for i in 0..n { + let pos = (i as u64).wrapping_mul(2_654_435_761) % 5_000_000; + let name = format!("r{i:08}"); + let body = make_bam_bytes(0, pos as i32, 0, name.as_bytes(), &[], 100, -1, -1, &[]); + let block_size = u32::try_from(body.len()).expect("record fits u32"); + buf.extend_from_slice(&block_size.to_le_bytes()); + buf.extend_from_slice(&body); + } + buf +} + +fn bench_serial_ingest(c: &mut Criterion) { + const N: usize = 200_000; + let framed = synth_framed(N); + let n_ref = 1u32; + + // Sanity: both paths see the same record count (so the throughput is over the + // same work). Computed once, untimed. + { + let mut bs = BoundaryState::new_no_header(); + let (offsets, _range) = bs.scan(&framed).expect("scan"); + assert_eq!(offsets.len().saturating_sub(1), N, "scan must yield N records"); + } + + let mut group = c.benchmark_group("serial_ingest"); + group.throughput(Throughput::Elements(N as u64)); + group.sample_size(20); + + // OLD path: scan + CoordinateChunkSorter::push — `push` copies each body into + // the sorter's arena (plus key-extract + internal ref-push). + group.bench_function("scan_copy_push", |b| { + b.iter_batched( + || { + RawExternalSorter::new(SortOrder::Coordinate) + .threads(1) + .into_coordinate_chunk_sorter(&Header::default()) + .expect("build coordinate chunk sorter") + }, + |mut sorter| { + let mut bs = BoundaryState::new_no_header(); + let (offsets, range) = bs.scan(&framed).expect("scan"); + let recs = bs.records_bytes(range); + for w in offsets.windows(2) { + let body = &recs[w[0] + 4..w[1]]; + sorter.push(body).expect("push"); + } + std::hint::black_box(&mut sorter); + }, + BatchSize::PerIteration, + ); + }); + + // NEW path: scan + key-extract + ref-build — NO body copy. This is the + // per-record work the redesign's serial FindBoundariesAndSort step does. + group.bench_function("scan_refbuild", |b| { + b.iter_batched( + || Vec::::with_capacity(N), + |mut refs| { + refs.clear(); + let mut bs = BoundaryState::new_no_header(); + let (offsets, range) = bs.scan(&framed).expect("scan"); + let recs = bs.records_bytes(range); + for w in offsets.windows(2) { + let body = &recs[w[0] + 4..w[1]]; + let key = extract_coordinate_key_inline(body, n_ref); + // offset/len point at the body (prefix skipped) — the redesign's + // arena ref representation; here the offset is illustrative. + let offset = u64::try_from(w[0] + 4).expect("offset fits u64"); + let len = u32::try_from(w[1] - w[0] - 4).expect("len fits u32"); + refs.push(RecordRef::new(key, offset, len)); + } + std::hint::black_box(&mut refs); + }, + BatchSize::PerIteration, + ); + }); + + group.finish(); +} + +criterion_group!(benches, bench_serial_ingest); +criterion_main!(benches); diff --git a/crates/fgumi-pipeline-io/src/boundaries.rs b/crates/fgumi-pipeline-io/src/boundaries.rs new file mode 100644 index 000000000..dc289f826 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/boundaries.rs @@ -0,0 +1,642 @@ +//! `BoundaryState` / `BoundaryBatch`: the BAM record-boundary state machine. +//! +//! Lives in `fgumi-pipeline-io` so both `FindBamBoundaries` (in the `fgumi` +//! crate, via a re-export shim) and the fused `SortBuffer` ingest can share it. +//! +//! Scans decompressed BGZF block data for BAM record boundaries (header skip, +//! cross-block record carryover, EOF validation) without decoding records. +//! Driven by the `FindBamBoundaries` step (`super::bam`). +//! +//! Relocated from the legacy `bam.rs` (deleted in the issue #330 migration); +//! the boundary-finding logic is reused verbatim so the new framework's +//! boundary semantics match the legacy pipeline's exactly. + +use std::io; + +/// Output of `FindBoundaries` step: buffer + record offsets for parallel decoding. +/// +/// This struct enables parallel BAM record decoding by pre-computing where +/// each record starts in the decompressed data. The actual parsing/decoding +/// can then be parallelized across multiple threads. +#[derive(Debug, Clone)] +pub struct BoundaryBatch { + /// The decompressed bytes (with leftover prepended, suffix removed). + pub buffer: Vec, + /// Byte offsets where each record starts (offsets into buffer). + /// Length = `num_records` + 1 (last entry is `buffer.len()` for easy slicing). + pub offsets: Vec, +} + +/// State for the `FindBoundaries` step (sequential). +/// +/// This state maintains leftover bytes from incomplete records that span +/// across BGZF block boundaries. The boundary finding is very fast (~0.1μs +/// per block) since it only reads 4-byte integers without decoding records. +/// +/// Uses a reusable work buffer to minimize allocations on the hot path. +pub struct BoundaryState { + /// Leftover bytes from previous block (incomplete record at end). + leftover: Vec, + /// Reusable working buffer to avoid per-call allocations. + work_buffer: Vec, + /// Whether the BAM header has been skipped. + header_skipped: bool, + /// Length of the previous call's `offsets` Vec, used to pre-size the next + /// one. Adjacent BGZF blocks hold near-identical record counts, so this + /// collapses the per-block push-regrowth (~8 reallocations) to ~1. The + /// returned `offsets` Vec is moved into `BoundaryBatch`, so it cannot be a + /// reused buffer; pre-sizing is the cheap, correctness-neutral alternative. + prev_offsets_len: usize, +} + +/// Return the byte length of the BAM header at the start of `data`. +/// +/// The BAM header consists of: +/// - 4-byte magic (`BAM\x01`) +/// - 4-byte `l_text` (little-endian u32) +/// - `l_text` bytes of plain-text header +/// - 4-byte `n_ref` (little-endian u32) +/// - for each of the `n_ref` references: 4-byte `l_name` + `l_name` bytes of name + 4-byte `l_ref` +/// +/// Returns: +/// - `Ok(Some(offset))` — `offset` bytes consume the complete header; the first record starts there. +/// - `Ok(None)` — `data` is too short to contain a complete header; more bytes are needed. +/// - `Err(InvalidData)` — `data` does not begin with the BAM magic. This path skips a BAM header, +/// so a wrong magic means the stream is not what the caller declared; fail closed rather than +/// silently treating arbitrary bytes as headerless records. Genuinely headerless streams must use +/// [`BoundaryState::new_no_header`], which never calls this. +/// +/// # Errors +/// +/// Returns `InvalidData` when `data` is long enough to check the magic but does not start with it. +pub fn bam_header_len(data: &[u8]) -> io::Result> { + // BAM header structure: + // - magic: 4 bytes ("BAM\1") + // - l_text: 4 bytes (header text length) + // - text: l_text bytes + // - n_ref: 4 bytes (number of references) + // - for each reference: + // - l_name: 4 bytes + // - name: l_name bytes + // - l_ref: 4 bytes + + if data.len() < 8 { + return Ok(None); + } + + // Check magic. This function is only reached on the header-skipping path + // (`header_skipped == false`); a wrong magic there means the stream is not + // the BAM the caller declared, so fail closed instead of misinterpreting the + // bytes as headerless records. + if &data[0..4] != fgumi_raw_bam::BAM_MAGIC { + return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid BAM magic")); + } + + let l_text = u32::from_le_bytes([data[4], data[5], data[6], data[7]]) as usize; + let mut offset = 8 + l_text; + + if data.len() < offset + 4 { + return Ok(None); + } + + let n_ref = + u32::from_le_bytes([data[offset], data[offset + 1], data[offset + 2], data[offset + 3]]) + as usize; + offset += 4; + + // Parse each reference + for _ in 0..n_ref { + if data.len() < offset + 4 { + return Ok(None); + } + let l_name = u32::from_le_bytes([ + data[offset], + data[offset + 1], + data[offset + 2], + data[offset + 3], + ]) as usize; + offset += 4 + l_name + 4; // l_name + name + l_ref + + if data.len() < offset { + return Ok(None); + } + } + + Ok(Some(offset)) +} + +impl BoundaryState { + /// Create a new boundary state. + #[must_use] + pub fn new() -> Self { + Self { + leftover: Vec::new(), + work_buffer: Vec::new(), + header_skipped: false, + prev_offsets_len: 0, + } + } + + /// Create a new boundary state that doesn't skip the header. + /// Use this when the input stream is already positioned past the header. + #[must_use] + pub fn new_no_header() -> Self { + Self { + leftover: Vec::new(), + work_buffer: Vec::new(), + header_skipped: true, + prev_offsets_len: 0, + } + } + + /// Find record boundaries in decompressed data. + /// + /// This is FAST (~0.1μs per block) because it only scans 4-byte integers + /// to find where records start - no actual record decoding is performed. + /// + /// # Arguments + /// + /// * `decompressed` - Decompressed bytes from one or more BGZF blocks + /// + /// # Returns + /// + /// A `BoundaryBatch` containing the complete records and their offsets. + /// Any incomplete record at the end is saved as leftover for the next call. + /// + /// # Errors + /// + /// Returns an I/O error if the BAM header is malformed. + /// + /// # Record-level validation + /// + /// This function does NOT validate individual record `block_size` values + /// against a malformed (but self-consistent) BAM stream. The per-record + /// cross-check below (offset delta vs. the stored prefix) is a + /// `debug_assertions`-only regression tripwire for this scanner's own + /// arithmetic — it re-reads the same `block_size` bytes the scan already + /// trusted, so it can only catch an internal bookkeeping bug, never input + /// corruption. Authoritative release-build validation of record structure + /// (out-of-bounds record end, trailing partial record) is performed + /// downstream by `parse_records` / `parse_record_ranges` on the same bytes, + /// which hard-error in all build modes. The `offsets` vector this returns + /// is not consumed in release builds (`FindBamBoundaries` forwards only + /// `buffer`), so promoting the cross-check to release would re-validate a + /// tautology at a per-record cost for no correctness benefit. + pub fn find_boundaries(&mut self, decompressed: &[u8]) -> io::Result { + let (offsets, range) = self.scan(decompressed)?; + // Owning copy of the complete records — for callers that need an owned + // buffer (`FindBamBoundaries`). The zero-copy ingest path + // ([`scan`](Self::scan) + [`records_bytes`](Self::records_bytes)) skips + // this allocation + copy entirely. + let buffer = self.work_buffer[range].to_vec(); + + // Debug-only regression tripwire (NOT input validation): cross-check + // each record's stored block_size prefix against the offset delta this + // scan just computed. Both derive from the same bytes with no + // intervening mutation, so this only catches an internal arithmetic / + // indexing bug in the scan above — a corrupt-but-self-consistent + // block_size passes trivially. Authoritative release validation lives + // in parse_records / parse_record_ranges downstream (see the + // `find_boundaries` doc comment). + #[cfg(debug_assertions)] + for i in 0..offsets.len().saturating_sub(1) { + let start = offsets[i]; + let end = offsets[i + 1]; + if end > start + 4 { + let stored = u32::from_le_bytes([ + buffer[start], + buffer[start + 1], + buffer[start + 2], + buffer[start + 3], + ]) as usize; + let expected = end - start - 4; + debug_assert_eq!( + stored, expected, + "find_boundaries: block_size mismatch at record {i}: stored={stored}, expected={expected}" + ); + } + } + + Ok(BoundaryBatch { buffer, offsets }) + } + + /// Zero-copy core of [`find_boundaries`](Self::find_boundaries): combine + /// leftover + `decompressed` into the reusable `work_buffer`, skip the header + /// (first call), scan record boundaries, and stash the trailing partial + /// record as leftover — **without** copying the complete records out. + /// + /// Returns `(offsets, range)` where the complete records live in + /// `self.work_buffer[range]` (valid until the next `scan`/`find_boundaries` + /// call) and `offsets[i] .. offsets[i+1]` slices record `i` *relative to + /// `range.start`* (so record `i`'s bytes are + /// `records_bytes()[offsets[i] .. offsets[i+1]]`). The caller must consume + /// the records before the next call. A header-only / incomplete-header block + /// yields `(vec![0], 0..0)`. + /// + /// # Errors + /// + /// Returns an I/O error if the BAM header is malformed. + pub fn scan( + &mut self, + decompressed: &[u8], + ) -> io::Result<(Vec, std::ops::Range)> { + // Step 1: Combine leftover with new data into reusable work_buffer. + self.work_buffer.clear(); + if !self.leftover.is_empty() { + self.work_buffer.append(&mut self.leftover); + } + self.work_buffer.extend_from_slice(decompressed); + + // Step 2: Skip header if not already done. + let mut cursor = 0usize; + if !self.header_skipped { + let Some(header_size) = bam_header_len(&self.work_buffer)? else { + // Not enough data to parse header; save as leftover, empty range. + std::mem::swap(&mut self.leftover, &mut self.work_buffer); + return Ok((vec![0], 0..0)); + }; + cursor = header_size; + self.header_skipped = true; + } + + // Step 3: Scan for record boundaries (FAST - just read integers). + let start_cursor = cursor; + let mut offsets = Vec::with_capacity(self.prev_offsets_len.max(1)); + offsets.push(0usize); + while cursor + 4 <= self.work_buffer.len() { + let block_size = u32::from_le_bytes([ + self.work_buffer[cursor], + self.work_buffer[cursor + 1], + self.work_buffer[cursor + 2], + self.work_buffer[cursor + 3], + ]) as usize; + let record_end = cursor + 4 + block_size; + if record_end > self.work_buffer.len() { + break; // Incomplete record - becomes leftover. + } + cursor = record_end; + offsets.push(cursor - start_cursor); + } + self.prev_offsets_len = offsets.len(); + + // Step 4: Save trailing partial record as leftover (small — at most one + // record). The complete records stay in `work_buffer[start_cursor..cursor]` + // for the caller to read borrowed, with no copy. + self.leftover.clear(); + self.leftover.extend_from_slice(&self.work_buffer[cursor..]); + + Ok((offsets, start_cursor..cursor)) + } + + /// Borrow the complete records produced by the most recent [`scan`](Self::scan), + /// given the `range` it returned. Valid until the next `scan`/`find_boundaries`. + #[must_use] + pub fn records_bytes(&self, range: std::ops::Range) -> &[u8] { + &self.work_buffer[range] + } + + /// Call at EOF to get any remaining leftover. + /// + /// This validates that any remaining bytes form complete records. + /// If there are incomplete bytes at EOF, an error is returned. + /// + /// # Errors + /// + /// Returns an I/O error if there are incomplete BAM records at EOF. + pub fn finish(&mut self) -> io::Result> { + if self.leftover.is_empty() { + return Ok(None); + } + + // Try to parse remaining leftover + let mut offsets = vec![0usize]; + let mut cursor = 0usize; + + while cursor + 4 <= self.leftover.len() { + let block_size = u32::from_le_bytes([ + self.leftover[cursor], + self.leftover[cursor + 1], + self.leftover[cursor + 2], + self.leftover[cursor + 3], + ]) as usize; + + let record_end = cursor + 4 + block_size; + if record_end > self.leftover.len() { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!( + "Incomplete BAM record at EOF: need {} bytes, have {}", + record_end - cursor, + self.leftover.len() - cursor + ), + )); + } + + cursor = record_end; + offsets.push(cursor); + } + + // The loop only advances `cursor` by whole records. If it stops with + // bytes still unconsumed (`cursor < leftover.len()`), those 1-3 trailing + // bytes are too short to even hold a 4-byte block-size prefix — i.e. a + // truncated BAM record. Surface it as an error rather than dropping the + // bytes and masking corruption. + if cursor < self.leftover.len() { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!( + "Incomplete BAM record at EOF: {} trailing byte(s) cannot form a complete record", + self.leftover.len() - cursor + ), + )); + } + + Ok(Some(BoundaryBatch { buffer: std::mem::take(&mut self.leftover), offsets })) + } +} + +impl Default for BoundaryState { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + /// Build a BAM header: magic, `l_text` + text, `n_ref` + one entry per `(name, l_ref)`. + /// + /// Every test header is constructed here rather than committed as a fixture, so a + /// header's declared lengths and its actual bytes cannot drift apart. + fn header(text: &str, refs: &[(&str, u32)]) -> Vec { + let mut h = Vec::new(); + h.extend_from_slice(fgumi_raw_bam::BAM_MAGIC); + h.extend_from_slice(&u32::try_from(text.len()).unwrap().to_le_bytes()); + h.extend_from_slice(text.as_bytes()); + h.extend_from_slice(&u32::try_from(refs.len()).unwrap().to_le_bytes()); + for (name, l_ref) in refs { + // l_name counts the trailing NUL, matching the BAM spec. + let name_bytes = format!("{name}\0"); + h.extend_from_slice(&u32::try_from(name_bytes.len()).unwrap().to_le_bytes()); + h.extend_from_slice(name_bytes.as_bytes()); + h.extend_from_slice(&l_ref.to_le_bytes()); + } + h + } + + /// Build one BAM record: a 4-byte little-endian `block_size` followed by + /// `payload_len` bytes of `fill`. The scanner only reads the length prefix, so + /// the payload just has to be the declared size and be identifiable. + fn record(payload_len: usize, fill: u8) -> Vec { + let mut r = u32::try_from(payload_len).unwrap().to_le_bytes().to_vec(); + r.extend(std::iter::repeat_n(fill, payload_len)); + r + } + + /// The expected total on-disk size of a record with `payload_len` bytes. + fn record_len(payload_len: usize) -> usize { + payload_len + 4 + } + + // --------------------------------------------------------------------- + // bam_header_len + // --------------------------------------------------------------------- + + /// What a `bam_header_len` case expects: a parsed length (or `None` for + /// "need more bytes"), or a hard `InvalidData` rejection. + #[derive(Debug, Clone, Copy)] + enum Expect { + Header(Option), + InvalidMagic, + } + + #[rstest] + #[case::empty(vec![], Expect::Header(None))] + #[case::shorter_than_magic(b"BAM".to_vec(), Expect::Header(None))] + #[case::magic_only_no_room_for_n_ref( + // 8 bytes: magic + l_text=0. Needs offset+4 == 12 to read n_ref. + [&fgumi_raw_bam::BAM_MAGIC[..], &0u32.to_le_bytes()[..]].concat(), + Expect::Header(None) + )] + #[case::no_refs(header("", &[]), Expect::Header(Some(12)))] + #[case::with_header_text(header("@HD\tVN:1.6\n", &[]), Expect::Header(Some(12 + 11)))] + // 12 + (4 l_name + 3 name + 4 l_ref) == 23 + #[case::one_ref(header("", &[("r1", 100)]), Expect::Header(Some(23)))] + #[case::two_refs(header("", &[("r1", 100), ("chr2", 200)]), Expect::Header(Some(23 + 4 + 5 + 4)))] + #[case::truncated_mid_ref_name(header("", &[("r1", 100)])[..20].to_vec(), Expect::Header(None))] + #[case::truncated_before_l_name(header("", &[("r1", 100)])[..14].to_vec(), Expect::Header(None))] + #[case::bad_magic(b"NOT\x01\x00\x00\x00\x00\x00\x00\x00\x00".to_vec(), Expect::InvalidMagic)] + fn bam_header_len_cases(#[case] data: Vec, #[case] expected: Expect) { + match expected { + Expect::Header(len) => assert_eq!(bam_header_len(&data).unwrap(), len), + Expect::InvalidMagic => { + let err = bam_header_len(&data).expect_err("expected an error"); + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + } + } + } + + // --------------------------------------------------------------------- + // Constructors + // --------------------------------------------------------------------- + + #[test] + fn new_skips_the_header_and_new_no_header_does_not() { + let hdr = header("", &[]); + let rec = record(8, 0xAB); + + // `new` consumes the header, so the first record starts after it. + let mut with_header = BoundaryState::new(); + let (offsets, range) = with_header.scan(&[hdr.clone(), rec.clone()].concat()).unwrap(); + assert_eq!(range, hdr.len()..hdr.len() + record_len(8)); + assert_eq!(offsets, vec![0, record_len(8)]); + assert_eq!(with_header.records_bytes(range), rec.as_slice()); + + // `new_no_header` treats byte 0 as the first record. + let mut headerless = BoundaryState::new_no_header(); + let (offsets, range) = headerless.scan(&rec).unwrap(); + assert_eq!(range, 0..record_len(8)); + assert_eq!(offsets, vec![0, record_len(8)]); + } + + #[test] + fn default_matches_new_and_still_skips_the_header() { + let data = [header("", &[]), record(4, 0x11)].concat(); + let mut from_default = BoundaryState::default(); + let mut from_new = BoundaryState::new(); + assert_eq!(from_default.scan(&data).unwrap(), from_new.scan(&data).unwrap()); + } + + // --------------------------------------------------------------------- + // scan + // --------------------------------------------------------------------- + + #[rstest] + #[case::single(vec![8])] + #[case::several_same_size(vec![4, 4, 4])] + #[case::mixed_sizes(vec![1, 16, 3, 32])] + #[case::zero_length_payload(vec![0, 0])] + fn scan_finds_every_complete_record(#[case] payloads: Vec) { + let mut data = header("", &[]); + for (i, len) in payloads.iter().enumerate() { + data.extend(record(*len, u8::try_from(i).unwrap())); + } + + let mut state = BoundaryState::new(); + let (offsets, range) = state.scan(&data).unwrap(); + + // One offset per record plus the terminating end offset. + assert_eq!(offsets.len(), payloads.len() + 1); + let mut running = 0usize; + for (i, len) in payloads.iter().enumerate() { + assert_eq!(offsets[i], running, "record {i} start"); + running += record_len(*len); + } + assert_eq!(*offsets.last().unwrap(), running); + assert_eq!(range.len(), running); + assert!(state.leftover.is_empty(), "no leftover when every record is complete"); + } + + #[test] + fn scan_holds_back_a_trailing_partial_record_as_leftover() { + let complete = record(8, 0x01); + let partial = &record(64, 0x02)[..10]; // declares 64 bytes, supplies 6 + let data = [header("", &[]), complete.clone(), partial.to_vec()].concat(); + + let mut state = BoundaryState::new(); + let (offsets, range) = state.scan(&data).unwrap(); + + // Only the complete record is emitted. + assert_eq!(offsets, vec![0, record_len(8)]); + assert_eq!(state.records_bytes(range), complete.as_slice()); + assert_eq!(state.leftover, partial, "the partial record is carried forward verbatim"); + } + + #[test] + fn scan_reassembles_a_record_split_across_two_blocks() { + let rec = record(32, 0x7E); + let data = [header("", &[]), rec.clone()].concat(); + let split = data.len() - 20; // cut mid-record + + let mut state = BoundaryState::new(); + + // First block ends mid-record: nothing complete yet. + let (offsets, range) = state.scan(&data[..split]).unwrap(); + assert_eq!(offsets, vec![0], "no complete record in the first block"); + assert_eq!(range.len(), 0); + assert!(!state.leftover.is_empty()); + + // Second block completes it, and the bytes match the original record. + let (offsets, range) = state.scan(&data[split..]).unwrap(); + assert_eq!(offsets, vec![0, record_len(32)]); + assert_eq!(state.records_bytes(range), rec.as_slice()); + assert!(state.leftover.is_empty()); + } + + #[test] + fn scan_defers_when_the_header_itself_is_incomplete() { + let hdr = header("@HD\tVN:1.6\n", &[("r1", 100)]); + let mut state = BoundaryState::new(); + + // A prefix too short to hold the whole header yields an empty batch and + // stashes everything for the next call. + let (offsets, range) = state.scan(&hdr[..10]).unwrap(); + assert_eq!(offsets, vec![0]); + assert_eq!(range, 0..0); + assert_eq!(state.leftover, hdr[..10]); + assert!(!state.header_skipped, "header must not be marked skipped yet"); + + // The rest of the header plus a record then parses normally. + let rec = record(8, 0x5A); + let (offsets, range) = state.scan(&[&hdr[10..], rec.as_slice()].concat()).unwrap(); + assert!(state.header_skipped); + assert_eq!(offsets, vec![0, record_len(8)]); + assert_eq!(state.records_bytes(range), rec.as_slice()); + } + + #[test] + fn scan_propagates_a_bad_magic_as_invalid_data() { + let mut state = BoundaryState::new(); + let err = state.scan(b"NOPE\x00\x00\x00\x00\x00\x00\x00\x00").expect_err("bad magic"); + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + } + + // --------------------------------------------------------------------- + // find_boundaries (owned-buffer wrapper over scan) + // --------------------------------------------------------------------- + + #[test] + fn find_boundaries_returns_the_same_bytes_scan_would_borrow() { + let recs = [record(8, 0x01), record(16, 0x02), record(2, 0x03)].concat(); + let data = [header("", &[("r1", 10)]), recs.clone()].concat(); + + let mut owned = BoundaryState::new(); + let batch = owned.find_boundaries(&data).unwrap(); + + let mut borrowed = BoundaryState::new(); + let (offsets, range) = borrowed.scan(&data).unwrap(); + + assert_eq!(batch.offsets, offsets); + assert_eq!(batch.buffer, borrowed.records_bytes(range)); + assert_eq!(batch.buffer, recs, "the owned copy is the record bytes, header excluded"); + + // Offsets slice the buffer into the original records. + for i in 0..batch.offsets.len() - 1 { + let rec = &batch.buffer[batch.offsets[i]..batch.offsets[i + 1]]; + let declared = u32::from_le_bytes(rec[..4].try_into().unwrap()) as usize; + assert_eq!(declared, rec.len() - 4, "record {i} length prefix matches its slice"); + } + } + + #[test] + fn find_boundaries_on_a_header_only_block_yields_an_empty_batch() { + let mut state = BoundaryState::new(); + let batch = state.find_boundaries(&header("", &[])).unwrap(); + assert!(batch.buffer.is_empty()); + assert_eq!(batch.offsets, vec![0]); + } + + // --------------------------------------------------------------------- + // finish + // --------------------------------------------------------------------- + + #[test] + fn finish_returns_none_when_nothing_is_pending() { + let mut state = BoundaryState::new_no_header(); + assert!(state.finish().unwrap().is_none()); + + // Also none after a scan that consumed every record. + let (_, _) = state.scan(&record(4, 0x09)).unwrap(); + assert!(state.finish().unwrap().is_none()); + } + + #[test] + fn finish_emits_leftover_that_forms_complete_records() { + // `scan` never leaves a *complete* record behind, so the pending buffer is + // seeded directly to exercise the success branch. + let mut state = BoundaryState::new_no_header(); + state.leftover = [record(4, 0xA1), record(8, 0xA2)].concat(); + + let batch = state.finish().unwrap().expect("complete records must be emitted"); + assert_eq!(batch.offsets, vec![0, record_len(4), record_len(4) + record_len(8)]); + assert_eq!(batch.buffer.len(), record_len(4) + record_len(8)); + assert!(state.leftover.is_empty(), "finish takes the pending bytes"); + } + + #[rstest] + // Declares a 64-byte payload but only supplies part of it. + #[case::truncated_payload(record(64, 0x02)[..10].to_vec())] + // Fewer than 4 bytes cannot even hold a block-size prefix. + #[case::one_trailing_byte(vec![0x00])] + #[case::three_trailing_bytes(vec![0x00, 0x01, 0x02])] + // A whole record followed by an unusable tail. + #[case::complete_then_stray_bytes([record(4, 0x03), vec![0xFF, 0xFF]].concat())] + fn finish_rejects_incomplete_trailing_bytes(#[case] pending: Vec) { + let mut state = BoundaryState::new_no_header(); + state.leftover = pending; + let err = state.finish().expect_err("truncated input must fail closed"); + assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof); + } +} diff --git a/crates/fgumi-pipeline-io/src/lib.rs b/crates/fgumi-pipeline-io/src/lib.rs new file mode 100644 index 000000000..5e6b3293f --- /dev/null +++ b/crates/fgumi-pipeline-io/src/lib.rs @@ -0,0 +1,30 @@ +#![deny(unsafe_code)] + +//! BAM-pipeline I/O layer for the fgumi typed-step pipeline. +//! +//! Provides the source, sink, and sort typed-step building blocks plus the +//! record-batch and BGZF-block buffer types shared across the fgumi pipeline: +//! +//! * [`source`] — BAM ingest steps ([`ReadBgzfBlocks`] and the +//! `read_bam*` helpers) that turn a reader into decompressed blocks. +//! * [`sink`] — BAM output steps ([`WriteBgzfFile`]). +//! * [`sort`] — in-pipeline sort steps ([`SortBuffer`], [`CompressSpill`], +//! [`SortSpillDecompress`], [`SortMerge`]). +//! * [`types`] — the record-batch / decompressed-block buffers +//! ([`RecordBatch`], [`DecompressedBlock`], [`BgzfBlock`], …) exchanged +//! between steps. + +pub mod boundaries; +pub mod sink; +pub mod sort; +pub mod source; +pub mod types; + +pub use fgumi_pipeline_core::HeaderHandle; +pub use sink::write_bgzf::WriteBgzfFile; +pub use sort::{CompressSpill, SortBuffer, SortMerge, SortSpillDecompress}; +pub use source::read_bam::{ + DEFAULT_BLOCKS_PER_BATCH, ReadBgzfBlocks, read_bam, read_bam_auto, read_bam_from_reader, + read_bam_stdin, +}; +pub use types::{BgzfBlock, DecompressedBlock, RecordBatch, RecordBatchBuilder}; diff --git a/crates/fgumi-pipeline-io/src/sink/mod.rs b/crates/fgumi-pipeline-io/src/sink/mod.rs new file mode 100644 index 000000000..56d141d03 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sink/mod.rs @@ -0,0 +1,2 @@ +pub mod write_bgzf; +pub mod write_raw; diff --git a/crates/fgumi-pipeline-io/src/sink/write_bgzf.rs b/crates/fgumi-pipeline-io/src/sink/write_bgzf.rs new file mode 100644 index 000000000..e98dab7bf --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sink/write_bgzf.rs @@ -0,0 +1,427 @@ +//! `WriteBgzfFile` sink step. `Serial + Affinity::Writer` by default, or +//! `StepKind::Detached` (its own dedicated thread, off the pool) when built +//! via [`WriteBgzfFile::with_detached`] — used only on the standalone-sort +//! terminal (lever 2, legacy "N + 2"). Receives pre-compressed `BgzfBlock`s +//! from `BgzfCompress` and writes them directly to disk. + +use std::fs::File; +use std::io::{self, BufWriter, Write}; +use std::path::Path; + +use fgumi_bgzf::{BGZF_EOF, InlineBgzfCompressor}; +use noodles::sam::Header; +use parking_lot::Mutex; + +use crate::types::BgzfBlock; +use fgumi_pipeline_core::{ + header::HeaderHandle, + step::{Affinity, DetachedGroup, Step, StepCtx, StepKind, StepOutcome, StepProfile}, +}; + +/// `Serial + sticky` BAM sink (or `Detached` — see [`Self::with_detached`]) +/// that consumes pre-compressed `BgzfBlock`s. +pub struct WriteBgzfFile { + state: Mutex>, + name: &'static str, + /// When `Some`, advertise `StepKind::Detached` so the framework drives this + /// sink on its own dedicated driver thread (off the work-stealing pool) in the + /// given [`DetachedGroup`], instead of `Serial + Affinity::Writer`. The caller + /// (e.g. the sort chain) supplies the group via [`Self::with_detached`], so + /// this generic sink carries no chain-specific grouping. `None` (the default) + /// keeps the pool-scheduled writer that every other chain uses. + detached_group: Option, +} + +struct WriterState { + out: BufWriter, + pending_header: Option, +} + +/// Transform applied to the aligner's runtime-resolved header before it is +/// written. Lets a post-`Align` stage (sort-order rewrite, consensus header, +/// clip) re-apply its header change on top of the resolved header — which +/// carries the aligner's runtime `@PG`/`@RG`/`@CO` — instead of discarding that +/// provenance by writing a build-time header. +pub type ResolvedHeaderTransform = Box io::Result
+ Send + Sync>; + +struct PendingHeader { + handle: HeaderHandle, + compression_level: u32, + transform: Option, +} + +impl WriteBgzfFile { + /// Open `path`, BGZF-compress and write the BAM header bytes, return + /// the sink ready to receive `BgzfBlock`s. + /// + /// # Errors + /// + /// Returns I/O errors from path open or header write. + pub fn new>( + path: P, + header: &Header, + compression_level: u32, + ) -> io::Result { + let file = File::create(path.as_ref())?; + let mut out = BufWriter::with_capacity(256 * 1024, file); + + let mut header_bytes = Vec::new(); + fgumi_bam_io::write_bam_header(&mut header_bytes, header) + .map_err(|e| io::Error::other(format!("write_bam_header: {e}")))?; + + let mut hc = InlineBgzfCompressor::new(compression_level); + hc.write_all(&header_bytes)?; + hc.flush()?; + hc.write_blocks_to(&mut out)?; + + Ok(Self { + state: Mutex::new(Some(WriterState { out, pending_header: None })), + name: "WriteBgzfFile", + detached_group: None, + }) + } + + /// Run this sink on a dedicated `StepKind::Detached` driver thread (in the + /// given [`DetachedGroup`]) instead of as a pool-scheduled `Serial + + /// Affinity::Writer` step. Used ONLY on the standalone-sort terminal (lever 2): + /// it frees a pool worker for the compression-bound work, matching the legacy + /// sort's dedicated writer thread. The caller chooses the group so this generic + /// sink stays chain-agnostic. The `try_run` body and the bytes it writes are + /// unchanged — the driver pops blocks in the same (reorder-stage-ordered) + /// sequence — so the output BAM is byte-identical to the pool-scheduled writer. + /// Affinity is ignored for `Detached`. + #[must_use] + pub fn with_detached(mut self, group: DetachedGroup) -> Self { + self.detached_group = Some(group); + self + } + + /// Open `path` and return the sink with the BAM header write + /// deferred until an upstream step resolves `handle`. + /// + /// `transform`, when `Some`, is applied to the resolved header before it is + /// written — see [`ResolvedHeaderTransform`] for why this exists; pass `None` + /// when the resolved header should be written unchanged. + /// + /// # Errors + /// + /// Returns I/O errors from path open. Header-write errors are + /// surfaced from `try_run` once the handle resolves. + pub fn new_with_handle>( + path: P, + handle: HeaderHandle, + compression_level: u32, + transform: Option, + ) -> io::Result { + let file = File::create(path.as_ref())?; + let out = BufWriter::with_capacity(256 * 1024, file); + Ok(Self { + state: Mutex::new(Some(WriterState { + out, + pending_header: Some(PendingHeader { handle, compression_level, transform }), + })), + name: "WriteBgzfFile", + detached_group: None, + }) + } + + fn try_write_pending_header(state: &mut WriterState) -> io::Result { + let Some(pending) = state.pending_header.as_ref() else { + return Ok(true); + }; + let header_clone = match pending.handle.try_get() { + None => return Ok(false), + Some(Err(e)) => return Err(e), + Some(Ok(h)) => h.clone(), + }; + let level = pending.compression_level; + // Re-apply the post-align header change (sort order / consensus / clip) + // on top of the aligner's runtime-resolved header so its @PG/@RG/@CO + // survive into the written header. + let header_clone = match &pending.transform { + Some(transform) => transform(header_clone)?, + None => header_clone, + }; + + let mut header_bytes = Vec::new(); + fgumi_bam_io::write_bam_header(&mut header_bytes, &header_clone) + .map_err(|e| io::Error::other(format!("write_bam_header: {e}")))?; + let mut hc = InlineBgzfCompressor::new(level); + hc.write_all(&header_bytes)?; + hc.flush()?; + hc.write_blocks_to(&mut state.out)?; + + state.pending_header = None; + Ok(true) + } +} + +impl Step for WriteBgzfFile { + type Input = BgzfBlock; + type Outputs = (); + + fn profile(&self) -> StepProfile { + StepProfile { + name: self.name, + // Detached (own driver thread) when a group was set via + // `with_detached`; otherwise the default pool-scheduled Serial + sticky + // writer. `sticky` is irrelevant for Detached (it never enters a + // worker's worklist). + kind: if self.detached_group.is_some() { StepKind::Detached } else { StepKind::Serial }, + sticky: true, + output_queues: vec![], + branch_ordering: vec![], + } + } + + fn detached_group(&self) -> DetachedGroup { + // The caller-supplied group (the sort chain passes `SORT_IO_GROUP`); + // `PerStep` fallback is unreachable for a Serial (non-detached) writer + // since `detached_group()` is only consulted for `Detached` steps. + self.detached_group.unwrap_or(DetachedGroup::PerStep) + } + + fn affinity(&self) -> Affinity { + // Ignored for `Detached` (no pool worker drives it); kept for the + // default Serial path where it pins the writer to the last worker. + Affinity::Writer + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + let mut guard = self.state.lock(); + let Some(state) = guard.as_mut() else { + return Ok(StepOutcome::Finished); + }; + + let header_ready = Self::try_write_pending_header(state)?; + if header_ready && let Some(block) = ctx.input.pop() { + state.out.write_all(&block.bytes)?; + return Ok(StepOutcome::Progress); + } + + if ctx.input.is_drained() { + if !header_ready { + return Err(io::Error::other( + "WriteBgzfFile: input drained before HeaderHandle was resolved", + )); + } + state.out.write_all(&BGZF_EOF)?; + state.out.flush()?; + let _ = guard.take(); + return Ok(StepOutcome::Finished); + } + Ok(StepOutcome::NoProgress) + } +} + +impl Drop for WriteBgzfFile { + /// Cleanup-only drop. The BGZF EOF marker is written **exclusively** by the + /// drained-finish path in `try_run` (which then takes the state so this drop + /// is a no-op for a normally-finished sink). If state is still present here + /// the sink was dropped before that path ran — i.e. an aborted/partial + /// stream — so we deliberately do **not** append `BGZF_EOF`: stamping the + /// EOF marker onto a truncated BAM would make it look like a complete stream + /// and hide the truncation from downstream readers. We only flush whatever + /// bytes were already buffered so the on-disk file reflects what was written + /// (and stays detectably truncated). A still-pending header means nothing + /// valid was written, so leave the file empty. + fn drop(&mut self) { + let mut guard = self.state.lock(); + if let Some(mut state) = guard.take() { + if state.pending_header.is_some() { + return; + } + let _ = state.out.flush(); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn empty_header() -> Header { + Header::default() + } + + #[test] + fn profile_advertises_serial_writer_sink() { + let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); + let header = empty_header(); + let step = WriteBgzfFile::new(&path, &header, 1).unwrap(); + let profile = step.profile(); + assert_eq!(profile.name, "WriteBgzfFile"); + assert_eq!(profile.kind, StepKind::Serial); + assert!(profile.sticky); + assert_eq!(step.affinity(), Affinity::Writer); + assert_eq!(profile.output_queues.len(), 0); + assert_eq!(profile.branch_ordering.len(), 0); + } + + /// L2.6: `with_detached()` flips the profile kind to `Detached` (its own + /// thread, off the pool) while leaving everything else — name, the + /// (now-irrelevant) sticky flag, the empty output edges — unchanged. The + /// default constructors stay `Serial`, so only the standalone-sort terminal + /// that opts in is affected. + #[test] + fn with_detached_advertises_detached_kind() { + let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); + let header = empty_header(); + let step = WriteBgzfFile::new(&path, &header, 1) + .unwrap() + .with_detached(DetachedGroup::Shared("test-io")); + let profile = step.profile(); + assert_eq!(profile.name, "WriteBgzfFile"); + assert_eq!(profile.kind, StepKind::Detached); + assert_eq!(step.detached_group(), DetachedGroup::Shared("test-io")); + assert_eq!(profile.output_queues.len(), 0); + } + + #[test] + fn header_only_round_trip() { + let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); + let header = empty_header(); + let step = WriteBgzfFile::new(&path, &header, 1).unwrap(); + let mut guard = step.state.lock(); + let mut state = guard.take().expect("state present"); + state.out.write_all(&BGZF_EOF).unwrap(); + state.out.flush().unwrap(); + drop(guard); + + let bytes = std::fs::read(&path).unwrap(); + assert!(bytes.len() >= 28, "BGZF EOF + header should be at least 28 bytes"); + assert_eq!(&bytes[0..2], &[0x1f, 0x8b], "BGZF/gzip magic at start"); + let tail = &bytes[bytes.len() - 28..]; + assert_eq!(tail, &BGZF_EOF, "file ends with BGZF EOF marker"); + } + + #[test] + fn new_with_handle_defers_header_until_resolved() { + let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); + let handle = HeaderHandle::new(); + let step = WriteBgzfFile::new_with_handle(&path, handle.clone(), 1, None).unwrap(); + + let bytes_before = std::fs::read(&path).unwrap(); + assert_eq!(bytes_before.len(), 0, "no bytes written until header resolves"); + + { + let mut guard = step.state.lock(); + let state = guard.as_mut().expect("state present"); + assert!(state.pending_header.is_some(), "handle still pending"); + let wrote = WriteBgzfFile::try_write_pending_header(state).unwrap(); + assert!(!wrote, "unresolved handle should yield without writing"); + assert!(state.pending_header.is_some(), "still pending after no-op probe"); + } + let bytes_mid = std::fs::read(&path).unwrap(); + assert_eq!(bytes_mid.len(), 0, "still nothing on disk after no-op probe"); + + handle.set(empty_header()).expect("first set"); + { + let mut guard = step.state.lock(); + let state = guard.as_mut().expect("state present"); + let wrote = WriteBgzfFile::try_write_pending_header(state).unwrap(); + assert!(wrote, "resolved handle should write"); + assert!(state.pending_header.is_none(), "pending slot cleared"); + state.out.flush().unwrap(); + } + let bytes_after = std::fs::read(&path).unwrap(); + assert!(bytes_after.len() >= 2, "header BGZF block emitted"); + assert_eq!(&bytes_after[0..2], &[0x1f, 0x8b], "BGZF/gzip magic at start"); + } + + #[test] + fn resolved_header_transform_runs_on_the_resolved_header() { + use std::sync::{Arc, Mutex as StdMutex}; + + let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); + let handle = HeaderHandle::new(); + + // Capture the header the transform is handed, to prove it is the + // runtime-resolved (aligner) header — not a build-time header. + let seen: Arc>> = Arc::new(StdMutex::new(None)); + let seen_for_closure = Arc::clone(&seen); + let transform: ResolvedHeaderTransform = Box::new(move |resolved: Header| { + *seen_for_closure.lock().unwrap() = Some(resolved.clone()); + Ok(resolved) + }); + + let step = + WriteBgzfFile::new_with_handle(&path, handle.clone(), 1, Some(transform)).unwrap(); + + // Nothing is written (and the transform must not run) until the handle + // resolves. + { + let mut guard = step.state.lock(); + let state = guard.as_mut().expect("state present"); + assert!(!WriteBgzfFile::try_write_pending_header(state).unwrap()); + } + assert!(seen.lock().unwrap().is_none(), "transform must not run before resolution"); + + // Resolve with a distinctive header standing in for the aligner's + // runtime-resolved header. + let resolved = Header::builder().add_comment("ALIGNER-PROVENANCE").build(); + handle.set(resolved.clone()).expect("set"); + { + let mut guard = step.state.lock(); + let state = guard.as_mut().expect("state present"); + assert!(WriteBgzfFile::try_write_pending_header(state).unwrap(), "resolved -> written"); + state.out.flush().unwrap(); + } + + // The transform ran, and it received the resolved header — so a + // post-align stage's transform is applied on top of the aligner header. + assert_eq!( + seen.lock().unwrap().as_ref(), + Some(&resolved), + "transform must receive the resolved header, preserving aligner provenance", + ); + assert!(std::fs::read(&path).unwrap().len() >= 2, "header block was written"); + } + + #[test] + fn new_with_handle_propagates_poison() { + let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); + let handle = HeaderHandle::new(); + let step = WriteBgzfFile::new_with_handle(&path, handle.clone(), 1, None).unwrap(); + + handle.poison(io::Error::new(io::ErrorKind::BrokenPipe, "aligner died")).unwrap(); + let mut guard = step.state.lock(); + let state = guard.as_mut().expect("state present"); + let err = WriteBgzfFile::try_write_pending_header(state).expect_err("poison"); + assert_eq!(err.kind(), io::ErrorKind::BrokenPipe); + assert_eq!(err.to_string(), "aligner died"); + } + + #[test] + fn drop_before_finish_does_not_append_eof_marker() { + // A sink dropped before the drained-finish path in `try_run` (e.g. a + // pipeline abort) must NOT append the BGZF EOF marker. Appending it + // would stamp a "complete stream" signature onto a truncated BAM, + // hiding the truncation from downstream readers. See the `Drop` doc. + let path = tempfile::NamedTempFile::new().unwrap(); + let path_buf = path.path().to_path_buf(); + let header = empty_header(); + // `new` eagerly writes the header (pending_header is None), so the + // only thing standing between this state and a valid EOF marker is + // the `try_run` drained-finish path, which we never reach. + let step = WriteBgzfFile::new(&path_buf, &header, 1).unwrap(); + drop(step); + + let bytes = std::fs::read(&path_buf).unwrap(); + assert!(bytes.len() >= 28, "header bytes should be on disk"); + let tail = &bytes[bytes.len() - 28..]; + assert_ne!(tail, &BGZF_EOF, "aborted output must not end with a valid BGZF EOF marker"); + } + + #[test] + fn drop_with_unresolved_handle_leaves_empty_file() { + let path = tempfile::NamedTempFile::new().unwrap(); + let path_buf = path.path().to_path_buf(); + let handle = HeaderHandle::new(); + let step = WriteBgzfFile::new_with_handle(&path_buf, handle, 1, None).unwrap(); + drop(step); + + let bytes = std::fs::read(&path_buf).unwrap(); + assert_eq!(bytes.len(), 0, "Drop with unresolved handle must skip EOF — see Drop doc"); + } +} diff --git a/crates/fgumi-pipeline-io/src/sink/write_raw.rs b/crates/fgumi-pipeline-io/src/sink/write_raw.rs new file mode 100644 index 000000000..4f01b08c2 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sink/write_raw.rs @@ -0,0 +1,187 @@ +//! `WriteRawFile` sink step: writes a byte stream verbatim to a file or stdout, +//! with **no** container header and **no** BGZF EOF marker. +//! +//! Unlike [`super::write_bgzf::WriteBgzfFile`] (which is BAM-specific — it emits +//! a BAM header on open and a BGZF EOF on drain), this sink just concatenates +//! the `bytes` of each block it receives. It backs FASTQ output: the chain's +//! FASTQ-encode step produces `DecompressedBlock`s of FASTQ text, which either +//! go straight here (plain output / stdout) or through `BgzfCompress` first +//! (`.gz`/`.bgz` output, producing `BgzfBlock`s — still just bytes to write). +//! +//! `Serial + Affinity::Writer + sticky`, matching `WriteBgzfFile`: exactly one +//! shared instance drains the (reorder-ordered) block stream to the sink. + +use std::fs::File; +use std::io::{self, BufWriter, Write}; +use std::path::Path; + +use parking_lot::Mutex; + +use crate::types::{BgzfBlock, DecompressedBlock}; +use fgumi_pipeline_core::{ + item::HeapSize, + step::{Affinity, Step, StepCtx, StepKind, StepOutcome, StepProfile}, +}; + +/// A pipeline block whose payload is a run of bytes to write verbatim. +pub trait RawBytesBlock: Send + HeapSize + 'static { + /// The bytes to write for this block. + fn bytes(&self) -> &[u8]; +} + +impl RawBytesBlock for DecompressedBlock { + fn bytes(&self) -> &[u8] { + &self.bytes + } +} + +impl RawBytesBlock for BgzfBlock { + fn bytes(&self) -> &[u8] { + &self.bytes + } +} + +/// `Serial + sticky` sink that writes each block's bytes verbatim to a file or +/// stdout. Generic over the block type so it serves both the plain +/// (`DecompressedBlock`) and BGZF-compressed (`BgzfBlock`) FASTQ tails. +pub struct WriteRawFile { + state: Mutex>>>, + /// Bytes appended once, after the last block, on a clean drain. Empty for + /// plain output; the 28-byte BGZF EOF marker for BGZF output so the `.gz` + /// stream is a complete, non-truncated BGZF file. + trailer: &'static [u8], + _marker: std::marker::PhantomData, +} + +impl WriteRawFile { + /// Open `path` for writing (`-` selects stdout), appending `trailer` once + /// after the final block on clean completion. No header is written. + /// + /// # Errors + /// + /// Returns I/O errors from opening the file. + pub fn new>(path: P, trailer: &'static [u8]) -> io::Result { + let inner: Box = if path.as_ref().as_os_str() == "-" { + Box::new(io::stdout()) + } else { + Box::new(File::create(path.as_ref())?) + }; + Ok(Self { + state: Mutex::new(Some(BufWriter::with_capacity(256 * 1024, inner))), + trailer, + _marker: std::marker::PhantomData, + }) + } +} + +impl Step for WriteRawFile { + type Input = B; + type Outputs = (); + + fn profile(&self) -> StepProfile { + StepProfile { + name: "WriteRawFile", + kind: StepKind::Serial, + sticky: true, + output_queues: vec![], + branch_ordering: vec![], + } + } + + fn affinity(&self) -> Affinity { + Affinity::Writer + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + let mut guard = self.state.lock(); + let Some(out) = guard.as_mut() else { + return Ok(StepOutcome::Finished); + }; + + if let Some(block) = ctx.input.pop() { + out.write_all(block.bytes())?; + return Ok(StepOutcome::Progress); + } + + if ctx.input.is_drained() { + // Clean end-of-stream: append the trailer (e.g. the BGZF EOF marker) + // exactly once, then flush and retire the sink. + if !self.trailer.is_empty() { + out.write_all(self.trailer)?; + } + out.flush()?; + let _ = guard.take(); + return Ok(StepOutcome::Finished); + } + Ok(StepOutcome::NoProgress) + } +} + +impl Drop for WriteRawFile { + /// Cleanup-only drop: flush buffered bytes but deliberately do **not** write + /// `trailer`. The trailer (e.g. the BGZF EOF marker for `.gz` output) is + /// written exclusively by the drained-finish path in `try_run`, which then + /// takes the state so this drop is a no-op for a cleanly-finished sink. If + /// state is still present here the stream was aborted mid-way, and stamping + /// the BGZF EOF onto a truncated `.gz` would make it look complete and hide + /// the truncation from readers — so we withhold it, exactly as + /// `WriteBgzfFile::drop` does. + fn drop(&mut self) { + if let Some(mut out) = self.state.lock().take() { + let _ = out.flush(); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn raw_bytes_block_exposes_payload() { + let plain = DecompressedBlock { batch_serial: 0, bytes: b"ACGT".to_vec() }; + assert_eq!(RawBytesBlock::bytes(&plain), b"ACGT"); + let bgzf = BgzfBlock { batch_serial: 0, bytes: b"\x1f\x8b".to_vec(), uncompressed_size: 4 }; + assert_eq!(RawBytesBlock::bytes(&bgzf), b"\x1f\x8b"); + } + + #[test] + fn profile_is_serial_writer_sink() { + let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); + let step = WriteRawFile::::new(&path, b"").unwrap(); + let profile = step.profile(); + assert_eq!(profile.name, "WriteRawFile"); + assert_eq!(profile.kind, StepKind::Serial); + assert!(profile.sticky); + assert_eq!(step.affinity(), Affinity::Writer); + } + + /// A sink dropped before the drained-finish path (an aborted stream) must + /// NOT append its trailer — stamping the BGZF EOF onto a truncated `.gz` + /// would hide the truncation. Mirrors + /// `WriteBgzfFile::drop_before_finish_does_not_append_eof_marker`. + #[test] + fn drop_before_finish_omits_trailer() { + let tmp = tempfile::NamedTempFile::new().unwrap(); + let path = tmp.path().to_path_buf(); + let trailer = b"\x1f\x8bTRAILER"; + let step = WriteRawFile::::new(&path, trailer).unwrap(); + // Write some payload bytes directly, then drop without draining (abort). + { + let mut guard = step.state.lock(); + guard.as_mut().unwrap().write_all(b"PAYLOAD").unwrap(); + } + drop(step); + + let bytes = std::fs::read(&path).unwrap(); + assert_eq!(bytes, b"PAYLOAD", "aborted stream must contain payload only, no trailer"); + assert!(!bytes.ends_with(trailer), "aborted stream must not end with the trailer"); + } + + #[test] + fn dash_path_selects_stdout_without_error() { + // `-` must construct a stdout-backed sink without touching the filesystem. + let step = WriteRawFile::::new("-", b"").unwrap(); + assert!(step.state.lock().is_some()); + } +} diff --git a/crates/fgumi-pipeline-io/src/sort/arena_ingest.rs b/crates/fgumi-pipeline-io/src/sort/arena_ingest.rs new file mode 100644 index 000000000..28cc840cb --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/arena_ingest.rs @@ -0,0 +1,2623 @@ +//! `ReadBlocks` serial arena-admit step, `InflateToArena` parallel step, and +//! `FindBoundariesAndSort` serial step. +//! +//! `ReadBlocks`: serial step that consumes raw `BgzfBlock`s, uses an [`ArenaPool`] +//! (capacity 1) to acquire/reuse arenas, reserves a fixed front region +//! `[0, FRONT_REGION)` per run for the straddler carry (Task 2), reserves an +//! uninit slot per block via `grow_uninit` at offsets `>= FRONT_REGION`, and +//! seals a run when its cumulative uncompressed bytes reach `run_cap = memory_limit`. +//! Sealed mid-stream runs emit [`ArenaBlock`]s with `seals_to_spill = true`; +//! the final residual run emits with `seals_to_spill = false`. +//! This is the second `unsafe` site in `fgumi-pipeline-io`; see CLAUDE.md +//! §"Approved hot-path unsafe (parallel-inflate sort ingest, fgumi-pipeline-io)" +//! for the full justification. +//! +//! `InflateToArena`: each worker decompresses one BGZF block directly into its +//! disjoint arena slot via [`fgumi_bgzf::decompress_into_slice`]. This is the first +//! `unsafe` site in `fgumi-pipeline-io`; see CLAUDE.md §"Approved hot-path +//! unsafe (parallel-inflate sort ingest, fgumi-pipeline-io)" for the +//! full justification and SAFETY invariant. +//! +//! `FindBoundariesAndSort`: serial step that scans a run's contiguous arena span +//! directly (no copy), skips the BAM header via [`crate::boundaries::bam_header_len`], +//! builds `(body_offset, block_size)` refs, sorts via +//! [`fgumi_sort::coordinate_chunk_from_arena_refs`], and emits +//! `SortChunkEvent::Residual` + `AllAnnounced`. Single-run / no-spill scope; +//! multi-run seal and straddler handling are 3b.4. +//! +//! Together these three steps form the BAM sort ingest front for the block-input +//! path. The chain builder's `add_sort` (`src/lib/pipeline/chains/builder.rs`) +//! wires them as `ReadBlocks → InflateToArena → FindBoundariesAndSort` (feeding +//! the Phase-1 spill/merge tail) when the sort source is raw BGZF blocks. + +use std::collections::VecDeque; +use std::io; +use std::sync::Arc; + +use fgumi_bgzf::{Decompressor, decompress_into_slice}; +use fgumi_sort::{ + ArenaPool, InMemoryChunk, PooledSegmentedBuf, RawSortKey, RecordRef, + coordinate_chunk_from_refs, extract_coordinate_key_inline, queryname_chunk_from_arena_refs, +}; + +use crate::types::BgzfBlock; + +use fgumi_pipeline_core::held::HeldSlot; +use fgumi_pipeline_core::item::{HeapSize, Ordered}; +use fgumi_pipeline_core::outputs::{OrderedBytesSingle, Single}; +use fgumi_pipeline_core::queues::QueueSpec; +use fgumi_pipeline_core::reorder::BranchOrdering; +use fgumi_pipeline_core::step::{DetachedGroup, Step, StepCtx, StepKind, StepOutcome, StepProfile}; +use fgumi_pipeline_core::{HeldRetry, Unpushed}; + +use crate::boundaries::bam_header_len; +use crate::sort::protocol::{MemoryChunkErased, SortChunkEvent}; + +// ============================================================================ +// Constants +// ============================================================================ + +/// Reserved front region at the start of every run's arena, in bytes. +/// +/// Block 0 always lands at offset `FRONT_REGION`. The region `[0, FRONT_REGION)` +/// is left uninitialized by `ReadBlocks` and filled by `FindBoundariesAndSort` +/// when it writes a carried straddler from the previous run (Task 2, 3b.4). +/// Must be larger than any single BAM record; 8 MiB covers all realistic records. +pub const FRONT_REGION: usize = 8 * 1024 * 1024; + +/// Maximum uncompressed size of a single BGZF block: the BGZF spec bounds an +/// uncompressed block at 2^16 = 65536 bytes. Used to cap a run's seal budget +/// (`run_cap`) so the single-segment-per-run invariant +/// `FRONT_REGION + Σ(block ISIZE) ≤ segment_size` holds with a true upper bound +/// of margin (see `ReadBlocks::new`). This must be `>=` any real block's ISIZE; +/// 65536 is the exact spec maximum, so the headroom is a genuine bound rather +/// than relying on the `< run_cap` cumulative check being off by one. +const MAX_BGZF_BLOCK: usize = 1 << 16; + +/// Bytes ahead of the current scan cursor to software-prefetch in the +/// `FindBoundariesAndSort` boundary+key scan. Chosen from a microbench over a +/// cold ~2.6 GiB arena (matching the ~220 B/record production density): 2 KiB +/// gave the best speedup (~15%); ≤1 KiB was negligible (too little lead time), +/// 4 KiB matched 2 KiB. At ~220 B/record this is ~9 records of lead. +const SCAN_PREFETCH_DISTANCE: usize = 2048; + +/// Max blocks `ReadBlocks` admits per `try_run` dispatch. +/// +/// `ReadBlocks` runs on the coordination driver (`StepKind::Detached`), whose +/// drain-first `round_robin_dispatch` restarts the whole downstream walk on every +/// `Progress` (see `runtime::driver`): admitting one block per dispatch made this +/// serial admit the input-side throughput bottleneck — the raw-block reader +/// upstream backed up while the parallel `InflateToArena` pool workers starved +/// (empty pops), leaving the pool short of full CPU occupancy. Admitting a bounded +/// batch per dispatch amortises the per-dispatch queue/reorder overhead so the +/// admit keeps the inflaters saturated, while the cap bounds how long the driver +/// dwells on admit (fairness vs the group's other steps) and how far it runs +/// ahead of the byte/count-bounded output queue (which backpressures early anyway). +const ADMIT_BATCH: usize = 64; + +/// Software-prefetch (read, into L1, temporal) the cache line containing `byte`. +/// The `FindBoundariesAndSort` scan walks the run's arena cold (it was written by +/// the parallel `InflateToArena` workers long before this serial scan runs), so it +/// is latency-bound on cache misses; prefetching [`SCAN_PREFETCH_DISTANCE`] ahead +/// hides them. `cfg`-gated to the supported architectures; a no-op elsewhere. +/// +/// SAFETY note: this is the only place `fgumi-pipeline-io` uses an architecture +/// intrinsic. Both `prfm` (`aarch64`) and `_mm_prefetch` (`x86_64`) are +/// *non-faulting hints* — they never read or write observable memory and never +/// trap, even on an unmapped address. `byte` is a live `&u8` (the caller +/// bounds-checks the index), so the pointer is valid to name. See CLAUDE.md +/// §"Approved hot-path unsafe (parallel-inflate sort ingest, fgumi-pipeline-io)". +#[inline] +fn prefetch_read_l1(byte: &u8) { + let ptr: *const u8 = byte; + #[cfg(target_arch = "aarch64")] + #[allow(unsafe_code)] + // SAFETY: `prfm pldl1keep` is a non-faulting prefetch hint over a valid pointer. + unsafe { + core::arch::asm!( + "prfm pldl1keep, [{p}]", + p = in(reg) ptr, + options(nostack, readonly, preserves_flags), + ); + } + #[cfg(target_arch = "x86_64")] + #[allow(unsafe_code)] + // SAFETY: `_mm_prefetch` is a non-faulting prefetch hint over a valid pointer. + unsafe { + core::arch::x86_64::_mm_prefetch::<{ core::arch::x86_64::_MM_HINT_T0 }>(ptr.cast()); + } + #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))] + { + let _ = ptr; // no portable stable prefetch; hint is a no-op on other arches + } +} + +// ============================================================================ +// ReadBlocks: serial arena-admit step +// ============================================================================ + +/// `Serial` step that consumes raw [`BgzfBlock`]s, acquires arenas from an +/// [`ArenaPool`] (capacity 1), grows the whole arena segment to full capacity +/// once at acquire (so block slots can be sliced by inflate workers without +/// further `&mut` growth), reserves a front region `[0, FRONT_REGION)` per run, +/// assigns each block a slot at an arithmetic offset (`FRONT_REGION` + prefix-sum +/// of ISIZE), seals a run when cumulative uncompressed bytes reach `memory_limit`, +/// and emits one [`ArenaBlock`] per block. +/// +/// **Eager (streaming) emission.** Blocks are emitted *as they are read*, not +/// buffered until the run seals: the `Arc` is created at +/// acquire and the slot offset is known immediately (from the block's ISIZE, +/// without inflating), so an inflate worker can start on block 0 while +/// `ReadBlocks` is still reading block 5000 — the Read‖Inflate overlap. Exactly +/// ONE block is withheld at a time (`last_block`): `is_last_of_run` and +/// `seals_to_spill` are run-level facts known only when the run seals, so the +/// most-recently-admitted block is held until either a successor in the same run +/// arrives (confirming it is *not* last → emit it tagged `false`) or the run +/// seals (it *is* last → stamp `is_last_of_run`/`seals_to_spill` and emit). This +/// single held block coincides with the deferred-seal block, so the two are one +/// mechanism. +/// +/// Mid-stream seals set `seals_to_spill = true`; the final (residual) run sets +/// `seals_to_spill = false`. Only the LAST block of a run carries a meaningful +/// `is_last_of_run`/`seals_to_spill`; downstream (`FindBoundariesAndSort`) reads +/// those fields only on the `is_last_of_run` block, so non-last blocks emit +/// `is_last_of_run = false, seals_to_spill = false`. +/// +/// With pool capacity 1, the next run's arena cannot be acquired until the +/// prior run's [`Arc`] drops (i.e. after `CompressSpill` +/// consumes and releases it). If acquisition fails, `try_run` returns +/// `NoProgress` (backpressure) and retries on the next call. +pub struct ReadBlocks { + /// Bounded pool (capacity 1) that owns arena storage and reuses it across runs. + pool: Arc, + /// Size of each arena segment (`FRONT_REGION + run_cap + MAX_BGZF_BLOCK`); the + /// whole segment is grown live once at acquire. + segment_size: usize, + /// `Arc`-shared arena for the current run, grown to full capacity at acquire so + /// blocks (and their inflate slots) can be handed out eagerly without further + /// `&mut` growth; `None` until the first block of the current run is admitted, + /// and again after the run seals (dropping this step's handle so the arena can + /// return to the pool once all emitted clones release it). + arena: Option>, + /// Next slot offset within the current run's arena (set to `FRONT_REGION` on a + /// fresh run, then advanced by each block's ISIZE). + run_offset: u64, + /// The most-recently-admitted block of the current run, withheld from emission + /// until we learn whether another block joins this run, so `is_last_of_run` / + /// `seals_to_spill` (run-level facts known only at seal) can be stamped on it. + last_block: Option, + /// A block popped from the input but not yet admitted because a deferred + /// seal fired first and the pool is momentarily exhausted (the just-sealed + /// run's arena has not returned yet). Re-admitted on a later `try_run` once + /// the pool yields an arena. Holding it here prevents losing the block. + deferred_block: Option, + /// Emitted `ArenaBlock`s ready to push to the output queue. + emit: VecDeque, + /// Held output slot for the backpressure retry path. + held: HeldSlot>, + /// Monotonically-increasing ordinal assigned to each admitted block (across all runs). + next_ordinal: u64, + /// Current run sequence number. + run_seq: u32, + /// Cumulative uncompressed bytes in the current run (resets on each seal). + run_cumulative: usize, + /// Budget per run: seal when `run_cumulative >= run_cap`. + run_cap: usize, + /// Armed when the current run reached `run_cap`, but the seal is DEFERRED + /// until the next block arrives. Deferring by one block guarantees that a + /// spill seal is only ever fired when at least one more block follows (so + /// the straddler carry from a spilled run always has a subsequent run to + /// complete it). If the input drains while a seal is armed, the run is + /// sealed as the RESIDUAL instead (the final run's last record always + /// completes at EOF, so the carry is empty). Without this, the budget seal + /// could fire on the very last input block, emitting a trailing *spill* + /// whose carry has no following run — surfacing as a spurious "truncated + /// BAM" error. + seal_pending: bool, + /// Set to `true` after the final run has been emitted. + finished: bool, + /// Byte limit for the output queue. + output_byte_limit: u64, +} + +impl ReadBlocks { + /// Create a new `ReadBlocks` step. + /// + /// - `memory_limit`: bytes of record data a run holds before sealing — the + /// FULL in-memory budget (`--max-memory × threads`, via + /// `resolve_memory_budget`). This is the SAME spill trigger legacy uses: + /// the front spills to disk ONLY when the data exceeds this budget. For + /// data that fits the budget the result is exactly ONE in-memory run, ZERO + /// disk spills, and one radix sort over the full budget — legacy's + /// in-memory algorithm and footprint, but with the no-copy arena ingest. + /// - `output_byte_limit`: byte-bounds the output queue. + /// + /// The arena segment is sized so one full run always fits within ONE segment + /// (`FRONT_REGION` + the run budget + one deferred-seal tipping block), which + /// keeps the `FindBoundariesAndSort` contiguous-slice scan valid WITHOUT + /// capping the run below the budget. (The earlier design fixed the segment at + /// 256 MiB and capped `run_cap` to fit it, which forced the front to spill + /// even when the data fit `--max-memory` — a wall-clock regression vs legacy. + /// Sizing the segment to the budget instead removes that forced spill.) + #[must_use] + pub fn new(memory_limit: usize, output_byte_limit: u64) -> Self { + // Seal a run only when its record data reaches the full budget (legacy's + // trigger); never cap below it. + let run_cap = memory_limit; + // One run = front region + up to `run_cap` of data + the single block that + // tips `run_cumulative` over `run_cap` (the deferred seal holds the NEXT + // block for the following run, so the current run overshoots by at most one + // block). Sizing the segment to that upper bound guarantees a run is + // gap-free within one segment, so the scan's single `arena.slice(..)` over + // `[scan_start, run_end)` never spans a segment boundary. + let segment_size = FRONT_REGION + run_cap + MAX_BGZF_BLOCK; + let pool = ArenaPool::new(1, segment_size); + Self { + pool, + segment_size, + arena: None, + run_offset: 0, + last_block: None, + deferred_block: None, + emit: VecDeque::new(), + held: HeldSlot::new(), + next_ordinal: 0, + run_seq: 0, + run_cumulative: 0, + run_cap, + seal_pending: false, + finished: false, + output_byte_limit, + } + } + + /// Ensure the current run's arena is acquired, grown to full capacity, and + /// `Arc`-shared, with `run_offset` reset to `FRONT_REGION`. + /// + /// Returns `true` if the arena is ready, `false` if the pool is exhausted + /// (backpressure: the previous run's chunk has not been consumed yet). + /// + /// # Side effects + /// + /// Acquires an arena from the pool, calls `reserve_full_capacity`, then grows + /// the ENTIRE segment to `segment_size` live bytes via a single `unsafe` + /// `grow_uninit` call (see the `// SAFETY:` comment inside), and wraps it in an + /// `Arc`. This `grow_uninit` is the second `unsafe` site in + /// `fgumi-pipeline-io`; see CLAUDE.md §"Approved hot-path unsafe (parallel-inflate + /// sort ingest, fgumi-pipeline-io)" — the `ReadBlocks::ensure_arena` grow-once + /// bullet — for the full justification. Growing once, before sharing, is the + /// soundness keystone: every block's slot offset is then computed arithmetically + /// (`FRONT_REGION` + prefix-sum of ISIZE) so no further `&mut` growth is needed + /// while inflate workers hold disjoint `slice_mut` views. + fn ensure_arena(&mut self) -> bool { + if self.arena.is_some() { + return true; + } + let Some(mut arena) = self.pool.try_acquire() else { + return false; + }; + arena.reserve_full_capacity(); + // Grow the whole segment to `segment_size` live bytes in ONE call, here on + // the serial admit path BEFORE the arena is wrapped in an `Arc` or any slice + // is handed out — so no concurrent borrow can exist during the grow. + // SAFETY: sole writer, no live borrow (the `Arc` is created only after this + // returns). `segment_size` is exactly the pool's segment size and the + // capacity `reserve_full_capacity` just guaranteed, so `grow_uninit` performs + // no realloc and the slot stays in one segment. `u8` has no validity + // invariant; every live byte is written exactly once — block slots by their + // inflate worker, the `[0, FRONT_REGION)` carry region by `FindBoundariesAndSort` + // — before any read. Bytes never read (the unused tail past a run's end, and + // the front region when there is no straddler) are never observed: `FBS` only + // ever slices `[scan_start, run_end)`. + #[allow(unsafe_code)] + let _all = unsafe { arena.grow_uninit(self.segment_size) }; + // `ArenaPool::try_acquire` already returns the `PooledSegmentedBuf` wrapper + // (it wraps internally so a caller cannot orphan a pooled arena — the hang + // that motivated making `PooledSegmentedBuf::pooled` private), so the arena + // is wrapped by construction here and only needs sharing. + self.arena = Some(Arc::new(arena)); + self.run_offset = FRONT_REGION as u64; + true + } + + /// Admit one [`BgzfBlock`]: ensure the arena exists, assign the block an + /// arithmetic slot offset, and emit it eagerly (the PRIOR withheld block, now + /// confirmed to have a same-run successor, is pushed to the output queue; THIS + /// block becomes the new withheld `last_block`). + /// + /// Returns `Ok(())` on success, or `Err(b)` handing the block back if the pool + /// is exhausted (the arena could not be acquired) — so the caller never loses + /// the block. + /// + /// No `unsafe` here: the arena was grown to full capacity once by `ensure_arena`, + /// so the slot `(offset, len)` is already a live region of the shared arena; + /// the inflate worker writes it via `slice_mut`. + pub(crate) fn admit_block(&mut self, b: BgzfBlock) -> Result<(), BgzfBlock> { + if !self.ensure_arena() { + return Err(b); + } + let arena = self.arena.as_ref().expect("arena present after ensure_arena"); + + let len = b.uncompressed_size; + let offset = self.run_offset; + self.run_offset += u64::from(len); + + let ordinal = self.next_ordinal; + self.next_ordinal += 1; + self.run_cumulative += len as usize; + + // `is_last_of_run` / `seals_to_spill` are run-level facts known only at seal, + // so the just-admitted block is withheld as the tentative last block. The + // PRIOR withheld block now has a same-run successor → it is NOT last → emit it. + let block = ArenaBlock { + arena: Arc::clone(arena), + ordinal, + offset, + len, + block: b.bytes, + is_last_of_run: false, + run_seq: self.run_seq, + seals_to_spill: false, + }; + if let Some(prev) = self.last_block.replace(block) { + self.emit.push_back(prev); + } + Ok(()) + } + + /// Seal the current run: stamp the withheld `last_block` as the run's final + /// block (`is_last_of_run = true` plus the run's `seals_to_spill`), emit it, + /// drop this step's arena handle (so the arena can return to the pool once all + /// emitted clones release it), and advance `run_seq`. + /// + /// All non-last blocks of the run were already emitted eagerly by `admit_block`; + /// only the single withheld block remains. + /// + /// `seals_to_spill`: `true` for mid-stream seals (the run becomes a disk spill), + /// `false` for the final residual run. + /// + /// A `seal_run` call always follows at least one admit for the run, so + /// `last_block` is `Some`; the `try_run` EOF guard returns before sealing an + /// arena-less, blockless state. + fn seal_run(&mut self, seals_to_spill: bool) { + if let Some(mut last) = self.last_block.take() { + last.is_last_of_run = true; + last.seals_to_spill = seals_to_spill; + self.emit.push_back(last); + } + // Drop our handle to the sealed run's arena; the emitted blocks (and their + // downstream consumers) keep it alive until they release their `Arc` clones, + // at which point `PooledSegmentedBuf::drop` returns it to the capacity-1 pool. + self.arena = None; + self.run_seq += 1; + self.run_cumulative = 0; + self.run_offset = 0; + } + + /// Test seam: admit blocks, then freeze the arena and drain all emitted + /// [`ArenaBlock`]s into a `Vec` for direct inspection, bypassing the pipeline + /// framework. Uses `seals_to_spill = false` (residual) for the single flush. + #[cfg(test)] + pub(crate) fn seal_and_drain_for_test(&mut self) -> Vec { + self.seal_run(false); + self.emit.drain(..).collect() + } +} + +impl Step for ReadBlocks { + type Input = BgzfBlock; + type Outputs = OrderedBytesSingle; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "ReadBlocks", + // Off-pool on the coordination driver (N+2): the serial arena-admit + // runs on a dedicated thread instead of stealing a pool worker slot + // from the parallel inflaters. Detached collapses the `ByItemOrdinal` + // output to `None` exactly as `Serial` did (transport-identical). + kind: StepKind::Detached, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], + branch_ordering: vec![BranchOrdering::ByItemOrdinal], + } + } + + fn detached_group(&self) -> DetachedGroup { + DetachedGroup::Shared(crate::sort::SORT_COORD_GROUP) + } + + // A single cohesive deferred-seal state machine: held-output retry, seal + // arming/execution, batched admit, and staged-emit drain are tightly coupled + // by `seal_pending` / `deferred_block` / `held` and are clearer read top to + // bottom than split across helpers that would each need the same state. + #[allow(clippy::too_many_lines)] + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + // 1. Retry any held output first (backpressure path). + if let Some(unpushed) = self.held.take() { + match ctx.outputs.retry(unpushed) { + Ok(()) => {} + Err(again) => { + self.held.put(again); + return Ok(StepOutcome::Contention); + } + } + } + + // 2. Drain any staged emitted ArenaBlocks before accepting new input. + if let Some(block) = self.emit.pop_front() { + match ctx.outputs.push(block) { + Ok(()) => return Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held.put(unpushed); + return Ok(StepOutcome::Progress); + } + } + } + + // 3. If finished and emit is empty, we are done. + if self.finished { + return Ok(StepOutcome::Finished); + } + + // 4. Re-admit a block held over from a prior deferred seal, if any. The + // seal already fired; we just need an arena to land this block in. + if self.deferred_block.is_some() { + // Drain any still-staged emit blocks first so the pool can recycle. + if let Some(out) = self.emit.pop_front() { + match ctx.outputs.push(out) { + Ok(()) => return Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held.put(unpushed); + return Ok(StepOutcome::Progress); + } + } + } + // Acquire the arena BEFORE consuming the held block, so a pool + // exhaustion does not drop it. + if !self.ensure_arena() { + // Pool still exhausted (the just-sealed run's arena has not + // returned yet): keep holding the block and backpressure. + return Ok(StepOutcome::NoProgress); + } + let block = self.deferred_block.take().expect("deferred_block present"); + match self.admit_block(block) { + Ok(()) => {} + Err(b) => { + // ensure_arena just returned true, so admit cannot fail; restore + // the block defensively rather than drop it. + self.deferred_block = Some(b); + debug_assert!(false, "admit_block must succeed after ensure_arena"); + return Ok(StepOutcome::NoProgress); + } + } + if self.run_cumulative >= self.run_cap { + self.seal_pending = true; + } + return Ok(StepOutcome::Progress); + } + + // 5. Pop and admit BgzfBlocks — BATCHED (up to `ADMIT_BATCH`) to keep the + // parallel inflaters fed. See `ADMIT_BATCH`: one-per-dispatch made this + // admit the input-side bottleneck under the coordination driver's + // drain-first restart. Only the clean admit path is batched; a seal + // boundary and + // output backpressure both break the batch and return, preserving the + // exact single-dispatch semantics of the delicate deferred-seal state + // machine (a seal is rare — ~one per spill). + let mut admitted_any = false; + for _ in 0..ADMIT_BATCH { + let Some(block) = ctx.input.pop() else { break }; + // A seal is armed from a previous block hitting the budget. The + // arrival of THIS block proves a following run exists, so it is safe + // to seal the previous run as a disk spill (its straddler carry will + // be completed by the run this block opens). End the batch here: the + // just-sealed run's arena must be consumed before the next admit. + if self.seal_pending { + self.seal_run(true); // deferred mid-stream seal → spill + self.seal_pending = false; + // The just-sealed run's arena is held (as an Arc) in `emit` until + // it is pushed downstream and consumed, so the capacity-1 pool is + // momentarily empty. Hold this block and re-admit it once the + // arena returns (step 4 on a later call). Drain `emit` now. + self.deferred_block = Some(block); + if let Some(out) = self.emit.pop_front() { + match ctx.outputs.push(out) { + Ok(()) => return Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held.put(unpushed); + return Ok(StepOutcome::Progress); + } + } + } + return Ok(StepOutcome::Progress); + } + match self.admit_block(block) { + Ok(()) => {} + Err(b) => { + // Pool exhausted on a fresh run's first block: hold the block and + // backpressure until the prior run's chunk is consumed and its + // arena returns. (Mid-run admits never fail — the arena is + // already acquired — so this is the cross-run boundary case.) + // If we already admitted this batch, that IS progress. + self.deferred_block = Some(b); + return Ok(if admitted_any { + StepOutcome::Progress + } else { + StepOutcome::NoProgress + }); + } + } + admitted_any = true; + // Check if we hit the run budget; arm the deferred seal. The next + // loop iteration observes `seal_pending` and runs the seal path above. + if self.run_cumulative >= self.run_cap { + self.seal_pending = true; + } + // Drain staged emit eagerly within the batch so the inflaters see + // blocks promptly and `emit` stays bounded. On output backpressure, + // hold and return — the batch ends; the rest is picked up next + // dispatch (step 1 retries the held item first). + while let Some(out) = self.emit.pop_front() { + match ctx.outputs.push(out) { + Ok(()) => {} + Err(unpushed) => { + self.held.put(unpushed); + return Ok(StepOutcome::Progress); + } + } + } + } + if admitted_any { + return Ok(StepOutcome::Progress); + } + + // 6. No input available. + if !ctx.input.is_drained() { + return Ok(StepOutcome::NoProgress); + } + + // 7. Input fully drained: seal the final (residual) run. A pending seal + // is DOWNGRADED to a residual here — there is no following run to + // complete a straddler carry, and the final run's last record always + // completes at EOF, so the carry is empty. `seal_run(false)` emits + // the residual. + self.seal_pending = false; + // Only seal if a block is withheld; if no block is held and the arena is + // None, there was no data at all after the previous seal — i.e. truly + // empty input (no blocks ever admitted). NOTE: the deferred-seal design + // guarantees the *final* run is always a residual (step 5 only seals a + // spill once a following block proves a next run exists; otherwise this + // step downgrades the pending seal to a residual), so this branch is NOT + // reached after any data — it is the empty-input case only. + if self.arena.is_none() && self.last_block.is_none() { + self.finished = true; + return Ok(StepOutcome::Finished); + } + self.seal_run(false); // final run → residual + self.finished = true; + + if let Some(block) = self.emit.pop_front() { + match ctx.outputs.push(block) { + Ok(()) => return Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held.put(unpushed); + return Ok(StepOutcome::Progress); + } + } + } + + Ok(StepOutcome::Finished) + } + + fn new_worker_copy(&self) -> Self { + // Serial steps are never cloned by the framework; this is unreachable. + panic!("ReadBlocks is Serial — new_worker_copy should never be called") + } +} + +// ============================================================================ +// Item types +// ============================================================================ + +/// Input to `InflateToArena`. +/// +/// Carries a grown-but-uninit slot `(offset, len)` in `arena` plus the full +/// raw BGZF `block` bytes to inflate into it. The slot was reserved by +/// `grow_uninit` on the serial admit path; it must be fully written by the +/// inflate worker before any read. +pub struct ArenaBlock { + /// The shared arena into which this block's bytes will be inflated. + pub arena: Arc, + /// Global ordinal, used by `ByItemOrdinal` reordering so downstream steps + /// receive blocks in the original file order. + pub ordinal: u64, + /// Byte offset of this block's slot within the arena (returned by + /// `grow_uninit`). + pub offset: u64, + /// Uncompressed size (ISIZE from the BGZF footer == slot length). + pub len: u32, + /// Complete raw BGZF block bytes (header + deflate payload + footer). + pub block: Vec, + /// `true` if this is the last block of the current run (e.g. a BAM file + /// segment); used by downstream steps to detect run boundaries. + pub is_last_of_run: bool, + /// Run sequence number; increments with each new run. + pub run_seq: u32, + /// `true` if this run seals to a disk spill (mid-stream seal); `false` if + /// this is the final residual run (in-memory, not spilled). + pub seals_to_spill: bool, +} + +impl HeapSize for ArenaBlock { + fn heap_size(&self) -> usize { + self.block.len() + } +} + +impl Ordered for ArenaBlock { + fn ordinal(&self) -> u64 { + self.ordinal + } +} + +/// Completion token emitted by `InflateToArena` after successfully inflating +/// one block. The decompressed bytes now live in `arena` at the byte range +/// `offset..offset + len`. The token carries no heap data of its own +/// (`heap_size == 0`). +pub struct InflatedBlock { + /// The arena holding the decompressed bytes. + pub arena: Arc, + /// Global ordinal (same value as the originating `ArenaBlock`). + pub ordinal: u64, + /// Byte offset of the decompressed data within the arena. + pub offset: u64, + /// Byte length of the decompressed data. + pub len: u32, + /// Forwarded from `ArenaBlock`. + pub is_last_of_run: bool, + /// Forwarded from `ArenaBlock`. + pub run_seq: u32, + /// Forwarded from `ArenaBlock`: `true` if this run seals to a disk spill. + pub seals_to_spill: bool, +} + +impl HeapSize for InflatedBlock { + fn heap_size(&self) -> usize { + 0 + } +} + +impl Ordered for InflatedBlock { + fn ordinal(&self) -> u64 { + self.ordinal + } +} + +// ============================================================================ +// Step +// ============================================================================ + +/// `Parallel + ByItemOrdinal` step that decompresses each `ArenaBlock`'s BGZF +/// bytes directly into its pre-reserved arena slot. +/// +/// Each parallel worker clone holds its own [`Decompressor`] (allocated via +/// `new_worker_copy`) so there is no contention on the decompression state. +pub struct InflateToArena { + decompressor: Decompressor, + held: HeldSlot>, + output_byte_limit: u64, +} + +impl InflateToArena { + /// Create a new `InflateToArena` step. + /// + /// `output_byte_limit` bounds the byte-counted output queue (tokens carry + /// `heap_size == 0`, so this mainly controls queue depth via the framework's + /// backpressure mechanism). + #[must_use] + pub fn new(output_byte_limit: u64) -> Self { + Self { decompressor: Decompressor::new(), held: HeldSlot::new(), output_byte_limit } + } + + /// Decompress `item.block` into the arena slot `(item.offset, item.len)`, + /// returning an [`InflatedBlock`] completion token on success. + /// + /// # Errors + /// + /// Returns an `io::Error` if BGZF decompression fails or if the + /// decompressed length does not match `item.len`. + /// + /// # Safety (caller contract) + /// + /// See the `#[allow(unsafe_code)]` block inside — the caller must have + /// reserved the slot via `grow_uninit` before constructing the `ArenaBlock`. + // `pub(crate)` so the unit test below can call it directly without wiring up + // the full pipeline framework. + pub(crate) fn inflate_one(&mut self, item: ArenaBlock) -> io::Result { + let ArenaBlock { + arena, + ordinal, + offset, + len, + block, + is_last_of_run, + run_seq, + seals_to_spill, + } = item; + + // SAFETY: `(offset, len)` was reserved by `grow_uninit` on the serial + // ReadBlocks admit path before this `ArenaBlock` was enqueued, so the + // slot is live within the arena's allocated storage. The ISIZE + // prefix-sum partitions the arena into non-overlapping slots, so this + // `&mut [u8]` aliases no other concurrent inflate worker's slice. + // `u8` has no validity invariant, so writing before reading is the only + // required contract — `decompress_into_slice` below fills every byte. + // `offset` is a u64 byte offset into the arena; on a 64-bit platform + // this always fits in usize — the arena itself cannot exceed + // `isize::MAX` bytes, which is the Rust allocation bound. + let offset_usize = + usize::try_from(offset).expect("arena offset must fit in usize on this platform"); + let len_usize = len as usize; // u32 always fits in usize + + #[allow(unsafe_code)] + let slot = unsafe { arena.slice_mut(offset_usize, len_usize) }; + + let n = decompress_into_slice(&block, &mut self.decompressor, slot)?; + debug_assert_eq!(n, len_usize, "decompressed length must match ISIZE"); + + Ok(InflatedBlock { arena, ordinal, offset, len, is_last_of_run, run_seq, seals_to_spill }) + } +} + +impl Step for InflateToArena { + type Input = ArenaBlock; + type Outputs = OrderedBytesSingle; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "InflateToArena", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], + branch_ordering: vec![BranchOrdering::ByItemOrdinal], + } + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + // Retry any held output first (backpressure path — mirror BgzfDecompress). + if let Some(unpushed) = self.held.take() { + match ctx.outputs.retry(unpushed) { + Ok(()) => {} + Err(again) => { + self.held.put(again); + // `Contention` keeps the worker alive for retry — `NoProgress` + // would let the framework silently drop the held item if input + // is also drained. + return Ok(StepOutcome::Contention); + } + } + } + + let Some(block) = ctx.input.pop() else { + // No input this call. If upstream is fully drained the held slot was + // already flushed by the Contention preamble above, so every item has + // been processed. For a Parallel step only the last clone to finish + // closes the shared output (gated by StepDrainCounter in the driver). + if ctx.input.is_drained() { + return Ok(StepOutcome::Finished); + } + return Ok(StepOutcome::NoProgress); + }; + + let inflated = self.inflate_one(block)?; + + match ctx.outputs.push(inflated) { + Ok(()) => Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held.put(unpushed); + Ok(StepOutcome::Progress) + } + } + } + + fn new_worker_copy(&self) -> Self { + Self::new(self.output_byte_limit) + } +} + +// ============================================================================ +// Arena sort strategy (per-order key extraction + seal) +// ============================================================================ + +/// Per-sort-order policy for the arena-front [`FindBoundariesAndSort`] scan. +/// +/// The strategy owns the accumulated per-record refs and knows how to (a) extract +/// a record's sort key from its arena-resident body and (b) at run seal, sort +/// those refs and wrap the shared arena into the correctly-typed +/// [`MemoryChunkErased`] — with zero record-body copies. `FindBoundariesAndSort` +/// is generic over this trait and monomorphised per order, so the per-record +/// [`push_record`](Self::push_record) call inlines into the boundary scan hot +/// loop with no dynamic dispatch. +pub trait ArenaSortStrategy: Send + 'static { + /// Reserve capacity for approximately `est_records` refs at the start of a + /// run, so the incremental per-record pushes do not reallocate a multi-GB ref + /// buffer mid-run. + fn reserve_for_run(&mut self, est_records: usize); + + /// Extract the sort key from `body` (the record's BAM body, `block_size` + /// prefix excluded) and accumulate a ref pointing at `(body_off, len)` in the + /// shared inflate arena. Called once per record on the boundary scan hot path. + /// + /// # Errors + /// + /// Returns an error if the record is invalid for this order (e.g. a + /// template-coordinate dropped-lane violation). Coordinate never errors. + fn push_record(&mut self, body: &[u8], body_off: u64, len: u32) -> io::Result<()>; + + /// Sort the refs accumulated for this run and wrap `arena` into the erased + /// chunk (zero record copies), resetting the accumulator for the next run. + fn seal(&mut self, arena: Arc, sort_threads: usize) -> MemoryChunkErased; + + /// A fresh, empty strategy carrying the same configuration — used to build a + /// worker copy of the step. `FindBoundariesAndSort` is `Serial`, so this only + /// ever constructs its single working instance. + #[must_use] + fn fresh(&self) -> Self + where + Self: Sized; +} + +/// Coordinate-order strategy: extracts the fixed `u64` coordinate key inline and +/// accumulates plain [`RecordRef`]s; seals to [`MemoryChunkErased::Coordinate`] +/// via [`coordinate_chunk_from_refs`]. +pub struct CoordinateStrategy { + /// BAM header reference-sequence count, used by [`extract_coordinate_key_inline`]. + n_ref: u32, + /// Coordinate-key refs accumulated across the current run's blocks. Filled + /// incrementally by the scan and consumed (via `mem::take`) at seal, leaving + /// an empty `Vec` for the next run. + refs: Vec, +} + +impl CoordinateStrategy { + /// Create a coordinate strategy for a header with `n_ref` reference sequences. + #[must_use] + pub fn new(n_ref: u32) -> Self { + Self { n_ref, refs: Vec::new() } + } +} + +impl ArenaSortStrategy for CoordinateStrategy { + #[inline] + fn reserve_for_run(&mut self, est_records: usize) { + self.refs.reserve(est_records); + } + + #[inline] + fn push_record(&mut self, body: &[u8], body_off: u64, len: u32) -> io::Result<()> { + let sort_key = extract_coordinate_key_inline(body, self.n_ref); + self.refs.push(RecordRef::new(sort_key, body_off, len)); + Ok(()) + } + + fn seal(&mut self, arena: Arc, sort_threads: usize) -> MemoryChunkErased { + let refs = std::mem::take(&mut self.refs); + MemoryChunkErased::Coordinate(coordinate_chunk_from_refs(arena, refs, sort_threads)) + } + + fn fresh(&self) -> Self { + Self::new(self.n_ref) + } +} + +/// Template-coordinate strategy: wraps a [`fgumi_sort::TemplateArenaAccumulator`], +/// which owns the library / cell-barcode / MI and `--key-types` narrowed-lane +/// machinery the template key needs; accumulates arena-pointing refs and seals to +/// [`MemoryChunkErased::TemplateCoordinate`] (an arena-backed +/// `InMemoryChunk`), byte-identical to the owned `TemplateChunkSorter`. +pub struct TemplateStrategy { + acc: fgumi_sort::TemplateArenaAccumulator, +} + +impl TemplateStrategy { + /// Wrap a template accumulator (built from the header via + /// [`fgumi_sort::TemplateArenaAccumulator::from_header`]). + #[must_use] + pub fn new(acc: fgumi_sort::TemplateArenaAccumulator) -> Self { + Self { acc } + } +} + +impl ArenaSortStrategy for TemplateStrategy { + #[inline] + fn reserve_for_run(&mut self, est_records: usize) { + self.acc.reserve(est_records); + } + + #[inline] + fn push_record(&mut self, body: &[u8], body_off: u64, len: u32) -> io::Result<()> { + self.acc + .push(body, body_off, len) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, format!("{e:#}"))) + } + + fn seal(&mut self, arena: Arc, sort_threads: usize) -> MemoryChunkErased { + MemoryChunkErased::TemplateCoordinate(self.acc.seal(arena, sort_threads)) + } + + fn fresh(&self) -> Self { + Self { acc: self.acc.fresh() } + } +} + +/// Queryname-order strategy (lexicographic or natural, selected by the key type +/// `K` + `wrap`): extracts the embedded read-name key from each record and +/// accumulates `(key, offset, len)` refs pointing into the shared arena; at seal +/// it comparator-sorts the refs (variable-length names are not radix-able) into +/// an arena-backed [`InMemoryChunk`], wrapped into the matching erased variant. +/// +/// The record bodies stay in the arena (zero-copy); only the small name bytes are +/// owned by each key, exactly as the legacy queryname sort. Queryname's +/// tie order is unspecified (the integration parity gate is name-order, not +/// byte-identity), so one globally-sorted chunk per run is correct. +pub struct QuerynameStrategy { + /// `(key, body_offset, len)` refs accumulated across the current run's blocks. + refs: Vec<(K, u64, u32)>, + /// Bounded rayon pool (sized to `sort_threads`) the per-run comparator sort + /// installs into, so the parallel sort does not oversubscribe the pipeline's + /// worker pool on a spill. Built once on the first seal, reused; `None` on a + /// fresh copy. Mirrors [`TemplateArenaAccumulator`](fgumi_sort::TemplateArenaAccumulator)'s pool. + sort_pool: Option, + /// Erases the sorted `InMemoryChunk` into the correct `MemoryChunkErased` + /// arm (`QuerynameLex` for the lex key, `QuerynameNatural` for the natural + /// key). A `fn` pointer so [`fresh`](ArenaSortStrategy::fresh) can copy it. + wrap: fn(InMemoryChunk) -> MemoryChunkErased, +} + +impl QuerynameStrategy { + /// Build a queryname strategy that erases its sealed chunk via `wrap`. + #[must_use] + pub fn new(wrap: fn(InMemoryChunk) -> MemoryChunkErased) -> Self { + Self { refs: Vec::new(), sort_pool: None, wrap } + } + + /// Test-only observation of the bounded sort pool's thread count, so a test + /// can assert the Phase-1 `sort_threads` value actually SIZED the worker pool + /// (the runtime effect), not merely that it was plumbed. `None` until the + /// first [`seal`](ArenaSortStrategy::seal) builds the pool. + #[cfg(test)] + pub(crate) fn sort_pool_threads(&self) -> Option { + self.sort_pool.as_ref().map(rayon::ThreadPool::current_num_threads) + } +} + +impl ArenaSortStrategy for QuerynameStrategy { + #[inline] + fn reserve_for_run(&mut self, est_records: usize) { + self.refs.reserve(est_records); + } + + #[inline] + fn push_record(&mut self, body: &[u8], body_off: u64, len: u32) -> io::Result<()> { + // Queryname keys are EMBEDDED_IN_RECORD: the name lives in `body`, so the + // key is reconstructed straight from it (no SortContext needed). + self.refs.push((K::extract_from_record(body), body_off, len)); + Ok(()) + } + + fn seal(&mut self, arena: Arc, sort_threads: usize) -> MemoryChunkErased { + let refs = std::mem::take(&mut self.refs); + let pool = self.sort_pool.get_or_insert_with(|| { + rayon::ThreadPoolBuilder::new() + .num_threads(sort_threads.max(1)) + .thread_name(|i| format!("qname-sort-{i}")) + .build() + .expect("build bounded queryname-sort rayon pool") + }); + let wrap = self.wrap; + pool.install(move || wrap(queryname_chunk_from_arena_refs(arena, refs))) + } + + fn fresh(&self) -> Self { + Self::new(self.wrap) + } +} + +// ============================================================================ +// FindBoundariesAndSort step +// ============================================================================ + +/// `Serial` step that scans each run's contiguous arena span directly (no copy), +/// skips the BAM header on run 0 via [`bam_header_len`], accumulates one ref per +/// record through the order's [`ArenaSortStrategy`], seals that into an arena-backed +/// chunk (each strategy picks its own builder — see [`CoordinateStrategy::seal`]), +/// and emits `SortChunkEvent::Spill` (for sealed mid-stream runs) or +/// `SortChunkEvent::Residual` (for the final run), followed by +/// `SortChunkEvent::AllAnnounced`. +/// +/// **Straddler handling (3b.4 front-region carry):** A BAM record that spans the +/// boundary between run k and run k+1 is handled as follows: +/// - While scanning run k's span, the trailing partial record (a 4-byte +/// `block_size` prefix whose body extends past the run's end) is detected when +/// `cur + 4 <= span.len()` but `cur + 4 + bs > span.len()`. Those partial bytes +/// (`span[cur..]`) are copied into `self.carry`. +/// - When the first block of run k+1 arrives (arena already frozen), the carry is +/// written right-aligned into that arena's front region at +/// `[FRONT_REGION - carry.len(), FRONT_REGION)` via `unsafe { arena.slice_mut(...) }`. +/// The scan then starts at `FRONT_REGION - carry.len()` so that `carry ++ block-0-head` +/// is physically contiguous, forming one complete record ref that is naturally +/// record 0 after sorting. +/// +/// **Parity invariant:** sealed mid-stream runs emit `Spill{seq: 0, 1, ...}`; +/// exactly one final run emits `Residual`. `AllAnnounced` carries +/// `slot_count = spilled_run_count, memory_chunk_count = 1`. +pub struct FindBoundariesAndSort { + /// Per-order key extraction + seal policy; owns the accumulated sort refs. + strategy: S, + /// Threads handed to the per-chunk coordinate sort (`coordinate_chunk_from_refs`); + /// `>1` enables the parallel radix on large chunks. Matches the pipeline's + /// configured thread count. + sort_threads: usize, + output_byte_limit: u64, + /// Held slot for backpressure retry on emitting events. + held: HeldSlot>, + /// The shared arena that holds the current run's decompressed bytes (set on + /// the first block of each run, cleared after the run is scanned). + arena: Option>, + /// Arena offset of the scan start for the current run. + /// + /// - Run 0 with empty carry: `FRONT_REGION + bam_header_len(...)`. + /// - Run k>0 with non-empty carry: `FRONT_REGION - carry.len()` (after + /// writing carry into the front region). + /// - Run k>0 with empty carry (no straddler): `FRONT_REGION`. + scan_start: u64, + /// Arena offset one past the last byte of the current run span (updated with + /// each block). + run_end: u64, + /// Arena offset of the next un-parsed record's first byte within the current + /// run. The incremental scan advances this as in-order blocks extend `run_end`, + /// parking it at the start of the first record that overruns the bytes inflated + /// so far; it resumes from here when the next block arrives. At seal, + /// `[scan_cursor, run_end)` is exactly the trailing partial record (the straddler + /// carry for a spill, or empty for a clean residual). + scan_cursor: u64, + /// `true` once run 0's BAM header has been skipped. Set `true` immediately for + /// runs `> 0` (no header). Stays `false` on run 0 until enough blocks have been + /// inflated for [`bam_header_len`] to parse the full header — a many-reference + /// header can exceed a single 64 KiB block, so the skip may take several blocks. + header_skipped: bool, + /// Pending events staged after each run scan; drained in subsequent `try_run` + /// calls. Holds `Spill`/`Residual` then `AllAnnounced` (only on the final run). + pending: VecDeque, + /// `true` once the final run has been scanned and `AllAnnounced` staged. + finalized: bool, + /// Trailing partial record bytes from the previous run (the carry buffer). + /// Non-empty when run k ended mid-record; written into run k+1's front region. + carry: Vec, + /// Sequence number for the next run's `Spill` event (monotonically increasing). + next_seq: u32, + /// Cumulative record count across all runs. + total_records: u64, + /// Number of runs that sealed to a disk spill (all runs except the final one). + spilled_run_count: u32, +} + +impl FindBoundariesAndSort { + /// Test-only borrow of the per-order strategy, so a test can observe + /// strategy-owned state (e.g. the bounded sort pool built at seal) after + /// driving the step — confirming the `sort_threads` handed to [`new`](Self::new) + /// is forwarded into `strategy.seal`. + #[cfg(test)] + pub(crate) fn strategy(&self) -> &S { + &self.strategy + } + + /// Create a new `FindBoundariesAndSort` step over the given per-order + /// `strategy` (e.g. [`CoordinateStrategy`], which carries the header's + /// reference-sequence count for key extraction). `sort_threads` is the thread + /// count handed to the per-chunk sort (the pipeline's configured threads). + /// `output_byte_limit` byte-bounds the output queue. + #[must_use] + pub fn new(strategy: S, sort_threads: usize, output_byte_limit: u64) -> Self { + Self { + strategy, + sort_threads, + output_byte_limit, + held: HeldSlot::new(), + arena: None, + scan_start: 0, + run_end: 0, + scan_cursor: 0, + header_skipped: false, + pending: VecDeque::new(), + finalized: false, + carry: Vec::new(), + next_seq: 0, + total_records: 0, + spilled_run_count: 0, + } + } + + /// Ingest one `InflatedBlock`, extending the current run's contiguous arena span + /// and then parsing every record made complete by it via + /// [`scan_available`](Self::scan_available) — so the scan overlaps the run's + /// still-inflating tail (Inflate‖Scan). When the block is the last of a run + /// (`is_last_of_run`), the run is sealed: the trailing partial record (if any) + /// becomes the carry, the accumulated refs are sorted, and a `Spill` or + /// `Residual` event is staged in `self.pending`. After the final run's + /// `Residual`, an `AllAnnounced` is also staged. + /// + /// On the first block of each run, the straddler carry from the previous run (if + /// any) is written right-aligned into the new arena's front region via an + /// `unsafe` `slice_mut` call (3rd `fgumi-pipeline-io` unsafe site — see + /// CLAUDE.md §"Approved hot-path unsafe (parallel-inflate sort ingest, + /// fgumi-pipeline-io)" for the full justification). + /// + /// # Errors + /// + /// Returns an `io::Error` if: + /// - The carry to write into the front region is longer than `FRONT_REGION`. + /// - A `block_size` value in the record stream overflows `u32`. + /// - At seal: run 0's header never fully arrived, the carry exceeds `FRONT_REGION`, + /// or the final residual run ends mid-record (truncated BAM). + pub(crate) fn ingest_block(&mut self, block: &InflatedBlock) -> io::Result<()> { + let block_end = block.offset + u64::from(block.len); + + if self.arena.is_none() { + // ---------------------------------------------------------------- + // First block of this run. + // ---------------------------------------------------------------- + + // Straddler carry: write carry right-aligned into the front region. + let scan_start = if !self.carry.is_empty() { + let l = self.carry.len(); + if l > FRONT_REGION { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "FindBoundariesAndSort: carry ({l} bytes) exceeds FRONT_REGION \ + ({FRONT_REGION}); record exceeds max straddler size — \ + raise --max-memory" + ), + )); + } + let write_offset = FRONT_REGION - l; + // SAFETY: `[FRONT_REGION - l, FRONT_REGION)` is the front region reserved + // by `ReadBlocks::ensure_arena` via `grow_uninit(FRONT_REGION)` before any + // block slot was allocated. Soundness rests on DISJOINTNESS, not on + // ordering: this range lies strictly below `FRONT_REGION`, while every + // inflate block slot lies at or above `FRONT_REGION` (block 0 starts there + // by construction in ReadBlocks), so the `&mut [u8]` synthesized here + // aliases no slot an `InflateToArena` worker is concurrently writing — and + // run k+1's later blocks MAY still be inflating when this carry write runs + // (FBS is serial but does not wait for the whole run to inflate). The + // arena's backing segment was frozen with `reserve_full_capacity` before + // being shared, so no `grow_uninit`/realloc can move these bytes underneath + // a live worker slice. This write happens exactly once per straddler per + // run; `u8` has no validity invariant. + #[allow(unsafe_code)] + let dst = unsafe { block.arena.slice_mut(write_offset, l) }; + dst.copy_from_slice(&self.carry); + self.carry.clear(); + #[allow(clippy::cast_possible_truncation)] + let start = write_offset as u64; + start + } else if self.next_seq == 0 { + // Run 0: BAM header starts at block.offset (the first byte of the + // run's data). In the real pipeline block.offset == FRONT_REGION + // (since ReadBlocks reserves the front region first), but in unit + // tests the arena may be laid out without the front region. + block.offset + } else { + // Non-first run, no carry: records start at FRONT_REGION. + #[allow(clippy::cast_possible_truncation)] + let start = FRONT_REGION as u64; + start + }; + + self.arena = Some(Arc::clone(&block.arena)); + self.scan_start = scan_start; + self.scan_cursor = scan_start; + // Runs > 0 have no BAM header (it lives only at the very start of input); + // their first record is the straddler/`FRONT_REGION` data, so nothing to + // skip. Run 0 must skip the header, which `scan_available` does once + // enough blocks have arrived. + self.header_skipped = self.next_seq != 0; + // Pre-size the ref buffer to the run's upper bound (the arena holds at + // most `segment_size` bytes of record data) so the incremental pushes do + // not trigger doubling reallocations of a multi-GB `Vec` mid-run. + let est = (block.arena.len() / 192).max(1024); + self.strategy.reserve_for_run(est); + } else { + debug_assert_eq!( + block.offset, self.run_end, + "FindBoundariesAndSort: non-contiguous arena block (expected offset {}, got {})", + self.run_end, block.offset + ); + } + self.run_end = block_end; + + // Parse every record made complete by the bytes inflated so far, overlapping + // the scan with the still-inflating tail of the run (Inflate‖Scan). + self.scan_available(&block.arena)?; + + if block.is_last_of_run { + self.seal_run(block.seals_to_spill, block.run_seq)?; + } + + Ok(()) + } + + /// Parse every record that is fully present in `[scan_cursor, run_end)`, pushing + /// a [`RecordRef`] (with its coordinate key extracted inline) for each, and park + /// `scan_cursor` at the first record that overruns the bytes inflated so far. + /// + /// Records are gap-free across block slots (the arena is one contiguous segment + /// with prefix-summed offsets), so the cursor walks straight through block + /// boundaries — a record that started in block k and continues into block k+1 is + /// simply parsed once k+1 has extended `run_end`. A trailing partial record at + /// the *run* boundary is not handled here; it is left in `[scan_cursor, run_end)` + /// for `seal_run` to carry (spill) or reject (truncated residual). + /// + /// For run 0, the BAM header is skipped first; if the header is not yet fully + /// inflated (`bam_header_len` returns `None`), the scan returns and retries on the + /// next block. + /// + /// # Errors + /// + /// Returns an `io::Error` on a `block_size` value that overflows `u32`. + fn scan_available(&mut self, arena: &PooledSegmentedBuf) -> io::Result<()> { + let scan_start = self.scan_start; + let scan_start_usize = usize::try_from(scan_start).expect("scan_start must fit in usize"); + let run_end_usize = usize::try_from(self.run_end).expect("run_end must fit in usize"); + let avail_len = run_end_usize.checked_sub(scan_start_usize).expect("run_end >= scan_start"); + + // Borrow the bytes inflated so far for this run directly from the arena — NO + // copy. The whole run lives in one segment, so this single slice spans + // `[scan_start, run_end)` and is re-taken (cheaply) as `run_end` grows. + let span = arena.slice(scan_start_usize, avail_len); + + // Run 0: skip the BAM header once it is fully present. A many-reference + // header can exceed one BGZF block, so `None` means "wait for more blocks", + // NOT an error (the error is raised at seal if the header never completes). + if !self.header_skipped { + // `?` surfaces a wrong-magic stream as an error here; `Ok(None)` still + // means "header not fully inflated yet — wait for more blocks". + match bam_header_len(span)? { + Some(h) => { + self.scan_cursor = scan_start + h as u64; + self.header_skipped = true; + } + None => return Ok(()), + } + } + + // Parse complete `[block_size(4)][body]` frames from the cursor, extracting + // the coordinate key inline (the body is already resident in `span`). + let mut cur = usize::try_from(self.scan_cursor - scan_start) + .expect("cursor offset within run fits in usize"); + let mut run_records: u64 = 0; + loop { + if cur + 4 > span.len() { + // Not even a full 4-byte length prefix present yet — wait. + break; + } + let bs = u32::from_le_bytes([span[cur], span[cur + 1], span[cur + 2], span[cur + 3]]) + as usize; + if cur + 4 + bs > span.len() { + // Record body not fully inflated yet — wait for the next block. + break; + } + // Software-prefetch a few records ahead to hide the cold-arena miss + // latency of this forward scan. + let pf = cur + SCAN_PREFETCH_DISTANCE; + if pf < span.len() { + prefetch_read_l1(&span[pf]); + } + #[allow(clippy::cast_possible_truncation)] + let body_arena_off = scan_start + cur as u64 + 4; + let bs_u32 = u32::try_from(bs).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("FindBoundariesAndSort: block_size {bs} overflows u32"), + ) + })?; + self.strategy.push_record(&span[cur + 4..cur + 4 + bs], body_arena_off, bs_u32)?; + cur += 4 + bs; + run_records += 1; + } + self.scan_cursor = scan_start + cur as u64; + self.total_records += run_records; + Ok(()) + } + + /// Seal the current run: finalize the trailing carry, sort the refs accumulated + /// incrementally by [`scan_available`](Self::scan_available), and stage a `Spill` + /// or `Residual` event. + /// + /// By the time this is called, `scan_available` has already run for the run's last + /// block (in `ingest_block`), so every complete record is in `self.refs` and + /// `scan_cursor` is parked at the start of the trailing partial record (if any). + /// `[scan_cursor, run_end)` is therefore exactly that partial record: + /// - for a mid-stream (spill) seal it becomes the straddler carry for the next run; + /// - for the final (residual) seal it MUST be empty — a non-empty tail means the + /// BAM ended mid-record (truncated/malformed), surfaced as a hard error. + /// + /// # Errors + /// + /// Returns an `io::Error` if run 0's header never fully arrived, on a truncated + /// final record, or if the carry exceeds `FRONT_REGION`. + fn seal_run(&mut self, seals_to_spill: bool, run_seq: u32) -> io::Result<()> { + let arena = self.arena.take().expect("seal_run called with no arena"); + + // Run 0's header must have been skipped by now; if `scan_available` never + // managed it, the input is too short to contain a valid BAM header. + if !self.header_skipped { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "FindBoundariesAndSort: input too short to contain a valid BAM header", + )); + } + + // The trailing partial record, if any, is `[scan_cursor, run_end)`. + let scan_cursor_usize = + usize::try_from(self.scan_cursor).expect("scan_cursor must fit in usize"); + let run_end_usize = usize::try_from(self.run_end).expect("run_end must fit in usize"); + let tail_len = + run_end_usize.checked_sub(scan_cursor_usize).expect("run_end >= scan_cursor"); + + if seals_to_spill { + // Mid-stream seal → the tail is the straddler carry for the next run. + if tail_len > 0 { + let tail = arena.slice(scan_cursor_usize, tail_len); + if tail_len > FRONT_REGION { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "FindBoundariesAndSort: carry ({tail_len} bytes) exceeds FRONT_REGION \ + ({FRONT_REGION}); record exceeds max straddler size — \ + raise --max-memory" + ), + )); + } + self.carry.extend_from_slice(tail); + } + } else if tail_len > 0 { + // Final (residual) run with a leftover tail → truncated/malformed BAM. + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "truncated BAM: final record incomplete at end of input", + )); + } + + // Sort the refs accumulated across this run's blocks and wrap the arena into + // the erased chunk — no record copies (keys were extracted during the scan). + let chunk = self.strategy.seal(Arc::clone(&arena), self.sort_threads); + + // Verify ReadBlocks and FindBoundariesAndSort are in lockstep on run + // sequencing. A desync here indicates a bug in the pipeline wiring (e.g. + // ReadBlocks emitting the wrong seq on an is_last_of_run block). + debug_assert_eq!( + run_seq, self.next_seq, + "ReadBlocks run_seq desynced from FindBoundariesAndSort seq counter" + ); + + let seq = self.next_seq; + self.next_seq += 1; + + if seals_to_spill { + // Mid-stream sealed run → disk spill. + debug_assert!( + !self.finalized, + "FindBoundariesAndSort: Spill emitted after finalization" + ); + self.pending.push_back(SortChunkEvent::Spill { + seq, + chunk, + records_ingested_so_far: self.total_records, + }); + self.spilled_run_count += 1; + } else { + // Final (residual) run. A non-empty carry here means the BAM ends + // mid-record — the file is truncated or malformed. We must surface + // this as a hard error in all build profiles: a debug_assert! would + // silently drop the partial record in release builds. + if !self.carry.is_empty() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "truncated BAM: final record incomplete at end of input", + )); + } + self.pending.push_back(SortChunkEvent::Residual { + chunk, + records_ingested_so_far: self.total_records, + }); + self.pending.push_back(SortChunkEvent::AllAnnounced { + slot_count: self.spilled_run_count, + memory_chunk_count: 1, + total_records: self.total_records, + }); + self.finalized = true; + } + + // Reset per-run span tracking. `self.refs` was emptied by `mem::take`; the + // next run's first block re-initializes `scan_cursor`/`header_skipped`. + self.scan_start = 0; + self.run_end = 0; + self.scan_cursor = 0; + self.header_skipped = false; + + Ok(()) + } + + /// Called when input is fully drained and no `is_last_of_run` block was received + /// (empty pipeline — no arena was ingested at all). Stages a well-formed + /// sentinel so downstream steps can complete. + /// + /// The "all runs sealed to spills, no residual" branch below is DEFENSIVE: the + /// `ReadBlocks` deferred-seal design guarantees the final run is always a + /// residual (it seals a spill only once a following block proves a next run + /// exists, and downgrades a pending seal to a residual at EOF), so for valid + /// input that branch is unreachable. It is retained as a hard error on a + /// non-empty carry so a genuinely truncated BAM (or a future regression in the + /// seal logic) surfaces loudly instead of silently dropping a record. + /// + /// Returns `Some(first_event)` popped from `self.pending`, or `None` if no + /// arena was ingested. + pub(crate) fn finalize(&mut self) -> io::Result> { + if self.finalized { + // Already handled by the last is_last_of_run block — nothing to do. + return Ok(self.pending.pop_front()); + } + if self.arena.is_none() && self.next_seq == 0 { + // No arena was ever ingested (empty input). + return Ok(None); + } + if self.arena.is_none() && self.next_seq > 0 { + // DEFENSIVE / unreachable for valid input: the deferred-seal design + // always makes the final run a residual, so we should never finalize + // with spills-only-and-no-residual. If we somehow do, a non-empty carry + // means the BAM ends mid-record (truncated input or a seal-logic + // regression) — surface it as a hard error rather than drop a record. + if !self.carry.is_empty() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "truncated BAM: final record incomplete at end of input (no-residual path)", + )); + } + // Emit AllAnnounced with memory_chunk_count=0 so SortMerge completes. + self.pending.push_back(SortChunkEvent::AllAnnounced { + slot_count: self.spilled_run_count, + memory_chunk_count: 0, + total_records: self.total_records, + }); + self.finalized = true; + return Ok(self.pending.pop_front()); + } + // Edge case: arena present but no is_last_of_run seen — treat as residual. + if self.arena.is_some() { + self.seal_run(false, self.next_seq)?; + } + Ok(self.pending.pop_front()) + } + + fn flush_held(&mut self, ctx: &mut StepCtx<'_, Self>) -> bool { + // `true` once the slot is clear (was empty, or the held event flushed); + // `false` while it's still held under backpressure. Uses the canonical + // re-hold helper so the put-back-on-reject invariant lives in one place. + !matches!(ctx.outputs.retry_held(&mut self.held), HeldRetry::StillHeld) + } + + fn emit_pending(&mut self, ctx: &mut StepCtx<'_, Self>) -> StepOutcome { + let Some(event) = self.pending.pop_front() else { + return StepOutcome::NoProgress; + }; + match ctx.outputs.push(event) { + Ok(()) => StepOutcome::Progress, + Err(unpushed) => { + self.held.put(unpushed); + StepOutcome::Progress + } + } + } +} + +impl Step for FindBoundariesAndSort { + type Input = InflatedBlock; + type Outputs = Single; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "FindBoundariesAndSort", + // Off-pool on the coordination driver (N+2): the serial boundary scan + // + radix sort + seal/spill framing runs on the dedicated coordination + // thread, keeping the pool on pure parallel inflate/compress. + kind: StepKind::Detached, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], + branch_ordering: vec![BranchOrdering::None], + } + } + + fn detached_group(&self) -> DetachedGroup { + DetachedGroup::Shared(crate::sort::SORT_COORD_GROUP) + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + // Retry any held output first (backpressure path). + if !self.flush_held(ctx) { + return Ok(StepOutcome::Contention); + } + + // Drain any staged pending events before processing new input. + if !self.pending.is_empty() { + return Ok(self.emit_pending(ctx)); + } + + // If already finalized, drain pending or finish. + if self.finalized { + return Ok(StepOutcome::Finished); + } + + // Try to pop and ingest one InflatedBlock. + if let Some(block) = ctx.input.pop() { + self.ingest_block(&block)?; + // If ingest_block sealed a run it may have staged events; drain them. + if !self.pending.is_empty() { + return Ok(self.emit_pending(ctx)); + } + return Ok(StepOutcome::Progress); + } + + // No input available. + if !ctx.input.is_drained() { + return Ok(StepOutcome::NoProgress); + } + + // Input fully drained: finalize (handles empty-input edge case). + if let Some(first_event) = self.finalize()? { + match ctx.outputs.push(first_event) { + Ok(()) => { + // The push landed, so `held` is empty: it is safe to drain the + // next staged event this call (e.g. `AllAnnounced` after + // `Residual`). `emit_pending` will itself hold on a full queue. + if !self.pending.is_empty() { + return Ok(self.emit_pending(ctx)); + } + } + Err(unpushed) => { + // The output queue is full: hold `first_event` and retry it on + // the next `try_run` via `flush_held`. Do NOT drain `pending` + // here — pushing a later event now would place it ahead of the + // still-held one, and a subsequent hold would clobber the held + // event (silent loss). Keep the one-event-per-call invariant + // exactly like the normal path above. + self.held.put(unpushed); + } + } + return Ok(StepOutcome::Progress); + } + + // No arena was ingested (empty input). + Ok(StepOutcome::Finished) + } + + fn new_worker_copy(&self) -> Self { + Self::new(self.strategy.fresh(), self.sort_threads, self.output_byte_limit) + } +} + +// ============================================================================ +// Tests +// ============================================================================ + +#[cfg(test)] +mod tests { + use super::*; + use fgumi_sort::{CoordinateChunkSorter, PooledSegmentedBuf, SegmentedBuf}; + use std::sync::Arc; + + // ----------------------------------------------------------------------- + // Test helpers shared across ReadBlocks and InflateToArena tests. + // ----------------------------------------------------------------------- + + /// Build a single BGZF block containing `payload` using `InlineBgzfCompressor` + /// at compression level 0 (stored / uncompressed). + /// Returns the raw block bytes (header + deflate payload + footer). + fn make_test_bgzf_block(payload: &[u8]) -> Vec { + let mut compressor = fgumi_bgzf::writer::InlineBgzfCompressor::new(0); + compressor.write_all(payload).expect("compress payload"); + compressor.flush().expect("flush compressor"); + let mut blocks = compressor.take_blocks(); + assert_eq!(blocks.len(), 1, "payload must fit in one BGZF block"); + blocks.remove(0).data + } + + /// Acquire a pooled arena for constructing block tokens in tests. + fn test_arena() -> Arc { + let pool = ArenaPool::new(1, 1024); + Arc::new(pool.try_acquire().expect("a fresh pool always has one arena")) + } + + // ----------------------------------------------------------------------- + // Block token accounting (HeapSize / Ordered) + // ----------------------------------------------------------------------- + + #[test] + fn arena_block_charges_only_its_compressed_bytes_to_the_heap() { + // The arena is shared via `Arc`, so it must NOT be counted per block — + // byte-bounded queues would otherwise over-charge by the arena size and + // stall the pipeline. + let block = ArenaBlock { + arena: test_arena(), + ordinal: 7, + offset: 64, + len: 128, + block: vec![0xAB; 40], + is_last_of_run: true, + run_seq: 2, + seals_to_spill: false, + }; + assert_eq!(block.heap_size(), 40, "only the owned BGZF bytes are charged"); + assert_eq!(block.ordinal(), 7, "ordinal drives ByItemOrdinal reordering"); + } + + #[test] + fn inflated_block_is_heap_free_because_its_bytes_live_in_the_arena() { + let token = InflatedBlock { + arena: test_arena(), + ordinal: 11, + offset: 0, + len: 256, + is_last_of_run: false, + run_seq: 1, + seals_to_spill: true, + }; + assert_eq!(token.heap_size(), 0, "the token owns no bytes of its own"); + assert_eq!(token.ordinal(), 11); + } + + // ----------------------------------------------------------------------- + // Step wiring + // ----------------------------------------------------------------------- + + #[test] + fn read_blocks_runs_off_pool_on_the_coordination_driver() { + let step = ReadBlocks::new(1 << 20, 4096); + let profile = step.profile(); + assert_eq!(profile.name, "ReadBlocks"); + // Detached keeps the serial arena-admit off a pool worker so it cannot + // steal a slot from the parallel inflaters. + assert_eq!(profile.kind, StepKind::Detached); + assert!(!profile.sticky); + assert_eq!(profile.branch_ordering, vec![BranchOrdering::ByItemOrdinal]); + match profile.output_queues.as_slice() { + [QueueSpec::ByteBounded { limit_bytes }] => assert_eq!(*limit_bytes, 4096), + other => panic!("expected one byte-bounded queue, got {other:?}"), + } + assert_eq!(step.detached_group(), DetachedGroup::Shared(crate::sort::SORT_COORD_GROUP)); + } + + #[test] + fn inflate_to_arena_is_parallel_and_worker_copies_share_only_config() { + let step = InflateToArena::new(8192); + let profile = step.profile(); + assert_eq!(profile.name, "InflateToArena"); + assert_eq!(profile.kind, StepKind::Parallel, "inflation fans out across workers"); + assert!(!profile.sticky); + assert_eq!(profile.branch_ordering, vec![BranchOrdering::ByItemOrdinal]); + + // Each worker gets an independent copy carrying the same queue bound. + let worker = step.new_worker_copy(); + match worker.profile().output_queues.as_slice() { + [QueueSpec::ByteBounded { limit_bytes }] => assert_eq!(*limit_bytes, 8192), + other => panic!("expected one byte-bounded queue, got {other:?}"), + } + } + + // ----------------------------------------------------------------------- + // Sort strategies + // ----------------------------------------------------------------------- + + #[test] + fn coordinate_strategy_accumulates_refs_and_seal_drains_them() { + let mut strategy = CoordinateStrategy::new(4); + strategy.reserve_for_run(8); + + // Two synthetic record bodies; only the leading tid/pos fields are read + // by the inline key extractor, so a minimal fixed-size body suffices. + let body = vec![0u8; 36]; + strategy.push_record(&body, 0, 36).unwrap(); + strategy.push_record(&body, 36, 36).unwrap(); + assert_eq!(strategy.refs.len(), 2, "each pushed record yields one arena ref"); + + let chunk = strategy.seal(test_arena(), 1); + assert!(matches!(chunk, MemoryChunkErased::Coordinate(_)), "coordinate seal"); + assert!(strategy.refs.is_empty(), "seal must drain the accumulator for the next run"); + } + + #[test] + fn coordinate_strategy_fresh_keeps_config_and_drops_accumulated_state() { + let mut strategy = CoordinateStrategy::new(9); + let body = vec![0u8; 36]; + strategy.push_record(&body, 0, 36).unwrap(); + assert!(!strategy.refs.is_empty()); + + let next = strategy.fresh(); + assert_eq!(next.n_ref, 9, "reference count is configuration and is carried over"); + assert!(next.refs.is_empty(), "a fresh strategy starts a new run with no refs"); + } + + /// Read the BGZF footer ISIZE field (last 4 bytes of the block), which is + /// the uncompressed size mod 2^32. Returns the value as `usize`. + fn uncompressed_size_of(block: &[u8]) -> usize { + assert!(block.len() >= 4, "block too short to contain BGZF footer"); + let n = block.len(); + u32::from_le_bytes([block[n - 4], block[n - 3], block[n - 2], block[n - 1]]) as usize + } + + /// Build a single BGZF block containing `payload` using `InlineBgzfCompressor` + /// at level 6, returning raw block bytes. Used by `InflateToArena` tests. + fn compress_one_block(payload: &[u8]) -> Vec { + let mut compressor = fgumi_bgzf::writer::InlineBgzfCompressor::new(6); + compressor.write_all(payload).expect("compress payload"); + compressor.flush().expect("flush compressor"); + let mut blocks = compressor.take_blocks(); + assert_eq!(blocks.len(), 1, "payload must fit in one BGZF block"); + blocks.remove(0).data + } + + // ----------------------------------------------------------------------- + // ReadBlocks unit tests + // ----------------------------------------------------------------------- + + /// Two synthetic BGZF blocks are admitted; after freeze the emitted + /// `ArenaBlock`s must carry contiguous offsets starting at `FRONT_REGION`, + /// correct `len`/`ordinal` values, and share the same arena `Arc`. + #[allow(unsafe_code)] + #[test] + fn read_blocks_admits_grows_and_emits_arena_blocks() { + let blk0 = make_test_bgzf_block(b"first-block-decompressed-payload"); + let blk1 = make_test_bgzf_block(b"second-block-payload-2"); + let isz0 = uncompressed_size_of(&blk0); + let isz1 = uncompressed_size_of(&blk1); + + // memory_limit larger than both blocks combined → single run + let mut step = ReadBlocks::new(64 * 1024 * 1024, 64 * 1024 * 1024); + assert!( + step.admit_block(BgzfBlock { + batch_serial: 0, + bytes: blk0.clone(), + uncompressed_size: u32::try_from(isz0).unwrap(), + }) + .is_ok() + ); + assert!( + step.admit_block(BgzfBlock { + batch_serial: 1, + bytes: blk1.clone(), + uncompressed_size: u32::try_from(isz1).unwrap(), + }) + .is_ok() + ); + let emitted: Vec = step.seal_and_drain_for_test(); + + assert_eq!(emitted.len(), 2); + // Block 0 lands at FRONT_REGION (after the reserved front region). + assert_eq!(emitted[0].offset, FRONT_REGION as u64, "block 0 must start at FRONT_REGION"); + // Block 1 is contiguous after block 0. + assert_eq!( + usize::try_from(emitted[1].offset).unwrap(), + FRONT_REGION + isz0, + "offsets must be contiguous after FRONT_REGION" + ); + assert!(emitted[1].is_last_of_run, "last block must carry is_last_of_run=true"); + assert!(!emitted[0].is_last_of_run, "first block must not carry is_last_of_run"); + assert_eq!(emitted[0].len as usize, isz0); + assert_eq!(emitted[1].len as usize, isz1); + assert_eq!(emitted[0].ordinal, 0); + assert_eq!(emitted[1].ordinal, 1); + assert_eq!(emitted[0].run_seq, 0); + assert_eq!(emitted[1].run_seq, 0); + assert!(!emitted[0].seals_to_spill, "residual run must not seal to spill"); + assert!(!emitted[1].seals_to_spill, "residual run must not seal to spill"); + // All blocks share the same arena Arc. + assert!(Arc::ptr_eq(&emitted[0].arena, &emitted[1].arena)); + } + + /// Feed enough blocks to force TWO runs. + /// + /// Asserts: + /// - Run 0 blocks: `run_seq == 0`, offsets start at `FRONT_REGION`, + /// last block has `is_last_of_run`, all have `seals_to_spill = true`. + /// - Run 1 blocks: `run_seq == 1`, offsets restart at `FRONT_REGION` + /// (fresh/reused arena from the pool), `seals_to_spill = false`. + /// - Ordinals are monotonically increasing across both runs. + /// - Run 0 and run 1 arena `Arc`s differ (or are the same reused physical + /// buffer returned via the pool after run 0's Arc drops). + #[allow(unsafe_code)] + #[allow(clippy::too_many_lines)] // exhaustive two-run state-machine assertions + #[test] + fn read_blocks_two_run_seal() { + // Make two blocks; each has isz bytes of uncompressed data. + let payload0 = b"run0-block0-payload".as_ref(); + let payload1 = b"run0-block1-payload".as_ref(); + let payload2 = b"run1-block0-payload".as_ref(); + let payload3 = b"run1-block1-payload".as_ref(); + + let blk0 = make_test_bgzf_block(payload0); + let blk1 = make_test_bgzf_block(payload1); + let blk2 = make_test_bgzf_block(payload2); + let blk3 = make_test_bgzf_block(payload3); + + let isz0 = uncompressed_size_of(&blk0); + let isz1 = uncompressed_size_of(&blk1); + let isz2 = uncompressed_size_of(&blk2); + let isz3 = uncompressed_size_of(&blk3); + + // memory_limit set to force a seal after run 0 (blk0+blk1): + // use isz0+isz1 as the cap so that after admitting blk1, run_cumulative + // >= run_cap and a seal fires. The arena segment is derived from this + // budget (FRONT_REGION + run_cap + one block), so all tiny test blocks + // still fit in one segment per run. + let run_cap = isz0 + isz1; + let mut step = ReadBlocks::new(run_cap, 64 * 1024 * 1024); + + // Admit run 0 blocks. + assert!( + step.admit_block(BgzfBlock { + batch_serial: 0, + bytes: blk0.clone(), + uncompressed_size: u32::try_from(isz0).unwrap(), + }) + .is_ok() + ); + assert!( + step.admit_block(BgzfBlock { + batch_serial: 1, + bytes: blk1.clone(), + uncompressed_size: u32::try_from(isz1).unwrap(), + }) + .is_ok() + ); + + // After run 0's blocks fill run_cap, seal run 0 explicitly and collect its + // blocks (simulating the seal that happens when run_cumulative >= run_cap). + // In the pipeline try_run, the seal fires after admit; here we do it manually + // via a dedicated seal so we can inspect both runs independently. + step.seal_run(true); // mid-stream → seals_to_spill = true + + // Collect run 0 blocks. + let run0_blocks: Vec = step.emit.drain(..).collect(); + assert_eq!(run0_blocks.len(), 2, "run 0 must emit 2 blocks"); + + // Run 0 invariants. Eager emission stamps `is_last_of_run` / + // `seals_to_spill` only on the run's LAST block (the withheld one); non-last + // blocks carry `false` (downstream reads these fields only on the last block). + for blk in &run0_blocks { + assert_eq!(blk.run_seq, 0, "run 0 blocks must have run_seq=0"); + } + assert!(!run0_blocks[0].is_last_of_run, "run 0 block 0 must not be last_of_run"); + assert!(run0_blocks[1].is_last_of_run, "run 0 block 1 must be last_of_run"); + assert!(!run0_blocks[0].seals_to_spill, "non-last block carries seals_to_spill=false"); + assert!(run0_blocks[1].seals_to_spill, "run 0 last block must seal to spill"); + assert_eq!( + run0_blocks[0].offset, FRONT_REGION as u64, + "run 0 block 0 must start at FRONT_REGION" + ); + assert_eq!( + usize::try_from(run0_blocks[1].offset).unwrap(), + FRONT_REGION + isz0, + "run 0 block 1 must be contiguous after block 0" + ); + assert_eq!(run0_blocks[0].ordinal, 0, "global ordinal monotonic"); + assert_eq!(run0_blocks[1].ordinal, 1, "global ordinal monotonic"); + + // Verify pool exhaustion: while run 0's Arc is still alive, admitting a new + // block must fail (pool capacity 1 → all arenas in flight). + { + let probe_blk = make_test_bgzf_block(b"pool-exhaustion-probe"); + let probe_isz = uncompressed_size_of(&probe_blk); + let result = step.admit_block(BgzfBlock { + batch_serial: 99, + bytes: probe_blk, + uncompressed_size: u32::try_from(probe_isz).unwrap(), + }); + assert!(result.is_err(), "pool must be exhausted while run 0 Arc is still in flight"); + // A failed admit (ensure_arena returned false) must not leak state: the + // block is handed back in the Err, no block is withheld, arena stays None. + assert!(step.last_block.is_none(), "failed admit must not withhold a block"); + assert!(step.arena.is_none(), "failed admit must leave arena as None"); + } + + // Drop run 0 blocks → their Arc refcount falls to 0 → + // PooledSegmentedBuf::drop releases the arena back to the pool. + drop(run0_blocks); + + // Now admit run 1 blocks — the pool has one free (reused) arena. + assert!( + step.admit_block(BgzfBlock { + batch_serial: 2, + bytes: blk2.clone(), + uncompressed_size: u32::try_from(isz2).unwrap(), + }) + .is_ok(), + "run 1 admit must succeed after pool release" + ); + assert!( + step.admit_block(BgzfBlock { + batch_serial: 3, + bytes: blk3.clone(), + uncompressed_size: u32::try_from(isz3).unwrap(), + }) + .is_ok(), + "run 1 second admit must succeed" + ); + + // Seal run 1 as the residual. + step.seal_run(false); + let run1_blocks: Vec = step.emit.drain(..).collect(); + assert_eq!(run1_blocks.len(), 2, "run 1 must emit 2 blocks"); + + // Run 1 invariants. + for blk in &run1_blocks { + assert_eq!(blk.run_seq, 1, "run 1 blocks must have run_seq=1"); + assert!(!blk.seals_to_spill, "run 1 (residual) blocks must have seals_to_spill=false"); + } + assert!(!run1_blocks[0].is_last_of_run, "run 1 block 0 must not be last_of_run"); + assert!(run1_blocks[1].is_last_of_run, "run 1 block 1 must be last_of_run"); + // Run 1 offsets restart at FRONT_REGION (fresh/reused arena). + assert_eq!( + run1_blocks[0].offset, FRONT_REGION as u64, + "run 1 block 0 must restart at FRONT_REGION" + ); + assert_eq!( + usize::try_from(run1_blocks[1].offset).unwrap(), + FRONT_REGION + isz2, + "run 1 block 1 must be contiguous after block 0" + ); + // Ordinals are globally monotonic across runs. + assert_eq!(run1_blocks[0].ordinal, 2, "global ordinal monotonic across runs"); + assert_eq!(run1_blocks[1].ordinal, 3, "global ordinal monotonic across runs"); + + // All run 1 blocks share the same arena Arc. + assert!(Arc::ptr_eq(&run1_blocks[0].arena, &run1_blocks[1].arena)); + + // Pool behavioral check: the run 1 arena was acquired after the pool round- + // tripped run 0's arena back. Verify the underlying storage was reused + // (not a fresh allocation) by comparing the allocated_capacity of the run 1 + // arena: the reused arena retains its segment allocations (reset_for_reuse + // keeps the Vec capacity), so its allocated_capacity is >= the derived + // segment size (FRONT_REGION + run_cap + one block, pre-reserved by + // reserve_full_capacity in ensure_arena). + let expected_segment_size = FRONT_REGION + run_cap + MAX_BGZF_BLOCK; + let run1_alloc_cap = run1_blocks[0].arena.allocated_capacity(); + assert!( + run1_alloc_cap >= expected_segment_size, + "reused arena must retain its segment capacity (got {run1_alloc_cap}, expected >= {expected_segment_size})" + ); + } + + /// Build a minimal valid BAM record body (no `block_size` prefix) with the given + /// ref\_id/pos and a one-character read name `name`. Returns the body bytes. + /// + /// Copied verbatim from `crates/fgumi-sort/src/ref_sort.rs`'s test module + /// (the `(i32, i32, u8)` version) so tests in this crate can construct + /// identical records without a cross-crate dependency on a `#[cfg(test)]` + /// helper. + fn coord_body(ref_id: i32, pos: i32, name: u8) -> Vec { + let mut b = Vec::new(); + b.extend_from_slice(&ref_id.to_le_bytes()); + b.extend_from_slice(&pos.to_le_bytes()); + b.push(2); // l_read_name (incl NUL): "\0" + b.push(0); // mapq + b.extend_from_slice(&0u16.to_le_bytes()); // bin + b.extend_from_slice(&0u16.to_le_bytes()); // n_cigar_op + b.extend_from_slice(&0u16.to_le_bytes()); // flag (forward) + b.extend_from_slice(&0u32.to_le_bytes()); // l_seq + b.extend_from_slice(&(-1i32).to_le_bytes()); // next_refID + b.extend_from_slice(&(-1i32).to_le_bytes()); // next_pos + b.extend_from_slice(&0i32.to_le_bytes()); // tlen + b.push(name); // read_name char + b.push(0); // read_name NUL terminator + b + } + + /// Build a minimal BAM binary header parseable by `bam_header_len`. + /// + /// Layout: `magic(4) + l_text=0(4) + n_ref(4)` followed by `n_ref` entries of + /// `l_name=2(4) + "r\0"(2) + l_ref=100(4)`. + fn minimal_bam_header(n_ref: u32) -> Vec { + let mut h = Vec::new(); + h.extend_from_slice(b"BAM\x01"); // magic + h.extend_from_slice(&0u32.to_le_bytes()); // l_text = 0 + h.extend_from_slice(&n_ref.to_le_bytes()); // n_ref + for _ in 0..n_ref { + h.extend_from_slice(&2u32.to_le_bytes()); // l_name = 2 + h.push(b'r'); + h.push(0); // name "r\0" + h.extend_from_slice(&100u32.to_le_bytes()); // l_ref = 100 + } + h + } + + /// Oracle test: `FindBoundariesAndSort` emitted chunk must be byte-identical to + /// the copy-based `CoordinateChunkSorter` over the same record bodies. + #[allow(unsafe_code)] + #[test] + fn find_boundaries_and_sort_single_run_matches_oracle() { + let n_ref = 4u32; + // Records deliberately out of coordinate order; distinct names witness stability. + let recs = vec![ + coord_body(2, 100, b'a'), + coord_body(0, 50, b'b'), + coord_body(2, 10, b'c'), + coord_body(0, 50, b'd'), // coordinate tie with b'b' → stable sort keeps b before d + coord_body(1, 999, b'e'), + ]; + + // Build an arena: [BAM header][[block_size][body]...] as one contiguous run. + let header = minimal_bam_header(n_ref); + let mut arena = SegmentedBuf::with_capacity(0, 1 << 20); + arena.reserve_full_capacity(); + + // SAFETY: every slot fully written before any read. + let h_off = unsafe { arena.grow_uninit(header.len()) }; + unsafe { arena.slice_mut(h_off, header.len()) }.copy_from_slice(&header); + // h_off is usize; on 64-bit targets usize fits in u64. + #[allow(clippy::cast_possible_truncation)] + let run_start = h_off as u64; + + for r in &recs { + let bs = u32::try_from(r.len()).unwrap(); + let po = unsafe { arena.grow_uninit(4) }; + unsafe { arena.slice_mut(po, 4) }.copy_from_slice(&bs.to_le_bytes()); + let bo = unsafe { arena.grow_uninit(r.len()) }; + unsafe { arena.slice_mut(bo, r.len()) }.copy_from_slice(r); + } + // arena.len() is usize; usize fits in u64 on 64-bit targets. + #[allow(clippy::cast_possible_truncation)] + let run_len = arena.len() as u64 - run_start; + let arena = Arc::new(PooledSegmentedBuf::unpooled(arena)); + + // Drive ingest_block + finalize directly (the test seam). + let mut step = + FindBoundariesAndSort::new(CoordinateStrategy::new(n_ref), 1, 64 * 1024 * 1024); + step.ingest_block(&InflatedBlock { + arena: Arc::clone(&arena), + ordinal: 0, + offset: run_start, + len: u32::try_from(run_len).unwrap(), + is_last_of_run: true, + run_seq: 0, + seals_to_spill: false, + }) + .expect("ingest_block must succeed"); + let ev = step.finalize().expect("finalize must succeed").expect("residual event"); + + let chunk_bytes: Vec> = match ev { + SortChunkEvent::Residual { chunk: MemoryChunkErased::Coordinate(c), .. } => { + (0..c.len()).map(|i| c.record_bytes(i).to_vec()).collect() + } + _ => panic!("expected Residual coordinate chunk, got something else"), + }; + + // Oracle: copy-based in-memory coordinate sort of the same record bodies. + let mut oracle = CoordinateChunkSorter::for_test(usize::MAX, n_ref); + for r in &recs { + let _ = oracle.push(r).unwrap(); + } + let oc = oracle.take_sorted_chunk(); + let oracle_bytes: Vec> = + (0..oc.len()).map(|i| oc.record_bytes(i).to_vec()).collect(); + + assert_eq!( + chunk_bytes, oracle_bytes, + "arena-scan single-run sort must be byte-identical to the copy-based sorter" + ); + } + + /// Two-run straddler test: a record that spans the boundary between run 0 and + /// run 1 must appear as record 0 of run 1's chunk with bytes equal to + /// `carry ++ head` (the full original record), and the merged union of both + /// runs (stable, lower-run first) must equal the oracle (`CoordinateChunkSorter` + /// over the same records in input order). + /// + /// Test layout: + /// - Run 0 arena (no `FRONT_REGION` prefix — built manually like the oracle test): + /// `[header][rec0_prefix+body][rec1_prefix+body][straddler_prefix+partial_body]` + /// The straddler record's `block_size` prefix PLUS the first few body bytes land + /// in run 0; the rest of the body starts at `FRONT_REGION` in run 1's arena. + /// - Run 1 arena: `FRONT_REGION` uninit bytes reserved, then `straddler_tail + + /// rec2_prefix + body`. + /// `FindBoundariesAndSort` writes the carry into `[FRONT_REGION - L, FRONT_REGION)`, + /// so `carry ++ straddler_tail` is contiguous and forms the complete record. + #[allow(unsafe_code)] + #[allow(clippy::too_many_lines)] + #[test] + fn find_boundaries_and_sort_two_run_straddler() { + let n_ref = 2u32; + let seg_size = 256 * 1024 * 1024usize; + + // Records: rec0 + rec1 go into run 0 (complete); straddler goes across + // the boundary; rec2 goes into run 1 (complete). + let rec0 = coord_body(0, 10, b'a'); + let rec1 = coord_body(1, 20, b'b'); + let straddler = coord_body(0, 5, b'c'); // will be straddled across runs + let rec2 = coord_body(1, 5, b'd'); + + let all_recs = vec![&rec0, &rec1, &straddler, &rec2]; + + // ----------------------------------------------------------------------- + // Build run 0 arena: [header][rec0][rec1][straddler_prefix+partial_body] + // We use the same manual layout as the oracle test (no FRONT_REGION prefix). + // ----------------------------------------------------------------------- + let header = minimal_bam_header(n_ref); + + let mut arena0 = SegmentedBuf::with_capacity(0, seg_size); + arena0.reserve_full_capacity(); + + // Write header. + let h_off = unsafe { arena0.grow_uninit(header.len()) }; + unsafe { arena0.slice_mut(h_off, header.len()) }.copy_from_slice(&header); + + // Write rec0 (prefix + body) — complete. + let bs0 = u32::try_from(rec0.len()).unwrap(); + let po = unsafe { arena0.grow_uninit(4) }; + unsafe { arena0.slice_mut(po, 4) }.copy_from_slice(&bs0.to_le_bytes()); + let bo = unsafe { arena0.grow_uninit(rec0.len()) }; + unsafe { arena0.slice_mut(bo, rec0.len()) }.copy_from_slice(&rec0); + + // Write rec1 (prefix + body) — complete. + let bs1 = u32::try_from(rec1.len()).unwrap(); + let po = unsafe { arena0.grow_uninit(4) }; + unsafe { arena0.slice_mut(po, 4) }.copy_from_slice(&bs1.to_le_bytes()); + let bo = unsafe { arena0.grow_uninit(rec1.len()) }; + unsafe { arena0.slice_mut(bo, rec1.len()) }.copy_from_slice(&rec1); + + // Write straddler prefix (4-byte block_size) + partial body. + // Split: put the 4-byte prefix + first 3 bytes of body in run 0. + let bs_str = u32::try_from(straddler.len()).unwrap(); + let partial_len = 3usize; // bytes of straddler body in run 0 + assert!(partial_len < straddler.len(), "partial_len must be < straddler body size"); + let po = unsafe { arena0.grow_uninit(4) }; + unsafe { arena0.slice_mut(po, 4) }.copy_from_slice(&bs_str.to_le_bytes()); + let bp = unsafe { arena0.grow_uninit(partial_len) }; + unsafe { arena0.slice_mut(bp, partial_len) }.copy_from_slice(&straddler[..partial_len]); + + // run 0 span: from header offset (0) to end of arena. + let run0_start = h_off as u64; + let run0_end = arena0.len() as u64; + let run0_len = usize::try_from(run0_end - run0_start).unwrap(); + + let arena0 = Arc::new(PooledSegmentedBuf::unpooled(arena0)); + + // ----------------------------------------------------------------------- + // Build run 1 arena: [FRONT_REGION uninit][straddler_tail][rec2 prefix+body] + // ----------------------------------------------------------------------- + let mut arena1 = SegmentedBuf::with_capacity(0, seg_size); + arena1.reserve_full_capacity(); + + // Reserve the FRONT_REGION prefix (uninit — the carry will be written here + // by FindBoundariesAndSort). + let _front = unsafe { arena1.grow_uninit(FRONT_REGION) }; + assert_eq!(arena1.len(), FRONT_REGION); + + // Write straddler tail (remaining body bytes) at FRONT_REGION. + let tail_len = straddler.len() - partial_len; + let st_off = unsafe { arena1.grow_uninit(tail_len) }; + assert_eq!(st_off, FRONT_REGION, "straddler tail must start at FRONT_REGION"); + unsafe { arena1.slice_mut(st_off, tail_len) }.copy_from_slice(&straddler[partial_len..]); + + // Write rec2 (prefix + body) — complete. + let bs2 = u32::try_from(rec2.len()).unwrap(); + let po = unsafe { arena1.grow_uninit(4) }; + unsafe { arena1.slice_mut(po, 4) }.copy_from_slice(&bs2.to_le_bytes()); + let bo = unsafe { arena1.grow_uninit(rec2.len()) }; + unsafe { arena1.slice_mut(bo, rec2.len()) }.copy_from_slice(&rec2); + + // run 1 spans from FRONT_REGION (the inflate data, not the front region). + let run1_data_start = FRONT_REGION as u64; + let run1_data_end = arena1.len() as u64; + let run1_data_len = usize::try_from(run1_data_end - run1_data_start).unwrap(); + + let arena1 = Arc::new(PooledSegmentedBuf::unpooled(arena1)); + + // ----------------------------------------------------------------------- + // Drive FindBoundariesAndSort across both runs. + // ----------------------------------------------------------------------- + let mut fbs = + FindBoundariesAndSort::new(CoordinateStrategy::new(n_ref), 1, 64 * 1024 * 1024); + + // Run 0: one block, is_last_of_run = true, seals_to_spill = true. + fbs.ingest_block(&InflatedBlock { + arena: Arc::clone(&arena0), + ordinal: 0, + offset: run0_start, + len: u32::try_from(run0_len).unwrap(), + is_last_of_run: true, + run_seq: 0, + seals_to_spill: true, + }) + .expect("run 0 ingest must succeed"); + + // After run 0's last block, a Spill event must be staged. + assert_eq!(fbs.pending.len(), 1, "run 0 must stage exactly one Spill event"); + let ev0 = fbs.pending.pop_front().unwrap(); + let run0_chunk = match ev0 { + SortChunkEvent::Spill { seq, chunk: MemoryChunkErased::Coordinate(c), .. } => { + assert_eq!(seq, 0, "first Spill must have seq=0"); + c + } + _ => panic!("expected Spill(Coordinate) for run 0, got something else"), + }; + + // Assert run 0's carry: 4 bytes (prefix) + partial_len body bytes. + let expected_carry_len = 4 + partial_len; + assert_eq!( + fbs.carry.len(), + expected_carry_len, + "carry after run 0 must be {expected_carry_len} bytes (prefix + partial body)" + ); + + // Run 0 chunk must NOT contain the straddler — only rec0 and rec1. + assert_eq!( + run0_chunk.len(), + 2, + "run 0 chunk must have exactly 2 complete records (rec0, rec1)" + ); + + // Run 1: one block, is_last_of_run = true, seals_to_spill = false. + fbs.ingest_block(&InflatedBlock { + arena: Arc::clone(&arena1), + ordinal: 1, + offset: run1_data_start, + len: u32::try_from(run1_data_len).unwrap(), + is_last_of_run: true, + run_seq: 1, + seals_to_spill: false, + }) + .expect("run 1 ingest must succeed"); + + // After run 1's last block, pending has Residual + AllAnnounced. + assert_eq!(fbs.pending.len(), 2, "run 1 must stage Residual + AllAnnounced"); + assert!(fbs.carry.is_empty(), "carry must be empty after the final run"); + + let ev_residual = fbs.pending.pop_front().unwrap(); + let SortChunkEvent::Residual { chunk: MemoryChunkErased::Coordinate(run1_chunk), .. } = + ev_residual + else { + panic!("expected Residual(Coordinate) for run 1, got something else") + }; + + let ev_announced = fbs.pending.pop_front().unwrap(); + match ev_announced { + SortChunkEvent::AllAnnounced { slot_count, memory_chunk_count, total_records } => { + assert_eq!(slot_count, 1, "AllAnnounced: slot_count must be 1 (one spill)"); + assert_eq!(memory_chunk_count, 1, "AllAnnounced: memory_chunk_count must be 1"); + assert_eq!(total_records, 4, "AllAnnounced: total_records must be 4"); + } + _ => panic!("expected AllAnnounced, got something else"), + } + + // Run 1 chunk: straddler (record 0) + rec2 (record 1) — 2 records. + assert_eq!(run1_chunk.len(), 2, "run 1 chunk must have 2 records (straddler + rec2)"); + + // Assert straddler is record 0 of run 1's sorted chunk. + // The straddler has key (ref_id=0, pos=5) and rec2 has (ref_id=1, pos=5); + // coordinate sort orders by ref_id first, so straddler (ref_id=0) comes before + // rec2 (ref_id=1). + let straddler_bytes = run1_chunk.record_bytes(0).to_vec(); + assert_eq!( + straddler_bytes, straddler, + "run 1 record 0 must be the straddler (carry ++ head = full original record)" + ); + let rec2_bytes = run1_chunk.record_bytes(1).to_vec(); + assert_eq!(rec2_bytes, rec2, "run 1 record 1 must be rec2"); + + // ----------------------------------------------------------------------- + // Oracle check: stable merge of run0 + run1 by coordinate key, lower + // source index first (run0 < run1) must equal the oracle over all 4 records. + // + // Oracle sort order: (ref_id=0,pos=5)=straddler, (ref_id=0,pos=10)=rec0, + // (ref_id=1,pos=5)=rec2, (ref_id=1,pos=20)=rec1. + // + // Merge gives: run0 records sorted = [rec0(0,10), rec1(1,20)]; + // run1 records sorted = [straddler(0,5), rec2(1,5)]. + // Interleaved lower-run-first stable merge: + // Compare run0[0]=(0,10) vs run1[0]=(0,5): run1 wins → straddler + // Compare run0[0]=(0,10) vs run1[1]=(1,5): run0 wins → rec0 + // Compare run0[1]=(1,20) vs run1[1]=(1,5): run1 wins → rec2 + // run1 exhausted → run0[1] = rec1 + // Merge result: [straddler, rec0, rec2, rec1] + // ----------------------------------------------------------------------- + let mut oracle = CoordinateChunkSorter::for_test(usize::MAX, n_ref); + for r in &all_recs { + let _ = oracle.push(r).unwrap(); + } + let oracle_chunk = oracle.take_sorted_chunk(); + let oracle_bytes: Vec> = + (0..oracle_chunk.len()).map(|i| oracle_chunk.record_bytes(i).to_vec()).collect(); + // Expected oracle order: [(0,5)=straddler, (0,10)=rec0, (1,5)=rec2, (1,20)=rec1] + assert_eq!(oracle_bytes[0], straddler, "oracle[0] must be straddler"); + assert_eq!(oracle_bytes[1], rec0, "oracle[1] must be rec0"); + assert_eq!(oracle_bytes[2], rec2, "oracle[2] must be rec2"); + assert_eq!(oracle_bytes[3], rec1, "oracle[3] must be rec1"); + + // Collect sorted bytes from both run chunks (run0 first, then run1) via a + // manual stable merge that mirrors the merge engine's lower-source-index-first + // tie-break. + // + // run0 sorted: [rec0(0,10), rec1(1,20)] + // run1 sorted: [straddler(0,5), rec2(1,5)] + let run0_sorted: Vec> = + (0..run0_chunk.len()).map(|i| run0_chunk.record_bytes(i).to_vec()).collect(); + let run1_sorted: Vec> = + (0..run1_chunk.len()).map(|i| run1_chunk.record_bytes(i).to_vec()).collect(); + + // Build a merged sequence via coordinate key comparison. + // We use the oracle order as the expected merged order (they must match). + let mut merged: Vec> = Vec::with_capacity(4); + let mut i0 = 0usize; + let mut i1 = 0usize; + while i0 < run0_sorted.len() || i1 < run1_sorted.len() { + if i0 >= run0_sorted.len() { + merged.push(run1_sorted[i1].clone()); + i1 += 1; + } else if i1 >= run1_sorted.len() { + merged.push(run0_sorted[i0].clone()); + i0 += 1; + } else { + // Extract coordinate key: ref_id (first i32) and pos (second i32). + let key0 = { + let b = &run0_sorted[i0]; + let r = i32::from_le_bytes(b[0..4].try_into().unwrap()); + let p = i32::from_le_bytes(b[4..8].try_into().unwrap()); + (r, p) + }; + let key1 = { + let b = &run1_sorted[i1]; + let r = i32::from_le_bytes(b[0..4].try_into().unwrap()); + let p = i32::from_le_bytes(b[4..8].try_into().unwrap()); + (r, p) + }; + // Stable: on tie, run0 (lower source index) wins. + if key1 < key0 { + merged.push(run1_sorted[i1].clone()); + i1 += 1; + } else { + merged.push(run0_sorted[i0].clone()); + i0 += 1; + } + } + } + + assert_eq!( + merged, oracle_bytes, + "stable merge of run0 + run1 chunks must equal the oracle over all 4 records" + ); + } + + #[allow(unsafe_code)] + #[test] + fn inflate_writes_decompressed_bytes_into_the_slot() { + // Build a payload that is comfortably under one BGZF block (< 64 KiB). + let payload = b"PARALLEL-INFLATE-ARENA-TEST".repeat(50); + let block = compress_one_block(&payload); + let isize = u32::try_from(payload.len()).unwrap(); + + // Construct an arena with enough capacity for the payload, reserve it, + // then carve out an uninit slot for the inflate worker to fill. + let mut arena = SegmentedBuf::with_capacity(0, 1 << 20); + arena.reserve_full_capacity(); + // SAFETY: slot is fully written by `inflate_one` below before any read. + let offset = unsafe { arena.grow_uninit(payload.len()) } as u64; + let arena = Arc::new(PooledSegmentedBuf::unpooled(arena)); + + let item = ArenaBlock { + arena: Arc::clone(&arena), + ordinal: 0, + offset, + len: isize, + block, + is_last_of_run: true, + run_seq: 0, + seals_to_spill: false, + }; + + let mut step = InflateToArena::new(64 * 1024 * 1024); + let inflated = step.inflate_one(item).expect("inflate must succeed"); + + assert_eq!(inflated.offset, offset, "offset must be forwarded unchanged"); + assert_eq!(inflated.len, isize, "len must be forwarded unchanged"); + assert!(!inflated.seals_to_spill, "seals_to_spill must be forwarded"); + // Confirm the decompressed bytes are in the arena at the reserved slot. + assert_eq!( + arena.slice( + usize::try_from(offset).unwrap(), + isize as usize, // u32 always fits in usize + ), + &payload[..], + "arena slot must contain the original payload after inflate" + ); + } + + /// Regression coverage for the `FindBoundariesAndSort` `try_run` + /// finalize/drained-branch held-overwrite bug. + /// + /// The bug: on the drained path `try_run` did `push(first_event)` and then, + /// in the SAME call, unconditionally drained the rest of `pending` via + /// `emit_pending` WITHOUT re-checking `held`. If the first push is rejected + /// (full output queue → `first_event` goes into `held`), the follow-on + /// `emit_pending` would push a LATER event (`AllAnnounced`) ahead of the + /// still-held `Residual` and then `held.put(...)` it — clobbering the held + /// `Residual` (`HeldSlot::put` asserts on a double-put, so it would panic / + /// lose the record). The fix drains `pending` only when the first push + /// SUCCEEDED, preserving the one-event-per-call discipline of the normal path. + /// + /// This test pins the ordering contract `finalize()` hands `try_run`, which is + /// exactly what the one-event-per-call fix must emit in order: `Residual` + /// FIRST, then `AllAnnounced`, each staged exactly once. If `finalize()` + /// staged them in the wrong order (or duplicated one), the held/retry path in + /// `try_run` could not emit a correct stream no matter how careful it is. + /// + /// Coverage limit: the held-overwrite lived inside `try_run`, which needs a + /// backpressuring `StepCtx` (an `OutputHandles>` whose queue rejects + /// the first push). `OutputHandles::new` is `pub(crate)` to + /// `fgumi-pipeline-core`, so a rejecting output context cannot be built from + /// this crate without editing another crate; the full `try_run` held/retry path + /// is exercised end-to-end by the framework-driven tests in `sort/tests.rs`. + /// Here we assert the finalize-branch invariant as directly as the in-crate + /// seams (`finalize`, `pending`) allow. + #[allow(unsafe_code)] + #[test] + fn finalize_stages_residual_before_all_announced_each_once() { + let n_ref = 2u32; + // A complete run (header + two whole records, no trailing partial), so + // sealing it as a residual leaves an empty carry. + let recs = [coord_body(0, 10, b'a'), coord_body(1, 20, b'b')]; + let header = minimal_bam_header(n_ref); + let mut arena = SegmentedBuf::with_capacity(0, 1 << 20); + arena.reserve_full_capacity(); + // SAFETY: every slot fully written before any read. + let h_off = unsafe { arena.grow_uninit(header.len()) }; + unsafe { arena.slice_mut(h_off, header.len()) }.copy_from_slice(&header); + #[allow(clippy::cast_possible_truncation)] + let run_start = h_off as u64; + for r in &recs { + let bs = u32::try_from(r.len()).unwrap(); + let po = unsafe { arena.grow_uninit(4) }; + unsafe { arena.slice_mut(po, 4) }.copy_from_slice(&bs.to_le_bytes()); + let bo = unsafe { arena.grow_uninit(r.len()) }; + unsafe { arena.slice_mut(bo, r.len()) }.copy_from_slice(r); + } + #[allow(clippy::cast_possible_truncation)] + let run_len = arena.len() as u64 - run_start; + let arena = Arc::new(PooledSegmentedBuf::unpooled(arena)); + + let mut step = + FindBoundariesAndSort::new(CoordinateStrategy::new(n_ref), 1, 64 * 1024 * 1024); + // Ingest the run WITHOUT `is_last_of_run`, so `ingest_block` does NOT seal: + // the seal (and the `Residual` + `AllAnnounced` staging) then happens in + // `finalize()` — the path `try_run`'s drained branch drives. + step.ingest_block(&InflatedBlock { + arena: Arc::clone(&arena), + ordinal: 0, + offset: run_start, + len: u32::try_from(run_len).unwrap(), + is_last_of_run: false, + run_seq: 0, + seals_to_spill: false, + }) + .expect("ingest_block must succeed"); + assert!(step.pending.is_empty(), "no events staged before finalize"); + + // finalize() returns the FIRST event (Residual) and leaves the remainder + // staged in `pending`. This is the one-event-at-a-time hand-off that + // `try_run`'s drained branch must respect: emit `first_event`, and only + // drain `pending` if that push landed. + let first = step.finalize().expect("finalize must succeed").expect("first event present"); + assert!( + matches!(first, SortChunkEvent::Residual { .. }), + "finalize must return Residual as the first event" + ); + // Exactly one event remains staged, and it is AllAnnounced (never emitted + // ahead of the Residual). + assert_eq!(step.pending.len(), 1, "exactly one event remains staged after the Residual"); + match step.pending.front().expect("AllAnnounced staged") { + SortChunkEvent::AllAnnounced { slot_count, memory_chunk_count, .. } => { + assert_eq!(*slot_count, 0, "no spills → slot_count 0"); + assert_eq!(*memory_chunk_count, 1, "one in-memory residual chunk"); + } + _ => panic!("second staged event must be AllAnnounced"), + } + + // A subsequent finalize() (mirroring a later drained `try_run` after the + // Residual flushed) hands back AllAnnounced, then nothing — proving each + // event is produced exactly once and in order. + let second = step.finalize().expect("finalize must succeed").expect("AllAnnounced present"); + assert!( + matches!(second, SortChunkEvent::AllAnnounced { .. }), + "second finalize must return AllAnnounced" + ); + assert!(step.pending.is_empty(), "no further events staged"); + assert!( + step.finalize().expect("finalize must succeed").is_none(), + "no more events after Residual + AllAnnounced" + ); + } + + /// Runtime proof that `--sort-threads` (Phase 1) controls the actual sort + /// worker count, not just that it parses. + /// + /// The chain builder resolves `--sort-threads` (falling back to `--threads`) + /// into `num_phase1_threads` and hands it to `FindBoundariesAndSort::new`, + /// which forwards it to `strategy.seal(arena, self.sort_threads)`. The + /// queryname strategy builds a bounded rayon pool sized to that value and runs + /// the per-run comparator sort inside it — so the pool's thread count IS the + /// effective Phase-1 concurrency. (The D1.1 regression handed the strategy the + /// raw global `--threads` instead, silently ignoring `--sort-threads`.) + /// + /// `ThreadPool::current_num_threads()` is the fixed pool size, so this is + /// deterministic — unlike counting how many workers a given input happens to + /// keep busy. Sort output is byte-identical across thread counts, so this pool + /// observation is the only way to assert the flag's runtime effect. + #[allow(unsafe_code)] + #[test] + fn phase1_sort_threads_sizes_the_queryname_worker_pool() { + use fgumi_sort::RawQuerynameLexKey; + + for sort_threads in [1usize, 3] { + let n_ref = 2u32; + let recs = [coord_body(0, 10, b'a'), coord_body(1, 20, b'b')]; + let header = minimal_bam_header(n_ref); + let mut arena = SegmentedBuf::with_capacity(0, 1 << 20); + arena.reserve_full_capacity(); + // SAFETY: every slot fully written before any read. + let h_off = unsafe { arena.grow_uninit(header.len()) }; + unsafe { arena.slice_mut(h_off, header.len()) }.copy_from_slice(&header); + #[allow(clippy::cast_possible_truncation)] + let run_start = h_off as u64; + for r in &recs { + let bs = u32::try_from(r.len()).unwrap(); + let po = unsafe { arena.grow_uninit(4) }; + unsafe { arena.slice_mut(po, 4) }.copy_from_slice(&bs.to_le_bytes()); + let bo = unsafe { arena.grow_uninit(r.len()) }; + unsafe { arena.slice_mut(bo, r.len()) }.copy_from_slice(r); + } + #[allow(clippy::cast_possible_truncation)] + let run_len = arena.len() as u64 - run_start; + let arena = Arc::new(PooledSegmentedBuf::unpooled(arena)); + + let mut step = FindBoundariesAndSort::new( + QuerynameStrategy::::new(MemoryChunkErased::QuerynameLex), + sort_threads, + 64 * 1024 * 1024, + ); + assert_eq!( + step.strategy().sort_pool_threads(), + None, + "sort pool is not built until the first seal" + ); + step.ingest_block(&InflatedBlock { + arena: Arc::clone(&arena), + ordinal: 0, + offset: run_start, + len: u32::try_from(run_len).unwrap(), + is_last_of_run: false, + run_seq: 0, + seals_to_spill: false, + }) + .expect("ingest_block must succeed"); + // finalize() seals the residual run, which builds + installs the + // bounded sort pool sized to `sort_threads`. + let _ = step.finalize().expect("finalize must succeed"); + assert_eq!( + step.strategy().sort_pool_threads(), + Some(sort_threads), + "Phase-1 queryname sort pool must be sized to sort_threads={sort_threads} \ + (the value FindBoundariesAndSort was constructed with)" + ); + } + } +} diff --git a/crates/fgumi-pipeline-io/src/sort/compress_spill.rs b/crates/fgumi-pipeline-io/src/sort/compress_spill.rs new file mode 100644 index 000000000..4381ede46 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/compress_spill.rs @@ -0,0 +1,228 @@ +//! `CompressSpill` — the second step of the P6 Phase-1 split (`SortBuffer` → +//! `CompressSpill` → `SortSpillDecompress` → `SortMerge`). +//! +//! `SortBuffer` (Serial) emits already-sorted chunks; `CompressSpill` +//! (`Parallel`) compresses each spill chunk to disk **inline on its framework +//! worker** (via [`fgumi_sort::write_sorted_chunk_inmem`], or the chunk's own +//! `write_spill` for the template-arena variant, retiring the private +//! `SortWorkerPool` compress path) and forwards the result as the existing +//! [`SortPhase1Event`], so `SortSpillDecompress` / `SortMerge` are unchanged. +//! +//! # Why a `Parallel` step is safe here +//! +//! `SortMerge` collects setup events and gates on **counts** (`slot_count` / +//! `memory_chunk_count` from `AllAnnounced`), not on event arrival order, so +//! multiple `CompressSpill` workers may emit `SpillReady` / `MemoryChunk` events +//! in any order. The one ordering-sensitive concern — the `LoserTree` tie-break +//! for equal sort keys — is handled by stamping each spill slot's `file_id` with +//! the chunk's **logical** spill index (`SortChunkEvent::Spill::seq`, assigned by +//! `SortBuffer`), so the tie-break is independent of which worker writes first. + +use std::io; +use std::path::Path; +use std::sync::Arc; + +use fgumi_sort::{SpillCodec, TmpDirAllocator}; +use parking_lot::Mutex; +use tempfile::TempDir; + +use crate::sort::protocol::{MemoryChunkErased, SortChunkEvent, SortPhase1Event}; +use fgumi_pipeline_core::{ + HeldRetry, Unpushed, + held::HeldSlot, + outputs::Single, + queues::QueueSpec, + reorder::BranchOrdering, + step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}, +}; + +/// `Parallel` step that compresses sorted spill chunks to disk and forwards +/// residual chunks, emitting [`SortPhase1Event`]s to `SortSpillDecompress`. +/// +/// Not to be confused with the similarly-named +/// [`SpillBlockCompress`](super::SpillBlockCompress): this `CompressSpill` is the +/// **composite compress-and-write-to-disk** step of the coarser +/// `SortBuffer → CompressSpill → SortSpillDecompress → SortMerge` chain, whereas +/// `SpillBlockCompress` is the **pure block-compression** middle step of the finer +/// `SpillGather → SpillBlockCompress → SpillWrite` split (where the disk write is a +/// separate `SpillWrite` step). +pub struct CompressSpill { + /// Shared temp-directory allocator (free-space-aware round-robin). Behind a + /// `Mutex` because the step is `Parallel`; the lock is held only for the + /// brief base-directory pick, never across the (expensive) compress+write. + alloc: Arc>, + /// Spill codec for chunk files (bgzf or zstd). + codec: SpillCodec, + /// Temp-file compression level (`0` = uncompressed bgzf; zstd level for zstd). + compression: u32, + /// At most one not-yet-pushable output event, parked on downstream + /// backpressure. This is the framework's standard backpressure idiom (there + /// is no peek-before-pop), identical to the sibling Parallel step + /// [`SortSpillDecompress`](super::SortSpillDecompress), and like it this one + /// held event sits **outside** `output_byte_limit`: total retained memory is + /// `queue_bytes + (≤1 event) × workers`. For a `SpillReady` the held event is + /// tiny (an `Arc` + path); only a `MemoryChunk` (residual) + /// holds records, and the residual must transit the pipeline occupying + /// ~`memory_limit` regardless of whether it sits in the queue or this slot, so + /// the held slot adds no peak beyond what the residual already costs. + held: HeldSlot>, + output_byte_limit: u64, + /// RAII temp-dir handles, shared across `Parallel` clones. Held for the + /// step's lifetime so spill files survive while being written; on the last + /// clone's drop (after the step finishes — i.e. every spill is written and + /// every slot has an open fd) the dirs are removed. `SortMerge` then reads + /// each slot via its already-open fd, matching the legacy `SortAndSpill` + /// unlink-after-emit lifetime (correct on the Unix targets). Empty in tests + /// that hold their own `TempDir`. Held purely for RAII (`Arc`-cloned to each + /// worker), never read for its value. + temp_dirs: Arc>, +} + +impl CompressSpill { + /// Build a `CompressSpill` step. + /// + /// `alloc` names spill files across the configured temp dirs; `codec` / + /// `compression` select the on-disk spill format. `temp_dirs` holds the RAII + /// handles for those dirs alive for the step's lifetime. `output_byte_limit` + /// byte-bounds the forwarded-event output queue (its `MemoryChunk` variant + /// retains sorted records, so the queue must budget on bytes, not count). + #[must_use] + pub fn new( + alloc: Arc>, + codec: SpillCodec, + compression: u32, + output_byte_limit: u64, + temp_dirs: Arc>, + ) -> Self { + Self { alloc, codec, compression, held: HeldSlot::new(), output_byte_limit, temp_dirs } + } + + fn flush_held(&mut self, ctx: &mut StepCtx<'_, Self>) -> bool { + // `true` once the slot is clear (was empty, or the held event flushed); + // `false` while it's still held under backpressure. Uses the canonical + // re-hold helper so the put-back-on-reject invariant lives in one place. + !matches!(ctx.outputs.retry_held(&mut self.held), HeldRetry::StillHeld) + } + + /// Allocate a spill path for chunk `seq`. Names the file by the logical spill + /// index so the path is unique without a shared counter, then draws a base + /// directory from the shared allocator (the only locked section). + fn spill_path(&self, seq: u32) -> io::Result { + let base = self.alloc.lock().next().map_err(|e| { + io::Error::other(format!("CompressSpill: temp-dir allocation failed: {e:#}")) + })?; + Ok(base.join(format!("chunk_{seq:04}.keyed"))) + } + + /// Compress one input event into the forwarded [`SortPhase1Event`]. The + /// `Spill` arm does the file write inline; the others are passthroughs. This + /// is `StepCtx`-free so it is unit-testable on synthetic chunks. + fn compress_event(&self, event: SortChunkEvent) -> io::Result { + match event { + SortChunkEvent::Spill { seq, chunk, records_ingested_so_far } => { + let path = self.spill_path(seq)?; + write_chunk(&chunk, &path, self.codec, self.compression)?; + let slot = fgumi_sort::open_spill_slot(&path, seq).map_err(|e| { + io::Error::other(format!( + "CompressSpill: failed to open spill slot {}: {e:#}", + path.display() + )) + })?; + Ok(SortPhase1Event::SpillReady { slot, path, records_ingested_so_far }) + } + SortChunkEvent::Residual { chunk, records_ingested_so_far } => { + // Wrap in a fresh, uniquely-owned `Arc`: the chunk is only ever + // moved (never cloned) onward, so `SortMerge`'s `Arc::try_unwrap` + // invariant holds. + Ok(SortPhase1Event::MemoryChunk { chunk: Arc::new(chunk), records_ingested_so_far }) + } + SortChunkEvent::AllAnnounced { slot_count, memory_chunk_count, total_records } => { + Ok(SortPhase1Event::AllAnnounced { slot_count, memory_chunk_count, total_records }) + } + } + } +} + +/// Dispatch a sorted [`MemoryChunkErased`] to [`fgumi_sort::write_sorted_chunk_inmem`] +/// (or the chunk's own `write_spill`, for the template-arena variant) for the +/// concrete key variant. +fn write_chunk( + chunk: &MemoryChunkErased, + path: &Path, + codec: SpillCodec, + compression: u32, +) -> io::Result<()> { + let result = match chunk { + MemoryChunkErased::Coordinate(c) => { + fgumi_sort::write_sorted_chunk_inmem(path, codec, compression, c) + } + MemoryChunkErased::QuerynameLex(c) => { + fgumi_sort::write_sorted_chunk_inmem(path, codec, compression, c) + } + MemoryChunkErased::QuerynameNatural(c) => { + fgumi_sort::write_sorted_chunk_inmem(path, codec, compression, c) + } + MemoryChunkErased::TemplateCoordinate(c) => c.write_spill(path, codec, compression), + }; + result.map_err(|e| { + io::Error::other(format!("CompressSpill: chunk write to {} failed: {e:#}", path.display())) + }) +} + +impl Clone for CompressSpill { + fn clone(&self) -> Self { + Self { + alloc: Arc::clone(&self.alloc), + codec: self.codec, + compression: self.compression, + held: HeldSlot::new(), + output_byte_limit: self.output_byte_limit, + temp_dirs: Arc::clone(&self.temp_dirs), + } + } +} + +impl Step for CompressSpill { + type Input = SortChunkEvent; + type Outputs = Single; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "CompressSpill", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], + branch_ordering: vec![BranchOrdering::None], + } + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + // 1. Drain held output first (at most one event is ever held, since we + // only pop a new input after `flush_held` clears the slot). + if !self.flush_held(ctx) { + return Ok(StepOutcome::Contention); + } + + // 2. Pop one input event, compress/forward it, hold on a full output. + if let Some(event) = ctx.input.pop() { + let forwarded = self.compress_event(event)?; + if let Err(unpushed) = ctx.outputs.push(forwarded) { + self.held.put(unpushed); + } + return Ok(StepOutcome::Progress); + } + + // 3. No input available. + if ctx.input.is_drained() { + return Ok(StepOutcome::Finished); + } + Ok(StepOutcome::NoProgress) + } + + fn new_worker_copy(&self) -> Self { + self.clone() + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/fgumi-pipeline-io/src/sort/compress_spill/tests.rs b/crates/fgumi-pipeline-io/src/sort/compress_spill/tests.rs new file mode 100644 index 000000000..789e91822 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/compress_spill/tests.rs @@ -0,0 +1,325 @@ +//! Unit tests for `CompressSpill::compress_event` (the `StepCtx`-free core). +//! +//! These drive the per-event compression directly on synthetic sorted chunks, +//! covering the three event arms: `Spill` (compress to disk, slot `file_id` = +//! logical `seq`), `Residual` (pass through as a uniquely-owned `MemoryChunk`), +//! and `AllAnnounced` (verbatim forward). The end-to-end chain parity vs the +//! legacy oracle is exercised later, once `SortBuffer` + `add_sort` are wired +//! (P6 increments 3–4). + +use super::*; + +use fgumi_raw_bam::RawRecord; +use fgumi_sort::{RawCoordinateKey, TemplateKey}; +use rstest::rstest; +use tempfile::TempDir; + +use crate::sort::protocol::MemoryChunkErased; + +/// Which easily-constructed sort-key variant a parameterized case exercises. +/// (The two queryname variants are the same generic `write_sorted_chunk::` +/// dispatch; they get real coverage from the inc-4 chain parity test, where +/// genuine queryname keys come from BAM data rather than fragile hand +/// construction.) +#[derive(Clone, Copy)] +enum SpillVariant { + Coordinate, + Template, +} + +type CoordRecords = Vec<(RawCoordinateKey, RawRecord)>; + +/// `n` coordinate records with distinct, sized payloads so the spill file has +/// real content (and multiple zstd frames for large `n`). +#[allow(clippy::cast_possible_truncation)] // payload byte is `% 251`, always fits u8 +fn coord_records(n: usize) -> CoordRecords { + (0..n) + .map(|i| { + // Distinct, ascending sort keys so the byte-identity checks exercise + // key serialization (not just record-byte order). + let key = RawCoordinateKey { sort_key: i as u64 }; + (key, RawRecord::from(vec![(i % 251) as u8; 100 + i % 64])) + }) + .collect() +} + +/// Deterministic coordinate records keyed off `seq`, so a concurrently-written +/// spill file can be checked against an independent reference write. +#[allow(clippy::cast_possible_truncation)] // payload byte is `% 251`, always fits u8 +fn coord_records_for(seq: u32) -> CoordRecords { + let n = 40 + (seq as usize % 24); + (0..n) + .map(|i| { + let byte = (seq as usize + i) % 251; + // Distinct keys per (seq, i) so key serialization is exercised. + let key = RawCoordinateKey { sort_key: (u64::from(seq) << 32) | i as u64 }; + (key, RawRecord::from(vec![byte as u8; 80 + i % 32])) + }) + .collect() +} + +/// Pack owned coordinate records into the zero-copy arena-backed +/// [`InMemoryChunk`] the `Coordinate` protocol variant now carries. The spill +/// file this produces must be byte-identical to a `write_sorted_chunk` of the +/// same owned records (the tests assert exactly that). +fn coord_chunk(records: CoordRecords) -> fgumi_sort::InMemoryChunk { + fgumi_sort::InMemoryChunk::from_owned_records( + records.into_iter().map(|(k, r)| (k, r.into_inner())).collect(), + ) +} + +/// `n` template-coordinate records (the other easily-constructed key variant). +#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] +fn template_records(n: usize) -> Vec<(TemplateKey, RawRecord)> { + (0..n) + .map(|i| { + let k = TemplateKey::new( + i as i32, + i as i32, + false, + i32::MAX, + i32::MAX, + false, + 0, + 0, + (0, false), + i as u64, + false, + ); + (k, RawRecord::from(vec![(i % 251) as u8; 120 + i % 48])) + }) + .collect() +} + +/// Wrap owned template records into the arena-backed `InMemoryChunk` +/// the `TemplateCoordinate` protocol variant now carries — the template analogue +/// of [`coord_chunk`]. The spill file this produces must be byte-identical to a +/// `write_sorted_chunk` of the same owned records (the test asserts exactly that). +fn template_chunk( + records: Vec<(TemplateKey, RawRecord)>, +) -> fgumi_sort::InMemoryChunk { + fgumi_sort::InMemoryChunk::from_owned_records( + records.into_iter().map(|(k, r)| (k, r.as_ref().to_vec())).collect(), + ) +} + +/// A `CompressSpill` over a single temp dir with an always-ample free-space +/// probe (deterministic, no dependency on the host's real free space), plus the +/// `TempDir` guard that must outlive the opened spill slots. +fn make_step(codec: SpillCodec, compression: u32) -> (CompressSpill, TempDir) { + let dir = TempDir::new().expect("temp dir"); + let alloc = + TmpDirAllocator::with_probe(vec![dir.path().to_path_buf()], Box::new(|_| Ok(u64::MAX)), 0) + .expect("allocator builds"); + // The test holds `dir` alive itself, so the step's RAII set is empty. + let step = CompressSpill::new( + Arc::new(Mutex::new(alloc)), + codec, + compression, + 8 * 1024 * 1024, + Arc::new(Vec::new()), + ); + (step, dir) +} + +/// A `Spill` event compresses the chunk to disk, opens a slot whose `file_id` +/// equals the logical `seq` (not write order), and writes bytes byte-identical +/// to a direct `write_sorted_chunk` of the same records — for every sort-key +/// variant the dispatch handles. +#[rstest] +#[case::coordinate(SpillVariant::Coordinate, 5, 1234)] +#[case::template(SpillVariant::Template, 2, 300)] +fn spill_event_writes_chunk_with_seq_file_id_byte_identical( + #[case] variant: SpillVariant, + #[case] seq: u32, + #[case] records_ingested: u64, +) { + let (step, dir) = make_step(SpillCodec::Zstd, 3); + + // Write an independent reference from the same (deterministic) records, then + // move those records into the event chunk — borrow-then-move avoids cloning + // the chunk (`MemoryChunkErased` is deliberately not `Clone`). + let reference = dir.path().join("reference.keyed"); + let chunk = match variant { + SpillVariant::Coordinate => { + let recs = coord_records(500); + fgumi_sort::write_sorted_chunk(&reference, SpillCodec::Zstd, 3, &recs).unwrap(); + MemoryChunkErased::Coordinate(coord_chunk(recs)) + } + SpillVariant::Template => { + let recs = template_records(300); + fgumi_sort::write_sorted_chunk(&reference, SpillCodec::Zstd, 3, &recs).unwrap(); + MemoryChunkErased::TemplateCoordinate(fgumi_sort::TemplateMemChunk::K40( + template_chunk(recs), + )) + } + }; + + let event = SortChunkEvent::Spill { seq, chunk, records_ingested_so_far: records_ingested }; + let forwarded = step.compress_event(event).expect("compress Spill event"); + + let SortPhase1Event::SpillReady { slot, path, records_ingested_so_far } = forwarded else { + panic!("Spill must forward as SpillReady"); + }; + assert_eq!(records_ingested_so_far, records_ingested, "records_ingested must be propagated"); + assert_eq!(slot.file_id, seq, "slot file_id must equal the logical spill seq"); + assert_eq!(slot.codec, SpillCodec::Zstd, "codec must be detected from the written magic"); + assert!(path.exists(), "spill file must exist"); + assert!(path.starts_with(dir.path()), "spill file must live under the allocated temp dir"); + assert!( + path.file_name().unwrap().to_str().unwrap().contains(&format!("{seq:04}")), + "spill file should be named by its seq" + ); + + // Byte-identical to a direct write of the same records (the step must add no + // framing of its own — it delegates straight to write_sorted_chunk). + assert_eq!( + std::fs::read(&path).unwrap(), + std::fs::read(&reference).unwrap(), + "CompressSpill output must match write_sorted_chunk byte-for-byte" + ); +} + +/// Distinct `seq` values yield distinct slots/paths, so concurrent workers never +/// collide and the merge can tie-break by `file_id`. +#[test] +fn distinct_seqs_yield_distinct_file_ids_and_paths() { + let (step, _dir) = make_step(SpillCodec::Bgzf, 1); + let mut paths = Vec::new(); + for seq in [0u32, 1, 7, 42] { + let event = SortChunkEvent::Spill { + seq, + chunk: MemoryChunkErased::Coordinate(coord_chunk(coord_records(50))), + records_ingested_so_far: u64::from(seq), + }; + let SortPhase1Event::SpillReady { slot, path, .. } = + step.compress_event(event).expect("compress") + else { + panic!("expected SpillReady"); + }; + assert_eq!(slot.file_id, seq); + paths.push(path); + } + let unique: std::collections::HashSet<_> = paths.iter().collect(); + assert_eq!(unique.len(), paths.len(), "every spill path must be unique"); +} + +/// A `Residual` event passes through as a `MemoryChunk` wrapping a uniquely-owned +/// `Arc` (the invariant `SortMerge`'s `Arc::try_unwrap` relies on). +#[test] +fn residual_event_passes_through_as_unique_memory_chunk() { + let (step, _dir) = make_step(SpillCodec::Zstd, 3); + let records = coord_records(120); + + let event = SortChunkEvent::Residual { + chunk: MemoryChunkErased::Coordinate(coord_chunk(records)), + records_ingested_so_far: 99, + }; + let forwarded = step.compress_event(event).expect("compress Residual event"); + + let SortPhase1Event::MemoryChunk { chunk, records_ingested_so_far } = forwarded else { + panic!("Residual must forward as MemoryChunk"); + }; + assert_eq!(records_ingested_so_far, 99); + let inner = Arc::try_unwrap(chunk).unwrap_or_else(|_| { + panic!("MemoryChunk Arc must be uniquely owned (SortMerge unwraps it)") + }); + assert_eq!(inner.len(), 120, "residual chunk content must pass through intact"); +} + +/// `AllAnnounced` forwards verbatim — the counts `SortBuffer` computed are the +/// completion target `SortMerge` keys off of. +#[test] +fn all_announced_passes_through_verbatim() { + let (step, _dir) = make_step(SpillCodec::Zstd, 3); + let event = + SortChunkEvent::AllAnnounced { slot_count: 4, memory_chunk_count: 1, total_records: 5000 }; + let forwarded = step.compress_event(event).expect("compress AllAnnounced"); + let SortPhase1Event::AllAnnounced { slot_count, memory_chunk_count, total_records } = forwarded + else { + panic!("AllAnnounced must forward as AllAnnounced"); + }; + assert_eq!((slot_count, memory_chunk_count, total_records), (4, 1, 5000)); +} + +/// Multiple cloned workers (the `new_worker_copy` fan-out) sharing one +/// allocator can compress spill chunks concurrently: every file is unique, +/// carries the right `file_id`, and is byte-identical to an independent +/// reference write — i.e. the shared `Mutex` is the only +/// cross-worker state and it serializes cleanly. +#[test] +fn parallel_workers_share_allocator_without_collision() { + let (base, dir) = make_step(SpillCodec::Zstd, 3); + let total_seqs: u32 = 64; + let num_workers = 8; + + // Each worker gets its own clone (fresh `held`), all sharing `base`'s + // allocator Arc — exactly how the framework materializes Parallel workers. + let results = std::thread::scope(|scope| { + let handles: Vec<_> = (0..num_workers) + .map(|w| { + let worker = base.clone(); + scope.spawn(move || { + let mut produced = Vec::new(); + let mut seq = w; + while seq < total_seqs { + let records = coord_records_for(seq); + let event = SortChunkEvent::Spill { + seq, + chunk: MemoryChunkErased::Coordinate(coord_chunk(records)), + records_ingested_so_far: u64::from(seq), + }; + let SortPhase1Event::SpillReady { slot, path, .. } = + worker.compress_event(event).expect("worker compress") + else { + panic!("expected SpillReady"); + }; + produced.push((seq, slot.file_id, path)); + seq += num_workers; + } + produced + }) + }) + .collect(); + handles.into_iter().map(|h| h.join().expect("worker thread")).collect::>() + }); + + let mut all: Vec<(u32, u32, std::path::PathBuf)> = results.into_iter().flatten().collect(); + all.sort_by_key(|(seq, _, _)| *seq); + + assert_eq!(all.len(), total_seqs as usize, "every seq must produce exactly one spill file"); + + let unique_paths: std::collections::HashSet<_> = all.iter().map(|(_, _, p)| p).collect(); + assert_eq!(unique_paths.len(), all.len(), "no two workers may collide on a spill path"); + + for (seq, file_id, path) in &all { + assert_eq!(file_id, seq, "file_id must equal the logical seq, not write order"); + // Independent reference write of this seq's deterministic records. + let reference = dir.path().join(format!("ref_{seq:04}.keyed")); + fgumi_sort::write_sorted_chunk(&reference, SpillCodec::Zstd, 3, &coord_records_for(*seq)) + .unwrap(); + assert_eq!( + std::fs::read(path).unwrap(), + std::fs::read(&reference).unwrap(), + "concurrently-written chunk {seq} must match its reference byte-for-byte" + ); + } +} + +/// An empty residual chunk is still a valid passthrough (zero-record fast path). +#[test] +fn empty_residual_passes_through() { + let (step, _dir) = make_step(SpillCodec::Zstd, 3); + let event = SortChunkEvent::Residual { + chunk: MemoryChunkErased::Coordinate(coord_chunk(Vec::new())), + records_ingested_so_far: 0, + }; + let SortPhase1Event::MemoryChunk { chunk, .. } = + step.compress_event(event).expect("compress empty residual") + else { + panic!("expected MemoryChunk"); + }; + let inner = + Arc::try_unwrap(chunk).unwrap_or_else(|_| panic!("MemoryChunk Arc must be uniquely owned")); + assert!(inner.is_empty()); +} diff --git a/crates/fgumi-pipeline-io/src/sort/merge.rs b/crates/fgumi-pipeline-io/src/sort/merge.rs new file mode 100644 index 000000000..2abdae65c --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/merge.rs @@ -0,0 +1,1128 @@ +//! `SortMerge` — third step of the runall-sort three-step chain. + +use std::collections::HashMap; +use std::io; +use std::sync::Arc; + +use fgumi_sort::{ + CbKey32, InMemoryChunk, MemorySources, MergeDriver, MergeDriverDyn, MergeStep, + QuerynameComparator, RawCoordinateKey, RawQuerynameKey, RawQuerynameLexKey, SortMergeSlot, + SortOrder, TemplateKey, TemplateKey24, TemplateMemChunk, TertKey32, +}; + +use crate::sort::protocol::{MemoryChunkErased, SortPhase2Event}; +use crate::types::{DecompressedBlock, RecordBatch, RecordBatchBuilder}; +use fgumi_pipeline_core::{ + HeapSize, HeldRetry, Ordered, Unpushed, + held::HeldSlot, + outputs::OrderedBytesSingle, + queues::QueueSpec, + reorder::BranchOrdering, + step::{DetachedGroup, Step, StepCtx, StepKind, StepOutcome, StepProfile}, +}; + +/// Default output batch size: 1024 records per emitted `RecordBatch`. +pub const DEFAULT_TARGET_BATCH_COUNT: usize = 1024; + +/// Max output batches emitted per `try_run` invocation in `Merging`. +const MAX_DRAIN_BATCHES_PER_LOCK: usize = 8; + +/// Initial reservation for an output-batch byte buffer, before any batch has +/// been emitted to size the next one from. Kept modest on purpose: most batches +/// fill on the record-count cap well below the output-queue byte budget, so +/// reserving the full budget for every buffer chronically over-allocates (and +/// inflates the byte-bounded queue's capacity-based accounting). Buffers grow +/// on demand via `extend_from_slice`, so under-reserving only costs a few +/// startup reallocations. +const INITIAL_OUTPUT_BUFFER_BYTES: usize = 64 * 1024; + +// ───────────────────────────────────────────────────────────────────────────── +// MergeOutput — the framing strategy the merge accumulates winners into. +// ───────────────────────────────────────────────────────────────────────────── + +/// The output-framing strategy `SortMerge` accumulates merged winner records +/// into. The merge state machine, `LoserTree` driver, source ordering and +/// tie-break are identical for every strategy; only the per-record framing and +/// the emitted item type differ. +/// +/// Two implementations exist: +/// +/// - [`RecordBatchOutput`] (the default) — accumulates raw record bodies into a +/// [`RecordBatch`] (flat backing buffer + per-record `(start, end)` ranges). +/// This is the **intermediate** sort output, consumed by `DecodeFromRecords` +/// downstream in a fused `--start-from sort` chain. +/// - [`BlockOutput`] — frames each record as `[u32 LE block_size][body]` +/// directly into a [`DecompressedBlock`], byte-for-byte identical to the +/// `SerializeRecordBatch` step it replaces (lever 1). This is the +/// **terminal** standalone-sort output, wired straight to `BgzfCompress`, +/// folding the former `SortMerge → SerializeRecordBatch → BgzfCompress` +/// triple into `SortMerge → BgzfCompress` (one fewer pool step, one fewer +/// reorder stage, one fewer memcpy per record). +pub trait MergeOutput: Send + 'static { + /// The emitted batch item type. + type Item: Send + HeapSize + Ordered + 'static; + /// The per-batch accumulator. + type Builder: MergeBatchBuilder; +} + +/// A per-batch accumulator for a [`MergeOutput`] strategy. Mirrors the +/// [`RecordBatchBuilder`] surface the merge loop already drives, so the merge +/// state machine is strategy-agnostic. +pub trait MergeBatchBuilder: Send + 'static { + /// The finalized batch item this builder produces. + type Item; + + /// Create a builder for batch `batch_serial`, reserving `bytes_cap` bytes of + /// payload and room for `records_cap` records. + fn with_capacity(batch_serial: u64, bytes_cap: usize, records_cap: usize) -> Self; + + /// Append one merged winner record's raw BAM body. + /// + /// # Errors + /// + /// Returns an error if the record cannot be framed (e.g. a body whose + /// length does not fit the strategy's length prefix). + fn push_record_bytes(&mut self, body: &[u8]) -> io::Result<()>; + + /// Number of records appended so far. + fn len(&self) -> usize; + + /// `true` iff no records have been appended. + fn is_empty(&self) -> bool; + + /// Total payload bytes accumulated so far (used to size the next buffer and + /// to enforce the per-batch byte cap). + fn total_bytes(&self) -> usize; + + /// Finalize and produce the batch item, consuming the builder. + fn build(self) -> Self::Item; +} + +/// Intermediate-sort output: raw record bodies into a [`RecordBatch`]. +pub struct RecordBatchOutput; + +impl MergeOutput for RecordBatchOutput { + type Item = RecordBatch; + type Builder = RecordBatchBuilder; +} + +impl MergeBatchBuilder for RecordBatchBuilder { + type Item = RecordBatch; + + fn with_capacity(batch_serial: u64, bytes_cap: usize, records_cap: usize) -> Self { + RecordBatchBuilder::with_capacity(batch_serial, bytes_cap, records_cap) + } + + fn push_record_bytes(&mut self, body: &[u8]) -> io::Result<()> { + RecordBatchBuilder::push_record_bytes(self, body); + Ok(()) + } + + fn len(&self) -> usize { + RecordBatchBuilder::len(self) + } + + fn is_empty(&self) -> bool { + RecordBatchBuilder::is_empty(self) + } + + fn total_bytes(&self) -> usize { + RecordBatchBuilder::total_bytes(self) + } + + fn build(self) -> RecordBatch { + RecordBatchBuilder::build(self) + } +} + +/// Terminal standalone-sort output: each record framed as +/// `[u32 LE block_size][body]` directly into a [`DecompressedBlock`], ready for +/// `BgzfCompress`. This is byte-for-byte identical to the framing the former +/// `SerializeRecordBatch` step produced (lever 1). +pub struct BlockOutput; + +impl MergeOutput for BlockOutput { + type Item = DecompressedBlock; + type Builder = BlockBuilder; +} + +/// Accumulates merged winner records as BAM on-disk framing +/// (`[u32 LE block_size][body]` per record) into a single byte buffer that +/// becomes a [`DecompressedBlock`]. This is the canonical BAM record layout; +/// the `fgumi` crate's `serialize::frame_record_into` is the sibling +/// implementation (a separate crate, so the two cannot share code) and the two +/// MUST stay byte-for-byte in sync — each has a layout test pinning it. +pub struct BlockBuilder { + batch_serial: u64, + bytes: Vec, + /// Record count — tracked separately because the framed byte buffer mixes + /// length prefixes with bodies, so it cannot be recovered from `bytes`. + records: usize, +} + +impl MergeBatchBuilder for BlockBuilder { + type Item = DecompressedBlock; + + fn with_capacity(batch_serial: u64, bytes_cap: usize, _records_cap: usize) -> Self { + // `_records_cap` sizes the `RecordBatch` ranges vector; the framed-block + // builder has no separate per-record allocation to reserve. + Self { batch_serial, bytes: Vec::with_capacity(bytes_cap), records: 0 } + } + + fn push_record_bytes(&mut self, body: &[u8]) -> io::Result<()> { + // `[u32 LE block_size][body]`, byte-identical to + // `SerializeRecordBatch::frame_record_into`. + let block_size = u32::try_from(body.len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("record exceeds u32 BAM block_size: {}", body.len()), + ) + })?; + self.bytes.extend_from_slice(&block_size.to_le_bytes()); + self.bytes.extend_from_slice(body); + self.records += 1; + Ok(()) + } + + fn len(&self) -> usize { + self.records + } + + fn is_empty(&self) -> bool { + self.records == 0 + } + + fn total_bytes(&self) -> usize { + self.bytes.len() + } + + fn build(self) -> DecompressedBlock { + DecompressedBlock { batch_serial: self.batch_serial, bytes: self.bytes } + } +} + +/// Same-variant collector for template-coordinate residual chunks. +/// +/// A single sort chooses its `--key-types` narrowed lane variant exactly once +/// (globally, on the first record) and reuses it for every run, so all template +/// chunks in one merge share one arm. The first push fixes the arm; subsequent +/// pushes assert-match it (a variant change mid-sort is impossible by +/// construction and would be a bug). +#[derive(Default)] +enum TemplateChunks { + /// No template chunks pushed yet — the variant is not yet known. + #[default] + Empty, + /// 24-byte core-only lane. + K24(Vec>), + /// 32-byte lane carrying `cb_hash`. + Cb32(Vec>), + /// 32-byte lane carrying the tertiary word. + Tert32(Vec>), + /// Full 40-byte key (all lanes) — the legacy owned path and full variant. + K40(Vec>), +} + +impl TemplateChunks { + /// Name the narrowed-lane variant for diagnostics. + fn variant_name(&self) -> &'static str { + match self { + Self::Empty => "empty", + Self::K24(_) => "K24", + Self::Cb32(_) => "Cb32", + Self::Tert32(_) => "Tert32", + Self::K40(_) => "K40", + } + } + + /// Accumulate one template chunk, which must keep the lane variant fixed. + /// + /// The `--key-types` narrowed-lane variant is chosen once per sort and is + /// global to the run, so every template chunk reaching the merge must carry + /// the same one. A variant change means the phase-1 producer and the merge + /// consumer disagree about the key width, and merging on would compare keys + /// of different layouts and emit silently mis-ordered output. + /// + /// # Errors + /// + /// Returns `InvalidData` if `chunk`'s variant differs from the accumulated + /// one. This is the same fail-closed treatment the sibling protocol + /// violations get (`ensure_single_lane`, `build_driver`) rather than a + /// panic, so a corrupt stream aborts the sort with a diagnosable error. + fn push(&mut self, chunk: TemplateMemChunk) -> io::Result<()> { + /// Build the mismatch error, naming both variants. + fn mismatch(found: &str, have: &str) -> io::Error { + io::Error::new( + io::ErrorKind::InvalidData, + format!( + "SortMerge: template chunk variant changed mid-sort \ + (accumulated {have}, got {found}); the --key-types lane \ + variant is global to a sort and must not change" + ), + ) + } + match chunk { + TemplateMemChunk::K24(c) => match self { + Self::Empty => *self = Self::K24(vec![c]), + Self::K24(v) => v.push(c), + other => return Err(mismatch("K24", other.variant_name())), + }, + TemplateMemChunk::Cb32(c) => match self { + Self::Empty => *self = Self::Cb32(vec![c]), + Self::Cb32(v) => v.push(c), + other => return Err(mismatch("Cb32", other.variant_name())), + }, + TemplateMemChunk::Tert32(c) => match self { + Self::Empty => *self = Self::Tert32(vec![c]), + Self::Tert32(v) => v.push(c), + other => return Err(mismatch("Tert32", other.variant_name())), + }, + TemplateMemChunk::K40(c) => match self { + Self::Empty => *self = Self::K40(vec![c]), + Self::K40(v) => v.push(c), + other => return Err(mismatch("K40", other.variant_name())), + }, + } + Ok(()) + } + + fn len(&self) -> usize { + match self { + Self::Empty => 0, + Self::K24(v) => v.len(), + Self::Cb32(v) => v.len(), + Self::Tert32(v) => v.len(), + Self::K40(v) => v.len(), + } + } + + /// Pop the sole chunk (caller guarantees exactly one) and re-erase it. + fn pop_single(self) -> TemplateMemChunk { + match self { + Self::K24(mut v) => TemplateMemChunk::K24(v.pop().expect("one chunk")), + Self::Cb32(mut v) => TemplateMemChunk::Cb32(v.pop().expect("one chunk")), + Self::Tert32(mut v) => TemplateMemChunk::Tert32(v.pop().expect("one chunk")), + Self::K40(mut v) => TemplateMemChunk::K40(v.pop().expect("one chunk")), + Self::Empty => unreachable!("pop_single called with no chunk"), + } + } +} + +#[derive(Default)] +struct MemoryChunksByKind { + coordinate: Vec>, + queryname_lex: Vec>, + queryname_natural: Vec>, + template_coordinate: TemplateChunks, +} + +impl MemoryChunksByKind { + /// Accumulate one erased chunk into its per-order bucket. + /// + /// # Errors + /// + /// Propagates the template lane-variant mismatch from + /// [`TemplateChunks::push`]; the other orders are infallible. + fn push(&mut self, chunk: MemoryChunkErased) -> io::Result<()> { + match chunk { + MemoryChunkErased::Coordinate(v) => self.coordinate.push(v), + MemoryChunkErased::QuerynameLex(v) => self.queryname_lex.push(v), + MemoryChunkErased::QuerynameNatural(v) => self.queryname_natural.push(v), + MemoryChunkErased::TemplateCoordinate(v) => { + return self.template_coordinate.push(v); + } + } + Ok(()) + } + + fn total_len(&self) -> usize { + self.coordinate.len() + + self.queryname_lex.len() + + self.queryname_natural.len() + + self.template_coordinate.len() + } + + /// Fail closed if any lane other than the one `sort_order` selects holds a + /// chunk. `build_driver` (and the single-chunk fast path) consume only the + /// selected lane, so a chunk in another lane — a Phase-2 protocol violation + /// emitting the wrong `MemoryChunkErased` variant — would be silently dropped + /// even though `total_len()` counted it toward setup completion. Reject it + /// rather than merge a partial result. + fn ensure_single_lane(&self, sort_order: SortOrder) -> io::Result<()> { + let selected_len = match sort_order { + SortOrder::Coordinate => self.coordinate.len(), + SortOrder::Queryname(QuerynameComparator::Lexicographic) => self.queryname_lex.len(), + SortOrder::Queryname(QuerynameComparator::Natural) => self.queryname_natural.len(), + SortOrder::TemplateCoordinate => self.template_coordinate.len(), + }; + let stray = self.total_len() - selected_len; + if stray > 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "SortMerge: {stray} residual memory chunk(s) in a lane not matching the \ + {sort_order:?} sort order — Phase-2 emitted a mismatched MemoryChunkErased \ + variant; failing closed rather than silently dropping records", + ), + )); + } + Ok(()) + } + + /// Consume the single chunk held across all kinds, re-erased. + /// + /// # Panics + /// + /// Panics if `total_len() != 1` (the single-source fast path's precondition). + fn into_single(mut self) -> MemoryChunkErased { + debug_assert_eq!(self.total_len(), 1, "into_single requires exactly one chunk"); + if let Some(c) = self.coordinate.pop() { + MemoryChunkErased::Coordinate(c) + } else if let Some(c) = self.queryname_lex.pop() { + MemoryChunkErased::QuerynameLex(c) + } else if let Some(c) = self.queryname_natural.pop() { + MemoryChunkErased::QuerynameNatural(c) + } else if self.template_coordinate.len() == 1 { + MemoryChunkErased::TemplateCoordinate( + std::mem::take(&mut self.template_coordinate).pop_single(), + ) + } else { + unreachable!("into_single called with no chunk") + } + } +} + +fn build_driver( + sort_order: SortOrder, + slots: Vec>, + chunks: MemoryChunksByKind, + total_records: u64, +) -> io::Result> { + Ok(match sort_order { + SortOrder::Coordinate => Box::new(MergeDriver::::from_slots( + slots, + MemorySources::Shared(chunks.coordinate), + total_records, + )), + SortOrder::Queryname(QuerynameComparator::Lexicographic) => { + Box::new(MergeDriver::::from_slots( + slots, + MemorySources::Shared(chunks.queryname_lex), + total_records, + )) + } + SortOrder::Queryname(QuerynameComparator::Natural) => { + Box::new(MergeDriver::::from_slots( + slots, + MemorySources::Shared(chunks.queryname_natural), + total_records, + )) + } + SortOrder::TemplateCoordinate => match chunks.template_coordinate { + // `Empty` means no residual chunk identified the `--key-types` lane. + // For valid input this only happens with empty input (no spill files + // either) — Phase-1's deferred seal always emits a variant-tagged + // residual otherwise. So `Empty` WITH spill slots can only arise from + // the documented "defensive/unreachable" no-residual finalize branch + // (a seal-logic regression). Defaulting to `TemplateKey` (K40) there + // would decode narrow (K24/Cb32/Tert32) spill files at the wrong key + // width and silently corrupt output, so fail closed instead of + // guessing the width. With no slots, any K is safe (nothing to merge). + TemplateChunks::Empty => { + if !slots.is_empty() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "SortMerge: template-coordinate spill slots present but no residual \ + chunk to identify the key-types lane — refusing to guess the key \ + width (would mis-decode narrow spill files). This indicates a \ + Phase-1 seal-logic regression.", + )); + } + Box::new(MergeDriver::::from_slots( + slots, + MemorySources::Shared(Vec::new()), + total_records, + )) + } + TemplateChunks::K24(v) => Box::new(MergeDriver::::from_slots( + slots, + MemorySources::Shared(v), + total_records, + )), + TemplateChunks::Cb32(v) => Box::new(MergeDriver::::from_slots( + slots, + MemorySources::Shared(v), + total_records, + )), + TemplateChunks::Tert32(v) => Box::new(MergeDriver::::from_slots( + slots, + MemorySources::Shared(v), + total_records, + )), + TemplateChunks::K40(v) => Box::new(MergeDriver::::from_slots( + slots, + MemorySources::Shared(v), + total_records, + )), + }, + }) +} + +enum NextBatch { + Batch(I), + Stalled(Option), + Done(Option, u64), +} + +enum SortMergeState { + WaitingForSetup { + slots: Vec>, + slot_index: HashMap, + memory_chunks: MemoryChunksByKind, + total_records: u64, + expected_slot_count: Option, + expected_memory_chunk_count: Option, + }, + Merging { + driver: Box, + builder: B, + next_ordinal: u64, + }, + /// Single-source fast path: 0 spill slots and exactly one in-memory chunk, so + /// the chunk is already globally sorted and no k-way merge is needed. Gather + /// its records in order straight into output blocks (the dominant in-memory + /// cost — a single-threaded k = 1 loser-tree walk — is pure overhead here). + FastPath { + chunk: MemoryChunkErased, + /// Index of the next record to gather. + cursor: usize, + /// Total record count (`chunk.len()`), cached to avoid re-dispatching. + total: usize, + builder: B, + next_ordinal: u64, + }, + Done, +} + +fn absorb_phase2_event( + event: SortPhase2Event, + slots: &mut Vec>, + slot_index: &mut HashMap, + memory_chunks: &mut MemoryChunksByKind, + total_records: &mut u64, + expected_slot_count: &mut Option, + expected_memory_chunk_count: &mut Option, +) -> io::Result<()> { + match event { + SortPhase2Event::SpillReady { slot, path: _, records_ingested_so_far } => { + if let std::collections::hash_map::Entry::Vacant(e) = slot_index.entry(slot.file_id) { + e.insert(slots.len()); + slots.push(slot); + } + *total_records = (*total_records).max(records_ingested_so_far); + } + SortPhase2Event::MemoryChunk { chunk, records_ingested_so_far } => { + // The MemoryChunk `Arc` is created uniquely in the Phase-1 producer + // (`CompressSpill` / `SpillWrite`) and only ever *moved* (never cloned) + // through `SortSpillDecompress` to here, so + // it must be uniquely owned at this single consumer. Fail closed on a + // shared `Arc` rather than silently deep-cloning a potentially large + // record vector on the merge setup path. + let inner = Arc::try_unwrap(chunk).map_err(|_| { + io::Error::other( + "SortMerge: MemoryChunk Arc unexpectedly shared at the merge consumer \ + (protocol invariant: memory chunks are moved, never cloned)", + ) + })?; + memory_chunks.push(inner)?; + *total_records = (*total_records).max(records_ingested_so_far); + } + SortPhase2Event::AllAnnounced { + slot_count, + memory_chunk_count, + total_records: ar_total, + } => { + // The Phase-2 protocol emits exactly one `AllAnnounced` (the last event + // from the Phase-1 producer, `CompressSpill` / `SpillWrite`). A second + // one is a protocol violation; fail + // closed rather than overwrite the prior expectations and risk masking + // the bug behind a silently-different completion target. + if expected_slot_count.is_some() || expected_memory_chunk_count.is_some() { + return Err(io::Error::other(format!( + "SortMerge: duplicate AllAnnounced — prior {expected_slot_count:?}/\ + {expected_memory_chunk_count:?}, new {slot_count}/{memory_chunk_count}; \ + the Phase-2 protocol emits exactly one AllAnnounced", + ))); + } + *expected_slot_count = Some(slot_count); + *expected_memory_chunk_count = Some(memory_chunk_count); + *total_records = (*total_records).max(ar_total); + } + } + Ok(()) +} + +fn slot_set_complete( + slots_len: usize, + memory_chunks_total_len: usize, + expected_slot_count: Option, + expected_memory_chunk_count: Option, +) -> bool { + matches!( + (expected_slot_count, expected_memory_chunk_count), + (Some(want_slots), Some(want_chunks)) + if u32::try_from(slots_len).unwrap_or(u32::MAX) == want_slots + && u32::try_from(memory_chunks_total_len).unwrap_or(u32::MAX) == want_chunks + ) +} + +/// `Detached + ByItemOrdinal` terminal merge: the final of the three sort +/// steps, producing the sorted output stream consumed by the sink. +/// +/// Generic over the output-framing strategy `O` (see [`MergeOutput`]): +/// [`RecordBatchOutput`] (the default) emits [`RecordBatch`] for a fused +/// intermediate sort, and [`BlockOutput`] frames records directly into +/// [`DecompressedBlock`]s for the standalone-sort terminal so the chain can +/// wire `SortMerge → BgzfCompress` with no intervening serialize step +/// (lever 1). The merge core — `LoserTree` driver, source ordering, tie-break, +/// cooperative `try_run` body — is identical for both. +pub struct SortMerge { + state: SortMergeState, + held: HeldSlot>, + sort_order: SortOrder, + target_batch_count: usize, + output_byte_limit: u64, + /// Optional sink for the end-of-run sort summary, filled when the merge + /// reaches `Done`. The standalone-sort summary finalize hook reads it to + /// log records processed/written and the spill-chunk count. + stats_slot: Option>>>, + /// Total records ingested, captured at the merge transition (the summary's + /// "records processed"). + processed: u64, + /// Number of spill files, captured at the merge transition (the summary's + /// "temporary chunks"). Zero for a fully in-memory sort. + chunk_count: usize, + /// INSTRUMENTATION (lever-2 merge-stall diagnosis; `RUST_LOG=info` at Done). + /// `SortMerge` runs on a single dedicated `Detached` thread (one instance, + /// never `new_worker_copy`'d), so plain `&mut self` counters are sound — no + /// atomics needed. + dbg: MergeDiag, +} + +/// Lever-2 diagnostic counters: is the serial merge starved on decompress +/// (`input-empty`/`stalls`) or blocked on the downstream writer +/// (`output_full`), and how much does its worker spin (`contention`)? +#[derive(Default, Clone, Copy)] +struct MergeDiag { + /// Merge-loop passes that ended `Stalled` — the winning source's next block + /// was not yet decompressed (INPUT-STARVED: the lever-2 hypothesis). + stalls: u64, + /// `ctx.outputs.push` returned `Err` — downstream (compress/write) full + /// (OUTPUT-BACKPRESSURE). + output_full: u64, + /// `try_run` returned `Contention` — the merge worker had nothing to do this + /// dispatch and spun/yielded (pure under-utilization). + contention: u64, + /// `try_run` calls that delivered ≥1 batch (PROGRESS dispatches). + progress_dispatches: u64, +} + +impl SortMerge { + /// Build a `SortMerge` step with default batch size. + #[must_use] + pub fn new(sort_order: SortOrder, output_byte_limit: u64) -> Self { + Self::with_target_batch_count(sort_order, output_byte_limit, DEFAULT_TARGET_BATCH_COUNT) + } + + /// Build a `SortMerge` step with a custom output batch size. + #[must_use] + pub fn with_target_batch_count( + sort_order: SortOrder, + output_byte_limit: u64, + target_batch_count: usize, + ) -> Self { + Self { + state: SortMergeState::WaitingForSetup { + slots: Vec::new(), + slot_index: HashMap::new(), + memory_chunks: MemoryChunksByKind::default(), + total_records: 0, + expected_slot_count: None, + expected_memory_chunk_count: None, + }, + held: HeldSlot::new(), + sort_order, + target_batch_count: target_batch_count.max(1), + output_byte_limit, + stats_slot: None, + processed: 0, + chunk_count: 0, + dbg: MergeDiag::default(), + } + } + + /// Attach a slot to receive the end-of-run [`fgumi_sort::SortStats`] when + /// the merge completes (records processed/written + spill-chunk count). Used + /// by the standalone-sort summary finalize hook; runall leaves it unset. + #[must_use] + pub fn with_stats_slot( + mut self, + slot: Arc>>, + ) -> Self { + self.stats_slot = Some(slot); + self + } + + fn flush_held(&mut self, ctx: &mut StepCtx<'_, Self>) -> bool { + // `true` once the slot is clear (was empty, or the held event flushed); + // `false` while it's still held under backpressure. Uses the canonical + // re-hold helper so the put-back-on-reject invariant lives in one place. + !matches!(ctx.outputs.retry_held(&mut self.held), HeldRetry::StillHeld) + } + + /// Drains every currently-available input event into the setup state and + /// returns the number absorbed. The drain is intentionally unbounded — the + /// upstream queue is byte-bounded, so memory is gated on the producer side. + /// + /// # Panics + /// + /// Panics if `self.state` is not `WaitingForSetup`. + /// + /// # Errors + /// + /// Returns an error on a Phase-2 protocol violation (a duplicate + /// `AllAnnounced`, or a `MemoryChunk` whose `Arc` is unexpectedly shared). + fn absorb_events_into_setup(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + let SortMergeState::WaitingForSetup { + slots, + slot_index, + memory_chunks, + total_records, + expected_slot_count, + expected_memory_chunk_count, + } = &mut self.state + else { + unreachable!("absorb_events_into_setup called outside WaitingForSetup state"); + }; + let mut absorbed = 0usize; + while let Some(event) = ctx.input.pop() { + absorb_phase2_event( + event, + slots, + slot_index, + memory_chunks, + total_records, + expected_slot_count, + expected_memory_chunk_count, + )?; + absorbed += 1; + } + Ok(absorbed) + } + + fn is_ready_to_merge(&self) -> bool { + let SortMergeState::WaitingForSetup { + slots, + memory_chunks, + expected_slot_count, + expected_memory_chunk_count, + .. + } = &self.state + else { + return false; + }; + slot_set_complete( + slots.len(), + memory_chunks.total_len(), + *expected_slot_count, + *expected_memory_chunk_count, + ) + } + + /// # Panics + /// + /// Panics if `self.state` is not `Merging`. + fn next_batch(&mut self) -> io::Result> { + let target = self.target_batch_count; + let byte_limit = self.output_byte_limit; + let bytes_cap = usize::try_from(byte_limit).unwrap_or(usize::MAX); + let SortMergeState::Merging { driver, builder, next_ordinal } = &mut self.state else { + unreachable!("next_batch called outside Merging state"); + }; + + let buffer_floor = INITIAL_OUTPUT_BUFFER_BYTES.min(bytes_cap); + let flush = |builder: &mut O::Builder, next_ordinal: &mut u64| { + *next_ordinal += 1; + // Size the next buffer to the batch we just filled, clamped to + // `[buffer_floor, bytes_cap]`. Count-bound batches stay small; a + // byte-bound batch carries ~`bytes_cap` forward. This avoids + // reserving the full byte budget for every (typically count-bound) + // batch — see `INITIAL_OUTPUT_BUFFER_BYTES`. + let hint = builder.total_bytes().clamp(buffer_floor, bytes_cap); + let next_builder = O::Builder::with_capacity(*next_ordinal, hint, target); + std::mem::replace(builder, next_builder).build() + }; + let flush_partial = |builder: &mut O::Builder, next_ordinal: &mut u64| { + if builder.is_empty() { None } else { Some(flush(builder, next_ordinal)) } + }; + + loop { + match driver + .try_step() + .map_err(|e| io::Error::other(format!("SortMerge: merge step failed: {e:#}")))? + { + MergeStep::Produced(bytes) => { + builder.push_record_bytes(bytes)?; + let count_full = builder.len() >= target; + let bytes_full = (builder.total_bytes() as u64) >= byte_limit; + if count_full || bytes_full { + return Ok(NextBatch::Batch(flush(builder, next_ordinal))); + } + } + MergeStep::Stalled => { + return Ok(NextBatch::Stalled(flush_partial(builder, next_ordinal))); + } + MergeStep::Done => { + return Ok(NextBatch::Done( + flush_partial(builder, next_ordinal), + driver.records_merged(), + )); + } + } + } + } + + /// # Panics + /// + /// Panics if `self.state` is not `Merging`. + fn emit_batches_cooperative(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + let mut delivered = 0usize; + loop { + match self.next_batch()? { + NextBatch::Batch(batch) => { + if let Err(unpushed) = ctx.outputs.push(batch) { + self.dbg.output_full += 1; + self.held.put(unpushed); + return Ok(StepOutcome::Progress); + } + delivered += 1; + if delivered >= MAX_DRAIN_BATCHES_PER_LOCK { + self.dbg.progress_dispatches += 1; + return Ok(StepOutcome::Progress); + } + } + NextBatch::Stalled(partial) => { + // INPUT-STARVED: the driver couldn't advance because the + // winning source's next block isn't decompressed yet. + self.dbg.stalls += 1; + if let Some(batch) = partial { + if let Err(unpushed) = ctx.outputs.push(batch) { + self.dbg.output_full += 1; + self.held.put(unpushed); + return Ok(StepOutcome::Progress); + } + delivered += 1; + } + return Ok(if delivered > 0 { + self.dbg.progress_dispatches += 1; + StepOutcome::Progress + } else { + // Pure under-utilization: this dispatch did nothing. + self.dbg.contention += 1; + StepOutcome::Contention + }); + } + NextBatch::Done(partial, merged) => { + if let Some(batch) = partial { + if let Err(unpushed) = ctx.outputs.push(batch) { + self.dbg.output_full += 1; + self.held.put(unpushed); + return Ok(StepOutcome::Progress); + } + delivered += 1; + } + log::info!("Sort merge complete: {merged} records merged"); + // INSTRUMENTATION (lever-2): is the serial merge starved on + // decompress (stalls/contention high) or blocked on the + // writer (output_full high)? `stalls` counts merge-loop + // passes that ended input-starved; `contention` counts + // dispatches that produced nothing (pure idle spin); + // `output_full` counts downstream-backpressure events. + let d = self.dbg; + log::info!( + "Sort merge diag: stalls={} contention={} output_full={} \ + progress_dispatches={} ({} records, {} sources)", + d.stalls, + d.contention, + d.output_full, + d.progress_dispatches, + merged, + self.chunk_count, + ); + if let Some(slot) = &self.stats_slot { + *slot.lock() = Some(fgumi_sort::SortStats { + total_records: self.processed, + output_records: merged, + chunks_written: self.chunk_count, + }); + } + self.state = SortMergeState::Done; + return Ok(if delivered > 0 { + StepOutcome::Progress + } else { + StepOutcome::NoProgress + }); + } + } + } + } + + /// Build the next output batch for the single-source fast path: gather records + /// from the sorted chunk into the builder until the count/byte cap, or `Done` + /// when the chunk is exhausted. Mirrors [`next_batch`](Self::next_batch)'s + /// framing and buffer-sizing exactly, so the output is byte-identical to a + /// (k = 1) loser-tree merge of the same chunk — only the record SOURCE differs + /// (a direct cursor instead of `driver.try_step()`). + /// + /// # Panics + /// + /// Panics if `self.state` is not `FastPath`. + fn next_fast_batch(&mut self) -> io::Result> { + let target = self.target_batch_count; + let byte_limit = self.output_byte_limit; + let bytes_cap = usize::try_from(byte_limit).unwrap_or(usize::MAX); + let buffer_floor = INITIAL_OUTPUT_BUFFER_BYTES.min(bytes_cap); + let SortMergeState::FastPath { chunk, cursor, total, builder, next_ordinal } = + &mut self.state + else { + unreachable!("next_fast_batch called outside FastPath state"); + }; + + let flush = |builder: &mut O::Builder, next_ordinal: &mut u64| { + *next_ordinal += 1; + let hint = builder.total_bytes().clamp(buffer_floor, bytes_cap); + let next_builder = O::Builder::with_capacity(*next_ordinal, hint, target); + std::mem::replace(builder, next_builder).build() + }; + let flush_partial = |builder: &mut O::Builder, next_ordinal: &mut u64| { + if builder.is_empty() { None } else { Some(flush(builder, next_ordinal)) } + }; + + loop { + if *cursor >= *total { + return Ok(NextBatch::Done(flush_partial(builder, next_ordinal), *total as u64)); + } + builder.push_record_bytes(chunk.record_bytes(*cursor))?; + *cursor += 1; + let count_full = builder.len() >= target; + let bytes_full = (builder.total_bytes() as u64) >= byte_limit; + if count_full || bytes_full { + return Ok(NextBatch::Batch(flush(builder, next_ordinal))); + } + } + } + + /// Cooperative emit loop for the single-source fast path. Mirrors + /// [`emit_batches_cooperative`](Self::emit_batches_cooperative) but never + /// `Stalled` (every record is already in memory). + /// + /// # Panics + /// + /// Panics if `self.state` is not `FastPath`. + fn emit_fast_batches(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + let mut delivered = 0usize; + loop { + match self.next_fast_batch()? { + NextBatch::Batch(batch) => { + if let Err(unpushed) = ctx.outputs.push(batch) { + self.dbg.output_full += 1; + self.held.put(unpushed); + return Ok(StepOutcome::Progress); + } + delivered += 1; + if delivered >= MAX_DRAIN_BATCHES_PER_LOCK { + self.dbg.progress_dispatches += 1; + return Ok(StepOutcome::Progress); + } + } + NextBatch::Stalled(_) => unreachable!("FastPath never stalls (all in memory)"), + NextBatch::Done(partial, count) => { + if let Some(batch) = partial { + if let Err(unpushed) = ctx.outputs.push(batch) { + self.dbg.output_full += 1; + self.held.put(unpushed); + return Ok(StepOutcome::Progress); + } + delivered += 1; + } + log::info!( + "Sort in-memory fast path complete: {count} records (single source, \ + no merge)" + ); + if let Some(slot) = &self.stats_slot { + *slot.lock() = Some(fgumi_sort::SortStats { + total_records: self.processed, + output_records: count, + chunks_written: 0, + }); + } + self.state = SortMergeState::Done; + return Ok(if delivered > 0 { + StepOutcome::Progress + } else { + StepOutcome::NoProgress + }); + } + } + } + } + + fn transition_to_merging(&mut self) -> io::Result<()> { + if !matches!(&self.state, SortMergeState::WaitingForSetup { .. }) { + return Ok(()); + } + let SortMergeState::WaitingForSetup { + mut slots, + slot_index: _, + memory_chunks, + total_records, + expected_slot_count: _, + expected_memory_chunk_count: _, + } = std::mem::replace(&mut self.state, SortMergeState::Done) + else { + unreachable!("just matched WaitingForSetup") + }; + // Fail closed before consuming only the selected lane (fast path or + // build_driver): a chunk stranded in a non-selected lane would otherwise + // be dropped silently. + memory_chunks.ensure_single_lane(self.sort_order)?; + slots.sort_by_key(|s| s.file_id); + // Capture summary inputs before `slots` is consumed by the driver: + // total records ingested and the spill-file count. + self.processed = total_records; + self.chunk_count = slots.len(); + let bytes_cap_for_init = usize::try_from(self.output_byte_limit).unwrap_or(usize::MAX); + let initial_bytes_for_init = INITIAL_OUTPUT_BUFFER_BYTES.min(bytes_cap_for_init); + // FAST PATH: zero spill slots + exactly one in-memory chunk → the chunk is + // already globally sorted, so skip the (k = 1) loser-tree merge and gather + // it directly. This is the in-memory regime's dominant cost. + if slots.is_empty() && memory_chunks.total_len() == 1 { + let chunk = memory_chunks.into_single(); + let total = chunk.len(); + let builder = + O::Builder::with_capacity(0, initial_bytes_for_init, self.target_batch_count); + self.state = + SortMergeState::FastPath { chunk, cursor: 0, total, builder, next_ordinal: 0 }; + return Ok(()); + } + let driver = build_driver(self.sort_order, slots, memory_chunks, total_records)?; + let bytes_cap = usize::try_from(self.output_byte_limit).unwrap_or(usize::MAX); + // Seed the first buffer modestly; subsequent buffers are sized from the + // prior batch's actual byte length (see `next_batch`). + let initial_bytes = INITIAL_OUTPUT_BUFFER_BYTES.min(bytes_cap); + let builder = O::Builder::with_capacity(0, initial_bytes, self.target_batch_count); + self.state = SortMergeState::Merging { driver, builder, next_ordinal: 0 }; + Ok(()) + } +} + +impl Step for SortMerge { + type Input = SortPhase2Event; + type Outputs = OrderedBytesSingle; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "SortMerge", + // The merge runs off the work-stealing pool, on the sort's shared + // COORDINATION driver thread (N+2) — the same thread that ran the + // phase-1 admit/sort/frame steps, which have Finished and left the + // driver's live set by the time the phase-2 merge runs, so the merge + // effectively gets a dedicated thread in phase 2 (mirrors main's main + // thread). Its cooperative `try_run` body is UNCHANGED — + // `run_detached_driver` drives it with the same `run_worker_loop` the + // pool uses (Park backoff), parking on `Contention`/`NoProgress` + // (winner-slot momentarily empty / output full) instead of the pool + // re-dispatching it. `Detached` collapses the declared `ByItemOrdinal` + // output to `None` exactly as `Serial` would (see + // `effective_branch_orderings`), so the output transport — a direct + // byte-bounded queue, no reorder stage — is byte-for-byte identical; + // the LoserTree core, source order (`slots.sort_by_key(file_id)` + + // residual last), and tie-break are untouched. SortMerge is only ever + // built by the sort chain's `add_sort`, so this is sort-chain-only. + kind: StepKind::Detached, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], + branch_ordering: vec![BranchOrdering::ByItemOrdinal], + } + } + + fn detached_group(&self) -> DetachedGroup { + // Co-located with the phase-1 coordination steps on ONE driver thread — + // phase 1 and phase 2 are temporally disjoint, so this is the true N+2. + DetachedGroup::Shared(crate::sort::SORT_COORD_GROUP) + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if !self.flush_held(ctx) { + return Ok(StepOutcome::Contention); + } + + if matches!(&self.state, SortMergeState::WaitingForSetup { .. }) { + // Drain the input queue unbounded: the setup absorb is cheap (it just + // moves `Arc`s/`Vec`s into the setup state) and the upstream queue is + // already byte-bounded, so backpressure belongs on the producer, not + // on a consumer-side drain cap. A cap here would cycle a full upstream + // queue through repeated partial drains and add producer contention. + let absorbed = self.absorb_events_into_setup(ctx)?; + if !self.is_ready_to_merge() { + if absorbed > 0 { + return Ok(StepOutcome::Progress); + } + if !ctx.input.is_drained() { + return Ok(StepOutcome::NoProgress); + } + // Input is drained but setup never completed. Fail closed for any + // setup that saw payload or a (mismatched) `AllAnnounced`, so an + // incomplete setup can never silently merge a partial result. The + // only legitimate drained-but-not-ready case is a wholly empty + // input (no slots, no chunks, no announcement), which merges to an + // empty output. + let SortMergeState::WaitingForSetup { + slots, + memory_chunks, + expected_slot_count, + expected_memory_chunk_count, + .. + } = &self.state + else { + unreachable!("state matched WaitingForSetup above"); + }; + let saw_payload = !slots.is_empty() || memory_chunks.total_len() > 0; + let saw_expectations = + expected_slot_count.is_some() || expected_memory_chunk_count.is_some(); + if saw_payload || saw_expectations { + return Err(io::Error::other(format!( + "SortMerge: setup incomplete at input drain \ + (slots={}, chunks={}, expected_slots={expected_slot_count:?}, \ + expected_chunks={expected_memory_chunk_count:?})", + slots.len(), + memory_chunks.total_len(), + ))); + } + } + self.transition_to_merging()?; + } + + match &self.state { + SortMergeState::Merging { .. } => self.emit_batches_cooperative(ctx), + SortMergeState::FastPath { .. } => self.emit_fast_batches(ctx), + SortMergeState::Done => Ok(StepOutcome::Finished), + SortMergeState::WaitingForSetup { .. } => { + unreachable!("Phase 1 must have left state non-WaitingForSetup") + } + } + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/fgumi-pipeline-io/src/sort/merge/tests.rs b/crates/fgumi-pipeline-io/src/sort/merge/tests.rs new file mode 100644 index 000000000..a54937c9a --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/merge/tests.rs @@ -0,0 +1,97 @@ +// End-to-end merge behavior is covered by the integration tests in sort/tests.rs; +// the unit tests here pin the memory-lane fail-closed guard. + +use super::*; + +/// A residual chunk stranded in a lane that does not match the sort order must +/// fail closed: `build_driver` and the single-chunk fast path consume only the +/// selected lane, so such a chunk would otherwise be dropped silently while +/// `total_len()` still counted it toward setup completion. +#[test] +fn mismatched_memory_lane_fails_closed() { + let mut chunks = MemoryChunksByKind::default(); + let chunk = + InMemoryChunk::from_owned_records(vec![(RawCoordinateKey { sort_key: 1 }, vec![9u8; 8])]); + chunks.push(MemoryChunkErased::Coordinate(chunk)).expect("coordinate lane never mismatches"); + + // The coordinate lane matches a Coordinate sort → accepted. + chunks.ensure_single_lane(SortOrder::Coordinate).expect("matching lane is accepted"); + + // The same chunk under a Queryname sort is a lane mismatch → fail closed. + let err = chunks + .ensure_single_lane(SortOrder::Queryname(QuerynameComparator::Natural)) + .expect_err("stray coordinate chunk under a queryname sort must error"); + assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); +} + +/// Template-coordinate spill slots present but no residual chunk to identify the +/// `--key-types` lane must fail closed: defaulting to K40 would mis-decode narrow +/// (K24/Cb32/Tert32) spill files. Unreachable for valid input (Phase-1 always +/// emits a variant-tagged residual), so this guards against a seal-logic +/// regression. With no slots, the empty-input case is still accepted. +#[test] +fn empty_template_lane_with_spill_slots_fails_closed() { + let slot = Arc::new(SortMergeSlot::new( + 0, + std::io::BufReader::new(tempfile::tempfile().unwrap()), + fgumi_sort::SpillCodec::Bgzf, + )); + // `Box` isn't `Debug`, so match rather than `expect_err`. + match build_driver(SortOrder::TemplateCoordinate, vec![slot], MemoryChunksByKind::default(), 1) + { + Err(e) => assert_eq!(e.kind(), std::io::ErrorKind::InvalidData), + Ok(_) => panic!("empty template lane with spill slots must fail closed"), + } + + // No slots → empty input; any key width is safe (nothing to merge). + assert!( + build_driver(SortOrder::TemplateCoordinate, Vec::new(), MemoryChunksByKind::default(), 0) + .is_ok(), + "empty template lane with no slots is valid", + ); +} + +/// The `--key-types` narrowed-lane variant is chosen once per sort and is global +/// to the run. A template chunk arriving with a different variant means phase 1 +/// and the merge disagree about the key width; merging on would compare keys of +/// different layouts and emit silently mis-ordered output. It must fail closed, +/// like the sibling `ensure_single_lane` / `build_driver` violations — not panic. +#[test] +fn template_variant_change_mid_sort_fails_closed() { + use fgumi_sort::{TemplateKey24, TemplateMemChunk, TertKey32}; + + let mut chunks = MemoryChunksByKind::default(); + + let k24 = InMemoryChunk::from_owned_records(vec![(TemplateKey24::default(), vec![1u8; 8])]); + chunks + .push(MemoryChunkErased::TemplateCoordinate(TemplateMemChunk::K24(k24))) + .expect("the first template chunk establishes the variant"); + + // A second chunk in a different lane width is the protocol violation. + let tert = InMemoryChunk::from_owned_records(vec![(TertKey32::default(), vec![2u8; 8])]); + let err = chunks + .push(MemoryChunkErased::TemplateCoordinate(TemplateMemChunk::Tert32(tert))) + .expect_err("a variant change must be rejected"); + + assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); + let msg = err.to_string(); + assert!(msg.contains("variant changed mid-sort"), "unexpected message: {msg}"); + // Both variants are named so the failure is diagnosable from the log alone. + assert!(msg.contains("K24"), "error names the accumulated variant: {msg}"); + assert!(msg.contains("Tert32"), "error names the offending variant: {msg}"); +} + +/// Repeated chunks of the SAME variant are the normal path and must keep working. +#[test] +fn repeated_template_chunks_of_one_variant_accumulate() { + use fgumi_sort::{TemplateKey24, TemplateMemChunk}; + + let mut chunks = MemoryChunksByKind::default(); + for i in 0..3u8 { + let c = InMemoryChunk::from_owned_records(vec![(TemplateKey24::default(), vec![i; 8])]); + chunks + .push(MemoryChunkErased::TemplateCoordinate(TemplateMemChunk::K24(c))) + .expect("same-variant chunks accumulate"); + } + assert_eq!(chunks.total_len(), 3, "all three chunks are retained"); +} diff --git a/crates/fgumi-pipeline-io/src/sort/mod.rs b/crates/fgumi-pipeline-io/src/sort/mod.rs new file mode 100644 index 000000000..c79e5d23b --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/mod.rs @@ -0,0 +1,41 @@ +//! Sort typed-steps for the unified pipeline. + +/// `DetachedGroup::Shared` label for the sort's **coordination** driver thread: +/// the serial phase-1 coordination steps (`ReadBlocks` admit, `FindBoundariesAndSort` +/// sort/seal, `SpillGather` framing) plus the phase-2 `SortMerge`. One dedicated +/// thread runs all of them off the pool — the true N+2 model (mirrors main's +/// single main thread). Phase 1 and phase 2 are temporally disjoint, so the +/// coordination steps finish and leave the driver's live set before the merge +/// runs, giving the merge a dedicated thread in phase 2. +pub const SORT_COORD_GROUP: &str = "sort-coord"; + +/// `DetachedGroup::Shared` label for the sort's **I/O writer** driver thread: +/// `SpillWrite` (phase 1) and `WriteBgzfFile` (phase 2). Isolated from the +/// coordination driver so a write flush never stalls coordination (main's reason +/// for the second dedicated thread). +pub const SORT_IO_GROUP: &str = "sort-io"; + +pub mod arena_ingest; +pub mod compress_spill; +pub mod merge; +pub mod protocol; +pub mod sort_buffer; +pub mod spill_block_compress; +pub mod spill_decompress; +pub mod spill_gather; +pub mod spill_write; + +pub use arena_ingest::{ + ArenaBlock, ArenaSortStrategy, CoordinateStrategy, FindBoundariesAndSort, InflateToArena, + InflatedBlock, QuerynameStrategy, ReadBlocks, TemplateStrategy, +}; +pub use compress_spill::CompressSpill; +pub use merge::{BlockOutput, MergeBatchBuilder, MergeOutput, RecordBatchOutput, SortMerge}; +pub use sort_buffer::SortBuffer; +pub use spill_block_compress::SpillBlockCompress; +pub use spill_decompress::{SortDecompressTuning, SortSpillDecompress}; +pub use spill_gather::SpillGather; +pub use spill_write::SpillWrite; + +#[cfg(test)] +pub mod tests; diff --git a/crates/fgumi-pipeline-io/src/sort/protocol.rs b/crates/fgumi-pipeline-io/src/sort/protocol.rs new file mode 100644 index 000000000..1fd8be082 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/protocol.rs @@ -0,0 +1,477 @@ +//! Typed-event protocol between the three sort steps in the runall-sort +//! fused chain. + +use std::path::PathBuf; +use std::sync::Arc; + +use fgumi_sort::{ + InMemoryChunk, RawCoordinateKey, RawQuerynameKey, RawQuerynameLexKey, SortMergeSlot, + TemplateMemChunk, +}; + +use fgumi_pipeline_core::item::HeapSize; + +/// Approximate fixed per-record index overhead of an arena-backed +/// [`InMemoryChunk`], in bytes, added once per record in +/// [`MemoryChunkErased::approx_heap_bytes`] on top of the variable record payload +/// (`InMemoryChunk::payload_bytes`). +/// +/// This covers one `(K, offset, len)` index slot — the sort key `K` (largest +/// variant: `TemplateKey`) plus the offset/len into the shared buffer — together +/// with allocator-bucket slack the payload byte count does not capture. It is an +/// intentionally conservative constant, not a `size_of` expression, so the queue +/// accounting over-counts rather than under-counts memory. +const PER_MEMORY_RECORD_OVERHEAD: usize = 354; + +/// In-memory sorted residual chunk produced by the Phase-1 sort head +/// (`SortBuffer` or `FindBoundariesAndSort`), type-erased over the sort-key +/// variant `K`. +pub enum MemoryChunkErased { + /// Coordinate-sort residual. `K = RawCoordinateKey`. Zero-copy arena-backed + /// chunk (shares the sort buffer's `Arc`). + Coordinate(InMemoryChunk), + /// Queryname-sort residual, lexicographic comparator. `K = RawQuerynameLexKey`. + /// Arena-backed like [`Coordinate`](Self::Coordinate): the record bodies live + /// in the chunk's shared buffer; the key owns its (small) name bytes. + QuerynameLex(InMemoryChunk), + /// Queryname-sort residual, natural comparator. `K = RawQuerynameKey`. + QuerynameNatural(InMemoryChunk), + /// Template-coordinate-sort residual, carried as a variant-tagged + /// [`TemplateMemChunk`] so the chosen `--key-types` narrow lane rides through + /// merge and spill like every other order's `K`. Zero-copy arena-backed chunk + /// (shares the sort buffer's `Arc`), like + /// [`Coordinate`](Self::Coordinate). + TemplateCoordinate(TemplateMemChunk), +} + +impl MemoryChunkErased { + /// Number of records in this chunk. + #[must_use] + pub fn len(&self) -> usize { + match self { + Self::Coordinate(v) => v.len(), + Self::QuerynameLex(v) => v.len(), + Self::QuerynameNatural(v) => v.len(), + Self::TemplateCoordinate(v) => v.len(), + } + } + + /// `true` iff the chunk has zero records. + #[must_use] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Borrow the `i`th record's raw BAM body bytes, in this chunk's sorted order. + /// + /// Zero-copy for every variant — each is arena-backed and slices its shared + /// `SegmentedBuf`. Used by `SortMerge`'s single-source fast path to gather a + /// fully-sorted in-memory chunk into output blocks without a (k = 1) + /// loser-tree merge. + /// + /// # Panics + /// + /// Panics if `i >= self.len()`. + #[must_use] + pub fn record_bytes(&self, i: usize) -> &[u8] { + match self { + Self::Coordinate(v) => v.record_bytes(i), + Self::QuerynameLex(v) => v.record_bytes(i), + Self::QuerynameNatural(v) => v.record_bytes(i), + Self::TemplateCoordinate(v) => v.record_bytes(i), + } + } + + /// Approximate heap footprint in bytes. + #[must_use] + pub fn approx_heap_bytes(&self) -> usize { + let (count, payload): (usize, usize) = match self { + Self::Coordinate(v) => (v.len(), v.payload_bytes()), + Self::QuerynameLex(v) => (v.len(), v.payload_bytes()), + Self::QuerynameNatural(v) => (v.len(), v.payload_bytes()), + Self::TemplateCoordinate(v) => (v.len(), v.payload_bytes()), + }; + count * PER_MEMORY_RECORD_OVERHEAD + payload + } +} + +/// Events from `SortBuffer` → `CompressSpill` (the P6 Phase-1 split). +/// +/// `SortBuffer` (Serial) sorts each filled buffer and emits the sorted records +/// as a chunk; `CompressSpill` (Parallel) then compresses spill chunks to disk +/// or passes the in-memory residual through. The seam carries already-sorted +/// `MemoryChunkErased`s, so it is byte-bounded by [`HeapSize`] just like the +/// downstream `SortPhase1Event` queue. +pub enum SortChunkEvent { + /// A sorted chunk to be compressed and written to disk by `CompressSpill`. + /// + /// `seq` is the chunk's logical spill index, assigned monotonically by + /// `SortBuffer`. `CompressSpill` uses it as the opened slot's `file_id` so + /// the merge tie-break order matches the legacy spill order regardless of + /// which Parallel worker writes the file. + Spill { seq: u32, chunk: MemoryChunkErased, records_ingested_so_far: u64 }, + /// A sorted in-memory residual chunk to pass straight through as a + /// `SortPhase1Event::MemoryChunk` — no disk round-trip (the fast path). + Residual { chunk: MemoryChunkErased, records_ingested_so_far: u64 }, + /// Terminal sentinel carrying the final counts (number of `Spill` chunks = + /// `slot_count`, number of `Residual` chunks = `memory_chunk_count`). + /// `CompressSpill` forwards it verbatim as `SortPhase1Event::AllAnnounced`. + AllAnnounced { slot_count: u32, memory_chunk_count: u32, total_records: u64 }, +} + +impl HeapSize for SortChunkEvent { + fn heap_size(&self) -> usize { + // Mirror `SortPhase1Event::heap_size`: a fixed per-event base so the + // byte-bounded transport queue cannot absorb an unbounded count of + // near-zero-cost control events (`AllAnnounced`). + let base = std::mem::size_of::(); + match self { + Self::Spill { chunk, .. } | Self::Residual { chunk, .. } => { + base + chunk.approx_heap_bytes() + } + Self::AllAnnounced { .. } => base, + } + } +} + +impl SortChunkEvent { + /// Running snapshot of records ingested at the moment this event was emitted. + #[must_use] + pub fn records_ingested_so_far(&self) -> u64 { + match self { + Self::Spill { records_ingested_so_far, .. } + | Self::Residual { records_ingested_so_far, .. } => *records_ingested_so_far, + Self::AllAnnounced { total_records, .. } => *total_records, + } + } +} + +/// Events from `SpillGather` → `SpillBlockCompress` → `SpillWrite` (the +/// block-parallel spill-write split that replaces the monolithic single-worker +/// `CompressSpill`). +/// +/// `SpillGather` (Serial) fans each `SortChunkEvent::Spill` chunk into +/// record-aligned raw [`Block`](Self::Block)s and forwards `Residual` / +/// `AllAnnounced`. `SpillBlockCompress` (Parallel) compresses each `Block`'s `bytes` +/// in place. `SpillWrite` (Serial) demultiplexes blocks back to per-`file_id` +/// spill files and emits the existing [`SortPhase1Event`]. +/// +/// **Ordinal contract:** `SpillGather` mints `ordinal` monotonically across +/// **every** emitted item (every block of every file, plus the `Residual` / +/// `AllAnnounced` passthroughs), so the stream is dense and gap-free for the +/// framework's single-cursor `ByItemOrdinal` reorder. Because `SortBuffer` +/// (Serial) emits `Spill` events one-at-a-time in `seq` order and +/// `SpillGather` (Serial) drains them in order, each file's blocks are +/// **contiguous** in the ordinal stream — so `SpillWrite` only ever has one +/// spill file open at a time. +pub enum SpillBlockEvent { + /// One raw (pre-compression) or compressed (post-`SpillBlockCompress`) block of a + /// spill file. `file_id` is the spill `seq` (the eventual slot `file_id`), + /// `is_last_in_file` marks the final block so `SpillWrite` can finalize the + /// file and emit `SpillReady`. `SpillWrite` (Serial) owns the path allocation, + /// so the block carries no path — only the `file_id` that names the file. + Block { + ordinal: u64, + file_id: u32, + is_last_in_file: bool, + records_ingested_so_far: u64, + bytes: Vec, + }, + /// A sorted in-memory residual chunk, passed straight through to + /// `SortPhase1Event::MemoryChunk` (no disk round-trip). + Residual { ordinal: u64, chunk: MemoryChunkErased, records_ingested_so_far: u64 }, + /// Terminal sentinel forwarded verbatim as `SortPhase1Event::AllAnnounced`. + AllAnnounced { ordinal: u64, slot_count: u32, memory_chunk_count: u32, total_records: u64 }, +} + +impl SpillBlockEvent { + /// The dense ordinal that drives `ByItemOrdinal` reordering into `SpillWrite`. + #[must_use] + pub fn ordinal(&self) -> u64 { + match self { + Self::Block { ordinal, .. } + | Self::Residual { ordinal, .. } + | Self::AllAnnounced { ordinal, .. } => *ordinal, + } + } +} + +impl HeapSize for SpillBlockEvent { + fn heap_size(&self) -> usize { + // Fixed per-event base (so the byte-bounded queue can't absorb an + // unbounded count of near-empty control events) plus the variable + // payload: a block's bytes, or a residual chunk's records. + let base = std::mem::size_of::(); + match self { + Self::Block { bytes, .. } => base + bytes.capacity(), + Self::Residual { chunk, .. } => base + chunk.approx_heap_bytes(), + Self::AllAnnounced { .. } => base, + } + } +} + +impl fgumi_pipeline_core::item::Ordered for SpillBlockEvent { + fn ordinal(&self) -> u64 { + SpillBlockEvent::ordinal(self) + } +} + +/// Events from the Phase-1 producers (`CompressSpill` / `SpillWrite`) → +/// `SortSpillDecompress`. +pub enum SortPhase1Event { + /// A spill chunk file has been closed and is ready for Phase 2 decompression. + SpillReady { slot: Arc, path: PathBuf, records_ingested_so_far: u64 }, + /// A par-sorted residual in-memory chunk passed through by the Phase-1 + /// producer (`CompressSpill` / `SpillWrite`) on its drained-completion path. + MemoryChunk { chunk: Arc, records_ingested_so_far: u64 }, + /// Sentinel emitted as the LAST event by the Phase-1 producer + /// (`CompressSpill` / `SpillWrite`) on its drained-completion path. + AllAnnounced { slot_count: u32, memory_chunk_count: u32, total_records: u64 }, +} + +/// Events from `SortSpillDecompress` → `SortMerge`. +pub enum SortPhase2Event { + /// Forwarded `SortPhase1Event::SpillReady`. + SpillReady { slot: Arc, path: PathBuf, records_ingested_so_far: u64 }, + /// Forwarded `SortPhase1Event::MemoryChunk`. + MemoryChunk { chunk: Arc, records_ingested_so_far: u64 }, + /// Forwarded `SortPhase1Event::AllAnnounced`. + AllAnnounced { slot_count: u32, memory_chunk_count: u32, total_records: u64 }, +} + +/// Generates the identical `HeapSize` impl and `records_ingested_so_far` accessor for a +/// spill-phase event enum. +/// +/// [`SortPhase1Event`] and [`SortPhase2Event`] are structurally identical (the Phase-2 +/// event is a verbatim forward of the Phase-1 event) but are kept as distinct types so +/// the typed-step pipeline cannot wire a Phase-1 producer output straight into a +/// Phase-2 (`SortMerge`) input. This macro removes the duplicated impl bodies without +/// collapsing the two types. +macro_rules! impl_spill_phase_event { + ($ty:ty) => { + impl HeapSize for $ty { + fn heap_size(&self) -> usize { + // Charge a fixed per-event base so the byte-bounded transport queues + // cannot accept an unbounded count of near-zero-cost control events + // (`SpillReady` with an empty path, `AllAnnounced`). Memory stays a + // function of configuration rather than event count. + let base = std::mem::size_of::(); + match self { + Self::SpillReady { path, .. } => base + path.as_os_str().len(), + Self::MemoryChunk { chunk, .. } => base + chunk.approx_heap_bytes(), + Self::AllAnnounced { .. } => base, + } + } + } + + impl $ty { + /// Running snapshot of records ingested at the moment this event was emitted. + #[must_use] + pub fn records_ingested_so_far(&self) -> u64 { + match self { + Self::SpillReady { records_ingested_so_far, .. } + | Self::MemoryChunk { records_ingested_so_far, .. } => *records_ingested_so_far, + Self::AllAnnounced { total_records, .. } => *total_records, + } + } + } + }; +} + +impl_spill_phase_event!(SortPhase1Event); +impl_spill_phase_event!(SortPhase2Event); + +#[cfg(test)] +mod tests { + use super::*; + + /// Build the arena-backed coordinate chunk the `Coordinate` variant carries, + /// from raw record payloads (keys are `default()`; irrelevant to these tests). + fn coord(payloads: Vec>) -> InMemoryChunk { + InMemoryChunk::from_owned_records( + payloads.into_iter().map(|b| (RawCoordinateKey::default(), b)).collect(), + ) + } + + #[test] + fn memory_chunk_len_and_is_empty() { + let chunk: MemoryChunkErased = MemoryChunkErased::Coordinate(coord(Vec::new())); + assert_eq!(chunk.len(), 0); + assert!(chunk.is_empty()); + + let chunk = MemoryChunkErased::Coordinate(coord(vec![vec![0xAA; 16], vec![0xBB; 32]])); + assert_eq!(chunk.len(), 2); + assert!(!chunk.is_empty()); + } + + #[test] + fn memory_chunk_approx_heap_bytes_counts_overhead_plus_payload() { + let chunk = MemoryChunkErased::Coordinate(coord(vec![vec![0xAA; 100], vec![0xBB; 200]])); + assert_eq!(chunk.approx_heap_bytes(), 2 * PER_MEMORY_RECORD_OVERHEAD + 300); + } + + #[test] + fn records_ingested_so_far_accessors() { + let dummy_path = std::path::PathBuf::from("/tmp/x"); + let slot = Arc::new(SortMergeSlot::new( + 0, + std::io::BufReader::new(tempfile::tempfile().unwrap()), + fgumi_sort::SpillCodec::Bgzf, + )); + let ev1 = SortPhase1Event::SpillReady { + slot: Arc::clone(&slot), + path: dummy_path.clone(), + records_ingested_so_far: 100, + }; + assert_eq!(ev1.records_ingested_so_far(), 100); + + let chunk = Arc::new(MemoryChunkErased::Coordinate(coord(Vec::new()))); + let ev2 = SortPhase1Event::MemoryChunk { + chunk: Arc::clone(&chunk), + records_ingested_so_far: 250, + }; + assert_eq!(ev2.records_ingested_so_far(), 250); + + let ev3 = + SortPhase2Event::SpillReady { slot, path: dummy_path, records_ingested_so_far: 300 }; + assert_eq!(ev3.records_ingested_so_far(), 300); + + let ev4 = SortPhase2Event::MemoryChunk { chunk, records_ingested_so_far: 400 }; + assert_eq!(ev4.records_ingested_so_far(), 400); + + let ev5 = SortPhase1Event::AllAnnounced { + slot_count: 4, + memory_chunk_count: 1, + total_records: 500, + }; + assert_eq!(ev5.records_ingested_so_far(), 500); + let ev6 = SortPhase2Event::AllAnnounced { + slot_count: 4, + memory_chunk_count: 1, + total_records: 600, + }; + assert_eq!(ev6.records_ingested_so_far(), 600); + } + + #[test] + fn sort_chunk_event_heap_size_and_accessors() { + let chunk = MemoryChunkErased::Coordinate(coord(vec![vec![0xAA; 100], vec![0xBB; 200]])); + let payload = chunk.approx_heap_bytes(); + + let spill = SortChunkEvent::Spill { seq: 3, chunk, records_ingested_so_far: 42 }; + assert_eq!(spill.heap_size(), std::mem::size_of::() + payload); + assert_eq!(spill.records_ingested_so_far(), 42); + + let residual = SortChunkEvent::Residual { + chunk: MemoryChunkErased::Coordinate(coord(Vec::new())), + records_ingested_so_far: 7, + }; + assert_eq!(residual.records_ingested_so_far(), 7); + + // Control events carry no heap payload but still cost a fixed base so the + // byte-bounded queue cannot accept an unbounded count of them. + let announced = SortChunkEvent::AllAnnounced { + slot_count: 4, + memory_chunk_count: 1, + total_records: 500, + }; + assert_eq!(announced.heap_size(), std::mem::size_of::()); + assert_eq!(announced.records_ingested_so_far(), 500); + } + + #[test] + fn all_announced_heap_size_charges_base_cost() { + // Control events carry no heap payload but must still cost a fixed, + // non-zero amount so the byte-bounded queues cannot accept an unbounded + // count of them. + let ev1 = SortPhase1Event::AllAnnounced { + slot_count: 16, + memory_chunk_count: 4, + total_records: 1_000_000, + }; + assert_eq!(ev1.heap_size(), std::mem::size_of::()); + let ev2 = SortPhase2Event::AllAnnounced { + slot_count: 16, + memory_chunk_count: 4, + total_records: 1_000_000, + }; + assert_eq!(ev2.heap_size(), std::mem::size_of::()); + } + + // ── SpillBlockEvent ───────────────────────────────────────────────────── + + /// The `Ordered` impl delegates to the inherent `SpillBlockEvent::ordinal`. + /// If that inherent method is ever removed or renamed, the call silently + /// resolves to the trait method itself and recurses until the stack blows. + /// Asserting through the trait for every variant pins the delegation. + #[test] + fn spill_block_event_ordinal_delegates_for_every_variant() { + use fgumi_pipeline_core::item::Ordered; + + let block = SpillBlockEvent::Block { + ordinal: 3, + file_id: 0, + is_last_in_file: false, + records_ingested_so_far: 0, + bytes: vec![0u8; 4], + }; + let residual = SpillBlockEvent::Residual { + ordinal: 4, + chunk: MemoryChunkErased::Coordinate(coord(vec![vec![1u8; 8]])), + records_ingested_so_far: 1, + }; + let announced = SpillBlockEvent::AllAnnounced { + ordinal: 5, + slot_count: 1, + memory_chunk_count: 1, + total_records: 1, + }; + + assert_eq!(::ordinal(&block), 3); + assert_eq!(::ordinal(&residual), 4); + assert_eq!(::ordinal(&announced), 5); + } + + /// `heap_size` charges a fixed per-event base plus the variable payload. The + /// base is what stops a byte-bounded queue absorbing an unbounded number of + /// near-empty control events, so a `Block` must scale with its bytes while + /// `AllAnnounced` stays at the base. + #[test] + fn spill_block_event_heap_size_charges_base_plus_payload() { + let base = std::mem::size_of::(); + + let announced = SpillBlockEvent::AllAnnounced { + ordinal: 0, + slot_count: 1, + memory_chunk_count: 0, + total_records: 0, + }; + assert_eq!(announced.heap_size(), base, "a control event costs only the base"); + + let small = SpillBlockEvent::Block { + ordinal: 0, + file_id: 0, + is_last_in_file: false, + records_ingested_so_far: 0, + bytes: Vec::with_capacity(64), + }; + let large = SpillBlockEvent::Block { + ordinal: 0, + file_id: 0, + is_last_in_file: false, + records_ingested_so_far: 0, + bytes: Vec::with_capacity(4096), + }; + assert_eq!(small.heap_size(), base + 64, "a block charges its byte capacity"); + assert_eq!(large.heap_size(), base + 4096); + assert!(large.heap_size() > small.heap_size(), "cost tracks payload size"); + + // A residual charges the chunk's records, so it too exceeds the base. + let residual = SpillBlockEvent::Residual { + ordinal: 0, + chunk: MemoryChunkErased::Coordinate(coord(vec![vec![7u8; 128]])), + records_ingested_so_far: 1, + }; + assert!(residual.heap_size() > base, "a residual charges its retained records"); + } +} diff --git a/crates/fgumi-pipeline-io/src/sort/sort_buffer.rs b/crates/fgumi-pipeline-io/src/sort/sort_buffer.rs new file mode 100644 index 000000000..7b97ff0a4 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/sort_buffer.rs @@ -0,0 +1,667 @@ +//! `SortBuffer` — first step of the P6 Phase-1 split (`SortBuffer` → +//! `CompressSpill` → `SortSpillDecompress` → `SortMerge`). +//! +//! `SortBuffer` (`Serial`) ingests `RecordBatch`es into an in-memory arena via +//! the order-erased `ChunkSorter` (each variant drives the same per-order +//! [`ArenaSortStrategy`] as the block-input `FindBoundariesAndSort` front), +//! sorts each filled arena, and emits the sorted records as a [`SortChunkEvent`] +//! — **without touching disk**. Mid-stream spill +//! chunks (`Spill { seq, .. }`) and the final residual (`Residual`) flow to the +//! `Parallel` `CompressSpill` step, which compresses + writes spills and passes +//! the residual through. This replaces the monolithic `SortAndSpill`, which +//! drove a private `SortWorkerPool` to compress inline. +//! +//! Spill chunks are emitted **as the buffer fills** (not accumulated to the +//! end), and `try_run` drains the staged-events queue before popping the next +//! input batch, so staged chunks never accumulate *across* batches. Within a +//! single batch, `ingest_one_batch` seals one chunk each time the arena reaches +//! `memory_limit`. In production this fires at most once per batch — each input +//! `RecordBatch` is block-bounded (`ParseBamRecords` emits one batch per +//! decompressed BGZF block, orders of magnitude below the 512 MB default +//! `memory_limit`) — so peak memory stays at ~one spill chunk plus the live +//! buffer. If `memory_limit` is configured far below the batch size (as some +//! tests do to force spilling), a single batch may seal several chunks into +//! `pending` in one `ingest_one_batch` call; this only raises transient peak +//! memory — every sealed chunk is still emitted in order and no records are +//! dropped. +//! +//! All four sort orders (coordinate, template-coordinate, queryname lex + +//! natural) route through this step via the `ChunkSorter` order enum. + +use std::collections::VecDeque; +use std::io; +use std::sync::Arc; + +use anyhow::{Result, anyhow}; +use fgumi_bam_io::ProgressTracker; +use fgumi_sort::{ + PooledSegmentedBuf, QuerynameComparator, RawExternalSorter, RawQuerynameKey, + RawQuerynameLexKey, SegmentedBuf, SortOrder, TemplateArenaAccumulator, +}; +use noodles::sam::Header; + +use crate::sort::protocol::{MemoryChunkErased, SortChunkEvent}; +use crate::sort::{ArenaSortStrategy, CoordinateStrategy, QuerynameStrategy, TemplateStrategy}; +use crate::types::RecordBatch; +use fgumi_pipeline_core::{ + HeldRetry, Unpushed, + held::HeldSlot, + outputs::Single, + queues::QueueSpec, + reorder::BranchOrdering, + step::{Affinity, Step, StepCtx, StepKind, StepOutcome, StepProfile}, +}; + +/// Max staged events flushed to the output per `try_run` invocation. +const MAX_EVENTS_PER_LOCK: usize = 8; + +/// Per-record memory overhead added to the arena byte count so the seal +/// threshold accounts for the strategy's `(key, offset, len)` ref alongside the +/// record body. Conservative; the exact value only shifts spill boundaries, not +/// the merged output (coordinate/template are globally stable; queryname's tie +/// order is unspecified). +const PER_RECORD_REF_OVERHEAD: usize = 64; + +/// Record-input arena accumulator: copies each pushed record into a growing +/// [`SegmentedBuf`] and drives an [`ArenaSortStrategy`] over the arena refs, so +/// the record-input (SAM / fused) path uses the SAME per-order sort engine as the +/// block-input arena front ([`FindBoundariesAndSort`](crate::sort::FindBoundariesAndSort)). +/// At seal the filled arena is wrapped in an `Arc` and handed to the strategy (no +/// further record copies), and a fresh arena starts the next run. +struct ArenaAccum { + strategy: S, + arena: SegmentedBuf, + memory_limit: usize, + memory_used: usize, + total_records: u64, + sort_threads: usize, +} + +impl ArenaAccum { + fn new(strategy: S, memory_limit: usize, sort_threads: usize) -> Self { + Self { + strategy, + arena: SegmentedBuf::new(), + memory_limit, + memory_used: 0, + total_records: 0, + sort_threads, + } + } + + /// Copy one record into the arena and accumulate its sort ref. Returns `true` + /// once the run's byte budget is reached. + #[allow(clippy::cast_possible_truncation)] // offset/len fit usize on all supported (64-bit) targets + fn push(&mut self, bam_bytes: &[u8]) -> Result { + let len = + u32::try_from(bam_bytes.len()).map_err(|_| anyhow!("record length exceeds u32"))?; + let offset = self.arena.extend_from_slice(bam_bytes) as u64; + let body = self.arena.slice(offset as usize, len as usize); + // The strategy reads `body` to extract the sort key and stores only + // `(key, offset, len)`; it does not retain `body`, so the arena may grow + // (and this slice's borrow end) freely afterwards. + self.strategy.push_record(body, offset, len).map_err(|e| anyhow!("{e:#}"))?; + self.memory_used += bam_bytes.len() + PER_RECORD_REF_OVERHEAD; + self.total_records += 1; + Ok(self.memory_used >= self.memory_limit) + } + + /// Seal the current run: wrap the filled arena in an `Arc`, sort + materialize + /// via the strategy, and reset for the next run. Empty if nothing was pushed + /// since the last seal. + fn take_sorted_chunk(&mut self) -> MemoryChunkErased { + let arena = std::mem::replace(&mut self.arena, SegmentedBuf::new()); + self.memory_used = 0; + let arc = Arc::new(PooledSegmentedBuf::unpooled(arena)); + self.strategy.seal(arc, self.sort_threads) + } + + fn total_records(&self) -> u64 { + self.total_records + } +} + +/// Order-erased record-input arena sorter. Each variant pairs an [`ArenaAccum`] +/// with the concrete [`ArenaSortStrategy`] for its order, so `SortBuffer` drives +/// the same per-order sort engine as the block-input `FindBoundariesAndSort`. +enum ChunkSorter { + Coordinate(ArenaAccum), + Template(ArenaAccum), + QuerynameLex(ArenaAccum>), + QuerynameNatural(ArenaAccum>), +} + +impl ChunkSorter { + /// Build the arena sorter matching `sorter.sort_order()`, provisioning each + /// order's strategy exactly as the block-input arena front does. + #[allow(clippy::needless_pass_by_value)] // by-value keeps the caller's move-in; only read here + fn from_sorter(sorter: RawExternalSorter, header: &Header) -> Result { + let memory_limit = sorter.memory_limit_bytes(); + // Phase-1 count honors the `--sort-threads` override (falls back to + // `--threads`); `num_threads()` would drop the override silently. + let sort_threads = sorter.phase1_threads(); + Ok(match sorter.sort_order() { + SortOrder::Coordinate => { + let n_ref = u32::try_from(header.reference_sequences().len()) + .map_err(|_| anyhow!("reference sequence count overflows u32"))?; + Self::Coordinate(ArenaAccum::new( + CoordinateStrategy::new(n_ref), + memory_limit, + sort_threads, + )) + } + SortOrder::TemplateCoordinate => { + let acc = TemplateArenaAccumulator::from_header( + header, + sorter.cell_tag_value(), + sorter.key_types_spec(), + ); + Self::Template(ArenaAccum::new( + TemplateStrategy::new(acc), + memory_limit, + sort_threads, + )) + } + SortOrder::Queryname(QuerynameComparator::Lexicographic) => { + Self::QuerynameLex(ArenaAccum::new( + QuerynameStrategy::new(MemoryChunkErased::QuerynameLex), + memory_limit, + sort_threads, + )) + } + SortOrder::Queryname(QuerynameComparator::Natural) => { + Self::QuerynameNatural(ArenaAccum::new( + QuerynameStrategy::new(MemoryChunkErased::QuerynameNatural), + memory_limit, + sort_threads, + )) + } + }) + } + + /// Push one record; `true` once the run hits the memory limit. + fn push(&mut self, bam_bytes: &[u8]) -> Result { + match self { + Self::Coordinate(s) => s.push(bam_bytes), + Self::Template(s) => s.push(bam_bytes), + Self::QuerynameLex(s) => s.push(bam_bytes), + Self::QuerynameNatural(s) => s.push(bam_bytes), + } + } + + /// Seal + materialize the current run into one erased chunk (empty if nothing + /// was pushed), resetting for the next run. + fn take_sorted_chunk(&mut self) -> MemoryChunkErased { + match self { + Self::Coordinate(s) => s.take_sorted_chunk(), + Self::Template(s) => s.take_sorted_chunk(), + Self::QuerynameLex(s) => s.take_sorted_chunk(), + Self::QuerynameNatural(s) => s.take_sorted_chunk(), + } + } + + /// Take the final residual as zero-or-one erased chunk. Every order seals one + /// globally-sorted chunk per run (coordinate/template are stable; queryname's + /// tie order is unspecified), so the residual is at most one chunk — the + /// legacy multi-chunk `par_chunks_mut` split is no longer needed. + /// + /// An empty residual is normally dropped. The one EXCEPTION is + /// template-coordinate when there were prior spills: the empty chunk still + /// carries the chosen `--key-types` narrow-lane variant, which is the only + /// signal `SortMerge`'s `build_driver` has to pick the spill files' key width. + /// Without it, a run whose records all spilled (empty final residual) would + /// leave the merge to default to the full 40-byte key and misread the + /// narrow-key spills. Coordinate and queryname have a fixed key type, so their + /// empty residual is dropped as before; with no spills there is nothing to + /// disambiguate, so the fast path (single non-empty chunk) is preserved. + fn take_residual_chunks(&mut self, had_spills: bool) -> Vec { + residual_chunks_for(self.take_sorted_chunk(), had_spills) + } + + fn total_records(&self) -> u64 { + match self { + Self::Coordinate(s) => s.total_records(), + Self::Template(s) => s.total_records(), + Self::QuerynameLex(s) => s.total_records(), + Self::QuerynameNatural(s) => s.total_records(), + } + } +} + +/// Decide whether a run's final (residual) chunk is emitted. +/// +/// Split out of `ChunkSorter::take_residual_chunks` so the rule can be tested +/// without constructing a whole accumulator: the interesting behaviour is a pure +/// function of the chunk and whether the run spilled. +/// +/// An empty residual is normally dropped. The one EXCEPTION is +/// template-coordinate when there were prior spills: the empty chunk is the only +/// carrier of the chosen `--key-types` narrow-lane variant, which is the only +/// signal `SortMerge`'s `build_driver` has to pick the spill files' key width. +/// Dropping it leaves the merge defaulting to the full 40-byte key and misreading +/// the narrow-key spills — wrong output, not a crash. +fn residual_chunks_for(chunk: MemoryChunkErased, had_spills: bool) -> Vec { + let keep_empty_for_variant = + had_spills && matches!(chunk, MemoryChunkErased::TemplateCoordinate(_)); + if chunk.is_empty() && !keep_empty_for_variant { Vec::new() } else { vec![chunk] } +} + +/// Push every record in `batch` into `sorter`, staging a sealed `Spill` chunk +/// into `pending` (and bumping `next_seq`) each time the arena fills. Returns +/// the number of records traversed and the first push failure, if any. +/// +/// The sorter is BORROWED, so a failure cannot leave the caller's sorter slot +/// empty — a state that would be indistinguishable from "finalized". The error +/// is returned rather than propagated with `?` so the caller can record it and +/// still account for the records it already consumed. +/// +/// Ingest stops at the first failing record: the pushes that preceded it are +/// retained (they are already in the arena), but continuing would ingest records +/// *after* a rejected one and quietly sort a subset of the input. +fn ingest_batch_records( + sorter: &mut ChunkSorter, + batch: &RecordBatch, + pending: &mut VecDeque, + next_seq: &mut u32, +) -> (u64, Option) { + let mut batch_records = 0u64; + for record in batch.iter_record_bytes() { + batch_records += 1; + let buffer_full = match sorter.push(record) { + Ok(full) => full, + Err(e) => return (batch_records, Some(format!("SortBuffer: push failed: {e:#}"))), + }; + if buffer_full { + // Seal the filled arena and stage it. We keep draining the rest + // of the batch rather than stopping early — breaking here would + // strand the batch's remaining records. Staging multiple chunks + // per batch is correct (each is emitted in order downstream); in + // production a block-bounded `RecordBatch` is far below + // `memory_limit`, so this fires at most once per batch (see the + // module docs). + let chunk = sorter.take_sorted_chunk(); + if !chunk.is_empty() { + pending.push_back(SortChunkEvent::Spill { + seq: *next_seq, + chunk, + records_ingested_so_far: sorter.total_records(), + }); + *next_seq += 1; + } + } + } + (batch_records, None) +} + +/// `Serial` step that buffers, sorts, and emits sorted chunks for `CompressSpill`. +pub struct SortBuffer { + /// In-memory buffering sorter. `Some` while ingesting; `None` after the + /// residual has been taken (finalized). + sorter: Option, + /// Sorted chunks awaiting output (spill chunks during ingest; the residual + + /// `AllAnnounced` after finalize). Drained before the next batch is ingested. + pending: VecDeque, + /// Monotonic spill index. Each `Spill` event's `seq` (= the eventual slot + /// `file_id`) makes the merge tie-break order independent of which + /// `CompressSpill` worker writes the file. Also the final `slot_count`. + next_seq: u32, + /// First ingest failure, if any. Once set the step is poisoned and every + /// later `try_run` re-raises instead of ingesting or finalizing — a step + /// that failed must never go on to report `Finished`, which downstream + /// would read as a complete sort (see `try_run`). Holds the message rather + /// than the `io::Error` because `io::Error` is not `Clone`. + failed: Option, + held: HeldSlot>, + output_byte_limit: u64, + affinity: Affinity, + /// Read+ingest progress, logged every 1M records (mirrors the legacy + /// "Read records" tracker) so the ingest rate over time is visible under + /// `RUST_LOG=info` — distinguishes a slow read path from a stalled one. + ingest_progress: ProgressTracker, +} + +impl SortBuffer { + /// Build a `SortBuffer` from a configured `RawExternalSorter` (any of the + /// four sort orders) and the output `Header`. + /// + /// `output_byte_limit` byte-bounds the output event queue (its chunk-bearing + /// variants retain sorted records, so the queue budgets on bytes, not count). + /// + /// # Errors + /// + /// Returns an error if the header's reference-sequence count does not fit in + /// a `u32` (the coordinate key's reference field). That conversion is the + /// only fallible step: the template path's `TemplateArenaAccumulator::from_header` + /// is infallible here. + pub fn from_sorter( + sorter: RawExternalSorter, + header: &Header, + output_byte_limit: u64, + ) -> Result { + let chunk_sorter = ChunkSorter::from_sorter(sorter, header)?; + Ok(Self { + sorter: Some(chunk_sorter), + pending: VecDeque::new(), + next_seq: 0, + failed: None, + held: HeldSlot::new(), + output_byte_limit, + affinity: Affinity::None, + ingest_progress: ProgressTracker::new("Sort ingest records").with_interval(1_000_000), + }) + } + + /// Override the affinity hint. + #[must_use] + pub fn with_affinity(mut self, affinity: Affinity) -> Self { + self.affinity = affinity; + self + } + + /// Re-raise the first ingest failure once the step has been poisoned. + /// + /// # Errors + /// + /// Returns an error naming the original failure if `failed` is set. + fn check_not_failed(&self) -> io::Result<()> { + match &self.failed { + None => Ok(()), + Some(message) => Err(io::Error::other(format!( + "SortBuffer: refusing to continue after an earlier failure: {message}" + ))), + } + } + + fn flush_held(&mut self, ctx: &mut StepCtx<'_, Self>) -> bool { + // `true` once the slot is clear (was empty, or the held event flushed); + // `false` while it's still held under backpressure. Uses the canonical + // re-hold helper so the put-back-on-reject invariant lives in one place. + !matches!(ctx.outputs.retry_held(&mut self.held), HeldRetry::StillHeld) + } + + /// Push up to `MAX_EVENTS_PER_LOCK` staged events to the output, parking the + /// first that can't be pushed in `held`. Caller guarantees `held` is empty. + fn emit_pending(&mut self, ctx: &mut StepCtx<'_, Self>) -> StepOutcome { + let mut emitted = 0usize; + while emitted < MAX_EVENTS_PER_LOCK { + let Some(event) = self.pending.pop_front() else { break }; + if let Err(unpushed) = ctx.outputs.push(event) { + self.held.put(unpushed); + return StepOutcome::Progress; + } + emitted += 1; + } + if emitted > 0 { StepOutcome::Progress } else { StepOutcome::NoProgress } + } + + /// Pop one input batch (if any) and push its records into the sorter, staging + /// a `Spill` chunk into `pending` whenever the buffer fills. Returns `true` + /// if a batch was consumed. + /// + /// # Errors + /// + /// Returns the first `ChunkSorter::push` failure, after poisoning the step so + /// no later `try_run` can finalize (see the `failed` field). + /// + /// # Panics + /// + /// Panics if called after finalize (`self.sorter` is `None`). + fn ingest_one_batch(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + let Some(batch) = ctx.input.pop() else { + return Ok(false); + }; + // Borrow the sorter rather than taking it: an error path that left + // `self.sorter` as `None` would be indistinguishable from "finalized". + // `pending` / `next_seq` are passed alongside as disjoint borrows. + let sorter = self.sorter.as_mut().expect("ingest_one_batch after finalize"); + let (batch_records, push_error) = + ingest_batch_records(sorter, &batch, &mut self.pending, &mut self.next_seq); + // Log ingest progress (every 1M records) so the read+ingest rate over + // wall time is visible — a steady rate means the read path is the cost; + // a bursty/stalling rate means downstream backpressure. + self.ingest_progress.log_if_needed(batch_records); + if let Some(message) = push_error { + self.failed = Some(message.clone()); + return Err(io::Error::other(message)); + } + Ok(true) + } + + /// Take the residual chunk and enqueue it (if non-empty) followed by the + /// terminal `AllAnnounced`. Consumes the sorter (`self.sorter` becomes + /// `None`), freeing the buffer + rayon pool. + /// + /// # Panics + /// + /// Panics if called twice (`self.sorter` already `None`). + fn finalize(&mut self) { + let mut sorter = self.sorter.take().expect("finalize called twice"); + let total_records = sorter.total_records(); + // `had_spills` keeps an EMPTY template-coordinate residual alive: it is + // the only carrier of the `--key-types` narrowed-lane variant, without + // which `SortMerge` falls back to the full 40-byte key and misreads the + // spills. See `take_residual_chunks` for the authoritative rule. + let residual_chunks = sorter.take_residual_chunks(self.next_seq > 0); + let memory_chunk_count = + u32::try_from(residual_chunks.len()).expect("residual chunk count fits u32"); + for chunk in residual_chunks { + self.pending.push_back(SortChunkEvent::Residual { + chunk, + records_ingested_so_far: total_records, + }); + } + self.pending.push_back(SortChunkEvent::AllAnnounced { + slot_count: self.next_seq, + memory_chunk_count, + total_records, + }); + // `sorter` dropped here (releases the buffer + private rayon pool). + } +} + +impl Step for SortBuffer { + type Input = RecordBatch; + type Outputs = Single; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "SortBuffer", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], + branch_ordering: vec![BranchOrdering::None], + } + } + + fn affinity(&self) -> Affinity { + self.affinity + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + // Fail closed. Today both runtimes stop dispatching a step that returned + // `Err` (`record_error` marks the pipeline done), so this is unreachable + // — but a scheduler that re-dispatched instead would otherwise fall + // through to `finalize()` and publish a residual + `AllAnnounced` for a + // sort that never ingested the rest of its input: a truncated result + // that is structurally indistinguishable from a complete one. + self.check_not_failed()?; + + if !self.flush_held(ctx) { + return Ok(StepOutcome::Contention); + } + + // Drain staged events before ingesting more — keeps peak memory bounded + // to ~one spill chunk plus the live buffer. + if !self.pending.is_empty() { + return Ok(self.emit_pending(ctx)); + } + + if self.sorter.is_some() { + if self.ingest_one_batch(ctx)? { + // Emit anything the batch just staged. + if !self.pending.is_empty() { + return Ok(self.emit_pending(ctx)); + } + return Ok(StepOutcome::Progress); + } + // No batch available right now. + if !ctx.input.is_drained() { + return Ok(StepOutcome::NoProgress); + } + // Input fully drained — produce the residual + AllAnnounced. + self.finalize(); + return Ok(self.emit_pending(ctx)); + } + + // Finalized and all events drained. + Ok(StepOutcome::Finished) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::sort::tests::record_with_mi; + use crate::types::RecordBatchBuilder; + use fgumi_sort::{ + InMemoryChunk, KeyTypesSpec, RawCoordinateKey, TemplateKey24, TemplateMemChunk, + }; + use rstest::rstest; + + fn coordinate(payloads: Vec>) -> MemoryChunkErased { + MemoryChunkErased::Coordinate(InMemoryChunk::from_owned_records( + payloads.into_iter().map(|b| (RawCoordinateKey::default(), b)).collect(), + )) + } + + fn template(payloads: Vec>) -> MemoryChunkErased { + MemoryChunkErased::TemplateCoordinate(TemplateMemChunk::K24( + InMemoryChunk::from_owned_records( + payloads.into_iter().map(|b| (TemplateKey24::default(), b)).collect(), + ), + )) + } + + /// The empty template-coordinate residual must survive `had_spills`, because + /// it is the only carrier of the narrowed-lane variant. Every other empty + /// residual is dropped, and every non-empty residual is kept. + #[rstest] + #[case::empty_coordinate_no_spills(coordinate(vec![]), false, 0)] + #[case::empty_coordinate_with_spills(coordinate(vec![]), true, 0)] + #[case::empty_template_no_spills(template(vec![]), false, 0)] + // The exception: kept purely to carry the key-width variant to the merge. + #[case::empty_template_with_spills(template(vec![]), true, 1)] + #[case::nonempty_coordinate_no_spills(coordinate(vec![vec![1u8; 8]]), false, 1)] + #[case::nonempty_coordinate_with_spills(coordinate(vec![vec![1u8; 8]]), true, 1)] + #[case::nonempty_template_no_spills(template(vec![vec![1u8; 8]]), false, 1)] + #[case::nonempty_template_with_spills(template(vec![vec![1u8; 8]]), true, 1)] + fn residual_chunk_emission_rule( + #[case] chunk: MemoryChunkErased, + #[case] had_spills: bool, + #[case] expected_len: usize, + ) { + assert_eq!(residual_chunks_for(chunk, had_spills).len(), expected_len); + } + + // ── Ingest failure handling ───────────────────────────────────────────── + + fn batch_of(records: &[Vec]) -> RecordBatch { + let total: usize = records.iter().map(Vec::len).sum(); + let mut builder = RecordBatchBuilder::with_capacity(0, total, records.len()); + for record in records { + builder.push_record_bytes(record); + } + builder.build() + } + + /// Template-coordinate sorter with every optional lane dropped + /// (`--key-types none`), so the first record fixes the narrowed variant and + /// any later record whose MI differs is rejected — the one reachable + /// `ChunkSorter::push` failure. + fn template_sorter_dropping_mi(memory_limit: usize) -> ChunkSorter { + let sorter = RawExternalSorter::new(SortOrder::TemplateCoordinate) + .memory_limit(memory_limit) + .threads(1) + .key_types(KeyTypesSpec::None); + ChunkSorter::from_sorter(sorter, &Header::default()).expect("build template chunk sorter") + } + + /// A push failure must stop ingest at the offending record and leave the + /// caller's sorter intact. `SortBuffer` distinguishes "ingesting" from + /// "finalized" solely by `sorter.is_some()`, so a failure that consumed the + /// sorter would let a later `try_run` report `Finished` — a truncated sort + /// that looks complete. + #[test] + fn ingest_batch_records_stops_at_the_failure_and_keeps_the_sorter() { + let mut sorter = template_sorter_dropping_mi(256 * 1024 * 1024); + let mut pending = VecDeque::new(); + let mut next_seq = 0u32; + let records = vec![ + record_with_mi(10, b"r1", 1), + record_with_mi(20, b"r2", 2), // differing MI — rejected + record_with_mi(30, b"r3", 1), + ]; + + let (traversed, error) = + ingest_batch_records(&mut sorter, &batch_of(&records), &mut pending, &mut next_seq); + + let message = error.expect("a differing MI under --key-types none must be rejected"); + assert!(message.starts_with("SortBuffer: push failed"), "unexpected message: {message}"); + assert_eq!(traversed, 2, "ingest stops at the offending record, not after the batch"); + assert_eq!(sorter.total_records(), 1, "only the accepted record reached the arena"); + // Borrowed, never moved out: the sorter is still usable afterwards. + assert_eq!(sorter.take_sorted_chunk().len(), 1); + assert!(pending.is_empty()); + assert_eq!(next_seq, 0); + } + + /// The clean path: each time the arena reaches `memory_limit` the sealed + /// chunk is staged and `next_seq` advances. A 1-byte limit seals per record. + #[test] + fn ingest_batch_records_stages_a_spill_chunk_each_time_the_arena_fills() { + let mut sorter = template_sorter_dropping_mi(1); + let mut pending = VecDeque::new(); + let mut next_seq = 0u32; + let records = vec![record_with_mi(10, b"r1", 1), record_with_mi(20, b"r2", 1)]; + + let (traversed, error) = + ingest_batch_records(&mut sorter, &batch_of(&records), &mut pending, &mut next_seq); + + assert!(error.is_none(), "constant MI is not a dropped-lane violation"); + assert_eq!(traversed, 2); + assert_eq!(pending.len(), 2, "one staged spill chunk per seal"); + assert_eq!(next_seq, 2); + } + + /// A poisoned step re-raises on every later `try_run` instead of falling + /// through to `finalize()`, and names the original failure so the re-raise + /// is not mistaken for a second, unrelated error. + #[rstest] + #[case::clean(None, None)] + #[case::poisoned( + Some("SortBuffer: push failed: dropped lane MI"), + Some( + "refusing to continue after an earlier failure: SortBuffer: push failed: dropped lane MI" + ) + )] + fn check_not_failed_re_raises_the_original_failure( + #[case] failed: Option<&str>, + #[case] expected_message: Option<&str>, + ) { + let mut step = SortBuffer::from_sorter( + RawExternalSorter::new(SortOrder::Coordinate).memory_limit(1 << 20).threads(1), + &Header::default(), + 1 << 20, + ) + .expect("build SortBuffer"); + step.failed = failed.map(str::to_string); + + match expected_message { + None => step.check_not_failed().expect("a step that never failed continues"), + Some(expected) => { + let err = step.check_not_failed().expect_err("a poisoned step must re-raise"); + assert!(err.to_string().contains(expected), "unexpected error: {err}"); + } + } + } +} diff --git a/crates/fgumi-pipeline-io/src/sort/spill_block_compress.rs b/crates/fgumi-pipeline-io/src/sort/spill_block_compress.rs new file mode 100644 index 000000000..7f698cbee --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/spill_block_compress.rs @@ -0,0 +1,152 @@ +//! `SpillBlockCompress` — middle step of the block-parallel spill-write split +//! (`SpillGather` → `SpillBlockCompress` → `SpillWrite`). +//! +//! `SpillBlockCompress` (`Parallel + ByItemOrdinal`) compresses each raw +//! [`SpillBlockEvent::Block`] from `SpillGather` into a self-contained +//! compressed unit (framed BGZF block(s) for bgzf, or a `[u32 len][zstd frame]` +//! for zstd) via the shared [`SpillBlockCompressor`] kernel, replacing the +//! `bytes` in place. `Residual` and `AllAnnounced` pass straight through. The +//! `ordinal` is preserved on every event, so the framework's `ByItemOrdinal` +//! reorder hands `SpillWrite` a dense, in-order stream regardless of which worker +//! compressed which block — exactly the output path's `BgzfCompress` idiom. +//! +//! Each `Parallel` worker holds its own [`SpillBlockCompressor`], built lazily on +//! the first block (the zstd compressor's construction is fallible, so it is +//! surfaced through `try_run`'s `io::Result` rather than the infallible +//! `new_worker_copy`). + +use std::io; + +use fgumi_sort::{SpillBlockCompressor, SpillCodec}; + +use crate::sort::protocol::SpillBlockEvent; +use fgumi_pipeline_core::{ + HeldRetry, Unpushed, + held::HeldSlot, + outputs::OrderedBytesSingle, + queues::QueueSpec, + reorder::BranchOrdering, + step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}, +}; + +/// `Parallel + ByItemOrdinal` block compressor for the spill-write split. +/// +/// Not to be confused with the similarly-named +/// [`CompressSpill`](super::CompressSpill): this `SpillBlockCompress` is the +/// **pure block-compression** middle step of the finer +/// `SpillGather → SpillBlockCompress → SpillWrite` split (the disk write is the +/// separate `SpillWrite` step), whereas `CompressSpill` is the **composite +/// compress-and-write-to-disk** step of the coarser +/// `SortBuffer → CompressSpill → SortSpillDecompress → SortMerge` chain. +pub struct SpillBlockCompress { + codec: SpillCodec, + compression: u32, + /// Per-worker compressor, built lazily on the first block. `None` until then + /// (and on fresh `new_worker_copy` clones). + compressor: Option, + held: HeldSlot>, + output_byte_limit: u64, +} + +impl SpillBlockCompress { + /// Build a `SpillBlockCompress` for `codec` at `compression`. `output_byte_limit` + /// byte-bounds the compressed-block output queue. + #[must_use] + pub fn new(codec: SpillCodec, compression: u32, output_byte_limit: u64) -> Self { + Self { codec, compression, compressor: None, held: HeldSlot::new(), output_byte_limit } + } + + fn flush_held(&mut self, ctx: &mut StepCtx<'_, Self>) -> bool { + !matches!(ctx.outputs.retry_held(&mut self.held), HeldRetry::StillHeld) + } + + /// Compress one event's payload (the `Block` arm) or pass it through. The + /// `ordinal` and routing fields are preserved. `StepCtx`-free for unit tests. + /// + /// # Errors + /// + /// Propagates compressor-init or compression errors. + fn compress_event(&mut self, event: SpillBlockEvent) -> io::Result { + match event { + SpillBlockEvent::Block { + ordinal, + file_id, + is_last_in_file, + records_ingested_so_far, + bytes, + } => { + if self.compressor.is_none() { + self.compressor = + Some(SpillBlockCompressor::new(self.codec, self.compression)?); + } + let compressor = self.compressor.as_mut().expect("compressor built above"); + let compressed = compressor.compress_block(&bytes)?; + Ok(SpillBlockEvent::Block { + ordinal, + file_id, + is_last_in_file, + records_ingested_so_far, + bytes: compressed, + }) + } + // Passthrough variants carry no compressible payload. + other @ (SpillBlockEvent::Residual { .. } | SpillBlockEvent::AllAnnounced { .. }) => { + Ok(other) + } + } + } +} + +impl Clone for SpillBlockCompress { + fn clone(&self) -> Self { + // Fresh per-worker compressor + held slot; shared config copied. + Self { + codec: self.codec, + compression: self.compression, + compressor: None, + held: HeldSlot::new(), + output_byte_limit: self.output_byte_limit, + } + } +} + +impl Step for SpillBlockCompress { + type Input = SpillBlockEvent; + type Outputs = OrderedBytesSingle; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "SpillBlockCompress", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], + branch_ordering: vec![BranchOrdering::ByItemOrdinal], + } + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if !self.flush_held(ctx) { + return Ok(StepOutcome::Contention); + } + + if let Some(event) = ctx.input.pop() { + let forwarded = self.compress_event(event)?; + if let Err(unpushed) = ctx.outputs.push(forwarded) { + self.held.put(unpushed); + } + return Ok(StepOutcome::Progress); + } + + if ctx.input.is_drained() { + return Ok(StepOutcome::Finished); + } + Ok(StepOutcome::NoProgress) + } + + fn new_worker_copy(&self) -> Self { + self.clone() + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/fgumi-pipeline-io/src/sort/spill_block_compress/tests.rs b/crates/fgumi-pipeline-io/src/sort/spill_block_compress/tests.rs new file mode 100644 index 000000000..6210470e6 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/spill_block_compress/tests.rs @@ -0,0 +1,186 @@ +//! Unit tests for `SpillBlockCompress::compress_event` (the `StepCtx`-free core). + +use super::*; +use crate::sort::protocol::MemoryChunkErased; +use fgumi_sort::{InMemoryChunk, RawCoordinateKey, SpillBlockDecompressor}; +use rstest::rstest; + +fn raw_block( + ordinal: u64, + file_id: u32, + is_last: bool, + records_ingested_so_far: u64, + bytes: Vec, +) -> SpillBlockEvent { + SpillBlockEvent::Block { + ordinal, + file_id, + is_last_in_file: is_last, + records_ingested_so_far, + bytes, + } +} + +#[rstest] +#[case(SpillCodec::Zstd)] +#[case(SpillCodec::Bgzf)] +fn block_payload_is_compressed_and_routing_preserved(#[case] codec: SpillCodec) { + let mut step = SpillBlockCompress::new(codec, 1, 1 << 20); + let raw = vec![0xABu8; 4096]; + let out = step.compress_event(raw_block(7, 2, true, 123, raw.clone())).unwrap(); + let SpillBlockEvent::Block { + ordinal, + file_id, + is_last_in_file, + records_ingested_so_far, + bytes, + } = out + else { + panic!("expected Block"); + }; + assert_eq!(ordinal, 7, "ordinal must be preserved ({codec:?})"); + assert_eq!(file_id, 2, "file_id must be preserved ({codec:?})"); + assert!(is_last_in_file, "is_last must be preserved ({codec:?})"); + assert_eq!( + records_ingested_so_far, 123, + "records_ingested_so_far must pass through compression unchanged ({codec:?})" + ); + assert_ne!(bytes, raw, "payload must change (be compressed) ({codec:?})"); + assert!(!bytes.is_empty(), "compressed payload non-empty ({codec:?})"); + + // Round-trip through the matching decoder (independent oracle): a compressor + // that silently corrupts still changes the bytes, so `assert_ne!` alone is + // too weak. `read_raw` parses the block framing (zstd length prefix / BGZF + // block) and `decompress_one` inverts the codec; the recovered payload must + // equal the exact input. + let mut dec = SpillBlockDecompressor::new(); + let mut cursor = std::io::Cursor::new(&bytes[..]); + let frames = dec.read_raw(&mut cursor, codec, 64).unwrap(); + let mut round = Vec::new(); + for frame in &frames { + round.extend_from_slice(&dec.decompress_one(codec, frame).unwrap()); + } + assert_eq!(round, raw, "compressed block must round-trip back to input ({codec:?})"); +} + +#[test] +fn residual_and_announced_pass_through_unchanged() { + let mut step = SpillBlockCompress::new(SpillCodec::Zstd, 1, 1 << 20); + + let chunk = MemoryChunkErased::Coordinate(InMemoryChunk::from_owned_records(vec![( + RawCoordinateKey { sort_key: 1 }, + vec![9u8; 8], + )])); + let residual = SpillBlockEvent::Residual { ordinal: 5, chunk, records_ingested_so_far: 1 }; + let out = step.compress_event(residual).unwrap(); + let SpillBlockEvent::Residual { ordinal, records_ingested_so_far, .. } = out else { + panic!("expected Residual"); + }; + assert_eq!(ordinal, 5, "residual ordinal must pass through unchanged"); + assert_eq!( + records_ingested_so_far, 1, + "residual records_ingested_so_far must pass through unchanged (not reset)" + ); + + let announced = SpillBlockEvent::AllAnnounced { + ordinal: 6, + slot_count: 2, + memory_chunk_count: 1, + total_records: 3, + }; + let out = step.compress_event(announced).unwrap(); + assert!(matches!( + out, + SpillBlockEvent::AllAnnounced { + ordinal: 6, + slot_count: 2, + memory_chunk_count: 1, + total_records: 3, + } + )); +} + +#[test] +fn clone_starts_with_fresh_lazy_compressor() { + let mut step = SpillBlockCompress::new(SpillCodec::Zstd, 1, 1 << 20); + // Force the original to build its compressor. + let _ = step.compress_event(raw_block(0, 0, true, 0, vec![1u8; 16])).unwrap(); + assert!(step.compressor.is_some(), "original built its compressor"); + let fresh = step.clone(); + assert!(fresh.compressor.is_none(), "clone must start with no compressor"); +} + +#[test] +fn new_worker_copy_is_independent_of_the_template() { + let mut template = SpillBlockCompress::new(SpillCodec::Bgzf, 3, 4096); + let _ = template.compress_event(raw_block(0, 0, true, 0, vec![2u8; 32])).unwrap(); + + let worker = template.new_worker_copy(); + // Config is inherited... + assert_eq!(worker.codec, SpillCodec::Bgzf); + assert_eq!(worker.compression, 3); + assert_eq!(worker.output_byte_limit, 4096); + // ...but the compressor is per-worker, so workers cannot share codec state. + assert!(worker.compressor.is_none(), "each worker builds its own compressor lazily"); +} + +#[test] +fn profile_advertises_parallel_byordinal_with_a_byte_bounded_queue() { + let step = SpillBlockCompress::new(SpillCodec::Zstd, 1, 8192); + let profile = step.profile(); + assert_eq!(profile.name, "SpillBlockCompress"); + // Parallel + ByItemOrdinal is what lets any worker compress any block while + // `SpillWrite` still receives a dense, in-order stream. + assert_eq!(profile.kind, StepKind::Parallel); + assert!(!profile.sticky); + assert_eq!(profile.branch_ordering, vec![BranchOrdering::ByItemOrdinal]); + match profile.output_queues.as_slice() { + [QueueSpec::ByteBounded { limit_bytes }] => assert_eq!(*limit_bytes, 8192), + other => panic!("expected a single byte-bounded queue, got {other:?}"), + } +} + +#[rstest] +#[case::zstd(SpillCodec::Zstd)] +#[case::bgzf(SpillCodec::Bgzf)] +fn an_empty_block_round_trips_to_empty(#[case] codec: SpillCodec) { + let mut step = SpillBlockCompress::new(codec, 1, 1 << 20); + let out = step.compress_event(raw_block(0, 0, true, 0, Vec::new())).unwrap(); + let SpillBlockEvent::Block { bytes, .. } = out else { panic!("expected Block") }; + + let mut dec = SpillBlockDecompressor::new(); + let mut cursor = std::io::Cursor::new(&bytes[..]); + let frames = dec.read_raw(&mut cursor, codec, 64).unwrap(); + let mut round = Vec::new(); + for frame in &frames { + round.extend_from_slice(&dec.decompress_one(codec, frame).unwrap()); + } + assert!(round.is_empty(), "an empty block must decompress back to empty ({codec:?})"); +} + +#[rstest] +#[case::zstd(SpillCodec::Zstd)] +#[case::bgzf(SpillCodec::Bgzf)] +fn consecutive_blocks_reuse_one_compressor_and_stay_independent(#[case] codec: SpillCodec) { + // The compressor is built once and reused across blocks; each block must still + // decode standalone, since `SpillWrite` may interleave files. + let mut step = SpillBlockCompress::new(codec, 1, 1 << 20); + let payloads = [vec![0x11u8; 512], vec![0x22u8; 1024], vec![0x33u8; 64]]; + + for (i, payload) in payloads.iter().enumerate() { + let out = step + .compress_event(raw_block(i as u64, 0, i == payloads.len() - 1, 0, payload.clone())) + .unwrap(); + let SpillBlockEvent::Block { bytes, .. } = out else { panic!("expected Block") }; + + let mut dec = SpillBlockDecompressor::new(); + let mut cursor = std::io::Cursor::new(&bytes[..]); + let frames = dec.read_raw(&mut cursor, codec, 64).unwrap(); + let mut round = Vec::new(); + for frame in &frames { + round.extend_from_slice(&dec.decompress_one(codec, frame).unwrap()); + } + assert_eq!(&round, payload, "block {i} must decode standalone ({codec:?})"); + } + assert!(step.compressor.is_some(), "the compressor is built once and retained"); +} diff --git a/crates/fgumi-pipeline-io/src/sort/spill_decompress.rs b/crates/fgumi-pipeline-io/src/sort/spill_decompress.rs new file mode 100644 index 000000000..770598485 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/spill_decompress.rs @@ -0,0 +1,565 @@ +//! `SortSpillDecompress` — Parallel typed step that reads spill chunk +//! files, decompresses their blocks, and pushes the decompressed bytes +//! into per-slot bounded queues on `SortMergeSlot`. +//! +//! # Two decompression granularities +//! +//! The step supports two strategies, selected by [`SortDecompressTuning`]: +//! +//! - **file-granularity (`file_granularity == true`, the fallback):** one worker +//! owns a file's decompression at a time. Under the per-slot reader lock it +//! reads AND decompresses up to `block_batch` blocks inline, in read order, +//! and pushes them to the slot's FIFO. No reorder buffer is needed — a plain +//! FIFO suffices because read-and-decompress is a single inline operation. This +//! is the proven path (see the `merge_slots` module header, "What used to live +//! here, and why it's gone (v4 vs v3.1)"). +//! +//! - **block-parallel (`file_granularity == false`):** multiple workers +//! decompress different blocks of the SAME file concurrently. Each `try_run` +//! holds the reader lock only for the READ (sequence-tagging each raw block via +//! `SortMergeReader::next_seq`), releases it, then decompresses its own batch +//! OUTSIDE the lock and reassembles via the slot's `ReorderBuffer`. The read +//! and decompression of a given block still happen within a single `try_run` +//! of a single worker — the lock is merely released between them. Parallelism +//! comes from multiple workers each grabbing the lock briefly, reading their +//! own batch, and decompressing concurrently — NOT from splitting read and +//! decompress across dispatches (which would re-open the v3 Skip-wedge +//! deadlock). +//! +//! # HARD INVARIANT +//! +//! A spill block must be read AND decompressed within a single `try_run` by a +//! single worker. Both paths uphold this. +//! +//! # Memory note +//! +//! `--max-memory` does NOT bound Phase-2 decompressed memory today: the FIFO is +//! count-bounded (`PHASE2_DECOMP_CAP`). The block-parallel path's reorder window +//! is the additional decompressed-memory surface a slow straggler could grow, so +//! it is explicitly bounded per-slot by `window_budget` (derived from the step's +//! `output_byte_limit`) via [`SortMergeSlot::bp_reorder_admits`]. + +use std::io; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use fgumi_sort::{SortMergeSlot, SpillBlockDecompressor}; +use parking_lot::Mutex; + +use crate::sort::protocol::{SortPhase1Event, SortPhase2Event}; +use fgumi_pipeline_core::{ + Unpushed, + held::HeldSlot, + outputs::Single, + queues::QueueSpec, + reorder::BranchOrdering, + step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}, +}; + +/// Default reorder-window byte budget substituted when the caller's +/// `output_byte_limit` is `0`. Mirrors the legacy pipeline's `effective_limit` +/// (`unified_pipeline`), which normalizes a `0` memory limit to a fixed cap +/// rather than treating it as "unlimited": `SortMergeSlot::bp_reorder_admits` +/// (via `ReorderBuffer::would_accept`) reads `window_budget == 0` as *no bound*, +/// so passing a resolved-to-zero budget straight through would remove the only +/// byte cap on decompressed stragglers. Matches +/// `fgumi_pipeline_core::reorder::DEFAULT_REORDER_OVERFLOW_BYTES` (256 MiB). +const DEFAULT_REORDER_WINDOW_BYTES: u64 = 256 * 1024 * 1024; + +/// Tuning for the Phase-2 spill decompression granularity. +/// +/// Threaded from `SortOptions` (`--sort::file-granularity` / +/// `--sort::block-batch`) through `ChainBuilder::add_sort` into +/// [`SortSpillDecompress::new`]. +#[derive(Debug, Clone, Copy)] +pub struct SortDecompressTuning { + /// `false` (default) — block-level parallel: hold the reader lock only for + /// the READ (sequence-tagged), release, decompress OUTSIDE the lock; multiple + /// workers decompress one file's blocks concurrently, reassembled by a + /// `ReorderBuffer`. The hardened production default (loom + soak matrix). + /// + /// `true` — one worker owns a file's decompression at a time (inline under + /// the reader lock, in-order, plain FIFO): the older single-worker-per-file + /// fallback. + pub file_granularity: bool, + /// Number of raw blocks claimed per reader-lock acquisition (replaces the + /// formerly-hardcoded batch size). Default `4` (restores the original + /// `MAX_BATCH_PER_CALL`; a fleet decompress-throughput bench will pick the + /// final value). Must be `>= 1`; [`SortSpillDecompress::new`] clamps lower + /// values and the CLI rejects them. + pub block_batch: usize, +} + +impl Default for SortDecompressTuning { + fn default() -> Self { + // Block-parallel is the production default: it cleared the hardening gate + // (loom over the real `SortMergeSlot` + the external-watchdog soak matrix + // + the reorder-window byte cap). `file_granularity = true` is the + // single-worker-per-file fallback. `block_batch` default is 4 (the + // original `MAX_BATCH_PER_CALL`); a fleet bench will tune it. Matches + // `SortOptions::default`. + Self { file_granularity: false, block_batch: 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, +} + +/// Parallel step that reads + decompresses spill chunk files and +/// pushes results into per-slot queues. Forwards `SortPhase1Event`s +/// verbatim to `SortMerge`. +pub struct SortSpillDecompress { + registry: Arc>>, + 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. Set from + /// `--merge-threads` (Phase-2 CPU control). + max_decompress: Option, + tuning: SortDecompressTuning, + /// Per-slot reorder-window byte budget for the block-parallel path. Derived + /// from `output_byte_limit` (one per-step byte budget per slot). Bounds the + /// reorder buffer so a slow straggler can't balloon decompressed memory. + window_budget: u64, +} + +impl SortSpillDecompress { + /// Construct a fresh step with an empty registry. + /// + /// `output_byte_limit` byte-bounds the forwarded-event output queue. + /// The forwarded `SortPhase2Event::MemoryChunk` variant retains sorted + /// record chunks, so this queue must budget on bytes (`HeapSize`), not + /// event count, to keep retained memory a function of configuration. + /// + /// `tuning` selects the decompression granularity (see + /// [`SortDecompressTuning`]). The block-parallel path derives its per-slot + /// reorder-window budget from `output_byte_limit`. + #[must_use] + pub fn new(output_byte_limit: u64, tuning: SortDecompressTuning) -> Self { + // Clamp `block_batch` to >= 1. A value of 0 reads zero blocks per + // acquisition, which on the inline path declares a phantom EOF after + // reading nothing (silent record loss) and on the block-parallel path + // never sets `reader_eof`/`queue_eof` (the merge livelocks). 0 is + // nonsensical for a "blocks per read" knob, so we normalize rather than + // propagate it. This is the single construction chokepoint for the step, + // so it defends every entry point (CLI, runall, direct construction). + let tuning = SortDecompressTuning { block_batch: tuning.block_batch.max(1), ..tuning }; + // Normalize a zero reorder-window budget to a sane default rather than + // propagating it: `bp_reorder_admits` treats `window_budget == 0` as + // "unlimited", so a resolved-to-zero budget (e.g. `--max-memory 0`) would + // silently remove the byte cap on decompressed stragglers. This mirrors + // the legacy `effective_limit` 0-normalization and is applied at the same + // construction chokepoint as the `block_batch` clamp, defending every + // entry point (CLI, runall, direct construction). + let window_budget = + if output_byte_limit == 0 { DEFAULT_REORDER_WINDOW_BYTES } else { output_byte_limit }; + Self { + registry: Arc::new(Mutex::new(Vec::new())), + block_dec: SpillBlockDecompressor::new(), + held: HeldSlot::new(), + output_byte_limit, + active: Arc::new(AtomicUsize::new(0)), + max_decompress: None, + tuning, + window_budget, + } + } + + /// 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; + }; + match ctx.outputs.retry(unpushed) { + Ok(()) => true, + Err(again) => { + self.held.put(again); + false + } + } + } + + fn push_or_hold(&mut self, ctx: &mut StepCtx<'_, Self>, event: SortPhase2Event) -> bool { + match ctx.outputs.push(event) { + Ok(()) => true, + Err(unpushed) => { + self.held.put(unpushed); + false + } + } + } + + fn snapshot_registry(&self) -> Vec> { + let registry = self.registry.lock(); + registry.iter().map(|e| Arc::clone(&e.slot)).collect() + } + + /// Slot indices ordered by ascending FIFO block count (most-starved first) — the + /// emptiest-first refill forecaster (see the budget-refill design spec §4.1). + /// Snapshots each slot's [`SortMergeSlot::fifo_len`] once (O(N) brief locks), then + /// sorts the indices, so the per-slot lock is taken exactly once per dispatch — not + /// inside the sort comparator. For the typical spill count (tens) this is negligible + /// next to the decompression work; if N grows into the hundreds, profile (a relaxed + /// cached length on the slot is the fallback) before keeping it. + /// + /// Slots that have already signalled `queue_eof` are dropped before the FIFO + /// lengths are read. The registry is append-only, so a drained slot would + /// otherwise stay in the scan for the rest of the run — taking its FIFO lock + /// on every dispatch only to be rejected immediately by + /// `try_fill_inline_slot` / `try_fill_block_parallel_slot`. With many spill + /// files that drained tail dominates the scan. Filtering is scheduling-only: + /// an EOF slot can never make progress, so skipping it changes no output. + #[must_use] + pub(crate) fn emptiest_first_order(slots: &[Arc]) -> Vec { + use std::sync::atomic::Ordering; + + let mut order: Vec = + (0..slots.len()).filter(|&i| !slots[i].queue_eof.load(Ordering::Acquire)).collect(); + // `sort_by_cached_key`, not `sort_by_key`: the key takes the slot's FIFO + // lock, and `sort_by_key` would re-take it O(n log n) times. This way each + // surviving slot is locked exactly once, and EOF slots not at all. + order.sort_by_cached_key(|&i| slots[i].fifo_len()); + order + } + + fn try_fill_some_slot(&mut self) -> io::Result { + let slots = self.snapshot_registry(); + // Refill the most-starved slot first so a free worker tops up the slot the merge + // will exhaust soonest, rather than the first in registry order. Scheduling-only: + // admission and the read-and-decompress-in-one-`try_run` invariant are unchanged, + // so this cannot affect output or wedge progress (a non-progressing slot returns + // `false` fast and the loop falls through to the next). + for i in Self::emptiest_first_order(&slots) { + let slot = &slots[i]; + let progressed = if self.tuning.file_granularity { + self.try_fill_inline_slot(slot)? + } else { + self.try_fill_block_parallel_slot(slot)? + }; + if progressed { + return Ok(true); + } + } + Ok(false) + } + + /// Inline (file-granularity) fill: read AND decompress up to `block_batch` + /// blocks under the reader lock, push them to the FIFO in read order. One + /// worker owns a slot at a time; no reorder buffer needed. + fn try_fill_inline_slot(&mut self, slot: &Arc) -> io::Result { + use std::sync::atomic::Ordering; + + if slot.queue_eof.load(Ordering::Acquire) { + return Ok(false); + } + + let mut reader_guard = match slot.reader.try_lock() { + Ok(guard) => guard, + // Contended (another worker owns the slot): a normal skip. + Err(std::sync::TryLockError::WouldBlock) => return Ok(false), + // Poisoned: a fill worker panicked mid-read. Fail the slot CLOSED so + // `SortMerge` surfaces the failure; swallowing it as a skip would + // leave `queue_eof` unset and spin `Contention` forever (deadlock). + Err(std::sync::TryLockError::Poisoned(_)) => { + Self::mark_slot_failed(slot); + return Err(io::Error::other( + "spill reader mutex poisoned: a decompress fill worker panicked", + )); + } + }; + + let room = { + let dec = slot.decompressed.lock().expect("decompressed mutex poisoned"); + fgumi_sort::PHASE2_DECOMP_CAP.saturating_sub(dec.len()) + }; + if room == 0 { + return Ok(false); + } + let want = room.min(self.tuning.block_batch); + + let decompressed_batch = + match self.block_dec.read_blocks(&mut reader_guard.inner, slot.codec, want) { + Ok(b) => b, + Err(e) => { + // Centralized in `mark_slot_failed` so failure semantics stay + // in one place (see the block-parallel path's use of it). + Self::mark_slot_failed(slot); + drop(reader_guard); + return Err(e); + } + }; + let got = decompressed_batch.len(); + let hit_eof = got < want; + + if got == 0 { + { + let _g = slot.decompressed.lock().expect("decompressed mutex poisoned"); + slot.queue_eof.store(true, Ordering::Release); + } + drop(reader_guard); + return Ok(true); + } + + { + let mut dec = slot.decompressed.lock().expect("decompressed mutex poisoned"); + for b in decompressed_batch { + dec.push_back(b); + } + if hit_eof { + slot.queue_eof.store(true, Ordering::Release); + } + } + drop(reader_guard); + Ok(true) + } + + /// Block-parallel fill: under the reader lock read (only) up to `block_batch` + /// raw blocks, sequence-tag them, release the lock, decompress OUTSIDE the + /// lock, then reassemble via the slot's reorder buffer and drain in-order + /// blocks into the FIFO. Multiple workers run this concurrently on the same + /// slot. + fn try_fill_block_parallel_slot(&mut self, slot: &Arc) -> io::Result { + use std::sync::atomic::Ordering; + + if slot.queue_eof.load(Ordering::Acquire) { + return Ok(false); + } + + // Phase A: read a fresh batch if the reader is still live and the + // FIFO / reorder window admit more. + if !slot.reader_eof.load(Ordering::Acquire) { + let acquired = match slot.reader.try_lock() { + Ok(guard) => Some(guard), + // Contended: fall through to the drain-only phase below. + Err(std::sync::TryLockError::WouldBlock) => None, + // Poisoned: a fill worker panicked mid-read. Fail closed so + // `SortMerge` surfaces it instead of spinning forever. + Err(std::sync::TryLockError::Poisoned(_)) => { + Self::mark_slot_failed(slot); + return Err(io::Error::other( + "spill reader mutex poisoned: a decompress fill worker panicked", + )); + } + }; + if let Some(mut reader_guard) = acquired { + // Re-check under the lock: another worker may have hit EOF. + if !slot.reader_eof.load(Ordering::Acquire) { + let next_seq = reader_guard.next_seq; + let fifo_room = slot.bp_fifo_room(); + let admit = + fifo_room > 0 && slot.bp_reorder_admits(next_seq, self.window_budget); + // NB: the reorder-window budget is checked once here (for + // `next_seq`), then up to `want` (≤ `block_batch`) blocks are + // inserted below without a per-block re-check. So the reorder + // window can transiently exceed `window_budget` by up to + // `block_batch - 1` blocks. This overshoot is bounded and by + // design: `block_batch` is small (default 4) and configurable, + // so worst-case resident bytes stay `O(window_budget + + // block_batch × block_size)` — not the unbounded growth the + // window guards against. Per-block admission is intentionally + // avoided to keep the reader-lock hold short (read the whole + // batch, release, decompress outside the lock). + if admit { + // Bound the read by FIFO room (as the inline path does): + // reading `block_batch` when only `fifo_room < block_batch` + // slots can drain would over-admit the surplus into the + // reorder window. `want >= 1` since `fifo_room > 0`. + let want = self.tuning.block_batch.min(fifo_room); + let start_seq = reader_guard.next_seq; + let raw = match self.block_dec.read_raw( + &mut reader_guard.inner, + slot.codec, + want, + ) { + Ok(r) => r, + Err(e) => { + Self::mark_slot_failed(slot); + drop(reader_guard); + return Err(e); + } + }; + let got = raw.len(); + // EOF only when the reader returned fewer than we asked + // for (`want`); a FIFO-limited short read is not EOF. + let hit_eof = got < want; + // Stamp the read range and account for it BEFORE releasing + // the lock, so a concurrent worker observing EOF cannot + // race ahead of this batch's in-flight accounting. The + // publish order (reserve `in_flight` before setting + // `reader_eof`) is the loom-verified protocol; it lives in + // `SortMergeSlot::bp_commit_read` as the single source of + // truth, so this call site and the loom model share it (see + // that method's doc and fgumi-sort tests/loom_merge_slots.rs). + reader_guard.next_seq += got as u64; + slot.bp_commit_read(got, hit_eof); + drop(reader_guard); + + // Decompress OUTSIDE the reader lock (still this try_run). + let mut blocks = Vec::with_capacity(got); + for raw_block in &raw { + match self.block_dec.decompress_one(slot.codec, raw_block) { + Ok(d) => blocks.push(d), + Err(e) => { + Self::mark_slot_failed(slot); + return Err(e); + } + } + } + slot.bp_insert_drain_finalize(start_seq, blocks, got); + return Ok(true); + } + } + } + } + + // Phase B: drain-only. Flush any now-in-order blocks the FIFO can accept + // (it may have freed up, or another worker delivered a straggler) and + // finalize EOF if fully delivered. + Ok(slot.bp_drain_and_finalize()) + } + + /// Mark a slot as failed (decompression / read error): set `decomp_error` + /// and `queue_eof` under the `decompressed` mutex so the consumer surfaces + /// the error in preference to a clean EOF. + fn mark_slot_failed(slot: &Arc) { + use std::sync::atomic::Ordering; + let _g = slot.decompressed.lock().expect("decompressed mutex poisoned"); + slot.decomp_error.store(true, Ordering::Release); + slot.queue_eof.store(true, Ordering::Release); + } +} + +impl Clone for SortSpillDecompress { + fn clone(&self) -> Self { + Self { + registry: Arc::clone(&self.registry), + 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, + tuning: self.tuning, + window_budget: self.window_budget, + } + } +} + +impl Step for SortSpillDecompress { + type Input = SortPhase1Event; + type Outputs = Single; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "SortSpillDecompress", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], + branch_ordering: vec![BranchOrdering::None], + } + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + // 1. Drain held output first. + if !self.flush_held(ctx) { + return Ok(StepOutcome::Contention); + } + + // 2. Pop one input event, register if SpillReady, forward all. + if let Some(event) = ctx.input.pop() { + let forwarded = match event { + SortPhase1Event::SpillReady { slot, path, records_ingested_so_far } => { + self.registry.lock().push(RegisteredSpill { slot: Arc::clone(&slot) }); + SortPhase2Event::SpillReady { slot, path, records_ingested_so_far } + } + SortPhase1Event::MemoryChunk { chunk, records_ingested_so_far } => { + SortPhase2Event::MemoryChunk { chunk, records_ingested_so_far } + } + SortPhase1Event::AllAnnounced { slot_count, memory_chunk_count, total_records } => { + SortPhase2Event::AllAnnounced { slot_count, memory_chunk_count, total_records } + } + }; + let _ = self.push_or_hold(ctx, forwarded); + return Ok(StepOutcome::Progress); + } + + // 3. Greedy slot-fill, admission-controlled by --merge-threads. + // `_permit` stays live across the `try_fill_some_slot` call — let-chain + // bindings are in scope for the conditions that follow them — and drops at + // the end of this `if` whichever way the fill goes (and on the `?` early + // return), releasing the count. + if let Some(_permit) = try_acquire(&self.active, self.max_decompress) + && self.try_fill_some_slot()? + { + return Ok(StepOutcome::Progress); + } + + // 4. No fill work. + let any_alive = self + .snapshot_registry() + .iter() + .any(|slot| !slot.queue_eof.load(std::sync::atomic::Ordering::Acquire)); + if any_alive { + return Ok(StepOutcome::Contention); + } + + if ctx.input.is_drained() { + return Ok(StepOutcome::Finished); + } + Ok(StepOutcome::NoProgress) + } + + fn new_worker_copy(&self) -> Self { + self.clone() + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/fgumi-pipeline-io/src/sort/spill_decompress/tests.rs b/crates/fgumi-pipeline-io/src/sort/spill_decompress/tests.rs new file mode 100644 index 000000000..f98a0d2e5 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/spill_decompress/tests.rs @@ -0,0 +1,205 @@ +use super::*; +use std::io::BufReader; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use fgumi_sort::{SortMergeSlot, SpillCodec}; + +/// A resolved-to-zero output budget must not disable the reorder-window byte cap: +/// `bp_reorder_admits` treats `window_budget == 0` as unlimited, so `new()` +/// substitutes the default cap (mirroring the legacy `effective_limit` +/// 0-normalization). A nonzero budget passes through unchanged. +#[test] +fn zero_output_byte_limit_normalizes_reorder_window() { + let zero = SortSpillDecompress::new(0, SortDecompressTuning::default()); + assert_eq!(zero.window_budget, DEFAULT_REORDER_WINDOW_BYTES); + assert_ne!(zero.window_budget, 0, "reorder window must stay bounded on a zero budget"); + + let nonzero = SortSpillDecompress::new(4 * 1024 * 1024, SortDecompressTuning::default()); + assert_eq!(nonzero.window_budget, 4 * 1024 * 1024, "nonzero budget passes through unchanged"); +} + +#[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"); +} + +// Most coverage for the decompress step lives in sort/tests.rs (the oracle parity +// suite drives the whole chain). This unit test pins the emptiest-first refill +// ordering in isolation. + +#[test] +fn emptiest_first_order_sorts_by_fifo_len_ascending() { + let mk = |file_id: u32, nblocks: usize| { + let s = Arc::new(SortMergeSlot::new( + file_id, + BufReader::new(tempfile::tempfile().expect("tempfile")), + SpillCodec::Bgzf, + )); + for _ in 0..nblocks { + s.decompressed.lock().expect("decompressed lock").push_back(vec![0u8]); + } + s + }; + // FIFO depths 5, 1, 3 ⇒ most-starved-first visit order is indices 1, 2, 0. + let slots = vec![mk(0, 5), mk(1, 1), mk(2, 3)]; + assert_eq!(SortSpillDecompress::emptiest_first_order(&slots), vec![1, 2, 0]); +} + +/// A poisoned `reader` mutex (a fill worker panicked while holding the lock) +/// must fail the slot CLOSED — `try_fill_*_slot` returns `Err` and sets +/// `decomp_error`/`queue_eof` — rather than being swallowed as `WouldBlock` and +/// skipped forever. If it were skipped, `queue_eof` would never be set and +/// `SortMerge` would spin on `Contention` and deadlock instead of surfacing the +/// panic. Regression test for the poisoned-vs-would-block conflation. +#[test] +fn poisoned_reader_lock_fails_closed_rather_than_hanging() { + let make_poisoned_slot = || { + let slot = Arc::new(SortMergeSlot::new( + 0, + BufReader::new(tempfile::tempfile().expect("tempfile")), + SpillCodec::Bgzf, + )); + let holder = Arc::clone(&slot); + // Panic while holding the reader lock; joining the panicked thread leaves + // the mutex poisoned (mirrors a fill worker dying mid-read). + let _ = std::thread::spawn(move || { + let _guard = holder.reader.lock().expect("acquire reader lock"); + panic!("simulated fill-worker panic under the reader lock"); + }) + .join(); + assert!(slot.reader.is_poisoned(), "precondition: reader mutex is poisoned"); + slot + }; + + let mut dec = SortSpillDecompress::new(4 * 1024 * 1024, SortDecompressTuning::default()); + + // Inline path. + let inline_slot = make_poisoned_slot(); + let inline = dec.try_fill_inline_slot(&inline_slot); + assert!(inline.is_err(), "inline path: poisoned reader must return Err, not Ok(false)"); + assert!(inline_slot.decomp_error.load(Ordering::Acquire), "inline: decomp_error set"); + assert!(inline_slot.queue_eof.load(Ordering::Acquire), "inline: queue_eof set"); + + // Block-parallel path. + let bp_slot = make_poisoned_slot(); + let bp = dec.try_fill_block_parallel_slot(&bp_slot); + assert!(bp.is_err(), "block-parallel path: poisoned reader must return Err, not skip"); + assert!(bp_slot.decomp_error.load(Ordering::Acquire), "bp: decomp_error set"); + assert!(bp_slot.queue_eof.load(Ordering::Acquire), "bp: queue_eof set"); +} + +/// Slots that already signalled `queue_eof` are dropped from the refill scan. +/// +/// The registry is append-only, so without this filter a drained slot stays in +/// the scan for the rest of the run: every dispatch clones its `Arc`, takes its +/// FIFO lock to sort, and is then rejected immediately by the `queue_eof` guard +/// in `try_fill_*_slot`. With many spill files that drained tail dominates the +/// scan. Filtering is scheduling-only — an EOF slot can never progress — so it +/// changes no output. +#[test] +fn emptiest_first_order_skips_slots_that_reached_eof() { + use std::sync::atomic::Ordering; + + let mk = |file_id: u32, nblocks: usize, eof: bool| { + let s = Arc::new(SortMergeSlot::new( + file_id, + BufReader::new(tempfile::tempfile().expect("tempfile")), + SpillCodec::Bgzf, + )); + for _ in 0..nblocks { + s.decompressed.lock().expect("decompressed lock").push_back(vec![0u8]); + } + if eof { + s.queue_eof.store(true, Ordering::Release); + } + s + }; + + // Slot 1 is the emptiest but has drained; it must not appear at all. + let slots = vec![mk(0, 5, false), mk(1, 0, true), mk(2, 3, false)]; + assert_eq!( + SortSpillDecompress::emptiest_first_order(&slots), + vec![2, 0], + "drained slots are skipped; the rest stay most-starved-first" + ); + + // Every slot drained ⇒ nothing to scan. + let all_done = vec![mk(0, 0, true), mk(1, 0, true)]; + assert!( + SortSpillDecompress::emptiest_first_order(&all_done).is_empty(), + "a fully drained registry yields an empty scan order" + ); + + // No slot drained ⇒ unchanged from the pre-filter behaviour. + let none_done = vec![mk(0, 5, false), mk(1, 1, false), mk(2, 3, false)]; + assert_eq!(SortSpillDecompress::emptiest_first_order(&none_done), vec![1, 2, 0]); +} + +/// The `block_batch` clamp is the sibling of the `window_budget` normalization +/// already pinned by `zero_output_byte_limit_normalizes_reorder_window`. A +/// `block_batch` of 0 would declare a phantom EOF after reading nothing on the +/// inline path — silent record loss — so it is clamped to at least one block. +#[test] +fn zero_block_batch_is_clamped_to_one() { + let tuning = SortDecompressTuning { block_batch: 0, ..Default::default() }; + let clamped = SortSpillDecompress::new(4 * 1024 * 1024, tuning); + assert_eq!(clamped.tuning.block_batch, 1, "a zero block_batch must be clamped, not honoured"); + + // A sane value passes through untouched. + let tuning = SortDecompressTuning { block_batch: 8, ..Default::default() }; + let passthrough = SortSpillDecompress::new(4 * 1024 * 1024, tuning); + assert_eq!(passthrough.tuning.block_batch, 8, "a nonzero block_batch is honoured"); +} diff --git a/crates/fgumi-pipeline-io/src/sort/spill_gather.rs b/crates/fgumi-pipeline-io/src/sort/spill_gather.rs new file mode 100644 index 000000000..4235bed12 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/spill_gather.rs @@ -0,0 +1,337 @@ +//! `SpillGather` — first step of the block-parallel spill-write split +//! (`SpillGather` → `SpillBlockCompress` → `SpillWrite`), replacing the monolithic +//! single-worker `CompressSpill`. +//! +//! `SpillGather` (`Serial`) consumes the [`SortChunkEvent`]s `SortBuffer` +//! emits and fans each `Spill` chunk into record-aligned **raw** (uncompressed) +//! [`SpillBlockEvent::Block`]s of ≤`BGZF_MAX_BLOCK_SIZE`, so the downstream +//! `Parallel` `SpillBlockCompress` can compress them across the framework pool. The +//! in-memory `Residual` and the terminal `AllAnnounced` pass straight through. +//! +//! # Ordinal minting +//! +//! The step mints `ordinal` monotonically across **every** emitted item — every +//! block of every file plus the passthrough `Residual` / `AllAnnounced` — so the +//! output stream is dense and gap-free. That is what lets the framework's +//! single-cursor `ByItemOrdinal` reorder deliver the blocks to the `Serial` +//! `SpillWrite` in order without any per-file ordering primitive. Because +//! `SortBuffer` (Serial) emits `Spill` events one-at-a-time in `seq` order and +//! this step (Serial) drains them in order, each file's blocks are contiguous in +//! the ordinal stream. +//! +//! # Bounded memory (incremental framing) +//! +//! A spill chunk is large (≈ the per-thread sort budget). Framing it **all** into +//! `pending` at once would duplicate the whole chunk in memory and then drain +//! that copy slowly through the byte-bounded queue while `SortBuffer` races ahead +//! filling the next multi-GB buffer — a ~2× peak-RSS blow-up. Instead the chunk +//! is held in `active` and framed **incrementally**: each `try_run` frames at most +//! `MAX_EVENTS_PER_LOCK` blocks into `pending`, drains them, and only frees the +//! source chunk once its last record is framed. `pending` therefore holds ≤ a +//! handful of 64 KiB blocks (~½ MiB) regardless of chunk size, and the chunk +//! drains as fast as `SpillBlockCompress` consumes blocks. + +use std::collections::VecDeque; +use std::io; + +use fgumi_bgzf::BGZF_MAX_BLOCK_SIZE; +use fgumi_sort::frame_keyed_record_into; + +use crate::sort::protocol::{MemoryChunkErased, SortChunkEvent, SpillBlockEvent}; +use fgumi_pipeline_core::{ + HeldRetry, Unpushed, + held::HeldSlot, + outputs::OrderedBytesSingle, + queues::QueueSpec, + reorder::BranchOrdering, + step::{DetachedGroup, Step, StepCtx, StepKind, StepOutcome, StepProfile}, +}; + +/// Max staged events flushed to the output per `try_run` invocation. Matches +/// `SortBuffer::MAX_EVENTS_PER_LOCK` so a single fanned chunk drains in bounded +/// slices rather than holding the step lock for the whole spill. +const MAX_EVENTS_PER_LOCK: usize = 8; + +/// Frame the `i`th record of a type-erased chunk into `out` in the spill layout +/// `[key?][u32 LE len][record]`, dispatching over the sort-key variant. +fn frame_record_at(chunk: &MemoryChunkErased, i: usize, out: &mut Vec) -> io::Result<()> { + match chunk { + MemoryChunkErased::Coordinate(c) => { + frame_keyed_record_into(out, c.key_at(i), c.record_bytes(i)) + } + MemoryChunkErased::QuerynameLex(c) => { + frame_keyed_record_into(out, c.key_at(i), c.record_bytes(i)) + } + MemoryChunkErased::QuerynameNatural(c) => { + frame_keyed_record_into(out, c.key_at(i), c.record_bytes(i)) + } + MemoryChunkErased::TemplateCoordinate(c) => c.frame_record_into(i, out), + } +} + +/// Frame records `[start..)` of `chunk` into a single block `out` of at most +/// `block_size` bytes, returning the index of the next un-framed record (== +/// `chunk.len()` when the chunk is exhausted). +/// +/// A record is never split across blocks: a record that would push a **non-empty** +/// block past `block_size` is left for the next block. The sole exception is a +/// single record larger than `block_size`, which forms its own oversized block +/// (BAM long reads can exceed 64 KiB; the codec re-blocks/​frames it safely, +/// matching the streaming `SyncSpillWriter`, so this is not an error). +fn frame_one_block( + chunk: &MemoryChunkErased, + start: usize, + block_size: usize, + out: &mut Vec, +) -> io::Result { + let len = chunk.len(); + let mut i = start; + while i < len { + let before = out.len(); + frame_record_at(chunk, i, out)?; + if before != 0 && out.len() > block_size { + // This record overflowed a non-empty block — roll it back so it + // starts the next block, and finish this one. + out.truncate(before); + break; + } + i += 1; + if out.len() >= block_size { + // Block full (the record fit exactly, or was a lone oversized record + // framed into an empty block). + break; + } + } + Ok(i) +} + +/// Frame an entire chunk into blocks (convenience for tests; production frames +/// incrementally via [`SpillGather::produce_blocks`]). Returns one `Vec` +/// per block; an empty chunk yields no blocks. +#[cfg(test)] +fn frame_chunk_into_blocks( + chunk: &MemoryChunkErased, + block_size: usize, +) -> io::Result>> { + let mut blocks: Vec> = Vec::new(); + let len = chunk.len(); + let mut idx = 0; + while idx < len { + let mut block = Vec::with_capacity(block_size + 1024); + idx = frame_one_block(chunk, idx, block_size, &mut block)?; + if !block.is_empty() { + blocks.push(block); + } + } + Ok(blocks) +} + +/// One spill chunk being framed incrementally into blocks across `try_run` calls. +struct ActiveSpill { + /// The source sorted chunk, held until its last record is framed. + chunk: MemoryChunkErased, + /// Index of the next un-framed record. + next_idx: usize, + /// Logical spill index (the eventual slot `file_id`). + file_id: u32, + records_ingested_so_far: u64, +} + +/// `Serial` step that fans sorted spill chunks into raw blocks for `SpillBlockCompress`. +pub struct SpillGather { + /// The spill chunk currently being framed incrementally (`None` between + /// chunks). Holding it here — rather than materializing all its blocks into + /// `pending` — is what keeps peak memory bounded. + active: Option, + /// Staged block / passthrough events awaiting output. Bounded to ≤ + /// `MAX_EVENTS_PER_LOCK` blocks because framing only tops it up when it is + /// already drained. + pending: VecDeque, + /// Monotonic ordinal minted across every emitted item (dense, gap-free). + next_ordinal: u64, + held: HeldSlot>, + /// Raw-block size threshold (records are cut into blocks of ≤ this many bytes). + block_size: usize, + output_byte_limit: u64, +} + +impl SpillGather { + /// Build a `SpillGather`. `output_byte_limit` byte-bounds the block-event + /// output queue (its `Block` / `Residual` variants retain bytes/records). + #[must_use] + pub fn new(output_byte_limit: u64) -> Self { + Self { + active: None, + pending: VecDeque::new(), + next_ordinal: 0, + held: HeldSlot::new(), + block_size: BGZF_MAX_BLOCK_SIZE, + output_byte_limit, + } + } + + /// Take the next dense ordinal. + fn next_ordinal(&mut self) -> u64 { + let o = self.next_ordinal; + self.next_ordinal += 1; + o + } + + /// `StepCtx`-free core: stage one input event. A `Spill` chunk is parked in + /// `active` for incremental framing (no blocks produced yet); `Residual` / + /// `AllAnnounced` are cheap and pushed straight to `pending`. Caller + /// guarantees `active` is `None` (the previous chunk fully framed). + fn stage_event(&mut self, event: SortChunkEvent) { + match event { + SortChunkEvent::Spill { seq, chunk, records_ingested_so_far } => { + // `SortBuffer` only emits `Spill` for a non-empty buffer; skip an + // (unexpected) empty chunk rather than open a zero-block file with + // no `is_last_in_file` terminator. + if chunk.is_empty() { + return; + } + self.active = + Some(ActiveSpill { chunk, next_idx: 0, file_id: seq, records_ingested_so_far }); + } + SortChunkEvent::Residual { chunk, records_ingested_so_far } => { + let ordinal = self.next_ordinal(); + self.pending.push_back(SpillBlockEvent::Residual { + ordinal, + chunk, + records_ingested_so_far, + }); + } + SortChunkEvent::AllAnnounced { slot_count, memory_chunk_count, total_records } => { + let ordinal = self.next_ordinal(); + self.pending.push_back(SpillBlockEvent::AllAnnounced { + ordinal, + slot_count, + memory_chunk_count, + total_records, + }); + } + } + } + + /// Frame up to `MAX_EVENTS_PER_LOCK` blocks of the `active` chunk into + /// `pending`, minting dense ordinals and flagging the final block + /// `is_last_in_file`. Frees the chunk (`active = None`) once its last record + /// is framed. No-op when there is no active chunk. + /// + /// # Errors + /// + /// Propagates framing errors (e.g. a record too large for the `u32` length + /// prefix). + fn produce_blocks(&mut self) -> io::Result<()> { + let block_size = self.block_size; + while self.pending.len() < MAX_EVENTS_PER_LOCK { + // Frame one block, scoping the borrow of `active` so `next_ordinal` + // (a `&mut self` method) can run afterward. + let framed = { + let Some(active) = self.active.as_ref() else { return Ok(()) }; + let len = active.chunk.len(); + let mut bytes = Vec::with_capacity(block_size + 1024); + let next = frame_one_block(&active.chunk, active.next_idx, block_size, &mut bytes)?; + (bytes, next, next >= len, active.file_id, active.records_ingested_so_far) + }; + let (bytes, next_idx, is_last, file_id, records_ingested_so_far) = framed; + let ordinal = self.next_ordinal(); + self.pending.push_back(SpillBlockEvent::Block { + ordinal, + file_id, + is_last_in_file: is_last, + records_ingested_so_far, + bytes, + }); + if is_last { + self.active = None; + return Ok(()); + } + // Advance the cursor for the next block. + self.active.as_mut().expect("active present (not last)").next_idx = next_idx; + } + Ok(()) + } + + fn flush_held(&mut self, ctx: &mut StepCtx<'_, Self>) -> bool { + !matches!(ctx.outputs.retry_held(&mut self.held), HeldRetry::StillHeld) + } + + /// Push up to `MAX_EVENTS_PER_LOCK` staged events, parking the first that + /// can't be pushed in `held`. Caller guarantees `held` is empty. + fn emit_pending(&mut self, ctx: &mut StepCtx<'_, Self>) -> StepOutcome { + let mut emitted = 0usize; + while emitted < MAX_EVENTS_PER_LOCK { + let Some(event) = self.pending.pop_front() else { break }; + if let Err(unpushed) = ctx.outputs.push(event) { + self.held.put(unpushed); + return StepOutcome::Progress; + } + emitted += 1; + } + if emitted > 0 { StepOutcome::Progress } else { StepOutcome::NoProgress } + } +} + +impl Step for SpillGather { + type Input = SortChunkEvent; + type Outputs = OrderedBytesSingle; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "SpillGather", + // Off-pool on the coordination driver (N+2): the serial spill framing + // + monotonic ordinal minting runs on the dedicated coordination + // thread instead of a pool worker, so it never starves the parallel + // spill compressors. Detached collapses `ByItemOrdinal` to `None` + // exactly as `Serial` did (transport-identical). + kind: StepKind::Detached, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], + branch_ordering: vec![BranchOrdering::ByItemOrdinal], + } + } + + fn detached_group(&self) -> DetachedGroup { + DetachedGroup::Shared(crate::sort::SORT_COORD_GROUP) + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if !self.flush_held(ctx) { + return Ok(StepOutcome::Contention); + } + + // Drain staged events before producing/ingesting more — bounds peak memory. + if !self.pending.is_empty() { + return Ok(self.emit_pending(ctx)); + } + + // Continue framing the active chunk (incrementally), if any. + if self.active.is_some() { + self.produce_blocks()?; + return Ok(self.emit_pending(ctx)); + } + + // No active chunk and nothing pending — ingest the next event. + if let Some(event) = ctx.input.pop() { + self.stage_event(event); + // Frame the first blocks of a freshly-parked chunk so we make progress. + if self.active.is_some() { + self.produce_blocks()?; + } + if !self.pending.is_empty() { + return Ok(self.emit_pending(ctx)); + } + return Ok(StepOutcome::Progress); + } + + // Drained: pending is empty and no chunk is mid-framing (checked above). + if ctx.input.is_drained() { + return Ok(StepOutcome::Finished); + } + Ok(StepOutcome::NoProgress) + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/fgumi-pipeline-io/src/sort/spill_gather/tests.rs b/crates/fgumi-pipeline-io/src/sort/spill_gather/tests.rs new file mode 100644 index 000000000..8a5a37cb1 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/spill_gather/tests.rs @@ -0,0 +1,323 @@ +//! Unit tests for `SpillGather`'s `StepCtx`-free core: chunk fan-out into raw +//! blocks (`frame_chunk_into_blocks`) and event staging (`stage_event`). + +use super::*; +use crate::sort::protocol::MemoryChunkErased; +use fgumi_sort::{InMemoryChunk, RawCoordinateKey}; +use proptest::prelude::*; + +/// Build a coordinate chunk from raw payloads with distinct keys, so the framed +/// blocks exercise key serialization (`RawCoordinateKey` is embedded, so no key +/// prefix is written — but the chunk path is still the production one). +fn coord_chunk(payloads: Vec>) -> MemoryChunkErased { + let recs = payloads + .into_iter() + .enumerate() + .map(|(i, b)| (RawCoordinateKey { sort_key: i as u64 }, b)) + .collect(); + MemoryChunkErased::Coordinate(InMemoryChunk::from_owned_records(recs)) +} + +/// Concatenate the framed-record bytes of all blocks (drops the per-block +/// boundaries) — the decompressed-stream-equivalent the readers see. +fn concat(blocks: &[Vec]) -> Vec { + blocks.iter().flat_map(|b| b.iter().copied()).collect() +} + +/// Drive `SpillGather` over a sequence of input events exactly as `try_run` +/// does — stage each event, then fully frame any active chunk (draining produced +/// blocks) before staging the next — and return every emitted event in order. +/// This preserves the dense-ordinal invariant (a chunk is fully framed, and its +/// block ordinals minted, before the next event is staged). +fn drive(step: &mut SpillGather, events: Vec) -> Vec { + let mut out = Vec::new(); + for event in events { + step.stage_event(event); + // Frame the active chunk to completion, collecting blocks as we go. + while step.active.is_some() { + step.produce_blocks().unwrap(); + while let Some(ev) = step.pending.pop_front() { + out.push(ev); + } + } + // Drain any passthrough (Residual / AllAnnounced) events. + while let Some(ev) = step.pending.pop_front() { + out.push(ev); + } + } + out +} + +#[test] +fn frame_empty_chunk_yields_no_blocks() { + let chunk = coord_chunk(Vec::new()); + let blocks = frame_chunk_into_blocks(&chunk, BGZF_MAX_BLOCK_SIZE).unwrap(); + assert!(blocks.is_empty(), "empty chunk must yield zero blocks"); +} + +#[test] +fn frame_cuts_blocks_at_threshold_without_splitting_records() { + // 10 records of 100 payload bytes each; embedded key adds only the 4-byte + // length prefix → 104 framed bytes/record. A 250-byte block threshold cuts + // *before* the record that would exceed 250, so each block holds 2 records + // (2×104=208 ≤ 250; a 3rd would be 312 > 250) → 5 blocks. + let chunk = coord_chunk((0u8..10).map(|i| vec![i; 100]).collect()); + let blocks = frame_chunk_into_blocks(&chunk, 250).unwrap(); + assert_eq!(blocks.len(), 5, "250-byte threshold must pack 2 records/block over 10 records"); + // Every block stays within the threshold and holds a whole number of + // 104-byte records (no record split). + for b in &blocks { + assert!(b.len() <= 250, "block exceeded threshold: {} bytes", b.len()); + assert_eq!(b.len() % 104, 0, "block boundary split a record: {} bytes", b.len()); + } + // Reassembling the blocks reproduces a single 10-record stream. + assert_eq!(concat(&blocks).len(), 10 * 104); +} + +proptest! { + /// The "no record split, ≤ threshold, exact reconstruction" packing invariant + /// of `frame_chunk_into_blocks` holds for arbitrary record sizes and + /// thresholds — far more of the space than the hand-picked cases above. + #[test] + fn frame_chunk_into_blocks_packing_invariant( + payloads in prop::collection::vec(prop::collection::vec(any::(), 0..40), 0..30), + threshold in 1usize..512, + ) { + let chunk = coord_chunk(payloads.clone()); + let blocks = frame_chunk_into_blocks(&chunk, threshold).unwrap(); + + // Content is threshold-independent: block boundaries move, bytes do not. + // Compare against both the all-in-one framing and the one-record-per-block + // framing (threshold 1), pinning exact reconstruction of the framed stream. + let one_block = frame_chunk_into_blocks(&chunk, BGZF_MAX_BLOCK_SIZE).unwrap(); + let per_record = frame_chunk_into_blocks(&chunk, 1).unwrap(); + prop_assert_eq!(concat(&blocks), concat(&one_block)); + prop_assert_eq!(concat(&blocks), concat(&per_record)); + + // Empty chunk → no blocks; any record (even a zero-length payload, which + // still frames its length prefix) → at least one, and never an empty one. + prop_assert_eq!(blocks.is_empty(), payloads.is_empty()); + prop_assert!(blocks.iter().all(|b| !b.is_empty()), "no empty blocks"); + + // No record is split: a block may exceed `threshold` only when a single + // record is itself larger than it (the "cut before exceeding" rule can't + // shrink one record), so every block is bounded by + // `max(threshold, largest single framed record)`. + let max_record = per_record.iter().map(Vec::len).max().unwrap_or(0); + let bound = threshold.max(max_record); + for b in &blocks { + prop_assert!(b.len() <= bound, "block {} exceeds bound {bound}", b.len()); + } + } +} + +#[test] +fn frame_one_block_when_under_threshold() { + let chunk = coord_chunk((0u8..5).map(|i| vec![i; 100]).collect()); + let blocks = frame_chunk_into_blocks(&chunk, BGZF_MAX_BLOCK_SIZE).unwrap(); + assert_eq!(blocks.len(), 1, "5 small records fit one block"); + assert_eq!(blocks[0].len(), 5 * 104); +} + +#[test] +fn stage_spill_mints_dense_ordinals_and_marks_last_block() { + let mut step = SpillGather::new(1 << 20); + step.block_size = 250; // force several blocks + let chunk = coord_chunk((0u8..10).map(|i| vec![i; 100]).collect()); + let events = drive( + &mut step, + vec![SortChunkEvent::Spill { seq: 3, chunk, records_ingested_so_far: 10 }], + ); + + assert!(events.len() > 1, "expected multiple block events"); + // Ordinals are dense 0..n. + for (i, ev) in events.iter().enumerate() { + assert_eq!(ev.ordinal(), i as u64, "ordinal not dense at index {i}"); + } + // Exactly the final block is flagged is_last_in_file; all carry file_id=3. + for (i, ev) in events.iter().enumerate() { + let SpillBlockEvent::Block { file_id, is_last_in_file, records_ingested_so_far, .. } = ev + else { + panic!("expected Block event"); + }; + assert_eq!(*file_id, 3); + assert_eq!(*records_ingested_so_far, 10); + assert_eq!(*is_last_in_file, i == events.len() - 1, "is_last wrong at {i}"); + } + // The chunk was freed once fully framed. + assert!(step.active.is_none(), "active chunk must be cleared after framing"); +} + +#[test] +fn produce_blocks_keeps_pending_bounded() { + // With many small records and a tiny block size, a single produce_blocks call + // must not materialize the whole chunk — it tops up at most + // MAX_EVENTS_PER_LOCK blocks. + let mut step = SpillGather::new(1 << 20); + step.block_size = 120; // ~1 record/block over 104-byte records + let chunk = coord_chunk((0u8..100).map(|i| vec![i; 100]).collect()); + step.stage_event(SortChunkEvent::Spill { seq: 0, chunk, records_ingested_so_far: 100 }); + step.produce_blocks().unwrap(); + assert_eq!(step.pending.len(), 8, "one produce call tops up to MAX_EVENTS_PER_LOCK blocks"); + assert!(step.active.is_some(), "chunk still mid-framing (not fully drained in one call)"); +} + +#[test] +fn stage_residual_and_announced_pass_through_with_ordinals() { + let mut step = SpillGather::new(1 << 20); + // A spill (2 small records → 1 block, ordinal 0), then residual, then + // AllAnnounced — ordinals must stay dense across the variant boundary, with + // the spill's block ordinal minted *before* the later events (the chunk is + // fully framed before the next event is staged). + let events = drive( + &mut step, + vec![ + SortChunkEvent::Spill { + seq: 0, + chunk: coord_chunk(vec![vec![1u8; 10], vec![2u8; 10]]), + records_ingested_so_far: 2, + }, + SortChunkEvent::Residual { + chunk: coord_chunk(vec![vec![3u8; 10]]), + records_ingested_so_far: 3, + }, + SortChunkEvent::AllAnnounced { slot_count: 1, memory_chunk_count: 1, total_records: 3 }, + ], + ); + + let ords: Vec = events.iter().map(SpillBlockEvent::ordinal).collect(); + assert_eq!(ords, vec![0, 1, 2], "ordinals must be dense across variants"); + + assert!(matches!(events[0], SpillBlockEvent::Block { is_last_in_file: true, .. })); + assert!(matches!(events[1], SpillBlockEvent::Residual { .. })); + assert!(matches!( + events[2], + SpillBlockEvent::AllAnnounced { + slot_count: 1, + memory_chunk_count: 1, + total_records: 3, + .. + } + )); +} + +// ── Step wiring + ordinal minting ──────────────────────────────────────────── + +#[test] +fn profile_runs_off_pool_on_the_coordination_driver() { + let step = SpillGather::new(4096); + let profile = step.profile(); + assert_eq!(profile.name, "SpillGather"); + // Detached keeps the serial framing + ordinal minting off a pool worker, so + // it cannot starve the parallel spill compressors downstream. + assert_eq!(profile.kind, StepKind::Detached); + assert!(!profile.sticky); + assert_eq!(profile.branch_ordering, vec![BranchOrdering::ByItemOrdinal]); + match profile.output_queues.as_slice() { + [QueueSpec::ByteBounded { limit_bytes }] => assert_eq!(*limit_bytes, 4096), + other => panic!("expected one byte-bounded queue, got {other:?}"), + } + assert_eq!(step.detached_group(), DetachedGroup::Shared(crate::sort::SORT_COORD_GROUP)); +} + +#[test] +fn ordinals_are_dense_and_monotonic_across_event_kinds() { + // Downstream `ByItemOrdinal` reordering depends on a gap-free sequence, and + // the counter is shared by every emitted event, not per-variant. + let mut step = SpillGather::new(1 << 20); + + step.stage_event(SortChunkEvent::Residual { + chunk: coord_chunk(vec![vec![1u8; 8]]), + records_ingested_so_far: 1, + }); + step.stage_event(SortChunkEvent::AllAnnounced { + slot_count: 1, + memory_chunk_count: 1, + total_records: 1, + }); + step.stage_event(SortChunkEvent::Residual { + chunk: coord_chunk(vec![vec![2u8; 8]]), + records_ingested_so_far: 2, + }); + + let ordinals: Vec = step + .pending + .iter() + .map(|e| match e { + SpillBlockEvent::Block { ordinal, .. } + | SpillBlockEvent::Residual { ordinal, .. } + | SpillBlockEvent::AllAnnounced { ordinal, .. } => *ordinal, + }) + .collect(); + assert_eq!(ordinals, vec![0, 1, 2], "ordinals must be dense across variants"); + assert_eq!(step.next_ordinal, 3); +} + +#[test] +fn an_empty_spill_chunk_is_skipped_rather_than_opening_a_file() { + // A zero-block file would never receive an `is_last_in_file` terminator, so + // `SpillWrite` would be left with a dangling open file forever. + let mut step = SpillGather::new(1 << 20); + step.stage_event(SortChunkEvent::Spill { + seq: 0, + chunk: coord_chunk(Vec::new()), + records_ingested_so_far: 0, + }); + assert!(step.active.is_none(), "an empty chunk must not become the active spill"); + assert!(step.pending.is_empty(), "and must emit nothing"); + assert_eq!(step.next_ordinal, 0, "and must not consume an ordinal"); +} + +#[test] +fn a_non_empty_spill_chunk_becomes_active_without_emitting_yet() { + let mut step = SpillGather::new(1 << 20); + step.stage_event(SortChunkEvent::Spill { + seq: 7, + chunk: coord_chunk(vec![vec![3u8; 32], vec![4u8; 32]]), + records_ingested_so_far: 2, + }); + let active = step.active.as_ref().expect("chunk must be parked for incremental framing"); + assert_eq!(active.file_id, 7, "file_id comes from the spill seq, not write order"); + assert_eq!(active.next_idx, 0); + assert!(step.pending.is_empty(), "framing happens in produce_blocks, not stage_event"); +} + +/// The `TemplateCoordinate` variant is framed by a *different* function than the +/// other three: `frame_record_at` routes it to `c.frame_record_into`, while +/// `Coordinate` / `QuerynameLex` / `QuerynameNatural` share `frame_keyed_record_into`. +/// Every other test here uses `coord_chunk`, so that branch never ran. A layout +/// divergence produces spill files that `SortMerge` misreads — wrong output, not a +/// crash — so the fourth variant needs its own framing coverage. +#[test] +fn template_coordinate_chunks_frame_through_their_own_path() { + use fgumi_sort::{TemplateKey24, TemplateMemChunk}; + + let payloads = [vec![0xA1u8; 24], vec![0xB2u8; 40], vec![0xC3u8; 8]]; + let recs = payloads.iter().map(|b| (TemplateKey24::default(), b.clone())).collect::>(); + let chunk = MemoryChunkErased::TemplateCoordinate(TemplateMemChunk::K24( + InMemoryChunk::from_owned_records(recs), + )); + + // Frame the whole chunk; every record must be emitted exactly once. + let mut blocks = Vec::new(); + let mut next = 0usize; + while next < chunk.len() { + let mut out = Vec::new(); + let consumed = + frame_one_block(&chunk, next, BGZF_MAX_BLOCK_SIZE, &mut out).expect("template framing"); + assert!(consumed > 0, "framing must make progress on every call"); + next += consumed; + blocks.push(out); + } + assert_eq!(next, payloads.len(), "every template record is framed"); + + // Each payload appears in the framed stream, so the template layout carries + // the record bodies through unchanged. + let framed = concat(&blocks); + for (i, p) in payloads.iter().enumerate() { + assert!( + framed.windows(p.len()).any(|w| w == p.as_slice()), + "record {i}'s body must survive template framing" + ); + } +} diff --git a/crates/fgumi-pipeline-io/src/sort/spill_write.rs b/crates/fgumi-pipeline-io/src/sort/spill_write.rs new file mode 100644 index 000000000..56b80525b --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/spill_write.rs @@ -0,0 +1,292 @@ +//! `SpillWrite` — final step of the block-parallel spill-write split +//! (`SpillGather` → `SpillBlockCompress` → `SpillWrite`). +//! +//! `SpillWrite` (`Serial + Affinity::Writer`) receives the compressed +//! [`SpillBlockEvent`]s in dense `ordinal` order (the framework's `ByItemOrdinal` +//! reorder feeds it like `WriteBgzfFile`), demultiplexes `Block`s back to +//! per-`file_id` spill files, and emits the existing [`SortPhase1Event`] so +//! `SortSpillDecompress` / `SortMerge` are unchanged. +//! +//! Because `SortBuffer` (Serial) emits spill chunks one-at-a-time in `seq` order, +//! `SpillGather` (Serial) fans them in order, and the reorder preserves that +//! order, each file's blocks arrive **contiguously** — so `SpillWrite` only ever +//! holds **one** open spill file at a time (`current`). It opens the file on the +//! first block (writing the codec magic), appends each compressed block, and on +//! `is_last_in_file` writes the codec trailer, opens the merge slot, and emits +//! `SpillReady`. This step owns the `TmpDirAllocator` (Serial ⇒ the pick is +//! uncontended) and the RAII temp-dir handles, matching the retired +//! `CompressSpill`'s lifetime contract. + +use std::fs::{File, OpenOptions}; +use std::io::{self, BufWriter, Write}; +use std::path::PathBuf; +use std::sync::Arc; + +use fgumi_bam_io::ProgressTracker; +use fgumi_sort::{SpillCodec, TmpDirAllocator, spill_magic, spill_trailer}; +use parking_lot::Mutex; +use tempfile::TempDir; + +use crate::sort::protocol::{SortPhase1Event, SpillBlockEvent}; +use fgumi_pipeline_core::{ + HeldRetry, Unpushed, + held::HeldSlot, + outputs::Single, + queues::QueueSpec, + reorder::BranchOrdering, + step::{Affinity, DetachedGroup, Step, StepCtx, StepKind, StepOutcome, StepProfile}, +}; + +/// The one spill file currently being written (open from its first block until +/// its `is_last_in_file` block). +struct OpenSpill { + file_id: u32, + path: PathBuf, + writer: BufWriter, +} + +/// `Serial + Affinity::Writer` step that writes per-`file_id` spill files from +/// the compressed block stream and emits `SortPhase1Event`s. +pub struct SpillWrite { + /// Shared temp-directory allocator (free-space-aware round-robin). `Serial`, + /// so the lock is effectively uncontended (one writer worker). + alloc: Arc>, + /// Spill codec for chunk files (bgzf or zstd). + codec: SpillCodec, + /// The currently-open spill file, if any. + current: Option, + held: HeldSlot>, + output_byte_limit: u64, + /// Compressed spill bytes written, logged every 256 MiB under `RUST_LOG=info` + /// so the spill-write rate over wall time is visible alongside ingest. + spill_progress: ProgressTracker, + /// RAII temp-dir handles, held for the step's lifetime so spill files survive + /// while being read by `SortMerge`. Matches `CompressSpill`'s lifetime. + #[allow(dead_code)] + temp_dirs: Arc>, + /// When `true`, advertise `StepKind::Detached` so the framework drives this + /// writer on its own dedicated thread (off the pool) instead of as a + /// pool-scheduled `Serial + Affinity::Writer` step. Set only on the + /// standalone-sort spill path via [`Self::with_detached`] — the exact + /// Phase-1 analogue of Lever 2's detached terminal writer — so the single + /// serial write stream stops consuming a compute worker that could be + /// compressing. Every other chain leaves it `false`. + detached: bool, +} + +impl SpillWrite { + /// Build a `SpillWrite`. `alloc` names spill files across the configured temp + /// dirs; `codec` selects the on-disk format; `temp_dirs` holds the RAII + /// handles alive for the step's lifetime. `output_byte_limit` byte-bounds the + /// forwarded-event output queue. + #[must_use] + pub fn new( + alloc: Arc>, + codec: SpillCodec, + output_byte_limit: u64, + temp_dirs: Arc>, + ) -> Self { + Self { + alloc, + codec, + current: None, + held: HeldSlot::new(), + output_byte_limit, + spill_progress: ProgressTracker::new("Spill bytes written") + .with_interval(256 * 1024 * 1024), + temp_dirs, + detached: false, + } + } + + /// Run this spill writer on its own dedicated `StepKind::Detached` thread + /// instead of as a pool-scheduled `Serial + Affinity::Writer` step. Used + /// ONLY on the standalone-sort spill path (the Phase-1 analogue of Lever 2's + /// detached terminal writer): it frees a pool worker for the + /// compression-bound `SpillBlockCompress` work, matching feat-runall's dedicated + /// spill-I/O thread — but as a single persistent thread for the whole run, + /// not one per spill chunk. + /// + /// The `try_run` body and the bytes it writes are unchanged: the + /// dedicated-thread driver pops blocks in the same `ByItemOrdinal` + /// reorder-stage-ordered sequence, so each spill file's blocks still arrive + /// contiguously (the one-open-file-at-a-time invariant holds) and every + /// spill file is byte-identical to the pool-scheduled writer's output. + /// Affinity is ignored for `Detached`. + #[must_use] + pub fn with_detached(mut self) -> Self { + self.detached = true; + self + } + + fn flush_held(&mut self, ctx: &mut StepCtx<'_, Self>) -> bool { + !matches!(ctx.outputs.retry_held(&mut self.held), HeldRetry::StillHeld) + } + + /// Allocate a spill path for `file_id` (named by the logical spill index so + /// the merge tie-break is independent of write order) and create the file, + /// writing the codec magic prologue. + fn open_file(&self, file_id: u32) -> io::Result { + let base = self.alloc.lock().next().map_err(|e| { + io::Error::other(format!("SpillWrite: temp-dir allocation failed: {e:#}")) + })?; + let path = base.join(format!("chunk_{file_id:04}.keyed")); + // `create_new` fails closed on a duplicate/stale path: a reused `file_id` + // (or a leftover file) must surface as an error rather than truncate an + // existing spill and silently corrupt merge input. + let file = OpenOptions::new().write(true).create_new(true).open(&path)?; + let mut writer = BufWriter::with_capacity(256 * 1024, file); + writer.write_all(spill_magic(self.codec))?; + Ok(OpenSpill { file_id, path, writer }) + } + + /// Process one input event, performing any disk writes and returning the + /// `SortPhase1Event` to emit (a `Block` only emits on `is_last_in_file`). + /// `StepCtx`-free for unit testing. + /// + /// # Errors + /// + /// Propagates file-create / write / slot-open errors. Also errors if a block + /// arrives for a different `file_id` than the open file while one is open + /// without an intervening `is_last_in_file` — a framework-ordering invariant + /// violation that must fail loud rather than corrupt a spill. + fn process_event(&mut self, event: SpillBlockEvent) -> io::Result> { + match event { + SpillBlockEvent::Block { + file_id, + is_last_in_file, + records_ingested_so_far, + bytes, + .. + } => { + // Open the file on its first block; otherwise the open file must + // match (blocks for one file are contiguous in the ordinal stream). + if self.current.is_none() { + self.current = Some(self.open_file(file_id)?); + } + let open = self.current.as_mut().expect("open file set above"); + if open.file_id != file_id { + return Err(io::Error::other(format!( + "SpillWrite: block for file_id {file_id} arrived while file_id {} \ + was still open (blocks must be contiguous per file)", + open.file_id + ))); + } + open.writer.write_all(&bytes)?; + self.spill_progress.log_if_needed(bytes.len() as u64); + + if is_last_in_file { + let OpenSpill { file_id, path, mut writer } = + self.current.take().expect("open file present"); + writer.write_all(spill_trailer(self.codec))?; + writer.flush()?; + drop(writer); // close the fd before opening the read slot + let slot = fgumi_sort::open_spill_slot(&path, file_id).map_err(|e| { + io::Error::other(format!( + "SpillWrite: failed to open spill slot {}: {e:#}", + path.display() + )) + })?; + Ok(Some(SortPhase1Event::SpillReady { slot, path, records_ingested_so_far })) + } else { + Ok(None) + } + } + SpillBlockEvent::Residual { chunk, records_ingested_so_far, .. } => { + self.ensure_no_open_file("residual")?; + // Wrap in a fresh, uniquely-owned `Arc`: the chunk is only ever + // moved (never cloned) onward, so `SortMerge`'s `Arc::try_unwrap` + // invariant holds. + Ok(Some(SortPhase1Event::MemoryChunk { + chunk: Arc::new(chunk), + records_ingested_so_far, + })) + } + SpillBlockEvent::AllAnnounced { + slot_count, memory_chunk_count, total_records, .. + } => { + self.ensure_no_open_file("AllAnnounced")?; + Ok(Some(SortPhase1Event::AllAnnounced { + slot_count, + memory_chunk_count, + total_records, + })) + } + } + } + + /// Error if a spill file is still open. A `Residual` / `AllAnnounced` event, + /// or end-of-stream, while `current` holds an unterminated file means a spill + /// lost its `is_last_in_file` block (a `SpillGather` framing bug) — failing + /// loud avoids dropping a spill or publishing `AllAnnounced` before its + /// `SpillReady`. + fn ensure_no_open_file(&self, at: &str) -> io::Result<()> { + if let Some(open) = &self.current { + return Err(io::Error::other(format!( + "SpillWrite: {at} arrived while spill file_id {} was still open \ + (missing is_last_in_file block)", + open.file_id + ))); + } + Ok(()) + } +} + +impl Step for SpillWrite { + type Input = SpillBlockEvent; + type Outputs = Single; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "SpillWrite", + // Detached (own thread) on the standalone-sort spill path; otherwise + // the default pool-scheduled Serial + sticky writer. `sticky` is + // irrelevant for Detached (it never enters a worker's worklist). + kind: if self.detached { StepKind::Detached } else { StepKind::Serial }, + sticky: true, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], + branch_ordering: vec![BranchOrdering::None], + } + } + + fn detached_group(&self) -> DetachedGroup { + // When detached (standalone-sort spill path), share the sort's I/O + // writer driver thread with the terminal `WriteBgzfFile` — phase-1 spill + // and phase-2 output writes are temporally disjoint (true N+2). Consulted + // only when the step is Detached. + DetachedGroup::Shared(crate::sort::SORT_IO_GROUP) + } + + fn affinity(&self) -> Affinity { + // Ignored for `Detached` (no pool worker drives it); kept for the + // default Serial path where it pins the writer to the last worker. + Affinity::Writer + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if !self.flush_held(ctx) { + return Ok(StepOutcome::Contention); + } + + if let Some(event) = ctx.input.pop() { + if let Some(out) = self.process_event(event)? + && let Err(unpushed) = ctx.outputs.push(out) + { + self.held.put(unpushed); + } + return Ok(StepOutcome::Progress); + } + + if ctx.input.is_drained() { + // End-of-stream with a file still open means the final spill never + // got its `is_last_in_file` block — fail loud rather than leave a + // truncated, unterminated spill on disk. + self.ensure_no_open_file("input drained")?; + return Ok(StepOutcome::Finished); + } + Ok(StepOutcome::NoProgress) + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/fgumi-pipeline-io/src/sort/spill_write/tests.rs b/crates/fgumi-pipeline-io/src/sort/spill_write/tests.rs new file mode 100644 index 000000000..3ddee1596 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/spill_write/tests.rs @@ -0,0 +1,280 @@ +//! Unit tests for `SpillWrite::process_event` (the `StepCtx`-free core): per-file +//! demux, codec magic/trailer bracketing, finalization on `is_last_in_file`, and +//! event mapping. Byte-exact readback of the assembled file is gated end-to-end +//! by the full-sort parity test (the production reader is crate-private to +//! `fgumi-sort`). + +use super::*; +use crate::sort::protocol::MemoryChunkErased; +use fgumi_sort::{InMemoryChunk, RawCoordinateKey, SpillBlockCompressor}; +use rstest::rstest; +use tempfile::TempDir; + +/// Build a `SpillWrite` writing into a fresh temp dir; returns it plus the dir +/// (kept alive by the caller) so written files can be inspected. +fn make_writer(codec: SpillCodec) -> (SpillWrite, TempDir) { + let dir = TempDir::new().unwrap(); + let alloc = TmpDirAllocator::new(vec![dir.path().to_path_buf()]).unwrap(); + let writer = SpillWrite::new(Arc::new(Mutex::new(alloc)), codec, 1 << 20, Arc::new(Vec::new())); + (writer, dir) +} + +/// Kernel-compress one raw block for `codec`, mirroring `SpillBlockCompress`. +fn compress(codec: SpillCodec, raw: &[u8]) -> Vec { + SpillBlockCompressor::new(codec, 1).unwrap().compress_block(raw).unwrap() +} + +fn block(codec: SpillCodec, file_id: u32, is_last: bool, raw: &[u8]) -> SpillBlockEvent { + SpillBlockEvent::Block { + ordinal: 0, + file_id, + is_last_in_file: is_last, + records_ingested_so_far: 42, + bytes: compress(codec, raw), + } +} + +#[rstest] +#[case(SpillCodec::Zstd)] +#[case(SpillCodec::Bgzf)] +fn non_last_block_opens_file_emits_nothing_last_block_emits_spill_ready(#[case] codec: SpillCodec) { + let (mut w, dir) = make_writer(codec); + // First (non-last) block: file opens, no event. + let out = w.process_event(block(codec, 5, false, &[1u8; 32])).unwrap(); + assert!(out.is_none(), "non-last block emits no event ({codec:?})"); + assert!(w.current.is_some(), "file must be open after first block ({codec:?})"); + + // Last block: trailer written, slot opened, SpillReady emitted. + let out = w.process_event(block(codec, 5, true, &[2u8; 32])).unwrap(); + let Some(SortPhase1Event::SpillReady { slot, path, records_ingested_so_far }) = out else { + panic!("expected SpillReady ({codec:?})"); + }; + assert!(w.current.is_none(), "file closed after last block ({codec:?})"); + assert_eq!(slot.file_id, 5, "slot file_id == logical seq ({codec:?})"); + assert_eq!(records_ingested_so_far, 42); + assert!(path.exists(), "spill file exists ({codec:?})"); + assert!(path.starts_with(dir.path()), "spill file under temp dir ({codec:?})"); + assert_eq!(slot.codec, codec, "codec detected from written magic ({codec:?})"); +} + +#[test] +fn distinct_file_ids_produce_distinct_files() { + let codec = SpillCodec::Zstd; + let (mut w, _dir) = make_writer(codec); + // File 0 (single block), then file 1 (single block) — contiguous per file. + let r0 = w.process_event(block(codec, 0, true, &[7u8; 16])).unwrap().unwrap(); + let r1 = w.process_event(block(codec, 1, true, &[8u8; 16])).unwrap().unwrap(); + let ( + SortPhase1Event::SpillReady { path: p0, slot: s0, .. }, + SortPhase1Event::SpillReady { path: p1, slot: s1, .. }, + ) = (r0, r1) + else { + panic!("expected two SpillReady events"); + }; + assert_ne!(p0, p1, "distinct file_ids must yield distinct paths"); + assert_eq!(s0.file_id, 0); + assert_eq!(s1.file_id, 1); +} + +#[test] +fn block_for_wrong_file_id_while_open_errors() { + let codec = SpillCodec::Zstd; + let (mut w, _dir) = make_writer(codec); + // Open file 0 with a non-last block, then feed a block for file 1 — a + // contiguity violation that must fail loud, not silently corrupt file 0. + w.process_event(block(codec, 0, false, &[1u8; 16])).unwrap(); + // `SortPhase1Event` is not `Debug`, so match instead of `unwrap_err`. + match w.process_event(block(codec, 1, false, &[2u8; 16])) { + Err(err) => { + assert!( + err.to_string().contains("contiguous"), + "expected contiguity error, got: {err}" + ); + } + Ok(_) => panic!("a block for a different open file_id must error"), + } +} + +#[test] +fn residual_while_file_open_errors() { + let codec = SpillCodec::Zstd; + let (mut w, _dir) = make_writer(codec); + // Open a file with a non-last block, then feed a Residual — a missing + // is_last_in_file terminator must fail loud, not drop the open spill. + w.process_event(block(codec, 0, false, &[1u8; 16])).unwrap(); + let chunk = MemoryChunkErased::Coordinate(InMemoryChunk::from_owned_records(vec![( + RawCoordinateKey { sort_key: 1 }, + vec![9u8; 8], + )])); + match w.process_event(SpillBlockEvent::Residual { + ordinal: 1, + chunk, + records_ingested_so_far: 1, + }) { + Err(err) => { + assert!(err.to_string().contains("still open"), "expected open-file error, got: {err}"); + } + Ok(_) => panic!("residual while a spill file is open must error"), + } +} + +#[test] +fn default_is_serial_writer_and_with_detached_flips_to_detached() { + use fgumi_pipeline_core::step::{Affinity, Step, StepKind}; + let (w, _dir) = make_writer(SpillCodec::Zstd); + // Default: pool-scheduled Serial + Affinity::Writer (pinned to worker N-1). + assert_eq!(w.profile().kind, StepKind::Serial, "default spill writer is Serial"); + assert_eq!(w.affinity(), Affinity::Writer, "default spill writer pins to the writer worker"); + // `with_detached()` flips only the advertised kind to Detached (own thread, + // off the pool); the write body — and hence the bytes it writes — is + // unchanged, so full-sort parity still covers the on-disk format. + let wd = w.with_detached(); + assert_eq!(wd.profile().kind, StepKind::Detached, "with_detached flips kind to Detached"); +} + +#[test] +fn residual_maps_to_memory_chunk_and_announced_passes_through() { + let codec = SpillCodec::Zstd; + let (mut w, _dir) = make_writer(codec); + + let chunk = MemoryChunkErased::Coordinate(InMemoryChunk::from_owned_records(vec![( + RawCoordinateKey { sort_key: 1 }, + vec![9u8; 8], + )])); + let out = w + .process_event(SpillBlockEvent::Residual { ordinal: 0, chunk, records_ingested_so_far: 3 }) + .unwrap(); + let Some(SortPhase1Event::MemoryChunk { chunk, records_ingested_so_far }) = out else { + panic!("expected MemoryChunk"); + }; + assert_eq!(records_ingested_so_far, 3); + assert_eq!(Arc::strong_count(&chunk), 1, "residual chunk wrapped in a fresh unique Arc"); + + let out = w + .process_event(SpillBlockEvent::AllAnnounced { + ordinal: 1, + slot_count: 4, + memory_chunk_count: 1, + total_records: 500, + }) + .unwrap(); + assert!(matches!( + out, + Some(SortPhase1Event::AllAnnounced { + slot_count: 4, + memory_chunk_count: 1, + total_records: 500, + }) + )); +} + +// ── Step wiring: profile / affinity / detached group ───────────────────────── + +#[test] +fn profile_defaults_to_a_pool_scheduled_serial_writer() { + let (w, _dir) = make_writer(SpillCodec::Zstd); + let profile = w.profile(); + assert_eq!(profile.name, "SpillWrite"); + assert_eq!(profile.kind, StepKind::Serial, "default is the pool-scheduled writer"); + assert!(profile.sticky); + assert_eq!(profile.branch_ordering, vec![BranchOrdering::None]); + match profile.output_queues.as_slice() { + [QueueSpec::ByteBounded { limit_bytes }] => assert_eq!(*limit_bytes, 1 << 20), + other => panic!("expected one byte-bounded queue, got {other:?}"), + } + // Affinity pins the pool-scheduled writer; ignored once detached. + assert_eq!(w.affinity(), Affinity::Writer); +} + +#[test] +fn with_detached_flips_only_the_step_kind() { + let (w, _dir) = make_writer(SpillCodec::Zstd); + let before = w.profile(); + let detached = w.with_detached(); + let after = detached.profile(); + + assert_eq!(before.kind, StepKind::Serial); + assert_eq!(after.kind, StepKind::Detached, "detached runs on its own thread"); + // Everything else about the step is unchanged — the doc promises the + // `try_run` body and the bytes written are identical either way. + assert_eq!(after.name, before.name); + assert_eq!(after.sticky, before.sticky); + assert_eq!(after.branch_ordering, before.branch_ordering); + assert_eq!(detached.affinity(), Affinity::Writer); +} + +#[test] +fn detached_writer_shares_the_sort_io_group() { + // Phase-1 spill and phase-2 output writes are temporally disjoint, so both + // ride the same driver thread rather than each taking one. + let (w, _dir) = make_writer(SpillCodec::Zstd); + assert_eq!(w.detached_group(), DetachedGroup::Shared(crate::sort::SORT_IO_GROUP)); + let (w2, _dir2) = make_writer(SpillCodec::Bgzf); + assert_eq!( + w2.with_detached().detached_group(), + DetachedGroup::Shared(crate::sort::SORT_IO_GROUP) + ); +} + +// ── Open-file bookkeeping ──────────────────────────────────────────────────── + +#[test] +fn ensure_no_open_file_passes_when_idle_and_fails_while_a_file_is_open() { + let (mut w, _dir) = make_writer(SpillCodec::Zstd); + w.ensure_no_open_file("Residual").expect("idle writer has no open file"); + + // Opening a file without its is_last block leaves it dangling. + let out = w.process_event(block(SpillCodec::Zstd, 3, false, &[7u8; 16])).unwrap(); + assert!(out.is_none()); + assert!(w.current.is_some()); + + let err = w.ensure_no_open_file("AllAnnounced").expect_err("dangling file must fail closed"); + let msg = err.to_string(); + assert!(msg.contains("AllAnnounced"), "error names the offending event: {msg}"); + assert!(msg.contains("file_id 3"), "error names the open file: {msg}"); +} + +#[test] +fn open_file_refuses_to_reuse_an_existing_path() { + let (w, dir) = make_writer(SpillCodec::Zstd); + // First open succeeds and creates the file on disk. + let opened = w.open_file(9).expect("first open succeeds"); + drop(opened); + assert!(dir.path().join("chunk_0009.keyed").exists(), "spill file is created eagerly"); + + // A reused file_id must fail closed rather than truncate the existing file: + // silently overwriting a spill would drop records from the merge. + // `OpenSpill` is not `Debug`, so match instead of using `expect_err`. + match w.open_file(9) { + Ok(_) => panic!("reusing a file_id must fail"), + Err(e) => assert_eq!(e.kind(), io::ErrorKind::AlreadyExists), + } +} + +/// `AllAnnounced` arriving while a spill file is still open must fail closed. +/// +/// Without the guard, `AllAnnounced` reaches `SortMerge` before the matching +/// `SpillReady`, so the merge starts against an undercounted slot set and +/// silently drops a spill file's records. +#[test] +fn all_announced_while_a_file_is_open_fails_closed() { + let codec = SpillCodec::Zstd; + let (mut w, _dir) = make_writer(codec); + // Open file 0 and never terminate it with an is_last_in_file block. + w.process_event(block(codec, 0, false, &[1u8; 16])).unwrap(); + + match w.process_event(SpillBlockEvent::AllAnnounced { + ordinal: 1, + slot_count: 1, + memory_chunk_count: 0, + total_records: 1, + }) { + Err(err) => { + let msg = err.to_string(); + assert!(msg.contains("AllAnnounced"), "error names the event: {msg}"); + assert!(msg.contains("still open"), "error names the cause: {msg}"); + assert!(msg.contains("file_id 0"), "error names the open file: {msg}"); + } + Ok(_) => panic!("AllAnnounced while a spill file is open must error"), + } +} diff --git a/crates/fgumi-pipeline-io/src/sort/tests.rs b/crates/fgumi-pipeline-io/src/sort/tests.rs new file mode 100644 index 000000000..382859b6d --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/tests.rs @@ -0,0 +1,1930 @@ +//! Tests for the runall-sort chains. +//! +//! The record-input chain is `SortBuffer` → `CompressSpill` → +//! `SortSpillDecompress` → `SortMerge`; the legacy `SortAndSpill` Phase-1 head +//! it replaced was retired in P7. The block-input arena front (`ReadBlocks` → +//! `InflateToArena` → `FindBoundariesAndSort`) is covered at the end of this +//! module. +//! +//! `RawExternalSorter::sort` (driven here via [`sort_via_legacy`]) is retained +//! as the parity oracle both chains are validated against. + +use std::io; +use std::sync::Arc; + +use anyhow::Result; +use fgumi_raw_bam::RawRecord; +use fgumi_raw_bam::testutil::make_bam_bytes; +use fgumi_sort::{QuerynameComparator, RawExternalSorter, SortOrder, SpillCodec}; +use noodles::sam::Header; +use parking_lot::Mutex; +use rstest::rstest; + +use super::*; +use crate::sort::protocol::SortChunkEvent; +use crate::types::RecordBatch; +use fgumi_pipeline_core::{ + Unpushed, + builder::{Pipeline, PipelineConfig}, + held::HeldSlot, + outputs::OrderedBytesSingle, + queues::QueueSpec, + reorder::BranchOrdering, + step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}, +}; + +// ── In-memory source / sink test steps ────────────────────────────────────── + +/// `Exclusive` source that drains a `Vec` one batch per `try_run` call. +struct VecSource { + batches: Vec, + held: HeldSlot>, + output_byte_limit: u64, +} + +impl VecSource { + fn new(mut batches: Vec, output_byte_limit: u64) -> Self { + batches.reverse(); + Self { batches, held: HeldSlot::new(), output_byte_limit } + } +} + +impl Step for VecSource { + type Input = (); + type Outputs = OrderedBytesSingle; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "VecSource", + kind: StepKind::Exclusive, + sticky: true, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], + branch_ordering: vec![BranchOrdering::ByItemOrdinal], + } + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if let Some(unpushed) = self.held.take() { + match ctx.outputs.retry(unpushed) { + Ok(()) => {} + Err(again) => { + self.held.put(again); + return Ok(StepOutcome::Progress); + } + } + } + let Some(batch) = self.batches.pop() else { + return Ok(StepOutcome::Finished); + }; + match ctx.outputs.push(batch) { + Ok(()) => Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held.put(unpushed); + Ok(StepOutcome::Progress) + } + } + } +} + +/// Sink that appends every received batch into a shared `Vec`. +struct VecSink { + received: Arc>>, + kind: StepKind, +} + +impl Step for VecSink { + type Input = RecordBatch; + type Outputs = (); + + fn profile(&self) -> StepProfile { + StepProfile { + name: "VecSink", + kind: self.kind, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + match ctx.input.pop() { + Some(batch) => { + self.received.lock().push(batch); + Ok(StepOutcome::Progress) + } + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } +} + +/// One sort's output as a flat list of raw BAM record-byte payloads, in output +/// order. Both the streaming pipeline and the legacy oracle produce this shape. +type RecordBytes = Vec>; + +// ── Synthetic-record helpers ──────────────────────────────────────────────── + +fn synthesize_records(n: usize, seed: u64) -> (Header, Vec) { + synthesize_sized_records(n, seed, 0) +} + +fn pack_batches(records: &[RawRecord], batch_size: usize) -> Vec { + use crate::types::RecordBatchBuilder; + records + .chunks(batch_size) + .enumerate() + .map(|(i, chunk)| { + let total: usize = chunk.iter().map(RawRecord::len).sum(); + let mut b = RecordBatchBuilder::with_capacity(i as u64, total, chunk.len()); + for r in chunk { + b.push_record_bytes(r.as_ref()); + } + b.build() + }) + .collect() +} + +fn drive_sort_pipeline( + sorter: RawExternalSorter, + header: &Header, + batches: Vec, + output_byte_limit: u64, + threads: usize, + sink_kind: StepKind, +) -> Result>> { + drive_sort_pipeline_tuned( + sorter, + header, + batches, + output_byte_limit, + threads, + sink_kind, + SortDecompressTuning::default(), + SpillCodec::Zstd, + ) +} + +/// Drive the production sort chain (`VecSource` → `SortBuffer` → +/// `CompressSpill` → `SortSpillDecompress` → `SortMerge` → `VecSink`). The +/// legacy `SortAndSpill` Phase-1 head was retired in P7, so the Phase-2 tests +/// (decompress-granularity, out-of-order, soak) run through the same production +/// chain `fgumi sort` / `runall` build. +#[allow(clippy::too_many_arguments)] +fn drive_sort_pipeline_tuned( + sorter: RawExternalSorter, + header: &Header, + batches: Vec, + output_byte_limit: u64, + threads: usize, + sink_kind: StepKind, + decompress_tuning: SortDecompressTuning, + spill_codec: SpillCodec, +) -> Result>> { + drive_sort_buffer_pipeline( + sorter, + header, + batches, + output_byte_limit, + threads, + sink_kind, + decompress_tuning, + spill_codec, + ) +} + +/// Drive the P6 four-step buffer chain +/// (`VecSource` → `SortBuffer` → `CompressSpill` → `SortSpillDecompress` → +/// `SortMerge` → `VecSink`) and collect the merged record bytes. Exercises +/// every sort order `SortBuffer` supports (see +/// `sort_buffer_chain_matches_legacy_all_orders`). +#[allow(clippy::too_many_arguments)] +fn drive_sort_buffer_pipeline( + sorter: RawExternalSorter, + header: &Header, + batches: Vec, + output_byte_limit: u64, + threads: usize, + sink_kind: StepKind, + decompress_tuning: SortDecompressTuning, + spill_codec: SpillCodec, +) -> Result>> { + use fgumi_sort::TmpDirAllocator; + + let received: Arc>> = Arc::new(Mutex::new(Vec::new())); + let sort_order = sorter.sort_order(); + + // Temp dir + allocator for CompressSpill, held alive by the step. The + // deterministic always-ample probe avoids any dependency on host free space. + let dir = tempfile::TempDir::new()?; + let alloc = + TmpDirAllocator::with_probe(vec![dir.path().to_path_buf()], Box::new(|_| Ok(u64::MAX)), 0)?; + let temp_dirs = Arc::new(vec![dir]); + + let source = VecSource::new(batches, output_byte_limit); + let sort_buffer = SortBuffer::from_sorter(sorter, header, output_byte_limit)?; + // Codec/compression affect only intermediate spill bytes, not the final + // sorted records, so any codec yields output parity. The caller passes the + // codec so codec-specific tests (e.g. the BGZF block-parallel parity test) + // actually exercise their codec end-to-end. + let compress = CompressSpill::new( + Arc::new(Mutex::new(alloc)), + spill_codec, + 3, + output_byte_limit, + temp_dirs, + ); + let decompress = SortSpillDecompress::new(output_byte_limit, decompress_tuning); + let merge = + SortMerge::::with_target_batch_count(sort_order, output_byte_limit, 256); + let sink = VecSink { received: Arc::clone(&received), kind: sink_kind }; + + let builder = Pipeline::builder(); + builder + .chain(source) + .chain(sort_buffer) + .chain(compress) + .chain(decompress) + .chain(merge) + .chain(sink) + .into_sink_marker(); + let pipeline = builder.build()?; + pipeline.run(PipelineConfig { threads, ..Default::default() })?; + + let collected = std::mem::take(&mut *received.lock()); + let mut out = Vec::new(); + for batch in collected { + for bytes in batch.iter_record_bytes() { + out.push(bytes.to_vec()); + } + } + Ok(out) +} + +/// Terminal-path sink: collects the [`DecompressedBlock`]s emitted by +/// `SortMerge` (the framed-bytes terminal output, lever 1). +struct BlockSink { + received: Arc>>, + kind: StepKind, +} + +impl Step for BlockSink { + type Input = crate::types::DecompressedBlock; + type Outputs = (); + + fn profile(&self) -> StepProfile { + StepProfile { + name: "BlockSink", + kind: self.kind, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + match ctx.input.pop() { + Some(block) => { + self.received.lock().push(block); + Ok(StepOutcome::Progress) + } + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } +} + +/// Parse a `[u32 LE block_size][body]`-framed byte buffer (the terminal-sort +/// `BlockOutput` framing, identical to the former `SerializeRecordBatch`) into +/// its record bodies. Mirrors what `BgzfDecompress → FindBamBoundaries` does +/// downstream; kept local to the test so parity is checked against an +/// independent re-implementation of the layout. +fn unframe_block_records(bytes: &[u8]) -> Vec> { + let mut out = Vec::new(); + let mut i = 0usize; + while i < bytes.len() { + let len = u32::from_le_bytes(bytes[i..i + 4].try_into().unwrap()) as usize; + i += 4; + out.push(bytes[i..i + len].to_vec()); + i += len; + } + out +} + +/// Drive the production sort chain with the **terminal** merge +/// (`SortMerge` → `BlockSink`) and recover the merged record +/// bodies by un-framing each `DecompressedBlock`. Used to prove the framed +/// terminal output carries byte-identical records to the `RecordBatchOutput` +/// path / legacy oracle (lever 1 parity gate). +#[allow(clippy::too_many_arguments)] +fn drive_sort_block_pipeline( + sorter: RawExternalSorter, + header: &Header, + batches: Vec, + output_byte_limit: u64, + threads: usize, + sink_kind: StepKind, + decompress_tuning: SortDecompressTuning, + spill_codec: SpillCodec, +) -> Result>> { + use fgumi_sort::TmpDirAllocator; + + let received: Arc>> = + Arc::new(Mutex::new(Vec::new())); + let sort_order = sorter.sort_order(); + + let dir = tempfile::TempDir::new()?; + let alloc = + TmpDirAllocator::with_probe(vec![dir.path().to_path_buf()], Box::new(|_| Ok(u64::MAX)), 0)?; + let temp_dirs = Arc::new(vec![dir]); + + let source = VecSource::new(batches, output_byte_limit); + let sort_buffer = SortBuffer::from_sorter(sorter, header, output_byte_limit)?; + let compress = CompressSpill::new( + Arc::new(Mutex::new(alloc)), + spill_codec, + 3, + output_byte_limit, + temp_dirs, + ); + let decompress = SortSpillDecompress::new(output_byte_limit, decompress_tuning); + let merge = + SortMerge::::with_target_batch_count(sort_order, output_byte_limit, 256); + let sink = BlockSink { received: Arc::clone(&received), kind: sink_kind }; + + let builder = Pipeline::builder(); + builder + .chain(source) + .chain(sort_buffer) + .chain(compress) + .chain(decompress) + .chain(merge) + .chain(sink) + .into_sink_marker(); + let pipeline = builder.build()?; + pipeline.run(PipelineConfig { threads, ..Default::default() })?; + + let mut collected = std::mem::take(&mut *received.lock()); + // `Detached` collapses ordering to None, so blocks arrive in merge-emit + // order on the single sink; sort by serial defensively in case a future + // change makes the sink parallel. + collected.sort_by_key(|b| b.batch_serial); + let mut out = Vec::new(); + for block in &collected { + out.extend(unframe_block_records(&block.bytes)); + } + Ok(out) +} + +// ── Reference: RawExternalSorter::sort to bytes ───────────────────────────── + +fn sort_via_legacy( + sort_order: SortOrder, + header: &Header, + records: &[RawRecord], + memory_limit: usize, + threads: usize, +) -> Result>> { + let tmp_in = tempfile::NamedTempFile::new()?; + { + let mut writer = fgumi_bam_io::create_raw_bam_writer(tmp_in.path(), header, 1, 1)?; + for r in records { + writer.write_raw_record(r.as_ref())?; + } + writer.finish()?; + } + + let tmp_out = tempfile::NamedTempFile::new()?; + let sorter = RawExternalSorter::new(sort_order) + .memory_limit(memory_limit) + .threads(threads) + .output_compression(1) + .temp_compression(1); + sorter.sort(tmp_in.path(), tmp_out.path())?; + + let (mut reader, _hdr) = fgumi_bam_io::create_raw_bam_reader_with_opts( + tmp_out.path(), + 1, + fgumi_bam_io::PipelineReaderOpts::default(), + )?; + let mut record = RawRecord::default(); + let mut out = Vec::new(); + loop { + match reader.read_record(&mut record)? { + 0 => break, + _ => out.push(record.as_ref().to_vec()), + } + } + Ok(out) +} + +// ── Tests ─────────────────────────────────────────────────────────────────── +// +// (The former `three_step_chain_{in_memory,multi_spill}_path_matches_legacy` +// tests, which drove the retired `SortAndSpill` head, are subsumed by +// `sort_buffer_chain_matches_legacy_all_orders` below — same orders, same +// regimes, against the same `sort_via_legacy` oracle, but through the +// production `SortBuffer` → `CompressSpill` chain.) + +// ── P6 buffer-chain parity (SortBuffer → CompressSpill → decompress → merge) ─ + +/// The P6 four-step chain must produce byte-identical output to the legacy +/// oracle for EVERY sort order, across the in-memory and multi-spill regimes at +/// 2 and 4 threads. This is the inc-4/5 parity gate: it exercises `SortBuffer`'s +/// streaming-spill emission, `CompressSpill`'s inline writes + slot opens +/// (`file_id == seq`), and the single-residual fast path all the way through the +/// real `SortSpillDecompress` + `SortMerge` — for coordinate, template, and both +/// queryname comparators. +#[rstest] +#[case::coord_inmem_t2(SortOrder::Coordinate, 5_000, 256 * 1024 * 1024, 2)] +#[case::coord_inmem_t4(SortOrder::Coordinate, 5_000, 256 * 1024 * 1024, 4)] +#[case::coord_spill_t2(SortOrder::Coordinate, 20_000, 256 * 1024, 2)] +#[case::coord_spill_t4(SortOrder::Coordinate, 20_000, 256 * 1024, 4)] +#[case::template_inmem(SortOrder::TemplateCoordinate, 5_000, 256 * 1024 * 1024, 2)] +#[case::template_spill(SortOrder::TemplateCoordinate, 20_000, 256 * 1024, 2)] +#[case::template_spill_t4(SortOrder::TemplateCoordinate, 20_000, 256 * 1024, 4)] +// Many-spill regime (~20+ spill runs at a 128 KiB buffer): exercises the deeper +// per-slot FIFO (cap 32) and the emptiest-first refill order across many slots, where +// the merge must still emit byte-identical output to the oracle. +#[case::coord_manyspill_t4(SortOrder::Coordinate, 60_000, 128 * 1024, 4)] +#[case::template_manyspill_t4(SortOrder::TemplateCoordinate, 60_000, 128 * 1024, 4)] +#[case::qname_lex_inmem(SortOrder::Queryname(QuerynameComparator::Lexicographic), 5_000, 256 * 1024 * 1024, 2)] +#[case::qname_lex_spill(SortOrder::Queryname(QuerynameComparator::Lexicographic), 20_000, 256 * 1024, 2)] +#[case::qname_nat_inmem(SortOrder::Queryname(QuerynameComparator::Natural), 5_000, 256 * 1024 * 1024, 2)] +#[case::qname_nat_spill(SortOrder::Queryname(QuerynameComparator::Natural), 20_000, 256 * 1024, 2)] +fn sort_buffer_chain_matches_legacy_all_orders( + #[case] sort_order: SortOrder, + #[case] n: usize, + #[case] memory_limit: usize, + #[case] threads: usize, +) { + let (header, records) = synthesize_records(n, 0x5A17_C0DE); + let sorter = RawExternalSorter::new(sort_order) + .memory_limit(memory_limit) + .threads(threads) + .output_compression(1) + .temp_compression(1); + let new_out = drive_sort_buffer_pipeline( + sorter, + &header, + pack_batches(&records, 256), + 4 * 1024 * 1024, + // Run the pipeline with exactly the case's thread count — production + // sort uses `num_threads` workers (no `max(3)` floor), so the parity + // test must too. + threads, + StepKind::Exclusive, + SortDecompressTuning::default(), + SpillCodec::Zstd, + ) + .expect("buffer pipeline drives to completion"); + + let legacy_out = + sort_via_legacy(sort_order, &header, &records, memory_limit, threads).expect("legacy"); + + assert_eq!(new_out.len(), legacy_out.len(), "{sort_order:?} record count mismatch"); + if is_stable_for_equal_keys(sort_order) { + // Coordinate / template-coordinate use a stable radix sort: equal keys + // keep input order deterministically, so output is byte-for-byte equal. + for (i, (got, want)) in new_out.iter().zip(legacy_out.iter()).enumerate() { + assert_eq!(got, want, "{sort_order:?} record {i} bytes differ"); + } + } else { + // Queryname uses an UNSTABLE comparator sort: the order of equal keys + // (same name + segment flags) is unspecified and differs between the + // parallel buffer chain and the `.sort()` oracle (and run-to-run). The + // sound parity claim is multiset equality — same records, possibly in a + // different equal-key order. Sortedness is covered by the production + // `queryname_*_sort_matrix` integration tests. + let mut got = new_out.clone(); + let mut want = legacy_out.clone(); + got.sort_unstable(); + want.sort_unstable(); + assert_eq!(got, want, "{sort_order:?} record multiset differs from the oracle"); + } +} + +/// Raw BAM record carrying `MI:i:` (aux tag bytes + `i` type + i32 LE). +/// Shared with `sort_buffer`'s unit tests, which drive the same dropped-lane +/// rejection directly against `ingest_batch_records`. +pub(super) fn record_with_mi(pos: i32, name: &[u8], mi: i32) -> Vec { + let mut aux = vec![b'M', b'I', b'i']; + aux.extend_from_slice(&mi.to_le_bytes()); + make_bam_bytes(0, pos, 0, name, &[], 40, -1, -1, &aux) +} + +/// An ingest failure must fail the whole pipeline, never produce a partial sort. +/// `SortBuffer` tracks "still ingesting" as `sorter.is_some()`, so a failure that +/// left that slot empty would be indistinguishable from "finalized" and could +/// surface as a clean `Finished` — a truncated result with a valid structure. +/// Drives the one reachable `ChunkSorter::push` failure: under +/// `--key-types none` the first record fixes the narrowed lane set, so a later +/// record with a differing MI is rejected. +#[test] +fn sort_buffer_chain_fails_the_pipeline_on_a_dropped_lane_violation() { + use fgumi_sort::KeyTypesSpec; + + let records: Vec = (0..8i32) + .map(|i| { + // Record 0 fixes the variant with MI 1; record 4 violates it. + let mi = if i == 4 { 2 } else { 1 }; + RawRecord::from(record_with_mi(100 - i, format!("rd_{i}").as_bytes(), mi)) + }) + .collect(); + + let sorter = RawExternalSorter::new(SortOrder::TemplateCoordinate) + .memory_limit(256 * 1024 * 1024) + .threads(1) + .key_types(KeyTypesSpec::None); + let err = drive_sort_buffer_pipeline( + sorter, + &Header::default(), + pack_batches(&records, 256), + 4 * 1024 * 1024, + 1, + StepKind::Exclusive, + SortDecompressTuning::default(), + SpillCodec::Zstd, + ) + .expect_err("a dropped-lane violation must fail the pipeline, not truncate the sort"); + + let message = format!("{err:#}"); + assert!( + message.contains("SortBuffer: push failed"), + "the ingest failure must reach the caller: {message}" + ); +} + +/// Regression guard for the `SortBuffer` peak-memory invariant: when a single +/// input `RecordBatch` is larger than `memory_limit`, `ingest_one_batch` seals +/// several chunks in one call (staging multiple `Spill` events into `pending` +/// before `emit_pending` runs). That transiently exceeds the "~one spill chunk" +/// production bound, but it MUST stay correct — every record emitted in sorted +/// order, none stranded by the loop. Packs all records into ONE oversized batch +/// against a small `memory_limit` (the case the 256-record-per-batch parity +/// matrix never hits) and asserts byte-for-byte parity with the legacy oracle +/// (coordinate is stable, so equal keys keep input order deterministically). +#[test] +fn sort_buffer_single_oversized_batch_seals_multiple_chunks_without_dropping() { + let sort_order = SortOrder::Coordinate; + let n = 20_000; + // Far below the single batch's byte size, so the one batch seals many chunks. + let memory_limit = 128 * 1024; + let threads = 2; + let (header, records) = synthesize_records(n, 0x0B16_BA7C); + + // One batch holding every record — forces multiple seals per ingest call. + let batches = pack_batches(&records, records.len()); + assert_eq!(batches.len(), 1, "test must drive a single oversized batch"); + + let sorter = RawExternalSorter::new(sort_order) + .memory_limit(memory_limit) + .threads(threads) + .output_compression(1) + .temp_compression(1); + let new_out = drive_sort_buffer_pipeline( + sorter, + &header, + batches, + 4 * 1024 * 1024, + threads, + StepKind::Exclusive, + SortDecompressTuning::default(), + SpillCodec::Zstd, + ) + .expect("buffer pipeline drives to completion"); + + let legacy_out = + sort_via_legacy(sort_order, &header, &records, memory_limit, threads).expect("legacy"); + + assert_eq!(new_out.len(), records.len(), "every record must survive — none dropped"); + assert_eq!(new_out.len(), legacy_out.len(), "record count matches oracle"); + for (i, (got, want)) in new_out.iter().zip(legacy_out.iter()).enumerate() { + assert_eq!(got, want, "record {i} bytes differ from oracle"); + } +} + +/// Lever 1 parity gate: the terminal `SortMerge` path (framed +/// `DecompressedBlock`s, wired straight to `BgzfCompress`) must carry +/// byte-identical records to the legacy oracle for every sort order, across the +/// in-memory and spill regimes. The framed bytes are `[u32 LE len][body]` per +/// record (identical to the removed `SerializeRecordBatch`); un-framing them +/// recovers the same record bodies the `RecordBatchOutput` path emits. This is +/// the gate proving the merge-side framing did not change the output bytes. +#[rstest] +#[case::coord_inmem(SortOrder::Coordinate, 5_000, 256 * 1024 * 1024, 2)] +#[case::coord_spill(SortOrder::Coordinate, 20_000, 256 * 1024, 4)] +#[case::template_inmem(SortOrder::TemplateCoordinate, 5_000, 256 * 1024 * 1024, 2)] +#[case::template_spill(SortOrder::TemplateCoordinate, 20_000, 256 * 1024, 4)] +#[case::qname_lex_spill(SortOrder::Queryname(QuerynameComparator::Lexicographic), 20_000, 256 * 1024, 2)] +#[case::qname_nat_spill(SortOrder::Queryname(QuerynameComparator::Natural), 20_000, 256 * 1024, 2)] +fn sort_merge_block_output_matches_legacy( + #[case] sort_order: SortOrder, + #[case] n: usize, + #[case] memory_limit: usize, + #[case] threads: usize, +) { + let (header, records) = synthesize_records(n, 0x5A17_C0DE); + let make_sorter = || { + RawExternalSorter::new(sort_order) + .memory_limit(memory_limit) + .threads(threads) + .output_compression(1) + .temp_compression(1) + }; + + // Terminal framed-block path. + let block_out = drive_sort_block_pipeline( + make_sorter(), + &header, + pack_batches(&records, 256), + 4 * 1024 * 1024, + threads, + StepKind::Exclusive, + SortDecompressTuning::default(), + SpillCodec::Zstd, + ) + .expect("block pipeline drives to completion"); + + // Intermediate RecordBatch path — the framed bytes must un-frame to exactly + // the records this path emits (cross-check the two SortMerge framings agree). + let batch_out = drive_sort_buffer_pipeline( + make_sorter(), + &header, + pack_batches(&records, 256), + 4 * 1024 * 1024, + threads, + StepKind::Exclusive, + SortDecompressTuning::default(), + SpillCodec::Zstd, + ) + .expect("buffer pipeline drives to completion"); + + let legacy_out = + sort_via_legacy(sort_order, &header, &records, memory_limit, threads).expect("legacy"); + + assert_eq!(block_out.len(), legacy_out.len(), "{sort_order:?} record count vs legacy"); + assert_eq!(block_out.len(), batch_out.len(), "{sort_order:?} record count vs RecordBatch path"); + + if is_stable_for_equal_keys(sort_order) { + // Stable orders: byte-for-byte identical output, in order. + for (i, (got, want)) in block_out.iter().zip(legacy_out.iter()).enumerate() { + assert_eq!(got, want, "{sort_order:?} record {i} (block path vs legacy) differs"); + } + assert_eq!(block_out, batch_out, "{sort_order:?} block vs RecordBatch path bytes differ"); + } else { + // Queryname's comparator sort is unstable on equal keys; assert multiset + // equality (same records, possibly different equal-key order). + let mut got = block_out.clone(); + let mut want = legacy_out.clone(); + got.sort_unstable(); + want.sort_unstable(); + assert_eq!(got, want, "{sort_order:?} block-path record multiset differs from oracle"); + + // The two framings must also agree with EACH OTHER. Without this the + // unstable branch only length-checks `batch_out`, so a framing divergence + // between `SortMerge` and `SortMerge` + // that preserved record count would pass the queryname cases silently. + let mut batch_sorted = batch_out.clone(); + batch_sorted.sort_unstable(); + assert_eq!( + got, batch_sorted, + "{sort_order:?} block vs RecordBatch path record multiset differs" + ); + } +} + +/// `true` for sort orders whose sort is stable on equal keys (so byte-for-byte +/// output parity is deterministic). Queryname's comparator sort is unstable. +fn is_stable_for_equal_keys(order: SortOrder) -> bool { + matches!(order, SortOrder::Coordinate | SortOrder::TemplateCoordinate) +} + +/// L2.6: a `StepKind::Detached` SINK — the `WriteBgzfFile` analogue, i.e. the +/// detached-thread runtime driving a pure consumer — yields byte-identical +/// merged output to the legacy oracle, exactly like the pool-scheduled sink. +/// The chain's `SortMerge` is already `Detached`, so this drives the full chain +/// through `pipeline.run` with TWO off-pool threads (merge + sink) over a +/// multi-spill coordinate workload, pinning that `run_detached_driver` preserves +/// the output bytes for both the producing (merge) and consuming (sink) roles. +#[test] +fn detached_sink_chain_matches_legacy_coordinate() { + let memory_limit = 64 * 1024; // small → forces many real spill files + let threads = 4; + let (header, records) = synthesize_records(20_000, 0xD17A_C4ED); + let sorter = RawExternalSorter::new(SortOrder::Coordinate) + .memory_limit(memory_limit) + .threads(threads) + .output_compression(1) + .temp_compression(1); + let detached_out = drive_sort_buffer_pipeline( + sorter, + &header, + pack_batches(&records, 256), + 4 * 1024 * 1024, + threads, + StepKind::Detached, + SortDecompressTuning::default(), + SpillCodec::Zstd, + ) + .expect("detached-sink buffer pipeline drives to completion"); + + let legacy_out = + sort_via_legacy(SortOrder::Coordinate, &header, &records, memory_limit, threads) + .expect("legacy"); + + assert_eq!(detached_out.len(), legacy_out.len(), "record count mismatch"); + for (i, (got, want)) in detached_out.iter().zip(legacy_out.iter()).enumerate() { + assert_eq!(got, want, "detached-sink record {i} bytes differ from oracle"); + } +} + +/// The buffer chain must hold coordinate parity across BOTH decompress +/// granularities × block batches, with a multi-spill workload that forces real +/// spill files (so `CompressSpill`'s written chunks feed the block-parallel +/// reorder path). Guards against any spill-format / slot-ordering drift between +/// `CompressSpill` and the proven `SortSpillDecompress` reader. +#[rstest] +#[case::file_b1(true, 1)] +#[case::file_b4(true, 4)] +#[case::block_b1(false, 1)] +#[case::block_b4(false, 4)] +fn sort_buffer_chain_coordinate_matches_legacy_across_decompress_tunings( + #[case] file_granularity: bool, + #[case] block_batch: usize, +) { + let (header, records) = synthesize_records(20_000, 0xBADD_CAFE); + let memory_limit = 256 * 1024; + let threads = 2; + let sorter = RawExternalSorter::new(SortOrder::Coordinate) + .memory_limit(memory_limit) + .threads(threads) + .output_compression(1) + .temp_compression(1); + let new_out = drive_sort_buffer_pipeline( + sorter, + &header, + pack_batches(&records, 256), + 4 * 1024 * 1024, + // Exactly the case's thread count (see the sibling test): production + // runs the pipeline with `num_threads`, no `max(3)` floor. + threads, + StepKind::Exclusive, + SortDecompressTuning { file_granularity, block_batch }, + SpillCodec::Zstd, + ) + .expect("buffer pipeline drives to completion"); + + let legacy_out = + sort_via_legacy(SortOrder::Coordinate, &header, &records, memory_limit, threads) + .expect("legacy"); + + assert_eq!(new_out.len(), legacy_out.len(), "record count mismatch"); + for (i, (got, want)) in new_out.iter().zip(legacy_out.iter()).enumerate() { + assert_eq!(got, want, "record {i} bytes differ"); + } +} + +/// Equal-key stability across spill boundaries — the output-identity-critical +/// tie-break the single-residual design depends on. Every record shares one +/// coordinate (tid 0, pos 0), so a stable sort must emit them in input order; +/// the buffer chain must preserve that across multiple spill chunks (tie-broken +/// by `file_id == seq`) and the residual. Pinned both directly (output == input +/// order) and against the legacy oracle. +#[test] +fn sort_buffer_chain_preserves_equal_key_input_order_across_spills() { + let header = Header::default(); + // All at tid 0, pos 0 → identical coordinate key; names encode input order. + // 6_000 records against a tiny memory limit forces several spill chunks. + let records: Vec = (0..6_000u32) + .map(|i| { + let name = format!("r{i:06}"); + RawRecord::from(make_bam_bytes(0, 0, 0, name.as_bytes(), &[], 80, -1, -1, &[])) + }) + .collect(); + let input_bytes: RecordBytes = records.iter().map(|r| r.as_ref().to_vec()).collect(); + let memory_limit = 256 * 1024; + + let sorter = RawExternalSorter::new(SortOrder::Coordinate) + .memory_limit(memory_limit) + .threads(2) + .output_compression(1) + .temp_compression(1); + let out = drive_sort_buffer_pipeline( + sorter, + &header, + pack_batches(&records, 256), + 4 * 1024 * 1024, + 2, + StepKind::Exclusive, + SortDecompressTuning::default(), + SpillCodec::Zstd, + ) + .expect("buffer pipeline drives to completion"); + + assert_eq!(out.len(), input_bytes.len(), "record count mismatch"); + assert_eq!(out, input_bytes, "equal-key records must preserve input order across spills"); + + let legacy = + sort_via_legacy(SortOrder::Coordinate, &header, &records, memory_limit, 2).expect("legacy"); + assert_eq!(out, legacy, "equal-key order must also match the legacy oracle"); +} + +// ── Decompression-granularity parity (file-granularity × block-batch) ──────── + +/// The streaming sort must produce byte-identical output for BOTH decompression +/// granularities (`file_granularity ∈ {true, false}`) across `block_batch ∈ +/// {1, 4}`, validated against the legacy reference. The multi-spill workload +/// forces real spill files so the block-parallel reorder path is exercised (the +/// in-memory-only path never opens a slot reader). +#[rstest] +#[case::file_b1(true, 1)] +#[case::file_b4(true, 4)] +#[case::block_b1(false, 1)] +#[case::block_b4(false, 4)] +// block_batch == 0 is clamped to 1 in `SortSpillDecompress::new`. Without the +// clamp, the inline path declares a phantom EOF after reading zero blocks +// (silent record loss) and the block-parallel path livelocks (queue_eof never +// finalizes). These cases assert the clamp holds: identical to legacy, no hang. +#[case::file_b0(true, 0)] +#[case::block_b0(false, 0)] +fn three_step_chain_granularity_matrix_matches_legacy( + #[case] file_granularity: bool, + #[case] block_batch: usize, +) { + let sort_order = SortOrder::Coordinate; + let threads = 4; + let (header, records) = synthesize_sized_records(30_000, 0x5EED_1234, 120); + // Small per-thread memory ⇒ many spill files ⇒ many slot blocks. + let memory_limit = 256 * 1024; + + let sorter = RawExternalSorter::new(sort_order) + .memory_limit(memory_limit) + .threads(2) + .output_compression(1) + .temp_compression(1); + let new_out = drive_sort_pipeline_tuned( + sorter, + &header, + pack_batches(&records, 256), + 4 * 1024 * 1024, + threads, + StepKind::Exclusive, + SortDecompressTuning { file_granularity, block_batch }, + SpillCodec::Zstd, + ) + .expect("pipeline drives to completion"); + + let legacy_out = + sort_via_legacy(sort_order, &header, &records, memory_limit, 2).expect("legacy"); + + assert_eq!( + new_out.len(), + legacy_out.len(), + "record count mismatch (file_granularity={file_granularity}, block_batch={block_batch})" + ); + assert_eq!( + new_out, legacy_out, + "sorted bytes differ (file_granularity={file_granularity}, block_batch={block_batch})" + ); +} + +/// Block-parallel decompression over BGZF spill files (the non-default codec) +/// must also match the legacy sorter — the matrix/soak/proptest exercise the +/// default zstd spills, so this confirms the block-parallel reorder path is +/// codec-agnostic. (Output records are codec-independent: the spill codec only +/// affects temp files, not the sorted output.) +#[test] +fn block_parallel_bgzf_spill_matches_legacy() { + let sort_order = SortOrder::Coordinate; + let (header, records) = synthesize_sized_records(30_000, 0x5EED_BEEF, 120); + let memory_limit = 256 * 1024; + let sorter = RawExternalSorter::new(sort_order) + .memory_limit(memory_limit) + .threads(2) + .output_compression(1) + .temp_compression(1) + .spill_codec(fgumi_sort::SpillCodec::Bgzf); + let new_out = drive_sort_pipeline_tuned( + sorter, + &header, + pack_batches(&records, 256), + 4 * 1024 * 1024, + 8, + StepKind::Exclusive, + SortDecompressTuning { file_granularity: false, block_batch: 2 }, + SpillCodec::Bgzf, + ) + .expect("pipeline drives to completion"); + let legacy_out = + sort_via_legacy(sort_order, &header, &records, memory_limit, 2).expect("legacy"); + assert_eq!(new_out.len(), legacy_out.len(), "record count mismatch (bgzf block-parallel)"); + assert_eq!(new_out, legacy_out, "sorted bytes differ (bgzf block-parallel)"); +} + +/// Block-parallel decompression completes out of order (workers decompress one +/// file's blocks concurrently), yet the reassembled output must be byte- +/// identical to the in-order (file-granularity) result. Property test over a +/// range of record counts and `block_batch` sizes and a high pipeline-thread +/// count (more concurrent decompressors ⇒ more out-of-order completion). A +/// straggler worker hitting reader-EOF while another holds an in-flight block +/// must not truncate the output (record count is asserted equal). +#[cfg(test)] +// Soak/matrix/proptest suites: multi-minute, so gated off the default test +// target and run on the nightly `cargo ci-test-stress` job instead. +#[cfg(feature = "stress-tests")] +mod proptests { + use super::*; + use proptest::prelude::*; + + proptest! { + #![proptest_config(ProptestConfig { cases: 24, ..ProptestConfig::default() })] + + #[test] + fn block_parallel_matches_file_granularity( + n_records in 2_000usize..18_000, + block_batch in 1usize..=6, + seed in any::(), + ) { + let sort_order = SortOrder::Coordinate; + let (header, records) = synthesize_sized_records(n_records, seed, 100); + // Force spilling so slots (and the reorder path) are exercised. + let memory_limit = 256 * 1024; + let pipeline_threads = 6; + + let make_sorter = || RawExternalSorter::new(sort_order) + .memory_limit(memory_limit) + .threads(2) + .output_compression(1) + .temp_compression(1); + + let in_order = drive_sort_pipeline_tuned( + make_sorter(), + &header, + pack_batches(&records, 256), + 4 * 1024 * 1024, + pipeline_threads, + StepKind::Exclusive, + SortDecompressTuning { file_granularity: true, block_batch }, + SpillCodec::Zstd, + ).expect("file-granularity pipeline"); + + let out_of_order = drive_sort_pipeline_tuned( + make_sorter(), + &header, + pack_batches(&records, 256), + 4 * 1024 * 1024, + pipeline_threads, + StepKind::Exclusive, + SortDecompressTuning { file_granularity: false, block_batch }, + SpillCodec::Zstd, + ).expect("block-parallel pipeline"); + + prop_assert_eq!(out_of_order.len(), records.len(), "no truncation"); + prop_assert_eq!(out_of_order, in_order, "block-parallel diverges from in-order"); + } + } +} + +#[cfg(feature = "stress-tests")] +/// Maximum-contention soak for the block-parallel decompress path +/// (`file_granularity == false`). Drives the path repeatedly under the most +/// adversarial settings the knobs allow — many spill files, a tiny reorder +/// window (so stragglers continuously hit `bp_reorder_admits` backpressure and +/// the Phase-B drain-only path), `block_batch == 1` (maximum per-block churn and +/// the most frequent `reader_eof`/`in_flight` transitions), and far more +/// pipeline worker threads (12) than sorter threads (so many workers race to +/// decompress one file's blocks concurrently and finalize out of order). +/// +/// Each iteration uses a fresh random seed and is checked for *byte-identity* +/// against the legacy `RawExternalSorter::sort` oracle — so a lost, duplicated, +/// or reordered block (the truncation class the `reader_eof`/`in_flight` +/// protocol guards against) fails the assertion. Each iteration runs under a +/// per-iteration wall-clock watchdog (via `run_watchdogged_parity`): a livelock +/// (e.g. `queue_eof` never finalizing) trips the timeout and fails the test +/// instead of hanging CI. +/// +/// This complements the loom model (exhaustive but tiny) and the proptest +/// (random sizes, moderate threads) by hammering the REAL pipeline under +/// sustained high contention for many iterations. +#[test] +fn block_parallel_high_contention_soak_matches_legacy() { + use std::time::Duration; + + const ITERATIONS: usize = 40; + const RECORDS_PER_ITER: usize = 15_000; + const PIPELINE_THREADS: usize = 12; + const SORTER_THREADS: usize = 2; + // Tiny per-thread sort memory ⇒ many spill files ⇒ many slot readers. + const MEMORY_LIMIT: usize = 128 * 1024; + // Tiny output/reorder-window budget ⇒ the block-parallel reorder window is + // ~1 block, so `bp_reorder_admits` backpressures aggressively and workers + // are repeatedly forced through the Phase-B drain-only path. + const OUTPUT_BYTE_LIMIT: u64 = 64 * 1024; + const BLOCK_BATCH: usize = 1; + // Per-iteration watchdog: a livelock in any single iteration fails fast. + const WATCHDOG: Duration = Duration::from_secs(60); + + let sort_order = SortOrder::Coordinate; + for iter in 0..ITERATIONS { + let seed = 0xA5A5_0000_u64 ^ (iter as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15); + let (header, records) = synthesize_sized_records(RECORDS_PER_ITER, seed, 110); + run_watchdogged_parity( + &format!("hc-soak-i{iter}"), + sort_order, + header, + records, + MEMORY_LIMIT, + SORTER_THREADS, + PIPELINE_THREADS, + OUTPUT_BYTE_LIMIT, + SortDecompressTuning { file_granularity: false, block_batch: BLOCK_BATCH }, + WATCHDOG, + ); + } +} + +#[cfg(feature = "stress-tests")] +/// Maximum-churn soak over the `SortBuffer` chain: a *tiny* memory limit seals a +/// run after only a handful of records, so the chain cycles +/// seal → compress → spill → decompress → merge as fast as it can while eight +/// pipeline threads contend for the `Serial` steps. A wedge (a step parking +/// while a downstream holds the chunk it is waiting on) trips the per-iteration +/// watchdog and fails fast instead of hanging, and byte-parity against the +/// legacy oracle proves nothing is lost, duplicated, or reordered under churn. +/// +/// NOTE: this drives `drive_sort_pipeline_tuned` (the record-input `SortBuffer` +/// chain), NOT the block-input arena front. The front's capacity-1 arena cycle — +/// `ReadBlocks` acquire+admit and `FindBoundariesAndSort` seal+free on the +/// coordination driver with `InflateToArena` on the pool in between — is covered +/// by `arena_front_chain_seals_multiple_runs_without_losing_records`, which seals +/// several runs through the real runtime and so exercises acquire/seal/free +/// across runs. That test is not a *soak*: there is no watchdogged high-churn +/// coverage of the arena front yet. +#[test] +fn sort_buffer_chain_tight_memory_soak_no_deadlock() { + use std::time::Duration; + + const ITERATIONS: usize = 24; + const RECORDS_PER_ITER: usize = 12_000; + const PIPELINE_THREADS: usize = 8; + const SORTER_THREADS: usize = 4; + // Very tight sort memory ⇒ `SortBuffer` seals after only a handful of + // records ⇒ maximal seal/spill/merge churn across the pipeline threads. + const MEMORY_LIMIT: usize = 24 * 1024; + const OUTPUT_BYTE_LIMIT: u64 = 48 * 1024; + const WATCHDOG: Duration = Duration::from_secs(60); + + for iter in 0..ITERATIONS { + let seed = 0xC0FF_EE00_u64 ^ (iter as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15); + let (header, records) = synthesize_sized_records(RECORDS_PER_ITER, seed, 130); + run_watchdogged_parity( + &format!("tight-memory-soak-i{iter}"), + SortOrder::Coordinate, + header, + records, + MEMORY_LIMIT, + SORTER_THREADS, + PIPELINE_THREADS, + OUTPUT_BYTE_LIMIT, + SortDecompressTuning { file_granularity: false, block_batch: 1 }, + WATCHDOG, + ); + } +} + +/// Assert two record-byte streams are identical, reporting only the FIRST +/// mismatch (index + lengths + 16-byte prefixes). A blanket `assert_eq!` on the +/// two `Vec>` would dump tens of thousands of binary records into CI +/// logs on divergence; this keeps a failure readable while still catching any +/// lost / duplicated / reordered / corrupted record. +fn assert_record_parity(label: &str, actual: &[Vec], expected: &[Vec]) { + if let Some((idx, (a, b))) = actual.iter().zip(expected).enumerate().find(|(_, (a, b))| a != b) + { + panic!( + "{label}: output diverges from legacy at record {idx}: \ + actual_len={}, expected_len={}, actual_prefix={:?}, expected_prefix={:?}", + a.len(), + b.len(), + &a[..a.len().min(16)], + &b[..b.len().min(16)], + ); + } + // No record differs within the common prefix; a length delta is the only + // remaining divergence (truncation or duplication). + assert_eq!( + actual.len(), + expected.len(), + "{label}: record count mismatch (truncation/duplication)" + ); +} + +/// Run one watchdog-guarded streaming sort and assert byte-identity against the +/// legacy oracle. Owns `header`/`records` so the worker thread can take them. +/// +/// Both the legacy oracle AND the streaming pipeline run *inside* the worker, so +/// the `recv_timeout` watchdog covers both — a stall in either (the path under +/// test or, defensively, the reference sorter) fails the test fast instead of +/// hanging the test process. The merge sink is `Serial` (every watchdog'd parity +/// caller drives the streaming three-step chain). Panics (failing the test) on +/// divergence, pipeline/oracle error, watchdog timeout (livelock), or a worker +/// panic — never hangs. +#[allow(clippy::too_many_arguments)] +fn run_watchdogged_parity( + label: &str, + sort_order: SortOrder, + header: Header, + records: Vec, + memory_limit: usize, + sorter_threads: usize, + pipeline_threads: usize, + output_byte_limit: u64, + tuning: SortDecompressTuning, + watchdog: std::time::Duration, +) { + use std::sync::mpsc; + + let (tx, rx) = mpsc::channel(); + let worker = std::thread::Builder::new() + .name(label.to_string()) + .spawn(move || { + // (pipeline_out, legacy_out) — both computed under the watchdog. + let result = (|| -> Result<(RecordBytes, RecordBytes)> { + let legacy_out = + sort_via_legacy(sort_order, &header, &records, memory_limit, sorter_threads)?; + let batches = pack_batches(&records, 256); + let sorter = RawExternalSorter::new(sort_order) + .memory_limit(memory_limit) + .threads(sorter_threads) + .output_compression(1) + .temp_compression(1); + let out = drive_sort_pipeline_tuned( + sorter, + &header, + batches, + output_byte_limit, + pipeline_threads, + StepKind::Serial, + tuning, + SpillCodec::Zstd, + )?; + Ok((out, legacy_out)) + })(); + let _ = tx.send(result); + }) + .expect("spawn soak worker"); + + match rx.recv_timeout(watchdog) { + Ok(Ok((out, legacy_out))) => { + assert_record_parity(label, &out, &legacy_out); + worker.join().expect("soak worker panicked"); + } + Ok(Err(e)) => panic!("{label}: pipeline/oracle errored: {e:#}"), + Err(mpsc::RecvTimeoutError::Timeout) => { + panic!("{label}: DEADLOCK/LIVELOCK — sort did not complete within {watchdog:?}") + } + Err(mpsc::RecvTimeoutError::Disconnected) => { + panic!("{label}: soak worker dropped its sender (panicked?)") + } + } +} + +#[cfg(feature = "stress-tests")] +/// Spill-pressure regimes for the block-parallel soak matrix. Both force real +/// spill files (so slot readers and the Phase-2 reorder path run); they differ +/// in how the spilled data is shaped across files. +#[derive(Clone, Copy, Debug)] +enum SoakRegime { + /// Larger-than-budget workload: total spilled bytes vastly exceed the + /// in-memory sort budget, producing MANY small spill files (tens). This is + /// the mandatory larger-than-RAM run — the external merge over many slot + /// readers, with frequent cross-file `reader_eof`/`in_flight` transitions, + /// is the case the truncation protocol must survive. + ManySmallFiles, + /// A handful of LARGE spill files: the budget admits a big batch, so only a + /// few (but > 1) files spill, each with many blocks. Stresses long per-file + /// block runs and the per-slot reorder window rather than cross-file churn. + FewLargeFiles, +} + +#[cfg(feature = "stress-tests")] +struct SoakParams { + records: usize, + seq_len: usize, + memory_limit: usize, + output_byte_limit: u64, + block_batch: usize, +} + +#[cfg(feature = "stress-tests")] +impl SoakRegime { + /// Discriminant folded into the per-case seed so each regime sorts a + /// distinct record set. + fn seed_salt(self) -> u64 { + match self { + SoakRegime::ManySmallFiles => 0x1111_1111_1111_1111, + SoakRegime::FewLargeFiles => 0x2222_2222_2222_2222, + } + } + + fn params(self) -> SoakParams { + match self { + // ~40k records ≈ 10 MB spilled into many (~100) small files at a + // 96 KiB budget, with a tiny reorder window and block_batch == 1 + // (max per-block churn and the most `reader_eof`/`in_flight` events). + SoakRegime::ManySmallFiles => SoakParams { + records: 40_000, + seq_len: 150, + memory_limit: 96 * 1024, + output_byte_limit: 128 * 1024, + block_batch: 1, + }, + // ~12k × 150B ≈ 1.8 MB spilled into ~4 large files at a 512 KiB + // budget, with a roomy window and block_batch == 4. + SoakRegime::FewLargeFiles => SoakParams { + records: 12_000, + seq_len: 150, + memory_limit: 512 * 1024, + output_byte_limit: 4 * 1024 * 1024, + block_batch: 4, + }, + } + } +} + +#[cfg(feature = "stress-tests")] +/// External-watchdog soak MATRIX for the Phase-2 decompress path. Crosses +/// pipeline-thread count × decompress granularity × spill regime, so both the +/// block-parallel reorder/in-flight/EOF protocol and the file-granularity FIFO +/// are hammered across {1, 2, 8} workers, {many small, few large} spill-file +/// shapes, and both code paths. Each (case × iteration) runs under a wall-clock +/// watchdog and is checked byte-for-byte against the legacy oracle, so a +/// livelock fails fast and any lost / duplicated / reordered block is caught. +/// +/// `rstest` generates the full cross product (3 × 2 × 2 = 12 cases); nextest runs +/// them as independent parallel tests. The mandatory larger-than-budget run is +/// `ManySmallFiles` (~100 spill files at a 96 KiB budget). This is the real P5 +/// hardening gate the OFF-default flip rests on; it complements the loom model +/// (exhaustive but tiny), the proptest (random sizes), and the single-corner +/// high-contention soak (12 threads, 1-block window). +#[rstest] +fn block_parallel_soak_matrix_matches_legacy( + #[values(1, 2, 8)] pipeline_threads: usize, + #[values(true, false)] file_granularity: bool, + #[values(SoakRegime::ManySmallFiles, SoakRegime::FewLargeFiles)] regime: SoakRegime, +) { + use std::time::Duration; + + const ITERATIONS: usize = 4; + const SORTER_THREADS: usize = 2; + const WATCHDOG: Duration = Duration::from_secs(120); + + let sort_order = SortOrder::Coordinate; + let SoakParams { records, seq_len, memory_limit, output_byte_limit, block_batch } = + regime.params(); + + for iter in 0..ITERATIONS { + // Distinct seed per (regime, threads, granularity, iter). + let seed = 0x50A4_0000_u64 + .wrapping_add((iter as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)) + .wrapping_add((pipeline_threads as u64) << 40) + .wrapping_add(u64::from(file_granularity) << 32) + ^ regime.seed_salt(); + let (header, recs) = synthesize_sized_records(records, seed, seq_len); + let label = format!( + "soak-{regime:?}-t{pipeline_threads}-fg{file_granularity}-bb{block_batch}-i{iter}" + ); + run_watchdogged_parity( + &label, + sort_order, + header, + recs, + memory_limit, + SORTER_THREADS, + pipeline_threads, + output_byte_limit, + SortDecompressTuning { file_granularity, block_batch }, + WATCHDOG, + ); + } +} + +/// Number of reference sequences the synthetic header declares. +const N_TEST_REFS: usize = 4; + +fn synthesize_sized_records(n: usize, seed: u64, seq_len: usize) -> (Header, Vec) { + let mut state = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15).wrapping_add(1); + let mut next_u32 = || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + #[allow(clippy::cast_possible_truncation)] + let v = state as u32; + v + }; + + // Reference sequences must exist for the mapped records below to carry a + // meaningful coordinate key. + let header = { + use noodles::sam::header::record::value::Map; + use noodles::sam::header::record::value::map::ReferenceSequence; + use std::num::NonZeroUsize; + let len = NonZeroUsize::new(1_000_000).expect("nonzero"); + let mut builder = Header::builder(); + for i in 0..N_TEST_REFS { + builder = builder + .add_reference_sequence(format!("chr{i}"), Map::::new(len)); + } + builder.build() + }; + let mut records = Vec::with_capacity(n); + for i in 0..n { + let name = format!("rd_{}", next_u32() % 100_000); + let pos: i32 = (next_u32() % 1_000_000).cast_signed(); + let is_paired = i % 2 == 0; + // Most records are MAPPED across a handful of references. With + // `tid = -1` everywhere, `extract_coordinate_key_inline` returns + // `RawCoordinateKey::unmapped()` (`u64::MAX`) for every record, `pos` is + // never read, and the coordinate cases degenerate into one equal-key + // bucket — they would pass even if the key comparison were broken. Every + // eighth record stays unmapped so the equal-key/tie path is still covered. + let unmapped = i % 8 == 0; + #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] + let tid: i32 = if unmapped { -1 } else { (i % N_TEST_REFS) as i32 }; + let flags: u16 = if unmapped { 0x4 } else { 0 } | if is_paired { 0x1 | 0x8 } else { 0 }; + let bytes = make_bam_bytes(tid, pos, flags, name.as_bytes(), &[], seq_len, -1, -1, &[]); + records.push(RawRecord::from(bytes)); + } + (header, records) +} + +#[rstest] +#[case::t1(1)] +#[case::t2(2)] +#[case::t4(4)] +#[case::t8(8)] +fn three_step_chain_large_spill_completes(#[case] pipeline_threads: usize) { + use std::time::Duration; + + let (header, records) = synthesize_sized_records(60_000, 0xBADD_CAFE, 200); + // Default decompress tuning (block-parallel, block_batch 4) under a + // per-case watchdog: the regression this pins is a deadlock at high + // pipeline-thread counts, so the watchdog converts a hang into a failure. + run_watchdogged_parity( + &format!("large-spill-t{pipeline_threads}"), + SortOrder::Coordinate, + header, + records, + 1024 * 1024, // sort memory_limit + 2, // sorter_threads + pipeline_threads, + 256 * 1024 * 1024, // output queue limit + SortDecompressTuning::default(), + Duration::from_secs(90), + ); +} + +#[test] +fn three_step_chain_empty_input_drains_cleanly() { + let header = Header::default(); + let sorter = RawExternalSorter::new(SortOrder::Coordinate).memory_limit(1024 * 1024); + let out = drive_sort_pipeline(sorter, &header, Vec::new(), 64 * 1024, 3, StepKind::Exclusive) + .expect("empty pipeline"); + assert!(out.is_empty()); +} + +// ── SortMerge fail-closed regression tests ────────────────────────────────── + +/// `Exclusive` source that drains a `Vec` one event per +/// `try_run`, feeding `SortMerge` directly. Used to drive the merge into a +/// drained-but-incomplete-setup state without standing up the spill machinery. +struct Phase2EventSource { + events: Vec, + held: HeldSlot>, + output_byte_limit: u64, +} + +impl Phase2EventSource { + fn new( + mut events: Vec, + output_byte_limit: u64, + ) -> Self { + events.reverse(); + Self { events, held: HeldSlot::new(), output_byte_limit } + } +} + +impl Step for Phase2EventSource { + type Input = (); + type Outputs = fgumi_pipeline_core::outputs::Single; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "Phase2EventSource", + kind: StepKind::Exclusive, + sticky: true, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], + branch_ordering: vec![BranchOrdering::None], + } + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if let Some(unpushed) = self.held.take() { + match ctx.outputs.retry(unpushed) { + Ok(()) => {} + Err(again) => { + self.held.put(again); + return Ok(StepOutcome::Progress); + } + } + } + let Some(event) = self.events.pop() else { + return Ok(StepOutcome::Finished); + }; + match ctx.outputs.push(event) { + Ok(()) => Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held.put(unpushed); + Ok(StepOutcome::Progress) + } + } + } +} + +fn run_merge_over_events( + events: Vec, +) -> Result>> { + // Delegates to `collect_merge_batches` so the merge chain is wired in exactly + // one place; this helper only flattens the batches into per-record bytes. + let batches = collect_merge_batches(events, 1 << 20, 256)?; + let mut out = Vec::new(); + for batch in batches { + for bytes in batch.iter_record_bytes() { + out.push(bytes.to_vec()); + } + } + Ok(out) +} + +/// A wholly empty event stream (no payload, no `AllAnnounced`) is the one +/// legitimate drained-but-not-ready case: it merges to an empty output. +#[test] +fn test_sort_merge_empty_input_merges_to_empty() { + let out = run_merge_over_events(Vec::new()).expect("empty merge should succeed"); + assert!(out.is_empty(), "expected no records, got {}", out.len()); +} + +/// An `AllAnnounced` that promises a slot which never arrives leaves the setup +/// incomplete when the input drains; `SortMerge` must fail closed rather than +/// silently merge a partial result. +#[test] +fn test_sort_merge_fails_closed_on_incomplete_setup() { + let events = vec![crate::sort::protocol::SortPhase2Event::AllAnnounced { + slot_count: 1, + memory_chunk_count: 0, + total_records: 0, + }]; + let err = run_merge_over_events(events).expect_err("incomplete setup must error"); + let msg = err.to_string(); + assert!(msg.contains("setup incomplete at input drain"), "unexpected error message: {msg}"); +} + +// ── SortMerge output-buffer sizing + duplicate-AllAnnounced regression ─────── + +/// Drive `SortMerge` over `events` and return the *batches* it emits (not the +/// flattened records), so tests can inspect per-batch buffer capacity. +fn collect_merge_batches( + events: Vec, + output_byte_limit: u64, + target_batch_count: usize, +) -> Result> { + let received: Arc>> = Arc::new(Mutex::new(Vec::new())); + let source = Phase2EventSource::new(events, output_byte_limit); + let merge = SortMerge::::with_target_batch_count( + SortOrder::Coordinate, + output_byte_limit, + target_batch_count, + ); + let sink = VecSink { received: Arc::clone(&received), kind: StepKind::Serial }; + + let builder = Pipeline::builder(); + builder.chain(source).chain(merge).chain(sink).into_sink_marker(); + let pipeline = builder.build()?; + pipeline.run(PipelineConfig { threads: 1, ..Default::default() })?; + + Ok(std::mem::take(&mut *received.lock())) +} + +/// Wrap `records` as a single coordinate-sorted in-memory chunk event. All keys +/// are `default()` (equal) — order does not matter for the buffer-sizing and +/// duplicate-announcement assertions, only that the chunk merges cleanly. +fn coordinate_memory_chunk_event( + records: Vec, +) -> crate::sort::protocol::SortPhase2Event { + use crate::sort::protocol::MemoryChunkErased; + let total = records.len() as u64; + let chunk = fgumi_sort::InMemoryChunk::from_owned_records( + records + .into_iter() + .map(|r| (fgumi_sort::RawCoordinateKey::default(), r.into_inner())) + .collect(), + ); + crate::sort::protocol::SortPhase2Event::MemoryChunk { + chunk: Arc::new(MemoryChunkErased::Coordinate(chunk)), + records_ingested_so_far: total, + } +} + +/// Count-bound output batches must not each reserve the full output-queue byte +/// budget. Drive a many-small-record merge that emits several count-capped +/// batches and assert their total resident capacity stays well under one byte +/// budget — it would be ~`num_batches * byte_limit` if every buffer reserved +/// the full budget (the pre-fix behavior). +#[test] +fn test_sort_merge_does_not_over_reserve_output_buffers() { + use fgumi_pipeline_core::item::HeapSize; + + let (_header, records) = synthesize_records(600, 7); + let byte_limit: u64 = 1 << 20; // 1 MiB + let target = 256; + let events = vec![ + coordinate_memory_chunk_event(records), + crate::sort::protocol::SortPhase2Event::AllAnnounced { + slot_count: 0, + memory_chunk_count: 1, + total_records: 600, + }, + ]; + let batches = collect_merge_batches(events, byte_limit, target).expect("merge should succeed"); + + let emitted: usize = batches.iter().map(|b| b.iter_record_bytes().count()).sum(); + assert_eq!(emitted, 600, "all records must be emitted"); + assert!( + batches.len() >= 2, + "workload must span multiple count-bound batches, got {}", + batches.len() + ); + let total_heap: usize = batches.iter().map(HeapSize::heap_size).sum(); + assert!( + (total_heap as u64) < byte_limit, + "output buffers over-reserved: {total_heap} bytes across {} batches \ + (would exceed one byte budget if each reserved the full {byte_limit})", + batches.len(), + ); +} + +/// The Phase-2 protocol emits exactly one `AllAnnounced`. A second one is a +/// protocol violation; `SortMerge` must fail closed rather than overwrite its +/// completion expectations. The first announcement over-promises (2 chunks) so +/// setup never completes and the duplicate is still absorbed in setup. +#[test] +fn test_sort_merge_fails_closed_on_duplicate_all_announced() { + let (_header, records) = synthesize_records(1, 1); + let byte_limit: u64 = 1 << 20; + let events = vec![ + coordinate_memory_chunk_event(records), + crate::sort::protocol::SortPhase2Event::AllAnnounced { + slot_count: 0, + memory_chunk_count: 2, + total_records: 1, + }, + crate::sort::protocol::SortPhase2Event::AllAnnounced { + slot_count: 0, + memory_chunk_count: 2, + total_records: 1, + }, + ]; + let err = collect_merge_batches(events, byte_limit, 256) + .expect_err("duplicate AllAnnounced must error"); + let msg = err.to_string(); + assert!(msg.contains("duplicate AllAnnounced"), "unexpected error: {msg}"); +} + +/// The `Arc::try_unwrap` guard in `absorb_phase2_event` has no coverage from the +/// other tests: `coordinate_memory_chunk_event` always mints a fresh `Arc`, so +/// only the success path runs. The protocol moves memory chunks rather than +/// cloning them, so a shared `Arc` at the merge consumer means a producer kept a +/// handle — deep-cloning the record vector instead would silently double memory. +#[test] +fn test_sort_merge_fails_closed_on_shared_memory_chunk_arc() { + use crate::sort::protocol::{MemoryChunkErased, SortPhase2Event}; + + let (_header, records) = synthesize_records(8, 3); + let chunk = + Arc::new(MemoryChunkErased::Coordinate(fgumi_sort::InMemoryChunk::from_owned_records( + records + .into_iter() + .map(|r| (fgumi_sort::RawCoordinateKey::default(), r.into_inner())) + .collect(), + ))); + + // Two events sharing ONE Arc — the violation. `Arc::strong_count` is 2 when + // the merge tries to take ownership of the first. + let events = vec![ + SortPhase2Event::MemoryChunk { chunk: Arc::clone(&chunk), records_ingested_so_far: 8 }, + SortPhase2Event::MemoryChunk { chunk, records_ingested_so_far: 8 }, + SortPhase2Event::AllAnnounced { slot_count: 0, memory_chunk_count: 2, total_records: 8 }, + ]; + + let err = run_merge_over_events(events).expect_err("a shared chunk Arc must fail closed"); + let msg = err.to_string(); + assert!( + msg.contains("Arc unexpectedly shared"), + "expected the shared-Arc guard message, got: {msg}" + ); +} + +/// The buffer-sizing logic is duplicated in `next_batch` (the k-way `Merging` +/// path) and `next_fast_batch` (the single-chunk fast path). +/// `test_sort_merge_does_not_over_reserve_output_buffers` has zero slots and one +/// memory chunk, so it only ever exercises the fast path — a regression that +/// reintroduced full-budget reservation in `next_batch` would pass it. Two +/// memory chunks force `build_driver` and the real k-way merge. +#[test] +fn test_sort_merge_does_not_over_reserve_on_the_kway_path() { + use fgumi_pipeline_core::item::HeapSize; + + let (_header, first) = synthesize_records(300, 7); + let (_header2, second) = synthesize_records(300, 11); + let byte_limit: u64 = 1 << 20; // 1 MiB + let target = 256; + + let events = vec![ + coordinate_memory_chunk_event(first), + coordinate_memory_chunk_event(second), + crate::sort::protocol::SortPhase2Event::AllAnnounced { + slot_count: 0, + memory_chunk_count: 2, + total_records: 600, + }, + ]; + let batches = collect_merge_batches(events, byte_limit, target).expect("merge should succeed"); + + let emitted: usize = batches.iter().map(|b| b.iter_record_bytes().count()).sum(); + assert_eq!(emitted, 600, "all records from both chunks must be emitted"); + assert!( + batches.len() >= 2, + "workload must span multiple count-bound batches, got {}", + batches.len() + ); + let total_heap: usize = batches.iter().map(HeapSize::heap_size).sum(); + assert!( + (total_heap as u64) < byte_limit, + "k-way output buffers over-reserved: {total_heap} bytes across {} batches", + batches.len(), + ); +} + +// ── Arena block-input front (ReadBlocks → InflateToArena → FindBoundariesAndSort) ── + +/// `Exclusive` source draining a `Vec` one block per `try_run`. +struct BgzfBlockSource { + blocks: Vec, + held: HeldSlot>, + output_byte_limit: u64, +} + +impl BgzfBlockSource { + fn new(mut blocks: Vec, output_byte_limit: u64) -> Self { + blocks.reverse(); + Self { blocks, held: HeldSlot::new(), output_byte_limit } + } +} + +impl Step for BgzfBlockSource { + type Input = (); + type Outputs = OrderedBytesSingle; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "BgzfBlockSource", + kind: StepKind::Exclusive, + sticky: true, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], + branch_ordering: vec![BranchOrdering::ByItemOrdinal], + } + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if let Some(unpushed) = self.held.take() { + match ctx.outputs.retry(unpushed) { + Ok(()) => {} + Err(again) => { + self.held.put(again); + return Ok(StepOutcome::Progress); + } + } + } + let Some(block) = self.blocks.pop() else { + return Ok(StepOutcome::Finished); + }; + match ctx.outputs.push(block) { + Ok(()) => Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held.put(unpushed); + Ok(StepOutcome::Progress) + } + } + } +} + +/// What the arena front produced: one entry per emitted chunk (in emit order) +/// plus the terminal `AllAnnounced` counts. +#[derive(Default)] +struct ArenaFrontOutput { + chunks: Vec, + slot_count: u32, + total_records: u64, +} + +/// Sink that copies each chunk's record bodies out and DROPS the chunk in the +/// same `try_run`. Retaining the chunks instead would pin their arena `Arc`: +/// `ReadBlocks` owns a capacity-1 arena pool, so run *k+1* cannot start until +/// run *k*'s chunk is released, and a hoarding sink wedges the pipeline. +struct ChunkEventSink { + received: Arc>, +} + +impl Step for ChunkEventSink { + type Input = SortChunkEvent; + type Outputs = (); + + fn profile(&self) -> StepProfile { + StepProfile { + name: "ChunkEventSink", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + match ctx.input.pop() { + Some(event) => { + let mut out = self.received.lock(); + match event { + SortChunkEvent::Spill { chunk, .. } + | SortChunkEvent::Residual { chunk, .. } => { + out.chunks.push( + (0..chunk.len()).map(|i| chunk.record_bytes(i).to_vec()).collect(), + ); + } + SortChunkEvent::AllAnnounced { slot_count, total_records, .. } => { + out.slot_count = slot_count; + out.total_records = total_records; + } + } + Ok(StepOutcome::Progress) + } + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } +} + +/// Minimal BAM binary header (`magic + l_text=0 + n_ref` + one entry per ref). +/// Only `n_ref` matters to the front's `bam_header_len` scan and to the +/// coordinate key, so the reference names/lengths are placeholders. +fn minimal_binary_bam_header(n_ref: u32) -> Vec { + let mut header = Vec::new(); + header.extend_from_slice(b"BAM\x01"); + header.extend_from_slice(&0u32.to_le_bytes()); // l_text = 0 + header.extend_from_slice(&n_ref.to_le_bytes()); + for _ in 0..n_ref { + header.extend_from_slice(&2u32.to_le_bytes()); // l_name = 2 + header.extend_from_slice(b"r\0"); + header.extend_from_slice(&1_000_000u32.to_le_bytes()); // l_ref + } + header +} + +/// Serialize `[binary header][[block_size][body]...]` and cut it into BGZF +/// blocks of at most `payload_bytes` uncompressed each. Records straddle block +/// boundaries by construction, which is what exercises the front's carry path. +fn bgzf_blocks_for( + records: &[RawRecord], + n_ref: u32, + payload_bytes: usize, +) -> Vec { + let mut stream = minimal_binary_bam_header(n_ref); + for record in records { + let bytes = record.as_ref(); + let block_size = u32::try_from(bytes.len()).expect("record fits u32"); + stream.extend_from_slice(&block_size.to_le_bytes()); + stream.extend_from_slice(bytes); + } + stream + .chunks(payload_bytes) + .enumerate() + .map(|(i, payload)| { + let mut compressor = fgumi_bgzf::writer::InlineBgzfCompressor::new(1); + compressor.write_all(payload).expect("compress payload"); + compressor.flush().expect("flush compressor"); + let mut blocks = compressor.take_blocks(); + assert_eq!(blocks.len(), 1, "payload must fit one BGZF block"); + crate::types::BgzfBlock { + batch_serial: i as u64, + bytes: blocks.remove(0).data, + uncompressed_size: u32::try_from(payload.len()).expect("payload fits u32"), + } + }) + .collect() +} + +/// Drive `BgzfBlockSource → ReadBlocks → InflateToArena → FindBoundariesAndSort` +/// and return every emitted chunk's record bodies, chunk by chunk, plus the +/// terminal `AllAnnounced` counts. +fn drive_arena_front( + blocks: Vec, + n_ref: u32, + memory_limit: usize, + output_byte_limit: u64, + threads: usize, +) -> Result { + let received: Arc> = Arc::new(Mutex::new(ArenaFrontOutput::default())); + + let builder = Pipeline::builder(); + builder + .chain(BgzfBlockSource::new(blocks, output_byte_limit)) + .chain(ReadBlocks::new(memory_limit, output_byte_limit)) + .chain(InflateToArena::new(output_byte_limit)) + .chain(FindBoundariesAndSort::new(CoordinateStrategy::new(n_ref), 1, output_byte_limit)) + .chain(ChunkEventSink { received: Arc::clone(&received) }) + .into_sink_marker(); + let pipeline = builder.build()?; + pipeline.run(PipelineConfig { threads, ..Default::default() })?; + + Ok(std::mem::take(&mut *received.lock())) +} + +/// End-to-end parity for the block-input arena front, driven through the real +/// runtime rather than by poking `ingest_block`/`finalize` directly. A +/// `memory_limit` above the whole input yields exactly ONE run, so the single +/// residual chunk must be byte-identical to the legacy oracle's coordinate sort. +/// +/// `payload_bytes` varies where records fall relative to block boundaries: the +/// small case guarantees records straddle blocks (the front's carry path), the +/// large case puts the whole stream in one block. +#[rstest] +#[case::straddling_blocks(4096, 20)] +#[case::one_block_per_run(60_000, 2)] +fn arena_front_chain_single_run_matches_legacy_oracle( + #[case] payload_bytes: usize, + #[case] min_blocks: usize, +) { + const N_RECORDS: usize = 2_000; + let (header, records) = synthesize_records(N_RECORDS, 0xA2E4_A100); + let n_ref = u32::try_from(header.reference_sequences().len()).expect("n_ref fits u32"); + + let blocks = bgzf_blocks_for(&records, n_ref, payload_bytes); + assert!( + blocks.len() >= min_blocks, + "expected at least {min_blocks} blocks at {payload_bytes} B/block, got {}", + blocks.len() + ); + + let out = drive_arena_front(blocks, n_ref, 256 * 1024 * 1024, 4 * 1024 * 1024, 4) + .expect("arena front drives to completion"); + + assert_eq!(out.slot_count, 0, "a budget-sized run never spills"); + assert_eq!(out.total_records, N_RECORDS as u64); + assert_eq!(out.chunks.len(), 1, "one run ⇒ one residual chunk"); + + let legacy_out = + sort_via_legacy(SortOrder::Coordinate, &header, &records, 256 * 1024 * 1024, 1) + .expect("legacy oracle"); + assert_eq!(out.chunks[0].len(), legacy_out.len(), "record count mismatch"); + for (i, (got, want)) in out.chunks[0].iter().zip(legacy_out.iter()).enumerate() { + assert_eq!(got, want, "record {i} bytes differ from the oracle"); + } +} + +/// The same front under a `memory_limit` far below the input: `ReadBlocks` seals +/// several runs, so the front emits `Spill` chunks ahead of the residual. Each +/// run is independently sorted (the merge is a later stage), so the claim here +/// is that every record survives exactly once and each chunk is itself sorted. +#[test] +fn arena_front_chain_seals_multiple_runs_without_losing_records() { + const N_RECORDS: usize = 4_000; + let (header, records) = synthesize_records(N_RECORDS, 0x5EA1_5EA1); + let n_ref = u32::try_from(header.reference_sequences().len()).expect("n_ref fits u32"); + + let out = drive_arena_front( + bgzf_blocks_for(&records, n_ref, 8192), + n_ref, + 64 * 1024, // far below the input ⇒ several runs + 4 * 1024 * 1024, + 4, + ) + .expect("arena front drives to completion"); + + assert!(out.chunks.len() > 1, "a tiny memory limit must seal several runs"); + assert_eq!(out.slot_count as usize, out.chunks.len() - 1, "every run but the last spills"); + assert_eq!(out.total_records, N_RECORDS as u64); + + let mut got: Vec> = out.chunks.into_iter().flatten().collect(); + let mut want: Vec> = records.iter().map(|r| r.as_ref().to_vec()).collect(); + got.sort_unstable(); + want.sort_unstable(); + assert_eq!(got, want, "the sealed runs must carry every input record exactly once"); +} diff --git a/crates/fgumi-pipeline-io/src/source/mod.rs b/crates/fgumi-pipeline-io/src/source/mod.rs new file mode 100644 index 000000000..5d82941cb --- /dev/null +++ b/crates/fgumi-pipeline-io/src/source/mod.rs @@ -0,0 +1 @@ +pub mod read_bam; diff --git a/crates/fgumi-pipeline-io/src/source/read_bam.rs b/crates/fgumi-pipeline-io/src/source/read_bam.rs new file mode 100644 index 000000000..28d55ba90 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/source/read_bam.rs @@ -0,0 +1,446 @@ +//! `ReadBgzfBlocks` source step + `read_bam(path)` convenience helper. +//! +//! Reads raw BGZF blocks from a file (no decompression) and emits them as +//! `BgzfBlock` items with monotonically increasing `batch_serial`. The +//! header bytes are NOT skipped here — they pass through as part of the +//! first block(s); `FindBamBoundaries` strips them downstream. + +use std::collections::VecDeque; +use std::fs::File; +use std::io; +use std::path::Path; + +use fgumi_bam_io::PipelineReaderOpts; +use fgumi_bgzf::reader::read_raw_blocks; +use noodles::sam::Header; + +use crate::types::BgzfBlock; +use fgumi_pipeline_core::{ + Unpushed, + held::HeldSlot, + outputs::OrderedBytesSingle, + queues::QueueSpec, + reorder::BranchOrdering, + step::{Affinity, Step, StepCtx, StepKind, StepOutcome, StepProfile}, +}; + +/// Legacy default blocks-per-batch. +pub const DEFAULT_BLOCKS_PER_BATCH: usize = 16; + +/// `Exclusive + sticky` source step that reads raw BGZF blocks from a file. +/// +/// The reader and the finished flag are plain owned fields, not `Arc`/atomics: +/// this is a `Serial` step, so the runtime drives a single shared instance and +/// never calls `new_worker_copy` on it (only `Parallel` steps are cloned per +/// worker). There is no second owner to share them with. +pub struct ReadBgzfBlocks { + reader: Option>, + blocks_per_batch: usize, + next_serial: u64, + pending: VecDeque, + held: HeldSlot>, + output_byte_limit: u64, + finished: bool, +} + +impl ReadBgzfBlocks { + #[must_use] + pub fn new( + reader: Box, + blocks_per_batch: usize, + output_byte_limit: u64, + ) -> Self { + Self { + reader: Some(reader), + blocks_per_batch: blocks_per_batch.max(1), + next_serial: 0, + pending: VecDeque::new(), + held: HeldSlot::new(), + output_byte_limit, + finished: false, + } + } +} + +impl Step for ReadBgzfBlocks { + type Input = (); + type Outputs = OrderedBytesSingle; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "ReadBgzfBlocks", + kind: StepKind::Serial, + sticky: true, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], + branch_ordering: vec![BranchOrdering::ByItemOrdinal], + } + } + + fn affinity(&self) -> Affinity { + Affinity::Reader + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + // 1. Drain the held slot first. + if let Some(unpushed) = self.held.take() { + match ctx.outputs.retry(unpushed) { + Ok(()) => {} + Err(again) => { + self.held.put(again); + return Ok(StepOutcome::Contention); + } + } + } + + // 2. Drain pending blocks (one per call iteration). + if let Some(block) = self.pending.pop_front() { + match ctx.outputs.push(block) { + Ok(()) => return Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held.put(unpushed); + return Ok(StepOutcome::Progress); + } + } + } + + if self.finished { + return Ok(StepOutcome::Finished); + } + + // 3. Read up to `blocks_per_batch` raw BGZF blocks. The reader is taken + // only at end of stream, so `None` here means `try_run` was called again + // after it already returned `Finished`. + let raw_blocks = { + let reader = self + .reader + .as_mut() + .expect("ReadBgzfBlocks: try_run called after the source reported Finished"); + read_raw_blocks(reader.as_mut(), self.blocks_per_batch)? + }; + + if raw_blocks.is_empty() { + self.finished = true; + // Release the reader (and its 2 MiB BufReader) as soon as the stream + // is drained rather than holding it for the rest of the run. + self.reader = None; + return Ok(StepOutcome::Finished); + } + + for raw in raw_blocks { + let serial = self.next_serial; + self.next_serial += 1; + self.pending.push_back(BgzfBlock { + batch_serial: serial, + uncompressed_size: u32::try_from(raw.uncompressed_size()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidData, + format!( + "ReadBgzfBlocks: BGZF uncompressed_size out of range: {}", + raw.uncompressed_size() + ), + ) + })?, + bytes: raw.data, + }); + } + + if let Some(block) = self.pending.pop_front() { + match ctx.outputs.push(block) { + Ok(()) => Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held.put(unpushed); + Ok(StepOutcome::Progress) + } + } + } else { + Ok(StepOutcome::NoProgress) + } + } +} + +/// Build a [`ReadBgzfBlocks`] step from an already-prepared reader + header. +#[must_use] +pub fn read_bam_from_reader( + reader: Box, + header: Header, + blocks_per_batch: usize, + output_byte_limit: u64, +) -> (ReadBgzfBlocks, Header) { + (ReadBgzfBlocks::new(reader, blocks_per_batch, output_byte_limit), header) +} + +/// Convenience helper: open a BAM file, parse its header, return the +/// `(step, header)` pair. +/// +/// The file is deliberately opened twice: first via +/// `create_raw_bam_reader_with_opts` to parse and return the `Header`, then +/// re-opened with `File::open` and a fresh `BufReader` at offset 0 so the raw +/// BGZF stream — including the header blocks — is emitted in full. Those header +/// blocks are stripped downstream by `FindBamBoundaries`. Do not "optimize" this +/// by reusing the first reader's position: skipping the header blocks corrupts +/// the raw-block stream. +/// +/// # Errors +/// +/// Returns I/O errors from file open or BAM-header parse. +pub fn read_bam>( + path: P, + opts: PipelineReaderOpts, + blocks_per_batch: usize, + output_byte_limit: u64, +) -> io::Result<(ReadBgzfBlocks, Header)> { + let path = path.as_ref(); + let (_, header) = fgumi_bam_io::create_raw_bam_reader_with_opts(path, 1, opts) + .map_err(|e| io::Error::other(format!("create_raw_bam_reader_with_opts: {e}")))?; + + let file = File::open(path)?; + let reader: Box = + Box::new(io::BufReader::with_capacity(2 * 1024 * 1024, file)); + Ok(read_bam_from_reader(reader, header, blocks_per_batch, output_byte_limit)) +} + +/// Stdin counterpart to [`read_bam`]. +/// +/// # Errors +/// +/// Returns I/O errors from stdin read or BAM-header parse. +pub fn read_bam_stdin( + opts: PipelineReaderOpts, + blocks_per_batch: usize, + output_byte_limit: u64, +) -> io::Result<(ReadBgzfBlocks, Header)> { + let (reader, header) = + fgumi_bam_io::create_bam_reader_for_pipeline_with_opts(Path::new("-"), opts).map_err( + |e| io::Error::other(format!("create_bam_reader_for_pipeline_with_opts: {e}")), + )?; + Ok(read_bam_from_reader(reader, header, blocks_per_batch, output_byte_limit)) +} + +/// Path-aware dispatcher: routes to [`read_bam_stdin`] when `path` is a +/// stdin sentinel (`-` or `/dev/stdin`) and to [`read_bam`] otherwise. +/// +/// # Errors +/// +/// Returns I/O errors from file open, stdin read, or BAM-header parse. +pub fn read_bam_auto>( + path: P, + opts: PipelineReaderOpts, + blocks_per_batch: usize, + output_byte_limit: u64, +) -> io::Result<(ReadBgzfBlocks, Header)> { + if fgumi_bam_io::is_stdin_path(path.as_ref()) { + read_bam_stdin(opts, blocks_per_batch, output_byte_limit) + } else { + read_bam(path, opts, blocks_per_batch, output_byte_limit) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[test] + fn profile_advertises_serial_reader_byordinal() { + let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); + let header = noodles::sam::Header::default(); + let writer = fgumi_bam_io::create_raw_bam_writer(&path, &header, 1, 1).unwrap(); + writer.finish().unwrap(); + + let (step, _hdr) = + read_bam(&path, PipelineReaderOpts::default(), DEFAULT_BLOCKS_PER_BATCH, 1024 * 1024) + .unwrap(); + let profile = step.profile(); + assert_eq!(profile.name, "ReadBgzfBlocks"); + assert_eq!(profile.kind, StepKind::Serial); + assert!(profile.sticky); + assert_eq!(step.affinity(), Affinity::Reader); + assert_eq!(profile.branch_ordering, vec![BranchOrdering::ByItemOrdinal]); + assert!(matches!(profile.output_queues[0], QueueSpec::ByteBounded { .. })); + } + + // --------------------------------------------------------------------- + // Driving the step through a real pipeline + // --------------------------------------------------------------------- + + /// Sink that records every `BgzfBlock` the source emits, in arrival order. + struct BlockSink { + received: std::sync::Arc>>, + } + + impl Step for BlockSink { + type Input = BgzfBlock; + type Outputs = (); + + fn profile(&self) -> StepProfile { + StepProfile { + name: "BlockSink", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + match ctx.input.pop() { + Some(block) => { + self.received.lock().push(block); + Ok(StepOutcome::Progress) + } + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } + } + + /// Write `record_count` records to a temp BAM and return its `path` plus the + /// on-disk bytes. + fn temp_bam(record_count: usize) -> (tempfile::TempPath, Vec) { + let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); + let header = noodles::sam::Header::default(); + let mut writer = fgumi_bam_io::create_raw_bam_writer(&path, &header, 1, 1).unwrap(); + for i in 0..record_count { + let name = format!("q{i}"); + let bytes = fgumi_raw_bam::testutil::make_bam_bytes( + 0, + i32::try_from(i).unwrap(), + 0, + name.as_bytes(), + &[], + 10, + -1, + -1, + &[], + ); + writer.write_raw_record(&bytes).unwrap(); + } + writer.finish().unwrap(); + let on_disk = std::fs::read(&path).unwrap(); + (path, on_disk) + } + + /// Run a `ReadBgzfBlocks -> BlockSink` pipeline and return the emitted blocks. + fn drive(path: &Path, blocks_per_batch: usize, threads: usize) -> Vec { + let received = std::sync::Arc::new(parking_lot::Mutex::new(Vec::new())); + let (source, _hdr) = + read_bam(path, PipelineReaderOpts::default(), blocks_per_batch, 1024 * 1024).unwrap(); + let sink = BlockSink { received: std::sync::Arc::clone(&received) }; + + let builder = fgumi_pipeline_core::builder::Pipeline::builder(); + builder.chain(source).chain(sink).into_sink_marker(); + let pipeline = builder.build().unwrap(); + pipeline + .run(fgumi_pipeline_core::builder::PipelineConfig { threads, ..Default::default() }) + .unwrap(); + + // Returned in ARRIVAL order, deliberately unsorted: the step declares + // `BranchOrdering::ByItemOrdinal`, so the sink must already see blocks in + // serial order. Sorting here would normalize out-of-order delivery and let + // an ordering regression pass. + std::mem::take(&mut *received.lock()) + } + + #[rstest] + #[case::single_block_batches(1)] + #[case::default_batching(DEFAULT_BLOCKS_PER_BATCH)] + #[case::larger_than_the_file(1024)] + fn try_run_emits_every_block_with_dense_serials(#[case] blocks_per_batch: usize) { + const BGZF_EOF_LEN: usize = 28; + let (path, on_disk) = temp_bam(64); + let blocks = drive(&path, blocks_per_batch, 1); + + assert!(!blocks.is_empty(), "a non-empty BAM must yield at least one block"); + + // Serials are dense and start at zero. + for (i, block) in blocks.iter().enumerate() { + assert_eq!(block.batch_serial, i as u64, "serial {i} must be dense"); + // `bytes` is the raw *compressed* block (this step does not inflate), + // so it is not the same length as `uncompressed_size`; only that the + // declared inflated size is populated is checkable here. + assert!(block.uncompressed_size > 0, "block {i} must declare an inflated size"); + assert!(!block.bytes.is_empty(), "block {i} must carry its compressed bytes"); + } + + // Concatenating the payloads reproduces the file minus its BGZF EOF block, + // which is what `FindBamBoundaries` downstream expects to receive. + let concatenated: Vec = blocks.iter().flat_map(|b| b.bytes.clone()).collect(); + assert_eq!(concatenated, on_disk[..on_disk.len() - BGZF_EOF_LEN]); + } + + #[test] + fn try_run_emits_the_same_blocks_regardless_of_thread_count() { + let (path, _) = temp_bam(64); + let one = drive(&path, 4, 1); + let many = drive(&path, 4, 4); + assert_eq!(one.len(), many.len(), "block count must not depend on threads"); + for (a, b) in one.iter().zip(many.iter()) { + assert_eq!(a.batch_serial, b.batch_serial); + assert_eq!(a.bytes, b.bytes); + } + } + + #[test] + fn try_run_on_a_header_only_bam_still_emits_the_header_block() { + const BGZF_EOF_LEN: usize = 28; + let (path, on_disk) = temp_bam(0); + let blocks = drive(&path, DEFAULT_BLOCKS_PER_BATCH, 1); + let concatenated: Vec = blocks.iter().flat_map(|b| b.bytes.clone()).collect(); + assert_eq!(concatenated, on_disk[..on_disk.len() - BGZF_EOF_LEN]); + } + + // --------------------------------------------------------------------- + // Constructor + dispatch + // --------------------------------------------------------------------- + + #[rstest] + #[case::zero_clamps_to_one(0, 1)] + #[case::one_stays_one(1, 1)] + #[case::larger_is_preserved(32, 32)] + fn new_clamps_blocks_per_batch_to_at_least_one( + #[case] requested: usize, + #[case] expected: usize, + ) { + let reader: Box = Box::new(io::Cursor::new(Vec::new())); + let step = ReadBgzfBlocks::new(reader, requested, 1024); + assert_eq!(step.blocks_per_batch, expected); + } + + #[test] + fn read_bam_auto_routes_a_regular_path_to_the_file_reader() { + let (path, _) = temp_bam(4); + // `read_bam_auto` on a non-stdin path must behave exactly like `read_bam`: + // same header, and a step that reads the same file. + let (_, auto_hdr) = + read_bam_auto(&path, PipelineReaderOpts::default(), 4, 1024 * 1024).unwrap(); + let (_, direct_hdr) = + read_bam(&path, PipelineReaderOpts::default(), 4, 1024 * 1024).unwrap(); + assert_eq!(auto_hdr, direct_hdr); + assert!(!fgumi_bam_io::is_stdin_path(AsRef::::as_ref(&path))); + } + + #[rstest] + #[case::dash("-")] + #[case::dev_stdin("/dev/stdin")] + fn stdin_sentinels_are_recognised(#[case] sentinel: &str) { + // Guards the branch condition in `read_bam_auto` without consuming the + // process's real stdin, which a test must not do. + assert!(fgumi_bam_io::is_stdin_path(Path::new(sentinel))); + } + + #[test] + fn read_bam_errors_on_a_missing_file() { + // `ReadBgzfBlocks` is not `Debug`, so inspect the variant directly rather + // than via `expect_err`. + let result = read_bam( + Path::new("/nonexistent/definitely/not/here.bam"), + PipelineReaderOpts::default(), + 4, + 1024, + ); + match result { + Ok(_) => panic!("a missing file must not open"), + Err(e) => assert!(!e.to_string().is_empty()), + } + } +} diff --git a/crates/fgumi-pipeline-io/src/types.rs b/crates/fgumi-pipeline-io/src/types.rs new file mode 100644 index 000000000..05f683be7 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/types.rs @@ -0,0 +1,302 @@ +//! Concrete data types that flow through the BAM step library. +//! +//! Every flowing type carries an explicit `batch_serial: u64` field and +//! impls both [`HeapSize`] (so byte-bounded queues can budget memory) and +//! [`Ordered`] (so `BranchOrdering::ByItemOrdinal` reorder stages preserve +//! global ordering across multi-step Parallel transforms). + +use fgumi_pipeline_core::{HeapSize, Ordered}; +use fgumi_raw_bam::RawRecord; + +// ───────────────────────────────────────────────────────────────────────────── +// BgzfBlock — raw compressed BGZF block + read-order serial. +// ───────────────────────────────────────────────────────────────────────────── + +/// Raw compressed BGZF block + parsed metadata. Carries read-order serial. +/// +/// Sentinel/EOF blocks have `bytes` containing the 28-byte BGZF EOF marker +/// and `uncompressed_size = 0`. +#[derive(Debug)] +pub struct BgzfBlock { + /// Read-order serial. Set by `ReadBgzfBlocks` based on block read index. + pub batch_serial: u64, + pub bytes: Vec, + /// Decompressed size, parsed from the BGZF block header. + pub uncompressed_size: u32, +} + +impl HeapSize for BgzfBlock { + fn heap_size(&self) -> usize { + // Byte-bounded queues budget on resident heap, so account for the full + // allocation (`capacity`), not just the populated prefix (`len`). + self.bytes.capacity() + } +} + +impl Ordered for BgzfBlock { + fn ordinal(&self) -> u64 { + self.batch_serial + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// DecompressedBlock — raw bytes from a BGZF decompression, record-aligned +// or not. +// ───────────────────────────────────────────────────────────────────────────── + +/// Decompressed bytes from one or more BGZF blocks. Carries a serial for +/// ordering purposes; record alignment is the consumer's responsibility. +#[derive(Debug)] +pub struct DecompressedBlock { + pub batch_serial: u64, + pub bytes: Vec, +} + +impl HeapSize for DecompressedBlock { + fn heap_size(&self) -> usize { + // Account for the full allocation (`capacity`), not just `len` — see + // the `BgzfBlock` impl above. + self.bytes.capacity() + } +} + +impl Ordered for DecompressedBlock { + fn ordinal(&self) -> u64 { + self.batch_serial + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// RecordBatch — parsed BAM records grouped into a batch. +// ───────────────────────────────────────────────────────────────────────────── + +/// A batch of parsed BAM records, stored as a flat backing buffer + per-record +/// `(start, end)` ranges. +#[derive(Debug)] +pub struct RecordBatch { + batch_serial: u64, + /// All record bodies concatenated, in batch order. + backing: Vec, + /// (start, end) byte ranges into `backing`, one per record. + ranges: Vec<(u32, u32)>, +} + +impl RecordBatch { + /// Construct a batch from a pre-parsed backing buffer and `(start, end)` ranges. + #[must_use] + pub fn from_parsed(batch_serial: u64, backing: Vec, ranges: Vec<(u32, u32)>) -> Self { + Self { batch_serial, backing, ranges } + } + + /// Convenience constructor: serializes a slice of `RawRecord`s into the + /// flat representation. + /// + /// # Panics + /// + /// Panics if the concatenated record bodies exceed `u32::MAX` bytes. + #[must_use] + pub fn new(batch_serial: u64, records: &[RawRecord]) -> Self { + let total: usize = records.iter().map(RawRecord::len).sum(); + let mut backing = Vec::with_capacity(total); + let mut ranges = Vec::with_capacity(records.len()); + for rec in records { + let start = u32::try_from(backing.len()).expect("backing fits in u32"); + backing.extend_from_slice(rec.as_ref()); + let end = u32::try_from(backing.len()).expect("backing fits in u32"); + ranges.push((start, end)); + } + Self { batch_serial, backing, ranges } + } + + /// Self-managed ordinal. + #[must_use] + pub fn batch_serial(&self) -> u64 { + self.batch_serial + } + + /// Number of records in the batch. + #[must_use] + pub fn len(&self) -> usize { + self.ranges.len() + } + + /// `true` iff the batch contains zero records. + #[must_use] + pub fn is_empty(&self) -> bool { + self.ranges.is_empty() + } + + /// Total bytes across all record bodies. + #[must_use] + pub fn total_bytes(&self) -> usize { + self.backing.len() + } + + /// Iterate the record bodies as borrowed byte slices into the backing buffer. + pub fn iter_record_bytes(&self) -> impl Iterator + '_ { + let backing = &self.backing[..]; + self.ranges.iter().map(move |&(s, e)| &backing[s as usize..e as usize]) + } +} + +/// Builder for emit-side `RecordBatch` construction. +#[derive(Debug)] +pub struct RecordBatchBuilder { + batch_serial: u64, + backing: Vec, + ranges: Vec<(u32, u32)>, +} + +impl RecordBatchBuilder { + /// Build an empty builder with reserved capacity. + #[must_use] + pub fn with_capacity(batch_serial: u64, bytes_cap: usize, records_cap: usize) -> Self { + Self { + batch_serial, + backing: Vec::with_capacity(bytes_cap), + ranges: Vec::with_capacity(records_cap), + } + } + + /// Append one record's body bytes. + /// + /// # Panics + /// + /// Panics if accumulated bytes would exceed `u32::MAX`. + pub fn push_record_bytes(&mut self, bytes: &[u8]) { + let start = u32::try_from(self.backing.len()).expect("backing fits in u32"); + self.backing.extend_from_slice(bytes); + let end = u32::try_from(self.backing.len()).expect("backing fits in u32"); + self.ranges.push((start, end)); + } + + /// Number of records appended so far. + #[must_use] + pub fn len(&self) -> usize { + self.ranges.len() + } + + /// `true` iff no records have been appended. + #[must_use] + pub fn is_empty(&self) -> bool { + self.ranges.is_empty() + } + + /// Total record bytes appended so far. + #[must_use] + pub fn total_bytes(&self) -> usize { + self.backing.len() + } + + /// Finalize and produce the `RecordBatch`. Consumes the builder. + #[must_use] + pub fn build(self) -> RecordBatch { + RecordBatch { batch_serial: self.batch_serial, backing: self.backing, ranges: self.ranges } + } +} + +impl HeapSize for RecordBatch { + fn heap_size(&self) -> usize { + // Account for the full allocation (`capacity`) of both buffers, not + // just their populated prefixes — see the `BgzfBlock` impl above. + self.backing.capacity() + self.ranges.capacity() * std::mem::size_of::<(u32, u32)>() + } +} + +impl Ordered for RecordBatch { + fn ordinal(&self) -> u64 { + self.batch_serial + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn bgzf_block_heap_size_matches_bytes_capacity() { + let b = BgzfBlock { batch_serial: 0, bytes: vec![0u8; 1024], uncompressed_size: 4096 }; + assert_eq!(b.heap_size(), 1024); + assert_eq!(b.ordinal(), 0); + } + + #[test] + fn decompressed_block_heap_size_matches_bytes_capacity() { + let b = DecompressedBlock { batch_serial: 7, bytes: vec![0u8; 4096] }; + assert_eq!(b.heap_size(), 4096); + assert_eq!(b.ordinal(), 7); + } + + #[test] + fn heap_size_counts_allocated_capacity_not_logical_len() { + // A buffer with spare capacity (e.g. after `with_capacity`) holds more + // resident heap than its `len` — byte-bounded queues must budget the + // full allocation or they undercount and bypass configured limits. + let mut bytes = Vec::with_capacity(8192); + bytes.extend_from_slice(&[0u8; 100]); + assert!(bytes.capacity() >= 8192 && bytes.len() == 100); + let cap = bytes.capacity(); + + let block = BgzfBlock { batch_serial: 0, bytes, uncompressed_size: 0 }; + assert_eq!(block.heap_size(), cap); + + let mut backing = Vec::with_capacity(4096); + backing.extend_from_slice(&[0u8; 10]); + let mut ranges = Vec::with_capacity(64); + ranges.push((0u32, 10u32)); + let backing_cap = backing.capacity(); + let ranges_cap = ranges.capacity(); + let batch = RecordBatch::from_parsed(0, backing, ranges); + assert_eq!(batch.heap_size(), backing_cap + ranges_cap * std::mem::size_of::<(u32, u32)>()); + } + + #[test] + fn record_batch_total_bytes_sums_record_lengths() { + let r1: RawRecord = vec![0u8; 100].into(); + let r2: RawRecord = vec![0u8; 200].into(); + let batch = RecordBatch::new(3, &[r1, r2]); + assert_eq!(batch.len(), 2); + assert_eq!(batch.total_bytes(), 300); + // `Vec::with_capacity` may over-allocate, so assert against the actual + // allocated capacities rather than the requested sizes (see the sibling + // `heap_size_counts_allocated_capacity_not_logical_len` test). + assert_eq!( + batch.heap_size(), + batch.backing.capacity() + batch.ranges.capacity() * std::mem::size_of::<(u32, u32)>() + ); + assert_eq!(batch.ordinal(), 3); + } + + #[test] + fn record_batch_from_parsed_round_trips_ranges() { + let backing = b"AAABBBBCC".to_vec(); + let ranges = vec![(0u32, 3u32), (3u32, 7u32), (7u32, 9u32)]; + let batch = RecordBatch::from_parsed(11, backing, ranges); + let got: Vec<&[u8]> = batch.iter_record_bytes().collect(); + assert_eq!(got, vec![&b"AAA"[..], &b"BBBB"[..], &b"CC"[..]]); + assert_eq!(batch.len(), 3); + assert_eq!(batch.total_bytes(), 9); + assert_eq!(batch.batch_serial(), 11); + } + + #[test] + fn record_batch_builder_collects_records() { + let mut b = RecordBatchBuilder::with_capacity(0, 64, 4); + b.push_record_bytes(&[0u8; 50]); + b.push_record_bytes(&[0u8; 75]); + assert_eq!(b.len(), 2); + assert!(!b.is_empty()); + assert_eq!(b.total_bytes(), 125); + let batch = b.build(); + let got: Vec = batch.iter_record_bytes().map(<[u8]>::len).collect(); + assert_eq!(got, vec![50, 75]); + // `heap_size` budgets allocated capacity (not logical length): the + // builder was seeded with a 64-byte backing buffer but 125 bytes were + // pushed, so `backing` reallocated and its capacity now exceeds 125. + assert_eq!( + batch.heap_size(), + batch.backing.capacity() + batch.ranges.capacity() * std::mem::size_of::<(u32, u32)>() + ); + assert!(batch.heap_size() >= 125 + 2 * std::mem::size_of::<(u32, u32)>()); + } +} diff --git a/crates/fgumi-sort/src/memory_probe.rs b/crates/fgumi-sort/src/memory_probe.rs index b4164bd59..ce7ef2c82 100644 --- a/crates/fgumi-sort/src/memory_probe.rs +++ b/crates/fgumi-sort/src/memory_probe.rs @@ -78,7 +78,6 @@ mod platform_ffi { /// `mi_stats_print_out(None, null_mut())` uses mimalloc's internal synchronization, /// making it safe to call concurrently with allocation/deallocation on other threads. #[cfg(feature = "memory-debug")] - #[allow(dead_code)] // consumed by main fgumi's unified_pipeline via the crate-root re-export pub fn print_mi_stats() { // SAFETY: mimalloc synchronizes stats collection internally. unsafe { diff --git a/crates/fgumi-sort/src/segmented_buf.rs b/crates/fgumi-sort/src/segmented_buf.rs index 84223f0be..1f03aee67 100644 --- a/crates/fgumi-sort/src/segmented_buf.rs +++ b/crates/fgumi-sort/src/segmented_buf.rs @@ -48,7 +48,6 @@ pub struct SegmentedBuf { cur: usize, } -#[allow(dead_code)] // Some methods are only exercised from tests for now. impl SegmentedBuf { /// Create a new buffer with the given initial capacity hint and segment size. /// diff --git a/crates/fgumi-sort/src/tmp_dir_alloc.rs b/crates/fgumi-sort/src/tmp_dir_alloc.rs index 610f0ba44..18360dc01 100644 --- a/crates/fgumi-sort/src/tmp_dir_alloc.rs +++ b/crates/fgumi-sort/src/tmp_dir_alloc.rs @@ -135,7 +135,6 @@ impl TmpDirAllocator { /// Override the periodic recheck interval (primarily for testing). #[must_use] - #[allow(dead_code)] pub fn with_recheck_interval(mut self, interval: usize) -> Self { self.recheck_interval = interval.max(1); self @@ -175,7 +174,6 @@ impl TmpDirAllocator { /// Drop a directory from rotation (e.g. after `ENOSPC` during a spill write). /// /// Matches by path equality. A no-op if the path isn't currently active. - #[allow(dead_code)] pub fn mark_full(&mut self, dir: &Path) { if let Some(pos) = self.active.iter().position(|d| d == dir) { self.active.remove(pos); diff --git a/scripts/publish-crates.sh b/scripts/publish-crates.sh index 91386f3aa..9da54aed9 100755 --- a/scripts/publish-crates.sh +++ b/scripts/publish-crates.sh @@ -33,6 +33,7 @@ CRATES=( fgumi-sam fgumi-metrics fgumi-sort + fgumi-pipeline-io fgumi-consensus fgumi )