diff --git a/.cargo/config.toml b/.cargo/config.toml index 3d8136b4b..85a8cbb0c 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -14,7 +14,14 @@ ci-publish-order = "run --package xtask -- check-publish-order" ci-doctest = "test --doc --workspace --features compare,simulate,profile-adjacency --locked" # Build the docs and fail on any rustdoc warning (e.g. broken intra-doc links). # RUSTDOCFLAGS="-D warnings" is set by the caller (CI job / pre-push). -ci-doc = "doc --no-deps --workspace --features compare,simulate,profile-adjacency --locked" +# +# `--document-private-items` is load-bearing, not cosmetic. Without it rustdoc +# only checks links on items it renders — i.e. public ones — so a broken link in +# the doc comment of a private fn, a private field, or a `#[cfg(test)]` helper +# passes silently. That is not hypothetical: several such links accumulated +# undetected and had to be found by hand. The flag makes the docs job check every +# doc comment in the workspace, which is what the job's name implies it does. +ci-doc = "doc --no-deps --workspace --document-private-items --features compare,simulate,profile-adjacency --locked" # Run only the #[ignore]-d integration tests (the sort-correctness suites in # tests/integration/{test_sort_correctness,test_async_reader,test_sort_write_index}.rs). # They require the `samtools` binary on PATH — samtools builds the BAM fixtures, @@ -34,7 +41,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/.coderabbit.yaml b/.coderabbit.yaml index 340662f87..a64c51329 100644 --- a/.coderabbit.yaml +++ b/.coderabbit.yaml @@ -164,22 +164,28 @@ reviews: # constructs. Re-read and update them whenever those modules are refactored or # renamed — a stale instruction silently misdirects the reviewer, and the # reviewer will never tell you it is chasing a symbol that no longer exists. - # Every symbol named below was verified with `git grep` to resolve on `main` - # at the commit that introduced this revision; re-verify when touching these - # modules. + # Every symbol named below was verified with `git grep` to resolve at the + # commit that introduced this revision; re-verify when touching these modules. # - # Deliberately NOT listed, because they do not exist on `main` and a glob that - # matches nothing is exactly the failure mode above — add an entry in the same - # PR that lands the code: - # - `crates/fgumi-pipeline-core/**`, `crates/fgumi-pipeline-io/**` (on - # `feat-runall`; the byte-bound and unbounded-drain rules in the - # `unified_pipeline` entry apply to them verbatim once merged) + # The two legacy homes share one entry: issue #330 renames + # `src/lib/unified_pipeline` to `src/lib/pipeline`, so the same instructions + # apply to whichever is present. `crates/fgumi-pipeline-core` is the extracted + # typed-step engine and gets its OWN entry below: it is a different design and + # names different symbols (`queues.rs` / `ByteBoundedQueue`, not `queue.rs` / + # `OrderedQueue`), so folding it into the legacy entry would point the reviewer + # at symbols that do not exist there. + # + # Not yet listed, and a glob that matches nothing is exactly the failure mode + # above — add an entry in the same PR that lands the code: + # - `crates/fgumi-pipeline-io/**` — does not exist here yet; the byte-bound + # and unbounded-drain rules in the pipeline entry apply to it verbatim once + # it lands. # - `crates/fgumi-fmt/**`, `crates/fgumi-cli-common/**`, - # `crates/fgumi-cli-macros/**` (in flight) - # `src/lib/pipeline/**` is listed alongside `unified_pipeline` only because - # issue #330 renames the latter to the former; today it matches nothing. + # `crates/fgumi-cli-macros/**` — these DO exist here but still have no + # instructions of their own, so they are reviewed by the generic themes + # only. path_instructions: - - path: "src/lib/{unified_pipeline,pipeline}/**/*.rs" + - path: "{src/lib/unified_pipeline,src/lib/pipeline}/**/*.rs" instructions: >- This is the hand-rolled concurrent step pipeline; its bugs are deadlocks, lost output, unbounded memory, and panics on malformed input — not style. @@ -207,6 +213,35 @@ reviews: a real, previously-shipped crash on corruption-controlled input. This module contains no `unsafe`; treat newly introduced `unsafe` here as out of policy and require it to be justified in CLAUDE.md first. + - path: "crates/fgumi-pipeline-core/src/**/*.rs" + instructions: >- + This is the extracted typed-step pipeline engine (the successor to the + `src/lib/*_pipeline` design above); its bugs are deadlocks, lost output, + unbounded memory, and panics on malformed input — not style. It names + DIFFERENT symbols than the legacy entry: queues live in `queues.rs` + (`ItemQueue` trait, `QueueSpec`, `CountBoundedQueue`, `ByteBoundedQueue`, + `UnboundedQueue`), NOT `queue.rs`/`OrderedQueue`. Require that no queue or + reorder buffer grows without a bound: byte-bounded transports go through + `ByteBoundedQueue` (steady-state memory a function of config, + not input size), and consumers/drains must stay unbounded — bound the + producers instead. Flag a cap checked on only one sub-condition so another + path bypasses it, and byte accounting that measures logical length where + the memory held is allocation capacity. Cancellation and error propagation + run through `PipelineSignal` / `CancelHandle` (`signal.rs`): require every + worker loop and the fused driver to observe `signal.is_done()` promptly and + never block forever on a channel whose peer has exited. Liveness for the + deadlock monitor is the per-worker `LivenessCounter` (`liveness.rs`), bumped + only on productive `StepOutcome::Progress`/`Finished`; the fused + single-thread path (`runtime/fused.rs`) is NOT monitored and carries its own + stall budget instead. The ONLY approved `unsafe` in this crate is the typed- + handle dispatch cache in `erased.rs` (`TypedStep`/`TypedStep2` + `resolve_input`/`resolve_outputs`), whose soundness rests on the invariants + documented in CLAUDE.md — chiefly that every step instance is dropped before + the `ChainContexts` it cached from; flag any change that weakens those + invariants or adds `unsafe` elsewhere without a CLAUDE.md allowlist entry. + Byte-level framing that parses length/offset fields out of the input stream + must validate against the remaining buffer with checked arithmetic before + slicing. - path: "crates/fgumi-sort/**/*.rs" instructions: >- This is the sort engine, including approved `unsafe` hot paths (LSD radix diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 57765d4ee..637c0937e 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -61,6 +61,39 @@ jobs: - name: Clippy check run: cargo ci-lint + # `merge_slots.rs` swaps `std::sync` for `loom::sync` under `--cfg loom`, and + # `tests/loom_merge_slots.rs` is `#![cfg(loom)]`. Neither is built by any other + # job, so without this one a broken model — or a `cfg(loom)` build that stopped + # compiling — is invisible: the test target compiles to an empty binary under a + # normal build and reports success. That silence is the whole reason this job + # exists. + # + # Per-PR rather than on a schedule, unlike `miri.yml` and `stress.yml`. Those + # two are scheduled because they are non-deterministic in ways unrelated to a + # given change (a nightly-toolchain regression; timing sensitivity). Loom + # explores a bounded state space deterministically — four of the five models + # under a preemption bound, see the test's "Preemption-bounded exploration" + # note — so the same input yields the same verdict with no flake, which is + # what lets it gate a PR, and a gate is what stops it rotting. + loom: + runs-on: ubuntu-latest + timeout-minutes: 20 + steps: + - name: Checkout code + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable + - name: Set up compilation cache (sccache) + uses: mozilla-actions/sccache-action@fc920bf0ec8de6ee65d409111f7ec508035751ba # v0.0.11 + # Release: the models explore enough interleavings that a debug build takes + # several times as long for the same verdict. + - name: Loom model check (SortMergeSlot) + run: cargo test -p fgumi-sort --test loom_merge_slots --release + env: + RUSTFLAGS: --cfg loom + coverage: runs-on: ubuntu-latest timeout-minutes: 20 @@ -84,7 +117,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/.github/workflows/miri.yml b/.github/workflows/miri.yml index 4828258b0..f848b8c3b 100644 --- a/.github/workflows/miri.yml +++ b/.github/workflows/miri.yml @@ -5,12 +5,19 @@ name: Miri (undefined-behavior check) # undefined behavior (out-of-bounds, invalid pointer use, aliasing violations), # turning those invariants into a machine-checked gate. # -# Scope: `fgumi-raw-bam`, which contains the raw-pointer queryname comparator -# (`natural_compare` / `natural_compare_nul`, via `get_unchecked` and `*const u8` -# walks) and has no FFI. `fgumi-sort` is intentionally not covered yet: its -# `memory_probe` module calls the mimalloc / mach2 FFI, which Miri cannot execute, -# so covering its radix-sort unsafe needs `#[cfg(not(miri))]` guards first -# (tracked as a follow-up). +# Scope: the two crates whose approved `unsafe` is pure Rust (no FFI), each +# narrowed to the module that carries it — +# - `fgumi-raw-bam::sort` — the raw-pointer queryname comparator +# (`natural_compare` / `natural_compare_nul`, via `get_unchecked` and +# `*const u8` walks). +# - `fgumi-pipeline-core::erased` — the typed-handle dispatch cache, four +# `mem::transmute`s that widen `&'a Handle` to `&'static` for storage and +# narrow it back on read. Aliasing/lifetime UB is exactly what Stacked +# Borrows checks, so this is the one place Miri adds signal the type system +# cannot. +# `fgumi-sort` is intentionally not covered yet: its `memory_probe` module calls +# the mimalloc / mach2 FFI, which Miri cannot execute, so covering its radix-sort +# unsafe needs `#[cfg(not(miri))]` guards first (tracked as a follow-up). # # Two triggers, two jobs to do. # @@ -117,7 +124,80 @@ jobs: # discarding); and reading — on failure, writing — the # `.proptest-regressions` file the lookup resolves to. MIRIFLAGS: -Zmiri-disable-isolation - run: cargo +nightly miri test -p fgumi-raw-bam sort + # + # `--list` first and fail on an empty selection: `cargo test ` + # exits 0 when the filter matches nothing, so renaming or moving the + # module would silently retire this gate while CI stayed green — the same + # zero-match hazard `compile_fail.rs`'s `EXPECTED_FIXTURES` guards. + run: | + set -euo pipefail + # The zero-match guard below proves the filter selects tests; it cannot + # prove the filter still COVERS the crate's `unsafe`. Pin that too: if a + # NEW `#[allow(unsafe_code)]` site appears outside the modules known to + # carry approved unsafe, the `sort` filter keeps matching and Miri + # silently stops checking the moved/new site — so flag it. + # + # Two modules are excluded because both carry documented, approved + # unsafe (see CLAUDE.md → "Unsafe Code"): + # - `sort.rs` — the natural-order queryname comparator, which the + # `sort` filter below DOES exercise under Miri. + # - `raw_bam_record.rs` — `read_raw_record`'s spare-capacity read. + # The `sort` filter does not run its tests, so this unsafe is not + # Miri-checked here (matching `main`, whose miri job runs only the + # `sort` filter too). Bringing it under Miri needs a separate + # `read_raw_record`-filtered run; tracked as a follow-up. Excluded + # here so this PRE-EXISTING approved site does not fail the guard. + stray=$(grep -rlE '^[[:space:]]*#!?\[[^]]*allow\(unsafe_code\)' crates/fgumi-raw-bam/src \ + | grep -vE '^crates/fgumi-raw-bam/src/(sort|raw_bam_record)\.rs$' || true) + if [ -n "${stray}" ]; then + echo "::error::unsafe_code outside the Miri-scoped fgumi-raw-bam modules: ${stray}" + exit 1 + fi + cargo +nightly miri test -p fgumi-raw-bam sort -- --list > "${RUNNER_TEMP}/miri-raw-bam-sort.list" + n=$(grep -c ': test$' "${RUNNER_TEMP}/miri-raw-bam-sort.list" || true) + echo "miri: the 'sort' filter matched ${n} test(s) in fgumi-raw-bam" + if [ "${n}" -eq 0 ]; then + echo "::error::the 'sort' filter matched no tests in fgumi-raw-bam — the Miri scope is stale" + exit 1 + fi + cargo +nightly miri test -p fgumi-raw-bam sort + - name: Miri — fgumi-pipeline-core typed-handle dispatch cache + # Scope to the `erased` module for the same reason: all four + # `#[allow(unsafe_code)]` sites in the crate live there + # (`TypedStep::resolve_input`/`resolve_outputs` and the `TypedStep2` + # pair). These 25 tests run clean under Miri in ~6s. The rest of the + # crate is deliberately excluded: the `builder` / `runtime` tests spawn + # worker threads and run a full pipeline, which takes Miri well over ten + # minutes, and three of them (`detached_two_sided_no_deadlock`, + # `driver_round_robins_all_live_before_parking`, + # `sticky_holding_source_yields_to_its_draining_consumer`) guard against + # a wedge with a WALL-CLOCK watchdog that `process::abort()`s — under + # Miri's slowdown that fires on a healthy run. Widening this scope means + # giving those watchdogs a `#[cfg(miri)]` budget first. + # + # Zero-match guard, as on the `fgumi-raw-bam` step above: this filter is + # the only thing pointing Miri at the crate's four `unsafe` sites, and an + # empty selection would pass silently. + run: | + set -euo pipefail + # Location guard, as on the `fgumi-raw-bam` step above: the step comment + # claims all four `#[allow(unsafe_code)]` sites live in `erased`, and + # nothing else enforces it. A site that moves elsewhere would leave the + # filter matching and the moved site unchecked. + stray=$(grep -rlE '^[[:space:]]*#!?\[[^]]*allow\(unsafe_code\)' crates/fgumi-pipeline-core/src \ + | grep -v '^crates/fgumi-pipeline-core/src/erased.rs$' || true) + if [ -n "${stray}" ]; then + echo "::error::unsafe_code outside the Miri-scoped 'erased' module: ${stray}" + exit 1 + fi + cargo +nightly miri test -p fgumi-pipeline-core erased -- --list > "${RUNNER_TEMP}/miri-pipeline-core-erased.list" + n=$(grep -c ': test$' "${RUNNER_TEMP}/miri-pipeline-core-erased.list" || true) + echo "miri: the 'erased' filter matched ${n} test(s) in fgumi-pipeline-core" + if [ "${n}" -eq 0 ]; then + echo "::error::the 'erased' filter matched no tests in fgumi-pipeline-core — the Miri scope is stale" + exit 1 + fi + cargo +nightly miri test -p fgumi-pipeline-core erased # Nothing watches a cron's result, so a failure has to announce itself. See the # action for why it reuses one issue per workflow. diff --git a/CLAUDE.md b/CLAUDE.md index 4a3b6566f..c64473b28 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -187,16 +187,70 @@ of records) and a safe rewrite measurably regresses sort throughput. `RawQuerynameKey::new`). - **`crates/fgumi-sort/src/radix.rs`** — internal radix-sort helpers; see file comments for the `SAFETY:` invariants. -- **`crates/fgumi-sort/src/phase1_keys.rs`** — one `#[allow(unsafe_code)]` site in - `prefetch_read_l1`, a software-prefetch hint issued while a deferred key batch - scans record bodies out of a sealed arena segment. Those bytes were written up - to a segment ago and are cold by the time a worker reads them, so the scan is - latency-bound on cache misses. SAFETY: both `prfm pldl1keep` (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 — and the - argument is a live `&u8`, so the pointer is valid to name. A no-op on other - architectures. Approved for the same reason as the other sort hot paths: it - runs once per record over hundreds of millions of records. +- **`crates/fgumi-sort/src/ref_sort.rs`** — one `#[allow(unsafe_code)]` site in + `sort_coordinate_refs`: a pointer cast reinterpreting `&mut [RecordRef]` as + `&mut [CoordSortRef]` to feed the parallel `voracious_mt_sort` radix on the + large-input / multi-thread coordinate path. SAFETY: `CoordSortRef` is + `#[repr(transparent)]` over `RecordRef`, so the two have identical size, + alignment, and layout; the pointer cast (clippy rejects a ref-to-ref + `transmute`) is sound. This is approved for the same reason as the other sort + hot paths — it runs once per coordinate sort over millions of record refs, and + the parallel radix is measurably (~4×) faster than the safe single-threaded + fallback. Two tests cover this, and they establish *different* things — do not + read either as proving both: + - `prop_ref_sort_matches_copy_sorter` runs `sort_threads = 1`, so it proves + **byte parity** of the *serial* path against the copy sorter. It never + reaches the cast. + - `parallel_coordinate_sort_matches_serial_radix_at_threshold` is the one that + exercises the cast — 4 threads at the parallel cutoff over deliberately tied + keys — but it asserts **ref order**, comparing `(sort_key, offset)` against + the serial radix. It does not compare serialized bytes. + + Byte-identity of the parallel path therefore follows by construction rather + than by direct assertion: the chunk is materialized from the refs in order, so + identical ref order yields identical output bytes. That is a sound inference, + but it is an inference, and a change that broke the materialization step would + not be caught by either test. +- **`crates/fgumi-sort/src/segmented_buf.rs`** — two `#[allow(unsafe_code)]` + sites backing the arena record buffer (a segmented, append-only `Vec` the + sort engine decompresses/frames records into without per-record allocation): + - `grow_uninit` — grows a segment's live `len` over + reserved-but-uninitialized bytes via `Vec::set_len` (after `reserve`), to skip + zero-filling the slot on the ingest hot path. SAFETY: after `reserve(additional)`, + `capacity() >= new_len`, so `set_len(new_len)` only extends `len` over + already-allocated bytes; `u8` has no drop glue and no validity invariant, so + growing `len` over uninitialized bytes is not itself UB — the documented contract + requires the caller to fully write the grown `[old_len, new_len)` region (via + `slice_mut`) before any read, and a read-before-write is the only UB, which the + contract forbids. Pointer stability is a *separate* invariant: this `reserve` + MAY reallocate and move the segment, so no `slice_mut` borrow may be live across + a `grow_uninit` — the borrow checker enforces this in the single-threaded case + (`&mut self`), and `reserve_full_capacity`, called once per fresh segment, + makes the per-slot `reserve` a no-op so it cannot move the inner buffer. + That covers the inner buffer ONLY. It is NOT sufficient to run `grow_uninit` + concurrently with a live `slice_mut` — that overlap is exactly what + `slice_mut`'s third precondition forbids, and Miri reports UB on it. + `grow_uninit` may reach `advance_segment`, which pushes to `self.segments`; + that push reallocates the OUTER `Vec>` whenever the vector is at + capacity, freeing the buffer a concurrent reader is indexing + (use-after-free), while its `set_len` races that reader's bounds check on + every call. Neither hazard is closed by pre-sizing, and neither can be: + the two methods take `&mut self` and `&self`, so overlapping them is an + aliasing violation whatever the caller does. That is why this is stated as + a precondition rather than engineered around. The supported concurrent shape + is to reserve every slot for a segment before handing any of them to workers. + - `slice_mut` — synthesizes a `&mut [u8]` slot from a shared `&self` + (so disjoint slots can be written concurrently by different workers). SAFETY: + the asserted `seg_offset + len <= seg.len()` keeps the range inside the + segment's live region (in-bounds pointer + length); the `&mut` is sound only + under the caller's disjointness contract — each call's range must not overlap + any other concurrently-borrowed range — so the produced `&mut` aliases no + other `&`/`&mut`, AND under the separate requirement that no `&mut self` + method runs while the borrow is live (see the `grow_uninit` note above). + Approved for the same reason as the other sort hot paths: + ingest runs once per record over millions-to-billions of records, and a safe + rewrite (zero-fill on grow, or an owning `Vec` per slot) measurably regresses + sort throughput. ### Approved natural-order comparator (fgumi-raw-bam) @@ -208,16 +262,17 @@ path: the comparator runs once per sort-key comparison, and the safe form (re-bounds-checking every byte / re-validating the null terminator) measurably regresses `samtools sort -n`–style throughput. -- **`crates/fgumi-raw-bam/src/sort.rs`** — four `#[allow(unsafe_code)]` sites: - - `natural_compare` (line ~80) — `get_unchecked` over `&[u8]` in the digit-run +- **`crates/fgumi-raw-bam/src/sort.rs`** — two production `#[allow(unsafe_code)]` + sites plus two test sites (the `#[cfg(test)]` boundary sits between them): + - `natural_compare` — `get_unchecked` over `&[u8]` in the digit-run hot loop. SAFETY: indices are bounded by the loop invariants `pa < alen` / `pb < blen`, asserted in `debug_assert!` for the `at` helper. - - `natural_compare_nul` (line ~180) — raw `*const u8` walk that mirrors + - `natural_compare_nul` — raw `*const u8` walk that mirrors samtools' `strnum_cmp`. SAFETY: caller guarantees both pointers are null-terminated; `RawQuerynameKey::new` enforces this for every production call site. - - `compare_nul` test helper and the `proptest` agreement test (lines ~273 and - ~300) — push an explicit NUL into a `Vec` then take `as_ptr()`. + - `compare_nul` test helper and the `proptest` agreement test — push an + explicit NUL into a `Vec` then take `as_ptr()`. SAFETY: the buffers are `to_vec()` + push, so the pointer is valid and null-terminated for the call's lifetime. @@ -274,6 +329,67 @@ contract rather than a weakened assertion: `ENCODE_LUT` maps non-ACGT bytes to a don't-care value, which the SIMD paths write through unconditionally while the scalar path leaves those lanes zero. +### Approved typed-handle cache (fgumi-pipeline-core) + +The pipeline runtime hands each step its input/output handles as +`&dyn Any`, so the `ErasedStep` adapter must `downcast_ref` them back to +their concrete types. That downcast sits on the per-item dispatch path: a +4-thread CODEC 8M run spends ≈1.6% of samples on `TypeId` compares alone. +The adapter resolves the handles once and caches them, which requires +storing a reference whose real lifetime the struct cannot name. + +- **`crates/fgumi-pipeline-core/src/erased.rs`** — four `#[allow(unsafe_code)]` + sites, two on `TypedStep` (`resolve_input`, `resolve_outputs`) and two on + `TypedStep2` (`resolve_inputs`, `resolve_outputs`). Each is a + `std::mem::transmute` that extends a `&'a Handle` to `&'static Handle` for + storage in the cache slot, and narrows it back to `&'a` on read. No pointer + is dereferenced through the `'static` form. SAFETY rests on three invariants + documented on `TypedStep` itself: the handle boxes are owned by + `ChainContexts` (alive for the whole `Pipeline::run`), every step instance is + dropped before those contexts are, and every dispatch passes the *same* box + for a given `step_idx`. The third is the one the compiler cannot check, so + each cache slot stores the address of the erased box it was resolved from and + every cached hit `assert!`s it — **unconditionally, release included**, so a + step reused across two *live* `ChainContexts` panics instead of reading + through the wrong one. The check is one load and one compare, not the + `downcast_ref` `TypeId` probe the cache exists to elide. A debug-only + `debug_assert!` keeps re-resolving the typed pointer as a second diagnostic. + Treat the assert as defence in depth, not as the safety argument: it compares + data addresses only, so an allocator that reuses a freed box's address defeats + it. Soundness rests on the second invariant — every step instance is dropped + before the contexts it cached from — which is what must be preserved by any + future change. `run_fused_single_thread` (`runtime/fused.rs`) is the one path + where the natural drop order inverts it: `steps` is a by-value parameter and + `ChainContexts` is a local, and parameters drop *after* locals. It re-binds + `steps` as a local declared *after* the contexts, so reverse-declaration order + drops the steps before the contexts on every exit — the normal return and a + panic unwinding out of a step — rather than relying on an explicit drop that an + unwind would skip. + +### Test-only `#[allow(unsafe_code)]` sites + +The counts above are **production** sites. Tests carry their own +`#[allow(unsafe_code)]` where they exercise an already-approved `unsafe` API +directly. In `fgumi-sort` those are `segmented_buf.rs` (8), `ref_sort.rs` (3), +and `chunk_sorter.rs` (2), all of them driving `SegmentedBuf::grow_uninit` / +`slice_mut` to check a reserve-then-write round-trip, or that a sorted chunk +matches its oracle. In `fgumi-raw-bam` they are the `compare_nul` helper and the +`proptest` agreement test in `src/sort.rs`, both already named above. They add no +new `unsafe` *surface*: each calls a function already justified above, under that +function's documented contract. + +A raw `grep -c 'allow(unsafe_code)'` over either crate therefore reports more +sites than this document names, and the difference is tests plus prose. When +auditing, count **attributes**, not matches: exclude the production sites listed +above, and exclude occurrences inside `///` / `//!` comments that merely *name* +the attribute (`segmented_buf.rs` has one such mention, which is why a naive +`grep -c` there over-counts the test sites by one). Then check each entry's own +wording for whether it is quoting a production-only count. + +Entries name the function rather than a line number on purpose: approximate +line numbers go stale the moment anything above them moves, and a confidently +wrong pointer is worse than none. + Any new `unsafe` site must extend this list and explain why the safe alternative is unacceptable. Do not introduce `unsafe` outside the crates listed in this section. diff --git a/Cargo.lock b/Cargo.lock index 6a4027d9a..2223a7fdb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -134,6 +134,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "arrayvec" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" + [[package]] name = "autocfg" version = "1.5.1" @@ -726,6 +732,8 @@ dependencies = [ "fgumi-consensus", "fgumi-dna", "fgumi-metrics", + "fgumi-pipeline-core", + "fgumi-pipeline-io", "fgumi-raw-bam", "fgumi-sam", "fgumi-simd-fastq", @@ -745,7 +753,6 @@ dependencies = [ "nix", "noodles", "noodles-bgzf", - "num_cpus", "parking_lot", "proptest", "rand 0.10.2", @@ -764,6 +771,7 @@ dependencies = [ name = "fgumi-bam-io" version = "0.7.0" dependencies = [ + "ahash", "anyhow", "bgzf", "bstr 1.13.1", @@ -794,6 +802,34 @@ dependencies = [ "rstest", ] +[[package]] +name = "fgumi-cli-common" +version = "0.7.0" +dependencies = [ + "anyhow", + "bytesize", + "clap", + "fgumi-fmt", + "log", + "rstest", + "sysinfo", + "tempfile", + "thiserror 2.0.20", +] + +[[package]] +name = "fgumi-cli-macros" +version = "0.7.0" +dependencies = [ + "anyhow", + "clap", + "proc-macro2", + "quote", + "rstest", + "syn 2.0.119", + "trybuild", +] + [[package]] name = "fgumi-consensus" version = "0.7.0" @@ -823,6 +859,14 @@ dependencies = [ "rstest", ] +[[package]] +name = "fgumi-fmt" +version = "0.7.0" +dependencies = [ + "num-format", + "rstest", +] + [[package]] name = "fgumi-metrics" version = "0.7.0" @@ -839,6 +883,41 @@ dependencies = [ "tempfile", ] +[[package]] +name = "fgumi-pipeline-core" +version = "0.7.0" +dependencies = [ + "ahash", + "anyhow", + "crossbeam-queue", + "log", + "noodles", + "parking_lot", + "proptest", + "rstest", + "trybuild", +] + +[[package]] +name = "fgumi-pipeline-io" +version = "0.7.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.7.0" @@ -908,6 +987,7 @@ dependencies = [ "libdeflater", "libmimalloc-sys", "log", + "loom", "mach2", "nix", "noodles", @@ -917,7 +997,9 @@ dependencies = [ "rayon", "rstest", "rustix", + "smallvec", "tempfile", + "voracious_radix_sort", "zstd", ] @@ -1023,6 +1105,21 @@ dependencies = [ "slab", ] +[[package]] +name = "generator" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3b854b0e584ead1a33f18b2fcad7cf7be18b3875c78816b753639aa501513ae" +dependencies = [ + "cc", + "cfg-if", + "libc", + "log", + "rustversion", + "windows-link", + "windows-result", +] + [[package]] name = "getrandom" version = "0.2.17" @@ -1116,12 +1213,6 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" -[[package]] -name = "hermit-abi" -version = "0.5.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" - [[package]] name = "iana-time-zone" version = "0.1.65" @@ -1397,6 +1488,19 @@ version = "0.4.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" +[[package]] +name = "loom" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "419e0dc8046cb947daa77eb95ae174acfbddb7673b4151f56d1eed8e93fbfaca" +dependencies = [ + "cfg-if", + "generator", + "scoped-tls", + "tracing", + "tracing-subscriber", +] + [[package]] name = "lru" version = "0.18.2" @@ -1412,6 +1516,15 @@ version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dae608c151f68243f2b000364e1f7b186d9c29845f7d2d85bd31b9ad77ad552b" +[[package]] +name = "matchers" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata 0.4.18", +] + [[package]] name = "matrixmultiply" version = "0.3.11" @@ -1612,6 +1725,15 @@ dependencies = [ "winapi", ] +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys", +] + [[package]] name = "num-bigint" version = "0.4.8" @@ -1631,6 +1753,16 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-format" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a652d9771a63711fd3c3deb670acfbe5c30a4072e664d7a3bf5a9e1056ac72c3" +dependencies = [ + "arrayvec", + "itoa 1.0.18", +] + [[package]] name = "num-integer" version = "0.1.47" @@ -1661,16 +1793,6 @@ dependencies = [ "libm", ] -[[package]] -name = "num_cpus" -version = "1.17.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" -dependencies = [ - "hermit-abi", - "libc", -] - [[package]] name = "objc2-core-foundation" version = "0.3.2" @@ -2206,6 +2328,12 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "scoped-tls" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1cf6437eb19a8f4a6cc0f7dca544973b0b78843adbfeb3683d1a94a0024a294" + [[package]] name = "scoped_threadpool" version = "0.1.9" @@ -2303,6 +2431,24 @@ dependencies = [ "zmij", ] +[[package]] +name = "serde_spanned" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26" +dependencies = [ + "serde_core", +] + +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + [[package]] name = "shlex" version = "2.0.1" @@ -2445,6 +2591,12 @@ dependencies = [ "windows", ] +[[package]] +name = "target-triple" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3a6bfce3d99adfa72d24750a61f782f3036a81e7f86d8841ee1326deaebd171" + [[package]] name = "tempfile" version = "3.27.0" @@ -2458,6 +2610,15 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "termcolor" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06794f8f6c5c898b3275aebefa6b8a1cb24cd2c6c79397ab15774837a0bc5755" +dependencies = [ + "winapi-util", +] + [[package]] name = "thiserror" version = "1.0.69" @@ -2504,6 +2665,15 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3bf63baf9f5039dadc247375c29eb13706706cfde997d0330d05aa63a77d8820" +[[package]] +name = "thread_local" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070" +dependencies = [ + "cfg-if", +] + [[package]] name = "tinytemplate" version = "1.2.1" @@ -2514,6 +2684,21 @@ dependencies = [ "serde_json", ] +[[package]] +name = "toml" +version = "1.1.4+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3aace63f4bbcdfc2c965b059de67119c89c4017a70d633be6c104910f67056f5" +dependencies = [ + "indexmap", + "serde_core", + "serde_spanned", + "toml_datetime", + "toml_parser", + "toml_writer", + "winnow", +] + [[package]] name = "toml_datetime" version = "1.1.1+spec-1.1.0" @@ -2544,6 +2729,76 @@ dependencies = [ "winnow", ] +[[package]] +name = "toml_writer" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d56353a2a665ad0f41a421187180aab746c8c325620617ad883a99a1cbe66d2" + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "pin-project-lite", + "tracing-core", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex-automata 0.4.18", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", +] + +[[package]] +name = "trybuild" +version = "1.0.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06649c6f63d86604ba0c8950d5a1829fc9a17afd70fc6629f481d75b6a624c78" +dependencies = [ + "glob", + "serde", + "serde_derive", + "serde_json", + "target-triple", + "termcolor", + "toml", +] + [[package]] name = "typenum" version = "1.20.1" @@ -2568,6 +2823,12 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + [[package]] name = "vcpkg" version = "0.2.15" @@ -2580,6 +2841,15 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "voracious_radix_sort" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "446e7ffcb6c27a71d05af7e51ef2ee5b71c48424b122a832f2439651e1914899" +dependencies = [ + "rayon", +] + [[package]] name = "wait-timeout" version = "0.2.1" diff --git a/Cargo.toml b/Cargo.toml index ee2d8d986..20cf6ec61 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-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/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] @@ -22,9 +22,14 @@ flate2 = { version = "1.1", features = ["zlib-rs"] } # Internal crates. Versions are bumped in lockstep by release-plz. fgumi-bam-io = { version = "0.7.0", path = "crates/fgumi-bam-io" } fgumi-bgzf = { version = "0.7.0", path = "crates/fgumi-bgzf" } +fgumi-cli-common = { version = "0.7.0", path = "crates/fgumi-cli-common" } +fgumi-cli-macros = { version = "0.7.0", path = "crates/fgumi-cli-macros" } +fgumi-fmt = { version = "0.7.0", path = "crates/fgumi-fmt" } fgumi-consensus = { version = "0.7.0", path = "crates/fgumi-consensus", default-features = false } fgumi-dna = { version = "0.7.0", path = "crates/fgumi-dna" } fgumi-metrics = { version = "0.7.0", path = "crates/fgumi-metrics" } +fgumi-pipeline-core = { version = "0.7.0", path = "crates/fgumi-pipeline-core" } +fgumi-pipeline-io = { version = "0.7.0", path = "crates/fgumi-pipeline-io" } fgumi-raw-bam = { version = "0.7.0", path = "crates/fgumi-raw-bam" } fgumi-sam = { version = "0.7.0", path = "crates/fgumi-sam" } fgumi-simd-fastq = { version = "0.7.0", path = "crates/fgumi-simd-fastq" } @@ -53,7 +58,7 @@ fgoxide = "0.6.0" fs4 = { version = "1.1", default-features = false, features = ["sync"] } libdeflater = "1.22" libmimalloc-sys = { version = "0.1.44", features = ["extended"] } -log = "0" +log = "0.4" mach2 = "0.6" memchr = "2" nix = { version = "0.31", default-features = false, features = ["fs"] } @@ -66,8 +71,10 @@ nix = { version = "0.31", default-features = false, features = ["fs"] } noodles = "0.115.0" noodles-bgzf = { version = "0.51.0", features = ["libdeflate"] } noodles-csi = "0.60.0" +num-format = "0.4" parking_lot = "0.12" proptest = "1.10" +quote = "1" rand = "0.10" # Chosen over the two syscall wrappers already in `fgumi-sort` on purpose: # `libc::getrlimit` is `unsafe` and the crate root is `#![deny(unsafe_code)]`, @@ -84,8 +91,13 @@ rustix = { version = "1.1", default-features = false, features = ["std"] } rayon = "1.10" rstest = "0.26" serde = { version = "1.0.228", features = ["derive"] } +# Held at 0.38 because sysinfo >=0.39 requires rustc 1.95 > our MSRV (rust-toolchain.toml pins +# 1.93). Centralized so `fgumi` and `fgumi-cli-common` cannot drift onto different versions and +# silently raise the effective MSRV. +sysinfo = { version = "0.38", default-features = false, features = ["system"] } tempfile = "3.3.0" thiserror = "2" +trybuild = "1.0" wide = "1.5" [package] @@ -134,9 +146,8 @@ csv = "1.1" memchr = { workspace = true } statrs = "0.18" bytesize = { workspace = true } -sysinfo = { version = "0.38", default-features = false, features = ["system"] } # unconditionally used by validate_against_system_memory; held at 0.38 because sysinfo >=0.39 requires rustc 1.95 > our MSRV (rust-toolchain.toml pins 1.93) +sysinfo = { workspace = true } # unconditionally used by validate_against_system_memory fs4 = { workspace = true } # cross-platform filesystem free-space queries -num_cpus = "1.16" rayon = { workspace = true } approx = { workspace = true } ahash = { workspace = true } @@ -154,6 +165,8 @@ fgumi-umi = { workspace = true } fgumi-consensus = { workspace = true } fgumi-bam-io = { workspace = true } fgumi-sort = { workspace = true } +fgumi-pipeline-core = { workspace = true } +fgumi-pipeline-io = { workspace = true } [target.'cfg(target_os = "macos")'.dependencies] mach2 = { workspace = true } @@ -203,6 +216,12 @@ harness = false # Exercises vanilla consensus / overlap helpers that live behind `simplex`. required-features = ["simplex"] +[[bench]] +name = "pipeline_dispatch" +harness = false +# Dispatch-overhead A/B for the typed-step runtime; no command drives the +# pipeline yet, so this is the only way to measure a hot-path change. + [[bench]] name = "raw_bam_accessors" harness = false diff --git a/benches/pipeline_dispatch.rs b/benches/pipeline_dispatch.rs new file mode 100644 index 000000000..ebaa48512 --- /dev/null +++ b/benches/pipeline_dispatch.rs @@ -0,0 +1,223 @@ +//! Dispatch-throughput benchmark for the typed-step pipeline runtime. +//! +//! This measures the runtime's *per-dispatch overhead*, not the work steps do: +//! every step here is a near-no-op, so wall time is dominated by the scheduler +//! loop, queue push/pop, and the reorder stage. That is exactly the path any +//! change to the driver's inner loop lands on. +//! +//! It exists because the pipeline had no benchmark at all — no command on this +//! branch drives it yet (the chain builders and command rewiring arrive later), +//! so there was no way to tell whether a hot-path change cost anything. +//! +//! Read the numbers as *relative* only. Absolute throughput here is meaningless +//! as a product metric (real steps decompress BGZF and parse records, which +//! dwarfs dispatch); the point is A/B on the same host. + +use std::collections::VecDeque; +use std::hint::black_box; +use std::io; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use criterion::{Criterion, criterion_group, criterion_main}; +use fgumi_pipeline_core::Unpushed; +use fgumi_pipeline_core::builder::{Pipeline, PipelineConfig}; +use fgumi_pipeline_core::held::HeldSlot; +use fgumi_pipeline_core::item::{HeapSize, Ordered}; +use fgumi_pipeline_core::outputs::OrderedBytesSingle; +use fgumi_pipeline_core::queues::QueueSpec; +use fgumi_pipeline_core::reorder::BranchOrdering; +use fgumi_pipeline_core::step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}; + +/// Per-edge byte budget. Large enough that the queues never bind — this +/// benchmark measures dispatch cost, and backpressure stalls would swamp it +/// with scheduling noise. +const EDGE_LIMIT_BYTES: u64 = 64 * 1024 * 1024; + +/// Items pushed through the chain per iteration. +const N_ITEMS: u64 = 50_000; + +/// A minimal ordered item. `heap_size` is a small constant rather than 0 so the +/// byte-bounded accounting does real work per item, as it would in production. +#[derive(Clone, Copy)] +struct Item { + ordinal: u64, +} + +impl HeapSize for Item { + fn heap_size(&self) -> usize { + 64 + } +} + +impl Ordered for Item { + fn ordinal(&self) -> u64 { + self.ordinal + } +} + +/// Emits `n` items and finishes. `Exclusive` so it owns its cursor. +struct CountingSource { + remaining: VecDeque, + held: HeldSlot>, +} + +impl CountingSource { + fn new(n: u64) -> Self { + Self { remaining: (0..n).map(|ordinal| Item { ordinal }).collect(), held: HeldSlot::new() } + } +} + +impl Step for CountingSource { + type Input = (); + type Outputs = OrderedBytesSingle; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "BenchSource", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: EDGE_LIMIT_BYTES }], + 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::Contention); + } + } + } + let Some(item) = self.remaining.pop_front() else { + return Ok(StepOutcome::Finished); + }; + if let Err(unpushed) = ctx.outputs.push(item) { + self.held.put(unpushed); + } + Ok(StepOutcome::Progress) + } +} + +/// A `Parallel` pass-through. Cloned per worker, so at `threads > 1` this is +/// what puts several workers on the dispatch path at once. +struct PassThrough { + held: HeldSlot>, +} + +impl Step for PassThrough { + type Input = Item; + type Outputs = OrderedBytesSingle; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "BenchPassThrough", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: EDGE_LIMIT_BYTES }], + 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::Contention); + } + } + } + let Some(item) = ctx.input.pop() else { + if ctx.input.is_drained() { + return Ok(StepOutcome::Finished); + } + return Ok(StepOutcome::NoProgress); + }; + // A trivial amount of real work so the compiler cannot fold the step away. + let out = Item { ordinal: black_box(item).ordinal }; + if let Err(unpushed) = ctx.outputs.push(out) { + self.held.put(unpushed); + } + Ok(StepOutcome::Progress) + } + + fn new_worker_copy(&self) -> Self { + Self { held: HeldSlot::new() } + } +} + +/// Terminal sink; counts arrivals into a shared atomic so the run is verifiable. +struct CountingSink { + seen: Arc, +} + +impl Step for CountingSink { + type Input = Item; + type Outputs = (); + + fn profile(&self) -> StepProfile { + StepProfile { + name: "BenchSink", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + match ctx.input.pop() { + Some(item) => { + black_box(item); + self.seen.fetch_add(1, Ordering::Relaxed); + Ok(StepOutcome::Progress) + } + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } +} + +/// Run the chain once, asserting every item arrived. +/// +/// The assert is not decoration: a timing harness reports a wall figure for a +/// run that failed or short-circuited just as happily as for a correct one, and +/// a chain that drops items would look like a speedup. +fn run_chain(threads: usize) { + let seen = Arc::new(AtomicU64::new(0)); + let sink_handle = Arc::clone(&seen); + + let builder = Pipeline::builder(); + builder + .chain(CountingSource::new(N_ITEMS)) + .chain(PassThrough { held: HeldSlot::new() }) + .chain(CountingSink { seen: sink_handle }) + .into_sink_marker(); + let pipeline = builder.build().expect("bench chain builds"); + pipeline.run(PipelineConfig { threads, ..Default::default() }).expect("bench chain runs"); + + assert_eq!( + seen.load(Ordering::Relaxed), + N_ITEMS, + "every item must reach the sink — a short run is an invalid measurement", + ); +} + +fn bench_dispatch(c: &mut Criterion) { + let mut group = c.benchmark_group("pipeline_dispatch"); + group.throughput(criterion::Throughput::Elements(N_ITEMS)); + for threads in [1usize, 4, 8] { + group.bench_function(format!("threads_{threads}"), |b| { + b.iter(|| run_chain(threads)); + }); + } + group.finish(); +} + +criterion_group!(benches, bench_dispatch); +criterion_main!(benches); diff --git a/benches/umi_assigner_threshold.rs b/benches/umi_assigner_threshold.rs index 90211600d..a72e527bc 100644 --- a/benches/umi_assigner_threshold.rs +++ b/benches/umi_assigner_threshold.rs @@ -53,7 +53,7 @@ const SEED: u64 = 0xFEED_BEEF; /// Number of threads the parallel assigner pool spawns. Mirrors what /// `Command::execute` passes (`num_threads`) on a typical workstation. fn bench_threads() -> usize { - num_cpus::get().clamp(2, 8) + std::thread::available_parallelism().map_or(1, std::num::NonZeroUsize::get).clamp(2, 8) } /// Generate `n_templates` UMIs of length `UMI_LEN` with ~`COPIES_PER_UMI` diff --git a/crates/fgumi-bam-io/Cargo.toml b/crates/fgumi-bam-io/Cargo.toml index 889b07266..9d60eb086 100644 --- a/crates/fgumi-bam-io/Cargo.toml +++ b/crates/fgumi-bam-io/Cargo.toml @@ -14,6 +14,7 @@ keywords = ["bioinformatics", "bam", "sequencing", "ngs"] pedantic = { level = "deny", priority = -1 } [dependencies] +ahash = { workspace = true } fgumi-raw-bam = { workspace = true, features = ["noodles"] } fgumi-bgzf = { workspace = true } noodles = { workspace = true, features = ["bam", "fasta", "sam", "bgzf", "core", "vcf"] } diff --git a/crates/fgumi-bam-io/src/grouping.rs b/crates/fgumi-bam-io/src/grouping.rs new file mode 100644 index 000000000..f24946265 --- /dev/null +++ b/crates/fgumi-bam-io/src/grouping.rs @@ -0,0 +1,1058 @@ +//! Shared grouping and decoded-record domain types. +//! +//! Used by both the group command path (`grouper`, `mi_group`, +//! `commands::{group,dedup}` in the main crate) and the typed-step pipeline +//! (`pipeline::steps::group`, `steps::parse::decode`). +//! +//! These are BAM-record domain types — pre-computed grouping keys, the +//! decoded-record representation, the batching-weight and grouper traits. +//! They live here in `fgumi-bam-io`, next to [`crate::MemoryEstimate`] and +//! the [`fgumi_raw_bam`] raw-record helpers they operate on, rather than in +//! the pipeline crate that merely consumes them. + +use std::io; +use std::sync::Arc; + +use noodles::sam::alignment::record::data::field::Tag; + +use crate::library::LibraryIndex; +use fgumi_raw_bam::{RawRecord, RawRecordView}; + +pub use crate::mem_estimate::MemoryEstimate; + +// ============================================================================ +// GroupKey - Pre-computed grouping key for fast comparison in Group step +// ============================================================================ + +/// Pre-computed grouping key for fast comparison in Group step. +/// +/// All fields are integers/hashes for O(1) comparison. This is computed during +/// the parallel Decode step so the serial Group step only does integer comparisons. +/// +/// For paired-end reads, positions are normalized so the lower position comes first. +/// For single-end reads, the mate fields use `UNKNOWN_*` sentinel values. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct GroupKey { + // Position info (normalized: lower position first) + /// Reference sequence index for position 1 (lower). + pub ref_id1: i32, + /// Unclipped 5' position for position 1. + pub pos1: i32, + /// Strand for position 1 (0=forward, 1=reverse). + pub strand1: u8, + /// Reference sequence index for position 2 (higher or mate). + pub ref_id2: i32, + /// Unclipped 5' position for position 2. + pub pos2: i32, + /// Strand for position 2. + pub strand2: u8, + + // Grouping metadata + /// Library index (pre-computed from RG tag via header lookup). + pub library_idx: u16, + /// Hash of cell barcode (0 if none). + pub cell_hash: u64, + + // For name-based grouping within position groups + /// Hash of QNAME for fast name comparison. + pub name_hash: u64, +} + +impl GroupKey { + /// Sentinel value for unknown reference ID (unpaired reads). + pub const UNKNOWN_REF: i32 = i32::MAX; + /// Sentinel value for unknown position (unpaired reads). + pub const UNKNOWN_POS: i32 = i32::MAX; + /// Sentinel value for unknown strand (unpaired reads). + pub const UNKNOWN_STRAND: u8 = u8::MAX; + + /// Create a `GroupKey` for a paired-end read with mate info. + /// + /// Positions are automatically normalized so the lower position comes first. + #[must_use] + #[allow(clippy::too_many_arguments)] + pub fn paired( + ref_id: i32, + pos: i32, + strand: u8, + mate_ref_id: i32, + mate_pos: i32, + mate_strand: u8, + library_idx: u16, + cell_hash: u64, + name_hash: u64, + ) -> Self { + // Normalize: put lower position first (matching ReadInfo behavior) + let (ref_id1, pos1, strand1, ref_id2, pos2, strand2) = + if (ref_id, pos, strand) <= (mate_ref_id, mate_pos, mate_strand) { + (ref_id, pos, strand, mate_ref_id, mate_pos, mate_strand) + } else { + (mate_ref_id, mate_pos, mate_strand, ref_id, pos, strand) + }; + + Self { ref_id1, pos1, strand1, ref_id2, pos2, strand2, library_idx, cell_hash, name_hash } + } + + /// Create a `GroupKey` for a single-end/unpaired read. + #[must_use] + pub fn single( + ref_id: i32, + pos: i32, + strand: u8, + library_idx: u16, + cell_hash: u64, + name_hash: u64, + ) -> Self { + Self { + ref_id1: ref_id, + pos1: pos, + strand1: strand, + ref_id2: Self::UNKNOWN_REF, + pos2: Self::UNKNOWN_POS, + strand2: Self::UNKNOWN_STRAND, + library_idx, + cell_hash, + name_hash, + } + } + + /// Returns the position-only key for grouping by genomic position. + /// + /// This is used by `RecordPositionGrouper` to determine if records belong to + /// the same position group (ignoring name). + #[must_use] + pub fn position_key(&self) -> (i32, i32, u8, i32, i32, u8, u16, u64) { + ( + self.ref_id1, + self.pos1, + self.strand1, + self.ref_id2, + self.pos2, + self.strand2, + self.library_idx, + self.cell_hash, + ) + } + + /// Whether this key carries a mate position, i.e. it was built by + /// [`Self::paired`] rather than [`Self::single`]. + /// + /// `strand2` is set to [`Self::UNKNOWN_STRAND`] by the single-end + /// constructor, so it is the discriminator between the two shapes. + #[must_use] + pub fn has_mate_position(&self) -> bool { + self.strand2 != Self::UNKNOWN_STRAND + } +} + +impl PartialOrd for GroupKey { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for GroupKey { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.position_key() + .cmp(&other.position_key()) + .then_with(|| self.name_hash.cmp(&other.name_hash)) + } +} + +impl Default for GroupKey { + fn default() -> Self { + Self { + ref_id1: Self::UNKNOWN_REF, + pos1: Self::UNKNOWN_POS, + strand1: Self::UNKNOWN_STRAND, + ref_id2: Self::UNKNOWN_REF, + pos2: Self::UNKNOWN_POS, + strand2: Self::UNKNOWN_STRAND, + library_idx: 0, + cell_hash: 0, + name_hash: 0, + } + } +} + +// ============================================================================ +// DecodedRecord - Record with pre-computed grouping key +// ============================================================================ + +/// A decoded BAM record with its pre-computed grouping key. +/// +/// This is the output of the Decode step and input to the Group step. +/// The key is computed during the parallel Decode step so that the +/// serial Group step only needs to do fast integer comparisons. +/// +/// # Cached-UMI invariant +/// +/// `umi_value_offset` / `umi_value_len` cache the position of the UMI value +/// *within* `data`. The invariant is: **the cache is only valid while `data` is +/// unmodified.** Any mutation of `data` that could shift or overwrite the UMI +/// value bytes invalidates the cache. To enforce this cheaply, the sole +/// mutable-bytes accessor ([`Self::raw_bytes_mut`]) resets the cache to +/// [`Self::UMI_OFFSET_UNCACHED`] on hand-out, so a `cached_umi()` read after a +/// mutation falls back to a re-scan instead of slicing stale bytes. +#[derive(Debug)] +pub struct DecodedRecord { + /// Pre-computed grouping key. + pub key: GroupKey, + /// Raw BAM record bytes. + data: RawRecord, + /// Cached record-relative offset of the UMI tag's value bytes (i.e. the + /// first byte after the 2-byte tag header and the 1-byte type byte). The + /// slice `data[umi_value_offset..umi_value_offset + umi_value_len]` yields + /// the UMI bytes without the trailing NUL. + /// + /// Set to [`Self::UMI_OFFSET_UNCACHED`] when no UMI position was cached + /// during decode (UMI tag missing, not Z-typed, or caching disabled). + umi_value_offset: u32, + /// Cached UMI value length in bytes, paired with `umi_value_offset`. + umi_value_len: u16, +} + +impl DecodedRecord { + /// Sentinel value for `umi_value_offset` indicating no cached UMI position. + /// Chosen as `u32::MAX` so it can never collide with a real BAM record + /// offset (BAM records are bounded well under 4 GiB). + pub const UMI_OFFSET_UNCACHED: u32 = u32::MAX; + + /// Create a decoded record from raw bytes, skipping noodles decode. + /// + /// Accepts anything that converts `Into` (e.g. a bare `Vec` or + /// an already-constructed `RawRecord`). + #[must_use] + pub fn from_raw_bytes(raw: impl Into, key: GroupKey) -> Self { + Self { + key, + data: raw.into(), + umi_value_offset: Self::UMI_OFFSET_UNCACHED, + umi_value_len: 0, + } + } + + /// Attach a cached UMI value position to this decoded record. + /// + /// `umi_value_offset` is the record-relative offset of the first UMI value + /// byte (after the tag header), and `umi_value_len` is the value length + /// (excluding any trailing NUL). + pub fn set_cached_umi(&mut self, umi_value_offset: u32, umi_value_len: u16) { + self.umi_value_offset = umi_value_offset; + self.umi_value_len = umi_value_len; + } + + /// Returns the cached UMI bytes (without trailing NUL) if a position was + /// recorded during decode and still falls within the raw record bytes. + /// Returns `None` if no cache is set or the cached position is out of range. + #[must_use] + pub fn cached_umi(&self) -> Option<&[u8]> { + if self.umi_value_offset == Self::UMI_OFFSET_UNCACHED { + return None; + } + let start = self.umi_value_offset as usize; + let end = start.checked_add(self.umi_value_len as usize)?; + let value = self.data.as_ref().get(start..end)?; + // The cache only ever holds a Z-typed UMI value (NUL-free by the BAM + // spec; the trailing NUL is excluded by `set_cached_umi`). If a mutation + // shifted the record bytes such that this offset now lands on a different + // region, the slice would typically contain an interior NUL — a cheap, + // zero-release-cost canary for a stale cache that survived bounds checks. + debug_assert!( + !value.contains(&0), + "stale cached UMI: offset {start} sliced bytes containing an interior NUL \ + (the record was likely mutated without invalidating the UMI cache)" + ); + Some(value) + } + + /// Returns the cached UMI offset ([`Self::UMI_OFFSET_UNCACHED`] when absent) + /// and length. + #[must_use] + pub fn cached_umi_position(&self) -> (u32, u16) { + (self.umi_value_offset, self.umi_value_len) + } + + /// Returns the cached UMI `(offset, len)` position, or `None` when no + /// position was recorded during decode (the [`Self::UMI_OFFSET_UNCACHED`] + /// sentinel). Keeps the sentinel check at the owning layer so callers can + /// branch on a plain `Option` instead of comparing against the sentinel. + #[must_use] + pub fn cached_umi_position_opt(&self) -> Option<(u32, u16)> { + if self.umi_value_offset == Self::UMI_OFFSET_UNCACHED { + None + } else { + Some((self.umi_value_offset, self.umi_value_len)) + } + } + + /// Returns a reference to the raw bytes. + #[must_use] + pub fn raw_bytes(&self) -> &[u8] { + self.data.as_ref() + } + + /// Immutable access to the underlying [`RawRecord`], for read-only + /// consumers (e.g. the FASTQ-encode step) that need the typed accessors + /// (`flags()`, sequence, qualities, tags) rather than the raw byte slice. + #[must_use] + pub fn record(&self) -> &RawRecord { + &self.data + } + + /// Mutable access to the underlying [`RawRecord`]. Used by mid-chain + /// `Parallel` Process steps that need to mutate record bytes (e.g., + /// MQ bumping, tag rewriting) without rebuilding the `DecodedRecord` + /// or invalidating its pre-computed `GroupKey`. Caller must not + /// change the record's identity (qname, library bytes, cell barcode) + /// — those feed `key`, which we don't recompute here. + /// + /// Handing out mutable bytes resets the cached UMI position to + /// [`Self::UMI_OFFSET_UNCACHED`] (see the type-level cached-UMI invariant): + /// a length-changing edit before the UMI tag could shift the value while + /// leaving the old offset in-bounds, so a later `cached_umi()` would slice + /// stale-but-valid bytes. Clearing forces a re-scan instead. This is a + /// single field store and the cache is only repopulated by the decode step. + pub fn raw_bytes_mut(&mut self) -> &mut RawRecord { + self.umi_value_offset = Self::UMI_OFFSET_UNCACHED; + self.umi_value_len = 0; + &mut self.data + } + + /// Takes the [`RawRecord`] out. + #[must_use] + pub fn into_raw_bytes(self) -> RawRecord { + self.data + } +} + +impl MemoryEstimate for DecodedRecord { + fn estimate_heap_size(&self) -> usize { + // RawRecord::capacity() returns the inner Vec capacity. + self.data.capacity() + } +} +// Vec, Vec, Vec, RecordBuf, () — all provided +// by the blanket/foreign impls in fgumi_bam_io::mem_estimate. + +// ============================================================================ +// GroupKeyConfig - Configuration for computing `GroupKey` during Decode +// ============================================================================ + +/// Configuration for computing `GroupKey` during the Decode step. +/// +/// When this is provided to the pipeline, the Decode step will compute +/// full `GroupKey` values for each record. This moves expensive computations +/// (CIGAR parsing, tag extraction) from the serial Group step to the parallel +/// Decode step. +#[derive(Debug, Clone)] +pub struct GroupKeyConfig { + /// Library index for fast RG → library lookup. + pub library_index: Arc, + /// Tag used for cell barcode extraction. None skips cell extraction. + pub cell_tag: Option, + /// When `true`, the Decode step computes only the read-name hash and + /// leaves the rest of the [`GroupKey`] at its default. Use for stages + /// that group by queryname (e.g. `correct`) and read only + /// [`GroupKey::name_hash`] — it skips the CIGAR 5′-position walk and the + /// aux-tag (RG/CB/MC) extraction pass entirely. See [`name_hash_key`]. + pub name_hash_only: bool, + /// UMI tag (raw 2-byte form) whose value position should be cached on each + /// [`DecodedRecord`] during decode. `None` disables caching — downstream + /// code must fall back to scanning aux data. This is orthogonal to + /// `name_hash_only`: the UMI cache scan is gated solely on `umi_tag` being + /// set (in practice only the Group stage sets it). + pub umi_tag: Option<[u8; 2]>, +} + +impl GroupKeyConfig { + /// Create a new `GroupKeyConfig` that computes the full position/cell key. + #[must_use] + pub fn new(library_index: LibraryIndex, cell_tag: Tag) -> Self { + Self { + library_index: Arc::new(library_index), + cell_tag: Some(cell_tag), + name_hash_only: false, + umi_tag: None, + } + } + + /// Create a `GroupKeyConfig` without cell barcode extraction. + #[must_use] + pub fn new_raw_no_cell(library_index: LibraryIndex) -> Self { + Self { + library_index: Arc::new(library_index), + cell_tag: None, + name_hash_only: false, + umi_tag: None, + } + } + + /// Create a `GroupKeyConfig` that computes only the read-name hash. + /// + /// For queryname-grouping stages (e.g. `correct`) whose grouper reads only + /// [`GroupKey::name_hash`]; skips the CIGAR position walk and the aux-tag + /// extraction pass. `library_index` is retained (unused for the key) so + /// the config shape is uniform across the pipeline. + #[must_use] + pub fn name_hash_only(library_index: LibraryIndex) -> Self { + Self { + library_index: Arc::new(library_index), + cell_tag: None, + name_hash_only: true, + umi_tag: None, + } + } + + /// Enable UMI position caching for the given raw 2-byte UMI tag (e.g. `*b"RX"`). + /// + /// The Decode step will record the UMI value's record-relative position on + /// each [`DecodedRecord`] so downstream UMI lookups can slice it directly + /// via [`DecodedRecord::cached_umi`] without re-scanning aux data. + #[must_use] + pub fn with_umi_tag(mut self, umi_tag: [u8; 2]) -> Self { + self.umi_tag = Some(umi_tag); + self + } +} + +impl Default for GroupKeyConfig { + fn default() -> Self { + Self { + library_index: Arc::new(LibraryIndex::default()), + cell_tag: Some(Tag::from([b'C', b'B'])), // Default cell barcode tag (CB) + name_hash_only: false, + umi_tag: None, + } + } +} + +/// Hash a record's read name, mapping an empty name to `hash_name(None)` so the +/// raw path matches the noodles path (which sees `None` for an empty name). +/// +/// Shared by [`name_hash_key`] and [`compute_group_key_from_raw`]: both must +/// produce the same hash for the same record, and a second copy of this block +/// could drift out of parity silently. +fn raw_name_hash(raw: &[u8]) -> u64 { + let name = fgumi_raw_bam::read_name(raw); + if name.is_empty() { + LibraryIndex::hash_name(None) + } else { + LibraryIndex::hash_name(Some(name)) + } +} + +/// Compute a [`GroupKey`] containing only the read-name hash, leaving all +/// position/strand/library/cell fields at their default. +/// +/// Reproduces exactly the `name_hash` that [`compute_group_key_from_raw`] +/// computes (empty name → `hash_name(None)`), so a queryname grouper sees an +/// identical hash. Skips the CIGAR 5′-position walk and the RG/CB/MC aux-tag +/// extraction pass. +/// +/// # Panics +/// +/// Panics if `raw` is not a validated BAM record payload (same contract as +/// [`compute_group_key_from_raw`]). +#[must_use] +pub fn name_hash_key(raw: &[u8]) -> GroupKey { + GroupKey { name_hash: raw_name_hash(raw), ..GroupKey::default() } +} + +/// Compute a `GroupKey` directly from raw BAM bytes, matching `compute_group_key()` exactly. +/// +/// Uses 1-based coordinate helpers to produce identical keys to the noodles path. +/// +/// # Panics +/// +/// Panics if `raw` is not a validated BAM record payload (as produced by the BAM +/// reader / raw-record pipeline). Callers must not pass arbitrary external bytes; +/// raw-field accessors will panic on malformed or truncated input. +#[must_use] +pub fn compute_group_key_from_raw( + raw: &[u8], + library_index: &LibraryIndex, + cell_tag: Option, +) -> GroupKey { + // Extract name hash (match noodles path: empty name → None → hash 0) + let name_hash = raw_name_hash(raw); + + // Check secondary/supplementary + let flg = RawRecordView::new(raw).flags(); + let is_secondary = (flg & fgumi_raw_bam::flags::SECONDARY) != 0; + let is_supplementary = (flg & fgumi_raw_bam::flags::SUPPLEMENTARY) != 0; + if is_secondary || is_supplementary { + // A secondary/supplementary read cannot compute its own template + // coordinate (it lacks its own primary's position). When `fgumi zipper` + // has stamped the exact coordinate into `tc`, key on it so the read + // groups into the same position group as its primary — instead of an + // UNKNOWN position that relies on the read sorting adjacent to its + // primary. Library/cell are extracted the same way as primaries so the + // position key matches (they share their template's RG/CB). + let aux_data = fgumi_raw_bam::aux_data_slice(raw); + if let Some([tid1, pos1, neg1, tid2, pos2, neg2]) = + fgumi_raw_bam::read_tc_template_coordinate(aux_data) + { + let cell_tag_bytes = cell_tag.map_or([0u8; 2], |t| [t.as_ref()[0], t.as_ref()[1]]); + let aux_tags = fgumi_raw_bam::extract_aux_string_tags(aux_data, cell_tag_bytes, None); + let library_idx = + aux_tags.rg.map_or(0, |rg| library_index.get(LibraryIndex::hash_rg(rg))); + let cell_hash = aux_tags.cell.map_or(0, |cb| LibraryIndex::hash_cell_barcode(Some(cb))); + return GroupKey::paired( + tid1, + pos1, + u8::from(neg1 != 0), + tid2, + pos2, + u8::from(neg2 != 0), + library_idx, + cell_hash, + name_hash, + ); + } + return GroupKey { name_hash, ..GroupKey::default() }; + } + + // Own position (1-based, matching noodles) — zero-allocation CIGAR iteration + let reverse = (flg & fgumi_raw_bam::flags::REVERSE) != 0; + let own_pos = fgumi_raw_bam::unclipped_5prime_from_raw_bam(raw); + + // A mapped record with an empty or truncated CIGAR has no computable + // unclipped 5' position, so `unclipped_5prime_from_raw_bam` returns the + // `i32::MAX` sentinel. Fall back to the name-only key rather than letting the + // sentinel reach a position slot — otherwise distinct templates sharing + // ref/strand/library/cell would collide on `i32::MAX` (`position_key` + // excludes `name_hash`). Matches the secondary/supplementary fallback above. + if own_pos == i32::MAX { + return GroupKey { name_hash, ..GroupKey::default() }; + } + + let own_ref_id = fgumi_raw_bam::ref_id(raw); + let strand = u8::from(reverse); + + // Single-pass aux tag extraction (RG, cell barcode, MC) + let aux_data = fgumi_raw_bam::aux_data_slice(raw); + let cell_tag_bytes = cell_tag.map_or([0u8; 2], |t| [t.as_ref()[0], t.as_ref()[1]]); + let aux_tags = fgumi_raw_bam::extract_aux_string_tags(aux_data, cell_tag_bytes, None); + + let library_idx = if let Some(rg) = aux_tags.rg { + let rg_hash = LibraryIndex::hash_rg(rg); + library_index.get(rg_hash) + } else { + 0 + }; + + let cell_hash = + if let Some(cb) = aux_tags.cell { LibraryIndex::hash_cell_barcode(Some(cb)) } else { 0 }; + + // Check if paired + let is_paired = (flg & fgumi_raw_bam::flags::PAIRED) != 0; + if !is_paired { + return GroupKey::single(own_ref_id, own_pos, strand, library_idx, cell_hash, name_hash); + } + + // Mate info — guard against MATE_UNMAPPED (matching noodles path) + let mate_unmapped = (flg & fgumi_raw_bam::flags::MATE_UNMAPPED) != 0; + let mate_reverse = (flg & fgumi_raw_bam::flags::MATE_REVERSE) != 0; + let mate_strand = u8::from(mate_reverse); + let raw_mate_ref_id = fgumi_raw_bam::mate_ref_id(raw); + let raw_mate_pos = fgumi_raw_bam::mate_pos(raw); + + // Get mate unclipped 5' position via MC tag (skip if mate is unmapped) + let mate_pos_result = if mate_unmapped { + None + } else { + aux_tags + .mc + .map(|mc| fgumi_raw_bam::mate_unclipped_5prime_1based(raw_mate_pos, mate_reverse, mc)) + }; + + match mate_pos_result { + Some(mp) => GroupKey::paired( + own_ref_id, + own_pos, + strand, + raw_mate_ref_id, + mp, + mate_strand, + library_idx, + cell_hash, + name_hash, + ), + None => { + // No MC tag — fall back to single-end behavior + GroupKey::single(own_ref_id, own_pos, strand, library_idx, cell_hash, name_hash) + } + } +} + +/// Groups a stream of in-order [`DecodedRecord`]s into completed groups. +/// +/// Implementors maintain partial groups across `add_records` calls and emit +/// completed ones; `finish` flushes any trailing partial group at EOF. Used +/// by the pipeline's Group step and the standalone grouping commands. +pub trait Grouper: Send { + /// The type of group produced by this grouper. + type Group: Send; + + /// Add decoded records to the grouper. + /// + /// Records are guaranteed to be in order (from template-coordinate sorted BAM). + /// The grouper maintains partial groups waiting for more records. + /// + /// Each `DecodedRecord` contains the record plus a pre-computed `GroupKey` + /// for fast comparison (position, name hash, library, etc.). + /// + /// Returns completed groups (may be empty if more records are needed). + /// + /// # Errors + /// + /// Returns an I/O error if grouping logic encounters invalid data. + fn add_records(&mut self, records: Vec) -> io::Result>; + + /// Signal that no more input will arrive (EOF). + /// + /// Returns any remaining partial group. + /// + /// # Errors + /// + /// Returns an I/O error if finalizing the grouper fails. + fn finish(&mut self) -> io::Result>; + + /// Returns true if the grouper has a partial group. + fn has_pending(&self) -> bool; +} + +#[cfg(test)] +mod tests { + use super::*; + use fgumi_raw_bam::SamBuilder; + use fgumi_raw_bam::SamTag; + use fgumi_raw_bam::flags; + + // ======================================================================== + // compute_group_key_from_raw — primary (fully-populated) path + // ======================================================================== + + /// Read group `RG1` resolves to library `libA`; `RG2` to `libB`. + fn library_index_with_two_groups() -> LibraryIndex { + use noodles::sam::header::record::value::Map; + use noodles::sam::header::record::value::map::ReadGroup; + use noodles::sam::header::record::value::map::read_group::tag as rg_tag; + + let mut header = noodles::sam::Header::builder(); + for (id, library) in [("RG1", "libA"), ("RG2", "libB")] { + let rg = Map::::builder() + .insert(rg_tag::LIBRARY, String::from(library)) + .build() + .expect("read group builds"); + header = header.add_read_group(bstr::BString::from(id), rg); + } + LibraryIndex::from_header(&header.build()) + } + + /// Build one mapped mate of a pair carrying `RG`, `CB` and `MC`. + fn paired_record_with_tags( + pos: i32, + mate_pos: i32, + reverse: bool, + mate_reverse: bool, + ) -> fgumi_raw_bam::RawRecord { + let mut flag = flags::PAIRED; + if reverse { + flag |= flags::REVERSE; + } + if mate_reverse { + flag |= flags::MATE_REVERSE; + } + let mut b = SamBuilder::new(); + b.ref_id(0) + .pos(pos) + .flags(flag) + .mate_ref_id(0) + .mate_pos(mate_pos) + .read_name(b"pair_full") + .cigar_ops(&[cigar_m(50)]) + .sequence(b"ACGT") + .qualities(&[30, 30, 30, 30]) + .add_string_tag(SamTag::RG, b"RG1") + .add_string_tag(SamTag::CB, b"ACGTACGTA") + .add_string_tag(SamTag::MC, b"50M"); + b.build() + } + + /// The fully-populated branch: paired, mate mapped, `MC` present, `RG` + /// resolving through a non-default `LibraryIndex`, and a cell barcode read via + /// `cell_tag`. Every other test in this module lands on a fallback (secondary, + /// supplementary, or paired-without-`MC`) with `LibraryIndex::default()` and + /// `cell_tag: None`, so `library_idx`, `cell_hash` and the mate slot are + /// otherwise never exercised — and all three participate in key equality. + #[test] + fn paired_with_mc_populates_mate_library_and_cell_fields() { + let lib = library_index_with_two_groups(); + let cb = noodles::sam::alignment::record::data::field::Tag::CELL_BARCODE_ID; + + let rec = paired_record_with_tags(1000, 1200, false, true); + let key = compute_group_key_from_raw(rec.as_ref(), &lib, Some(cb)); + + // Mate slot is populated, not left at the single-end sentinels. + assert_ne!(key.ref_id2, GroupKey::UNKNOWN_REF, "MC must populate the mate slot"); + assert_ne!(key.strand2, GroupKey::UNKNOWN_STRAND); + + // Pin the mate coordinate: it must be the MC-derived unclipped 5' position, + // not the raw `mate_pos`. `!= UNKNOWN_POS` alone would accept a wrong + // CIGAR walk. + let expected_mate_pos = fgumi_raw_bam::mate_unclipped_5prime_1based( + fgumi_raw_bam::mate_pos(rec.as_ref()), + true, + b"50M", + ); + assert_eq!(key.pos2, expected_mate_pos, "mate slot must hold the MC-derived position"); + + // RG resolved through the index — not the unknown bucket. + assert_eq!(key.library_idx, lib.get(LibraryIndex::hash_rg(b"RG1"))); + assert_ne!(key.library_idx, 0, "RG1 must resolve to a real library"); + + // Cell barcode hashed because `cell_tag` named CB. + assert_eq!(key.cell_hash, LibraryIndex::hash_cell_barcode(Some(b"ACGTACGTA"))); + assert_ne!(key.cell_hash, 0); + + assert_eq!(key.name_hash, LibraryIndex::hash_name(Some(b"pair_full"))); + + // Own slot holds this record's own unclipped 5' position. + assert_eq!(key.pos1, fgumi_raw_bam::unclipped_5prime_from_raw_bam(rec.as_ref())); + assert_eq!(key.ref_id1, fgumi_raw_bam::ref_id(rec.as_ref())); + } + + /// Both mates of the same template must normalize to one key — this is what + /// makes them group together. The mate sees own/mate swapped and the strands + /// exchanged. + #[test] + fn paired_with_mc_normalizes_both_mates_to_the_same_key() { + let lib = library_index_with_two_groups(); + let cb = noodles::sam::alignment::record::data::field::Tag::CELL_BARCODE_ID; + + let forward = paired_record_with_tags(1000, 1200, false, true); + let reverse = paired_record_with_tags(1200, 1000, true, false); + + let key_forward = compute_group_key_from_raw(forward.as_ref(), &lib, Some(cb)); + let key_reverse = compute_group_key_from_raw(reverse.as_ref(), &lib, Some(cb)); + + assert_eq!(key_forward, key_reverse, "both mates of a template must share one key"); + } + + /// A secondary/supplementary read cannot derive its own template coordinate, + /// so `fgumi zipper` stamps the primary's into the `tc` aux array. Keying on + /// it puts the read in the SAME position group as its primary; without it the + /// read falls back to a name-only key and relies on sorting adjacency. + #[test] + fn secondary_with_tc_tag_keys_on_the_stamped_template_coordinate() { + let lib = library_index_with_two_groups(); + + let mut b = SamBuilder::new(); + b.ref_id(0) + .pos(5000) + .flags(flags::PAIRED | flags::SECONDARY) + .read_name(b"sec_with_tc") + .cigar_ops(&[cigar_m(50)]) + .sequence(b"ACGT") + .qualities(&[30, 30, 30, 30]) + .add_string_tag(SamTag::RG, b"RG1") + // tc = [tid1, pos1, neg1, tid2, pos2, neg2] — the primary's coordinate. + .add_array_i32(SamTag::TC, &[0, 1001, 0, 0, 1249, 1]); + let rec = b.build(); + + let key = compute_group_key_from_raw(rec.as_ref(), &lib, None); + + // Both slots come from `tc`, NOT from this record's own pos (5000). + assert_eq!((key.ref_id1, key.pos1, key.strand1), (0, 1001, 0)); + assert_eq!((key.ref_id2, key.pos2, key.strand2), (0, 1249, 1)); + assert_ne!(key.pos1, 5001, "the record's own position must not be used"); + + // Library still resolves, so the key matches its primary's. + assert_eq!(key.library_idx, lib.get(LibraryIndex::hash_rg(b"RG1"))); + assert_eq!(key.name_hash, LibraryIndex::hash_name(Some(b"sec_with_tc"))); + } + + /// Without a `tc` tag the secondary path still falls back to the name-only + /// key — including for a mapped record with an EMPTY cigar, where + /// `unclipped_5prime_from_raw_bam` returns `i32::MAX`. The fallback must not + /// leak that sentinel into a position slot. + #[test] + fn secondary_without_tc_falls_back_to_name_only_even_with_an_empty_cigar() { + let lib = library_index_with_two_groups(); + + let mut b = SamBuilder::new(); + b.ref_id(0) + .pos(5000) + .flags(flags::SUPPLEMENTARY) + .read_name(b"sup_no_tc") + .cigar_ops(&[]) + .sequence(b"ACGT") + .qualities(&[30, 30, 30, 30]); + let rec = b.build(); + + let key = compute_group_key_from_raw(rec.as_ref(), &lib, None); + + assert_eq!(key, GroupKey { name_hash: key.name_hash, ..GroupKey::default() }); + assert_eq!(key.name_hash, LibraryIndex::hash_name(Some(b"sup_no_tc"))); + assert_eq!(key.pos1, GroupKey::default().pos1, "no i32::MAX sentinel may leak in"); + } + + /// The PRIMARY path has the same sentinel hazard: a mapped, unpaired primary + /// with an empty CIGAR has no computable unclipped 5' position, so it must + /// fall back to the name-only key too. Without the guard the `i32::MAX` + /// sentinel would reach `pos1` and merge distinct templates sharing + /// ref/strand/library/cell (`position_key` excludes `name_hash`). + #[test] + fn primary_with_empty_cigar_falls_back_to_name_only() { + let lib = library_index_with_two_groups(); + + // Default flags (0) => mapped, unpaired, primary. + let mut b = SamBuilder::new(); + b.ref_id(0) + .pos(5000) + .read_name(b"primary_no_cigar") + .cigar_ops(&[]) + .sequence(b"ACGT") + .qualities(&[30, 30, 30, 30]); + let rec = b.build(); + + let key = compute_group_key_from_raw(rec.as_ref(), &lib, None); + + assert_eq!(key, GroupKey { name_hash: key.name_hash, ..GroupKey::default() }); + assert_eq!(key.name_hash, LibraryIndex::hash_name(Some(b"primary_no_cigar"))); + assert_eq!(key.pos1, GroupKey::default().pos1, "no i32::MAX sentinel may leak in"); + } + + /// `library_idx` is part of key equality: the same position read under a + /// different read group must NOT group with it. + #[test] + fn differing_library_splits_the_key() { + let lib = library_index_with_two_groups(); + + let mut b = SamBuilder::new(); + b.ref_id(0) + .pos(1000) + .flags(flags::PAIRED) + .mate_ref_id(0) + .mate_pos(1200) + .read_name(b"pair_full") + .cigar_ops(&[cigar_m(50)]) + .sequence(b"ACGT") + .qualities(&[30, 30, 30, 30]) + .add_string_tag(SamTag::RG, b"RG2") + .add_string_tag(SamTag::MC, b"50M"); + let other_library = b.build(); + + let key_a = compute_group_key_from_raw( + paired_record_with_tags(1000, 1200, false, true).as_ref(), + &lib, + None, + ); + let key_b = compute_group_key_from_raw(other_library.as_ref(), &lib, None); + + assert_ne!(key_a.library_idx, key_b.library_idx); + assert_ne!(key_a, key_b, "a different library must not group with the first"); + } + + /// CIGAR op `(len << 4) | op_code`; op 0 = `M`. + fn cigar_m(len: u32) -> u32 { + len << 4 + } + + // ======================================================================== + // GroupKey::paired normalization + // ======================================================================== + + #[test] + fn paired_normalizes_swapped_own_and_mate_to_equal_keys() { + // Two reads of the same template: one sees (own=A, mate=B), the other + // sees (own=B, mate=A). After normalization both must yield the same + // GroupKey, so they group together. + let a = GroupKey::paired(0, 100, 0, 0, 200, 1, 7, 99, 42); + let b = GroupKey::paired(0, 200, 1, 0, 100, 0, 7, 99, 42); + assert_eq!(a, b, "swapped own/mate positions must normalize to equal keys"); + + // The lower (ref_id, pos, strand) tuple is placed in slot 1. + assert_eq!((a.ref_id1, a.pos1, a.strand1), (0, 100, 0)); + assert_eq!((a.ref_id2, a.pos2, a.strand2), (0, 200, 1)); + } + + #[test] + fn paired_keys_order_by_normalized_position_then_name() { + // Position 1 < position 2 orders the keys; equal positions fall back to + // name_hash ordering. + let lower = GroupKey::paired(0, 100, 0, 0, 200, 0, 0, 0, 1); + let higher = GroupKey::paired(0, 150, 0, 0, 200, 0, 0, 0, 1); + assert!(lower < higher, "lower position-1 must sort first"); + + let same_pos_low_name = GroupKey::paired(0, 100, 0, 0, 200, 0, 0, 0, 1); + let same_pos_high_name = GroupKey::paired(0, 100, 0, 0, 200, 0, 0, 0, 2); + assert!(same_pos_low_name < same_pos_high_name, "equal positions must order by name_hash",); + } + + // ======================================================================== + // compute_group_key_from_raw: secondary / supplementary fallback + // ======================================================================== + + #[test] + fn secondary_and_supplementary_records_yield_name_hash_only_key() { + let lib = LibraryIndex::default(); + let expected_name_hash = LibraryIndex::hash_name(Some(b"rec1")); + + for flag in [flags::SECONDARY, flags::SUPPLEMENTARY] { + // Give the record a real mapped position so we can prove the + // position fields are *not* derived for secondary/supplementary. + let mut b = SamBuilder::new(); + b.ref_id(0) + .pos(500) + .flags(flag) + .read_name(b"rec1") + .cigar_ops(&[cigar_m(50)]) + .sequence(b"ACGT") + .qualities(&[30, 30, 30, 30]); + let rec = b.build(); + + let key = compute_group_key_from_raw(rec.as_ref(), &lib, None); + + // Only the name hash is set; every position field stays at default. + assert_eq!(key, GroupKey { name_hash: expected_name_hash, ..GroupKey::default() }); + assert_eq!(key.ref_id1, GroupKey::UNKNOWN_REF); + assert_eq!(key.pos1, GroupKey::UNKNOWN_POS); + assert_eq!(key.strand1, GroupKey::UNKNOWN_STRAND); + } + } + + // ======================================================================== + // compute_group_key_from_raw: paired record missing MC falls back to single + // ======================================================================== + + #[test] + fn paired_without_mc_tag_falls_back_to_single_semantics() { + let lib = LibraryIndex::default(); + + // Paired + mate mapped, but no MC tag: cannot compute the mate's + // unclipped 5' position, so the key must use single-end semantics + // (mate fields left at the UNKNOWN sentinels). + let mut b = SamBuilder::new(); + b.ref_id(0) + .pos(1000) + .flags(flags::PAIRED) + .mate_ref_id(0) + .mate_pos(1200) + .read_name(b"pair_no_mc") + .cigar_ops(&[cigar_m(50)]) + .sequence(b"ACGT") + .qualities(&[30, 30, 30, 30]); + let rec = b.build(); + + let key = compute_group_key_from_raw(rec.as_ref(), &lib, None); + + // Mate fields fall back to the single-end sentinels. + assert_eq!(key.ref_id2, GroupKey::UNKNOWN_REF); + assert_eq!(key.pos2, GroupKey::UNKNOWN_POS); + assert_eq!(key.strand2, GroupKey::UNKNOWN_STRAND); + + // The own position is still populated; the key matches the single-end + // key built from the same fields. + let expected = GroupKey::single( + key.ref_id1, + key.pos1, + key.strand1, + key.library_idx, + key.cell_hash, + LibraryIndex::hash_name(Some(b"pair_no_mc")), + ); + assert_eq!(key, expected); + } + + // ======================================================================== + // name_hash_key parity with compute_group_key_from_raw + // ======================================================================== + + #[test] + fn name_hash_key_matches_compute_group_key_name_hash() { + let lib = LibraryIndex::default(); + + // A representative spread of records: mapped single-end, secondary, + // paired-without-MC, and an empty-named record (exercises the + // hash_name(None) branch shared by both functions). + let mut single = SamBuilder::new(); + single + .ref_id(0) + .pos(10) + .read_name(b"alpha") + .cigar_ops(&[cigar_m(8)]) + .sequence(b"ACGTACGT") + .qualities(&[30; 8]); + + let mut secondary = SamBuilder::new(); + secondary.ref_id(0).pos(20).flags(flags::SECONDARY).read_name(b"beta"); + + let mut paired = SamBuilder::new(); + paired + .ref_id(0) + .pos(30) + .flags(flags::PAIRED) + .mate_ref_id(0) + .mate_pos(60) + .read_name(b"gamma"); + + let mut empty_name = SamBuilder::new(); + empty_name.ref_id(0).pos(40).read_name(b""); + + for mut builder in [single, secondary, paired, empty_name] { + let rec = builder.build(); + let raw = rec.as_ref(); + assert_eq!( + name_hash_key(raw).name_hash, + compute_group_key_from_raw(raw, &lib, None).name_hash, + "name_hash parity mismatch", + ); + } + } + + // ======================================================================== + // DecodedRecord cached-UMI invalidation (S6-002) + // ======================================================================== + + /// Build a `DecodedRecord` carrying an RX UMI tag with its value position + /// cached, exactly as the Decode step's `cache_umi_position` would. + fn decoded_with_cached_umi(umi: &[u8]) -> DecodedRecord { + use fgumi_raw_bam::SamTag; + let mut b = SamBuilder::new(); + b.read_name(b"read1").sequence(b"ACGT").qualities(&[30; 4]); + b.add_string_tag(SamTag::RX, umi); + let rec = b.build(); + let bytes = rec.as_ref(); + let aux_offset = + fgumi_raw_bam::aux_data_offset_from_record(bytes).expect("aux offset present"); + let aux = &bytes[aux_offset..]; + let (off_in_aux, len) = + fgumi_raw_bam::find_string_tag_position(aux, *SamTag::RX).expect("RX tag present"); + let offset = u32::try_from(aux_offset).expect("aux offset fits u32") + off_in_aux; + let mut d = DecodedRecord::from_raw_bytes(rec, GroupKey::default()); + d.set_cached_umi(offset, len); + d + } + + #[test] + fn cached_umi_returns_value_then_raw_bytes_mut_invalidates() { + // The cache resolves to the UMI value up front; handing out mutable bytes + // must reset the cache to the sentinel so a later read re-scans rather + // than slicing stale-but-in-range bytes (see the cached-UMI invariant). + let mut decoded = decoded_with_cached_umi(b"ACGTACGT"); + assert_eq!(decoded.cached_umi(), Some(b"ACGTACGT".as_ref()), "cache populated up front"); + assert_ne!(decoded.cached_umi_position().0, DecodedRecord::UMI_OFFSET_UNCACHED); + + let _ = decoded.raw_bytes_mut(); + + assert_eq!( + decoded.cached_umi_position().0, + DecodedRecord::UMI_OFFSET_UNCACHED, + "raw_bytes_mut resets the cache to the uncached sentinel", + ); + assert!(decoded.cached_umi().is_none(), "cached_umi returns None after mutation"); + } +} diff --git a/crates/fgumi-bam-io/src/lib.rs b/crates/fgumi-bam-io/src/lib.rs index 089632560..134c47c80 100644 --- a/crates/fgumi-bam-io/src/lib.rs +++ b/crates/fgumi-bam-io/src/lib.rs @@ -12,7 +12,9 @@ #![deny(unsafe_code)] pub mod format; +pub mod grouping; pub mod header; +pub mod library; pub mod mem_estimate; pub mod os_hints; pub mod paths; @@ -26,6 +28,10 @@ pub mod writer; pub(crate) mod vendored; pub use format::{FORMAT_PREFIX_LEN, InputFormat, classify_input}; +pub use grouping::{ + DecodedRecord, GroupKey, GroupKeyConfig, Grouper, compute_group_key_from_raw, name_hash_key, +}; +pub use library::{LibraryIndex, LibraryLookup, build_library_lookup, unknown_library}; pub use mem_estimate::MemoryEstimate; pub use paths::{is_stdin_path, is_stdout_path}; pub use progress::ProgressTracker; diff --git a/crates/fgumi-bam-io/src/library.rs b/crates/fgumi-bam-io/src/library.rs new file mode 100644 index 000000000..db0d2e19b --- /dev/null +++ b/crates/fgumi-bam-io/src/library.rs @@ -0,0 +1,295 @@ +//! Library lookup tables built from SAM `@RG` headers. +//! +//! [`LibraryLookup`] maps read-group IDs to library-name strings; +//! [`LibraryIndex`] is the hash-based hot-path variant returning numeric +//! library indices (used by `grouping::compute_group_key_from_raw`). +//! +//! Relocated here from the main crate's `read_info` module so the grouping +//! domain types (which depend on `LibraryIndex`) can live in this crate +//! alongside [`fgumi_raw_bam::RawRecord`]-adjacent helpers and `MemoryEstimate`. + +use std::collections::HashMap; +use std::sync::Arc; + +use bstr::ByteSlice; +use noodles::sam::header::Header; +use noodles::sam::header::record::value::map::read_group::tag as rg_tag; + +/// A lookup table mapping read group IDs to library names. +/// +/// This is built from the SAM header's `@RG` lines and used to resolve the library +/// name (`LB` field) from a record's `RG` tag. This matches fgbio's behavior where +/// grouping uses the library name, not the read group ID. +/// +/// Uses `Arc` for library names to avoid cloning strings for every read. +/// +/// # Note: `LibraryLookup` vs `LibraryIndex` +/// +/// Both `LibraryLookup` and [`LibraryIndex`] exist for different use cases: +/// - `LibraryLookup`: String-based lookup returning library names. Used where +/// the actual library-name string is needed (e.g. the main crate's +/// `ReadInfo::from`). +/// - [`LibraryIndex`]: Hash-based lookup returning numeric indices. Used by +/// [`compute_group_key_from_raw`](crate::compute_group_key_from_raw) in the +/// hot path where only equality comparison matters, avoiding string +/// allocations. +pub type LibraryLookup = Arc>>; + +/// Shared "unknown" library string to avoid repeated allocations. +static UNKNOWN_LIBRARY: std::sync::LazyLock> = + std::sync::LazyLock::new(|| Arc::from("unknown")); + +/// Returns the shared "unknown" library name (`Arc` of `"unknown"`), +/// used as the fallback when a read group has no `LB` field. +#[must_use] +pub fn unknown_library() -> Arc { + Arc::clone(&UNKNOWN_LIBRARY) +} + +/// Builds a library lookup table from a SAM header. +/// +/// Iterates through all `@RG` lines in the header and creates a mapping from +/// read group ID to library name. If a read group has no `LB` field, it maps +/// to "unknown" (matching fgbio's behavior). +/// +/// # Arguments +/// +/// * `header` - The SAM header containing `@RG` lines +/// +/// # Returns +/// +/// An `Arc` mapping read group IDs to library names +#[must_use] +pub fn build_library_lookup(header: &Header) -> LibraryLookup { + let mut lookup = HashMap::new(); + + for (id, rg) in header.read_groups() { + // Get the LB field from the read group's other_fields + let library: Arc = rg + .other_fields() + .get(&rg_tag::LIBRARY) + .map_or_else(|| Arc::clone(&UNKNOWN_LIBRARY), |s| Arc::from(s.to_string())); + lookup.insert(id.to_string(), library); + } + + Arc::new(lookup) +} + +// ============================================================================ +// LibraryIndex - Fast RG to library index mapping for GroupKey computation +// ============================================================================ + +/// Fast lookup from `RG` tag value to library index. +/// +/// This provides `O(1)` library lookup during Decode using string hashing, +/// replacing the `O(n)` string comparison in the original `LibraryLookup`. +#[derive(Debug, Clone)] +pub struct LibraryIndex { + /// Map from `RG` string hash to library index. + lookup: ahash::AHashMap, + /// Library names for each index (for output/debugging). + names: Vec>, + /// Unknown library index (always 0). + unknown_idx: u16, +} + +impl LibraryIndex { + /// Build a library index from a SAM header. + /// + /// Each unique library name gets a sequential index starting from 0. + /// Index 0 is reserved for "unknown" library. + /// + /// # Panics + /// + /// Panics if the header contains more than 65,535 distinct libraries. + #[must_use] + pub fn from_header(header: &Header) -> Self { + let mut lookup = ahash::AHashMap::new(); + let mut names = vec![Arc::clone(&UNKNOWN_LIBRARY)]; // Index 0 = unknown + let mut library_to_idx: ahash::AHashMap, u16> = ahash::AHashMap::new(); + library_to_idx.insert(Arc::clone(&UNKNOWN_LIBRARY), 0); + + for (id, rg) in header.read_groups() { + // Get library name from LB field + let library: Arc = rg + .other_fields() + .get(&rg_tag::LIBRARY) + .map_or_else(|| Arc::clone(&UNKNOWN_LIBRARY), |s| Arc::from(s.to_string())); + + // Get or create library index + let lib_idx = *library_to_idx.entry(library.clone()).or_insert_with(|| { + let idx: u16 = + names.len().try_into().expect("too many distinct libraries for u16 index"); + names.push(library); + idx + }); + + // Hash the RG string and map to library index + let rg_hash = Self::hash_rg(id.as_bytes()); + lookup.insert(rg_hash, lib_idx); + } + + Self { lookup, names, unknown_idx: 0 } + } + + /// Get the library index for a read group hash. + /// + /// Returns 0 (unknown) if the `RG` hash is not found. + #[must_use] + pub fn get(&self, rg_hash: u64) -> u16 { + *self.lookup.get(&rg_hash).unwrap_or(&self.unknown_idx) + } + + /// Get the library name for an index. + #[must_use] + pub fn library_name(&self, idx: u16) -> &Arc { + self.names.get(idx as usize).unwrap_or(&self.names[0]) + } + + /// Hash a byte slice using `AHash`. Returns 0 for `None`. + /// + /// This is the single hashing implementation used by all `hash_*` methods. + #[must_use] + pub fn hash_bytes(bytes: Option<&[u8]>) -> u64 { + use ahash::AHasher; + use std::hash::{Hash, Hasher}; + match bytes { + Some(b) => { + let mut hasher = AHasher::default(); + b.hash(&mut hasher); + hasher.finish() + } + None => 0, + } + } + + /// Hash an `RG` tag value for lookup. + #[must_use] + pub fn hash_rg(rg_bytes: &[u8]) -> u64 { + Self::hash_bytes(Some(rg_bytes)) + } + + /// Hash a cell barcode for `GroupKey`. + #[must_use] + pub fn hash_cell_barcode(cell_bytes: Option<&[u8]>) -> u64 { + Self::hash_bytes(cell_bytes) + } + + /// Hash a read name for `GroupKey`. + #[must_use] + pub fn hash_name(name_bytes: Option<&[u8]>) -> u64 { + Self::hash_bytes(name_bytes) + } +} + +impl Default for LibraryIndex { + fn default() -> Self { + Self { + lookup: ahash::AHashMap::new(), + names: vec![Arc::clone(&UNKNOWN_LIBRARY)], + unknown_idx: 0, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use noodles::sam::header::record::value::Map; + use noodles::sam::header::record::value::map::ReadGroup; + use noodles::sam::header::record::value::map::read_group::tag as rg_tag; + + /// Header with three read groups: two sharing a library, one with no `LB`. + fn header_with_read_groups() -> Header { + let mut header = Header::builder(); + // RG1 and RG3 deliberately share a library so the dedup path is exercised. + for (id, library) in [("RG1", "libA"), ("RG2", "libB"), ("RG3", "libA")] { + let rg = Map::::builder() + .insert(rg_tag::LIBRARY, String::from(library)) + .build() + .expect("read group builds"); + header = header.add_read_group(bstr::BString::from(id), rg); + } + // A fourth group with no LB at all, which must fall back to "unknown". + header = header.add_read_group( + bstr::BString::from("RG4"), + Map::::builder().build().expect("read group builds"), + ); + header.build() + } + + /// `library_idx` is the field `GroupKey` equality depends on, so read groups + /// sharing a library MUST collapse to one index and distinct libraries MUST + /// NOT — the first would merge unrelated UMI groups, the second would split + /// one group in half. + #[test] + fn from_header_assigns_one_index_per_distinct_library() { + let index = LibraryIndex::from_header(&header_with_read_groups()); + + let idx1 = index.get(LibraryIndex::hash_rg(b"RG1")); + let idx2 = index.get(LibraryIndex::hash_rg(b"RG2")); + let idx3 = index.get(LibraryIndex::hash_rg(b"RG3")); + + assert_eq!(idx1, idx3, "read groups sharing a library share an index"); + assert_ne!(idx1, idx2, "distinct libraries get distinct indices"); + assert_ne!(idx1, 0, "a resolved library is never the unknown index"); + assert_ne!(idx2, 0, "a resolved library is never the unknown index"); + + assert_eq!(index.library_name(idx1).as_ref(), "libA"); + assert_eq!(index.library_name(idx2).as_ref(), "libB"); + } + + /// A read group with no `LB` maps to the shared "unknown" library at index 0, + /// which is also the miss value — so an absent `LB` and an absent read group + /// group together rather than each forming a singleton. + #[test] + fn from_header_maps_missing_library_and_unknown_read_group_to_index_zero() { + let index = LibraryIndex::from_header(&header_with_read_groups()); + + assert_eq!(index.get(LibraryIndex::hash_rg(b"RG4")), 0, "no LB field → unknown"); + assert_eq!(index.get(LibraryIndex::hash_rg(b"ABSENT")), 0, "unseen RG → unknown"); + assert_eq!(index.library_name(0).as_ref(), "unknown"); + } + + /// Out-of-range indices fall back to "unknown" rather than panicking. + #[test] + fn library_name_saturates_to_unknown_for_an_out_of_range_index() { + let index = LibraryIndex::from_header(&header_with_read_groups()); + assert_eq!(index.library_name(u16::MAX).as_ref(), "unknown"); + } + + /// The string-keyed lookup carries the same LB resolution as the hashed one, + /// including the "unknown" fallback. + #[test] + fn build_library_lookup_maps_each_read_group_to_its_library() { + let lookup = build_library_lookup(&header_with_read_groups()); + + assert_eq!(lookup.len(), 4); + assert_eq!(lookup.get("RG1").expect("RG1 present").as_ref(), "libA"); + assert_eq!(lookup.get("RG2").expect("RG2 present").as_ref(), "libB"); + assert_eq!(lookup.get("RG3").expect("RG3 present").as_ref(), "libA"); + assert_eq!(lookup.get("RG4").expect("RG4 present").as_ref(), "unknown"); + } + + /// An empty header yields an index that resolves everything to unknown. + #[test] + fn from_header_with_no_read_groups_resolves_everything_to_unknown() { + let index = LibraryIndex::from_header(&Header::default()); + assert_eq!(index.get(LibraryIndex::hash_rg(b"RG1")), 0); + assert!(build_library_lookup(&Header::default()).is_empty()); + } + + /// `hash_bytes(None)` is the documented zero sentinel, and the typed wrappers + /// agree with it — `GroupKey`'s `cell_hash`/`name_hash` rely on that. + #[test] + fn hash_helpers_agree_and_treat_none_as_zero() { + assert_eq!(LibraryIndex::hash_bytes(None), 0); + assert_eq!(LibraryIndex::hash_cell_barcode(None), 0); + assert_eq!(LibraryIndex::hash_name(None), 0); + assert_eq!( + LibraryIndex::hash_name(Some(b"read1")), + LibraryIndex::hash_bytes(Some(b"read1")) + ); + assert_ne!(LibraryIndex::hash_name(Some(b"read1")), 0); + } +} 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-bam-io/src/reorder.rs b/crates/fgumi-bam-io/src/reorder.rs index 9c7081099..07384f415 100644 --- a/crates/fgumi-bam-io/src/reorder.rs +++ b/crates/fgumi-bam-io/src/reorder.rs @@ -48,7 +48,7 @@ pub struct ReorderBuffer { /// Sparse buffer: index (seq - `next_seq`) maps to `Option<(T, usize)>` where /// usize is the pre-computed heap size (0 if not tracked). buffer: VecDeque>, - /// Next sequence number to release (also the sequence number corresponding to buffer[0]). + /// Next sequence number to release (also the sequence number corresponding to `buffer[0]`). next_seq: u64, /// Number of items currently stored. count: usize, diff --git a/crates/fgumi-bgzf/src/lib.rs b/crates/fgumi-bgzf/src/lib.rs index 54e70aee8..667c330aa 100644 --- a/crates/fgumi-bgzf/src/lib.rs +++ b/crates/fgumi-bgzf/src/lib.rs @@ -17,9 +17,9 @@ pub use header::{ is_bgzf_header, }; pub use reader::{ - BGZF_EOF, BGZF_FOOTER_SIZE, RawBgzfBlock, decompress_block, decompress_block_into, - decompress_block_into_opts, decompress_block_slice_into, decompress_block_slice_into_opts, - read_raw_blocks, + BGZF_EOF, BGZF_FOOTER_SIZE, MAX_UNCOMPRESSED_BLOCK_SIZE, RawBgzfBlock, decompress_block, + decompress_block_into, decompress_block_into_opts, decompress_block_slice_into, + decompress_block_slice_into_opts, decompress_into_slice, read_raw_blocks, uncompressed_size, }; // Re-export the libdeflater decompressor so downstream crates can name the type // required by `decompress_block_into_opts` without depending on libdeflater diff --git a/crates/fgumi-bgzf/src/reader.rs b/crates/fgumi-bgzf/src/reader.rs index adcc13556..c46158500 100644 --- a/crates/fgumi-bgzf/src/reader.rs +++ b/crates/fgumi-bgzf/src/reader.rs @@ -545,6 +545,23 @@ fn decompress_and_verify( output: &mut Vec, verify_crc: bool, ) -> io::Result<()> { + // A corrupt footer can claim an arbitrary ISIZE — up to 4 GiB — while the + // BSIZE stays a valid u16. Reject any uncompressed size over the single-block + // maximum before it sizes an allocation, so neither the stored-copy path + // below nor the deflate `resize` zero-fills a multi-gigabyte buffer for a + // 4-byte corruption. Every real BGZF block decompresses to <= 64 KiB. This + // bounds the raw-`Vec` entry path (`decompress_block_into_opts`) to match the + // slice path, which reaches this same chokepoint. + if uncompressed_size > MAX_UNCOMPRESSED_BLOCK_SIZE { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "BGZF block claims {uncompressed_size} uncompressed bytes, over the \ + {MAX_UNCOMPRESSED_BLOCK_SIZE}-byte maximum for a single block (corrupt ISIZE)" + ), + )); + } + // Stored-block fast path. The deflate frame for a stored block is: // byte 0 : BFINAL | BTYPE | (5 padding bits to next byte boundary) // bytes 1-2: LEN (little-endian, u16) @@ -555,7 +572,7 @@ fn decompress_and_verify( // "stored". `payload_len == LEN + 5` is the structural guarantee that // there's exactly one stored sub-block spanning the BGZF payload — the // form every real level-0 producer emits. - if !compressed.is_empty() && compressed[0] & 0b110 == 0 { + if is_stored_block(compressed) { return copy_stored_and_verify( compressed, uncompressed_size, @@ -630,6 +647,132 @@ fn copy_stored_and_verify( output: &mut Vec, verify_crc: bool, ) -> io::Result<()> { + // Validate the stored framing (len >= 5, LEN + 5 == payload, LEN == ISIZE) + // and get the payload slice. Shared with the slice path via + // `parse_stored_frame` so the three framing checks live in one place. + let payload = parse_stored_frame(compressed, uncompressed_size)?; + let start = output.len(); + output.extend_from_slice(payload); + // `payload.len() == uncompressed_size` (enforced by `parse_stored_frame`), + // so `&output[start..]` is exactly the LEN payload bytes just appended. + let result = verify_decompression( + &output[start..], + uncompressed_size, + expected_crc, + block_len, + verify_crc, + ); + if result.is_err() { + output.truncate(start); + } + result +} + +// ============================================================================ +// Slice-fill decompression (arena buffer recycling) — grafted from #714 onto +// main's reader.rs (R10 sync). Fills a caller-provided &mut [u8] sized to the +// block ISIZE, verifying CRC unconditionally. +// ============================================================================ + +/// Largest uncompressed payload a single BGZF block can hold. +/// +/// A block's `BSIZE` is a `u16` holding the total size minus one, which caps a +/// block at 64 KiB; the SAM/BGZF spec and htslib's `BGZF_MAX_BLOCK_SIZE` both +/// use this bound, and the `bgzf` crate refuses to *write* a larger one +/// (`BlockSizeExceeded`). The footer's ISIZE is a `u32`, so a corrupt or +/// hostile block can claim up to 4 GiB; anything a caller sizes from that value +/// has to be bounded first. +/// +/// Note this is deliberately 64 KiB rather than [`crate::writer::BGZF_MAX_BLOCK_SIZE`] +/// (65280), which is the size *we* fill blocks to. Other writers legitimately +/// emit up to the spec limit, so validating against our own chunk size would +/// reject valid input. +pub const MAX_UNCOMPRESSED_BLOCK_SIZE: usize = 64 * 1024; + +/// The uncompressed size a BGZF block's footer claims, validated against the +/// per-block maximum. +/// +/// This is the size a caller must give [`decompress_into_slice`]'s `out`, and +/// it is the number a caller allocates from, so it is checked rather than +/// returned raw: ISIZE is a `u32` sitting in the file, and a corrupt or hostile +/// footer claiming 4 GiB would otherwise become a 4 GiB allocation in the +/// caller before this crate ever saw the block. +/// +/// Takes the block as a `&[u8]` so reading the footer costs nothing — +/// [`RawBgzfBlock::uncompressed_size`] answers the same question but needs an +/// owned `Vec`, which would mean copying the block to size a slot for it. +/// +/// # Errors +/// +/// Returns `io::ErrorKind::InvalidData` if `block` is too short to hold a BGZF +/// header + footer, or if the claimed size exceeds +/// [`MAX_UNCOMPRESSED_BLOCK_SIZE`]. +pub fn uncompressed_size(block: &[u8]) -> io::Result { + if block.len() < BGZF_HEADER_SIZE + BGZF_FOOTER_SIZE { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "BGZF block too short to contain header + footer", + )); + } + let claimed = uncompressed_size_from_slice(block); + if claimed > MAX_UNCOMPRESSED_BLOCK_SIZE { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "BGZF block claims an uncompressed size of {claimed} bytes, above the \ + {MAX_UNCOMPRESSED_BLOCK_SIZE}-byte maximum for a single block" + ), + )); + } + Ok(claimed) +} + +/// Whether a BGZF block's deflate payload is a **stored** (uncompressed) frame, +/// i.e. `BTYPE = 00`. +/// +/// `BTYPE` lives in bits 1-2 of the deflate stream's first byte, so masking with +/// `0b110` isolates it. An empty payload is not a stored frame — there is no +/// first byte to read, and the framing parser needs five. +/// +/// Shared by every decompress entry point so the dispatch predicate cannot +/// drift between them; [`parse_stored_frame`] validates the rest of the framing +/// once this returns `true`. +#[must_use] +fn is_stored_block(compressed: &[u8]) -> bool { + !compressed.is_empty() && compressed[0] & 0b110 == 0 +} + +/// Parse and validate a deflate **stored** sub-block frame, returning the +/// LEN-byte payload slice (`&compressed[5..]`). Shared by +/// [`copy_stored_and_verify`] and [`copy_stored_and_verify_slice`] so the +/// stored-framing checks cannot drift between the `Vec` and fixed-slice +/// decompress entry points. +/// +/// The caller is responsible for confirming that `compressed[0]` has +/// `BTYPE = 00` before calling. This validates the rest of the framing: +/// +/// * `compressed.len() >= 5` (room for the BFINAL/BTYPE byte + LEN + NLEN). +/// * `LEN + 5 == compressed.len()` — exactly one stored sub-block fills the +/// BGZF payload. Every real level-0 producer (`samtools view -u`, htsjdk, +/// [`InlineBgzfCompressor`](crate::writer::InlineBgzfCompressor)) emits +/// this shape. We intentionally do **not** fall back to libdeflater if +/// this check fails: the input is either malformed or uses a multi-sub- +/// block stored stream we have no real-world reason to support, and a +/// loud error beats silently papering over corruption. +/// * `LEN == expected_len` — the deflate frame's LEN agrees with the BGZF +/// footer's ISIZE (the caller passes the footer ISIZE, or the caller-sized +/// output length that equals it). +/// +/// NLEN (one's complement of LEN) is not checked. NLEN doesn't cover the +/// payload bytes, so a corrupt NLEN with an intact payload would pass the +/// framing check anyway; the BGZF footer's CRC32 is the authoritative +/// integrity check on the data itself. +/// +/// # Errors +/// +/// Returns `io::ErrorKind::InvalidData` if the stored framing is truncated, +/// spans more than one sub-block, or its LEN disagrees with `expected_len`. +fn parse_stored_frame(compressed: &[u8], expected_len: usize) -> io::Result<&[u8]> { // Deflate framing is 5 bytes (BFINAL/BTYPE byte + LEN + NLEN). if compressed.len() < 5 { return Err(io::Error::new( @@ -651,33 +794,166 @@ fn copy_stored_and_verify( format!("BGZF stored block size mismatch: LEN={len}, payload={}", compressed.len()), )); } - if len != uncompressed_size { + if len != expected_len { return Err(io::Error::new( io::ErrorKind::InvalidData, - format!("BGZF stored block ISIZE mismatch: footer={uncompressed_size}, LEN={len}"), + format!("BGZF stored block ISIZE mismatch: footer={expected_len}, LEN={len}"), )); } + // `len + 5 == compressed.len()`, so this is exactly the LEN payload bytes. + Ok(&compressed[5..]) +} - let start = output.len(); - // `len + 5 == compressed.len()` is checked above, so `&compressed[5..]` - // is exactly the LEN payload bytes. - output.extend_from_slice(&compressed[5..]); - let result = verify_decompression( - &output[start..start + len], - uncompressed_size, - expected_crc, - block_len, - verify_crc, - ); - if result.is_err() { - output.truncate(start); - } - result +/// Slice-writing sibling of [`copy_stored_and_verify`]: copy a deflate +/// **stored** sub-block's payload straight into a caller-sized `out` slice +/// (whose length is the block's ISIZE), skipping libdeflater, then verify the +/// BGZF footer CRC32 over the copied bytes. Returns the number of bytes written +/// (`== out.len()` on success). Used by [`decompress_into_slice`]'s fast path. +/// +/// The caller must have confirmed `compressed[0]` has `BTYPE = 00` first. +/// Unlike the `Vec` sibling there is nothing to roll back: `out` is the +/// caller's buffer, and — matching [`decompress_into_slice`]'s non-stored path +/// — a verification failure leaves `out` holding the (CRC-rejected) bytes, +/// which the caller discards along with the returned error. +fn copy_stored_and_verify_slice( + compressed: &[u8], + out: &mut [u8], + expected_crc: u32, + block_len: usize, +) -> io::Result { + // `out` is caller-sized to the footer's ISIZE, so passing `out.len()` also + // enforces LEN == ISIZE; the returned payload is then exactly `out.len()`. + let payload = parse_stored_frame(compressed, out.len())?; + out.copy_from_slice(payload); + verify_decompression(out, out.len(), expected_crc, block_len, true)?; + Ok(out.len()) } -// ============================================================================ -// Tests -// ============================================================================ +/// Inflate `compressed` into the whole of `out` and verify the result against +/// the BGZF footer, returning the number of bytes written. +/// +/// `out.len()` is taken as the expected uncompressed size, so +/// [`verify_decompression`] checks both the exact fill +/// (`bytes_written == out.len()`) and the CRC32. Shared by +/// [`decompress_into_slice`] and [`decompress_and_verify`]'s non-stored branch +/// so the inflate-then-verify invariant lives in one place — the same reason +/// [`parse_stored_frame`] exists for the stored branch. +fn deflate_into_slice_and_verify( + compressed: &[u8], + expected_crc: u32, + block_len: usize, + decompressor: &mut Decompressor, + out: &mut [u8], +) -> io::Result { + let bytes_written = decompressor.deflate_decompress(compressed, out).map_err(|e| { + io::Error::new(io::ErrorKind::InvalidData, format!("BGZF decompression failed: {e:?}")) + })?; + verify_decompression(&out[..bytes_written], out.len(), expected_crc, block_len, true)?; + Ok(bytes_written) +} + +/// Decompress a full BGZF block's DEFLATE payload into a caller-provided, +/// pre-sized slice. `out.len()` must equal the block's ISIZE (uncompressed +/// size), which callers should take from [`uncompressed_size`] — it reads the +/// footer straight off the same `&[u8]` and bounds the claim, so the slot a +/// caller allocates can never come from an unvalidated `u32`. +/// +/// This is the fixed-slice analogue of [`decompress_block_slice_into`] (which +/// appends to a `Vec`); it lets a caller decompress straight into an arena +/// slot rather than into a buffer it then has to copy out of. The two names are +/// close and their `slice` refers to opposite operands: in +/// [`decompress_block_slice_into`] it is the *input* block, given as a slice +/// rather than a [`RawBgzfBlock`]; here it is the *output*. +/// +/// Like the sibling decompressors ([`decompress_block_into`] / +/// [`decompress_block_slice_into`]), the decompressed payload is verified +/// against the BGZF footer's ISIZE (it must exactly fill `out`) and CRC32, so +/// a short fill or a silently-corrupt-but-decodable block is caught here +/// rather than fed into the arena. +/// +/// # Returns +/// +/// The number of bytes written, which on success is always `out.len()`: the +/// decompressing paths are held to it by the exact-fill check, and a block with +/// an ISIZE of zero (the BGZF EOF marker) returns `0` against the zero-length +/// slot its ISIZE requires. It is returned so the call reads like the +/// `Read`-style APIs it sits beside, not because a short result is reachable. +/// +/// # Errors +/// +/// Returns an `io::Error` if `block` is shorter than a BGZF header + footer, if +/// the footer's ISIZE exceeds [`MAX_UNCOMPRESSED_BLOCK_SIZE`], if `out` is not +/// sized to that ISIZE, if the DEFLATE stream is invalid, if it does not +/// exactly fill `out`, or if the CRC32 does not match the footer. +/// +/// Note this validates the block's *framing*, not its full header: a payload +/// whose header [`crate::header::validate`] would reject can still reach the +/// decompressor and fail there. Callers reading from a stream get the header +/// check from [`read_raw_blocks`]; this entry point is for callers that already +/// hold a framed block. +/// +/// On error `out` is left clobbered — unlike the `Vec` siblings, which roll +/// their output back. There is nothing to roll back to here: the buffer belongs +/// to the caller, who must treat its contents as undefined unless this returns +/// `Ok`. +pub fn decompress_into_slice( + block: &[u8], + decompressor: &mut Decompressor, + out: &mut [u8], +) -> io::Result { + // Same accessor the caller sizes `out` with, so the two cannot disagree + // about either the value or the bound. It carries the length and ISIZE + // checks, which is why neither is repeated here. + let uncompressed_size = uncompressed_size(block)?; + // Check the caller's slot against the footer up front rather than letting a + // mis-sized `out` surface downstream. Both paths below compare against + // `out.len()` on the assumption it *is* the ISIZE, so without this a wrongly + // sized slot is reported as a fault in the block: the stored path would say + // `ISIZE mismatch: footer=`, naming a value the footer does not + // contain, and send a reader looking for file corruption instead of at the + // arena that sized the slot. + if out.len() != uncompressed_size { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "BGZF output slice is {} bytes, but the block's ISIZE is {uncompressed_size}", + out.len() + ), + )); + } + // Only the *exact* BGZF EOF marker — the one every reader meets at end of + // stream — is waved through here, matching the `Vec` siblings + // (`decompress_block_slice_into` short-circuits on `data == BGZF_EOF`). + // Every other zero-ISIZE block, including a CRC-corrupted EOF marker whose + // bytes no longer match, falls through to `deflate_into_slice_and_verify` + // so its CRC32 is still checked. Gating on `uncompressed_size == 0` instead + // would let a corrupted zero-size block return `Ok(0)` unverified, which is + // the one thing the `Vec` path is careful not to do. libdeflater does + // return `Ok(0)` for the marker's `03 00` payload against a zero-length + // slice, so this is not repairing a failure; it makes the answer ours + // rather than resting on an undocumented edge of the C library. + // + // Placed *after* the size check on purpose: ahead of it, a caller passing a + // wrongly-sized slot for the marker would get a silent `Ok(0)` instead of + // being told the slot is wrong. + if block == BGZF_EOF { + return Ok(0); + } + let compressed = compressed_data_from_slice(block); + // Stored-block fast path — level-0 blocks (`samtools view -u`, htsjdk's + // level-0 writer, [`InlineBgzfCompressor::new(0)`]) skip the libdeflater + // round-trip and get the stored-framing-specific LEN/ISIZE diagnostics. + if is_stored_block(compressed) { + return copy_stored_and_verify_slice(compressed, out, crc32_from_slice(block), block.len()); + } + deflate_into_slice_and_verify( + compressed, + crc32_from_slice(block), + block.len(), + decompressor, + out, + ) +} #[cfg(test)] mod tests { @@ -1321,6 +1597,45 @@ mod tests { assert!(out.is_empty(), "output should be rolled back on failure"); } + /// A corrupt footer claiming a multi-gigabyte ISIZE must fail closed on the + /// raw-`Vec` entry path (`decompress_block_into_opts`) BEFORE `resize` + /// allocates for it — not after decompression, as the `+1` mismatch above + /// does. The `uncompressed_size` bound in `decompress_and_verify` is the only + /// guard on this path (the slice path reaches the same chokepoint). + #[test] + fn decompress_rejects_a_giant_isize_before_allocating() { + use crate::writer::InlineBgzfCompressor; + + let original = b"payload for the oversized-ISIZE rejection test, long enough to deflate"; + let mut compressor = InlineBgzfCompressor::new(6); + compressor.write_all(original).expect("write"); + compressor.flush().expect("flush"); + let blocks = compressor.take_blocks(); + assert_eq!(blocks.len(), 1); + assert_ne!( + blocks[0].data[BGZF_HEADER_SIZE] & 0b110, + 0, + "expected a deflate-coded block so the resize path (not the stored copy) runs" + ); + + // Set the footer ISIZE (last 4 bytes, little-endian) to ~4 GiB while the + // BSIZE stays a valid u16 — the exact shape a 4-byte corruption produces. + let mut giant = blocks[0].data.clone(); + let n = giant.len(); + giant[n - 4..].copy_from_slice(&u32::MAX.to_le_bytes()); + let giant_block = RawBgzfBlock { data: giant }; + + let mut decompressor = Decompressor::new(); + let mut out = Vec::new(); + let err = decompress_block_into_opts(&giant_block, &mut decompressor, &mut out, false) + .expect_err("a >64 KiB ISIZE claim must be rejected before allocating"); + assert!( + err.to_string().contains("maximum for a single block"), + "error should name the single-block maximum, got: {err}" + ); + assert!(out.is_empty(), "no buffer must be allocated/left on the rejected path"); + } + /// Slice-API twin of `decompress_opts_skips_crc_but_still_checks_size`. /// /// The FASTQ pipeline decompresses through the slice entry point @@ -1436,4 +1751,277 @@ mod tests { .expect("slice API: verify_crc=false must skip the CRC32 check"); assert!(out.is_empty()); } + + // ── #714 graft: arena slice-fill decompression tests (R10 sync) ── + + // ── `uncompressed_size` ───────────────────────────────────────────────── + + /// The public accessor callers size their slot with must agree with the + /// footer for a real block, and must reject the two inputs that would + /// otherwise become a bad allocation: a block too short to have a footer, + /// and a footer claiming more than a block can hold. + #[rstest] + #[case::level_0_stored(0)] + #[case::level_6_deflate(6)] + fn uncompressed_size_reports_the_footer(#[case] level: u32) { + let payload = one_block_payload(); + let block = first_block_at_level(&payload, level); + assert_eq!(uncompressed_size(&block).expect("valid block"), payload.len()); + } + + /// A claim above the per-block maximum is refused rather than handed back + /// for a caller to allocate from. ISIZE is a `u32`, so the untrusted range + /// runs to 4 GiB. + #[test] + fn uncompressed_size_rejects_a_claim_above_the_block_maximum() { + let mut block = first_block_at_level(&one_block_payload(), 6); + let len = block.len(); + let claimed = u32::try_from(MAX_UNCOMPRESSED_BLOCK_SIZE + 1).expect("fits"); + block[len - 4..].copy_from_slice(&claimed.to_le_bytes()); + let err = uncompressed_size(&block).expect_err("an oversized claim must be rejected"); + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + assert!(err.to_string().contains("above the"), "got: {err}"); + } + + /// Exactly the per-block maximum is legal — other writers may fill a block + /// to the spec limit even though this crate stops at `BGZF_MAX_BLOCK_SIZE`, + /// so the bound must not be off by one against them. + #[test] + fn uncompressed_size_accepts_the_block_maximum() { + let mut block = first_block_at_level(&one_block_payload(), 6); + let len = block.len(); + let claimed = u32::try_from(MAX_UNCOMPRESSED_BLOCK_SIZE).expect("fits"); + block[len - 4..].copy_from_slice(&claimed.to_le_bytes()); + assert_eq!( + uncompressed_size(&block).expect("the maximum itself is valid"), + MAX_UNCOMPRESSED_BLOCK_SIZE + ); + } + + /// Too short to hold a footer at all — must error rather than read the + /// ISIZE out of whatever bytes happen to be there. + #[test] + fn uncompressed_size_rejects_a_block_without_a_footer() { + let err = uncompressed_size(&[0u8; 10]).expect_err("a short block must be rejected"); + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + assert!(err.to_string().contains("too short to contain"), "got: {err}"); + } + + // ── `decompress_into_slice` ───────────────────────────────────────────── + + /// Compress `payload` into BGZF blocks at `level` and return the first + /// block's raw bytes. Level 0 emits deflate **stored** blocks (the fast + /// path); any other level emits a real deflate stream. + fn first_block_at_level(payload: &[u8], level: u32) -> Vec { + let mut compressor = crate::writer::InlineBgzfCompressor::new(level); + compressor.write_all(payload).expect("write payload"); + compressor.flush().expect("flush"); + compressor.take_blocks().into_iter().next().expect("payload fits in one block").data + } + + /// A payload that fits comfortably in one BGZF block (< 64 KiB). + fn one_block_payload() -> Vec { + b"the quick brown fox jumps over the lazy dog".repeat(100) + } + + /// Both decompress paths must fill a pre-sized `&mut [u8]` slot with + /// exactly the original bytes: level 0 through the stored-block copy, any + /// other level through libdeflater. `expect_stored` pins which path the + /// case actually takes, so a producer change that stops emitting stored + /// blocks fails here rather than silently reducing coverage to one path. + #[rstest] + #[case::level_0_stored(0, true)] + #[case::level_6_deflate(6, false)] + fn decompress_into_slice_fills_presized_slot(#[case] level: u32, #[case] expect_stored: bool) { + let payload = one_block_payload(); + let block = first_block_at_level(&payload, level); + + let compressed = compressed_data_from_slice(&block); + assert_eq!( + is_stored_block(compressed), + expect_stored, + "level {level} took the unexpected decompress path" + ); + + let mut out = vec![0u8; uncompressed_size_from_slice(&block)]; + let n = decompress_into_slice(&block, &mut Decompressor::new(), &mut out) + .expect("decompress into slot"); + assert_eq!(n, payload.len()); + assert_eq!(&out[..n], payload.as_slice()); + } + + /// The BGZF EOF marker carries no payload, and every stream ends with one, + /// so a caller decompressing each block of a stream in turn hits it. It has + /// to succeed against a zero-length slot, as the two `Vec` siblings do. + /// + /// Its payload is `03 00` — a fixed-Huffman empty frame — so + /// `is_stored_block` is false and it would otherwise reach libdeflater with + /// a zero-length output slice. That returns `Ok(0)` today; this pins the + /// result to the crate rather than to that behaviour. + #[test] + fn decompress_into_slice_accepts_the_eof_marker() { + let compressed = compressed_data_from_slice(&BGZF_EOF); + assert!( + !is_stored_block(compressed), + "the EOF payload is a fixed-Huffman frame, so the stored path must not claim it" + ); + assert_eq!(uncompressed_size(&BGZF_EOF).expect("the EOF marker is a valid block"), 0); + + let n = decompress_into_slice(&BGZF_EOF, &mut Decompressor::new(), &mut []) + .expect("the EOF marker must decompress to nothing"); + assert_eq!(n, 0); + } + + /// A zero-ISIZE block must still reject a wrongly-sized slot. The obvious + /// place to short-circuit the EOF marker is *before* the slot-size check, + /// which would turn this into a silent `Ok(0)` and hide the caller's sizing + /// bug — the same misreport the up-front check exists to prevent. + #[test] + fn decompress_into_slice_rejects_a_sized_slot_for_an_empty_block() { + let mut out = [0u8; 16]; + let err = decompress_into_slice(&BGZF_EOF, &mut Decompressor::new(), &mut out) + .expect_err("a 16-byte slot for a 0-byte block must be rejected"); + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + assert!(err.to_string().contains("output slice is"), "got: {err}"); + } + + /// A CRC-corrupted EOF marker has ISIZE 0 but is no longer the exact + /// [`BGZF_EOF`] bytes, so the slot path must run it through CRC verification + /// like the `Vec` siblings do, not wave it through on `uncompressed_size == + /// 0`. Regression: the short-circuit here once keyed off the ISIZE, so a + /// corrupted zero-size block returned `Ok(0)` against its (correctly sized) + /// zero-length slot without the CRC ever being checked. + #[test] + fn decompress_into_slice_verifies_a_corrupted_eof_marker() { + // Flip a CRC32 footer bit: still ISIZE 0, but no longer the exact marker. + let crc_off = BGZF_EOF.len() - BGZF_FOOTER_SIZE; + let mut corrupted = BGZF_EOF.to_vec(); + corrupted[crc_off] ^= 0x01; + assert_ne!(corrupted.as_slice(), &BGZF_EOF[..], "must not be the exact EOF marker"); + assert_eq!(uncompressed_size(&corrupted).expect("valid framing"), 0, "still zero ISIZE"); + + // Correctly sized (zero-length) slot, so the size check cannot fire — + // the rejection must come from CRC verification, not slot sizing. + let err = decompress_into_slice(&corrupted, &mut Decompressor::new(), &mut []) + .expect_err("a corrupted zero-size block must not be waved through"); + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + assert!(err.to_string().contains("CRC32"), "error should mention CRC32: {err}"); + } + + /// A block at `level` with a flipped footer CRC32 bit: it still produces + /// the right number of bytes, but the checksum no longer matches. Run at + /// both levels because the CRC is verified at two separate call sites — + /// `copy_stored_and_verify_slice` for level 0 and + /// `deflate_into_slice_and_verify` for the rest — and deleting either one + /// used to leave the whole suite green. + fn block_bad_crc(level: u32) -> (Vec, Vec) { + let mut block = first_block_at_level(&one_block_payload(), level); + let len = block.len(); + block[len - 8] ^= 0x01; // footer CRC32 is bytes [len - 8 .. len - 4] + let out = vec![0u8; uncompressed_size_from_slice(&block)]; + (block, out) + } + + /// A level-6 block whose footer ISIZE is 16 too high, with `out` sized to + /// that inflated ISIZE. libdeflater reports the true (shorter) length via + /// `Ok(n)` rather than erroring, so this is what reaches the exact-fill + /// check — the guarantee the docs single out, and the one branch that + /// stops the arena being handed a slot with a stale tail. + fn block_short_fill() -> (Vec, Vec) { + let mut block = first_block_at_level(&one_block_payload(), 6); + let len = block.len(); + let inflated = uncompressed_size_from_slice(&block) + 16; + block[len - 4..].copy_from_slice(&u32::try_from(inflated).expect("fits").to_le_bytes()); + (block, vec![0u8; inflated]) + } + + /// A footer claiming more than a single BGZF block can hold. ISIZE is a + /// `u32`, so an unbounded contract would have a caller allocate up to 4 GiB + /// from a corrupt footer before anything validated it. + fn block_isize_above_max() -> (Vec, Vec) { + let mut block = first_block_at_level(&one_block_payload(), 6); + let len = block.len(); + let claimed = u32::try_from(MAX_UNCOMPRESSED_BLOCK_SIZE + 1).expect("fits"); + block[len - 4..].copy_from_slice(&claimed.to_le_bytes()); + // Deliberately a *correctly* sized slot for the claim: the bound must + // fire on the claim itself, not merely because `out` disagrees with it. + (block, vec![0u8; MAX_UNCOMPRESSED_BLOCK_SIZE + 1]) + } + + /// A caller-oversized `out` (ISIZE + 16). The error must name the slice as + /// the thing that is wrong: before the up-front size check this surfaced as + /// a complaint about the block's own ISIZE, sending a reader after file + /// corruption that isn't there. The level is irrelevant — this check fires + /// before the stored/deflate dispatch — so one case covers it. + fn block_oversized_out() -> (Vec, Vec) { + let block = first_block_at_level(&one_block_payload(), 6); + let out = vec![0u8; uncompressed_size_from_slice(&block) + 16]; + (block, out) + } + + /// Synthesised stored block whose 4-byte payload can't hold the 5-byte + /// deflate stored frame (BFINAL/BTYPE + LEN + NLEN), so the stored fast + /// path errors before copying. `InlineBgzfCompressor` only emits valid + /// frames, so this is built byte by byte: 18-byte header + 4-byte payload + /// + 8-byte footer. + /// + /// The header is a **well-formed** one — `XLEN = 6` and the `BC` subfield's + /// `SLEN = 2` are set, not left zero — so the fixture isolates the stored + /// framing defect it is named for. An invalid header here would be rejected + /// for a different reason by anything that validates one, making the case + /// pass for the wrong cause. + fn block_truncated_stored_framing() -> (Vec, Vec) { + const BLOCK_SIZE: usize = BGZF_HEADER_SIZE + 4 + BGZF_FOOTER_SIZE; + let mut data = vec![0u8; BLOCK_SIZE]; + data[0] = 0x1f; // gzip magic + deflate method + FEXTRA flag + data[1] = 0x8b; + data[2] = 0x08; + data[3] = 0x04; + data[10] = 0x06; // XLEN = 6: the extra field is exactly the BC subfield + data[11] = 0x00; + data[12] = b'B'; // BC subfield ID + data[13] = b'C'; + data[14] = 0x02; // SLEN = 2: BC holds a two-byte BSIZE + data[15] = 0x00; + let bsize_bytes = u16::try_from(BLOCK_SIZE - 1).expect("block fits in u16").to_le_bytes(); + data[16] = bsize_bytes[0]; + data[17] = bsize_bytes[1]; + // Payload bytes 18..22 stay zero → BTYPE bits (1-2) are 00, the stored + // fast path. Footer ISIZE = 1 matches the caller-sized `out`; the CRC + // is never reached. + data[BLOCK_SIZE - 4] = 1; + debug_assert!(crate::header::is_bgzf_header(&data), "fixture header must be well-formed"); + (data, vec![0u8; 1]) + } + + /// A block shorter than the 26-byte header + footer minimum — must error + /// rather than panic on an out-of-bounds slice or a subtract overflow. + fn block_too_short() -> (Vec, Vec) { + (vec![0u8; 10], vec![0u8; 16]) + } + + /// Every malformed input is rejected as `InvalidData` with a message that + /// names what went wrong, so a corrupt block is diagnosable rather than a + /// bare "decompression failed". + #[rstest] + #[case::bad_crc_deflate(block_bad_crc(6), "CRC32")] + #[case::bad_crc_stored(block_bad_crc(0), "CRC32")] + #[case::short_fill(block_short_fill(), "size mismatch")] + #[case::isize_above_max(block_isize_above_max(), "above the")] + #[case::oversized_out(block_oversized_out(), "output slice is")] + #[case::truncated_stored_framing(block_truncated_stored_framing(), "stored block too small")] + #[case::too_short_block(block_too_short(), "too short to contain")] + fn decompress_into_slice_rejects_invalid( + #[case] block_and_out: (Vec, Vec), + #[case] expect_substr: &str, + ) { + let (block, mut out) = block_and_out; + let err = decompress_into_slice(&block, &mut Decompressor::new(), &mut out) + .expect_err("malformed block must be rejected"); + assert_eq!(err.kind(), io::ErrorKind::InvalidData, "expected InvalidData, got {err:?}"); + assert!( + err.to_string().contains(expect_substr), + "error should contain {expect_substr:?}, got: {err}" + ); + } } diff --git a/crates/fgumi-bgzf/src/writer.rs b/crates/fgumi-bgzf/src/writer.rs index d63e9133a..8de60e6c0 100644 --- a/crates/fgumi-bgzf/src/writer.rs +++ b/crates/fgumi-bgzf/src/writer.rs @@ -16,6 +16,28 @@ use std::io; /// Maximum uncompressed size for a BGZF block (64KB - header/footer overhead). pub const BGZF_MAX_BLOCK_SIZE: usize = bgzf::BGZF_BLOCK_SIZE; +/// Cap on the number of compression buffers [`InlineBgzfCompressor`] pools for +/// reuse. +/// +/// The pool only has to cover the *compress* side, which takes one buffer at a +/// time: `compress_current_buffer` pops exactly one per block. The queue on the +/// other side can be arbitrarily long — `write_all` appends a block per 64 KiB +/// written — so a drain of many blocks hands back more buffers than the +/// compressor will ever want at once. The surplus is dropped rather than parked +/// for the compressor's lifetime, which is what keeps the pool's memory a +/// function of this constant instead of of how much the caller buffered before +/// draining. +pub const MAX_POOLED_BUFFERS: usize = 4; + +/// Largest buffer [`InlineBgzfCompressor::recycle_buffer`] will pool. +/// +/// A block buffer is sized by `Compressor::compress` to +/// `header + deflate_compress_bound(input) + footer`. That bound exceeds the +/// input for incompressible data, so a legitimate buffer can run somewhat over +/// [`BGZF_MAX_BLOCK_SIZE`]; double it so no real block buffer is ever refused, +/// while still rejecting a wildly oversized `Vec` a caller hands back. +pub const MAX_POOLED_BUFFER_BYTES: usize = BGZF_MAX_BLOCK_SIZE * 2; + // ============================================================================ // Block types // ============================================================================ @@ -184,6 +206,36 @@ impl InlineBgzfCompressor { std::mem::take(&mut self.completed_blocks) } + /// Return a drained block buffer to the internal pool for reuse by a later + /// block compression. + /// + /// Consumers that drive the compressor with [`write_all`](Self::write_all) + + /// [`flush`](Self::flush) + [`take_blocks`](Self::take_blocks) (rather than + /// [`write_blocks_to`](Self::write_blocks_to), which recycles automatically) + /// otherwise leave `buffer_pool` empty, so every block compression allocates + /// a fresh output `Vec`. Such a consumer can hand back any block `Vec` it is + /// done with via this method to restore the recycling. + /// + /// The buffer is cleared — but not shrunk, since its capacity is the whole + /// point — before being pooled. A buffer is only pooled if the pool has + /// room ([`MAX_POOLED_BUFFERS`]) *and* the buffer is a plausible block + /// buffer ([`MAX_POOLED_BUFFER_BYTES`]); otherwise it is dropped. Bounding + /// the count alone would let one oversized `Vec` handed in by a caller sit + /// in the pool for the compressor's lifetime, so the pool's memory is + /// bounded on both axes. + /// + /// Not to be confused with the pipeline's `WorkerCoreState::recycle_buffer`, + /// which pools the *uncompressed* worker buffers under a different cap. + /// This one pools the compressor's own compressed-block output buffers. + pub fn recycle_buffer(&mut self, mut buffer: Vec) { + if self.buffer_pool.len() < MAX_POOLED_BUFFERS + && buffer.capacity() <= MAX_POOLED_BUFFER_BYTES + { + buffer.clear(); + self.buffer_pool.push(buffer); + } + } + /// Write all completed compressed blocks directly to output and recycle buffers. /// /// This is efficient for single-threaded use as it writes blocks directly @@ -191,14 +243,39 @@ impl InlineBgzfCompressor { /// /// # Errors /// - /// Returns an error if writing to the output fails. + /// Returns an error if writing to the output fails. The block that failed + /// to write, and every block after it, is left in the queue, so they are + /// available to **inspect** rather than being silently discarded. + /// + /// They are *not* safe to replay. [`io::Write::write_all`] loops over + /// [`write`](io::Write::write), advancing past each `Ok(n)`, so it can + /// commit part of the failing block to `output` before an error surfaces — + /// and the whole block is re-queued, not the unwritten remainder. Writing + /// the queue again would repeat those bytes inside a gzip member and + /// produce a stream no BGZF reader can decode. Treat `output` as + /// unrecoverable after this returns an error: the queued blocks are for + /// diagnostics, or for writing somewhere new. pub fn write_blocks_to(&mut self, output: &mut W) -> io::Result<()> { - for block in self.completed_blocks.drain(..) { - output.write_all(&block.data)?; - // Recycle the buffer for reuse - let mut buf = block.data; - buf.clear(); - self.buffer_pool.push(buf); + // Nothing queued is the common call (a flush with no pending blocks). + // Returning here keeps `completed_blocks`'s allocation, which the + // `mem::take` below would otherwise hand away and force a realloc on + // the next compression. + if self.completed_blocks.is_empty() { + return Ok(()); + } + // Drain into a temporary so `recycle_buffer` (which borrows `self` + // mutably) can be called per block without holding a borrow on + // `self.completed_blocks`. + let mut remaining = std::mem::take(&mut self.completed_blocks).into_iter(); + while let Some(block) = remaining.next() { + if let Err(e) = output.write_all(&block.data) { + self.completed_blocks.push(block); + self.completed_blocks.extend(remaining); + return Err(e); + } + // Route the drained buffer through the capped recycle path so the + // pool stays bounded, just like the steady-state recycle path. + self.recycle_buffer(block.data); } Ok(()) } @@ -209,9 +286,12 @@ impl InlineBgzfCompressor { return Ok(()); } - // Get buffer from pool or allocate new + // Get buffer from pool or allocate new. No clear is needed: the only + // thing this buffer is handed to is `Compressor::compress`, which + // resizes it from zero before writing (bgzf 0.4's `resize_uninit` + // begins with `Vec::clear`). Capacity is what the pool is for, and that + // survives. let mut compressed_data = self.buffer_pool.pop().unwrap_or_default(); - compressed_data.clear(); // Compress using bgzf crate's Compressor self.compressor @@ -443,4 +523,169 @@ mod tests { fn test_compress_level_out_of_range_panics() { let _ = InlineBgzfCompressor::new(13); } + + // ── Buffer recycling ──────────────────────────────────────────────────── + + /// Compress `payload` into one block and hand the compressor back. + fn compressor_with_one_block(payload: &[u8]) -> (InlineBgzfCompressor, Vec) { + let mut compressor = InlineBgzfCompressor::new(6); + compressor.write_all(payload).expect("writing data should succeed"); + compressor.flush().expect("flushing compressor should succeed"); + let blocks = compressor.take_blocks(); + (compressor, blocks) + } + + /// A `take_blocks` consumer keeps ownership of every block buffer, so the + /// pool is left empty and each subsequent compression allocates afresh. + /// Handing a drained buffer back must restore the recycling. + #[test] + fn test_recycle_buffer_repopulates_the_pool() { + let (mut compressor, blocks) = compressor_with_one_block(b"payload for one block"); + assert!( + compressor.buffer_pool.is_empty(), + "take_blocks should leave the pool empty, that is the gap recycle_buffer closes" + ); + + let buffer = blocks.into_iter().next().expect("one block").data; + let recycled_capacity = buffer.capacity(); + // The identity of the allocation, not just of the `Vec`. Everything + // below is about this exact heap block being handed back to the next + // compression rather than freed and replaced. + let recycled_ptr = buffer.as_ptr(); + compressor.recycle_buffer(buffer); + assert_eq!(compressor.buffer_pool.len(), 1); + // The allocation is what is being reused, so the capacity has to survive + // pooling -- a `shrink_to_fit` alongside the `clear` would satisfy every + // other assertion here while deleting the point of the method. + assert_eq!( + compressor.buffer_pool[0].capacity(), + recycled_capacity, + "pooling must retain the buffer's allocation, not just the buffer" + ); + + // The next compression must consume the pooled buffer rather than + // allocate. Asserting the pool merely drained would also pass for a + // pop-and-discard implementation, so pin the produced block to the + // recycled allocation instead. + compressor.write_all(b"a second block").expect("writing data should succeed"); + compressor.flush().expect("flushing compressor should succeed"); + assert!(compressor.buffer_pool.is_empty(), "the pooled buffer should have been taken"); + let next = compressor.take_blocks(); + assert_eq!( + next[0].data.as_ptr(), + recycled_ptr, + "the new block should be built in the recycled allocation, not a fresh one" + ); + } + + /// The pool is bounded on capacity as well as count: a caller handing back + /// an outsized `Vec` must not park that allocation for the compressor's + /// lifetime. A real block buffer (~64 KiB) is well under the limit, so this + /// rejects only what it is meant to. + #[test] + fn test_recycle_buffer_refuses_an_oversized_buffer() { + let (mut compressor, blocks) = compressor_with_one_block(b"payload for one block"); + let real = blocks.into_iter().next().expect("one block").data; + assert!( + real.capacity() <= MAX_POOLED_BUFFER_BYTES, + "a real block buffer ({} bytes) must still be poolable", + real.capacity() + ); + + compressor.recycle_buffer(Vec::with_capacity(MAX_POOLED_BUFFER_BYTES + 1)); + assert!(compressor.buffer_pool.is_empty(), "an oversized buffer should not be pooled"); + + compressor.recycle_buffer(real); + assert_eq!(compressor.buffer_pool.len(), 1, "a real block buffer should be pooled"); + } + + /// `write_blocks_to` is the other path into the pool, and the line this + /// change rewrote. It must recycle what it drains, and honour the same cap + /// as the explicit entry point -- a drain of many blocks cannot grow the + /// pool past it. + #[test] + fn test_write_blocks_to_recycles_up_to_the_cap() { + // More blocks than the cap, so both halves of the behaviour are visible. + let payload = vec![b'x'; BGZF_MAX_BLOCK_SIZE * (MAX_POOLED_BUFFERS + 2)]; + let mut compressor = InlineBgzfCompressor::new(6); + compressor.write_all(&payload).expect("writing data should succeed"); + compressor.flush().expect("flushing compressor should succeed"); + let queued = compressor.completed_blocks.len(); + assert!(queued > MAX_POOLED_BUFFERS, "need more than {MAX_POOLED_BUFFERS} blocks"); + assert!(compressor.buffer_pool.is_empty(), "pool starts empty"); + + let mut sink = Vec::new(); + compressor.write_blocks_to(&mut sink).expect("writing blocks should succeed"); + + assert!(compressor.completed_blocks.is_empty(), "every block should have been written"); + assert_eq!( + compressor.buffer_pool.len(), + MAX_POOLED_BUFFERS, + "draining {queued} blocks should fill the pool to the cap and drop the surplus" + ); + } + + /// A consumer that hands back more buffers than the compressor can use + /// must not grow the pool without limit; surplus buffers are dropped. + #[test] + fn test_recycle_buffer_bounds_the_pool() { + let (mut compressor, _blocks) = compressor_with_one_block(b"payload for one block"); + for _ in 0..32 { + compressor.recycle_buffer(vec![0u8; 1024]); + } + assert_eq!( + compressor.buffer_pool.len(), + MAX_POOLED_BUFFERS, + "pool should stop growing at the cap" + ); + } + + /// A sink that accepts `accept` bytes total and then fails, so a + /// multi-block `write_blocks_to` fails partway through. + struct ShortSink { + accept: usize, + } + + impl io::Write for ShortSink { + fn write(&mut self, buf: &[u8]) -> io::Result { + if buf.len() > self.accept { + return Err(io::Error::new(io::ErrorKind::WriteZero, "sink is full")); + } + self.accept -= buf.len(); + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + /// On a write failure, the unwritten blocks must stay queued. Draining them + /// into the void would lose compressed output with no way for the caller to + /// notice, since the error says nothing about how far the write got. + #[test] + fn test_write_blocks_to_keeps_unwritten_blocks_on_error() { + // Two full blocks, then a sink that has room for the first one only. + let payload = vec![b'x'; BGZF_MAX_BLOCK_SIZE * 2]; + let mut compressor = InlineBgzfCompressor::new(6); + compressor.write_all(&payload).expect("writing data should succeed"); + compressor.flush().expect("flushing compressor should succeed"); + let first_block_len = compressor.completed_blocks[0].data.len(); + let queued = compressor.completed_blocks.len(); + assert!(queued >= 2, "expected at least two blocks, got {queued}"); + + let mut sink = ShortSink { accept: first_block_len }; + let err = compressor + .write_blocks_to(&mut sink) + .expect_err("the sink should reject the second block"); + assert_eq!(err.kind(), io::ErrorKind::WriteZero); + // Identity, not just the count: a queue of the right length holding the + // wrong blocks (say, the tail restored but the failing block dropped) + // loses exactly the output this behaviour exists to preserve. Serials + // are contiguous from 0, so the retained queue must start at 1 and run + // to the end. + let retained: Vec = compressor.completed_blocks.iter().map(|b| b.serial).collect(); + let expected: Vec = (1..u64::try_from(queued).expect("block count fits")).collect(); + assert_eq!(retained, expected, "the failing block and its tail should still be queued"); + } } diff --git a/crates/fgumi-cli-common/Cargo.toml b/crates/fgumi-cli-common/Cargo.toml new file mode 100644 index 000000000..379ddf08b --- /dev/null +++ b/crates/fgumi-cli-common/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "fgumi-cli-common" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +description = "Shared CLI types and helpers for fgumi commands" +repository.workspace = true +license.workspace = true + +[dependencies] +anyhow = { workspace = true } +bytesize = { workspace = true } +clap = { workspace = true } +fgumi-fmt = { workspace = true } +log = { workspace = true } +sysinfo = { workspace = true } +thiserror = { workspace = true } + +[dev-dependencies] +rstest = { workspace = true } +tempfile = { workspace = true } + +[lints.clippy] +pedantic = { level = "deny", priority = -1 } diff --git a/crates/fgumi-cli-common/src/lib.rs b/crates/fgumi-cli-common/src/lib.rs new file mode 100644 index 000000000..04fe8dd83 --- /dev/null +++ b/crates/fgumi-cli-common/src/lib.rs @@ -0,0 +1,959 @@ +#![deny(unsafe_code)] +//! Shared CLI types and helpers for fgumi commands. +//! +//! Holds the pieces every command needs in common: the [`Command`] trait, the [`FgumiError`] +//! type, host-capacity detection, memory-budget parsing and resolution, and the clap argument +//! groups shared across subcommands. + +// ───────────────────────────────────────────────────────────────────────────── +// Command trait +// ───────────────────────────────────────────────────────────────────────────── + +/// Trait implemented by all fgumi CLI commands. +/// +/// Each command provides an `execute` method that runs the command's main logic. +/// The `command_line` parameter contains the full command invocation for @PG records. +/// +/// Commands are dispatched by a hand-written `match` in the binary rather than by +/// `#[enum_dispatch]`: that macro pairs a trait with its enum through a registry local to a +/// single proc-macro invocation, so it cannot link a trait defined here to an enum defined +/// in another crate. +pub trait Command { + #[allow(clippy::missing_errors_doc)] + fn execute(&self, command_line: &str) -> anyhow::Result<()>; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Error types +// ───────────────────────────────────────────────────────────────────────────── + +use thiserror::Error; + +/// Result type alias for fgumi operations (preferred in standalone crates). +pub type FgumiResult = std::result::Result; + +/// Unqualified result alias used by the umbrella crate's error/validation modules. +/// +/// Both `FgumiResult` and `Result` are the same type; the two names exist +/// so callers that shadow `std::result::Result` with `use crate::errors::Result` +/// (umbrella convention) resolve to the same `FgumiError`-based alias. +pub type Result = std::result::Result; + +/// Error type for fgumi operations +#[derive(Error, Debug)] +pub enum FgumiError { + /// Invalid parameter value provided + #[error("Invalid parameter '{parameter}': {reason}")] + InvalidParameter { + /// The parameter name + parameter: String, + /// Explanation of why it's invalid + reason: String, + }, + + /// Invalid frequency threshold + #[error("Invalid frequency threshold: {value} (must be between {min} and {max})")] + InvalidFrequency { + /// The invalid frequency value + value: f64, + /// Minimum valid value + min: f64, + /// Maximum valid value + max: f64, + }, + + /// Invalid quality threshold + #[error("Invalid quality threshold: {value} (must be between 0 and {max})")] + InvalidQuality { + /// The invalid quality value + value: u8, + /// Maximum valid value (usually 93 for SAM/BAM) + max: u8, + }, + + /// File format error + #[error("Invalid {file_type} file '{path}': {reason}")] + InvalidFileFormat { + /// Type of file (e.g., "BAM", "FASTQ") + file_type: String, + /// Path to the file + path: String, + /// Explanation of the problem + reason: String, + }, + + /// Required reference sequence not found + #[error("Reference sequence '{ref_name}' not found in header")] + ReferenceNotFound { + /// The reference sequence name + ref_name: String, + }, + + /// Invalid memory size string + #[error("Invalid memory size: {reason}")] + InvalidMemorySize { + /// Explanation of why the value is invalid + reason: String, + }, +} + +// ───────────────────────────────────────────────────────────────────────────── +// System detection +// ───────────────────────────────────────────────────────────────────────────── + +/// Returns the effective total memory available to this process in bytes. +/// +/// Checks cgroup memory limits for container environments, falling back to +/// physical RAM on bare-metal or macOS. Always returns `min(cgroup, physical)`. +/// +/// Call this once and reuse the result: it constructs a `sysinfo::System`, reads +/// `/proc/meminfo`, and probes the cgroup v1/v2 limit files, none of which is free. +/// [`resolve_memory_budget`] deliberately calls it exactly once. +#[must_use] +pub fn detect_total_memory() -> usize { + let mut system = sysinfo::System::new(); + system.refresh_memory(); + let physical = system.total_memory(); + let bytes = system.cgroup_limits().map_or(physical, |c| c.total_memory.min(physical)); + // Saturate at usize::MAX / 2 rather than usize::MAX on 32-bit platforms so + // the downstream `budget > total` overflow check in `resolve_memory_budget` + // can fire correctly (no value can exceed usize::MAX, so using it as the + // fallback renders the check dead). + usize::try_from(bytes).unwrap_or(usize::MAX / 2) +} + +/// Returns the number of logical CPUs available to this process. +/// +/// Honors cgroup CPU quotas (e.g. `--cpus` in Docker or Kubernetes resource limits), so a +/// container gets its quota rather than the host's core count. Returns 1 when the platform +/// cannot report parallelism. +#[must_use] +pub fn detect_cpu_count() -> usize { + std::thread::available_parallelism().map_or(1, std::num::NonZeroUsize::get) +} + +// ───────────────────────────────────────────────────────────────────────────── +// Formatting helpers +// ───────────────────────────────────────────────────────────────────────────── + +// Pure formatting with no CLI dependency, so it lives in the `fgumi-fmt` leaf crate that +// lower layers can reach too. Re-exported here for callers already depending on this crate. +pub use fgumi_fmt::{format_count, format_duration, format_rate}; + +// ───────────────────────────────────────────────────────────────────────────── +// Operation timer +// ───────────────────────────────────────────────────────────────────────────── + +/// Operation timing and summary helper. +/// +/// Tracks operation timing and provides formatted summary output. +pub struct OperationTimer { + operation: String, + start_time: std::time::Instant, +} + +impl OperationTimer { + /// Creates a new operation timer and logs the start. + #[must_use] + pub fn new(operation: &str) -> Self { + log::info!("{operation} ..."); + Self { operation: operation.to_string(), start_time: std::time::Instant::now() } + } + + /// Logs the completion with item count and rate. + pub fn log_completion(&self, count: u64) { + let duration = self.start_time.elapsed(); + log::info!( + "{} completed: {} in {} ({})", + self.operation, + format_count(count), + format_duration(duration), + format_rate(count, duration) + ); + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Validation helpers +// ───────────────────────────────────────────────────────────────────────────── + +use bytesize::ByteSize; +use std::path::Path; + +/// Validate that a file exists. +/// +/// # Errors +/// +/// Returns [`FgumiError::InvalidFileFormat`] if the file does not exist. +pub fn validate_file_exists>(path: P, description: &str) -> FgumiResult<()> { + let path_ref = path.as_ref(); + if !path_ref.exists() { + return Err(FgumiError::InvalidFileFormat { + file_type: description.to_string(), + path: path_ref.display().to_string(), + reason: "File does not exist".to_string(), + }); + } + Ok(()) +} + +/// Returns `true` if `s` contains a genuine decimal scientific-notation +/// mantissa-exponent (e.g. `1e5`, `1.5e10`, `2E+3`). +/// +/// Only an `e`/`E` that is preceded by a digit or `.` **and** followed by a +/// digit or a sign counts. This deliberately does *not* match the exabyte units +/// `EB`/`EiB` (where `E` is followed by `B`/`i`), which `bytesize` accepts — the +/// earlier `contains('e') || contains('E')` heuristic wrongly rejected them. +fn looks_like_scientific_notation(s: &str) -> bool { + let bytes = s.as_bytes(); + bytes.iter().enumerate().any(|(i, &b)| { + if b != b'e' && b != b'E' { + return false; + } + let prev_is_mantissa = i > 0 && (bytes[i - 1].is_ascii_digit() || bytes[i - 1] == b'.'); + let next_is_exponent = matches!( + bytes.get(i + 1), + Some(&c) if c.is_ascii_digit() || c == b'+' || c == b'-' + ); + prev_is_mantissa && next_is_exponent + }) +} + +/// Parses a memory size string into bytes. +/// +/// Accepts both plain numbers (interpreted as MiB) and human-readable formats like: +/// - "2GB", "2G" -> 2 gigabytes (decimal: 2,000,000,000) +/// - "1.5GB" -> 1.5 gigabytes +/// - "1024MB", "1024M" -> 1024 megabytes (decimal) +/// - "512MiB" -> 512 mebibytes (binary: 536,870,912) +/// - "1EB", "1EiB" -> 1 exabyte / exbibyte (decimal / binary; the largest units +/// accepted — values above `u64::MAX` saturate rather than wrapping) +/// - "768" -> 768 MiB (plain numbers are interpreted as mebibytes) +/// +/// # Errors +/// +/// Returns [`FgumiError::InvalidMemorySize`] if the string cannot be parsed. +pub fn parse_memory_size(size_str: &str) -> FgumiResult { + let trimmed = size_str.trim(); + if trimmed.is_empty() { + return Err(FgumiError::InvalidMemorySize { + reason: "Memory size cannot be empty".to_string(), + }); + } + + if trimmed.starts_with('-') { + return Err(FgumiError::InvalidMemorySize { + reason: format!("Memory size cannot be negative: '{trimmed}'"), + }); + } + + if let Ok(mb_value) = trimmed.parse::() { + if mb_value == 0 { + return Err(FgumiError::InvalidMemorySize { + reason: "Memory size cannot be zero".to_string(), + }); + } + if mb_value > 1_000_000 { + return Err(FgumiError::InvalidMemorySize { + reason: format!( + "Plain number memory size too large: {} MiB. Use human-readable format like '{}GiB' instead.", + mb_value, + mb_value / 1024 + ), + }); + } + return mb_value.checked_mul(1024 * 1024).ok_or_else(|| FgumiError::InvalidMemorySize { + reason: format!("Memory size calculation overflow for {mb_value} MiB"), + }); + } + + if looks_like_scientific_notation(trimmed) { + return Err(FgumiError::InvalidMemorySize { + reason: format!( + "Scientific notation not supported: '{trimmed}'. Use integer values or human-readable formats like '2GB'." + ), + }); + } + + if trimmed.contains('.') && trimmed.chars().all(|c| c.is_ascii_digit() || c == '.') { + return Err(FgumiError::InvalidMemorySize { + reason: format!( + "Plain decimal numbers not supported: '{trimmed}'. Use an integer for MiB (e.g. '768') or a human-readable format (e.g. '1.5GB')." + ), + }); + } + + match trimmed.parse::() { + Ok(size) => { + if size.0 == 0 { + return Err(FgumiError::InvalidMemorySize { + reason: format!("Memory size cannot be zero: '{trimmed}'"), + }); + } + Ok(size.0) + } + Err(_) => Err(FgumiError::InvalidMemorySize { + reason: format!( + "Invalid memory size '{trimmed}'. Valid formats:\n\ + - Plain numbers (interpreted as MiB): '768', '4096'\n\ + - Human-readable (decimal): '2GB', '1024MB'\n\ + - Human-readable (binary): '1GiB', '512MiB'" + ), + }), + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Memory/compression options +// ───────────────────────────────────────────────────────────────────────────── + +/// A memory limit, either auto-detected from the host or a fixed byte count. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MemoryLimit { + /// Detect the (cgroup-aware) host memory and subtract the reserve. + Auto, + /// Use a fixed memory limit in bytes. + Fixed(usize), +} + +impl Default for MemoryLimit { + fn default() -> Self { + Self::Fixed(768 * 1024 * 1024) + } +} + +impl std::fmt::Display for MemoryLimit { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Auto => f.write_str("auto"), + Self::Fixed(bytes) => format_binary_bytes(*bytes, f), + } + } +} + +/// How much memory to reserve for other processes when [`MemoryLimit::Auto`] is used. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum MemoryReserve { + /// Automatic: `min(10 GiB, 50% of host memory)`. + /// Matches the clap `default_value = "auto"` on `SortOptions::memory_reserve`. + #[default] + Auto, + /// Reserve a fixed number of bytes. + Fixed(usize), +} + +impl std::fmt::Display for MemoryReserve { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Auto => f.write_str("auto"), + Self::Fixed(bytes) => format_binary_bytes(*bytes, f), + } + } +} + +/// Format a byte count in the largest binary unit that divides cleanly. +/// +/// Emits the compact, unspaced form the CLI accepts as input (`768MiB`), so a default shown +/// in `--help` can be pasted straight back on the command line. `bytesize::ByteSize`'s own +/// `Display` also round-trips through [`parse_memory`], but renders as `768.0 MiB` — a space +/// and a redundant `.0`, which reads poorly as an argument default. +fn format_binary_bytes(bytes: usize, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + const K: usize = 1024; + const M: usize = K * 1024; + const G: usize = M * 1024; + if bytes >= G && bytes.is_multiple_of(G) { + write!(f, "{}GiB", bytes / G) + } else if bytes >= M && bytes.is_multiple_of(M) { + write!(f, "{}MiB", bytes / M) + } else if bytes >= K && bytes.is_multiple_of(K) { + write!(f, "{}KiB", bytes / K) + } else { + write!(f, "{bytes}B") + } +} + +/// The minimum per-thread memory budget (256 MiB). +pub const MIN_MEMORY_PER_THREAD: usize = 256 * 1024 * 1024; + +/// Default auto-reserve cap: 10 GiB. +/// +/// Computed with `saturating_mul` so it does not const-overflow `usize` on +/// 32-bit targets (where `usize == u32` cannot hold 10 GiB). There it saturates +/// to `usize::MAX`, so `resolve_reserve`'s `AUTO_RESERVE_CAP.min(total / 2)` +/// simply yields `total / 2` — consistent with `detect_total_memory`, which also +/// deliberately supports 32-bit by saturating. +const AUTO_RESERVE_CAP: usize = 10usize.saturating_mul(1024 * 1024 * 1024); + +/// Parse a memory-limit string (e.g. "512M", "1G", "768", "auto"). +/// +/// # Errors +/// +/// Returns an error string if parsing fails. +pub fn parse_memory(s: &str) -> std::result::Result { + let s = s.trim(); + if s.eq_ignore_ascii_case("auto") { + return Ok(MemoryLimit::Auto); + } + Ok(MemoryLimit::Fixed(parse_memory_bytes(s, "Memory size")?)) +} + +/// Parse a memory-reserve string (e.g. "10G", "auto"). +/// +/// # Errors +/// +/// Returns an error string if parsing fails. +pub fn parse_memory_reserve(s: &str) -> std::result::Result { + let s = s.trim(); + if s.eq_ignore_ascii_case("auto") { + return Ok(MemoryReserve::Auto); + } + Ok(MemoryReserve::Fixed(parse_memory_bytes(s, "Memory reserve")?)) +} + +/// Resolve a [`MemoryReserve`] to a concrete byte count given total host memory. +#[must_use] +pub fn resolve_reserve(reserve: MemoryReserve, total_memory: usize) -> usize { + match reserve { + MemoryReserve::Fixed(bytes) => bytes, + MemoryReserve::Auto => AUTO_RESERVE_CAP.min(total_memory / 2), + } +} + +/// Resolve a memory budget to a concrete byte count. +/// +/// # Errors +/// +/// Returns an error if `threads` is 0 or the multiplication overflows. +pub fn resolve_memory_budget( + limit: MemoryLimit, + reserve: MemoryReserve, + threads: usize, + per_thread: bool, +) -> anyhow::Result { + resolve_memory_budget_with_total(limit, reserve, threads, per_thread, detect_total_memory()) +} + +/// Pure resolver behind [`resolve_memory_budget`], with `total` injected for testability. +fn resolve_memory_budget_with_total( + limit: MemoryLimit, + reserve: MemoryReserve, + threads: usize, + per_thread: bool, + total: usize, +) -> anyhow::Result { + if threads == 0 { + anyhow::bail!("--threads must be at least 1"); + } + + let budget = match limit { + MemoryLimit::Fixed(bytes) => { + if per_thread { + bytes + .checked_mul(threads) + .ok_or_else(|| anyhow::anyhow!("memory limit × {threads} threads overflowed"))? + } else { + bytes + } + } + MemoryLimit::Auto => { + let margin = resolve_reserve(reserve, total); + let available = total.saturating_sub(margin); + let target = if per_thread { + (available / threads) + .max(MIN_MEMORY_PER_THREAD) + .checked_mul(threads) + .ok_or_else(|| anyhow::anyhow!("auto memory budget overflowed"))? + } else { + available.max(MIN_MEMORY_PER_THREAD) + }; + let budget = target.min(available); + if budget < target { + log::warn!( + "Auto memory: capping budget to host-available {} (minimum viable target {} \ + exceeds it after reserve {}); throughput may drop but the run stays within memory", + bytesize::ByteSize(budget as u64), + bytesize::ByteSize(target as u64), + bytesize::ByteSize(margin as u64), + ); + } + // Only the per-thread arm allocates `budget` as `threads` independent slices; + // otherwise it is one shared pool and `budget / threads` would misdescribe it. + if per_thread { + log::debug!( + "Auto memory: {} of {} ({}/thread × {} threads, reserve {})", + bytesize::ByteSize(budget as u64), + bytesize::ByteSize(total as u64), + bytesize::ByteSize((budget / threads) as u64), + threads, + bytesize::ByteSize(margin as u64), + ); + } else { + log::debug!( + "Auto memory: {} of {} (shared across {} threads, reserve {})", + bytesize::ByteSize(budget as u64), + bytesize::ByteSize(total as u64), + threads, + bytesize::ByteSize(margin as u64), + ); + } + budget + } + }; + + if budget > total { + log::warn!( + "Memory budget {} exceeds total host memory {}; this may cause OOM (or, for sort, earlier spill-to-disk)", + bytesize::ByteSize(budget as u64), + bytesize::ByteSize(total as u64), + ); + } + + Ok(budget) +} + +/// Parse a memory size string into `usize` bytes (private helper). +fn parse_memory_bytes(s: &str, label: &str) -> std::result::Result { + let bytes = parse_memory_size(s).map_err(|e| e.to_string())?; + usize::try_from(bytes).map_err(|_| format!("{label} too large: {bytes}")) +} + +/// Parses a boolean value from a string, accepting: true/false, yes/no, y/n, t/f +/// (case-insensitive). Matches sopt/fgbio behavior. +/// +/// # Errors +/// +/// Returns an error string if the input is not a recognized boolean. +pub fn parse_bool(s: &str) -> std::result::Result { + match s.to_ascii_lowercase().as_str() { + "true" | "t" | "yes" | "y" => Ok(true), + "false" | "f" | "no" | "n" => Ok(false), + _ => Err(format!("Invalid boolean value '{s}'. Expected: true|false|yes|no|y|n|t|f")), + } +} + +/// Options for output compression. +/// +/// Controls BGZF compression level for BAM output files. +#[derive(Debug, Clone, clap::Args)] +pub struct CompressionOptions { + /// Compression level for output BAM (0-12). + /// + /// Level 0 disables compression (uncompressed BGZF blocks). + /// Level 1 is fastest of the compressing levels with larger files; + /// level 12 produces the smallest files but is slowest. + #[arg(long, default_value_t = 1, value_parser = clap::value_parser!(u32).range(0..=12))] + pub compression_level: u32, +} + +impl Default for CompressionOptions { + /// Mirrors the clap `default_value_t = 1` so programmatic/default-constructed + /// callers emit level-1 compression rather than the `u32` default of `0` + /// (uncompressed BGZF). + fn default() -> Self { + Self { compression_level: 1 } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + use std::time::Duration; + + #[test] + fn test_detect_total_memory_nonzero() { + let total = detect_total_memory(); + assert!(total > 0, "expected non-zero total memory, got {total}"); + } + + #[test] + fn test_detect_cpu_count_at_least_one() { + assert!(detect_cpu_count() >= 1); + } + + // `format_count`, `format_duration` and `format_rate` are re-exports from `fgumi-fmt`, + // which owns their case tables and doctests. The `pub use` is compile-checked, so + // re-asserting their behavior here would only duplicate coverage. + + #[test] + fn test_parse_memory_size_plain() { + assert_eq!(parse_memory_size("768").unwrap(), 768 * 1024 * 1024); + } + + #[test] + fn test_parse_memory_size_human() { + assert_eq!(parse_memory_size("2GB").unwrap(), 2 * 1000 * 1000 * 1000); + assert_eq!(parse_memory_size("512MiB").unwrap(), 512 * 1024 * 1024); + } + + #[test] + fn test_parse_memory_size_errors() { + assert!(parse_memory_size("").is_err()); + assert!(parse_memory_size("-1").is_err()); + assert!(parse_memory_size("0").is_err()); + } + + /// Audit C3: the scientific-notation guard rejected any `e`/`E`, which also + /// rejected the exabyte units `EB`/`EiB` that `bytesize` accepts. It must + /// only fire for a genuine mantissa-exponent (a digit/`.`, then `e`/`E`, + /// then a digit or sign). + #[rstest] + #[case::exabyte("1EB", true)] + #[case::exbibyte("1EiB", true)] + #[case::exabyte_multi("2EB", true)] + #[case::gibibyte("512MiB", true)] + #[case::scientific_lower("1e5", false)] + #[case::scientific_upper("1E5", false)] + #[case::scientific_decimal("1.5e10", false)] + #[case::scientific_signed("2E+3", false)] + #[case::scientific_neg("3e-2", false)] + fn test_parse_memory_size_scientific_notation_guard(#[case] input: &str, #[case] ok: bool) { + assert_eq!( + parse_memory_size(input).is_ok(), + ok, + "parse_memory_size({input:?}) expected ok={ok}, got {:?}", + parse_memory_size(input) + ); + } + + #[test] + fn memory_limit_auto_displays_as_auto() { + assert_eq!(MemoryLimit::Auto.to_string(), "auto"); + } + + /// A rendered limit must survive a round trip through `parse_memory`, so a default shown + /// in `--help` can be pasted back on the command line unchanged. + #[rstest] + #[case::gibibytes(2 * 1024 * 1024 * 1024)] + #[case::mebibytes(768 * 1024 * 1024)] + #[case::kibibytes(4 * 1024)] + #[case::bare_bytes(100)] + #[case::not_a_clean_multiple(1536)] + fn memory_display_round_trips_through_parse_memory(#[case] bytes: usize) { + let rendered = MemoryLimit::Fixed(bytes).to_string(); + assert_eq!( + parse_memory(&rendered).unwrap(), + MemoryLimit::Fixed(bytes), + "{rendered} did not round-trip" + ); + } + + #[test] + fn test_resolve_reserve_auto() { + // Target-aware: a fixed 32 GiB `usize` literal cannot compile on 32-bit (`usize == u32`). + // Derive `total` from the cap so the test builds on every target, and assert the exact + // `resolve_reserve` formula (`AUTO_RESERVE_CAP.min(total / 2)`): on 64-bit the 10 GiB cap + // strictly binds (total/2 = 15 GiB > cap); on 32-bit the cap saturates to `usize::MAX`, so + // the expected value is `total / 2` — matching the documented 32-bit saturation behavior. + let total = AUTO_RESERVE_CAP.saturating_mul(3); + let r = resolve_reserve(MemoryReserve::Auto, total); + assert_eq!(r, AUTO_RESERVE_CAP.min(total / 2)); + } + + #[test] + fn test_resolve_memory_budget_fixed() { + let budget = resolve_memory_budget( + MemoryLimit::Fixed(512 * 1024 * 1024), + MemoryReserve::Auto, + 4, + true, + ) + .unwrap(); + assert_eq!(budget, 4 * 512 * 1024 * 1024); + } + + #[rstest] + #[case::true_word("true", Some(true))] + #[case::t_short("t", Some(true))] + #[case::yes_word("yes", Some(true))] + #[case::y_short("y", Some(true))] + #[case::false_word("false", Some(false))] + #[case::f_short("f", Some(false))] + #[case::no_word("no", Some(false))] + #[case::n_short("n", Some(false))] + #[case::mixed_case("TrUe", Some(true))] + #[case::garbage("maybe", None)] + #[case::empty("", None)] + fn parse_bool_accepts_the_documented_spellings( + #[case] input: &str, + #[case] expected: Option, + ) { + assert_eq!(parse_bool(input).ok(), expected); + } + + #[derive(clap::Parser)] + struct CompressionHarness { + #[command(flatten)] + compression: CompressionOptions, + } + + // In-range values (the 0-12 bounds; 0 = uncompressed) parse. + #[rstest] + #[case::uncompressed(0_u32)] + #[case::fastest(1_u32)] + #[case::midpoint(6_u32)] + #[case::maximum(12_u32)] + fn test_compression_level_accepts_in_range(#[case] level: u32) { + use clap::Parser; + + assert_eq!( + CompressionHarness::try_parse_from(["prog", "--compression-level", &level.to_string()]) + .unwrap() + .compression + .compression_level, + level + ); + } + + // The programmatic `Default` must match the clap `default_value_t = 1`; a + // derived `Default` would silently yield level 0 (uncompressed BGZF). + #[test] + fn test_compression_options_default_matches_cli_default() { + assert_eq!(CompressionOptions::default().compression_level, 1); + } + + #[test] + fn test_compression_level_default_and_rejects_out_of_range() { + use clap::Parser; + + // Default is 1 when the flag is omitted. + assert_eq!(CompressionHarness::parse_from(["prog"]).compression.compression_level, 1); + + // Out-of-range values are rejected at parse time rather than silently accepted. + assert!(CompressionHarness::try_parse_from(["prog", "--compression-level", "13"]).is_err()); + assert!(CompressionHarness::try_parse_from(["prog", "--compression-level", "99"]).is_err()); + } + + #[test] + fn test_resolve_memory_budget_auto_low_available() { + // When available memory per thread is below MIN_MEMORY_PER_THREAD, the + // budget is floored to MIN_MEMORY_PER_THREAD × threads, then capped at + // available (which is less), so the result equals available. + let total = 2 * MIN_MEMORY_PER_THREAD; // very tight: only 2 × floor per 4 threads + let reserve = 0; + let budget = resolve_memory_budget_with_total( + MemoryLimit::Auto, + MemoryReserve::Fixed(reserve), + 4, + true, + total, + ) + .unwrap(); + // available = total - 0 = 2×MIN; per-thread = 2×MIN/4 < MIN → floored to MIN; + // target = MIN × 4 = 4×MIN > available → capped at available. + assert_eq!(budget, total); + } + + #[test] + fn test_resolve_memory_budget_auto_margin_exceeds_total() { + // When the reserve margin >= total, saturating_sub → 0 available. + // Budget is then floored to MIN_MEMORY_PER_THREAD, capped at 0 + // (available), so result == 0 (the cap wins). + let total = 512 * 1024 * 1024_usize; // 512 MiB + let margin = total + 1; // margin exceeds total + let budget = resolve_memory_budget_with_total( + MemoryLimit::Auto, + MemoryReserve::Fixed(margin), + 1, + false, + total, + ) + .unwrap(); + // available = total.saturating_sub(margin) = 0; target = max(0, MIN) = MIN; + // budget = MIN.min(0) = 0. + assert_eq!(budget, 0); + } + + #[test] + fn test_resolve_memory_budget_threads_zero_bails() { + // `threads == 0` is an invalid input and must bail before any + // per-thread arithmetic. + let err = resolve_memory_budget_with_total( + MemoryLimit::Fixed(512 * 1024 * 1024), + MemoryReserve::Auto, + 0, + true, + 32 * 1024 * 1024 * 1024, + ) + .expect_err("threads == 0 must be rejected"); + let msg = format!("{err:#}"); + assert!(msg.contains("--threads must be at least 1"), "got: {msg}"); + } + + #[test] + fn test_resolve_memory_budget_fixed_per_thread_overflow() { + // Fixed(usize::MAX) × 2 threads overflows the `checked_mul`, which must + // surface an error rather than wrap. + let err = resolve_memory_budget_with_total( + MemoryLimit::Fixed(usize::MAX), + MemoryReserve::Auto, + 2, + true, + 32 * 1024 * 1024 * 1024, + ) + .expect_err("Fixed per-thread multiplication must overflow"); + let msg = format!("{err:#}"); + assert!(msg.contains("overflowed"), "got: {msg}"); + } + + #[test] + fn test_resolve_memory_budget_auto_per_thread_overflow() { + // Construct an Auto config whose per-thread floor × threads overflows: + // with `total == usize::MAX` and a tiny reserve, `available` is enormous, + // so `(available / threads).max(MIN) * threads` overflows `usize`. + let err = resolve_memory_budget_with_total( + MemoryLimit::Auto, + MemoryReserve::Fixed(0), + usize::MAX, + true, + usize::MAX, + ) + .expect_err("Auto per-thread multiplication must overflow"); + let msg = format!("{err:#}"); + assert!(msg.contains("auto memory budget overflowed"), "got: {msg}"); + } + + /// `format_rate` has three regimes: a sub-millisecond guard that avoids dividing by ~zero, + /// a normal items/s path, and an items/min path once the rate drops below one per second. + #[rstest] + #[case::sub_millisecond_guard(5, Duration::from_micros(100), "5 items/s")] + #[case::exactly_one_per_second(1, Duration::from_secs(1), "1 items/s")] + #[case::thousands_per_second(1000, Duration::from_secs(1), "1,000 items/s")] + #[case::below_one_per_second(30, Duration::from_secs(60), "30.0 items/min")] + #[case::far_below_one_per_second(1, Duration::from_secs(120), "0.5 items/min")] + fn format_rate_switches_units_below_one_per_second( + #[case] count: u64, + #[case] duration: Duration, + #[case] expected: &str, + ) { + assert_eq!(format_rate(count, duration), expected); + } + + #[test] + fn operation_timer_reports_completion_without_panicking() { + let timer = OperationTimer::new("test operation"); + timer.log_completion(1234); + assert!(timer.start_time.elapsed() >= Duration::ZERO); + assert_eq!(timer.operation, "test operation"); + } + + #[test] + fn validate_file_exists_accepts_a_real_file() { + let file = tempfile::NamedTempFile::new().expect("failed to create temp file"); + validate_file_exists(file.path(), "input BAM").expect("existing file must validate"); + } + + #[test] + fn validate_file_exists_rejects_a_missing_path() { + let dir = tempfile::tempdir().expect("failed to create temp dir"); + let missing = dir.path().join("definitely-not-here.bam"); + let err = validate_file_exists(&missing, "input BAM") + .expect_err("a missing path must not validate"); + let msg = err.to_string(); + assert!(msg.contains("input BAM"), "error should name the description; got: {msg}"); + assert!(msg.contains("File does not exist"), "got: {msg}"); + } + + /// The error arms of `parse_memory_size` that the happy-path tests do not reach. Each input + /// is rejected for a different documented reason, so the assertions pin the reason text. + #[rstest] + #[case::plain_number_too_large("2000000", "Plain number memory size too large")] + #[case::plain_decimal("1.5", "Plain decimal numbers not supported")] + #[case::human_readable_zero("0B", "Memory size cannot be zero")] + #[case::unparseable("not-a-size", "Invalid memory size")] + fn parse_memory_size_rejects_with_a_specific_reason( + #[case] input: &str, + #[case] expected_reason: &str, + ) { + let err = parse_memory_size(input).expect_err("input must be rejected"); + let msg = err.to_string(); + assert!(msg.contains(expected_reason), "expected {expected_reason:?}, got: {msg}"); + } + + #[test] + fn memory_limit_default_is_768_mib() { + assert_eq!(MemoryLimit::default(), MemoryLimit::Fixed(768 * 1024 * 1024)); + } + + /// `format_binary_bytes` picks the largest binary unit that divides cleanly, falling back to + /// bytes. It backs the `Display` impls of both `MemoryLimit` and `MemoryReserve`. + #[rstest] + #[case::gibibytes(2 * 1024 * 1024 * 1024, "2GiB")] + #[case::mebibytes(512 * 1024 * 1024, "512MiB")] + #[case::kibibytes(4 * 1024, "4KiB")] + #[case::bare_bytes(100, "100B")] + #[case::not_a_clean_multiple(1536, "1536B")] + fn memory_display_uses_the_largest_clean_binary_unit( + #[case] bytes: usize, + #[case] expected: &str, + ) { + assert_eq!(MemoryLimit::Fixed(bytes).to_string(), expected); + assert_eq!(MemoryReserve::Fixed(bytes).to_string(), expected); + } + + #[test] + fn memory_reserve_auto_displays_as_auto() { + assert_eq!(MemoryReserve::Auto.to_string(), "auto"); + assert_eq!(MemoryReserve::default(), MemoryReserve::Auto); + } + + #[rstest] + #[case::lowercase_auto("auto", MemoryLimit::Auto)] + #[case::uppercase_auto("AUTO", MemoryLimit::Auto)] + #[case::surrounding_whitespace(" auto ", MemoryLimit::Auto)] + #[case::human_readable("512MiB", MemoryLimit::Fixed(512 * 1024 * 1024))] + #[case::plain_number_is_mib("768", MemoryLimit::Fixed(768 * 1024 * 1024))] + fn parse_memory_accepts_auto_and_sizes(#[case] input: &str, #[case] expected: MemoryLimit) { + assert_eq!(parse_memory(input).unwrap(), expected); + } + + #[rstest] + #[case::lowercase_auto("auto", MemoryReserve::Auto)] + #[case::uppercase_auto("AUTO", MemoryReserve::Auto)] + #[case::human_readable("10GiB", MemoryReserve::Fixed(10 * 1024 * 1024 * 1024))] + fn parse_memory_reserve_accepts_auto_and_sizes( + #[case] input: &str, + #[case] expected: MemoryReserve, + ) { + assert_eq!(parse_memory_reserve(input).unwrap(), expected); + } + + #[test] + fn parse_memory_surfaces_the_underlying_parse_error() { + let err = parse_memory("not-a-size").expect_err("garbage must not parse"); + assert!(err.contains("Invalid memory size"), "got: {err}"); + let err = parse_memory_reserve("-1").expect_err("negative must not parse"); + assert!(err.contains("cannot be negative"), "got: {err}"); + } + + #[test] + fn resolve_memory_budget_fixed_without_per_thread_ignores_thread_count() { + // The non-per-thread Fixed arm returns the limit verbatim, regardless of `threads`. + let budget = resolve_memory_budget_with_total( + MemoryLimit::Fixed(512 * 1024 * 1024), + MemoryReserve::Auto, + 8, + false, + 16 * 1024 * 1024 * 1024, + ) + .unwrap(); + assert_eq!(budget, 512 * 1024 * 1024); + } + + #[test] + fn resolve_memory_budget_allows_a_limit_above_total_host_memory() { + // Over-committing is a warning, not an error: the caller may know better (e.g. swap), and + // for sort it just means spilling to disk earlier. + let total = 1024 * 1024 * 1024; + let budget = resolve_memory_budget_with_total( + MemoryLimit::Fixed(4 * 1024 * 1024 * 1024), + MemoryReserve::Auto, + 1, + false, + total, + ) + .unwrap(); + assert_eq!(budget, 4 * 1024 * 1024 * 1024); + assert!(budget > total, "this case must exercise the over-commit path"); + } +} diff --git a/crates/fgumi-cli-macros/Cargo.toml b/crates/fgumi-cli-macros/Cargo.toml new file mode 100644 index 000000000..29d862d72 --- /dev/null +++ b/crates/fgumi-cli-macros/Cargo.toml @@ -0,0 +1,27 @@ +[package] +name = "fgumi-cli-macros" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +description = "Proc-macro support crate for fgumi CLI option re-exposure (multi_options)" +repository.workspace = true +license.workspace = true + +[lib] +proc-macro = true + +[dependencies] +# `syn` is NOT centralized: this crate pins the 2.x line, while `xtask` pins 3.x +# with `extra-traits`. A single workspace entry cannot represent both majors. +syn = { version = "2", features = ["full"] } +quote = { workspace = true } +proc-macro2 = "1" + +[dev-dependencies] +clap = { workspace = true } +anyhow = { workspace = true } +rstest = { workspace = true } +trybuild = { workspace = true } + +[lints.clippy] +pedantic = { level = "deny", priority = -1 } diff --git a/crates/fgumi-cli-macros/src/lib.rs b/crates/fgumi-cli-macros/src/lib.rs new file mode 100644 index 000000000..01b294d26 --- /dev/null +++ b/crates/fgumi-cli-macros/src/lib.rs @@ -0,0 +1,1610 @@ +#![deny(unsafe_code)] + +//! Proc-macro support for fgumi CLI option re-exposure. +//! +//! Provides [`multi_options`] — an attribute macro that pairs a standalone +//! command's `clap::Args` options struct with a sibling `Multi` struct +//! whose CLI flags are prefixed and grouped under a help heading. Used by the +//! `runall` command to re-expose every per-stage option of `fgumi sort` / +//! `fgumi group` / `fgumi simplex` / `fgumi duplex` / `fgumi codec` without +//! hand-maintaining a parallel option set on `RunAll`. +//! +//! # Flag naming +//! +//! A re-exposed flag is named `--::` (e.g. `--sort::max-memory`). +//! The `::` separator is deliberate: `runall` flattens several stages into one +//! command, and every stage independently owns flags like `--threads` or +//! `--max-memory`. A plain kebab prefix (`--sort-max-memory`) would be +//! ambiguous with a stage that genuinely has a `--max` flag taking a `memory` +//! value, and would read as one flag name rather than as a stage-qualified one. +//! `::` cannot appear in a clap-derived flag name, so it can never collide with +//! a real flag, and it makes the stage qualifier obvious in `--help`. +//! +//! # Field-kind handling +//! +//! * `#[arg(skip)]` / `#[arg(skip = expr)]` fields are invisible to the CLI. +//! They are carried onto the Multi struct as skip fields too, so their values +//! survive a `From` + `validate()` round trip; when the Multi struct is +//! parsed from the command line they take the skip expression's value, or the +//! original struct's `Default` for a bare `skip`. +//! * `Option`, `Vec` and bare `bool` fields keep their type: clap already +//! treats each as absent-able (`None`, empty, `false`). +//! * Fields with any `default_value*` attribute keep their type and their +//! default — the attribute is copied verbatim, so the prefixed flag advertises +//! and applies exactly the default the standalone command does. +//! * Every other field is required: it becomes `Option` on the Multi struct, +//! and the generated `validate()` returns an error when the corresponding +//! `--::` is missing. This keeps required-ness staged — clap +//! never refuses to parse, so `runall` can report which stage is missing what. +//! +//! # What is and is not carried over +//! +//! Copied verbatim onto the generated field: `#[doc]` comments (each attribute +//! separately, so clap's short/long help split survives), `#[cfg]` and +//! `#[cfg_attr]` gates, every `default_value*`, and every other `#[arg(...)]` +//! key the macro does not classify (`value_parser`, `value_name`, `action`, +//! `num_args`, `hide`, `env`, …). A `#[cfg]` additionally gates the generated +//! conversion arms, so a field compiled out is absent from all three sites; +//! `#[cfg_attr]` never removes a field, so it lands on the field alone. +//! +//! Rewritten: `long` (re-prefixed) and long-form aliases (`alias`, +//! `aliases`, `visible_alias`, `visible_aliases` — each re-prefixed the same +//! way, so no un-namespaced flag ever reaches the parent command). +//! +//! Dropped: `short` and every short alias — one letter cannot be namespaced per +//! stage. A field-level `help_heading` is dropped for the same reason: the +//! companion files every field under its stage's heading, and a field-level one +//! is emitted after it and wins, so the field would escape its stage — and two +//! stages declaring the same heading would merge into one section naming +//! neither. `required` is also dropped from the clap side and enforced by +//! `validate()` instead, so a missing value is reported by `runall`'s staged +//! validation rather than by clap's parser. +//! +//! `help` and `long_help` are preserved verbatim, like any other unclassified +//! key: they are the field's own documentation and read identically on both +//! commands. One consequence is that an explicit `help` on a required field +//! replaces the generated "Required when `` is selected." line, since +//! clap prefers `help` over `#[doc]`; the staged `validate()` error still names +//! both the flag and the stage. +//! +//! Rejected at expansion time, with a spanned build error: cross-field +//! reference keys (`requires`, `conflicts_with`, …) whose arg ids would dangle +//! once fields are prefixed; `id` / `name` overrides; positional arguments, +//! whether declared by `index` or by declaring neither `long` nor `short` +//! (several stages' positionals would be mutually ambiguous, and clap panics +//! when a positional carries a `long`); clap's `key(value)` call form for any +//! key the macro classifies; a clap attribute hidden behind `#[cfg_attr]`, +//! which the macro cannot classify through; a `required` that could never be +//! enforced, because the field always holds a value (defaulted — including a +//! defaulted `Option` or `Vec` — skipped, or a bare `bool`); +//! `#[command(...)]` on a field; struct-level `#[command(...)]` / +//! `#[group(...)]`; generic structs, since the companion and its conversions +//! are emitted without generic parameters; and the legacy `#[clap(...)]` / +//! `#[structopt(...)]` spellings, which every classifier would otherwise +//! silently ignore. +//! +//! Note that `env` is preserved verbatim: two stages re-exposing the same +//! options struct therefore read the same environment variable, exactly as the +//! two standalone commands would. + +use proc_macro::TokenStream; +use proc_macro2::TokenStream as TokenStream2; +use quote::{format_ident, quote}; +use syn::{Lit, Meta}; + +/// Attribute macro that passes the original options struct through unchanged and +/// generates a `Multi` companion whose flags are named +/// `--::` and filed under `` in `--help`, plus a +/// `validate()` method, a `TryFrom>` impl and a +/// `From` impl. +/// +/// Usage: `#[multi_options("prefix", "Help Heading")]` +/// +/// The generated items inherit the annotated struct's visibility. +/// +/// See the crate-level docs for the field-kind classification rules and for the +/// full list of what is carried over, rewritten, dropped and rejected. +/// +/// # Errors +/// +/// Emits a spanned compile error — pointing at the offending field, attribute or +/// literal — rather than panicking. See the crate docs' "What is and is not +/// carried over" section for the rejected forms. +#[proc_macro_attribute] +pub fn multi_options(attr: TokenStream, item: TokenStream) -> TokenStream { + let item = TokenStream2::from(item); + match expand(TokenStream2::from(attr), &item) { + Ok(expanded) => expanded.into(), + Err(error) => { + // Re-emit the annotated item next to the diagnostic. Without it, every + // downstream reference to the struct raises its own "cannot find type" + // error and buries the one that actually explains the problem. + let compile_error = error.into_compile_error(); + quote! { #item #compile_error }.into() + } + } +} + +/// Expand `#[multi_options(...)]`, returning either the generated tokens or the +/// accumulated diagnostics. +fn expand(attr: TokenStream2, item: &TokenStream2) -> syn::Result { + let args = syn::parse2::(attr)?; + let input = parse_annotated_struct(item)?; + let fields = named_fields(&input)?; + + let context = ExpandContext { + struct_name: &input.ident, + vis: &input.vis, + prefix: &args.prefix, + heading: &args.heading, + }; + let generated = generate_fields(fields, &context)?; + Ok(render(&input, &generated, &context)) +} + +/// Parse the annotated item, rejecting anything that is not a struct. +fn parse_annotated_struct(item: &TokenStream2) -> syn::Result { + let parsed = syn::parse2::(item.clone())?; + let syn::Item::Struct(input) = parsed else { + return Err(syn::Error::new_spanned( + parsed, + "multi_options only supports structs with named fields", + )); + }; + reject_unsupported_struct_attrs(&input.attrs)?; + // The companion struct and both conversion impls are emitted without generic + // parameters, so a generic options struct expands into code that cannot + // compile — and every resulting error names the type parameter rather than + // this macro. + if !input.generics.params.is_empty() { + return Err(syn::Error::new_spanned(&input.generics, GENERICS_MSG)); + } + // `Generics`' `ToTokens` emits only the `<...>` params, so spanning on it + // when they are empty would collapse the span to the call site and lose the + // field-accurate diagnostics every other rejection here produces. + if let Some(where_clause) = &input.generics.where_clause { + return Err(syn::Error::new_spanned(where_clause, GENERICS_MSG)); + } + Ok(input) +} + +/// Borrow the struct's named fields, rejecting tuple and unit structs. +fn named_fields( + input: &syn::ItemStruct, +) -> syn::Result<&syn::punctuated::Punctuated> { + match &input.fields { + syn::Fields::Named(named) => Ok(&named.named), + other => Err(syn::Error::new_spanned( + other, + "multi_options only supports structs with named fields", + )), + } +} + +/// Classify and generate every field, reporting all bad fields at once rather +/// than making the author fix them one build at a time. +fn generate_fields( + fields: &syn::punctuated::Punctuated, + context: &ExpandContext<'_>, +) -> syn::Result> { + let mut errors: Option = None; + let mut generated = Vec::with_capacity(fields.len()); + for field in fields { + match ParsedField::from_field(field, context.prefix) { + Ok(parsed) => generated.push(parsed.generate(context)), + Err(error) => match &mut errors { + Some(accumulated) => accumulated.combine(error), + None => errors = Some(error), + }, + } + } + match errors { + Some(error) => Err(error), + None => Ok(generated), + } +} + +/// Assemble the original struct, the generated companion, and the conversions. +fn render( + input: &syn::ItemStruct, + generated: &[GeneratedField], + context: &ExpandContext<'_>, +) -> TokenStream2 { + let struct_name = context.struct_name; + let multi_struct_name = format_ident!("Multi{}", struct_name); + let multi_fields = generated.iter().map(|g| &g.multi_field); + let validate_arms = generated.iter().map(|g| &g.validate_arm); + let from_arms = generated.iter().map(|g| &g.from_arm); + let vis = context.vis; + + quote! { + #input + + /// Prefixed options struct generated by `#[multi_options]` for the + /// runall command. Carries the same fields as the original options + /// struct but exposes each via `--::`, filed under the + /// stage's help heading. + #[derive(::clap::Args, Debug, Clone)] + // clap adopts a flattened struct's doc comment as the parent command's + // description when the parent declares none. This struct's rustdoc is + // for docs.rs (and for crates that deny(missing_docs)) — it would be + // nonsense as a command description, and with several stages flattened + // clap would arbitrarily pick whichever came first. + #[command(about = None, long_about = None)] + #vis struct #multi_struct_name { + #(#multi_fields)* + } + + impl #multi_struct_name { + /// Validate required fields and convert to the original options + /// struct. Returns `Err` naming the missing `--::` + /// when a field the standalone command requires was not supplied. + #vis fn validate(self) -> ::anyhow::Result<#struct_name> { + <#struct_name as ::core::convert::TryFrom>::try_from(self) + } + } + + impl ::core::convert::TryFrom<#multi_struct_name> for #struct_name { + type Error = ::anyhow::Error; + + fn try_from(opts: #multi_struct_name) -> ::anyhow::Result { + Ok(Self { + #(#validate_arms)* + }) + } + } + + impl ::core::convert::From<#struct_name> for #multi_struct_name { + fn from(opts: #struct_name) -> Self { + Self { + #(#from_arms)* + } + } + } + } +} + +/// Everything the per-field generator needs from the annotated struct. +struct ExpandContext<'a> { + struct_name: &'a syn::Ident, + vis: &'a syn::Visibility, + prefix: &'a str, + heading: &'a str, +} + +// ───────────────────────────────────────────────────────────────────────────── +// Macro arguments +// ───────────────────────────────────────────────────────────────────────────── + +/// Parsed arguments for `#[multi_options("prefix", "heading")]`. +struct MultiOptionsArgs { + prefix: String, + heading: String, +} + +/// The one-line usage reminder attached to every argument-shape diagnostic. +const USAGE: &str = "multi_options requires two string literal arguments: \ + #[multi_options(\"prefix\", \"heading\")]"; + +/// Rejection message for a generic annotated struct, shared by the type-parameter +/// and `where`-clause arms so the two cannot drift apart. +const GENERICS_MSG: &str = "multi_options does not support generic structs: the generated Multi struct and its \ + conversions are emitted without generic parameters, so the expansion would not compile. \ + Use a concrete options struct."; + +impl syn::parse::Parse for MultiOptionsArgs { + fn parse(input: syn::parse::ParseStream) -> syn::Result { + let prefix_lit: syn::LitStr = + input.parse().map_err(|e| syn::Error::new(e.span(), USAGE))?; + input.parse::().map_err(|e| syn::Error::new(e.span(), USAGE))?; + let heading_lit: syn::LitStr = + input.parse().map_err(|e| syn::Error::new(e.span(), USAGE))?; + if !input.is_empty() { + return Err(syn::Error::new(input.span(), USAGE)); + } + + let prefix = prefix_lit.value(); + validate_prefix(&prefix).map_err(|msg| syn::Error::new(prefix_lit.span(), msg))?; + let heading = heading_lit.value(); + if heading.is_empty() { + return Err(syn::Error::new( + heading_lit.span(), + "multi_options: the help heading must not be empty", + )); + } + + Ok(Self { prefix, heading }) + } +} + +/// Check that a prefix can be spliced into both a flag name and an identifier. +/// +/// The prefix appears in the flag (`--::`) and in the generated +/// field identifier (`_`), so an empty or non-identifier prefix +/// would otherwise surface as `--::flag` or as an opaque `format_ident!` panic +/// with no mention of `multi_options`. +fn validate_prefix(prefix: &str) -> Result<(), String> { + let Some(first) = prefix.chars().next() else { + return Err("multi_options: the prefix must not be empty (it would generate flags named \ + `--::`)" + .to_string()); + }; + if !first.is_ascii_alphabetic() { + return Err(format!( + "multi_options: the prefix must start with an ASCII letter, got `{prefix}` — it also \ + becomes the leading segment of the generated field identifier `_`" + )); + } + if let Some(bad) = + prefix.chars().find(|c| !(c.is_ascii_alphanumeric() || *c == '_' || *c == '-')) + { + return Err(format!( + "multi_options: the prefix may only contain ASCII letters, digits, `_` and `-`, but \ + `{prefix}` contains `{bad}`" + )); + } + Ok(()) +} + +/// Reject struct-level clap configuration the macro cannot faithfully reproduce. +/// +/// The generated struct is built from the fields alone, so any struct-level clap +/// setting would apply to the standalone command and silently not to the +/// re-exposed one. +fn reject_unsupported_struct_attrs(attrs: &[syn::Attribute]) -> syn::Result<()> { + for attr in attrs { + let path = attr.path(); + // A struct-level `#[cfg_attr(...)]` needs no special handling here: rustc + // expands a *container* cfg_attr before this attribute macro runs, so an + // active predicate arrives as a literal `#[command(...)]`/`#[clap(...)]` + // (caught below) and an inactive one expands to nothing (correctly inert). + // This differs from a *field* cfg_attr, which reaches the macro verbatim + // as body tokens — hence the separate `reject_conditional_clap_attr` there. + if path.is_ident("clap") || path.is_ident("structopt") { + return Err(legacy_spelling_error(attr)); + } + if path.is_ident("group") { + return Err(syn::Error::new_spanned( + attr, + "multi_options does not support a struct-level #[group(...)]: the group names its \ + members by their unprefixed arg ids, which do not exist on the generated Multi \ + struct. Enforce the grouping in the command's validate()/resolve() instead.", + )); + } + if path.is_ident("command") { + return Err(syn::Error::new_spanned( + attr, + "multi_options does not carry a struct-level #[command(...)] onto the generated \ + Multi struct, so the standalone and re-exposed commands would silently diverge. \ + Move the setting onto the individual #[arg(...)] attributes, or drop it.", + )); + } + } + Ok(()) +} + +/// Reject a `#[cfg_attr(...)]` that hides a clap attribute behind a condition. +/// +/// Every classifier keys on a literal `#[arg(...)]` / `#[command(...)]`, so a +/// `#[cfg_attr(unix, arg(long, default_value_t = 3))]` is invisible: the field is +/// classified as required, wrapped in `Option`, and *then* the forwarded +/// attribute applies `default_value_t` to the wrapped type — a wall of type +/// errors that never mentions `multi_options`. Classifying through the condition +/// is not possible either, since the macro cannot evaluate `cfg` predicates. +fn reject_conditional_clap_attr(attr: &syn::Attribute, field: &syn::Field) -> syn::Result<()> { + let Ok(metas) = + attr.parse_args_with(syn::punctuated::Punctuated::::parse_terminated) + else { + // Not a shape we can inspect; leave it to rustc. + return Ok(()); + }; + // The first meta is the `cfg` predicate; the rest are the attributes it gates. + for meta in metas.iter().skip(1) { + let path = meta.path(); + if path.is_ident("arg") || path.is_ident("clap") || path.is_ident("command") { + let key = path.get_ident().map_or_else(|| "arg".to_string(), ToString::to_string); + return Err(syn::Error::new_spanned( + attr, + format!( + "multi_options: field `{}` hides a clap attribute behind #[cfg_attr(…, \ + {key}(…))]. The macro classifies fields from their literal #[arg(...)] \ + attributes and cannot evaluate a cfg predicate, so this one would be ignored \ + and the field misclassified. Apply #[cfg] to the field and write the \ + #[arg(...)] unconditionally.", + field_name(field) + ), + )); + } + } + Ok(()) +} + +/// Build the diagnostic for clap's legacy attribute spellings. +/// +/// Every classifier in this macro keys on `#[arg(...)]` / `#[command(...)]`, so a +/// `#[clap(skip)]` would be invisible — the field would be exposed as a required +/// CLI flag instead of being skipped. Reject rather than silently misclassify. +fn legacy_spelling_error(attr: &syn::Attribute) -> syn::Error { + syn::Error::new_spanned( + attr, + "multi_options does not support the legacy #[clap(...)] / #[structopt(...)] spelling: \ + every classifier keys on #[arg(...)] and #[command(...)], so this attribute would be \ + silently ignored and the field misclassified. Use the #[arg(...)] / #[command(...)] \ + spelling.", + ) +} + +// ───────────────────────────────────────────────────────────────────────────── +// clap `#[arg(...)]` key tables +// ───────────────────────────────────────────────────────────────────────────── + +/// clap `#[arg(...)]` keys that reference *another argument by its string id*. +/// +/// The Multi struct renames every field to `_` and never emits a +/// matching arg alias, so any of these ids would dangle on the Multi side and +/// clap would panic ("arg id `x` not defined") when it builds the `runall` +/// command. Reject them at macro-expansion time instead (D1). +const CROSS_REFERENCE_ARG_KEYS: &[&str] = &[ + "requires", + "requires_all", + "requires_if", + "requires_ifs", + "conflicts_with", + "conflicts_with_all", + "overrides_with", + "overrides_with_all", + "required_if_eq", + "required_if_eq_all", + "required_if_eq_any", + "required_unless_present", + "required_unless_present_any", + "required_unless_present_all", + "default_value_if", + "default_value_ifs", + "default_values_if", + "default_values_ifs", + "group", + "groups", +]; + +/// clap `#[arg(...)]` keys that rename the argument itself. +/// +/// The Multi struct derives every arg id from its own prefixed field name, so an +/// explicit id would reintroduce the unprefixed name and collide across stages. +const RENAMING_ARG_KEYS: &[&str] = &["id", "name"]; + +/// clap `#[arg(...)]` keys that declare the argument to be positional. +/// +/// A positional argument has no flag name to prefix, and clap panics outright +/// when one is given a `long` — which is exactly what the Multi struct emits. +const POSITIONAL_ARG_KEYS: &[&str] = &["index"]; + +/// clap `#[arg(...)]` keys that declare a default value, in every spelling. +/// +/// A field carrying any of these is optional on the standalone command, so it +/// must not be classified as required. Deliberately excludes +/// `default_missing_value`, which is the value used when a flag is passed +/// *without* one and says nothing about whether the flag may be omitted. +const DEFAULT_VALUE_ARG_KEYS: &[&str] = &[ + "default_value", + "default_value_t", + "default_value_os", + "default_value_os_t", + "default_values", + "default_values_t", + "default_values_os", + "default_values_os_t", +]; + +/// clap `#[arg(...)]` keys this macro classifies via `Meta::NameValue` / +/// `Meta::Path` (to strip, rewrite, or read them). +/// +/// clap also accepts the equivalent `key(value)` call form, which arrives as a +/// `Meta::List` and would slip past every classifier — silently changing the +/// flag name, the required-ness, the default, or leaking an un-prefixed alias. +/// Reject the call form for these keys (D2). +const CALL_FORM_SENSITIVE_ARG_KEYS: &[&str] = &[ + "long", + "short", + "required", + "skip", + "alias", + "aliases", + "visible_alias", + "visible_aliases", + "short_alias", + "short_aliases", + "visible_short_alias", + "visible_short_aliases", + "help_heading", + "default_value", + "default_value_t", + "default_value_os", + "default_value_os_t", + "default_values", + "default_values_t", + "default_values_os", + "default_values_os_t", +]; + +/// Long-form alias keys, which are re-prefixed exactly like `long`. +const LONG_ALIAS_ARG_KEYS: &[&str] = &["alias", "visible_alias"]; + +/// Long-form alias keys taking a list of aliases; every element is re-prefixed. +const LONG_ALIAS_LIST_ARG_KEYS: &[&str] = &["aliases", "visible_aliases"]; + +/// Short-form alias keys, dropped for the same reason as `short`. +const SHORT_ALIAS_ARG_KEYS: &[&str] = + &["short_alias", "short_aliases", "visible_short_alias", "visible_short_aliases"]; + +// ───────────────────────────────────────────────────────────────────────────── +// Per-field parsing +// ───────────────────────────────────────────────────────────────────────────── + +/// How a `#[arg(skip)]` field obtains its value. +enum Skip { + /// Bare `#[arg(skip)]` — falls back to the original struct's `Default`. + Bare, + /// `#[arg(skip = expr)]` — uses the expression verbatim. + Expr(syn::Expr), +} + +/// One field of the annotated struct, classified in a single pass over its +/// attributes. +/// +/// Every `#[arg(...)]` meta is visited exactly once and routed to the piece of +/// state it affects, so the key tables above are the single source of truth for +/// classification — earlier revisions re-walked `field.attrs` in six independent +/// helpers, which let the three key lists drift apart. +struct ParsedField<'a> { + ident: &'a syn::Ident, + ty: &'a syn::Type, + /// `#[doc]`, `#[cfg]` and `#[cfg_attr]` attributes, copied onto the + /// generated field verbatim. + forwarded: Vec<&'a syn::Attribute>, + /// The `#[cfg]` subset, which must additionally gate the `TryFrom` and + /// `From` arms so a gated-out field is absent from all three sites. + /// + /// `#[cfg_attr]` is deliberately excluded: it never removes a field, so the + /// arms do not need it, and only `#[cfg]` is valid on a struct-expression + /// field — forwarding `#[cfg_attr(unix, serde(skip))]` onto one would expand + /// to `#[serde(skip)]` on an expression, which does not compile. + cfgs: Vec<&'a syn::Attribute>, + /// `#[arg(...)]` entries carried onto the generated field, already rendered + /// as `, key = value` continuations ready to splice into a new `#[arg(...)]`. + preserved: Vec, + skip: Option, + /// The value of an explicit `#[arg(long = "...")]` override. + long_override: Option, + has_default: bool, + /// Whether the field carries `#[arg(required)]` / `#[arg(required = true)]`. + required: bool, + /// Whether the field declares a flag name — `long` or `short`, in any + /// spelling. A field declaring neither is positional to clap. + has_flag_name: bool, +} + +impl<'a> ParsedField<'a> { + /// Classify one field, or return every diagnostic it earns. + fn from_field(field: &'a syn::Field, prefix: &str) -> syn::Result { + let ident = field.ident.as_ref().ok_or_else(|| { + syn::Error::new_spanned(field, "multi_options only supports structs with named fields") + })?; + let mut parsed = Self { + ident, + ty: &field.ty, + forwarded: Vec::new(), + cfgs: Vec::new(), + preserved: Vec::new(), + skip: None, + long_override: None, + has_default: false, + required: false, + has_flag_name: false, + }; + + for attr in &field.attrs { + let path = attr.path(); + if path.is_ident("doc") { + parsed.forwarded.push(attr); + } else if path.is_ident("cfg") { + parsed.forwarded.push(attr); + parsed.cfgs.push(attr); + } else if path.is_ident("cfg_attr") { + reject_conditional_clap_attr(attr, field)?; + parsed.forwarded.push(attr); + } else if path.is_ident("clap") || path.is_ident("structopt") { + return Err(legacy_spelling_error(attr)); + } else if path.is_ident("command") { + return Err(field_command_error(attr, field)); + } else if path.is_ident("arg") { + for meta in parse_arg_metas(attr, field)? { + parsed.absorb_arg_meta(&meta, field, prefix)?; + } + } + } + + parsed.check_required_is_enforceable(field)?; + parsed.check_is_not_positional(field)?; + Ok(parsed) + } + + /// Reject a field clap would treat as a positional argument. + /// + /// In clap's derive a field with neither `long` nor `short` is positional. + /// The Multi struct always emits a `long`, which would silently convert + /// `fgumi sort ` into `--sort::input ` — and, when the field + /// carries an explicit `index`, makes clap panic outright while building the + /// runall command. A positional cannot be namespaced per stage anyway: with + /// several stages flattened into one command, their positionals would be + /// mutually ambiguous. + fn check_is_not_positional(&self, field: &syn::Field) -> syn::Result<()> { + if self.skip.is_some() || self.has_flag_name { + return Ok(()); + } + Err(syn::Error::new_spanned( + field, + format!( + "multi_options: field `{}` declares neither `long` nor `short`, so clap treats it \ + as a positional argument. A positional has no flag name to prefix, and several \ + stages flattened into one runall command would have mutually ambiguous \ + positionals. Add an explicit `#[arg(long)]` (or `#[arg(long = \"...\")]`) to the \ + field.", + field_name(field) + ), + )) + } + + /// Route one `#[arg(...)]` meta to the state it affects. + fn absorb_arg_meta( + &mut self, + meta: &Meta, + field: &syn::Field, + prefix: &str, + ) -> syn::Result<()> { + let Some(key) = meta.path().get_ident().map(ToString::to_string) else { + self.preserved.push(quote! { , #meta }); + return Ok(()); + }; + let key = key.as_str(); + let name = field_name(field); + reject_unsupported_arg_key(meta, field, &name, key)?; + + match key { + "skip" => { + self.skip = Some(match meta { + Meta::NameValue(nv) => Skip::Expr(nv.value.clone()), + _ => Skip::Bare, + }); + } + // The Multi field declares its own prefixed `long`; a bare `long` + // just means "kebab of the field name", which is the fallback. + "long" => { + self.has_flag_name = true; + if let Meta::NameValue(nv) = meta { + self.long_override = Some(string_literal(&nv.value).ok_or_else(|| { + non_literal_error(field, &name, "long", "a string literal") + })?); + } + } + // One letter cannot be namespaced per stage. + "short" => self.has_flag_name = true, + _ if SHORT_ALIAS_ARG_KEYS.contains(&key) => {} + // Neither can a help heading. The companion files every field under + // the stage's heading; a field-level one is emitted after it and + // wins, so the field escapes its stage — and two stages declaring + // the same heading would merge into one section that names neither. + "help_heading" => {} + // Required-ness is enforced by validate(), not by clap, so runall + // can name the stage that is missing a value. + "required" => { + self.required = match meta { + Meta::Path(_) => true, + Meta::NameValue(nv) => bool_literal(&nv.value).ok_or_else(|| { + non_literal_error(field, &name, "required", "`true` or `false`") + })?, + Meta::List(_) => unreachable!("call form rejected above"), + }; + } + _ if LONG_ALIAS_ARG_KEYS.contains(&key) => { + self.preserved.push(prefixed_alias(meta, field, &name, key, prefix)?); + } + _ if LONG_ALIAS_LIST_ARG_KEYS.contains(&key) => { + self.preserved.push(prefixed_alias_list(meta, field, &name, key, prefix)?); + } + _ => { + if DEFAULT_VALUE_ARG_KEYS.contains(&key) { + self.has_default = true; + } + self.preserved.push(quote! { , #meta }); + } + } + Ok(()) + } + + /// Reject `#[arg(required)]` on a field whose generated form always holds a + /// value. + /// + /// `validate()` enforces required-ness by observing absence — a `None`, an + /// empty `Vec`, or the `Option` wrapper the macro adds. A skipped field, a + /// defaulted field and a bare `bool` are never absent, so the requirement + /// would be silently unenforceable. + /// + /// `Option` and `Vec` are absent-able only while they carry no default: + /// clap applies a declared `default_value` to those types too (an omitted + /// `Option` with `default_value = "42"` parses as `Some(42)`, not + /// `None`), so the generated `is_none()` / `is_empty()` check would never + /// fire and the requirement would be lost in silence. + fn check_required_is_enforceable(&self, field: &syn::Field) -> syn::Result<()> { + if !self.required { + return Ok(()); + } + let name = field_name(field); + if self.skip.is_some() { + return Err(syn::Error::new_spanned( + field, + format!( + "multi_options: field `{name}` combines #[arg(required …)] with #[arg(skip)]. \ + A skipped field is never supplied on the command line, so the requirement \ + could never be satisfied." + ), + )); + } + let absent_able = (is_option_type(self.ty) || is_vec_type(self.ty)) && !self.has_default; + if self.is_required_kind() || absent_able { + return Ok(()); + } + Err(syn::Error::new_spanned( + field, + format!( + "multi_options: field `{name}` combines #[arg(required …)] with a default value or \ + a bare `bool`. The generated field always holds a value — clap applies a declared \ + default to `Option` and `Vec` as well — so the Multi side cannot distinguish \ + \"not supplied\" from \"supplied the default\" and the requirement would be \ + silently unenforceable. Drop `required`, or drop the default." + ), + )) + } + + /// Whether the field must be wrapped in `Option` on the Multi struct and + /// enforced by `validate()`. + /// + /// `Option`, `Vec` and bare `bool` are all absent-able as-is, and a + /// field with a default is never missing. + fn is_required_kind(&self) -> bool { + !self.has_default + && !is_option_type(self.ty) + && !is_vec_type(self.ty) + && !is_bool_type(self.ty) + } + + /// Emit the Multi-struct field plus the matching `TryFrom` and `From` arms. + fn generate(&self, ctx: &ExpandContext<'_>) -> GeneratedField { + let field_ident = self.ident; + let field_type = self.ty; + let prefixed_ident = format_ident!("{}_{}", ctx.prefix.replace('-', "_"), field_ident); + let forwarded = &self.forwarded; + let cfgs = &self.cfgs; + let preserved = &self.preserved; + let vis = ctx.vis; + + // Skipped fields are invisible to the CLI on both sides. Carrying them + // onto the Multi struct as skip fields keeps `From` + `validate()` a + // genuine round trip instead of re-deriving the value from `Default`. + if let Some(skip) = &self.skip { + let value = match skip { + Skip::Expr(expr) => quote! { #expr }, + Skip::Bare => { + let struct_name = ctx.struct_name; + quote! { #struct_name::default().#field_ident } + } + }; + return GeneratedField { + multi_field: quote! { + #(#forwarded)* + #[arg(skip = #value)] + #vis #prefixed_ident: #field_type, + }, + validate_arm: quote! { #(#cfgs)* #field_ident: opts.#prefixed_ident, }, + from_arm: quote! { #(#cfgs)* #prefixed_ident: opts.#field_ident, }, + }; + } + + // Honor an explicit `#[arg(long = "...")]` so the prefixed flag tracks the + // standalone command's flag name (`tmp_dirs` with `long = "tmp-dir"` + // becomes `--::tmp-dir`, not `--::tmp-dirs`). + let base_name = + self.long_override.clone().unwrap_or_else(|| field_ident.to_string().replace('_', "-")); + let long_name = format!("{}::{}", ctx.prefix, base_name); + let heading = ctx.heading; + let prefix = ctx.prefix; + let error_msg = format!("--{long_name} is required when {prefix} is selected"); + + if self.is_required_kind() { + let required_doc = format!("Required when {prefix} is selected."); + return GeneratedField { + multi_field: quote! { + #(#forwarded)* + // Blank doc line so clap keeps the required notice in long help + // (`--help`) rather than merging it into the field's short help, + // which is the first paragraph of the doc comment. + #[doc = ""] + #[doc = #required_doc] + #[arg(long = #long_name, help_heading = #heading #(#preserved)*)] + #vis #prefixed_ident: Option<#field_type>, + }, + validate_arm: quote! { + #(#cfgs)* + #field_ident: opts.#prefixed_ident + .ok_or_else(|| ::anyhow::anyhow!(#error_msg))?, + }, + from_arm: quote! { #(#cfgs)* #prefixed_ident: Some(opts.#field_ident), }, + }; + } + + // Absent-able or defaulted: the field keeps its type and its declared + // default. `required` on such a field is dropped from the clap side, so + // re-assert it here — otherwise a flag the standalone command demands + // would be quietly optional on the Multi side. + let value = if self.required && is_option_type(field_type) { + quote! {{ + let value = opts.#prefixed_ident; + if value.is_none() { + ::anyhow::bail!(#error_msg); + } + value + }} + } else if self.required && is_vec_type(field_type) { + quote! {{ + let value = opts.#prefixed_ident; + if value.is_empty() { + ::anyhow::bail!(#error_msg); + } + value + }} + } else { + quote! { opts.#prefixed_ident } + }; + + GeneratedField { + multi_field: quote! { + #(#forwarded)* + #[arg(long = #long_name, help_heading = #heading #(#preserved)*)] + #vis #prefixed_ident: #field_type, + }, + validate_arm: quote! { #(#cfgs)* #field_ident: #value, }, + from_arm: quote! { #(#cfgs)* #prefixed_ident: opts.#field_ident, }, + } + } +} + +/// The three token streams one source field contributes to the expansion. +struct GeneratedField { + multi_field: TokenStream2, + validate_arm: TokenStream2, + from_arm: TokenStream2, +} + +// ───────────────────────────────────────────────────────────────────────────── +// Attribute helpers +// ───────────────────────────────────────────────────────────────────────────── + +/// Render a field's name for a diagnostic. +/// +/// `multi_options` only accepts named-field structs, so the `None` arm is +/// unreachable in practice; it exists so a diagnostic never panics while +/// reporting another diagnostic. +fn field_name(field: &syn::Field) -> String { + field.ident.as_ref().map_or_else(|| "".to_string(), ToString::to_string) +} + +/// Parse the comma-separated metas out of one `#[arg(...)]` attribute. +/// +/// An unparseable `#[arg(...)]` is a hard build error: silently treating it as +/// "no attributes" would misclassify the field — a dropped `skip` exposes a +/// field that should have no CLI flag, a dropped `long` renames it, a dropped +/// `default_value` makes it required. +fn parse_arg_metas( + attr: &syn::Attribute, + field: &syn::Field, +) -> syn::Result> { + attr.parse_args_with(syn::punctuated::Punctuated::::parse_terminated) + .map_err(|e| { + syn::Error::new_spanned( + attr, + format!( + "multi_options: failed to parse #[arg(...)] on field `{}`: {e}", + field_name(field) + ), + ) + }) +} + +/// Build the diagnostic for a field-level `#[command(...)]`. +/// +/// `flatten` and `subcommand` nest another struct whose fields the macro cannot +/// reach to prefix; any other field-level `#[command(...)]` key would simply not +/// be carried over. +fn field_command_error(attr: &syn::Attribute, field: &syn::Field) -> syn::Error { + let name = field_name(field); + let metas = attr + .parse_args_with(syn::punctuated::Punctuated::::parse_terminated) + .ok(); + let nests = metas.is_some_and(|metas| { + metas + .iter() + .any(|meta| meta.path().is_ident("flatten") || meta.path().is_ident("subcommand")) + }); + if nests { + syn::Error::new_spanned( + attr, + format!( + "multi_options does not support #[command(flatten)] / #[command(subcommand)] on \ + field `{name}`. The nested struct's fields cannot be reached to prefix them; \ + inline them directly." + ), + ) + } else { + syn::Error::new_spanned( + attr, + format!( + "multi_options does not carry a field-level #[command(...)] onto the generated \ + Multi struct, so field `{name}` would behave differently on the two commands. \ + Move the setting onto #[arg(...)], or drop it." + ), + ) + } +} + +/// Reject an `#[arg(...)]` key the macro cannot faithfully re-expose. +/// +/// Covers three latent traps: cross-field reference keys whose arg-id strings +/// dangle once the field is prefixed (D1), explicit id overrides that would +/// reintroduce the un-prefixed name, and clap's `key(value)` call form for any +/// key the macro only recognizes in `key = value` / bare-`key` form (D2). +fn reject_unsupported_arg_key( + meta: &Meta, + field: &syn::Field, + name: &str, + key: &str, +) -> syn::Result<()> { + if CROSS_REFERENCE_ARG_KEYS.contains(&key) { + return Err(syn::Error::new_spanned( + field, + format!( + "multi_options: field `{name}` uses #[arg({key} …)], which references another \ + argument by its unprefixed id. The Multi struct renames fields to \ + `_`, so that id would dangle and clap would panic when it builds \ + the runall command. Enforce this coupling in the command's validate()/resolve() \ + instead (see AlignerOptions::resolve)." + ), + )); + } + if RENAMING_ARG_KEYS.contains(&key) { + return Err(syn::Error::new_spanned( + field, + format!( + "multi_options: field `{name}` uses #[arg({key} …)], which overrides the \ + argument's id. The Multi struct derives every id from its own prefixed field \ + name, so an explicit id would reintroduce the un-prefixed name and collide \ + between stages. Drop it." + ), + )); + } + if POSITIONAL_ARG_KEYS.contains(&key) { + return Err(syn::Error::new_spanned( + field, + format!( + "multi_options: field `{name}` uses #[arg({key} …)], which makes it a positional \ + argument. The Multi struct gives every field a prefixed `long`, and clap panics \ + when a positional has one (\"is a positional argument and can't have short or \ + long name versions\"). A positional cannot be namespaced per stage — expose it \ + as a flag with `#[arg(long)]` instead." + ), + )); + } + if matches!(meta, Meta::List(_)) && CALL_FORM_SENSITIVE_ARG_KEYS.contains(&key) { + return Err(syn::Error::new_spanned( + field, + format!( + "multi_options: field `{name}` uses the call form #[arg({key}(…))]. Use the \ + `{key} = …` name-value form (or bare `{key}`) — the macro only classifies those \ + spellings and would mishandle the call form." + ), + )); + } + Ok(()) +} + +/// Re-prefix a single long alias so it can never reach the parent command +/// un-namespaced. +fn prefixed_alias( + meta: &Meta, + field: &syn::Field, + name: &str, + key: &str, + prefix: &str, +) -> syn::Result { + let Meta::NameValue(nv) = meta else { + return Err(non_literal_error(field, name, key, "a string literal")); + }; + let alias = string_literal(&nv.value) + .ok_or_else(|| non_literal_error(field, name, key, "a string literal"))?; + let prefixed = format!("{prefix}::{alias}"); + let key = format_ident!("{}", key); + Ok(quote! { , #key = #prefixed }) +} + +/// Re-prefix every element of an alias list (`aliases = ["a", "b"]`). +fn prefixed_alias_list( + meta: &Meta, + field: &syn::Field, + name: &str, + key: &str, + prefix: &str, +) -> syn::Result { + const EXPECTED: &str = "an array of string literals"; + let Meta::NameValue(nv) = meta else { + return Err(non_literal_error(field, name, key, EXPECTED)); + }; + let syn::Expr::Array(array) = &nv.value else { + return Err(non_literal_error(field, name, key, EXPECTED)); + }; + let prefixed = array + .elems + .iter() + .map(|elem| { + string_literal(elem) + .map(|alias| format!("{prefix}::{alias}")) + .ok_or_else(|| non_literal_error(field, name, key, EXPECTED)) + }) + .collect::>>()?; + let key = format_ident!("{}", key); + Ok(quote! { , #key = [#(#prefixed),*] }) +} + +/// Build the diagnostic for a classified key whose value is not the literal form +/// the macro can read. +fn non_literal_error(field: &syn::Field, name: &str, key: &str, expected: &str) -> syn::Error { + syn::Error::new_spanned( + field, + format!( + "multi_options: field `{name}` uses #[arg({key} = …)] with a value that is not \ + {expected}. The macro reads this key at expansion time — to rewrite the flag name, \ + re-prefix the alias, or classify required-ness — so it cannot be an arbitrary \ + expression." + ), + ) +} + +/// Read a string literal out of an attribute value. +fn string_literal(expr: &syn::Expr) -> Option { + if let syn::Expr::Lit(expr_lit) = expr + && let Lit::Str(lit) = &expr_lit.lit + { + return Some(lit.value()); + } + None +} + +/// Read a boolean literal out of an attribute value. +fn bool_literal(expr: &syn::Expr) -> Option { + if let syn::Expr::Lit(expr_lit) = expr + && let Lit::Bool(lit) = &expr_lit.lit + { + return Some(lit.value); + } + None +} + +// ───────────────────────────────────────────────────────────────────────────── +// Type predicates +// ───────────────────────────────────────────────────────────────────────────── + +/// Check whether a type is `Vec`. +/// +/// clap collects a `Vec` naturally (empty when the flag is absent), so it +/// needs no default to be absent-able. +fn is_vec_type(ty: &syn::Type) -> bool { + if let syn::Type::Path(type_path) = ty + && let Some(segment) = type_path.path.segments.last() + { + return segment.ident == "Vec"; + } + false +} + +/// Check whether a type is a bare `bool`. +/// +/// clap gives a bare `bool` field `ArgAction::SetTrue` — a valueless flag +/// defaulting to `false` — so it needs no `default_value` to be absent-able. +fn is_bool_type(ty: &syn::Type) -> bool { + if let syn::Type::Path(type_path) = ty + && let Some(segment) = type_path.path.segments.last() + { + return segment.ident == "bool" && segment.arguments.is_none(); + } + false +} + +/// Check whether a type is `Option`. +fn is_option_type(ty: &syn::Type) -> bool { + if let syn::Type::Path(type_path) = ty + && let Some(segment) = type_path.path.segments.last() + { + return segment.ident == "Option"; + } + false +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + use syn::parse::Parser; + + /// Parse a single named struct field from tokens (e.g. + /// `#[arg(long, requires = "x")] pub y: u32`). + fn named_field(tokens: TokenStream2) -> syn::Field { + syn::Field::parse_named.parse2(tokens).expect("parse named field") + } + + /// Classify a field the way `expand` does, returning the diagnostic message + /// on rejection. + fn classify(tokens: TokenStream2) -> Result<(), String> { + let field = named_field(tokens); + ParsedField::from_field(&field, "p").map(|_| ()).map_err(|e| e.to_string()) + } + + /// Every `#[arg(...)]` spelling the macro classifies must be accepted in its + /// bare and name-value forms. + #[rstest] + #[case::short_and_string_default(quote! { #[arg(long, short = 'x', default_value = "7")] pub a: u32 })] + #[case::long_override_and_typed_default( + quote! { #[arg(long = "max-memory", default_value_t = 5)] pub b: usize } + )] + #[case::bare_long_on_option(quote! { #[arg(long)] pub c: Option })] + #[case::vec_with_required( + quote! { #[arg(long, value_delimiter = ',', required = true)] pub d: Vec } + )] + #[case::short_only(quote! { #[arg(short = 'x')] pub e: u32 })] + #[case::bare_skip(quote! { #[arg(skip)] pub f: u32 })] + #[case::skip_with_expr(quote! { #[arg(skip = 7u32)] pub g: u32 })] + #[case::long_alias(quote! { #[arg(long, alias = "ref")] pub h: u32 })] + #[case::alias_list(quote! { #[arg(long, aliases = ["ref", "fasta"])] pub i: u32 })] + #[case::short_alias(quote! { #[arg(long, short_alias = 'r')] pub j: u32 })] + #[case::os_default(quote! { #[arg(long, default_value_os = "/tmp/x")] pub k: PathBuf })] + #[case::doc_and_cfg(quote! { /// docs + #[cfg(unix)] + #[arg(long)] pub l: Option })] + fn accepts_supported_arg_forms(#[case] tokens: TokenStream2) { + assert_eq!(classify(tokens), Ok(()), "expected supported form to be accepted"); + } + + /// D1: every key that names another argument by id would dangle once the + /// field is prefixed. The macro classifies on the key alone, so the uniform + /// `key = "other"` spelling exercises each one. + #[rstest] + fn rejects_every_cross_reference_key( + #[values( + "requires", + "requires_all", + "requires_if", + "requires_ifs", + "conflicts_with", + "conflicts_with_all", + "overrides_with", + "overrides_with_all", + "required_if_eq", + "required_if_eq_all", + "required_if_eq_any", + "required_unless_present", + "required_unless_present_any", + "required_unless_present_all", + "default_value_if", + "default_value_ifs", + "default_values_if", + "default_values_ifs", + "group", + "groups" + )] + key: &str, + ) { + let key_ident = format_ident!("{}", key); + let err = classify(quote! { #[arg(long, #key_ident = "other")] pub a: u32 }) + .expect_err("cross-reference key must be rejected"); + assert!(err.contains(key), "message should name the key {key:?}: {err}"); + assert!(err.contains("field `a`"), "message should name the field: {err}"); + } + + /// Exhaustive by construction: drives the constant itself, so a key added to + /// `CROSS_REFERENCE_ARG_KEYS` is covered without touching the table above. + #[test] + fn every_cross_reference_key_in_the_table_is_rejected() { + for key in CROSS_REFERENCE_ARG_KEYS { + let key_ident = format_ident!("{}", key); + let err = classify(quote! { #[arg(long, #key_ident = "other")] pub a: u32 }) + .expect_err("every cross-reference key must be rejected"); + assert!(err.contains(key), "message should name the key {key:?}: {err}"); + } + } + + /// Every key the classifier reads or rewrites must also be rejected in + /// clap's `key(value)` call form — that form arrives as a `Meta::List` and + /// matches no classifier arm, so a key missing from + /// `CALL_FORM_SENSITIVE_ARG_KEYS` is silently preserved verbatim and the + /// rewrite is skipped. Adding a spelling to any table below without adding + /// it here reopens the D2 hole for that key. + #[test] + fn every_classified_key_is_call_form_sensitive() { + let classified = ["skip", "long", "short", "required", "help_heading"] + .iter() + .copied() + .chain(DEFAULT_VALUE_ARG_KEYS.iter().copied()) + .chain(LONG_ALIAS_ARG_KEYS.iter().copied()) + .chain(LONG_ALIAS_LIST_ARG_KEYS.iter().copied()) + .chain(SHORT_ALIAS_ARG_KEYS.iter().copied()); + for key in classified { + assert!( + CALL_FORM_SENSITIVE_ARG_KEYS.contains(&key), + "classified key {key:?} is missing from CALL_FORM_SENSITIVE_ARG_KEYS, so its \ + `{key}(…)` call form would slip past the classifier" + ); + } + } + + /// Every `default_value*` spelling must mark the field as defaulted; a + /// missed one is silently reclassified as required and gains a spurious + /// `Option` wrapper plus a false "Required when …" help line. + #[test] + fn every_default_value_spelling_makes_a_field_non_required() { + for key in DEFAULT_VALUE_ARG_KEYS { + let key_ident = format_ident!("{}", key); + let field = named_field(quote! { #[arg(long, #key_ident = "x")] pub a: u32 }); + let parsed = ParsedField::from_field(&field, "p") + .unwrap_or_else(|e| panic!("{key} should classify: {e}")); + assert!(parsed.has_default, "{key} must count as a default"); + assert!(!parsed.is_required_kind(), "{key} must not leave the field required"); + } + } + + /// A positional argument has no flag name to prefix, and clap panics when one + /// carries a `long` — which the companion always emits. + #[rstest] + #[case::no_attributes(quote! { pub a: u32 }, "positional")] + #[case::arg_without_long_or_short(quote! { #[arg(value_name = "FILE")] pub b: u32 }, "positional")] + #[case::explicit_index(quote! { #[arg(index = 1)] pub c: u32 }, "index")] + fn rejects_positional_fields(#[case] tokens: TokenStream2, #[case] needle: &str) { + let err = classify(tokens).expect_err("a positional field must be rejected"); + assert!(err.contains(needle), "message should mention {needle:?}: {err}"); + } + + /// A skipped field is never a CLI argument, so it needs no flag name. + #[test] + fn skipped_fields_are_exempt_from_the_positional_check() { + assert_eq!(classify(quote! { #[arg(skip)] pub a: u32 }), Ok(())); + } + + /// The macro cannot evaluate a `cfg` predicate, so a clap attribute behind + /// one would be ignored and the field misclassified. + #[rstest] + #[case::arg(quote! { #[cfg_attr(unix, arg(long, default_value_t = 3))] pub a: u32 })] + #[case::command(quote! { #[cfg_attr(unix, command(flatten))] pub b: u32 })] + #[case::legacy_clap(quote! { #[cfg_attr(unix, clap(long))] pub c: u32 })] + fn rejects_clap_attributes_hidden_behind_cfg_attr(#[case] tokens: TokenStream2) { + let err = classify(tokens).expect_err("a conditional clap attribute must be rejected"); + assert!(err.contains("cfg_attr"), "message should name the attribute: {err}"); + } + + /// A `#[cfg_attr]` carrying no clap attribute is forwarded untouched. + #[test] + fn allows_cfg_attr_without_a_clap_attribute() { + assert_eq!( + classify(quote! { #[cfg_attr(unix, doc = "unix only")] #[arg(long)] pub a: u32 }), + Ok(()) + ); + } + + /// D2: the `key(value)` call form arrives as `Meta::List` and would slip past + /// the `Meta::NameValue` / `Meta::Path` classifiers. + #[rstest] + #[case::long(quote! { #[arg(long("x"))] pub a: u32 }, "long")] + #[case::short(quote! { #[arg(long, short('x'))] pub b: u32 }, "short")] + #[case::default_value_t(quote! { #[arg(default_value_t(4))] pub c: u32 }, "default_value_t")] + #[case::required(quote! { #[arg(long, required(true))] pub d: u32 }, "required")] + #[case::alias(quote! { #[arg(long, alias("ref"))] pub e: u32 }, "alias")] + #[case::visible_alias( + quote! { #[arg(long, visible_alias("ref"))] pub f: u32 }, + "visible_alias" + )] + #[case::default_value_os( + quote! { #[arg(long, default_value_os("/tmp/x"))] pub g: PathBuf }, + "default_value_os" + )] + fn rejects_call_form_for_classified_keys(#[case] tokens: TokenStream2, #[case] needle: &str) { + let err = classify(tokens).expect_err("call form must be rejected"); + assert!(err.contains(needle), "message should name the key {needle:?}: {err}"); + } + + #[test] + fn allows_call_form_for_unclassified_keys() { + // `value_parser(...)` in call form is preserved verbatim and is not one of + // the classified keys, so it must not be rejected. + assert_eq!( + classify(quote! { #[arg(long, value_parser(clap::value_parser!(u32)))] pub a: u32 }), + Ok(()) + ); + } + + /// Keys the macro reads at expansion time cannot be arbitrary expressions. + #[rstest] + #[case::long(quote! { #[arg(long = SOME_CONST)] pub a: u32 }, "long")] + #[case::alias(quote! { #[arg(long, alias = SOME_CONST)] pub b: u32 }, "alias")] + #[case::aliases(quote! { #[arg(long, aliases = SOME_CONST)] pub c: u32 }, "aliases")] + #[case::aliases_of_non_literals( + quote! { #[arg(long, aliases = [SOME_CONST])] pub d: u32 }, + "aliases" + )] + #[case::required(quote! { #[arg(long, required = SOME_CONST)] pub e: u32 }, "required")] + fn rejects_non_literal_values_for_classified_keys( + #[case] tokens: TokenStream2, + #[case] needle: &str, + ) { + let err = classify(tokens).expect_err("non-literal value must be rejected"); + assert!(err.contains(needle), "message should name the key {needle:?}: {err}"); + } + + /// An explicit arg id would reintroduce the un-prefixed name. + #[rstest] + #[case::id(quote! { #[arg(long, id = "other")] pub a: u32 }, "id")] + #[case::name(quote! { #[arg(long, name = "other")] pub b: u32 }, "name")] + fn rejects_renaming_keys(#[case] tokens: TokenStream2, #[case] needle: &str) { + let err = classify(tokens).expect_err("renaming key must be rejected"); + assert!(err.contains(needle), "message should name the key {needle:?}: {err}"); + } + + /// The legacy spellings bypass every classifier, so they must not be ignored. + #[rstest] + #[case::clap_skip(quote! { #[clap(skip)] pub a: u32 })] + #[case::clap_long(quote! { #[clap(long, default_value = "3")] pub b: u32 })] + #[case::structopt(quote! { #[structopt(long)] pub c: u32 })] + fn rejects_legacy_attribute_spellings(#[case] tokens: TokenStream2) { + let err = classify(tokens).expect_err("legacy spelling must be rejected"); + assert!(err.contains("#[clap(...)]"), "message should name the spelling: {err}"); + } + + /// A field-level `#[command(...)]` either nests a struct the macro cannot + /// reach, or would silently not be carried over. + #[rstest] + #[case::flatten(quote! { #[command(flatten)] pub a: Inner }, "flatten")] + #[case::subcommand(quote! { #[command(subcommand)] pub b: Inner }, "subcommand")] + #[case::other(quote! { #[command(next_help_heading = "x")] pub c: u32 }, "field-level")] + fn rejects_field_level_command_attrs(#[case] tokens: TokenStream2, #[case] needle: &str) { + let err = classify(tokens).expect_err("field-level #[command(...)] must be rejected"); + assert!(err.contains(needle), "message should mention {needle:?}: {err}"); + } + + /// `#[command(flatten)]` is matched on the parsed meta, not on a substring of + /// the attribute's tokens, so a key that merely *contains* "flatten" is + /// reported as an ordinary unsupported field-level attribute. + #[test] + fn flatten_guard_does_not_substring_match() { + let err = classify(quote! { #[command(help_heading = "flatten me")] pub a: u32 }) + .expect_err("field-level #[command(...)] is rejected"); + assert!( + err.contains("field-level"), + "a value merely containing \"flatten\" must not be reported as a flatten: {err}" + ); + } + + /// `required` is only enforceable where absence is observable. + /// + /// The `Option` / `Vec` cases are the subtle ones: clap applies a + /// declared default even to those types, so the generated `is_none()` / + /// `is_empty()` check never fires and the requirement is silently lost. + #[rstest] + #[case::defaulted(quote! { #[arg(long, default_value_t = 3, required = true)] pub a: u32 })] + #[case::bare_bool(quote! { #[arg(long, required = true)] pub b: bool })] + #[case::skipped(quote! { #[arg(skip, required = true)] pub c: u32 })] + #[case::defaulted_option( + quote! { #[arg(long, default_value = "3", required = true)] pub d: Option } + )] + #[case::defaulted_vec( + quote! { #[arg(long, default_values_t = [1u32], required = true)] pub e: Vec } + )] + fn rejects_unenforceable_required(#[case] tokens: TokenStream2) { + let err = classify(tokens).expect_err("unenforceable `required` must be rejected"); + assert!(err.contains("required"), "message should name the key: {err}"); + } + + /// Without a default, absence IS observable on both types, so `required` is + /// enforceable and must be accepted. + #[rstest] + #[case::option(quote! { #[arg(long, required = true)] pub a: Option })] + #[case::vec(quote! { #[arg(long, required = true)] pub b: Vec })] + fn accepts_enforceable_required_on_absent_able_types(#[case] tokens: TokenStream2) { + assert_eq!(classify(tokens), Ok(())); + } + + /// A generic struct expands into a companion emitted without generics, so the + /// expansion cannot compile — and the resulting errors name `T`, never + /// `multi_options`. + #[rstest] + #[case::type_param(quote! { pub struct Opts { #[arg(long)] pub a: T } })] + #[case::lifetime(quote! { pub struct Opts<'a> { #[arg(long)] pub a: &'a str } })] + #[case::where_clause( + quote! { pub struct Opts where u32: Clone { #[arg(long)] pub a: u32 } } + )] + fn rejects_generic_structs(#[case] tokens: TokenStream2) { + let err = parse_annotated_struct(&tokens) + .expect_err("a generic struct must be rejected") + .to_string(); + assert!(err.contains("generic"), "message should mention generics: {err}"); + } + + /// The ordinary non-generic case must still pass the same guard. + #[test] + fn accepts_a_plain_struct() { + assert!(parse_annotated_struct("e! { pub struct Opts { pub a: u32 } }).is_ok()); + } + + /// `required = false` is a no-op and must not trip the enforceability check. + #[test] + fn accepts_required_false_on_a_defaulted_field() { + assert_eq!( + classify(quote! { #[arg(long, default_value_t = 3, required = false)] pub a: u32 }), + Ok(()) + ); + } + + /// `required` has a bare form as well as the name-value one, and only the + /// name-value `false` disables it. + #[rstest] + #[case::bare(quote! { #[arg(long, required)] pub a: Option }, true)] + #[case::explicit_true(quote! { #[arg(long, required = true)] pub a: Option }, true)] + #[case::explicit_false(quote! { #[arg(long, required = false)] pub a: Option }, false)] + #[case::absent(quote! { #[arg(long)] pub a: Option }, false)] + fn recognizes_every_required_spelling(#[case] tokens: TokenStream2, #[case] expected: bool) { + let field = named_field(tokens); + let parsed = ParsedField::from_field(&field, "p").expect("classify"); + assert_eq!(parsed.required, expected); + } + + /// A bare alias key carries no value to re-prefix, so it cannot be honored. + #[rstest] + #[case::bare_alias(quote! { #[arg(long, alias)] pub a: u32 }, "alias")] + #[case::bare_aliases(quote! { #[arg(long, aliases)] pub b: u32 }, "aliases")] + fn rejects_valueless_alias_keys(#[case] tokens: TokenStream2, #[case] needle: &str) { + let err = classify(tokens).expect_err("a valueless alias key must be rejected"); + assert!(err.contains(needle), "message should name the key {needle:?}: {err}"); + } + + /// Struct-level clap configuration is not reproduced on the companion, so it + /// must be rejected rather than silently applied to one command only. + #[rstest] + #[case::legacy_clap(quote! { #[clap(next_help_heading = "x")] }, "#[clap(...)]")] + #[case::legacy_structopt(quote! { #[structopt(name = "x")] }, "#[clap(...)]")] + #[case::group(quote! { #[group(required = true)] }, "#[group(...)]")] + #[case::command(quote! { #[command(next_help_heading = "x")] }, "#[command(...)]")] + fn rejects_unsupported_struct_level_attrs( + #[case] attr_tokens: TokenStream2, + #[case] needle: &str, + ) { + let item: syn::ItemStruct = syn::parse2(quote! { + #attr_tokens + pub struct Opts { pub a: u32 } + }) + .expect("parse struct"); + let err = reject_unsupported_struct_attrs(&item.attrs) + .expect_err("struct-level attr must be rejected") + .to_string(); + assert!(err.contains(needle), "message should name {needle:?}: {err}"); + } + + /// `#[derive(...)]` and doc comments ride along on every annotated struct and + /// must not be mistaken for clap configuration. + #[test] + fn accepts_derive_and_doc_attrs_on_the_struct() { + let item: syn::ItemStruct = syn::parse2(quote! { + /// Docs. + #[derive(clap::Args, Debug, Clone)] + pub struct Opts { pub a: u32 } + }) + .expect("parse struct"); + assert!(reject_unsupported_struct_attrs(&item.attrs).is_ok()); + } + + /// The attribute takes exactly two non-empty string literals. + #[rstest] + #[case::well_formed(quote! { "sort", "Sort Options" }, true)] + #[case::missing_heading(quote! { "sort" }, false)] + #[case::missing_comma(quote! { "sort" "Sort Options" }, false)] + #[case::extra_argument(quote! { "sort", "Sort Options", "extra" }, false)] + #[case::empty(quote! {}, false)] + #[case::empty_heading(quote! { "sort", "" }, false)] + #[case::non_literal(quote! { sort, "Sort Options" }, false)] + fn parses_only_two_string_literal_arguments( + #[case] tokens: TokenStream2, + #[case] expected_ok: bool, + ) { + assert_eq!(syn::parse2::(tokens).is_ok(), expected_ok); + } + + /// The prefix is spliced into both a flag name and an identifier. + #[rstest] + #[case::simple("sort", true)] + #[case::with_digit("codec2", true)] + #[case::with_underscore("read_group", true)] + #[case::with_dash("read-group", true)] + #[case::empty("", false)] + #[case::leading_digit("2fast", false)] + #[case::leading_dash("-sort", false)] + #[case::colon("so::rt", false)] + #[case::space("so rt", false)] + #[case::non_ascii("sørt", false)] + fn validate_prefix_accepts_only_identifier_fragments( + #[case] prefix: &str, + #[case] expected_ok: bool, + ) { + assert_eq!(validate_prefix(prefix).is_ok(), expected_ok, "prefix {prefix:?}"); + } + + fn parse_type(tokens: TokenStream2) -> syn::Type { + syn::parse2(tokens).expect("parse type") + } + + /// `is_vec_type` drives whether a field is absent-able, so a non-path type — + /// reference, array, tuple, slice — must answer `false` rather than panic or + /// match on the last path segment of something that has none. + #[rstest] + #[case::vec(quote! { Vec }, true)] + #[case::qualified_vec(quote! { std::vec::Vec }, true)] + #[case::option_not_vec(quote! { Option }, false)] + #[case::plain(quote! { u32 }, false)] + #[case::reference(quote! { &str }, false)] + #[case::array(quote! { [u8; 4] }, false)] + #[case::tuple(quote! { (u32, u32) }, false)] + #[case::unit(quote! { () }, false)] + fn is_vec_type_only_matches_path_types_named_vec( + #[case] tokens: TokenStream2, + #[case] expected: bool, + ) { + assert_eq!(is_vec_type(&parse_type(tokens)), expected); + } + + /// Same for `is_option_type`, which decides whether a field is treated as + /// optional. + #[rstest] + #[case::option(quote! { Option }, true)] + #[case::qualified_option(quote! { std::option::Option }, true)] + #[case::vec_not_option(quote! { Vec }, false)] + #[case::plain(quote! { u32 }, false)] + #[case::reference(quote! { &str }, false)] + #[case::array(quote! { [u8; 4] }, false)] + #[case::tuple(quote! { (u32, u32) }, false)] + #[case::unit(quote! { () }, false)] + fn is_option_type_only_matches_path_types_named_option( + #[case] tokens: TokenStream2, + #[case] expected: bool, + ) { + assert_eq!(is_option_type(&parse_type(tokens)), expected); + } + + /// `is_bool_type` keeps clap's valueless `SetTrue` idiom out of the required + /// classification. `Option` and `Vec` are *not* bare bools. + #[rstest] + #[case::bare_bool(quote! { bool }, true)] + #[case::qualified_bool(quote! { core::primitive::bool }, true)] + #[case::option_bool(quote! { Option }, false)] + #[case::vec_bool(quote! { Vec }, false)] + #[case::plain(quote! { u32 }, false)] + #[case::reference(quote! { &bool }, false)] + #[case::tuple(quote! { (bool, bool) }, false)] + #[case::unit(quote! { () }, false)] + fn is_bool_type_only_matches_the_bare_bool_path( + #[case] tokens: TokenStream2, + #[case] expected: bool, + ) { + assert_eq!(is_bool_type(&parse_type(tokens)), expected); + } +} diff --git a/crates/fgumi-cli-macros/tests/behavior.rs b/crates/fgumi-cli-macros/tests/behavior.rs new file mode 100644 index 000000000..c26a2f5d0 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/behavior.rs @@ -0,0 +1,721 @@ +//! Behavioral contract tests for `multi_options`. +//! +//! Where `smoke.rs` covers the three field-kind branches, this file pins the +//! properties that make the re-exposed flags *faithful* to the standalone +//! command: the declared clap default is the one that applies, help headings +//! do not leak onto the parent command, short/long help stay split, `cfg` +//! gating survives, aliases are namespaced, and the +//! `From` + `validate()` pair is a genuine round trip. +//! +//! Each fixture is built around an attribute shape a real fgumi options struct +//! actually uses; `real_world.rs` assembles those shapes into larger fixtures +//! modelled on the `Sort` and `GroupReadsByUmi` commands. + +use std::path::PathBuf; + +use clap::{Args, CommandFactory, Parser}; +use fgumi_cli_macros::multi_options; + +// ───────────────────────────────────────────────────────────────────────────── +// Declared clap defaults are authoritative +// ───────────────────────────────────────────────────────────────────────────── + +/// The `Default` impl below deliberately DISAGREES with every `default_value*` +/// attribute. A macro that re-derives defaults from `Struct::default()` would +/// hand the prefixed flag the `Default` value; preserving the original +/// attribute verbatim keeps the standalone and prefixed flags identical. +#[multi_options("drift", "Drift Options")] +#[derive(Args, Debug, Clone, PartialEq)] +pub struct DriftOptions { + /// String-form default. + #[arg(long, default_value = "42")] + pub string_defaulted: u32, + + /// Typed default. + #[arg(long, default_value_t = 7)] + pub typed_defaulted: u32, + + /// OS-string default — `default_value_os` must count as "has a default", + /// not be misclassified as a required field. + #[arg(long, default_value_os = "/tmp/os-default")] + pub os_defaulted: PathBuf, +} + +impl Default for DriftOptions { + fn default() -> Self { + // Every value here disagrees with the clap default above. + Self { + string_defaulted: 999, + typed_defaulted: 999, + os_defaulted: PathBuf::from("/tmp/wrong"), + } + } +} + +#[derive(Parser, Debug)] +struct DriftWrapper { + #[command(flatten)] + opts: MultiDriftOptions, +} + +#[test] +fn declared_clap_defaults_win_over_a_disagreeing_default_impl() { + let parsed = DriftWrapper::try_parse_from(["test-prog"]).expect("parse with all defaults"); + let opts = parsed.opts.validate().expect("validate"); + + assert_eq!(opts.string_defaulted, 42, "string-form default_value must be preserved verbatim"); + assert_eq!(opts.typed_defaulted, 7, "default_value_t must be preserved verbatim"); + assert_eq!( + opts.os_defaulted, + PathBuf::from("/tmp/os-default"), + "default_value_os must be preserved verbatim" + ); +} + +#[test] +fn default_value_os_field_is_not_treated_as_required() { + // A field misclassified as required would be wrapped in `Option` and + // rejected by `validate()` when omitted; it also would not carry the + // default. Parsing with no flags at all proves it is optional. + let parsed = DriftWrapper::try_parse_from(["test-prog"]).expect("parse"); + assert!(parsed.opts.validate().is_ok(), "default_value_os field must not be required"); +} + +/// Asserted per argument rather than against the rendered help text: a bare +/// `help.contains("42")` also passes when the value appears anywhere else in +/// the help, including on a different flag. +#[rstest::rstest] +#[case::string_form("drift::string-defaulted", "42")] +#[case::typed("drift::typed-defaulted", "7")] +#[case::os_string("drift::os-defaulted", "/tmp/os-default")] +fn prefixed_defaults_are_advertised_in_help(#[case] long: &str, #[case] expected: &str) { + let command = DriftWrapper::command(); + let arg = command + .get_arguments() + .find(|arg| arg.get_long() == Some(long)) + .unwrap_or_else(|| panic!("--{long} should be registered")); + let defaults: Vec = + arg.get_default_values().iter().map(|value| value.to_string_lossy().into_owned()).collect(); + assert_eq!( + defaults, + vec![expected.to_string()], + "--{long} must advertise its declared default" + ); +} + +// ───────────────────────────────────────────────────────────────────────────── +// Help headings do not leak onto the parent command +// ───────────────────────────────────────────────────────────────────────────── + +#[multi_options("stage", "Stage Options")] +#[derive(Args, Debug, Clone, PartialEq)] +pub struct StageOptions { + /// A stage knob. + #[arg(long, default_value_t = 1)] + pub knob: u32, +} + +impl Default for StageOptions { + fn default() -> Self { + Self { knob: 1 } + } +} + +#[test] +fn parent_args_declared_after_a_flattened_multi_keep_their_own_heading() { + /// A parent command that declares its own argument *after* flattening the + /// generated struct. With a struct-level `next_help_heading`, clap applies + /// the heading to every argument registered after it — silently filing + /// `--after-the-flatten` under "Stage Options". + #[derive(Parser, Debug)] + struct Wrapper { + #[command(flatten)] + stage: MultiStageOptions, + + /// Declared after the flatten; belongs to the parent, not the stage. + #[arg(long)] + after_the_flatten: Option, + } + + // Asserted on each argument's heading rather than on where it lands in the + // rendered help: help ordering is clap's to change, and a parent argument + // rendered before the heading may still carry it. + let command = Wrapper::command(); + let heading_of = |long: &str| { + command + .get_arguments() + .find(|arg| arg.get_long() == Some(long)) + .unwrap_or_else(|| panic!("--{long} should be registered")) + .get_help_heading() + .map(ToString::to_string) + }; + + assert_eq!( + heading_of("stage::knob"), + Some("Stage Options".to_string()), + "the stage's own argument must carry the stage heading" + ); + assert_eq!( + heading_of("after-the-flatten"), + None, + "--after-the-flatten belongs to the parent and must carry no stage heading" + ); +} + +// A parent that declares no doc comment of its own. The comment below is +// deliberately NOT a rustdoc comment: clap adopts a parent's own doc as its +// description, which would mask the very leak this test looks for. +// +// clap adopts a flattened `Args` struct's doc comment as the parent command's +// description whenever the parent has none. The generated companion carries +// rustdoc (crates that deny(missing_docs) require it), so the macro must reset +// `about` explicitly — otherwise an undocumented parent would describe itself +// with the companion's boilerplate, and with several stages flattened clap +// would arbitrarily pick whichever came first. +#[derive(Parser, Debug)] +struct UndocumentedParent { + #[command(flatten)] + opts: MultiStageOptions, +} + +#[test] +fn the_companions_own_docs_never_become_the_parents_description() { + // Asserted on the command's own metadata rather than on rendered help, so + // the test cannot be satisfied by clap merely laying the description out + // somewhere the substring check does not look. + let command = UndocumentedParent::command(); + let about = command.get_about().map(ToString::to_string); + let long_about = command.get_long_about().map(ToString::to_string); + + // The contract is that an undocumented parent stays undocumented, not + // merely that the companion's *current* boilerplate wording is absent — a + // substring check would start passing the moment that wording changed, + // leak and all. + for (label, description) in [("about", &about), ("long_about", &long_about)] { + assert_eq!( + *description, None, + "an undocumented parent must keep an unset {label}, got {description:?}" + ); + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Field-level help metadata +// ───────────────────────────────────────────────────────────────────────────── + +/// A field-level `help_heading` cannot be namespaced per stage. Emitted after +/// the companion's own heading it would win, filing the field outside its +/// stage — and two stages declaring the same heading would merge into one +/// section naming neither. `help` carries no such ambiguity: it is the field's +/// own documentation and reads identically on both commands. +#[multi_options("meta", "Meta Options")] +#[derive(Args, Debug, Clone, PartialEq)] +pub struct MetaOptions { + /// Doc-derived help. + #[arg(long, help_heading = "Escape Hatch")] + pub own_heading: Option, + + /// Doc-derived help that clap must not prefer. + #[arg(long, help = "explicit help wins")] + pub own_help: Option, +} + +#[derive(Parser, Debug)] +struct MetaWrapper { + #[command(flatten)] + opts: MultiMetaOptions, +} + +#[test] +fn a_field_level_help_heading_cannot_escape_the_stages_heading() { + let command = MetaWrapper::command(); + let arg = command + .get_arguments() + .find(|arg| arg.get_long() == Some("meta::own-heading")) + .expect("--meta::own-heading should be registered"); + + assert_eq!( + arg.get_help_heading().map(ToString::to_string), + Some("Meta Options".to_string()), + "the stage's heading must win over the field's own" + ); +} + +#[test] +fn an_explicit_help_is_preserved_verbatim() { + let command = MetaWrapper::command(); + let arg = command + .get_arguments() + .find(|arg| arg.get_long() == Some("meta::own-help")) + .expect("--meta::own-help should be registered"); + + assert_eq!( + arg.get_help().map(ToString::to_string), + Some("explicit help wins".to_string()), + "an explicit help is the field's own documentation and must survive re-exposure" + ); +} + +// ───────────────────────────────────────────────────────────────────────────── +// Short/long help split survives re-exposure +// ───────────────────────────────────────────────────────────────────────────── + +#[multi_options("docs", "Docs Options")] +#[derive(Args, Debug, Clone, PartialEq)] +pub struct DocsOptions { + /// Short summary line. + /// + /// Longer explanation that clap shows only for `--help`, never for `-h`. + /// Real fgumi options structs carry several such paragraphs per field. + #[arg(long, default_value_t = 1)] + pub documented: u32, +} + +impl Default for DocsOptions { + fn default() -> Self { + Self { documented: 1 } + } +} + +#[derive(Parser, Debug)] +struct DocsWrapper { + #[command(flatten)] + opts: MultiDocsOptions, +} + +/// Returns the `--docs::documented` argument of the prefixed companion. +/// +/// The short/long split is asserted on `Arg::get_help()`/`get_long_help()` +/// rather than on `render_help()`: rendered help is clap's layout concern and +/// wraps to the terminal width once the `wrap_help` feature is anywhere in the +/// workspace feature graph, so a substring check there can fail while the +/// contract holds — and pass while it does not, since the text of one argument +/// is indistinguishable from any other's in the rendered block. +fn documented_arg(command: &clap::Command) -> &clap::Arg { + command + .get_arguments() + .find(|arg| arg.get_long() == Some("docs::documented")) + .expect("--docs::documented should be registered") +} + +#[test] +fn short_help_shows_only_the_first_doc_paragraph() { + let command = DocsWrapper::command(); + // clap trims the trailing period when it derives short help from a doc + // comment, so match the sentence without it. + let short = documented_arg(&command).get_help().map(ToString::to_string); + let short = short.expect("the documented field must carry short help"); + assert!(short.contains("Short summary line"), "short help should show the summary: {short}"); + assert!( + !short.contains("Longer explanation"), + "short help must not carry the long explanation: {short}" + ); +} + +#[test] +fn long_help_shows_every_doc_paragraph() { + let command = DocsWrapper::command(); + let long = documented_arg(&command).get_long_help().map(ToString::to_string); + let long = long.expect("the documented field must carry long help"); + assert!(long.contains("Short summary line."), "long help should show the summary: {long}"); + assert!( + long.contains("Longer explanation"), + "long help should show the full explanation: {long}" + ); +} + +// ───────────────────────────────────────────────────────────────────────────── +// cfg-gated fields +// ───────────────────────────────────────────────────────────────────────────── + +/// `#[cfg(any())]` is never satisfied and `#[cfg(all())]` always is, so this +/// fixture exercises both sides of cfg forwarding without needing a feature +/// flag. A macro that drops `#[cfg]` emits a Multi field (and `validate()` / +/// `From` arms) for a field that does not exist on the original struct, which +/// does not compile — so merely building this file is most of the test. +#[multi_options("gated", "Gated Options")] +#[derive(Args, Debug, Clone, PartialEq)] +pub struct GatedOptions { + /// Always present. + #[arg(long, default_value_t = 1)] + pub always: u32, + + /// Compiled out. + #[cfg(any())] + #[arg(long, default_value_t = 2)] + pub never: u32, +} + +impl Default for GatedOptions { + fn default() -> Self { + Self { always: 1 } + } +} + +#[derive(Parser, Debug)] +struct GatedWrapper { + #[command(flatten)] + opts: MultiGatedOptions, +} + +#[test] +fn cfg_gated_out_field_is_absent_from_the_multi_struct() { + // Constructing the Multi struct with only the surviving field proves the + // gated-out field was not emitted. + let multi = MultiGatedOptions { gated_always: 5 }; + let opts = multi.validate().expect("validate"); + assert_eq!(opts.always, 5); + + assert!( + GatedWrapper::command().get_arguments().all(|arg| arg.get_long() != Some("gated::never")), + "the cfg-ed out field must not be registered as an argument" + ); +} + +// ───────────────────────────────────────────────────────────────────────────── +// Aliases are namespaced; short flags are never propagated +// ───────────────────────────────────────────────────────────────────────────── + +/// Mirrors `fgumi clip`'s `--ref` / `-r` shape: a long override, a short flag +/// and a long alias on one required field. `multi_aliased` additionally covers +/// the list spellings, which take a different code path to the single-valued +/// ones and would otherwise only be checked for "does not error". +#[multi_options("alias", "Alias Options")] +#[derive(Args, Debug, Clone, PartialEq)] +pub struct AliasOptions { + /// Reference fasta, with a short flag and a long alias. + #[arg(long = "reference", short = 'r', alias = "ref")] + pub reference: PathBuf, + + /// Carries every long-alias spelling at once. + #[arg( + long = "multi-aliased", + aliases = ["one", "two"], + visible_alias = "vis", + visible_aliases = ["vis-one", "vis-two"], + short_alias = 's', + default_value_t = 0 + )] + pub multi_aliased: u32, +} + +#[derive(Parser, Debug)] +struct AliasWrapper { + #[command(flatten)] + opts: MultiAliasOptions, +} + +#[test] +fn long_alias_is_namespaced_under_the_prefix() { + let parsed = AliasWrapper::try_parse_from(["test-prog", "--alias::ref", "/tmp/a.fa"]) + .expect("the prefixed alias should parse"); + let opts = parsed.opts.validate().expect("validate"); + assert_eq!(opts.reference, PathBuf::from("/tmp/a.fa")); +} + +/// Assert a parse failed *because clap does not know the flag* — the contract +/// under test — rather than for any reason at all. +/// +/// A bare `is_err()` would pass by accident: every companion field is defaulted, +/// absent-able, or staged as `Option`, so clap never fails a parse for a +/// missing value today, and a later fixture change could satisfy the assertion +/// with an unrelated failure. +#[track_caller] +fn assert_unknown_argument(result: Result, flag: &str) { + let err = result.map(|_| ()).expect_err(&format!("{flag} must be rejected")); + assert_eq!( + err.kind(), + clap::error::ErrorKind::UnknownArgument, + "{flag} must be rejected as an unknown argument, got {:?}: {err}", + err.kind() + ); +} + +#[test] +fn unprefixed_alias_does_not_leak_onto_the_parent_command() { + assert_unknown_argument( + AliasWrapper::try_parse_from(["test-prog", "--ref", "/tmp/a.fa"]), + "--ref", + ); +} + +#[test] +fn short_flag_is_not_propagated() { + assert_unknown_argument(AliasWrapper::try_parse_from(["test-prog", "-r", "/tmp/a.fa"]), "-r"); + assert_unknown_argument( + AliasWrapper::try_parse_from(["test-prog", "--alias::reference", "/tmp/a.fa", "-s", "1"]), + "-s", + ); +} + +/// Every long-alias spelling — `aliases`, `visible_alias`, `visible_aliases` — +/// must reach clap re-prefixed, not merely survive classification. +#[rstest::rstest] +#[case::alias_list_first("--alias::one")] +#[case::alias_list_second("--alias::two")] +#[case::visible_alias("--alias::vis")] +#[case::visible_alias_list_first("--alias::vis-one")] +#[case::visible_alias_list_second("--alias::vis-two")] +fn every_long_alias_spelling_is_namespaced(#[case] flag: &str) { + let parsed = + AliasWrapper::try_parse_from(["test-prog", "--alias::reference", "/tmp/a.fa", flag, "7"]) + .unwrap_or_else(|e| panic!("{flag} should parse: {e}")); + assert_eq!(parsed.opts.validate().expect("validate").multi_aliased, 7); +} + +/// The un-prefixed spellings must not reach the parent command, where two +/// stages re-exposing the same options struct would collide on them. +#[rstest::rstest] +#[case::alias_list_first("--one")] +#[case::alias_list_second("--two")] +#[case::visible_alias("--vis")] +#[case::visible_alias_list_first("--vis-one")] +#[case::visible_alias_list_second("--vis-two")] +fn no_long_alias_spelling_leaks_unprefixed(#[case] flag: &str) { + assert_unknown_argument( + AliasWrapper::try_parse_from(["test-prog", "--alias::reference", "/tmp/a.fa", flag, "7"]), + flag, + ); +} + +/// Visible aliases must survive into `--help`, hidden ones must not. The parse +/// cases above accept every spelling regardless of visibility, so a macro that +/// re-emitted every alias as *hidden* would pass them while silently dropping the +/// visible aliases from the rendered help — a user-facing regression the parse +/// tests cannot see. Assert against rendered help directly (clap 4 exposes no +/// `Arg::get_visible_aliases` getter). +#[test] +fn visible_aliases_reach_help_and_hidden_ones_do_not() { + let help = AliasWrapper::command().render_long_help().to_string(); + for visible in ["alias::vis", "alias::vis-one", "alias::vis-two"] { + assert!( + help.contains(visible), + "visible alias `{visible}` should appear in --help, but did not:\n{help}" + ); + } + for hidden in ["alias::one", "alias::two"] { + assert!( + !help.contains(hidden), + "hidden alias `{hidden}` must not appear in --help, but did:\n{help}" + ); + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Bare `bool` stays a valueless flag +// ───────────────────────────────────────────────────────────────────────────── + +#[multi_options("flag", "Flag Options")] +#[derive(Args, Debug, Clone, PartialEq)] +pub struct FlagOptions { + /// Bare bool — clap's `SetTrue` idiom, absent means `false`. + #[arg(long)] + pub enabled: bool, +} + +#[derive(Parser, Debug)] +struct FlagWrapper { + #[command(flatten)] + opts: MultiFlagOptions, +} + +#[test] +fn bare_bool_defaults_to_false_without_being_required() { + let parsed = FlagWrapper::try_parse_from(["test-prog"]).expect("parse with no flags"); + let opts = parsed.opts.validate().expect("a bare bool must not be required"); + assert!(!opts.enabled); +} + +#[test] +fn bare_bool_takes_no_value() { + let parsed = + FlagWrapper::try_parse_from(["test-prog", "--flag::enabled"]).expect("valueless flag"); + assert!(parsed.opts.validate().expect("validate").enabled); +} + +// ───────────────────────────────────────────────────────────────────────────── +// `required` on absent-able types is enforced by validate() +// ───────────────────────────────────────────────────────────────────────────── + +#[multi_options("req", "Req Options")] +#[derive(Args, Debug, Clone, PartialEq)] +pub struct RequiredOptions { + /// Required despite being `Option` on the standalone command. + #[arg(long, required = true)] + pub needed: Option, + + /// Required despite being `Vec` on the standalone command. + #[arg(long, required = true, action = clap::ArgAction::Append)] + pub needed_many: Vec, +} + +#[derive(Parser, Debug)] +struct RequiredWrapper { + #[command(flatten)] + opts: MultiRequiredOptions, +} + +#[test] +fn required_option_field_is_enforced_by_validate() { + // clap must not enforce it during parsing (staged validation owns that), + // but validate() must still refuse the missing value. + let parsed = RequiredWrapper::try_parse_from(["test-prog", "--req::needed-many", "1"]) + .expect("parse should succeed; required-ness is staged"); + let err = parsed.opts.validate().expect_err("validate must reject the missing Option field"); + let msg = format!("{err:#}"); + // `--req::needed` is a prefix of `--req::needed-many`, so match the whole + // " is required" phrase: a bare substring check would pass even if + // validate() named the wrong field. + assert!( + msg.contains("--req::needed is required"), + "error should name the Option field, not --req::needed-many: {msg}" + ); +} + +#[test] +fn required_vec_field_is_enforced_by_validate() { + let parsed = RequiredWrapper::try_parse_from(["test-prog", "--req::needed", "1"]) + .expect("parse should succeed; required-ness is staged"); + let err = parsed.opts.validate().expect_err("validate must reject the empty Vec field"); + let msg = format!("{err:#}"); + assert!( + msg.contains("--req::needed-many is required"), + "error should name the Vec field: {msg}" + ); +} + +#[test] +fn required_fields_validate_once_supplied() { + let parsed = RequiredWrapper::try_parse_from([ + "test-prog", + "--req::needed", + "3", + "--req::needed-many", + "4", + ]) + .expect("parse"); + let opts = parsed.opts.validate().expect("validate"); + assert_eq!(opts.needed, Some(3)); + assert_eq!(opts.needed_many, vec![4]); +} + +// ───────────────────────────────────────────────────────────────────────────── +// Vec fields pass through; the struct `Default` is not consulted +// ───────────────────────────────────────────────────────────────────────────── + +/// `Default` returns a non-empty Vec, but clap never consults a struct's +/// `Default` — the standalone command yields an empty Vec when the flag is +/// omitted, so the prefixed flag must too. +#[multi_options("vecs", "Vec Options")] +#[derive(Args, Debug, Clone, PartialEq)] +pub struct VecOptions { + /// Repeatable values. + #[arg(long, action = clap::ArgAction::Append)] + pub values: Vec, +} + +impl Default for VecOptions { + fn default() -> Self { + Self { values: vec![7, 8, 9] } + } +} + +#[test] +fn omitted_vec_matches_the_standalone_command_not_the_struct_default() { + #[derive(Parser, Debug)] + struct Wrapper { + #[command(flatten)] + opts: MultiVecOptions, + } + #[derive(Parser, Debug)] + struct Standalone { + #[command(flatten)] + opts: VecOptions, + } + + let standalone = Standalone::try_parse_from(["test-prog"]).expect("parse standalone"); + let prefixed = Wrapper::try_parse_from(["test-prog"]).expect("parse prefixed"); + + assert!(standalone.opts.values.is_empty(), "standalone yields an empty Vec"); + assert_eq!( + prefixed.opts.validate().expect("validate").values, + standalone.opts.values, + "the prefixed flag must agree with the standalone command" + ); +} + +// ───────────────────────────────────────────────────────────────────────────── +// From + validate() is a lossless round trip +// ───────────────────────────────────────────────────────────────────────────── + +/// Mirrors `GroupOptions`' shape: CLI fields plus `#[arg(skip)]` slots that the +/// command fills in itself. +#[multi_options("trip", "Trip Options")] +#[derive(Args, Debug, Clone, PartialEq)] +pub struct TripOptions { + /// A normal flag. + #[arg(long, default_value_t = 1)] + pub knob: u32, + + /// Skipped slot with an explicit expression. + #[arg(skip = 5u32)] + pub skipped_with_expr: u32, + + /// Bare skipped slot. + #[arg(skip)] + pub bare_skipped: u32, +} + +impl Default for TripOptions { + fn default() -> Self { + Self { knob: 1, skipped_with_expr: 5, bare_skipped: 11 } + } +} + +#[test] +fn round_trip_preserves_skip_field_values() { + // Values that match neither the skip expression nor the struct Default, so + // a re-derived value is distinguishable from a preserved one. + let original = TripOptions { knob: 2, skipped_with_expr: 100, bare_skipped: 200 }; + let multi: MultiTripOptions = original.clone().into(); + let back = multi.validate().expect("validate"); + assert_eq!(back, original, "From + validate must preserve skip-field values"); +} + +#[test] +fn parsed_skip_fields_fall_back_to_their_declared_value() { + #[derive(Parser, Debug)] + struct Wrapper { + #[command(flatten)] + opts: MultiTripOptions, + } + + let parsed = Wrapper::try_parse_from(["test-prog"]).expect("parse"); + let opts = parsed.opts.validate().expect("validate"); + assert_eq!(opts.skipped_with_expr, 5, "skip = expr supplies the parse-time value"); + assert_eq!(opts.bare_skipped, 11, "bare skip falls back to the struct Default"); +} + +#[test] +fn skip_fields_are_not_exposed_as_flags() { + #[derive(Parser, Debug)] + struct Wrapper { + #[command(flatten)] + opts: MultiTripOptions, + } + // Assert on the registered arguments, which is the actual contract; a help + // substring check would break on any doc comment that happens to use the word. + let command = Wrapper::command(); + let longs: Vec<&str> = command.get_arguments().filter_map(clap::Arg::get_long).collect(); + assert!(!longs.contains(&"trip::skipped-with-expr"), "skip fields must not become flags"); + assert!(!longs.contains(&"trip::bare-skipped"), "skip fields must not become flags"); + assert!(longs.contains(&"trip::knob"), "the non-skip field should still be registered"); +} + +#[test] +fn try_from_is_the_canonical_conversion() { + let original = TripOptions { knob: 3, skipped_with_expr: 4, bare_skipped: 5 }; + let multi = MultiTripOptions::from(original.clone()); + let back = TripOptions::try_from(multi).expect("TryFrom should succeed"); + assert_eq!(back, original); +} diff --git a/crates/fgumi-cli-macros/tests/compile_fail.rs b/crates/fgumi-cli-macros/tests/compile_fail.rs new file mode 100644 index 000000000..8a908dbc9 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/compile_fail.rs @@ -0,0 +1,13 @@ +//! Compile-fail coverage for the build-time diagnostics `multi_options` documents. +//! +//! The macro reports every unsupported input by returning a spanned `syn::Error` +//! and emitting it via `into_compile_error()`, so these surface as ordinary +//! compilation errors — they cannot be exercised with `#[should_panic]`, because +//! the macro runs while the test crate is being compiled, not while it runs. +//! `trybuild` compiles each case and diffs the emitted error against its +//! committed `.stderr`, which pins both the message and the span it points at. + +#[test] +fn documented_invalid_inputs_fail_to_compile() { + trybuild::TestCases::new().compile_fail("tests/ui/*.rs"); +} diff --git a/crates/fgumi-cli-macros/tests/real_world.rs b/crates/fgumi-cli-macros/tests/real_world.rs new file mode 100644 index 000000000..c2fd20144 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/real_world.rs @@ -0,0 +1,503 @@ +//! Parity tests over fixtures shaped like the fgumi options structs +//! `multi_options` exists to re-expose. +//! +//! The macro's classification rules only matter insofar as they hold for the +//! attribute *combinations* `runall` will actually annotate, so the fixtures +//! below are modelled on the two commands with the widest `#[arg(...)]` surface: +//! `Sort` (`src/lib/commands/sort.rs`) and `GroupReadsByUmi` +//! (`src/lib/commands/group.rs`). Between them they cover a string-form +//! `default_value` on a `Display`-less type, a `long` override paired with a +//! `short`, a repeatable `Vec`, a ranged `value_parser`, a hidden expert flag, +//! `#[arg(skip)]` slots with and without a declared value, a required +//! `value_enum`, and the `num_args = 0..=1` boolean form every fgumi flag uses. +//! +//! These are **fixtures, not mirrors**. `fgumi-cli-macros` sits below the fgumi +//! commands in the dependency graph — they will annotate their options structs +//! with `multi_options`, not the other way round — so its tests deliberately do +//! not reach into `fgumi_lib` for the production definitions: the value types are +//! local stand-ins, and the fields and defaults are not kept in lockstep (the +//! production `Sort`, for one, spells its memory default `"768M"`, not +//! `"768MiB"`). What is under test is the macro's contract, not fgumi's option +//! set — a production option change is expected to leave these tests untouched +//! and is covered by that command's own tests. What warrants a fixture here is a +//! genuinely new *attribute shape* reaching a struct the macro must re-expose. +//! +//! The central assertion is parity: parsing the standalone command and parsing +//! the prefixed `Multi*` companion with the same inputs must produce identical +//! option structs. That covers every classification rule at once — a dropped +//! default, a misclassified `bool`, a leaked alias or a lost `value_parser` all +//! surface as a field that disagrees. + +use std::path::PathBuf; + +use clap::{Args, CommandFactory, Parser, ValueEnum}; +use fgumi_cli_macros::multi_options; + +// ───────────────────────────────────────────────────────────────────────────── +// Local stand-ins for the fgumi value types +// ───────────────────────────────────────────────────────────────────────────── + +/// Stand-in for `fgumi_lib`'s `MemoryLimit`, which parses "768MiB"-style values +/// and deliberately does **not** implement `Display` — the reason the macro must +/// preserve the string-form `default_value` rather than rewriting it to +/// `default_value_t`. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct MemoryLimit(pub u64); + +fn parse_memory(value: &str) -> Result { + let (digits, multiplier) = match value.strip_suffix("MiB") { + Some(digits) => (digits, 1024 * 1024), + None => match value.strip_suffix("GiB") { + Some(digits) => (digits, 1024 * 1024 * 1024), + None => (value, 1), + }, + }; + let n = digits.parse::().map_err(|e| format!("invalid memory value {value:?}: {e}"))?; + // `checked_mul` so an oversized `GiB` value maps to a parser error rather than + // panicking on multiplication overflow in a debug build. + let bytes = n + .checked_mul(multiplier) + .ok_or_else(|| format!("invalid memory value {value:?}: overflows u64"))?; + Ok(MemoryLimit(bytes)) +} + +/// Stand-in for the `parse_bool` value parser shared by every fgumi boolean flag. +fn parse_bool(value: &str) -> Result { + match value { + "true" | "yes" | "t" | "y" => Ok(true), + "false" | "no" | "f" | "n" => Ok(false), + other => Err(format!("invalid boolean {other:?}")), + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, ValueEnum)] +pub enum SpillCodec { + #[default] + Zstd, + Bgzf, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, ValueEnum)] +pub enum Strategy { + #[default] + Identity, + Edit, + Adjacency, + Paired, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct SortOrderArg; + +// ───────────────────────────────────────────────────────────────────────────── +// `SortOptions` — string-form defaults, a `Display`-less value type, a long +// override with a short flag, a repeatable Vec, hidden expert flags, a skip slot +// ───────────────────────────────────────────────────────────────────────────── + +#[multi_options("sort", "Sort Options")] +#[derive(Args, Debug, Clone, PartialEq)] +pub struct SortOptions { + /// Maximum memory for in-memory sorting. + /// + /// Default is "768MiB" per thread (matching samtools' 768 MiB). Explicit + /// values like "512MiB", "1GiB", "4GiB" are per-thread when + /// --memory-per-thread is enabled (default). + /// + /// When the limit is reached, sorted chunks spill to temporary files. + #[arg(short = 'm', long = "max-memory", default_value = "768MiB", value_parser = parse_memory)] + pub max_memory: MemoryLimit, + + /// Scale memory limit by thread count (samtools behavior). + #[arg(long = "memory-per-thread", default_value = "true", num_args = 0..=1, default_missing_value = "true", action = clap::ArgAction::Set, value_parser = parse_bool)] + pub memory_per_thread: bool, + + /// Temporary directory for intermediate files. Repeatable. + #[arg(short = 'T', long = "tmp-dir", action = clap::ArgAction::Append)] + pub tmp_dirs: Vec, + + /// Compression level for temporary chunk files (0-9). + #[arg(long = "temp-compression", default_value = "1", value_parser = clap::value_parser!(u32).range(0..=9))] + pub temp_compression: u32, + + /// Codec for temporary spill files: `zstd` (default) or `bgzf`. + #[arg(long = "temp-codec", default_value = "zstd")] + pub temp_codec: SpillCodec, + + /// Worker threads for the accumulation/sort/spill phase (Phase 1). + #[arg(long = "sort-threads")] + pub sort_threads: Option, + + /// Phase-2 spill decompression granularity (expert tuning). + #[arg(long = "file-granularity", default_value = "false", num_args = 0..=1, default_missing_value = "true", action = clap::ArgAction::Set, value_parser = parse_bool, hide = true)] + pub file_granularity: bool, + + /// Sort order (chain-builder slot), populated by the command, never by clap. + #[arg(skip)] + pub order: SortOrderArg, +} + +impl Default for SortOptions { + fn default() -> Self { + Self { + max_memory: MemoryLimit(768 * 1024 * 1024), + memory_per_thread: true, + tmp_dirs: Vec::new(), + temp_compression: 1, + temp_codec: SpillCodec::Zstd, + sort_threads: None, + file_granularity: false, + order: SortOrderArg, + } + } +} + +#[derive(Parser, Debug)] +struct StandaloneSort { + #[command(flatten)] + opts: SortOptions, +} + +#[derive(Parser, Debug)] +struct PrefixedSort { + #[command(flatten)] + opts: MultiSortOptions, +} + +#[test] +fn sort_options_defaults_match_between_standalone_and_prefixed() { + let standalone = StandaloneSort::try_parse_from(["fgumi"]).expect("standalone defaults"); + let prefixed = PrefixedSort::try_parse_from(["fgumi"]).expect("prefixed defaults"); + + assert_eq!( + prefixed.opts.validate().expect("validate"), + standalone.opts, + "every re-exposed default must equal the standalone command's default" + ); +} + +#[test] +fn sort_options_supplied_values_match_between_standalone_and_prefixed() { + let standalone = StandaloneSort::try_parse_from([ + "fgumi", + "--max-memory", + "4GiB", + "--memory-per-thread", + "false", + "--tmp-dir", + "/scratch/a", + "--tmp-dir", + "/scratch/b", + "--temp-compression", + "9", + "--temp-codec", + "bgzf", + "--sort-threads", + "8", + ]) + .expect("standalone parse"); + + let prefixed = PrefixedSort::try_parse_from([ + "fgumi", + "--sort::max-memory", + "4GiB", + "--sort::memory-per-thread", + "false", + "--sort::tmp-dir", + "/scratch/a", + "--sort::tmp-dir", + "/scratch/b", + "--sort::temp-compression", + "9", + "--sort::temp-codec", + "bgzf", + "--sort::sort-threads", + "8", + ]) + .expect("prefixed parse"); + + assert_eq!(prefixed.opts.validate().expect("validate"), standalone.opts); +} + +/// Assert a parse failed for the stated reason, not merely that it failed. A bare +/// `is_err()` would be satisfied by any unrelated failure a later fixture change +/// introduced. +#[track_caller] +fn assert_parse_error_kind( + result: Result, + expected: clap::error::ErrorKind, + what: &str, +) { + let err = result.map(|_| ()).expect_err(&format!("{what} must be rejected")); + assert_eq!(err.kind(), expected, "{what}: expected {expected:?}, got {:?}: {err}", err.kind()); +} + +#[test] +fn sort_options_value_parsers_are_still_enforced_on_the_prefixed_side() { + // `temp_compression`'s ranged value_parser and `max_memory`'s custom parser + // both arrive as `#[arg(...)]` metas the macro copies verbatim. The failure + // must come from the value parser, not from clap failing to know the flag. + assert_parse_error_kind( + PrefixedSort::try_parse_from(["fgumi", "--sort::temp-compression", "10"]), + clap::error::ErrorKind::ValueValidation, + "the ranged value_parser must reject 10", + ); + assert_parse_error_kind( + PrefixedSort::try_parse_from(["fgumi", "--sort::max-memory", "not-a-size"]), + clap::error::ErrorKind::ValueValidation, + "the custom value_parser must reject a malformed size", + ); +} + +#[test] +fn sort_options_hidden_flag_stays_hidden_when_re_exposed() { + // The contract is the `hide` setting on the registered argument, not the + // absence of a substring from rendered text. + let command = PrefixedSort::command(); + let arg = command + .get_arguments() + .find(|arg| arg.get_long() == Some("sort::file-granularity")) + .expect("the prefixed flag should still be registered"); + assert!(arg.is_hide_set(), "hide = true must be preserved onto the prefixed flag"); +} + +#[test] +fn sort_options_short_flags_are_not_propagated() { + // `-m` and `-T` belong to the standalone command; on runall they would + // collide with every other stage's short flags. + assert!(StandaloneSort::try_parse_from(["fgumi", "-m", "1GiB"]).is_ok()); + assert_parse_error_kind( + PrefixedSort::try_parse_from(["fgumi", "-m", "1GiB"]), + clap::error::ErrorKind::UnknownArgument, + "-m", + ); + assert_parse_error_kind( + PrefixedSort::try_parse_from(["fgumi", "-T", "/scratch"]), + clap::error::ErrorKind::UnknownArgument, + "-T", + ); +} + +#[test] +fn every_standalone_long_flag_has_a_prefixed_counterpart() { + let standalone_flags: Vec = StandaloneSort::command() + .get_arguments() + .filter_map(|arg| arg.get_long().map(ToString::to_string)) + .filter(|long| long != "help") + .collect(); + assert!(!standalone_flags.is_empty(), "fixture should declare long flags"); + + let prefixed_flags: Vec = PrefixedSort::command() + .get_arguments() + .filter_map(|arg| arg.get_long().map(ToString::to_string)) + .collect(); + + // Compared as sets, not one-directionally: a containment check passes even + // when the prefixed command exposes an extra flag the standalone one never + // declared. + let mut expected: Vec = + standalone_flags.iter().map(|long| format!("sort::{long}")).collect(); + let mut actual: Vec = + prefixed_flags.iter().filter(|long| *long != "help").cloned().collect(); + expected.sort(); + actual.sort(); + + assert_eq!( + actual, expected, + "the prefixed command must expose exactly the prefixed counterparts, no more and no less" + ); +} + +#[test] +fn skip_slot_is_not_exposed_but_survives_the_round_trip() { + // The contract is that no `--sort::order` argument is registered — not that + // the word "order" is absent from the help text, which any future doc + // comment could break while the contract still holds. + assert!( + PrefixedSort::command().get_arguments().all(|arg| arg.get_long() != Some("sort::order")), + "the skip slot must not become a flag" + ); + + let original = SortOptions { order: SortOrderArg, ..SortOptions::default() }; + let multi: MultiSortOptions = original.clone().into(); + assert_eq!(multi.validate().expect("validate"), original); +} + +// ───────────────────────────────────────────────────────────────────────────── +// `GroupOptions` — a required `value_enum`, several skip slots +// ───────────────────────────────────────────────────────────────────────────── + +#[multi_options("group", "Group Options")] +#[derive(Args, Debug, Clone, PartialEq)] +pub struct GroupOptions { + /// Minimum mapping quality. + #[arg(short = 'm', long = "min-map-q", default_value = "1")] + pub min_map_q: u8, + + /// Include non-PF reads. + #[arg(short = 'n', long = "include-non-pf-reads", default_value = "false", num_args = 0..=1, default_missing_value = "true", action = clap::ArgAction::Set, value_parser = parse_bool)] + pub include_non_pf_reads: bool, + + /// UMI assignment strategy. Required on the standalone command. + #[arg(short = 's', long = "strategy", value_enum)] + pub strategy: Strategy, + + /// Minimum UMI length. + #[arg(short = 'l', long = "min-umi-length")] + pub min_umi_length: Option, + + /// Strategy actually used, resolved by the command. + #[arg(skip = Strategy::Identity)] + pub effective_strategy: Strategy, + + /// Edits actually used, resolved by the command. + #[arg(skip)] + pub effective_edits: u32, +} + +impl Default for GroupOptions { + fn default() -> Self { + Self { + min_map_q: 1, + include_non_pf_reads: false, + strategy: Strategy::Identity, + min_umi_length: None, + effective_strategy: Strategy::Identity, + effective_edits: 0, + } + } +} + +#[derive(Parser, Debug)] +struct StandaloneGroup { + #[command(flatten)] + opts: GroupOptions, +} + +#[derive(Parser, Debug)] +struct PrefixedGroup { + #[command(flatten)] + opts: MultiGroupOptions, +} + +#[test] +fn group_options_match_when_the_required_strategy_is_supplied() { + let standalone = StandaloneGroup::try_parse_from(["fgumi", "--strategy", "adjacency"]) + .expect("standalone parse"); + let prefixed = PrefixedGroup::try_parse_from(["fgumi", "--group::strategy", "adjacency"]) + .expect("prefixed parse"); + + assert_eq!(prefixed.opts.validate().expect("validate"), standalone.opts); +} + +#[test] +fn group_required_field_is_staged_not_enforced_by_clap() { + // The standalone command refuses to parse without --strategy. On the runall + // side clap must accept the parse so `validate()` can name the stage and the + // flag that is missing. + let err = StandaloneGroup::try_parse_from(["fgumi"]) + .map(|_| ()) + .expect_err("the standalone command requires --strategy"); + assert_eq!( + err.kind(), + clap::error::ErrorKind::MissingRequiredArgument, + "expected clap itself to demand --strategy, got {:?}: {err}", + err.kind() + ); + + let prefixed = PrefixedGroup::try_parse_from(["fgumi"]).expect("parse must be staged"); + let err = prefixed.opts.validate().expect_err("validate must reject the missing strategy"); + let msg = format!("{err:#}"); + assert!(msg.contains("--group::strategy"), "error should name the prefixed flag: {msg}"); + assert!(msg.contains("required when group is selected"), "error should name the stage: {msg}"); +} + +// ───────────────────────────────────────────────────────────────────────────── +// Two stages flattened side by side — the shape `runall` actually builds +// ───────────────────────────────────────────────────────────────────────────── + +/// `runall` flattens every stage into one command. This is where an un-prefixed +/// flag, a propagated short, or a leaked `next_help_heading` would surface as a +/// clap panic or a mis-filed argument. +#[derive(Parser, Debug)] +struct RunAllLike { + #[command(flatten)] + sort: MultiSortOptions, + + #[command(flatten)] + group: MultiGroupOptions, + + /// A flag belonging to runall itself, declared after both stages. + #[arg(long)] + threads: Option, +} + +#[test] +fn two_stages_and_a_parent_flag_coexist_in_one_command() { + let parsed = RunAllLike::try_parse_from([ + "fgumi", + "--sort::max-memory", + "2GiB", + "--group::strategy", + "paired", + "--threads", + "4", + ]) + .expect("a two-stage command must build and parse"); + + assert_eq!( + parsed.sort.validate().expect("sort").max_memory, + MemoryLimit(2 * 1024 * 1024 * 1024) + ); + assert_eq!(parsed.group.validate().expect("group").strategy, Strategy::Paired); + assert_eq!(parsed.threads, Some(4)); +} + +#[test] +fn each_stage_gets_its_own_help_heading_and_the_parent_keeps_its_own() { + // Asserted on each argument's own `help_heading` rather than on the byte + // offsets of headings within `render_long_help()`: rendered help is clap's + // layout concern — it depends on argument ordering and, with the + // `wrap_help` feature active anywhere in the workspace feature graph, on + // terminal width — so offset comparisons can hold while the headings are + // wrong, and break while they are right. + let command = RunAllLike::command(); + let heading_of = |long: &str| { + command + .get_arguments() + .find(|arg| arg.get_long() == Some(long)) + .unwrap_or_else(|| panic!("--{long} should be registered")) + .get_help_heading() + .map(ToString::to_string) + }; + + assert_eq!( + heading_of("sort::max-memory"), + Some("Sort Options".to_string()), + "a sort field must be filed under the sort stage's heading" + ); + assert_eq!( + heading_of("group::strategy"), + Some("Group Options".to_string()), + "a group field must be filed under the group stage's heading" + ); + assert_eq!( + heading_of("threads"), + None, + "the parent's own flag must not inherit a stage heading" + ); +} + +#[test] +fn stages_sharing_a_flag_name_do_not_collide() { + // Both stages declare a `-m` short and a min/max flag on the standalone side. + // Prefixing is what keeps them apart; a leak would make clap panic while + // building the command above, so reaching this assertion is the test. + let flags: Vec = RunAllLike::command() + .get_arguments() + .filter_map(|arg| arg.get_long().map(ToString::to_string)) + .collect(); + let mut sorted = flags.clone(); + sorted.sort(); + sorted.dedup(); + assert_eq!(sorted.len(), flags.len(), "no flag name may be declared twice: {flags:?}"); +} diff --git a/crates/fgumi-cli-macros/tests/smoke.rs b/crates/fgumi-cli-macros/tests/smoke.rs new file mode 100644 index 000000000..075316df8 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/smoke.rs @@ -0,0 +1,412 @@ +//! Smoke test for the `multi_options` attribute macro. +//! +//! Exercises the field-kind branches the macro generates code for: +//! skipped (`#[arg(skip)]`), absent-able or defaulted (`Option`, +//! `Vec`, bare `bool`, any `default_value*`), and required +//! (everything else). Verifies that: +//! +//! * The original struct compiles unchanged with its bare flags. +//! * The generated `Multi` struct exposes the same fields +//! under `--::` names. +//! * `Multi::validate()` succeeds when required fields are +//! supplied and fails (with a helpful error) when they are not. +//! * The round-trip `From<>` impl preserves field values. + +use clap::{Args, Parser}; +use fgumi_cli_macros::multi_options; + +/// Standalone options struct annotated with `multi_options`. The +/// generated `MultiFooOptions` is what the smoke test pokes at. +/// +/// `Default` is hand-rolled here only because `bare_skipped`-style +/// fields need it; the CLI defaults come from the `#[arg(...)]` +/// attributes, which the macro copies verbatim. See +/// `behavior.rs::declared_clap_defaults_win_over_a_disagreeing_default_impl` +/// for the test that pins the two apart. +#[multi_options("foo", "Foo Options")] +#[derive(Args, Debug, Clone, PartialEq)] +pub struct FooOptions { + /// Optional knob (None when not passed). + #[arg(long)] + pub optional_knob: Option, + + /// Defaulted knob — keeps its `default_value_t` on both sides. + #[arg(long, default_value_t = 7)] + pub defaulted_knob: u32, + + /// Required knob — missing on the prefixed side fails `validate()`. + #[arg(long)] + pub required_knob: u32, +} + +impl Default for FooOptions { + fn default() -> Self { + Self { optional_knob: None, defaulted_knob: 7, required_knob: 0 } + } +} + +#[test] +fn multi_struct_validate_round_trips_when_all_fields_supplied() { + // Build a MultiFooOptions by hand (avoiding clap parsing so the + // test doesn't depend on a full clap Command harness). + let multi = MultiFooOptions { + foo_optional_knob: Some(42), + foo_defaulted_knob: 3, + foo_required_knob: Some(100), + }; + + let opts = multi.validate().expect("validate should succeed with all fields set"); + assert_eq!(opts.optional_knob, Some(42)); + assert_eq!(opts.defaulted_knob, 3); + assert_eq!(opts.required_knob, 100); +} + +#[test] +fn multi_struct_validate_fails_when_required_missing() { + let multi = + MultiFooOptions { foo_optional_knob: None, foo_defaulted_knob: 7, foo_required_knob: None }; + + let err = multi.validate().expect_err("validate should fail when required field is None"); + let msg = format!("{err:#}"); + assert!( + msg.contains("--foo::required-knob"), + "expected '--foo::required-knob' in error, got: {msg}" + ); + assert!(msg.contains("required when foo is selected"), "got: {msg}"); +} + +#[test] +fn from_original_round_trips_field_values() { + let original = FooOptions { optional_knob: Some(11), defaulted_knob: 22, required_knob: 33 }; + let multi: MultiFooOptions = original.clone().into(); + assert_eq!(multi.foo_optional_knob, Some(11)); + assert_eq!(multi.foo_defaulted_knob, 22); + assert_eq!(multi.foo_required_knob, Some(33)); + + let back = multi.validate().expect("round-trip validate should succeed"); + assert_eq!(back, original); +} + +#[test] +fn multi_struct_parses_prefixed_flags_via_clap() { + // Use a wrapper struct with `#[command(flatten)]` because + // `Args`-derived structs can only be parsed via a parent + // `Parser`-derived command. + #[derive(Parser, Debug)] + struct Wrapper { + #[command(flatten)] + foo_opts: MultiFooOptions, + } + + let wrapper = Wrapper::try_parse_from([ + "test-prog", + "--foo::required-knob", + "5", + "--foo::optional-knob", + "9", + ]) + .expect("parse should succeed"); + let opts = wrapper.foo_opts.validate().expect("validate after parse"); + assert_eq!(opts.required_knob, 5); + assert_eq!(opts.optional_knob, Some(9)); + assert_eq!(opts.defaulted_knob, 7); // unchanged +} + +/// Coverage for the `default_value` (string-form) and `Vec` +/// (pass-through) handling extensions. Mirrors the structure of +/// `FooOptions` but exercises the two paths separately: +/// +/// * `string_defaulted` uses `default_value = "..."` (the string +/// form parsed via `value_parser`); the macro must treat this as +/// "defaulted" not "required". +/// * `repeated_paths` is a `Vec` — the macro must treat Vec +/// like Option: pass it through with its type intact, since +/// clap collects an omitted repeatable flag as an empty Vec. +#[multi_options("bar", "Bar Options")] +#[derive(Args, Debug, Clone, PartialEq)] +pub struct BarOptions { + /// Defaulted via the string form of `default_value`. + #[arg(long, default_value = "42")] + pub string_defaulted: u32, + + /// Repeatable Vec field. + #[arg(long, action = clap::ArgAction::Append)] + pub repeated_paths: Vec, +} + +impl Default for BarOptions { + fn default() -> Self { + Self { string_defaulted: 42, repeated_paths: Vec::new() } + } +} + +#[test] +fn macro_treats_default_value_string_form_as_defaulted_not_required() { + // No string_defaulted on the cli → uses the Default. + #[derive(Parser, Debug)] + struct Wrapper { + #[command(flatten)] + bar_opts: MultiBarOptions, + } + let wrapper = Wrapper::try_parse_from(["test-prog"]).expect("parse with all defaults"); + let opts = wrapper.bar_opts.validate().expect("validate"); + assert_eq!(opts.string_defaulted, 42, "string-form default should apply"); + assert!(opts.repeated_paths.is_empty(), "Vec default is empty"); +} + +#[test] +fn macro_treats_vec_t_as_repeating_pass_through() { + #[derive(Parser, Debug)] + struct Wrapper { + #[command(flatten)] + bar_opts: MultiBarOptions, + } + let wrapper = Wrapper::try_parse_from([ + "test-prog", + "--bar::repeated-paths", + "1", + "--bar::repeated-paths", + "2", + "--bar::repeated-paths", + "3", + ]) + .expect("parse with three --repeated-paths"); + let opts = wrapper.bar_opts.validate().expect("validate"); + assert_eq!(opts.repeated_paths, vec![1, 2, 3]); +} + +/// Coverage for four attribute-handling fixes: +/// +/// * `skip_with_expr` uses `#[arg(skip = expr)]`. The macro must use +/// the provided expression verbatim in `validate()` rather than +/// hard-wiring `BazOptions::default().skip_with_expr`. +/// * `bare_skipped` uses a bare `#[arg(skip)]` (no `= expr`). The +/// macro must fall back to `BazOptions::default().bare_skipped` +/// (the struct default), NOT `u32::default()`. +/// * `conditionally_required` carries `#[arg(required = true)]`. The +/// macro must strip the `required = ...` name-value form so the +/// generated Multi field (an `Option`) is not forced by clap +/// during parsing — required-ness is handled by `validate()`. +/// * `tmp_dirs` carries `#[arg(long = "tmp-dir")]`, a `long` override +/// that differs from the kebab of the field name. The macro must +/// honor the override so the Multi flag is `--baz::tmp-dir`, NOT +/// `--baz::tmp-dirs`. +#[multi_options("baz", "Baz Options")] +#[derive(Args, Debug, Clone, PartialEq)] +pub struct BazOptions { + /// Skipped field with an explicit expression value. Not exposed on + /// the CLI; `validate()` must use this expression, not `Default`. + #[arg(skip = 99u32)] + pub skip_with_expr: u32, + + /// Bare-skipped field. Not exposed on the CLI; `validate()` must + /// pull its value from `BazOptions::default()`, not `u32::default()`. + #[arg(skip)] + pub bare_skipped: u32, + + /// Required (non-default) field that also carries `required = true`. + /// The Multi side must wrap this as `Option` and not propagate + /// `required = true` to clap. + #[arg(long, required = true)] + pub conditionally_required: u32, + + /// Repeatable field with a `long` override (`--tmp-dir`, not + /// `--tmp-dirs`). The Multi side must derive `--baz::tmp-dir`. + #[arg(long = "tmp-dir", action = clap::ArgAction::Append)] + pub tmp_dirs: Vec, +} + +impl Default for BazOptions { + fn default() -> Self { + // `skip_with_expr`'s Default is deliberately non-zero (and different + // from the `skip = 99u32` expression) so the test proves the macro + // emits the skip expression, not this Default value. `bare_skipped`'s + // Default is deliberately non-zero (and different from `u32::default()` + // == 0) so the bare-skip test proves the macro pulls from this struct + // Default rather than the field type's Default. + Self { + skip_with_expr: 1, + bare_skipped: 77, + conditionally_required: 0, + tmp_dirs: Vec::new(), + } + } +} + +/// A parent command for the Baz fixture; skip fields are carried on the Multi +/// struct but never exposed as flags, so they take their declared value at parse +/// time. +#[derive(Parser, Debug)] +struct BazWrapper { + #[command(flatten)] + baz_opts: MultiBazOptions, +} + +#[test] +fn skip_with_expr_uses_provided_expression_not_default() { + // The skip expression is `99`, while `BazOptions::default().skip_with_expr` + // is `1`. Parsing must yield `99`, proving the expression is honored. + let parsed = BazWrapper::try_parse_from(["test-prog", "--baz::conditionally-required", "5"]) + .expect("parse should succeed"); + let opts = parsed.baz_opts.validate().expect("validate should succeed"); + assert_eq!(opts.skip_with_expr, 99, "skip = 99 expression must be used verbatim"); + assert_eq!(opts.conditionally_required, 5); +} + +#[test] +fn bare_skip_falls_back_to_struct_default_not_type_default() { + // `bare_skipped` is `#[arg(skip)]` (no `= expr`), so the macro must populate + // it from `BazOptions::default().bare_skipped` (== 77), NOT `u32::default()` + // (== 0). + let parsed = BazWrapper::try_parse_from(["test-prog", "--baz::conditionally-required", "1"]) + .expect("parse should succeed"); + let opts = parsed.baz_opts.validate().expect("validate should succeed"); + assert_eq!( + opts.bare_skipped, 77, + "bare #[arg(skip)] must use the struct Default (77), not u32::default() (0)" + ); +} + +#[test] +fn skip_fields_are_carried_on_the_multi_struct_for_a_lossless_round_trip() { + // Values that match neither the skip expression nor the struct Default, so a + // re-derived value is distinguishable from a preserved one. + let original = BazOptions { + skip_with_expr: 1000, + bare_skipped: 2000, + conditionally_required: 3, + tmp_dirs: vec![1], + }; + let multi: MultiBazOptions = original.clone().into(); + assert_eq!(multi.validate().expect("validate"), original); +} + +#[test] +fn long_override_renames_multi_flag() { + // `tmp_dirs` carries `#[arg(long = "tmp-dir")]`, so the Multi flag must be + // `--baz::tmp-dir` (honoring the override), and `--baz::tmp-dirs` (the + // kebab of the field name) must be rejected. + #[derive(Parser, Debug)] + struct Wrapper { + #[command(flatten)] + baz_opts: MultiBazOptions, + } + + // The overridden flag name parses. + let wrapper = Wrapper::try_parse_from([ + "test-prog", + "--baz::conditionally-required", + "1", + "--baz::tmp-dir", + "3", + "--baz::tmp-dir", + "4", + ]) + .expect("--baz::tmp-dir should parse"); + let opts = wrapper.baz_opts.validate().expect("validate after parse"); + assert_eq!(opts.tmp_dirs, vec![3, 4]); + + // The field-name kebab (`--baz::tmp-dirs`) must NOT be a valid flag, proving + // the macro honored the `long = "tmp-dir"` override. Assert clap rejects it + // as an *unknown* argument specifically — a bare `is_err()` would also be + // satisfied by an unrelated parse failure. + let err = Wrapper::try_parse_from([ + "test-prog", + "--baz::conditionally-required", + "1", + "--baz::tmp-dirs", + "3", + ]) + .expect_err("--baz::tmp-dirs must be rejected (override renamed it to --baz::tmp-dir)"); + assert_eq!( + err.kind(), + clap::error::ErrorKind::UnknownArgument, + "expected an unknown-argument rejection, got {:?}: {err}", + err.kind() + ); +} + +#[test] +fn required_name_value_does_not_force_clap_requirement() { + // `required = true` on the original field must be stripped: parsing + // the Multi side WITHOUT the prefixed flag must succeed (clap must not + // enforce it), and the conditional requirement is surfaced by validate(). + #[derive(Parser, Debug)] + struct Wrapper { + #[command(flatten)] + baz_opts: MultiBazOptions, + } + + // No --baz::conditionally-required supplied: parse must succeed because + // the Multi field is Option and `required = true` was stripped. + let wrapper = + Wrapper::try_parse_from(["test-prog"]).expect("parse should succeed without required flag"); + let err = wrapper + .baz_opts + .validate() + .expect_err("validate should fail when conditionally-required is missing"); + let msg = format!("{err:#}"); + assert!( + msg.contains("--baz::conditionally-required"), + "expected '--baz::conditionally-required' in error, got: {msg}" + ); + + // Supplying it parses and validates. + let wrapper = Wrapper::try_parse_from(["test-prog", "--baz::conditionally-required", "7"]) + .expect("parse should succeed when flag supplied"); + let opts = wrapper.baz_opts.validate().expect("validate after parse"); + assert_eq!(opts.conditionally_required, 7); + assert_eq!(opts.skip_with_expr, 99); +} + +// ───────────────────────────────────────────────────────────────────────────── +// Call-form (`Meta::List`) passthrough +// ───────────────────────────────────────────────────────────────────────────── + +/// `CALL_FORM_SENSITIVE_ARG_KEYS` rejects the `key(value)` call form only for the keys the +/// macro classifies (`long`, `short`, `required`, `skip`, the alias keys and every +/// `default_value*`). Every other call-form key is preserved verbatim onto the Multi struct — +/// this pins that it actually reaches clap and takes effect, rather than being silently dropped. +#[multi_options("qux", "Qux Options")] +#[derive(Args, Debug, Clone, PartialEq)] +pub struct QuxOptions { + /// Range-validated via the call form `value_parser(...)`, which arrives as a + /// `Meta::List` and is passed through untouched. + #[arg(long, default_value_t = 3, value_parser(clap::value_parser!(u32).range(0..=9)))] + pub bounded: u32, +} + +impl Default for QuxOptions { + fn default() -> Self { + Self { bounded: 3 } + } +} + +#[test] +fn call_form_value_parser_is_preserved_and_enforced_on_the_multi_struct() { + #[derive(Parser, Debug)] + struct Wrapper { + #[command(flatten)] + qux_opts: MultiQuxOptions, + } + + // In range: parses and round-trips. + let wrapper = Wrapper::try_parse_from(["test-prog", "--qux::bounded", "9"]) + .expect("in-range value should parse"); + assert_eq!(wrapper.qux_opts.validate().expect("validate").bounded, 9); + + // Out of range: the passed-through `value_parser` must still reject it. If the + // `Meta::List` were dropped during re-exposure, clap would happily accept 10. + let err = Wrapper::try_parse_from(["test-prog", "--qux::bounded", "10"]) + .expect_err("out-of-range value must be rejected by the preserved value_parser"); + assert_eq!( + err.kind(), + clap::error::ErrorKind::ValueValidation, + "10 must be rejected by the preserved value_parser, not merely mentioned in some other \ + error: {err}" + ); + + // The default survives too. + let wrapper = Wrapper::try_parse_from(["test-prog"]).expect("defaults should parse"); + assert_eq!(wrapper.qux_opts.validate().expect("validate").bounded, 3); +} diff --git a/crates/fgumi-cli-macros/tests/ui/bad_prefix.rs b/crates/fgumi-cli-macros/tests/ui/bad_prefix.rs new file mode 100644 index 000000000..194366db1 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/bad_prefix.rs @@ -0,0 +1,20 @@ +//! The prefix is spliced into both the flag name (`--::`) and the +//! generated field identifier (`_`), so a non-identifier prefix +//! must be rejected at the literal rather than panicking inside `format_ident!`. +use fgumi_cli_macros::multi_options; + +#[multi_options("", "Probe Options")] +#[derive(clap::Args, Debug, Clone)] +pub struct EmptyPrefixOpts { + #[arg(long)] + pub value: u32, +} + +#[multi_options("2fast", "Probe Options")] +#[derive(clap::Args, Debug, Clone)] +pub struct LeadingDigitOpts { + #[arg(long)] + pub value: u32, +} + +fn main() {} diff --git a/crates/fgumi-cli-macros/tests/ui/bad_prefix.stderr b/crates/fgumi-cli-macros/tests/ui/bad_prefix.stderr new file mode 100644 index 000000000..13f8bd27b --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/bad_prefix.stderr @@ -0,0 +1,11 @@ +error: multi_options: the prefix must not be empty (it would generate flags named `--::`) + --> tests/ui/bad_prefix.rs:6:17 + | +6 | #[multi_options("", "Probe Options")] + | ^^ + +error: multi_options: the prefix must start with an ASCII letter, got `2fast` — it also becomes the leading segment of the generated field identifier `_` + --> tests/ui/bad_prefix.rs:13:17 + | +13 | #[multi_options("2fast", "Probe Options")] + | ^^^^^^^ diff --git a/crates/fgumi-cli-macros/tests/ui/call_form_long.rs b/crates/fgumi-cli-macros/tests/ui/call_form_long.rs new file mode 100644 index 000000000..ade9030b4 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/call_form_long.rs @@ -0,0 +1,12 @@ +//! clap's `key(value)` call form arrives as `Meta::List` and slips past every classifier, +//! so it must be rejected for the keys the macro rewrites. +use fgumi_cli_macros::multi_options; + +#[multi_options("probe", "Probe Options")] +#[derive(clap::Args, Debug, Clone)] +pub struct CallFormOpts { + #[arg(long("renamed"))] + pub value: u32, +} + +fn main() {} diff --git a/crates/fgumi-cli-macros/tests/ui/call_form_long.stderr b/crates/fgumi-cli-macros/tests/ui/call_form_long.stderr new file mode 100644 index 000000000..91cb0e509 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/call_form_long.stderr @@ -0,0 +1,6 @@ +error: multi_options: field `value` uses the call form #[arg(long(…))]. Use the `long = …` name-value form (or bare `long`) — the macro only classifies those spellings and would mishandle the call form. + --> tests/ui/call_form_long.rs:8:5 + | +8 | / #[arg(long("renamed"))] +9 | | pub value: u32, + | |__________________^ diff --git a/crates/fgumi-cli-macros/tests/ui/command_flatten.rs b/crates/fgumi-cli-macros/tests/ui/command_flatten.rs new file mode 100644 index 000000000..21541a8fd --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/command_flatten.rs @@ -0,0 +1,18 @@ +//! Re-prefixing a clap-flattened struct's fields is out of scope, so +//! `#[command(flatten)]` must fail the build rather than silently drop the nesting. +use fgumi_cli_macros::multi_options; + +#[derive(clap::Args, Debug, Clone)] +pub struct Inner { + #[arg(long)] + pub inner_value: u32, +} + +#[multi_options("probe", "Probe Options")] +#[derive(clap::Args, Debug, Clone)] +pub struct FlattenOpts { + #[command(flatten)] + pub inner: Inner, +} + +fn main() {} diff --git a/crates/fgumi-cli-macros/tests/ui/command_flatten.stderr b/crates/fgumi-cli-macros/tests/ui/command_flatten.stderr new file mode 100644 index 000000000..09a5694b2 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/command_flatten.stderr @@ -0,0 +1,5 @@ +error: multi_options does not support #[command(flatten)] / #[command(subcommand)] on field `inner`. The nested struct's fields cannot be reached to prefix them; inline them directly. + --> tests/ui/command_flatten.rs:14:5 + | +14 | #[command(flatten)] + | ^^^^^^^^^^^^^^^^^^^ diff --git a/crates/fgumi-cli-macros/tests/ui/conditional_clap_attr.rs b/crates/fgumi-cli-macros/tests/ui/conditional_clap_attr.rs new file mode 100644 index 000000000..c7dc1a676 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/conditional_clap_attr.rs @@ -0,0 +1,13 @@ +//! The macro classifies fields from their literal `#[arg(...)]` attributes and +//! cannot evaluate a cfg predicate, so a clap attribute hidden behind +//! `#[cfg_attr]` would be ignored and the field silently misclassified. +use fgumi_cli_macros::multi_options; + +#[multi_options("probe", "Probe Options")] +#[derive(clap::Args, Debug, Clone)] +pub struct ConditionalOpts { + #[cfg_attr(unix, arg(long, default_value_t = 3))] + pub knob: u32, +} + +fn main() {} diff --git a/crates/fgumi-cli-macros/tests/ui/conditional_clap_attr.stderr b/crates/fgumi-cli-macros/tests/ui/conditional_clap_attr.stderr new file mode 100644 index 000000000..28ac1c420 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/conditional_clap_attr.stderr @@ -0,0 +1,5 @@ +error: multi_options: field `knob` hides a clap attribute behind #[cfg_attr(…, arg(…))]. The macro classifies fields from their literal #[arg(...)] attributes and cannot evaluate a cfg predicate, so this one would be ignored and the field misclassified. Apply #[cfg] to the field and write the #[arg(...)] unconditionally. + --> tests/ui/conditional_clap_attr.rs:9:5 + | +9 | #[cfg_attr(unix, arg(long, default_value_t = 3))] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/crates/fgumi-cli-macros/tests/ui/cross_reference_arg.rs b/crates/fgumi-cli-macros/tests/ui/cross_reference_arg.rs new file mode 100644 index 000000000..b04645297 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/cross_reference_arg.rs @@ -0,0 +1,14 @@ +//! `requires` names another argument by its unprefixed id, which dangles once the Multi +//! struct renames fields, so it must fail the build. +use fgumi_cli_macros::multi_options; + +#[multi_options("probe", "Probe Options")] +#[derive(clap::Args, Debug, Clone)] +pub struct CrossRefOpts { + #[arg(long)] + pub first: u32, + #[arg(long, requires = "first")] + pub second: u32, +} + +fn main() {} diff --git a/crates/fgumi-cli-macros/tests/ui/cross_reference_arg.stderr b/crates/fgumi-cli-macros/tests/ui/cross_reference_arg.stderr new file mode 100644 index 000000000..79dbda5db --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/cross_reference_arg.stderr @@ -0,0 +1,6 @@ +error: multi_options: field `second` uses #[arg(requires …)], which references another argument by its unprefixed id. The Multi struct renames fields to `_`, so that id would dangle and clap would panic when it builds the runall command. Enforce this coupling in the command's validate()/resolve() instead (see AlignerOptions::resolve). + --> tests/ui/cross_reference_arg.rs:10:5 + | +10 | / #[arg(long, requires = "first")] +11 | | pub second: u32, + | |___________________^ diff --git a/crates/fgumi-cli-macros/tests/ui/generic_struct.rs b/crates/fgumi-cli-macros/tests/ui/generic_struct.rs new file mode 100644 index 000000000..f8d3fcdb4 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/generic_struct.rs @@ -0,0 +1,20 @@ +//! The companion struct and both conversion impls are emitted without generic +//! parameters, so a generic options struct would expand into code that cannot +//! compile — with every error naming the type parameter rather than the macro. +//! +//! The bounds are chosen so clap's own derive is satisfied; the only diagnostic +//! left is the macro's. +use fgumi_cli_macros::multi_options; + +#[multi_options("probe", "Probe Options")] +#[derive(clap::Args, Debug, Clone)] +pub struct GenericOpts +where + T: Clone + Send + Sync + 'static + std::str::FromStr, + ::Err: std::error::Error + Send + Sync + 'static, +{ + #[arg(long)] + pub value: T, +} + +fn main() {} diff --git a/crates/fgumi-cli-macros/tests/ui/generic_struct.stderr b/crates/fgumi-cli-macros/tests/ui/generic_struct.stderr new file mode 100644 index 000000000..8557edc7a --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/generic_struct.stderr @@ -0,0 +1,5 @@ +error: multi_options does not support generic structs: the generated Multi struct and its conversions are emitted without generic parameters, so the expansion would not compile. Use a concrete options struct. + --> tests/ui/generic_struct.rs:11:23 + | +11 | pub struct GenericOpts + | ^^^ diff --git a/crates/fgumi-cli-macros/tests/ui/legacy_clap_attr.rs b/crates/fgumi-cli-macros/tests/ui/legacy_clap_attr.rs new file mode 100644 index 000000000..53217b854 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/legacy_clap_attr.rs @@ -0,0 +1,15 @@ +//! Every classifier keys on `#[arg(...)]`, so clap's legacy `#[clap(...)]` +//! spelling would be silently ignored — a `#[clap(skip)]` field would be exposed +//! as a required CLI flag. Reject it instead. +use fgumi_cli_macros::multi_options; + +#[multi_options("probe", "Probe Options")] +#[derive(clap::Args, Debug, Clone)] +pub struct LegacyOpts { + #[clap(skip)] + pub hidden: u32, + #[arg(long)] + pub value: u32, +} + +fn main() {} diff --git a/crates/fgumi-cli-macros/tests/ui/legacy_clap_attr.stderr b/crates/fgumi-cli-macros/tests/ui/legacy_clap_attr.stderr new file mode 100644 index 000000000..0b203837c --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/legacy_clap_attr.stderr @@ -0,0 +1,5 @@ +error: multi_options does not support the legacy #[clap(...)] / #[structopt(...)] spelling: every classifier keys on #[arg(...)] and #[command(...)], so this attribute would be silently ignored and the field misclassified. Use the #[arg(...)] / #[command(...)] spelling. + --> tests/ui/legacy_clap_attr.rs:9:5 + | +9 | #[clap(skip)] + | ^^^^^^^^^^^^^ diff --git a/crates/fgumi-cli-macros/tests/ui/non_struct.rs b/crates/fgumi-cli-macros/tests/ui/non_struct.rs new file mode 100644 index 000000000..996ef9254 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/non_struct.rs @@ -0,0 +1,11 @@ +//! `multi_options` generates a companion struct from named fields, so applying it +//! to any other item — an enum here — must fail the build. +use fgumi_cli_macros::multi_options; + +#[multi_options("probe", "Probe Options")] +pub enum NotAStruct { + First, + Second, +} + +fn main() {} diff --git a/crates/fgumi-cli-macros/tests/ui/non_struct.stderr b/crates/fgumi-cli-macros/tests/ui/non_struct.stderr new file mode 100644 index 000000000..a1e58a966 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/non_struct.stderr @@ -0,0 +1,8 @@ +error: multi_options only supports structs with named fields + --> tests/ui/non_struct.rs:6:1 + | +6 | / pub enum NotAStruct { +7 | | First, +8 | | Second, +9 | | } + | |_^ diff --git a/crates/fgumi-cli-macros/tests/ui/positional_field.rs b/crates/fgumi-cli-macros/tests/ui/positional_field.rs new file mode 100644 index 000000000..5639f8caa --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/positional_field.rs @@ -0,0 +1,15 @@ +//! A positional argument has no flag name to prefix, and clap panics outright when +//! one carries a `long` — which the generated companion always emits. Several +//! stages flattened into one runall command would have ambiguous positionals +//! anyway, so both spellings must fail the build. +use fgumi_cli_macros::multi_options; + +#[multi_options("probe", "Probe Options")] +#[derive(clap::Args, Debug, Clone)] +pub struct PositionalOpts { + pub implicit_positional: u32, + #[arg(index = 2)] + pub explicit_positional: u32, +} + +fn main() {} diff --git a/crates/fgumi-cli-macros/tests/ui/positional_field.stderr b/crates/fgumi-cli-macros/tests/ui/positional_field.stderr new file mode 100644 index 000000000..b56a28ead --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/positional_field.stderr @@ -0,0 +1,12 @@ +error: multi_options: field `implicit_positional` declares neither `long` nor `short`, so clap treats it as a positional argument. A positional has no flag name to prefix, and several stages flattened into one runall command would have mutually ambiguous positionals. Add an explicit `#[arg(long)]` (or `#[arg(long = "...")]`) to the field. + --> tests/ui/positional_field.rs:10:5 + | +10 | pub implicit_positional: u32, + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +error: multi_options: field `explicit_positional` uses #[arg(index …)], which makes it a positional argument. The Multi struct gives every field a prefixed `long`, and clap panics when a positional has one ("is a positional argument and can't have short or long name versions"). A positional cannot be namespaced per stage — expose it as a flag with `#[arg(long)]` instead. + --> tests/ui/positional_field.rs:11:5 + | +11 | / #[arg(index = 2)] +12 | | pub explicit_positional: u32, + | |________________________________^ diff --git a/crates/fgumi-cli-macros/tests/ui/private_visibility.rs b/crates/fgumi-cli-macros/tests/ui/private_visibility.rs new file mode 100644 index 000000000..49013fc00 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/private_visibility.rs @@ -0,0 +1,17 @@ +//! The generated companion inherits the annotated struct's visibility, so a +//! private options struct must not produce a `pub` companion that escapes its +//! module. +mod inner { + use fgumi_cli_macros::multi_options; + + #[multi_options("probe", "Probe Options")] + #[derive(clap::Args, Debug, Clone)] + struct PrivateOpts { + #[arg(long, default_value_t = 1)] + value: u32, + } +} + +fn main() { + let _ = inner::MultiPrivateOpts { probe_value: 1 }; +} diff --git a/crates/fgumi-cli-macros/tests/ui/private_visibility.stderr b/crates/fgumi-cli-macros/tests/ui/private_visibility.stderr new file mode 100644 index 000000000..1f1e6e378 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/private_visibility.stderr @@ -0,0 +1,12 @@ +error[E0603]: struct `MultiPrivateOpts` is private + --> tests/ui/private_visibility.rs:16:20 + | +16 | let _ = inner::MultiPrivateOpts { probe_value: 1 }; + | ^^^^^^^^^^^^^^^^ private struct + | +note: the struct `MultiPrivateOpts` is defined here + --> tests/ui/private_visibility.rs:7:5 + | + 7 | #[multi_options("probe", "Probe Options")] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + = note: this error originates in the attribute macro `multi_options` (in Nightly builds, run with -Z macro-backtrace for more info) diff --git a/crates/fgumi-cli-macros/tests/ui/struct_level_command.rs b/crates/fgumi-cli-macros/tests/ui/struct_level_command.rs new file mode 100644 index 000000000..9c0af1a02 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/struct_level_command.rs @@ -0,0 +1,14 @@ +//! Struct-level clap configuration is not carried onto the generated companion, +//! so it must fail the build rather than silently apply to only one of the two +//! commands. +use fgumi_cli_macros::multi_options; + +#[multi_options("probe", "Probe Options")] +#[derive(clap::Args, Debug, Clone)] +#[command(next_help_heading = "Somewhere Else")] +pub struct StructLevelOpts { + #[arg(long)] + pub value: u32, +} + +fn main() {} diff --git a/crates/fgumi-cli-macros/tests/ui/struct_level_command.stderr b/crates/fgumi-cli-macros/tests/ui/struct_level_command.stderr new file mode 100644 index 000000000..c159954d2 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/struct_level_command.stderr @@ -0,0 +1,5 @@ +error: multi_options does not carry a struct-level #[command(...)] onto the generated Multi struct, so the standalone and re-exposed commands would silently diverge. Move the setting onto the individual #[arg(...)] attributes, or drop it. + --> tests/ui/struct_level_command.rs:8:1 + | +8 | #[command(next_help_heading = "Somewhere Else")] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/crates/fgumi-cli-macros/tests/ui/tuple_struct.rs b/crates/fgumi-cli-macros/tests/ui/tuple_struct.rs new file mode 100644 index 000000000..3c45dc343 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/tuple_struct.rs @@ -0,0 +1,7 @@ +//! `multi_options` requires named fields; a tuple struct must fail the build. +use fgumi_cli_macros::multi_options; + +#[multi_options("probe", "Probe Options")] +pub struct TupleOpts(pub u32); + +fn main() {} diff --git a/crates/fgumi-cli-macros/tests/ui/tuple_struct.stderr b/crates/fgumi-cli-macros/tests/ui/tuple_struct.stderr new file mode 100644 index 000000000..c250af561 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/tuple_struct.stderr @@ -0,0 +1,5 @@ +error: multi_options only supports structs with named fields + --> tests/ui/tuple_struct.rs:5:21 + | +5 | pub struct TupleOpts(pub u32); + | ^^^^^^^^^ diff --git a/crates/fgumi-cli-macros/tests/ui/unenforceable_required.rs b/crates/fgumi-cli-macros/tests/ui/unenforceable_required.rs new file mode 100644 index 000000000..db1074ba8 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/unenforceable_required.rs @@ -0,0 +1,15 @@ +//! `validate()` enforces required-ness by observing absence, which a defaulted +//! field and a bare `bool` never exhibit — so the requirement would be silently +//! unenforceable on the re-exposed flag. +use fgumi_cli_macros::multi_options; + +#[multi_options("probe", "Probe Options")] +#[derive(clap::Args, Debug, Clone)] +pub struct UnenforceableOpts { + #[arg(long, default_value_t = 3, required = true)] + pub defaulted: u32, + #[arg(long, required = true)] + pub toggled: bool, +} + +fn main() {} diff --git a/crates/fgumi-cli-macros/tests/ui/unenforceable_required.stderr b/crates/fgumi-cli-macros/tests/ui/unenforceable_required.stderr new file mode 100644 index 000000000..894e36f72 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/unenforceable_required.stderr @@ -0,0 +1,13 @@ +error: multi_options: field `defaulted` combines #[arg(required …)] with a default value or a bare `bool`. The generated field always holds a value — clap applies a declared default to `Option` and `Vec` as well — so the Multi side cannot distinguish "not supplied" from "supplied the default" and the requirement would be silently unenforceable. Drop `required`, or drop the default. + --> tests/ui/unenforceable_required.rs:9:5 + | + 9 | / #[arg(long, default_value_t = 3, required = true)] +10 | | pub defaulted: u32, + | |______________________^ + +error: multi_options: field `toggled` combines #[arg(required …)] with a default value or a bare `bool`. The generated field always holds a value — clap applies a declared default to `Option` and `Vec` as well — so the Multi side cannot distinguish "not supplied" from "supplied the default" and the requirement would be silently unenforceable. Drop `required`, or drop the default. + --> tests/ui/unenforceable_required.rs:11:5 + | +11 | / #[arg(long, required = true)] +12 | | pub toggled: bool, + | |_____________________^ diff --git a/crates/fgumi-cli-macros/tests/ui/where_only.rs b/crates/fgumi-cli-macros/tests/ui/where_only.rs new file mode 100644 index 000000000..2e12d2c67 --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/where_only.rs @@ -0,0 +1,16 @@ +//! A `where` clause with no type parameters must still be rejected, and the +//! diagnostic must point at the clause: spanning on an empty `Generics` would +//! collapse to the call site. +use fgumi_cli_macros::multi_options; + +#[multi_options("probe", "Probe Options")] +#[derive(clap::Args, Debug, Clone)] +pub struct WhereOnlyOpts +where + u32: Clone, +{ + #[arg(long)] + pub value: u32, +} + +fn main() {} diff --git a/crates/fgumi-cli-macros/tests/ui/where_only.stderr b/crates/fgumi-cli-macros/tests/ui/where_only.stderr new file mode 100644 index 000000000..b0d697e8b --- /dev/null +++ b/crates/fgumi-cli-macros/tests/ui/where_only.stderr @@ -0,0 +1,6 @@ +error: multi_options does not support generic structs: the generated Multi struct and its conversions are emitted without generic parameters, so the expansion would not compile. Use a concrete options struct. + --> tests/ui/where_only.rs:9:1 + | + 9 | / where +10 | | u32: Clone, + | |_______________^ diff --git a/crates/fgumi-consensus/src/codec_caller.rs b/crates/fgumi-consensus/src/codec_caller.rs index ea295f8f0..ad7a86906 100644 --- a/crates/fgumi-consensus/src/codec_caller.rs +++ b/crates/fgumi-consensus/src/codec_caller.rs @@ -308,7 +308,7 @@ struct SingleStrandConsensus { /// Per-molecule duplex base tallies produced while building a duplex consensus. /// -/// These are returned rather than folded straight into [`CodecStats`] so the caller can +/// These are returned rather than folded straight into [`CodecConsensusStats`] so the caller can /// apply them only *after* the consensus record has been serialized. A molecule whose output /// fails to build — an over-long input-derived UMI makes `try_build_record` reject the read /// name — emits nothing, so its duplex bases must not be counted, exactly as diff --git a/crates/fgumi-fmt/Cargo.toml b/crates/fgumi-fmt/Cargo.toml new file mode 100644 index 000000000..b97a7e15c --- /dev/null +++ b/crates/fgumi-fmt/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "fgumi-fmt" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +description = "Human-readable formatting helpers (counts, durations) shared across the fgumi workspace" +repository.workspace = true +license.workspace = true + +[dependencies] +num-format = { workspace = true } + +[dev-dependencies] +rstest = { workspace = true } + +[lints.clippy] +pedantic = { level = "deny", priority = -1 } diff --git a/crates/fgumi-fmt/src/lib.rs b/crates/fgumi-fmt/src/lib.rs new file mode 100644 index 000000000..4cbe5d96c --- /dev/null +++ b/crates/fgumi-fmt/src/lib.rs @@ -0,0 +1,136 @@ +#![deny(unsafe_code)] + +//! Human-readable formatting helpers shared across the fgumi workspace. +//! +//! These are needed by crates at several layers that do not — and should not — depend on one +//! another, so they live in a leaf crate near the bottom of the graph. Anything that only +//! turns a number into a string for humans to read belongs here rather than in a crate that +//! would drag `clap`, `sysinfo`, or BAM I/O along with it. + +use num_format::{Locale, ToFormattedString}; +use std::time::Duration; + +/// Format an integer with comma thousands-separators (e.g. `1234567` → `"1,234,567"`). +/// +/// # Examples +/// +/// ``` +/// use fgumi_fmt::format_count; +/// +/// assert_eq!(format_count(1_234_567), "1,234,567"); +/// assert_eq!(format_count(123), "123"); +/// assert_eq!(format_count(0), "0"); +/// ``` +#[must_use] +pub fn format_count(n: u64) -> String { + n.to_formatted_string(&Locale::en) +} + +/// Formats a duration in human-readable form (e.g. `"45s"`, `"2m 15s"`, `"1h 30m"`). +/// +/// Renders at most two units (whole seconds below a minute, minutes-and-seconds below an +/// hour, hours-and-minutes above); the finer unit is dropped when it is zero. +/// +/// # Examples +/// +/// ``` +/// use fgumi_fmt::format_duration; +/// use std::time::Duration; +/// +/// assert_eq!(format_duration(Duration::from_secs(45)), "45s"); +/// assert_eq!(format_duration(Duration::from_secs(135)), "2m 15s"); +/// assert_eq!(format_duration(Duration::from_secs(5400)), "1h 30m"); +/// ``` +#[must_use] +pub fn format_duration(duration: Duration) -> String { + let secs = duration.as_secs(); + if secs < 60 { + format!("{secs}s") + } else if secs < 3600 { + let mins = secs / 60; + let remaining_secs = secs % 60; + if remaining_secs == 0 { format!("{mins}m") } else { format!("{mins}m {remaining_secs}s") } + } else { + let hours = secs / 3600; + let mins = (secs % 3600) / 60; + if mins == 0 { format!("{hours}h") } else { format!("{hours}h {mins}m") } + } +} + +/// Formats a throughput as items per second, falling back to items per minute below one +/// item per second. +/// +/// Durations under a millisecond are reported as `count` items/s rather than dividing by a +/// near-zero elapsed time, which would produce a meaningless rate. +/// +/// # Examples +/// +/// ``` +/// use fgumi_fmt::format_rate; +/// use std::time::Duration; +/// +/// assert_eq!(format_rate(600, Duration::from_secs(60)), "10 items/s"); +/// // Below one item per second the units switch to items/min. +/// assert_eq!(format_rate(30, Duration::from_secs(60)), "30.0 items/min"); +/// ``` +#[must_use] +#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation, clippy::cast_sign_loss)] +pub fn format_rate(count: u64, duration: Duration) -> String { + let secs = duration.as_secs_f64(); + if secs < 0.001 { + return format!("{} items/s", format_count(count)); + } + + let rate = count as f64 / secs; + if rate >= 1.0 { + format!("{} items/s", format_count(rate as u64)) + } else { + let items_per_min = count as f64 / (secs / 60.0); + format!("{items_per_min:.1} items/min") + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::zero(0, "0")] + #[case::single(1, "1")] + #[case::three_digits(123, "123")] + #[case::first_comma(1_234, "1,234")] + #[case::two_groups(1_234_567, "1,234,567")] + #[case::round_million(1_000_000, "1,000,000")] + #[case::max(u64::MAX, "18,446,744,073,709,551,615")] + fn format_count_inserts_commas(#[case] n: u64, #[case] expected: &str) { + assert_eq!(format_count(n), expected); + } + + #[rstest] + #[case::zero(0, "0s")] + #[case::sub_minute(45, "45s")] + #[case::exact_minute(60, "1m")] + #[case::minutes_seconds(135, "2m 15s")] + #[case::exact_hour(3600, "1h")] + #[case::hours_minutes(5400, "1h 30m")] + fn format_duration_human_readable(#[case] secs: u64, #[case] expected: &str) { + assert_eq!(format_duration(Duration::from_secs(secs)), expected); + } + + /// `format_rate` has three regimes: a sub-millisecond guard that avoids dividing by + /// ~zero, a normal items/s path, and an items/min path below one item per second. + #[rstest] + #[case::sub_millisecond_guard(5, Duration::from_micros(100), "5 items/s")] + #[case::exactly_one_per_second(1, Duration::from_secs(1), "1 items/s")] + #[case::thousands_per_second(1000, Duration::from_secs(1), "1,000 items/s")] + #[case::below_one_per_second(30, Duration::from_secs(60), "30.0 items/min")] + #[case::far_below_one_per_second(1, Duration::from_secs(120), "0.5 items/min")] + fn format_rate_switches_units_below_one_per_second( + #[case] count: u64, + #[case] duration: Duration, + #[case] expected: &str, + ) { + assert_eq!(format_rate(count, duration), expected); + } +} diff --git a/crates/fgumi-pipeline-core/Cargo.toml b/crates/fgumi-pipeline-core/Cargo.toml new file mode 100644 index 000000000..594424540 --- /dev/null +++ b/crates/fgumi-pipeline-core/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "fgumi-pipeline-core" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +description = "Typed-step pipeline framework core (steps, queues, reorder stage, scheduler) for fgumi" +repository.workspace = true +license.workspace = true + +[dependencies] +ahash = { workspace = true } +anyhow = { workspace = true } +crossbeam-queue = { workspace = true } +log = { workspace = true } +noodles = { workspace = true, features = ["sam"] } +parking_lot = { workspace = true } + +[dev-dependencies] +proptest = { workspace = true } +rstest = { workspace = true } +trybuild = { workspace = true } + +[lints.clippy] +pedantic = { level = "deny", priority = -1 } diff --git a/crates/fgumi-pipeline-core/src/builder.rs b/crates/fgumi-pipeline-core/src/builder.rs new file mode 100644 index 000000000..e834be0a0 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/builder.rs @@ -0,0 +1,4058 @@ +//! `Pipeline`, `PipelineBuilder`, `Chain`, `MultiChain2/3/4`, `BuildError`. +//! +//! The builder uses `RefCell` for the steps + graph so multiple `Chain` +//! handles (one per branch of a multi-output step) can coexist as `&self`. +//! Move semantics on `Chain` enforce single-consumer per branch at the +//! type-system level; the `ChainGraph` performs the runtime all-wired check. + +use crate::liveness::LivenessCounter; +use std::cell::RefCell; +use std::marker::PhantomData; +use std::sync::Arc; + +use super::erased::{ErasedStep, TypedStep, TypedStep2}; +use super::item::HeapSize; +use super::outputs::{Single, StepOutputs}; +use super::runtime::stats::PipelineStats; +use super::signal::{CancelHandle, PipelineSignal}; +use super::step::{Step, Step2}; +use super::topology::{BranchIdx, ChainGraph, StepIdx}; + +/// Errors from `PipelineBuilder::build()`. +#[derive(Debug)] +pub enum BuildError { + UnwiredOutput { + step: &'static str, + branch: &'static str, + }, + Empty, + /// A step whose `Input = ()` (a source) has an edge wired INTO it. Its + /// input is implicit, so chain-context construction hands it a dummy unit + /// input handle and nothing ever pops the wired edge — the producer's + /// output would be silently discarded. + WiredIntoSource { + step: &'static str, + producer: &'static str, + }, +} + +impl std::fmt::Display for BuildError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::UnwiredOutput { step, branch } => { + write!(f, "step {step:?} has unwired output branch {branch:?}") + } + Self::Empty => write!(f, "pipeline has no steps"), + Self::WiredIntoSource { step, producer } => write!( + f, + "step {step:?} is a source (Input = ()) but {producer:?} is wired into it; \ + a source's input is implicit, so those items would never be consumed" + ), + } + } +} + +impl std::error::Error for BuildError {} + +/// How much pipeline instrumentation to collect. Off by default (zero overhead: +/// queues stay non-instrumented and no sampler thread spawns). Each higher level +/// is a superset of the one below. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum InstrumentationLevel { + /// No edge instrumentation. The pool hot path takes its existing zero-cost + /// route; this is the production default. + #[default] + Off, + /// Per-edge throughput/occupancy/latency counters + the occupancy sampler + + /// the end-of-run edge table and bottleneck verdict. + Summary, + /// `Summary` plus a per-tick occupancy/throughput timeline TSV. + Timeline, + /// `Timeline` plus direct dwell/park-time latency (Detached edges). + Deep, +} + +impl InstrumentationLevel { + /// `true` for any level above `Off` (edge metrics are collected). + #[must_use] + pub fn is_on(self) -> bool { + !matches!(self, Self::Off) + } + + /// `true` if the background occupancy sampler should run. + #[must_use] + pub fn samples(self) -> bool { + self.is_on() + } + + /// `true` if the per-tick timeline TSV should be written. + #[must_use] + pub fn timeline(self) -> bool { + matches!(self, Self::Timeline | Self::Deep) + } + + /// `true` if direct dwell/park-time latency should be recorded. + #[must_use] + pub fn deep(self) -> bool { + matches!(self, Self::Deep) + } +} + +#[derive(Debug, Clone)] +pub struct PipelineConfig { + pub threads: usize, + /// Optional shared stats collector. Construct via `Pipeline::stats()`, + /// then pass into the config; the framework writes per-step counters + /// during the run, and the caller reads them after `run` returns. + /// `None` (the default) keeps the worker loop on its zero-cost path. + pub stats: Option>, + /// Deadlock-detection timeout in seconds. When > 0, `Pipeline::run` + /// spawns a background monitor that polls liveness every `timeout / 4` + /// seconds (clamped to ≥1s). If no step has progressed (or finished) for + /// `timeout` seconds, the monitor logs at warn level so the user has a + /// starting point for debugging the stall. `0` (default) disables the + /// monitor. + /// + /// `stats` is **optional** and independent of arming: liveness comes from + /// [`crate::liveness::LivenessCounter`], and a stats handle only adds the + /// per-step snapshot to the stall report. The helpers in + /// `commands/common.rs` auto-attach a stats handle when this is non-zero so + /// those reports are populated without callers pairing the two flags. + /// + /// Mirrors the legacy framework's `--deadlock-timeout` semantics + /// (default 10s, 0 = disabled). + pub deadlock_timeout_secs: u64, + /// Total byte budget for byte-bounded queues across the chain. + /// `None` keeps each queue at the per-step limit set by its + /// `QueueSpec::ByteBounded { limit_bytes }`. `Some(total)` + /// enables the queue-memory rebalancer: an initial pass evenly + /// distributes `total` across all byte-bounded queues, and a + /// background thread periodically reads queue fullness and + /// shifts budget toward consistently-full queues (producer-bound + /// bottlenecks) at the expense of consistently-empty ones. + /// + /// Floor of 1 MiB per queue is enforced regardless of the + /// rebalancer's decisions to prevent pathological starvation. + pub queue_memory_total: Option, + /// How much per-edge instrumentation to collect (default `Off` = zero + /// overhead). Threaded into queue construction (instrumented transports) and + /// the occupancy sampler. See [`InstrumentationLevel`]. + pub instrumentation: InstrumentationLevel, + /// Where to write the per-tick timeline TSV when + /// `instrumentation.timeline()`. `None` → a default `pipeline-trace.tsv`. + pub trace_path: Option, + /// Per-worker dispatch-order policy. The default + /// [`ChainOrderScheduler`](crate::runtime::ChainOrderScheduler) walks live + /// steps upstream-first (historical behaviour). A chain may opt into + /// [`DrainFirstScheduler`](crate::runtime::DrainFirstScheduler) to walk + /// downstream-first (drain buffered work before producing more) — e.g. the + /// sort chain, to overlap its serial boundary/key scan with the parallel + /// inflate instead of starving it behind inflate on the shared pool. + pub scheduler: Arc, +} + +impl Default for PipelineConfig { + fn default() -> Self { + Self { + threads: std::thread::available_parallelism().map_or(1, std::num::NonZero::get), + stats: None, + // Still disarmed by default — but for a different reason than + // before, and the remaining blocker is now the only one. + // + // It used to be unaffordable: liveness came from `PipelineStats`, + // so arming the monitor meant paying `Instant::now()` on every + // dispatch. That is fixed — liveness is now a sharded counter + // (`crate::liveness`), and the monitor's teardown no longer polls a + // flag in 25ms slices, which cost up to a full slice of dead time on + // every run (measured at +80% wall on a 27ms pipeline). Both costs + // are gone; benchmarking shows arming it is now free within noise. + // + // What still blocks flipping this to `DEFAULT_DEADLOCK_TIMEOUT_SECS` + // is `PipelineError::MonitorBlindTransport`: an armed monitor + // *rejects* any chain with a non-`ByteBounded` edge, because its + // stall verdict distinguishes "idle" from "wedged" by counting bytes + // in flight and cannot see an `Unbounded` or `CountBounded` queue. + // `Process2` uses `CountBounded`, so defaulting this on would fail + // those chains outright rather than merely watching them. Teaching + // the verdict to handle byte-blind edges is its own change. + // + // Both the arming and the transport requirement are properties of + // the *scheduled* path only: a run that fuses to a single thread + // returns before either, so it would neither gain the monitor nor + // be rejected for a byte-blind edge. It is bounded by the fused + // path's own stall budget instead. + deadlock_timeout_secs: 0, + queue_memory_total: None, + instrumentation: InstrumentationLevel::Off, + trace_path: None, + scheduler: Arc::new(super::runtime::ChainOrderScheduler), + } + } +} + +impl PipelineConfig { + /// Builder-style helper to attach a stats collector. + #[must_use] + pub fn with_stats(mut self, stats: Arc) -> Self { + self.stats = Some(stats); + self + } + + /// Builder-style helper to set the deadlock-detection timeout. + /// [`DEFAULT_DEADLOCK_TIMEOUT_SECS`] is the recommended value. + /// `0` disables it. + /// + /// **Scheduled path only.** A non-zero value arms the monitor, and arming + /// additionally requires every output transport to be `ByteBounded` — see + /// [`crate::signal::PipelineError::MonitorBlindTransport`]. A run that + /// fuses to a single thread returns before either the transport check or + /// the monitor spawn, so it accepts a non-zero timeout with `CountBounded` + /// or `Unbounded` edges and simply ignores it; the fused path is bounded by + /// its own stall budget instead. + /// + /// `stats` is **optional** and independent of arming: liveness comes from + /// [`crate::liveness::LivenessCounter`], and a stats handle only adds the + /// per-step snapshot to the stall report. + #[must_use] + pub fn with_deadlock_timeout(mut self, timeout_secs: u64) -> Self { + self.deadlock_timeout_secs = timeout_secs; + self + } + + /// Builder-style helper to set the per-worker dispatch-order policy. + /// Defaults to [`ChainOrderScheduler`](crate::runtime::ChainOrderScheduler) + /// (upstream-first); pass + /// [`DrainFirstScheduler`](crate::runtime::DrainFirstScheduler) to walk + /// downstream-first. + #[must_use] + pub fn with_scheduler(mut self, scheduler: Arc) -> Self { + self.scheduler = scheduler; + self + } + + /// Builder-style helper to set the total queue-memory budget. + /// `None` keeps per-step `ByteBounded` limits at their static + /// values; `Some(total)` enables the rebalancer. + #[must_use] + pub fn with_queue_memory_total(mut self, total: Option) -> Self { + self.queue_memory_total = total; + self + } + + /// Builder-style helper to set the instrumentation level (default `Off`). + #[must_use] + pub fn with_instrumentation(mut self, level: InstrumentationLevel) -> Self { + self.instrumentation = level; + self + } +} + +pub struct PipelineBuilder { + inner: RefCell, +} + +struct BuilderInner { + steps: Vec>, + graph: ChainGraph, +} + +impl Default for PipelineBuilder { + fn default() -> Self { + Self::new() + } +} + +impl PipelineBuilder { + #[must_use] + pub fn new() -> Self { + Self { inner: RefCell::new(BuilderInner { steps: Vec::new(), graph: ChainGraph::new() }) } + } + + /// Add the first chain link. Requires `step: Step` (i.e., a source). + #[must_use = "pipeline branches must be wired to a sink"] + pub fn chain(&self, step: S) -> Chain<'_, S::Outputs> + where + S: Step, + { + let mut inner = self.inner.borrow_mut(); + // Input arity 0, not `register_step`'s default of 1: a source's input is + // implicit, so the graph must reject any attempt to wire an edge INTO + // it. See `append_source` for the failure the default lets through. + let producer = + inner.graph.register_step_with_input_arity(step.profile().name, S::Outputs::arity(), 0); + inner.steps.push(Box::new(TypedStep::new(step))); + + Chain { builder: self, producer, branch: BranchIdx(0), _phantom: PhantomData } + } + + /// Add a source step (first link) without requiring the typed `Chain` + /// return value. Used by the `ChainBuilder` in the parent `fgumi` crate's + /// `pipeline::chains` to accumulate steps across `add_source` / + /// `add_` / `add_sink` method calls. + /// + /// Returns `(StepIdx, BranchIdx(0))` for the newly registered step. + /// The caller tracks this tail and passes it to + /// [`Self::append_step`] for all subsequent steps. + /// + /// This deliberately bypasses the typed `Chain<'_, S::Outputs>` API. + /// Type correctness is NOT validated at [`Self::build`] time — `build()` + /// only checks that the chain is non-empty (rejecting a zero-step chain with + /// `BuildError::Empty`) and that every output branch is wired. A mis-typed step + /// sequence (e.g. a step whose `Input=A` consumes an output of type `B`) + /// builds successfully and panics at the first dispatch in + /// `TypedStep::resolve_input` with "input handle downcast failed — + /// chain topology invariant". That panic is loud and immediate but it + /// is a runtime check, not a compile-time or build-time one. Callers + /// are responsible for maintaining type correctness. `pub` so the + /// `chains` layer in the parent `fgumi` crate can drive incremental + /// assembly across the crate boundary. + pub fn append_source(&self, step: S) -> (StepIdx, BranchIdx) + where + S: Step, + { + let mut inner = self.inner.borrow_mut(); + // Input arity 0 — see `chain`. Leaving this at `register_step`'s default + // of 1 makes a source look like it has a free input slot: because + // `HeapSize for ()` exists, a producer whose output shape is + // `Single<()>` type-checks against a `Step`, `wire_to_slot` + // accepts slot 0 against the bogus arity, `build()` reports the chain + // fully wired, and `build_chain_contexts_inner` still takes the + // `is_source()` branch and hands the step a dummy unit input handle — + // so the wired edge's items are never popped. Registering arity 0 makes + // the graph reject that edge at wire time instead. + let producer = + inner.graph.register_step_with_input_arity(step.profile().name, S::Outputs::arity(), 0); + inner.steps.push(Box::new(TypedStep::new(step))); + (producer, BranchIdx(0)) + } + + /// Append a step to the chain by wiring it to the current tail + /// `(prev_producer, prev_branch)`. Used by the parent crate's + /// `ChainBuilder` for the same reason as [`Self::append_source`] — + /// type-erased incremental chain assembly across method-boundary calls. + /// + /// Returns `(StepIdx, BranchIdx(0))` for the newly registered step. + pub fn append_step( + &self, + step: S, + prev: (StepIdx, BranchIdx), + ) -> (StepIdx, BranchIdx) { + let mut inner = self.inner.borrow_mut(); + let consumer = inner.graph.register_step(step.profile().name, S::Outputs::arity()); + inner.graph.wire(prev.0, prev.1, consumer); + inner.steps.push(Box::new(TypedStep::new(step))); + (consumer, BranchIdx(0)) + } + + /// Append a two-input [`Step2`] step, wiring `prev_a` into input slot 0 + /// and `prev_b` into input slot 1. The type-erased counterpart of + /// [`MultiChain2Ordered::join`] — used by the parent crate's + /// `ChainBuilder::add_zipper` to wire the unmapped and mapped source + /// chains into the zipper-merge step across method-boundary calls. + /// + /// Returns `(StepIdx, BranchIdx(0))` for the newly registered step. + pub fn append_step2( + &self, + step: S, + prev_a: (StepIdx, BranchIdx), + prev_b: (StepIdx, BranchIdx), + ) -> (StepIdx, BranchIdx) { + let mut inner = self.inner.borrow_mut(); + let consumer = + inner.graph.register_step_with_input_arity(step.profile().name, S::Outputs::arity(), 2); + inner.graph.wire_to_slot(prev_a.0, prev_a.1, consumer, 0); + inner.graph.wire_to_slot(prev_b.0, prev_b.1, consumer, 1); + inner.steps.push(Box::new(TypedStep2::new(step))); + (consumer, BranchIdx(0)) + } + + /// Finalize the chain. Returns `Err(UnwiredOutput)` if any output branch + /// is dangling, `Err(Empty)` if the chain has zero steps. + /// + /// # Errors + /// + /// See `BuildError`. + pub fn build(self) -> Result { + let inner = self.inner.into_inner(); + if inner.steps.is_empty() { + return Err(BuildError::Empty); + } + if let Some((producer, _branch, branch_name)) = inner.graph.first_unwired() { + return Err(BuildError::UnwiredOutput { + step: inner.graph.step_name(producer), + branch: branch_name, + }); + } + // Registering sources with `input_arity = 0` makes `wire_to_slot` reject + // an edge into a source built through `PipelineBuilder::{chain, + // append_source}`. It does NOT cover a source reached as a *consumer*: + // `Chain::chain` / `append_step` register any consumer with arity 1, so + // a `Step` appended there wires cleanly. `is_source()` is + // `Input == ()` regardless of how the step was registered, so check it + // here — this is the one place every wiring path converges. + for (idx, step) in inner.steps.iter().enumerate() { + if !step.is_source() { + continue; + } + let consumer = super::topology::StepIdx(idx); + if let Some(producer) = inner.graph.first_producer_into(consumer) { + return Err(BuildError::WiredIntoSource { + step: inner.graph.step_name(consumer), + producer: inner.graph.step_name(producer), + }); + } + } + Ok(Pipeline { steps: inner.steps, graph: inner.graph, signal: PipelineSignal::new() }) + } +} + +/// In-progress chain handle. Each `.chain()` call consumes self and returns +/// a fresh `Chain` rooted at the new tail (Rust move semantics enforce +/// single-consumer at the type-system level). +#[must_use = "pipeline branches must be wired to a sink"] +pub struct Chain<'b, O> { + builder: &'b PipelineBuilder, + producer: StepIdx, + branch: BranchIdx, + _phantom: PhantomData O>, +} + +impl<'b, T: Send + HeapSize + 'static> Chain<'b, Single> { + /// Extend the chain with a step accepting `T`. + #[must_use = "pipeline branches must be wired to a sink"] + pub fn chain(self, step: S) -> Chain<'b, S::Outputs> + where + S: Step, + { + let mut inner = self.builder.inner.borrow_mut(); + let consumer = inner.graph.register_step(step.profile().name, S::Outputs::arity()); + inner.graph.wire(self.producer, self.branch, consumer); + inner.steps.push(Box::new(TypedStep::new(step))); + + Chain { + builder: self.builder, + producer: consumer, + branch: BranchIdx(0), + _phantom: PhantomData, + } + } +} + +impl<'b, T: Send + super::item::HeapSize + super::item::Ordered + 'static> + Chain<'b, super::outputs::OrderedBytesSingle> +{ + /// Extend the chain with a step accepting `T`. Mirror of + /// `Chain>::chain` for the heap-aware ordered output shape + /// used by Phase 3 BAM steps. + #[must_use = "pipeline branches must be wired to a sink"] + pub fn chain(self, step: S) -> Chain<'b, S::Outputs> + where + S: Step, + { + let mut inner = self.builder.inner.borrow_mut(); + let consumer = inner.graph.register_step(step.profile().name, S::Outputs::arity()); + inner.graph.wire(self.producer, self.branch, consumer); + inner.steps.push(Box::new(TypedStep::new(step))); + + Chain { + builder: self.builder, + producer: consumer, + branch: BranchIdx(0), + _phantom: PhantomData, + } + } +} + +impl Chain<'_, ()> { + /// Convenience to drop a sink-tail chain handle without `let _ = ...`. + /// Sinks have `Outputs = ()` (zero branches) so the all-wired check + /// never flags them; this method exists only to suppress the + /// `#[must_use]` warning at the call site when a chain naturally ends + /// at a sink. + pub fn into_sink_marker(self) { + let _ = (self.builder, self.producer, self.branch); + } +} + +impl<'b, A: Send + HeapSize + 'static, B: Send + HeapSize + 'static> Chain<'b, (A, B)> { + /// Convert a 2-output chain into per-branch sub-chains. + #[must_use = "all chain branches must be wired to a sink"] + pub fn into_multi(self) -> MultiChain2<'b, A, B> { + MultiChain2 { + b0: Chain { + builder: self.builder, + producer: self.producer, + branch: BranchIdx(0), + _phantom: PhantomData, + }, + b1: Chain { + builder: self.builder, + producer: self.producer, + branch: BranchIdx(1), + _phantom: PhantomData, + }, + _phantom: PhantomData, + } + } +} + +impl<'b, A, B> Chain<'b, super::outputs::OrderedBytesTuple2> +where + A: Send + super::item::HeapSize + super::item::Ordered + 'static, + B: Send + super::item::HeapSize + super::item::Ordered + 'static, +{ + /// Convert a 2-output ordered + byte-bounded chain into per-branch + /// sub-chains. Each branch is exposed as + /// `Chain>` so downstream chained steps see + /// the byte-aware ordered shape (matching what + /// `Chain>::chain` accepts). + #[must_use = "all chain branches must be wired to a sink"] + pub fn into_multi(self) -> MultiChain2Ordered<'b, A, B> { + MultiChain2Ordered { + b0: Chain { + builder: self.builder, + producer: self.producer, + branch: BranchIdx(0), + _phantom: PhantomData, + }, + b1: Chain { + builder: self.builder, + producer: self.producer, + branch: BranchIdx(1), + _phantom: PhantomData, + }, + _phantom: PhantomData, + } + } +} + +impl<'b, A, B, C> Chain<'b, super::outputs::OrderedBytesTuple3> +where + A: Send + super::item::HeapSize + super::item::Ordered + 'static, + B: Send + super::item::HeapSize + super::item::Ordered + 'static, + C: Send + super::item::HeapSize + super::item::Ordered + 'static, +{ + /// Convert a 3-output ordered + byte-bounded chain into per-branch + /// sub-chains, the 3-way counterpart of the + /// [`OrderedBytesTuple2`](super::outputs::OrderedBytesTuple2) impl above. + /// + /// Without this, a step whose `Outputs` is `OrderedBytesTuple3` cannot be + /// wired through the typed builder at all: `append_step` exposes only + /// branch 0, and branch wiring is not otherwise public — so the step type + /// checks, builds, and is simply unreachable. + #[must_use = "all chain branches must be wired to a sink"] + pub fn into_multi(self) -> MultiChain3Ordered<'b, A, B, C> { + MultiChain3Ordered { + b0: Chain { + builder: self.builder, + producer: self.producer, + branch: BranchIdx(0), + _phantom: PhantomData, + }, + b1: Chain { + builder: self.builder, + producer: self.producer, + branch: BranchIdx(1), + _phantom: PhantomData, + }, + b2: Chain { + builder: self.builder, + producer: self.producer, + branch: BranchIdx(2), + _phantom: PhantomData, + }, + _phantom: PhantomData, + } + } +} + +impl<'b, A, B, C> Chain<'b, (A, B, C)> +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, + C: Send + HeapSize + 'static, +{ + #[must_use = "all chain branches must be wired to a sink"] + pub fn into_multi(self) -> MultiChain3<'b, A, B, C> { + MultiChain3 { + b0: Chain { + builder: self.builder, + producer: self.producer, + branch: BranchIdx(0), + _phantom: PhantomData, + }, + b1: Chain { + builder: self.builder, + producer: self.producer, + branch: BranchIdx(1), + _phantom: PhantomData, + }, + b2: Chain { + builder: self.builder, + producer: self.producer, + branch: BranchIdx(2), + _phantom: PhantomData, + }, + _phantom: PhantomData, + } + } +} + +impl<'b, A, B, C, D> Chain<'b, (A, B, C, D)> +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, + C: Send + HeapSize + 'static, + D: Send + HeapSize + 'static, +{ + #[must_use = "all chain branches must be wired to a sink"] + pub fn into_multi(self) -> MultiChain4<'b, A, B, C, D> { + MultiChain4 { + b0: Chain { + builder: self.builder, + producer: self.producer, + branch: BranchIdx(0), + _phantom: PhantomData, + }, + b1: Chain { + builder: self.builder, + producer: self.producer, + branch: BranchIdx(1), + _phantom: PhantomData, + }, + b2: Chain { + builder: self.builder, + producer: self.producer, + branch: BranchIdx(2), + _phantom: PhantomData, + }, + b3: Chain { + builder: self.builder, + producer: self.producer, + branch: BranchIdx(3), + _phantom: PhantomData, + }, + _phantom: PhantomData, + } + } +} + +#[must_use = "all chain branches must be wired to a sink"] +pub struct MultiChain2<'b, A: Send + HeapSize + 'static, B: Send + HeapSize + 'static> { + pub b0: Chain<'b, Single>, + pub b1: Chain<'b, Single>, + pub(crate) _phantom: PhantomData<&'b PipelineBuilder>, +} + +impl<'b, A: Send + HeapSize + 'static, B: Send + HeapSize + 'static> MultiChain2<'b, A, B> { + /// Construct a `MultiChain2` from two independent source-side + /// chains that each produce `Single` / `Single`. Used when + /// two parallel source subchains converge at a `Step2` consumer + /// (e.g. zipper's mapped + unmapped BAM source subchains, AAM's + /// aligner-output + original-record-buffer subchains). + /// + /// Mirrors [`Chain::into_multi`]'s output-side counterpart: that + /// method takes one producer with two output branches and splits + /// it into a `MultiChain2`; this method takes two distinct + /// producers (each with one output branch) and packages them. + /// + /// # Panics + /// + /// Panics if `a` and `b` are anchored at different + /// `PipelineBuilder` instances (i.e. the framework would have to + /// wire across pipelines, which is meaningless). + #[must_use = "pipeline branches must be wired to a sink"] + pub fn from_chains(a: Chain<'b, Single>, b: Chain<'b, Single>) -> Self { + assert!( + std::ptr::eq(a.builder, b.builder), + "MultiChain2::from_chains: both chains must be anchored at the same builder", + ); + Self { b0: a, b1: b, _phantom: PhantomData } + } + + /// Join two parallel sub-chains into a single [`Step2`] consumer. + /// + /// Wires `b0` into the consumer's input slot 0 + /// ([`StepCtx2`](crate::step::StepCtx2)'s `a`) and `b1` into input slot 1 + /// ([`StepCtx2`](crate::step::StepCtx2)'s `b`), registering the consumer with input arity + /// 2 in the chain graph. Each branch's typed producer-side + /// [`OutputQueueSet`](crate::handles::OutputQueueSet) is later (at chain-run time) drained into + /// the consumer's + /// [`crate::handles::TwoInputHandles`] + /// via [`TypedStep2::build_two_input_handles`]. + /// + /// Returns a single-branch downstream [`Chain`] typed by the + /// joined step's `S::Outputs`, ready to chain further single-input + /// steps onto. + /// + /// # Type bounds + /// + /// `S: Step2` — the joined step's input + /// types must match the two upstream branch types exactly. + /// + /// # Panics + /// + /// `wire_to_slot`'s defensive panic fires if either upstream + /// branch was already wired (the `Chain` move semantics prevent + /// this in well-formed code). + pub fn join(self, step: S) -> Chain<'b, S::Outputs> + where + S: Step2, + { + let builder = self.b0.builder; + let p0 = self.b0.producer; + let br0 = self.b0.branch; + let p1 = self.b1.producer; + let br1 = self.b1.branch; + + let mut inner = builder.inner.borrow_mut(); + let consumer = + inner.graph.register_step_with_input_arity(step.profile().name, S::Outputs::arity(), 2); + inner.graph.wire_to_slot(p0, br0, consumer, 0); + inner.graph.wire_to_slot(p1, br1, consumer, 1); + inner.steps.push(Box::new(TypedStep2::new(step))); + + Chain { builder, producer: consumer, branch: BranchIdx(0), _phantom: PhantomData } + } +} + +/// Ordered-bytes variant of `MultiChain2`. Each branch is typed as +/// `Chain>` so downstream chained steps see the +/// byte-aware ordered representation (matching the queue topology +/// `OrderedBytesTuple2` actually constructed). +#[must_use = "all chain branches must be wired to a sink"] +pub struct MultiChain2Ordered<'b, A, B> +where + A: Send + super::item::HeapSize + super::item::Ordered + 'static, + B: Send + super::item::HeapSize + super::item::Ordered + 'static, +{ + pub b0: Chain<'b, super::outputs::OrderedBytesSingle>, + pub b1: Chain<'b, super::outputs::OrderedBytesSingle>, + pub(crate) _phantom: PhantomData<&'b PipelineBuilder>, +} + +impl<'b, A, B> MultiChain2Ordered<'b, A, B> +where + A: Send + super::item::HeapSize + super::item::Ordered + 'static, + B: Send + super::item::HeapSize + super::item::Ordered + 'static, +{ + /// Construct a `MultiChain2Ordered` from two independent + /// source-side chains that each produce `OrderedBytesSingle` / + /// `OrderedBytesSingle`. The ordered/byte-bounded counterpart + /// of [`MultiChain2::from_chains`] — used when two parallel + /// source subchains coming out of ordered, byte-bounded BAM/FASTQ + /// step libraries (decompress → boundaries → decode → group) need + /// to converge at a [`Step2`] consumer. + /// + /// Anchored at the same `PipelineBuilder` as both input chains; + /// panics otherwise. + /// + /// # Panics + /// + /// Panics if `a` and `b` are anchored at different + /// `PipelineBuilder` instances. + #[must_use = "pipeline branches must be wired to a sink"] + pub fn from_chains( + a: Chain<'b, super::outputs::OrderedBytesSingle>, + b: Chain<'b, super::outputs::OrderedBytesSingle>, + ) -> Self { + assert!( + std::ptr::eq(a.builder, b.builder), + "MultiChain2Ordered::from_chains: both chains must be anchored at the same builder", + ); + Self { b0: a, b1: b, _phantom: PhantomData } + } + + /// Join two parallel ordered/byte-bounded sub-chains into a single + /// [`Step2`] consumer. Ordered counterpart of + /// [`MultiChain2::join`]: wires `b0` into the consumer's input + /// slot 0 ([`StepCtx2`](crate::step::StepCtx2)'s `a`) and `b1` into input slot 1 + /// ([`StepCtx2`](crate::step::StepCtx2)'s `b`), registering the consumer with input arity 2. + /// + /// Returns a single-branch downstream [`Chain`] typed by the + /// joined step's `S::Outputs`. + /// + /// # Type bounds + /// + /// `S: Step2` — the joined step's input + /// types must match the two upstream branch element types exactly. + /// The framework only requires `HeapSize` on `Step2::InputA` / + /// `Step2::InputB`; the `Ordered` bound carried by + /// `OrderedBytesSingle` is upstream-side typing and does not flow + /// into the consumer's input handle (steps see plain + /// `InputHandle` regardless of upstream queue topology). + /// + /// # Panics + /// + /// `wire_to_slot`'s defensive panic fires if either upstream + /// branch was already wired (the `Chain` move semantics prevent + /// this in well-formed code). + pub fn join(self, step: S) -> Chain<'b, S::Outputs> + where + S: Step2, + { + let builder = self.b0.builder; + let p0 = self.b0.producer; + let br0 = self.b0.branch; + let p1 = self.b1.producer; + let br1 = self.b1.branch; + + let mut inner = builder.inner.borrow_mut(); + let consumer = + inner.graph.register_step_with_input_arity(step.profile().name, S::Outputs::arity(), 2); + inner.graph.wire_to_slot(p0, br0, consumer, 0); + inner.graph.wire_to_slot(p1, br1, consumer, 1); + inner.steps.push(Box::new(TypedStep2::new(step))); + + Chain { builder, producer: consumer, branch: BranchIdx(0), _phantom: PhantomData } + } +} + +#[must_use = "all chain branches must be wired to a sink"] +pub struct MultiChain3<'b, A, B, C> +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, + C: Send + HeapSize + 'static, +{ + pub b0: Chain<'b, Single>, + pub b1: Chain<'b, Single>, + pub b2: Chain<'b, Single>, + pub(crate) _phantom: PhantomData<&'b PipelineBuilder>, +} + +/// Per-branch sub-chains of an ordered + byte-bounded 3-way fan-out. The +/// 3-way counterpart of [`MultiChain2Ordered`]; each branch is exposed as +/// `Chain>` so downstream chained steps see the +/// byte-aware ordered shape. +#[must_use = "all chain branches must be wired to a sink"] +pub struct MultiChain3Ordered<'b, A, B, C> +where + A: Send + super::item::HeapSize + super::item::Ordered + 'static, + B: Send + super::item::HeapSize + super::item::Ordered + 'static, + C: Send + super::item::HeapSize + super::item::Ordered + 'static, +{ + pub b0: Chain<'b, super::outputs::OrderedBytesSingle>, + pub b1: Chain<'b, super::outputs::OrderedBytesSingle>, + pub b2: Chain<'b, super::outputs::OrderedBytesSingle>, + pub(crate) _phantom: PhantomData<&'b PipelineBuilder>, +} + +#[must_use = "all chain branches must be wired to a sink"] +pub struct MultiChain4<'b, A, B, C, D> +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, + C: Send + HeapSize + 'static, + D: Send + HeapSize + 'static, +{ + pub b0: Chain<'b, Single>, + pub b1: Chain<'b, Single>, + pub b2: Chain<'b, Single>, + pub b3: Chain<'b, Single>, + pub(crate) _phantom: PhantomData<&'b PipelineBuilder>, +} + +/// A built pipeline ready to run. +/// +/// `steps` is consumed by `Pipeline::run`; `graph` is read by `run` and +/// `dag()`; `signal` backs both `cancel_handle` and `run`'s outcome +/// plumbing. +pub struct Pipeline { + pub(crate) steps: Vec>, + pub(crate) graph: ChainGraph, + pub(crate) signal: Arc, +} + +impl Pipeline { + #[must_use] + pub fn builder() -> PipelineBuilder { + PipelineBuilder::new() + } + + #[must_use] + pub fn cancel_handle(&self) -> CancelHandle { + CancelHandle::from_signal(Arc::clone(&self.signal)) + } + + /// Construct a fresh `PipelineStats` collector sized to this pipeline's + /// chain. Wrap the returned `Arc` and pass it into `PipelineConfig::stats` + /// (or via `PipelineConfig::with_stats`) before calling `run`. Counters + /// can be read at any time after the run completes. + #[must_use] + pub fn stats(&self) -> Arc { + let names: Vec<&'static str> = self.steps.iter().map(|s| s.profile().name).collect(); + Arc::new(PipelineStats::new(names)) + } + + /// Render the chain shape as a multi-line debug string. Lists each step + /// in chain order with its profile (kind, sticky, branch count) and + /// the consumer for each output branch. Used for diagnostics and for + /// runall-style `--explain` output. + /// + /// The output isn't a stable serialization format — it's a developer- + /// readable summary, intended to be `println!`'d during debugging or + /// embedded in error messages. + #[must_use] + pub fn dag(&self) -> String { + use std::fmt::Write as _; + + let mut s = String::new(); + let _ = writeln!( + s, + "Pipeline DAG ({} step{}):", + self.steps.len(), + if self.steps.len() == 1 { "" } else { "s" } + ); + for (idx, step) in self.steps.iter().enumerate() { + let profile = step.profile(); + let n_branches = self.graph.branch_count(super::topology::StepIdx(idx)); + // Render the *effective* (post-collapse) ordering so the diagnostic + // matches the transport actually built. Single-input Serial / + // Exclusive producers collapse declared `ByOrdinal` / `ByItemOrdinal` + // to `None` (no reorder stage); `Step2` producers (input_arity == 2) + // and `Parallel` producers keep their declared ordering verbatim. + // Using the shared `effective_branch_orderings` helper keeps `dag()` + // and `build_output_set` from drifting. + let effective_orderings = if step.input_arity() == 2 { + profile.branch_ordering.clone() + } else { + super::erased::effective_branch_orderings(profile.kind, &profile.branch_ordering) + }; + let _ = write!( + s, + " [{idx}] {name:<24} {kind:?} sticky={sticky} branches={n_branches}", + idx = idx, + name = profile.name, + kind = profile.kind, + sticky = profile.sticky, + n_branches = n_branches, + ); + if n_branches == 0 { + let _ = writeln!(s, " (sink)"); + } else { + let _ = writeln!(s); + for branch_usize in 0..n_branches { + let branch = super::topology::BranchIdx(branch_usize); + let consumer = self.graph.consumer(super::topology::StepIdx(idx), branch); + let consumer_name = match consumer { + Some(c) => self.graph.step_name(c), + None => "", + }; + let queue_spec = profile + .output_queues + .get(branch_usize) + .copied() + .unwrap_or(super::queues::QueueSpec::Unbounded); + let ordering = effective_orderings + .get(branch_usize) + .copied() + .unwrap_or(super::reorder::BranchOrdering::None); + let _ = writeln!( + s, + " .{branch_usize}: {queue_spec:?} {ordering:?} → {consumer_name}", + ); + } + } + } + s + } + + /// Run the pipeline to completion. + /// + /// Spawns `config.threads` worker threads, runs each step's `try_run` + /// until every step has reported `Finished`, joins on completion, and + /// returns `Ok(())` on clean exit or an `Err(PipelineError)` if any step + /// returned `Err` or the caller cancelled via the `CancelHandle`. + /// + /// # Errors + /// + /// Returns `PipelineError::NotEnoughThreads` if the chain has more + /// `Exclusive` steps than `config.threads`. Returns `PipelineError::Io` + /// if any step's `try_run` returned `Err`. + /// Returns `PipelineError::Cancelled` if `cancel_handle().cancel()` was + /// called during the run. + /// + /// # Panics + /// + /// Panics if a worker thread panics — the panic is propagated via + /// `JoinHandle::join`. Worker panics indicate a framework or step bug + /// (e.g., a contract violation that triggered a `debug_assert!`). + #[allow(clippy::too_many_lines)] + pub fn run(self, config: PipelineConfig) -> Result<(), super::signal::PipelineError> { + use std::thread; + + use super::runtime::{ + StepDrainCounter, WorkerCore, assign_exclusive_owners, assign_sticky_owners, + build_chain_contexts, build_worker_storage, extract_detached_steps, + run_detached_driver, run_fused_single_thread, run_worker_loop, + should_fuse_single_thread, + }; + use super::step::{DetachedGroup, StepKind}; + use super::topology::StepIdx; + + let Self { mut steps, graph, signal } = self; + let n_threads = config.threads; + let stats_arc = config.stats; + let deadlock_timeout_secs = config.deadlock_timeout_secs; + let scheduler = Arc::clone(&config.scheduler); + assert!(n_threads > 0, "PipelineConfig::threads must be > 0"); + if let Some(stats) = stats_arc.as_ref() { + assert_eq!( + stats.n_steps(), + steps.len(), + "PipelineConfig::stats was sized for {} steps but pipeline has {}; \ + obtain the stats handle from `Pipeline::stats()` after `build()`", + stats.n_steps(), + steps.len() + ); + } + // Always-on liveness signal for the deadlock monitor. Sized for the + // worker pool plus a slot per possible driver thread, so every bumper + // gets its own cache line (see `crate::liveness`). Deliberately NOT + // gated on `stats`: the monitor must be armable without paying for + // per-dispatch timing, which is exactly what kept it disarmed before. + let liveness = Arc::new(LivenessCounter::new(n_threads + steps.len())); + + // 0. Fused single-thread fast path (issue #330). A single-source + // source→sink chain at one worker is driven inline over direct buffers, + // skipping the scheduler's round-robin poll / contention / reorder + // overhead (~2/3 of `try_run` calls at t=1 are otherwise wasted). + // Fan-out is allowed (e.g. the `--rejects` kept/rejects split); only + // two-input `Step2` merges (zipper, align) and `--threads ≥ 2` fall + // through to the scheduled worker pool below. The deadlock monitor and + // queue rebalancer are not spawned: a single-worker inline drive cannot + // deadlock, and its edges are driven producer-then-consumer in one pass, + // so there is no cross-worker imbalance to rebalance. Those edges DO + // carry each step's profiled byte bound, so `queue_memory_total` is + // handed to `run_fused_single_thread` and applied to them there — the + // fused contexts are built inside that call, not here. + // Instrumentation forces the scheduled path (see `should_fuse_single_thread`): + // the fused path has no per-edge metrics / occupancy sampler / verdict. + if should_fuse_single_thread(n_threads, config.instrumentation, &steps, &graph) { + log::debug!( + "Using fused single-thread pipeline ({} steps, direct buffers)", + steps.len() + ); + let result = run_fused_single_thread( + steps, + &graph, + &signal, + stats_arc.as_ref(), + config.queue_memory_total, + deadlock_timeout_secs, + ); + // Same end-of-run stats snapshot the scheduled path emits, so + // `--pipeline-stats` shows the fused chain's per-step counters. + if let Some(stats) = stats_arc.as_ref() { + let snapshot = stats.snapshot(); + log::info!("Pipeline end-of-run stats:"); + for line in format!("{snapshot}").lines() { + log::info!("{line}"); + } + } + return result; + } + + // 1. Assign Exclusive owners; bail if too many. + let owners = assign_exclusive_owners(&steps, n_threads)?; + + // 1a. Compute per-worker sticky-driven step (Exclusive sticky union + // with Serial+sticky+Affinity targeting). Indexed by worker id. + let sticky_owners = assign_sticky_owners(&steps, &owners, n_threads); + + // 2. Build per-step contexts (input + output handles). + let contexts = Arc::new(build_chain_contexts(&steps, &graph, config.instrumentation)); + + // 2-pre-monitor invariant: if the deadlock monitor will be armed, every + // output transport must be ByteBounded so `in_flight_bytes` can see a + // wedge on it. A CountBounded/Unbounded edge is invisible to the probe + // and would silently disable fail-fast on that edge. Checked here in + // every build (not debug-only — it returns a real `PipelineError`) while + // `steps` is still alive, before it is consumed by + // `build_worker_storage` below. Test chains use CountBounded/Unbounded + // but do not arm the monitor, so this only fires for a fail-fast run. + // + // Keyed on the timeout ALONE, deliberately: liveness now comes from + // `LivenessCounter`, so the monitor arms on a non-zero timeout whether + // or not a stats handle is attached. Gating this on `stats.is_some()` + // too — as it was when liveness was read out of `PipelineStats` — would + // let an armed, stats-less run start on a blind edge, where + // `in_flight_bytes` reports 0, `classify_stall` reads that as + // `Starving`, and the stall clock resets on every poll, so the wedge + // never reaches the fatal timeout. The arming condition here must track + // the one at the monitor spawn below. + if deadlock_timeout_secs > 0 { + ensure_monitor_visible_transports(&steps, &graph)?; + } + + // 2a. If a total queue-memory budget was supplied, evenly + // redistribute it across all byte-bounded queues now (before + // workers start) so the initial state matches the user's + // budget instead of the per-step defaults baked into each + // step's `QueueSpec::ByteBounded { limit_bytes }`. Floor at + // 1 MiB per queue to prevent zero-budget queues that would + // always reject pushes. + if let Some(total) = config.queue_memory_total { + apply_initial_queue_budget(&contexts.bounded_queues, total); + } + + // 3. Per-step drain counter — init N for Parallel, 1 otherwise. + let drain_counters: Vec> = steps + .iter() + .map(|step| { + let initial = match step.profile().kind { + StepKind::Parallel => n_threads, + // `Detached` runs on a single dedicated thread, so (like + // `Serial`/`Exclusive`) exactly one finisher closes its + // output edge. + StepKind::Serial | StepKind::Exclusive | StepKind::Detached => 1, + }; + StepDrainCounter::new(initial) + }) + .collect(); + + // 3a. Extract `Detached` steps' real instances for their dedicated + // threads, leaving same-position placeholders so `build_worker_storage` + // (which Skips Detached on every worker) and the `step_idx`-aligned + // `ChainContexts` stay correct. The contexts were already built from + // `&steps` above, so each extracted step's input/output handles live in + // `contexts[step_idx]`. Done before `build_worker_storage` consumes + // `steps`. Empty for every non-sort chain (nothing declares Detached). + let detached_steps = extract_detached_steps(&mut steps); + + // 4. Build per-worker step storage (consumes `steps`). + let mut worker_entries = build_worker_storage(steps, &owners, n_threads); + + let signal_arc = Arc::clone(&signal); + + // 4a. Optional deadlock-detection monitor. Spawns a watcher + // thread that periodically samples the stats snapshot; if no + // step's `progress + finished` counter has advanced for + // `deadlock_timeout_secs` seconds, it logs the snapshot at + // `warn` level so the user has a starting point. Polls every + // `max(1, deadlock_timeout_secs / 4)` seconds. + let (monitor_stop, monitor_handle) = match deadlock_timeout_secs { + n if n > 0 => { + let stop = Arc::new(StopSignal::default()); + let stop_clone = Arc::clone(&stop); + let liveness_monitor = Arc::clone(&liveness); + let stats_for_monitor = stats_arc.as_ref().map(Arc::clone); + let contexts_clone = Arc::clone(&contexts); + let signal_clone = Arc::clone(&signal); + let warn_timeout = std::time::Duration::from_secs(deadlock_timeout_secs); + let fatal_timeout = std::time::Duration::from_secs( + deadlock_timeout_secs.saturating_mul(DEADLOCK_FATAL_MULTIPLE), + ); + let poll_interval = + std::time::Duration::from_secs(deadlock_timeout_secs.max(4) / 4); + let handle = thread::Builder::new() + .name("fgumi-deadlock-monitor".to_string()) + .spawn(move || { + run_deadlock_monitor( + &stop_clone, + &liveness_monitor, + stats_for_monitor.as_ref(), + &contexts_clone, + &signal_clone, + warn_timeout, + fatal_timeout, + poll_interval, + ); + }) + .expect("failed to spawn deadlock monitor thread"); + (Some(stop), Some(handle)) + } + _ => (None, None), + }; + + // 4b. Optional queue-memory rebalancer. Spawns a watcher + // thread that periodically samples each registered queue's + // `current_bytes / limit_bytes` ratio and shifts budget + // toward consistently-full queues at the expense of + // consistently-empty ones. Total budget is preserved. + let (rebalancer_stop, rebalancer_handle) = + if config.queue_memory_total.is_some() && !contexts.bounded_queues.is_empty() { + let stop = Arc::new(StopSignal::default()); + let stop_clone = Arc::clone(&stop); + // Capture handles into a contiguous Vec — the monitor + // doesn't need step indices/names beyond debug logging. + let handles: Vec> = + contexts.bounded_queues.iter().map(|rq| Arc::clone(&rq.handle)).collect(); + let names: Vec<&'static str> = + contexts.bounded_queues.iter().map(|rq| rq.producer_step_name).collect(); + let handle = thread::Builder::new() + .name("fgumi-queue-rebalancer".to_string()) + .spawn(move || { + run_queue_rebalancer(&stop_clone, &handles, &names); + }) + .expect("failed to spawn queue rebalancer thread"); + (Some(stop), Some(handle)) + } else { + (None, None) + }; + + // 4c. Optional occupancy sampler (`--pipeline-trace`). Polls each + // byte-bounded edge's depth into its `EdgeMetrics` histogram on a fixed + // cadence. Runs whenever instrumentation samples AND there are edges to + // sample (`contexts.edges` is empty at level `Off` and on the fused + // single-thread fast path, so this stays inert there). Read-only over + // the live queues — never perturbs the worker hot path. + let (sampler_stop, sampler_handle) = + if config.instrumentation.samples() && !contexts.edges.is_empty() { + let stop = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let stop_clone = Arc::clone(&stop); + let contexts_clone = Arc::clone(&contexts); + // Timeline TSV path only when the level requests it. + let trace_path = if config.instrumentation.timeline() { + Some( + config + .trace_path + .clone() + .unwrap_or_else(|| std::path::PathBuf::from("pipeline-trace.tsv")), + ) + } else { + None + }; + let handle = thread::Builder::new() + .name("fgumi-occupancy-sampler".to_string()) + .spawn(move || { + crate::runtime::sampler::run_occupancy_sampler( + &stop_clone, + &contexts_clone.edges, + crate::runtime::sampler::DEFAULT_SAMPLE_INTERVAL, + trace_path, + ); + }) + .expect("failed to spawn occupancy sampler thread"); + (Some(stop), Some(handle)) + } else { + (None, None) + }; + + // 4d. Spawn one dedicated OS thread per `Detached` driver GROUP, in chain + // order, BEFORE the workers — same lifecycle slot as the monitor / + // rebalancer / sampler. Each thread drives its group with the SAME + // `run_worker_loop` the pool uses (via `run_detached_driver`), off the + // work-stealing pool, so a Detached step never consumes a pool worker + // slot (the "N + 2" threading). A group with one step is the legacy + // one-thread-per-detached-step case (`DetachedGroup::PerStep`); a + // `Shared` group co-locates several steps on one driver thread. Joined + // after the workers (step 6d). Empty for every non-sort chain, so this + // is a no-op there. + let detached_handles: Vec> = detached_steps + .into_iter() + .map(|group| { + let contexts_clone = Arc::clone(&contexts); + let signal_clone = Arc::clone(&signal_arc); + let drain_counters_clone: Vec> = + drain_counters.iter().map(Arc::clone).collect(); + let stats_clone = stats_arc.as_ref().map(Arc::clone); + let liveness_clone = Arc::clone(&liveness); + let thread_name = match group.label() { + DetachedGroup::Shared(label) => format!("fgumi-driver-{label}"), + DetachedGroup::PerStep => { + format!("fgumi-detached-{}", group.primary_name()) + } + }; + thread::Builder::new() + .name(thread_name) + .spawn(move || { + // Catch a driver-thread panic so we can signal + // cancellation before unwinding — a wedged pool worker + // parked on this group's (now-dead) edge only exits on + // `is_done()`. Re-raise after signalling so the join in + // step 6d still collects the payload. + if let Err(panic) = + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + run_detached_driver( + group, + &contexts_clone, + &drain_counters_clone, + &signal_clone, + stats_clone.as_ref(), + &liveness_clone, + ); + })) + { + signal_clone.cancel(); + std::panic::resume_unwind(panic); + } + }) + .expect("failed to spawn detached driver thread") + }) + .collect(); + + // Holds the first worker panic payload; re-raised after helper threads + // are cleaned up so monitor/rebalancer shutdown always executes. + let mut worker_panic: Option> = None; + + if n_threads == 1 { + // Single-threaded fast path: run the worker loop directly on + // the caller's thread instead of spawning + joining a fresh + // OS thread. The framework machinery (`build_worker_storage`, + // drain counters, stats) is identical to the multi-threaded + // path — only the spawn/join is skipped. + // + // Savings: thread-spawn-and-join (a few ms one-time on Apple + // Silicon / Linux), and the caller's thread name / TLS is + // preserved (matters for log correlation in some tools). + // + // Per-call cost gap vs the legacy single-threaded path is + // dominated by the framework's per-record book-keeping + // (queue byte tracking, ordinal allocation, drain checks), + // not the spawn — see commit message benchmarks. + let entries = worker_entries + .pop() + .expect("build_worker_storage with n_threads=1 returns one entry vec"); + debug_assert!(worker_entries.is_empty()); + let exclusive_owner = owners + .iter() + .enumerate() + .find_map(|(idx, &own)| if own == Some(0) { Some(StepIdx(idx)) } else { None }); + let sticky_owner = sticky_owners[0]; + let mut worker = WorkerCore::new(0, exclusive_owner, sticky_owner); + let mut entries_local = entries; + // Defer a panic on the single-threaded fast path the same way the + // multi-worker join loop does: capture the payload, signal + // cancellation, and let the common monitor/rebalancer shutdown run + // before re-raising at step 7. Without this, a worker-loop panic + // unwinds straight through the caller and leaks the helper threads. + // `AssertUnwindSafe` is sound: after a panic we never touch `worker` + // or `entries_local` again — the run is shutting down. + if let Err(panic) = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + run_worker_loop( + &mut worker, + &mut entries_local, + &contexts, + &drain_counters, + &signal_arc, + stats_arc.as_ref(), + &liveness, + scheduler.as_ref(), + ); + })) { + signal_arc.cancel(); + worker_panic = Some(panic); + } + } else { + // 5. Spawn worker threads. + let mut handles = Vec::with_capacity(n_threads); + for (worker_id, entries) in worker_entries.into_iter().enumerate() { + let exclusive_owner = owners.iter().enumerate().find_map(|(idx, &own)| { + if own == Some(worker_id) { Some(StepIdx(idx)) } else { None } + }); + let sticky_owner = sticky_owners[worker_id]; + + let contexts_clone = Arc::clone(&contexts); + let signal_clone = Arc::clone(&signal_arc); + let drain_counters_clone: Vec> = + drain_counters.iter().map(Arc::clone).collect(); + let stats_clone = stats_arc.as_ref().map(Arc::clone); + let liveness_clone = Arc::clone(&liveness); + let scheduler_clone = Arc::clone(&scheduler); + + let handle = thread::Builder::new() + .name(format!("fgumi-worker-{worker_id}")) + .spawn(move || { + let mut worker = WorkerCore::new(worker_id, exclusive_owner, sticky_owner); + let mut entries_local = entries; + // Catch a worker-loop panic so we can signal cancellation + // *before* unwinding. A peer parked in its retry loop on a + // full/empty queue only exits when it observes + // `signal.is_done()`; without an early `cancel()` here, the + // join loop below could block forever on an earlier, + // now-wedged worker and never reach this thread's panic. + // We re-raise after signalling so the join still collects + // the payload (preserving the deferred re-raise at step 7). + // `AssertUnwindSafe` is sound: on panic the run is tearing + // down and neither local is used again. + if let Err(panic) = + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + run_worker_loop( + &mut worker, + &mut entries_local, + &contexts_clone, + &drain_counters_clone, + &signal_clone, + stats_clone.as_ref(), + &liveness_clone, + scheduler_clone.as_ref(), + ); + })) + { + signal_clone.cancel(); + std::panic::resume_unwind(panic); + } + }) + .expect("failed to spawn worker thread"); + handles.push(handle); + } + + // 6. Join workers. Capture the first worker panic payload so cleanup + // can proceed; re-raise after monitor/rebalancer threads are stopped. + for h in handles { + if let Err(panic) = h.join() + && worker_panic.is_none() + { + worker_panic = Some(panic); + } + } + } + + // 6d. Join the Detached step threads (after the workers). The sort + // writer (Detached, consuming the pool's Compress output) and merge + // (Detached, producing to the pool's Serialize input) finish their final + // flush only after their pool peers have drained, so joining them here — + // after the worker join, before the monitor stop — captures the full + // chain completion (and any Detached-thread panic) in the same deferred + // re-raise path the workers use. A no-op when no step is Detached. + for h in detached_handles { + if let Err(panic) = h.join() + && worker_panic.is_none() + { + worker_panic = Some(panic); + } + } + + // 6a. Stop and join the deadlock monitor (if spawned). We + // signal stop *after* workers join so the monitor sees the + // final stats state and doesn't fire spurious warnings during + // normal pipeline shutdown (where steps stop progressing + // because they're done, not stuck). + if let Some(stop) = monitor_stop { + stop.stop(); + } + if let Some(handle) = monitor_handle { + let _ = handle.join(); + } + + // 6b. Stop and join the queue rebalancer (if spawned). + if let Some(stop) = rebalancer_stop { + stop.stop(); + } + if let Some(handle) = rebalancer_handle { + let _ = handle.join(); + } + + if let Some(stop) = sampler_stop { + stop.store(true, std::sync::atomic::Ordering::Relaxed); + } + if let Some(handle) = sampler_handle { + let _ = handle.join(); + } + + // 6c. End-of-run stats snapshot — emitted at info-level so + // bench profiling can see which step accumulated the most + // wall-clock time inside `try_run`. Only logged when stats + // were enabled (`with_stats(...)` on the config); cheap + // either way. + if let Some(stats) = stats_arc.as_ref() { + // When instrumentation is on, fold the per-edge table + bottleneck + // verdict into the snapshot (the edges live in `contexts`, only + // reachable here inside `run`); otherwise the edge-less snapshot. + let snapshot = if config.instrumentation.is_on() { + stats.snapshot_with_edges(&contexts.edges, stats.elapsed_ns()) + } else { + stats.snapshot() + }; + log::info!("Pipeline end-of-run stats:"); + for line in format!("{snapshot}").lines() { + log::info!("{line}"); + } + } + + // 7. Re-raise worker panics after helper threads are cleaned up so + // monitor/rebalancer shutdown code always executes. + if let Some(panic) = worker_panic { + std::panic::resume_unwind(panic); + } + + // 8. Surface error or cancellation. PipelineError isn't Clone + // (io::Error isn't Clone); `to_result` reconstructs the recorded + // outcome and, for an external cancel whose payload isn't yet visible + // to this thread, synthesizes `Cancelled` from the terminal state. + let _ = graph; + signal.to_result() + } +} + +/// Default multiple of the warn window (`--deadlock-timeout`) after which a +/// persistent stall *with work still in flight* is treated as a fatal wedge. +/// Diverges from the legacy single-window kill: a single huge dispatch +/// (busy-locus group, large sort merge) can legitimately flatten progress for +/// many seconds, so we only fail after the stall persists well past the warn +/// window. A real deadlock hangs forever, so waiting longer to be sure is free. +/// Recommended stall patience for the deadlock monitor, in seconds, for callers +/// arming it via [`PipelineConfig::with_deadlock_timeout`]. +/// +/// Matches the fused path's `DEFAULT_STALL_BUDGET` so both runtimes bound a +/// wedge alike, and for the same reason: this catches a PERMANENT wedge, not a +/// slow source, and the costs are asymmetric — err long. +/// +/// Not yet the `PipelineConfig::default()` value; see the note there. +pub const DEFAULT_DEADLOCK_TIMEOUT_SECS: u64 = 60; + +const DEADLOCK_FATAL_MULTIPLE: u64 = 6; + +/// Total bytes currently held across the **byte-bounded** transport queues and +/// their reorder overflow stashes. Non-zero means work is stuck on a byte-bounded +/// edge; zero means those edges are idle (e.g. waiting on a slow upstream pipe). +/// This is the signal that distinguishes a real wedge from upstream starvation +/// on byte-bounded edges. +/// +/// LIMITATION: only `ByteBounded` branches register a queue handle (see +/// `build_chain_contexts_inner`), so a `CountBounded`/`Unbounded` transport — and +/// any reorder stash on such a branch — would not contribute here, making the +/// monitor blind to a wedge living entirely on such an edge. +/// +/// This is safe because **no production pipeline edge uses +/// `CountBounded`/`Unbounded`**: every production step profile declares +/// `QueueSpec::ByteBounded`, so the probe covers every production transport. +/// (The sort spill→merge path is fully byte-bounded too — its deadlock-free +/// backpressure is the internal `SortMergeSlot` slot table, not a pipeline +/// `CountBounded` edge.) [`ensure_monitor_visible_transports`] enforces this +/// invariant in **every** build (it returns +/// [`crate::signal::PipelineError::MonitorBlindTransport`], not a debug assertion) whenever the +/// monitor is armed, so a future step that declares a monitor-blind transport on +/// a fail-fast pipeline fails at startup rather than silently losing the wedge +/// verdict. Extending the probe to count-based queues for full defense-in-depth +/// is tracked separately. +/// +/// SECOND LIMITATION: byte accounting is `T::heap_size()`-only — +/// `ByteBoundedQueue` never counts `size_of::()` — so a `ByteBounded` edge +/// holding items whose `heap_size()` is 0 also reports zero here, however many +/// are queued. A wedge stranding only such items therefore classifies as +/// [`StallVerdict::Starving`], which resets the stall clock on every poll, so +/// `deadlock_timeout_secs` can never fail it. Most production item types +/// (record batches, BGZF blocks) carry real heap payloads, but at least one does +/// not: `InflatedBlock` (the arena-ingest completion token, whose payload lives +/// in the shared arena rather than on its own heap) reports `heap_size() == 0` by +/// design, so a `ByteBounded` edge carrying only inflated-block tokens is +/// invisible to this probe. That edge is **not** wired into a monitored +/// production chain today (the arena front is exercised only by its own +/// tests/benches). Before it is, close the gap at the source — give +/// `InflatedBlock` a `size_of::()` base (as the sort control events already +/// do) — or have the monitor consult queue *occupancy*, not just bytes, on the +/// starvation branch. Until then this remains a latent, unreachable gap. +fn in_flight_bytes(contexts: &crate::runtime::contexts::ChainContexts) -> u64 { + contexts + .bounded_queues + .iter() + .map(|rq| { + rq.handle.current_bytes() + + rq.reorder_cap.as_ref().map_or(0, |r| r.current_buffer_bytes()) + }) + .sum() +} + +/// Return the first `(step_name, QueueSpec)` whose output transport is invisible +/// to the deadlock monitor's [`in_flight_bytes`] probe — i.e. a `CountBounded` +/// or `Unbounded` branch, which registers no byte-probe handle. `None` means +/// every output transport is monitor-visible (`ByteBounded`). +/// +/// Used by [`ensure_monitor_visible_transports`] to pin the "no production edge +/// is monitor-blind" invariant the [`in_flight_bytes`] doc relies on. Iterates +/// every branch the graph declares for each step (`0..branch_count`) rather than +/// just the explicit `output_queues` entries: a step may declare fewer specs than +/// it has branches, and `dag()`/context-building resolve those missing branches +/// to `QueueSpec::Unbounded`. Using the same `unwrap_or(Unbounded)` fallback here +/// keeps the guard from overlooking an implicit (and therefore monitor-blind) +/// Unbounded branch. +fn first_monitor_blind_transport( + steps: &[Box], + graph: &super::topology::ChainGraph, +) -> Option<(&'static str, super::queues::QueueSpec)> { + use super::queues::QueueSpec; + for (step_idx, step) in steps.iter().enumerate() { + let profile = step.profile(); + for branch in 0..graph.branch_count(super::topology::StepIdx(step_idx)) { + let spec = profile.output_queues.get(branch).copied().unwrap_or(QueueSpec::Unbounded); + match spec { + QueueSpec::ByteBounded { .. } => {} + QueueSpec::CountBounded { .. } | QueueSpec::Unbounded => { + return Some((profile.name, spec)); + } + } + } + } + None +} + +/// Invariant check (run only when the deadlock monitor is armed): every +/// production output transport must be `ByteBounded` so the [`in_flight_bytes`] +/// probe can see a wedge on it. A `CountBounded`/`Unbounded` edge would be +/// invisible to the monitor, silently disabling fail-fast on that edge — exactly +/// the blind spot this guard exists to catch. The framework still permits +/// `CountBounded`/`Unbounded` for `#[cfg(test)]` chains (which do not arm the +/// monitor), so this only fires for a real fail-fast pipeline. +/// +/// Returns a [`crate::signal::PipelineError::MonitorBlindTransport`] (rather than panicking or +/// being a debug-only check) so the guard runs in release builds — where the +/// blind spot actually matters — yet a misconfigured chain fails gracefully at +/// startup, consistent with the other build/run-time validations (e.g. +/// [`crate::signal::PipelineError::NotEnoughThreads`]) rather than crashing the process. +fn ensure_monitor_visible_transports( + steps: &[Box], + graph: &super::topology::ChainGraph, +) -> Result<(), super::signal::PipelineError> { + if let Some((name, spec)) = first_monitor_blind_transport(steps, graph) { + return Err(super::signal::PipelineError::MonitorBlindTransport { + step: name, + spec: format!("{spec:?}"), + }); + } + Ok(()) +} + +/// Background deadlock monitor body. Polls `stats` every `poll_interval`, +/// tracking the cumulative `progress + finished` counter across all steps. +/// +/// On a stall (no advance), it consults [`in_flight_bytes`] to tell a real +/// wedge from upstream starvation (mirrors legacy `check_deadlock_and_restore`): +/// - idle with nothing in flight → starvation, reset the clock, keep watching; +/// - stuck work past `warn_timeout` → `warn` snapshot, once per warn window; +/// - stuck work past `fatal_timeout` → record [`crate::signal::PipelineError::TimedOut`] and +/// `cancel()`, so workers observe `is_done()` and the run fails fast instead +/// of hanging forever. +/// +/// Exits when `stop` is set (workers joined) or the pipeline is already done. +#[allow(clippy::too_many_arguments)] // one monitor's worth of shared state; a struct would only rename it +fn run_deadlock_monitor( + stop: &Arc, + liveness: &Arc, + stats: Option<&Arc>, + contexts: &Arc, + signal: &Arc, + warn_timeout: std::time::Duration, + fatal_timeout: std::time::Duration, + poll_interval: std::time::Duration, +) { + let mut mon_state = StallMonitorState { + last_total: liveness.total(), + stall_start: std::time::Instant::now(), + last_warn: None, + }; + while !stop.is_stopped() { + sleep_until_stop(stop, poll_interval); + if stop.is_stopped() || signal.is_done() { + break; + } + let now = std::time::Instant::now(); + let now_total = liveness.total(); + let progressed = now_total != mon_state.last_total; + let stall_secs = now.duration_since(mon_state.stall_start).as_secs(); + let stuck = in_flight_bytes(contexts); + let verdict = classify_stall( + progressed, + stall_secs, + warn_timeout.as_secs(), + fatal_timeout.as_secs(), + stuck, + ); + if apply_stall_verdict( + verdict, + now, + now_total, + stall_secs, + stuck, + warn_timeout, + stats.map(std::convert::AsRef::as_ref), + signal, + &mut mon_state, + ) { + break; + } + } +} + +/// Rolling state the deadlock monitor carries across poll iterations. +struct StallMonitorState { + last_total: u64, + stall_start: std::time::Instant, + last_warn: Option, +} + +/// React to one classified [`StallVerdict`], mutating the rolling monitor +/// `mon_state` and performing the verdict's side effect (a warn snapshot, or +/// recording a fatal [`crate::signal::PipelineError::TimedOut`] + `cancel`). Extracted from +/// [`run_deadlock_monitor`]'s poll loop so each per-verdict action is +/// unit-testable without spawning a thread or waiting on wall-clock time — +/// mirroring the pure `classify_stall` / `in_flight_bytes` split. Returns +/// `true` when a fatal wedge was recorded and the monitor should stop. +#[allow(clippy::too_many_arguments)] +fn apply_stall_verdict( + verdict: StallVerdict, + now: std::time::Instant, + now_total: u64, + stall_secs: u64, + stuck: u64, + warn_timeout: std::time::Duration, + // Optional: the monitor no longer requires instrumentation to run, so an + // uninstrumented run still reports the stall — just without the per-step + // snapshot. Losing the snapshot is far better than losing the detection, + // which is what requiring `stats` used to cost. + stats: Option<&PipelineStats>, + signal: &PipelineSignal, + mon_state: &mut StallMonitorState, +) -> bool { + use super::signal::PipelineError; + match verdict { + StallVerdict::Progressing => { + mon_state.last_total = now_total; + mon_state.stall_start = now; + mon_state.last_warn = None; + } + StallVerdict::Starving => { + // Idle, waiting on a slow upstream — not a deadlock. Reset the + // stall clock so an idle gap never accumulates toward a fatal. + mon_state.stall_start = now; + mon_state.last_warn = None; + } + StallVerdict::Watching => {} + StallVerdict::Stalled => { + // Warn at most once per warn window: a long stall emits periodic + // diagnostics without spamming every poll. + if mon_state.last_warn.is_none_or(|w| now.duration_since(w) >= warn_timeout) { + log::warn!( + "Pipeline stall: no progress for {stall_secs}s with {stuck} bytes \ + still in flight. Snapshot follows." + ); + if let Some(stats) = stats { + let snapshot = stats.snapshot(); + for line in format!("{snapshot}").lines() { + log::warn!("{line}"); + } + } else { + log::warn!( + "(no per-step snapshot: PipelineConfig::stats is not set; \ + attach a stats handle to see which step is stuck)" + ); + } + mon_state.last_warn = Some(now); + } + } + StallVerdict::Wedged => { + log::error!( + "Pipeline deadlock: no progress for {stall_secs}s with {stuck} bytes \ + stuck in flight; failing the pipeline. Snapshot follows." + ); + if let Some(stats) = stats { + let snapshot = stats.snapshot(); + for line in format!("{snapshot}").lines() { + log::error!("{line}"); + } + } else { + log::error!( + "(no per-step snapshot: PipelineConfig::stats is not set; \ + attach a stats handle to see which step is stuck)" + ); + } + signal.record_error(PipelineError::TimedOut { stalled_secs: stall_secs }); + signal.cancel(); + return true; + } + } + false +} + +/// Sum of `progress_count + finished_count` across all steps. The +/// monitor uses this as a single scalar progress watermark; a change +/// means *some* step is making forward progress. +/// Sleep up to `dur`, returning early as soon as `stop` is set. Polls the +/// flag in short slices so a background helper thread (deadlock monitor, +/// queue rebalancer) exits within tens of milliseconds at teardown instead +/// of blocking the main thread's `join()` for a full poll interval after +/// the pipeline has already finished. A plain `thread::sleep(poll_interval)` +/// here adds a fixed dead-time tail (up to `poll_interval`) to every run — +/// negligible on long jobs but a large *relative* regression on short ones +/// (e.g. FASTQ extract), since the worker pool is already idle and waiting. +fn sleep_until_stop(stop: &StopSignal, dur: std::time::Duration) { + // A condvar, not a polled sleep. The previous shape slept in 25ms slices and + // re-checked a flag, so teardown waited up to a full slice for the monitor to + // notice — dead time on the critical path of *every* run, since `Pipeline::run` + // joins this thread before returning. + // + // That is not a rounding error at the short end. With the monitor armed by + // default, a 27ms pipeline measured +80% wall on the dispatch benchmark, and + // the excess was almost exactly one slice. Waking the sleeper directly takes + // that to zero, which is what makes arming the monitor by default affordable. + let (lock, cvar) = (&stop.stopped, &stop.waker); + let mut stopped = lock.lock().expect("stop mutex not poisoned"); + if *stopped { + return; + } + // `Instant::now() + dur` panics when the deadline is not representable (a + // caller could pass an enormous `deadlock_timeout_secs`). `Pipeline::run` + // discards this thread's join error, so a panic here would silently disarm + // the monitor while the run reports success. Guard with `checked_add`: an + // unrepresentable deadline means "effectively never", so wait untimed until + // `stop()` wakes us — the only teardown path that matters. + let Some(deadline) = std::time::Instant::now().checked_add(dur) else { + while !*stopped { + stopped = cvar.wait(stopped).expect("stop condvar not poisoned"); + } + return; + }; + // `wait_timeout` can wake spuriously; the loop re-checks both the flag and + // the deadline, so a spurious wake just resumes waiting for the remainder. + while !*stopped { + let Some(remaining) = deadline.checked_duration_since(std::time::Instant::now()) else { + return; + }; + let (guard, _) = cvar.wait_timeout(stopped, remaining).expect("stop condvar not poisoned"); + stopped = guard; + } +} + +/// Stop flag for a background helper thread, with a condvar so a waiting thread +/// wakes the instant it is set rather than on its next poll tick. +#[derive(Debug, Default)] +pub(crate) struct StopSignal { + stopped: std::sync::Mutex, + waker: std::sync::Condvar, +} + +impl StopSignal { + /// Whether `stop` has been called. + pub(crate) fn is_stopped(&self) -> bool { + *self.stopped.lock().expect("stop mutex not poisoned") + } + + /// Set the flag and wake the sleeper immediately. + pub(crate) fn stop(&self) { + *self.stopped.lock().expect("stop mutex not poisoned") = true; + self.waker.notify_all(); + } +} + +/// What the deadlock monitor should do after one poll. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum StallVerdict { + /// Global progress advanced since the last poll — reset and keep watching. + Progressing, + /// No progress, but nothing is in flight: the pipeline is idle waiting on + /// a slow upstream (e.g. a stdin pipe), not deadlocked. Reset the stall + /// clock and keep watching. Mirrors legacy `check_deadlock_and_restore`'s + /// starvation guard (an empty pipeline is never a deadlock). + Starving, + /// No progress with work stuck, but below the warn threshold — no-op. + Watching, + /// No progress with work stuck past the warn (but below the fatal) + /// threshold — log a diagnostic snapshot and keep watching. + Stalled, + /// No progress with work stuck past the fatal threshold — a genuine wedge. + /// Fail the pipeline fast instead of hanging forever. + Wedged, +} + +/// Classify one deadlock-monitor poll. +/// +/// - `progressed`: did the global progress counter advance since the last poll? +/// - `stall_secs`: how long progress has been flat (0 if it just advanced). +/// - `warn_secs` / `fatal_secs`: the warn and fatal stall thresholds. +/// - `in_flight_bytes`: bytes currently held across transport queues **and** +/// reorder buffers. +/// +/// The starvation guard (no progress + nothing in flight ⇒ not a deadlock) and +/// fatal-on-stuck-work behavior mirror the legacy pipeline's +/// `check_deadlock_and_restore`. The split warn/fatal thresholds diverge from +/// legacy's single 10s kill: a single huge dispatch (busy-locus group, large +/// sort merge) can legitimately flatten progress for many seconds with work +/// stuck in queues, so we only fail after the stall persists well past the warn +/// window. +fn classify_stall( + progressed: bool, + stall_secs: u64, + warn_secs: u64, + fatal_secs: u64, + in_flight_bytes: u64, +) -> StallVerdict { + if progressed { + return StallVerdict::Progressing; + } + if in_flight_bytes == 0 { + return StallVerdict::Starving; + } + if stall_secs >= fatal_secs { + return StallVerdict::Wedged; + } + if stall_secs >= warn_secs { + return StallVerdict::Stalled; + } + StallVerdict::Watching +} + +/// Per-queue floor: never let the rebalancer take a queue below this +/// (`ByteBoundedQueue` panics if `limit_bytes == 0`, and very small +/// limits effectively wedge the producer). +const MIN_PER_QUEUE_BYTES: u64 = 1024 * 1024; + +/// Floor for the per-branch reorder overflow stash. Liveness needs no floor +/// (`next_serial` is always exempt — any cap ≥ 0 is deadlock-free), so this +/// is purely a throughput knob: keep enough lookahead headroom that a +/// reorder-heavy edge doesn't thrash one item per round-robin pass. +const MIN_REORDER_OVERFLOW_BYTES: u64 = 4 * 1024 * 1024; + +/// Initial-allocation pass for `queue_memory_total`. Distributes +/// `total` evenly across all byte-bounded queues. Floors each queue +/// at `MIN_PER_QUEUE_BYTES` even if the per-queue share would be +/// smaller — in that case the effective total exceeds the user's +/// budget, but starvation is the worse failure mode. +/// +/// Each ordered byte-bounded branch's reorder overflow stash is sized from +/// the SAME `per_queue` value (clamped to `[MIN_REORDER_OVERFLOW_BYTES, +/// DEFAULT_REORDER_OVERFLOW_BYTES]`), so the off-budget stash tracks the +/// transport budget instead of a fixed 256 MiB. The clamp ceiling is the +/// prior fixed value, so high thread counts (large `per_queue`) keep today's +/// reorder headroom — no `--threads N` regression — while low thread counts +/// (small `per_queue`) get a streaming-sized stash. +pub(crate) fn apply_initial_queue_budget( + queues: &[crate::runtime::contexts::RegisteredQueue], + total: u64, +) { + if queues.is_empty() { + return; + } + let per_queue = (total / (queues.len() as u64)).max(MIN_PER_QUEUE_BYTES); + let reorder_cap = reorder_cap_for(per_queue); + for rq in queues { + rq.handle.set_limit_bytes(per_queue); + if let Some(reorder) = &rq.reorder_cap { + reorder.set_max_overflow_bytes(reorder_cap); + } + } +} + +/// The reorder overflow cap for a branch whose transport budget is +/// `per_queue`: track the transport budget, clamped to +/// `[MIN_REORDER_OVERFLOW_BYTES, DEFAULT_REORDER_OVERFLOW_BYTES]`. The +/// ceiling is the prior fixed value, so a large `per_queue` (high thread +/// counts) keeps today's reorder headroom; a small `per_queue` (low thread +/// counts / lean budget) shrinks the off-budget stash to a streaming size. +fn reorder_cap_for(per_queue: u64) -> u64 { + per_queue.clamp(MIN_REORDER_OVERFLOW_BYTES, crate::reorder::DEFAULT_REORDER_OVERFLOW_BYTES) +} + +/// Background queue-memory rebalancer body. Polls each queue's +/// `current_bytes / limit_bytes` fullness ratio every 1 second. +/// Identifies the most-full producer (likely bottleneck) and the +/// least-full consumer (over-budget). Shifts a fraction of budget +/// from least to most full, preserving total budget. +/// +/// The algorithm is deliberately simple — incremental shifts (10% +/// of the source's limit per tick) converge gradually so transient +/// spikes don't overshoot. Floors each queue at `MIN_PER_QUEUE_BYTES`. +/// +/// Exits when `stop` is set (typically after workers join). +#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation, clippy::cast_sign_loss)] +fn run_queue_rebalancer( + stop: &Arc, + handles: &[Arc], + names: &[&'static str], +) { + if handles.len() < 2 { + // Nothing to rebalance with one or zero queues. + return; + } + let poll_interval = std::time::Duration::from_secs(1); + let shift_fraction: f64 = 0.10; + + while !stop.is_stopped() { + sleep_until_stop(stop, poll_interval); + if stop.is_stopped() { + break; + } + + // Snapshot fullness ratios. + let snapshot: Vec<(usize, u64, u64, f64)> = handles + .iter() + .enumerate() + .map(|(idx, h)| { + let cur = h.current_bytes(); + let lim = h.limit_bytes(); + let ratio = if lim == 0 { 0.0 } else { (cur as f64) / (lim as f64) }; + (idx, cur, lim, ratio) + }) + .collect(); + + // Find the most-full and least-full queues. + let max = snapshot + .iter() + .max_by(|a, b| a.3.partial_cmp(&b.3).unwrap_or(std::cmp::Ordering::Equal)) + .copied(); + let min = snapshot + .iter() + .min_by(|a, b| a.3.partial_cmp(&b.3).unwrap_or(std::cmp::Ordering::Equal)) + .copied(); + + let (Some((max_idx, _, max_lim, max_ratio)), Some((min_idx, _, min_lim, min_ratio))) = + (max, min) + else { + continue; + }; + if max_idx == min_idx { + continue; + } + // Only rebalance when the imbalance is meaningful: the + // fullest queue is ≥80% full AND the emptiest is ≤20% full. + // Otherwise the system is in steady state and we shouldn't + // perturb the limits. + if max_ratio < 0.80 || min_ratio > 0.20 { + continue; + } + + // Shift from min to max. + let to_shift = ((min_lim as f64) * shift_fraction) as u64; + if to_shift == 0 { + continue; + } + let new_min = min_lim.saturating_sub(to_shift).max(MIN_PER_QUEUE_BYTES); + if new_min == min_lim { + // Floor reached; can't shrink further. + continue; + } + let actual_shift = min_lim - new_min; + let new_max = max_lim.saturating_add(actual_shift); + handles[min_idx].set_limit_bytes(new_min); + handles[max_idx].set_limit_bytes(new_max); + log::debug!( + "queue rebalance: shift {} bytes {} ({} -> {}) -> {} ({} -> {})", + actual_shift, + names[min_idx], + min_lim, + new_min, + names[max_idx], + max_lim, + new_max + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io; + + use rstest::rstest; + + use crate::outputs::Single; + use crate::queues::QueueSpec; + use crate::reorder::BranchOrdering; + use crate::step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}; + + // Each case pins one level's (is_on/samples, timeline, deep) predicates, so a + // failure identifies the specific level that regressed. `samples()` tracks + // `is_on()`, so the two share the `is_on` column. + #[rstest] + #[case(InstrumentationLevel::Off, false, false, false)] + #[case(InstrumentationLevel::Summary, true, false, false)] + #[case(InstrumentationLevel::Timeline, true, true, false)] + #[case(InstrumentationLevel::Deep, true, true, true)] + fn instrumentation_level_predicates( + #[case] level: InstrumentationLevel, + #[case] is_on: bool, + #[case] timeline: bool, + #[case] deep: bool, + ) { + assert_eq!(level.is_on(), is_on); + assert_eq!(level.samples(), is_on); + assert_eq!(level.timeline(), timeline); + assert_eq!(level.deep(), deep); + } + + #[test] + fn instrumentation_defaults_are_off() { + assert_eq!(InstrumentationLevel::default(), InstrumentationLevel::Off); + assert_eq!(PipelineConfig::default().instrumentation, InstrumentationLevel::Off); + assert!(PipelineConfig::default().trace_path.is_none()); + } + + #[test] + fn apply_initial_queue_budget_sets_registered_reorder_cap() { + // End-to-end wiring: a registered ordered byte-bounded branch's reorder + // cap is re-sized by the budget pass (not left at its construction + // default). Guards the Pass-1.5 registration + the `set` call together. + use crate::queues::{BoundedQueueHandle, ByteBoundedQueue, ItemQueue}; + use crate::reorder::{ + DEFAULT_REORDER_OVERFLOW_BYTES, ReorderCapHandle, ReorderStage, Sequenced, + }; + use crate::runtime::contexts::RegisteredQueue; + use crate::topology::{BranchIdx, StepIdx}; + + // A reorder stage constructed at the 256 MiB fallback; keep a concrete + // handle so we can read the cap back after the budget pass. + let transport: Arc>> = + Arc::new(ByteBoundedQueue::>::new(1024 * 1024)); + let stage = Arc::new(ReorderStage::::with_max_overflow_bytes( + transport, + DEFAULT_REORDER_OVERFLOW_BYTES, + )); + assert_eq!(stage.current_max_overflow_bytes(), DEFAULT_REORDER_OVERFLOW_BYTES); + let reorder_dyn: Arc = stage.clone(); + + // A transport-limit handle for the `RegisteredQueue.handle` slot. + let transport_q = Arc::new(ByteBoundedQueue::::new(1024 * 1024)); + let transport_handle: Arc = transport_q; + + let registered = vec![RegisteredQueue { + producer_step_name: "TestStep", + producer_step: StepIdx(0), + branch: BranchIdx(0), + handle: transport_handle, + reorder_cap: Some(reorder_dyn), + }]; + + // Lean total → per_queue = 8 MiB (1 queue) → reorder clamp = 8 MiB. + apply_initial_queue_budget(®istered, 8 * 1024 * 1024); + assert_eq!( + stage.current_max_overflow_bytes(), + 8 * 1024 * 1024, + "budget pass must re-size the registered reorder cap to the clamped per_queue" + ); + + // Huge total → per_queue huge → reorder clamped back to the ceiling. + apply_initial_queue_budget(®istered, 100 * 1024 * 1024 * 1024); + assert_eq!( + stage.current_max_overflow_bytes(), + DEFAULT_REORDER_OVERFLOW_BYTES, + "high budget clamps the reorder cap to the 256 MiB ceiling (no t>1 regression)" + ); + } + + #[test] + fn apply_initial_queue_budget_floors_each_queue_at_min_when_budget_tiny() { + // When `total / n_queues < MIN_PER_QUEUE_BYTES`, the budget pass floors + // each transport at `MIN_PER_QUEUE_BYTES` even though the effective total + // then exceeds the user's budget — starvation (a zero/tiny-budget queue + // that always rejects pushes, wedging the producer) is the worse failure + // mode. A regression dropping the `.max(MIN_PER_QUEUE_BYTES)` would not be + // caught by the lean/huge cases the sibling test covers. + use crate::queues::{BoundedQueueHandle, ByteBoundedQueue}; + use crate::runtime::contexts::RegisteredQueue; + use crate::topology::{BranchIdx, StepIdx}; + + // Four registered queues, each constructed at the 1 MiB floor. + let handles: Vec>> = + (0..4).map(|_| Arc::new(ByteBoundedQueue::::new(MIN_PER_QUEUE_BYTES))).collect(); + let registered: Vec = handles + .iter() + .enumerate() + .map(|(i, h)| RegisteredQueue { + producer_step_name: "TestStep", + producer_step: StepIdx(i), + branch: BranchIdx(0), + handle: Arc::clone(h) as Arc, + reorder_cap: None, + }) + .collect(); + + // total = 1 byte over 4 queues → per_queue would be 0, floored to 1 MiB. + apply_initial_queue_budget(®istered, 1); + for h in &handles { + assert_eq!( + h.limit_bytes(), + MIN_PER_QUEUE_BYTES, + "each queue's transport limit must be floored to MIN_PER_QUEUE_BYTES \ + when the per-queue share underflows the floor" + ); + } + } + + #[test] + fn reorder_cap_tracks_per_queue_clamped_to_floor_and_ceiling() { + let ceiling = crate::reorder::DEFAULT_REORDER_OVERFLOW_BYTES; + // Mid-range per_queue passes through unchanged. + assert_eq!(reorder_cap_for(32 * 1024 * 1024), 32 * 1024 * 1024); + // Tiny per_queue (lean / low-thread) is floored — but stays small. + assert_eq!(reorder_cap_for(1024), MIN_REORDER_OVERFLOW_BYTES); + assert_eq!(reorder_cap_for(MIN_REORDER_OVERFLOW_BYTES - 1), MIN_REORDER_OVERFLOW_BYTES); + // Huge per_queue (high thread counts) is capped at today's ceiling → + // no `--threads N` regression. + assert_eq!(reorder_cap_for(4 * ceiling), ceiling); + assert_eq!(reorder_cap_for(ceiling), ceiling); + } + + // ───── Test stubs ───── + + #[derive(Clone)] + struct StubSource; + impl Step for StubSource { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "Source", + kind: StepKind::Exclusive, + sticky: true, + output_queues: vec![QueueSpec::CountBounded { capacity: 64 }], + branch_ordering: vec![BranchOrdering::ByOrdinal], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::Finished) + } + } + + #[derive(Clone)] + struct StubTransform; + impl Step for StubTransform { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "Transform", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 64 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + #[derive(Clone)] + struct StubFanOut2; + impl Step for StubFanOut2 { + type Input = u64; + type Outputs = (u32, String); + fn profile(&self) -> StepProfile { + StepProfile { + name: "FanOut2", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![ + QueueSpec::CountBounded { capacity: 32 }, + QueueSpec::CountBounded { capacity: 32 }, + ], + branch_ordering: vec![BranchOrdering::None, BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + #[derive(Clone)] + struct StubSinkU32; + impl Step for StubSinkU32 { + type Input = u32; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "SinkU32", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + #[derive(Clone)] + struct StubSinkString; + impl Step for StubSinkString { + type Input = String; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "SinkString", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + // ───── Tests ───── + + #[test] + fn empty_builder_returns_empty_error() { + let builder = PipelineBuilder::new(); + assert!(matches!(builder.build(), Err(BuildError::Empty))); + } + + #[test] + fn unwired_source_returns_unwired_output() { + let builder = PipelineBuilder::new(); + let _chain = builder.chain(StubSource); + let result = builder.build(); + assert!(matches!(result, Err(BuildError::UnwiredOutput { step: "Source", branch: "0" }))); + } + + /// Sources register with `input_arity = 0`, matching the contract + /// `ChainGraph::register_step_with_input_arity` documents. A source's input + /// is implicit, so the graph must have no slot to wire an edge into. + #[rstest] + #[case::chain(false)] + #[case::append_source(true)] + fn sources_register_with_zero_input_arity(#[case] via_append_source: bool) { + let builder = PipelineBuilder::new(); + let source = if via_append_source { + builder.append_source(StubSource).0 + } else { + let chain = builder.chain(StubSource); + // `Chain` does not expose its `StepIdx`; the source is step 0. + drop(chain); + StepIdx(0) + }; + assert_eq!( + builder.inner.borrow().graph.input_arity(source), + 0, + "a source's input is implicit — it must have no input slot" + ); + } + + /// Wiring a producer into a source is rejected at build time. + /// + /// Arity 0 makes `wire_to_slot` panic for a source registered through + /// `PipelineBuilder::{chain, append_source}`, but a source appended as a + /// *consumer* goes through `Chain::chain`, which registers every consumer + /// with arity 1 — so the wire succeeds and only `build()` can catch it. + /// Without the check the chain builds clean, `build_chain_contexts_inner` + /// hands the second source a dummy unit input handle because `is_source()` + /// is true, and `StubSource`'s output is silently discarded. + #[test] + fn build_rejects_a_producer_wired_into_a_source() { + /// Emits `()`, so its `Chain>` type-checks against any + /// `Step` — i.e. against another source. + #[derive(Clone)] + struct UnitSource; + impl Step for UnitSource { + type Input = (); + type Outputs = Single<()>; + fn profile(&self) -> StepProfile { + StepProfile { + name: "UnitSource", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + let builder = PipelineBuilder::new(); + // `StubSource: Step` — accepted here purely because + // `HeapSize for ()` exists, not because the wiring is meaningful. + builder.chain(UnitSource).chain(StubSource).chain(StubSinkU32).into_sink_marker(); + + assert!( + matches!( + builder.build(), + Err(BuildError::WiredIntoSource { step: "Source", producer: "UnitSource" }) + ), + "a producer wired into a source must be a build error, not silent data loss" + ); + } + + #[test] + fn source_to_transform_unwired_at_transform_returns_unwired() { + let builder = PipelineBuilder::new(); + let _chain = builder.chain(StubSource).chain(StubTransform); + let result = builder.build(); + assert!(matches!( + result, + Err(BuildError::UnwiredOutput { step: "Transform", branch: "0" }) + )); + } + + #[test] + fn fully_wired_source_to_sink_succeeds() { + let builder = PipelineBuilder::new(); + builder.chain(StubSource).chain(StubSinkU32).into_sink_marker(); + let result = builder.build(); + assert!(result.is_ok()); + let pipeline = result.unwrap(); + assert_eq!(pipeline.graph.n_steps(), 2); + } + + #[test] + fn unwired_fanout_branch_is_detected() { + let builder = PipelineBuilder::new(); + let after_fanout = builder.chain(StubSource).chain(StubTransform).chain(StubFanOut2); + let multi = after_fanout.into_multi(); + // Wire branch 0 to a sink, drop branch 1 (unwired). + multi.b0.chain(StubSinkU32).into_sink_marker(); + drop(multi.b1); + + let result = builder.build(); + assert!(matches!(result, Err(BuildError::UnwiredOutput { step: "FanOut2", branch: "1" }))); + } + + #[test] + fn fanout_with_both_branches_wired_succeeds() { + let builder = PipelineBuilder::new(); + let after_fanout = builder.chain(StubSource).chain(StubTransform).chain(StubFanOut2); + let multi = after_fanout.into_multi(); + multi.b0.chain(StubSinkU32).into_sink_marker(); + multi.b1.chain(StubSinkString).into_sink_marker(); + + let result = builder.build(); + assert!(result.is_ok()); + let pipeline = result.unwrap(); + assert_eq!(pipeline.graph.n_steps(), 5); + } + + #[test] + fn pipeline_config_default_uses_available_parallelism() { + let cfg = PipelineConfig::default(); + assert!(cfg.threads >= 1); + } + + #[derive(Clone)] + struct StubFanOut3; + impl Step for StubFanOut3 { + type Input = u64; + type Outputs = (u32, String, u64); + fn profile(&self) -> StepProfile { + StepProfile { + name: "FanOut3", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 8 }; 3], + branch_ordering: vec![BranchOrdering::None; 3], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + #[derive(Clone)] + struct StubFanOut4; + impl Step for StubFanOut4 { + type Input = u64; + type Outputs = (u32, String, u64, u32); + fn profile(&self) -> StepProfile { + StepProfile { + name: "FanOut4", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 8 }; 4], + branch_ordering: vec![BranchOrdering::None; 4], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + #[derive(Clone)] + struct StubSinkU64; + impl Step for StubSinkU64 { + type Input = u64; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "SinkU64", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + /// `into_multi` must hand back one `Chain` per declared branch, each + /// pinned to its own `BranchIdx`. If two sub-chains shared a branch index + /// the builder would report a phantom unwired branch even though the + /// caller wired every one — so wiring all of them must build cleanly. + #[test] + fn into_multi_3_exposes_one_chain_per_branch() { + let builder = PipelineBuilder::new(); + let multi = builder.chain(StubSource).chain(StubTransform).chain(StubFanOut3).into_multi(); + multi.b0.chain(StubSinkU32).into_sink_marker(); + multi.b1.chain(StubSinkString).into_sink_marker(); + multi.b2.chain(StubSinkU64).into_sink_marker(); + + let pipeline = builder.build().expect("all three branches wired"); + assert_eq!(pipeline.graph.n_steps(), 6, "source + transform + fanout + 3 sinks"); + } + + /// The 3-branch counterpart of `unwired_fanout_branch_is_detected`: the + /// dropped branch must be named by index, proving `into_multi` gave branch + /// 2 its own identity rather than aliasing an earlier one. + #[test] + fn into_multi_3_reports_the_branch_left_unwired() { + let builder = PipelineBuilder::new(); + let multi = builder.chain(StubSource).chain(StubTransform).chain(StubFanOut3).into_multi(); + multi.b0.chain(StubSinkU32).into_sink_marker(); + multi.b1.chain(StubSinkString).into_sink_marker(); + drop(multi.b2); + + assert!(matches!( + builder.build(), + Err(BuildError::UnwiredOutput { step: "FanOut3", branch: "2" }) + )); + } + + #[test] + fn into_multi_4_exposes_one_chain_per_branch() { + let builder = PipelineBuilder::new(); + let multi = builder.chain(StubSource).chain(StubTransform).chain(StubFanOut4).into_multi(); + multi.b0.chain(StubSinkU32).into_sink_marker(); + multi.b1.chain(StubSinkString).into_sink_marker(); + multi.b2.chain(StubSinkU64).into_sink_marker(); + multi.b3.chain(StubSinkU32).into_sink_marker(); + + let pipeline = builder.build().expect("all four branches wired"); + assert_eq!(pipeline.graph.n_steps(), 7, "source + transform + fanout + 4 sinks"); + } + + #[test] + fn into_multi_4_reports_the_branch_left_unwired() { + let builder = PipelineBuilder::new(); + let multi = builder.chain(StubSource).chain(StubTransform).chain(StubFanOut4).into_multi(); + multi.b0.chain(StubSinkU32).into_sink_marker(); + multi.b1.chain(StubSinkString).into_sink_marker(); + multi.b2.chain(StubSinkU64).into_sink_marker(); + drop(multi.b3); + + assert!(matches!( + builder.build(), + Err(BuildError::UnwiredOutput { step: "FanOut4", branch: "3" }) + )); + } + + #[derive(Clone, Copy)] + struct Ord32 { + ordinal: u64, + } + impl crate::item::HeapSize for Ord32 {} + impl crate::item::Ordered for Ord32 { + fn ordinal(&self) -> u64 { + self.ordinal + } + } + + #[derive(Clone)] + struct OrderedSource; + impl Step for OrderedSource { + type Input = (); + type Outputs = crate::outputs::OrderedBytesTuple2; + fn profile(&self) -> StepProfile { + StepProfile { + name: "OrderedSource", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: 1024 }; 2], + branch_ordering: vec![BranchOrdering::ByItemOrdinal; 2], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::Finished) + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + #[derive(Clone)] + struct OrderedSink; + impl Step for OrderedSink { + type Input = Ord32; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "OrderedSink", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + /// The ordered + byte-bounded 2-way fan-out — the shape real BAM source + /// chains produce — is the one `into_multi` variant with no coverage. Its + /// sub-chains are `Chain>`, not `Chain>`, + /// so a branch mix-up would surface as a type error at the call site; what + /// this pins is that both branches are exposed and separately wireable. + #[test] + fn into_multi_ordered_bytes_2_exposes_one_chain_per_branch() { + let builder = PipelineBuilder::new(); + let multi = builder.chain(OrderedSource).into_multi(); + let _: &Chain<'_, crate::outputs::OrderedBytesSingle> = &multi.b0; + multi.b0.chain(OrderedSink).into_sink_marker(); + multi.b1.chain(OrderedSink).into_sink_marker(); + + let pipeline = builder.build().expect("both ordered branches wired"); + assert_eq!(pipeline.graph.n_steps(), 3, "source + 2 sinks"); + } + + /// The unwired branch is reported BY INDEX, so the two sub-chains are + /// distinct edges rather than two views of branch 0. + #[test] + fn into_multi_ordered_bytes_2_reports_the_branch_left_unwired() { + let builder = PipelineBuilder::new(); + let multi = builder.chain(OrderedSource).into_multi(); + multi.b0.chain(OrderedSink).into_sink_marker(); + drop(multi.b1); + + assert!(matches!( + builder.build(), + Err(BuildError::UnwiredOutput { step: "OrderedSource", branch: "1" }) + )); + } + + #[derive(Clone)] + struct OrderedSource3; + impl Step for OrderedSource3 { + type Input = (); + type Outputs = crate::outputs::OrderedBytesTuple3; + fn profile(&self) -> StepProfile { + StepProfile { + name: "OrderedSource3", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: 1024 }; 3], + branch_ordering: vec![BranchOrdering::ByItemOrdinal; 3], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::Finished) + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + /// The 3-way counterpart of `into_multi_ordered_bytes_2_exposes_one_chain_per_branch`: + /// all three ordered + byte-bounded branches are exposed as + /// `Chain>` and separately wireable. + #[test] + fn into_multi_ordered_bytes_3_exposes_one_chain_per_branch() { + let builder = PipelineBuilder::new(); + let multi = builder.chain(OrderedSource3).into_multi(); + let _: &Chain<'_, crate::outputs::OrderedBytesSingle> = &multi.b0; + multi.b0.chain(OrderedSink).into_sink_marker(); + multi.b1.chain(OrderedSink).into_sink_marker(); + multi.b2.chain(OrderedSink).into_sink_marker(); + + let pipeline = builder.build().expect("all three ordered branches wired"); + assert_eq!(pipeline.graph.n_steps(), 4, "source + 3 sinks"); + } + + /// The 3-way counterpart of `into_multi_ordered_bytes_2_reports_the_branch_left_unwired`: + /// the dropped branch is reported BY INDEX, proving `b2` is a distinct edge + /// rather than an alias of `b0`/`b1`. + #[test] + fn into_multi_ordered_bytes_3_reports_the_branch_left_unwired() { + let builder = PipelineBuilder::new(); + let multi = builder.chain(OrderedSource3).into_multi(); + multi.b0.chain(OrderedSink).into_sink_marker(); + multi.b1.chain(OrderedSink).into_sink_marker(); + drop(multi.b2); + + assert!(matches!( + builder.build(), + Err(BuildError::UnwiredOutput { step: "OrderedSource3", branch: "2" }) + )); + } + + // Each case pins one `BuildError` variant's rendering, so a failure names + // the variant whose message drifted rather than a combined assert. + #[rstest] + #[case::empty(BuildError::Empty, "pipeline has no steps")] + #[case::unwired( + BuildError::UnwiredOutput { step: "FanOut2", branch: "1" }, + "step \"FanOut2\" has unwired output branch \"1\"" + )] + fn build_error_displays_actionably(#[case] err: BuildError, #[case] expected: &str) { + assert_eq!(err.to_string(), expected); + } + + /// A read of one `PipelineConfig` field, so each setter case can name the + /// field it targets and assert on a comparable value. + #[derive(Debug, PartialEq, Eq)] + enum ConfigProbe { + DeadlockTimeout(u64), + StatsPresent(bool), + QueueMemoryTotal(Option), + Instrumentation(InstrumentationLevel), + SchedulerName(&'static str), + } + + impl ConfigProbe { + /// Read the same field this probe names out of `cfg`. + fn read_from(&self, cfg: &PipelineConfig) -> Self { + match self { + Self::DeadlockTimeout(_) => Self::DeadlockTimeout(cfg.deadlock_timeout_secs), + Self::StatsPresent(_) => Self::StatsPresent(cfg.stats.is_some()), + Self::QueueMemoryTotal(_) => Self::QueueMemoryTotal(cfg.queue_memory_total), + Self::Instrumentation(_) => Self::Instrumentation(cfg.instrumentation), + Self::SchedulerName(_) => Self::SchedulerName(cfg.scheduler.name()), + } + } + } + + /// The builder-style setters are the only way a caller configures a + /// `PipelineConfig`. Each case applies one setter and asserts two things: + /// the field it targets took the new value, and a field it does NOT target + /// is still at its default — a copy-paste slip that assigned the wrong field + /// would silently disable instrumentation or the deadlock monitor. + /// + /// `scheduler` is a trait object, so its case asserts the observable + /// `name()` rather than allocation identity: pointer inequality would pass + /// even if the setter stored the wrong scheduler. + #[rstest] + #[case::with_deadlock_timeout( + &|c: PipelineConfig| c.with_deadlock_timeout(30), + ConfigProbe::DeadlockTimeout(30), + ConfigProbe::StatsPresent(false) + )] + #[case::with_stats( + &|c: PipelineConfig| c.with_stats(Arc::new(PipelineStats::new(vec!["Source"]))), + ConfigProbe::StatsPresent(true), + ConfigProbe::DeadlockTimeout(0) + )] + #[case::with_queue_memory_total_some( + &|c: PipelineConfig| c.with_queue_memory_total(Some(64 * 1024 * 1024)), + ConfigProbe::QueueMemoryTotal(Some(64 * 1024 * 1024)), + ConfigProbe::StatsPresent(false) + )] + // `None` is the documented "keep per-step static limits, rebalancer off" + // value, so it must round-trip as None rather than be treated as unset. + #[case::with_queue_memory_total_none( + &|c: PipelineConfig| c.with_queue_memory_total(None), + ConfigProbe::QueueMemoryTotal(None), + ConfigProbe::StatsPresent(false) + )] + #[case::with_instrumentation( + &|c: PipelineConfig| c.with_instrumentation(InstrumentationLevel::Deep), + ConfigProbe::Instrumentation(InstrumentationLevel::Deep), + ConfigProbe::DeadlockTimeout(0) + )] + #[case::with_scheduler( + &|c: PipelineConfig| c.with_scheduler(Arc::new(crate::runtime::DrainFirstScheduler)), + ConfigProbe::SchedulerName("drain-first"), + ConfigProbe::Instrumentation(InstrumentationLevel::Off) + )] + fn pipeline_config_setter_sets_only_its_own_field( + #[case] apply: &dyn Fn(PipelineConfig) -> PipelineConfig, + #[case] expected: ConfigProbe, + #[case] untouched: ConfigProbe, + ) { + let defaults = PipelineConfig::default(); + assert_eq!( + untouched.read_from(&defaults), + untouched, + "the case's `untouched` probe must state the actual default" + ); + + let cfg = apply(PipelineConfig::default()); + assert_eq!(expected.read_from(&cfg), expected, "setter set its own field"); + assert_eq!(untouched.read_from(&cfg), untouched, "setter left the other field alone"); + } + + /// The default scheduler must stay upstream-first; `with_scheduler`'s case + /// above is only meaningful against a known baseline. + #[test] + fn default_scheduler_walks_chain_order() { + let cfg = PipelineConfig::default(); + assert_eq!(cfg.scheduler.name(), "chain-order"); + assert_eq!(cfg.scheduler.walk(), crate::runtime::WalkDirection::Forward); + } + + /// `append_source` / `append_step` are the type-erased assembly API the + /// parent crate's `ChainBuilder` drives across method-boundary calls, where + /// the fluent `chain` API cannot be used because the chain type changes at + /// every step. Wiring by returned `(StepIdx, BranchIdx)` must produce the + /// same graph the fluent API does. + #[test] + fn append_source_and_append_step_build_the_same_graph_as_chaining() { + let builder = PipelineBuilder::new(); + let src = builder.append_source(StubSource); + let mid = builder.append_step(StubTransform, src); + let _sink = builder.append_step(StubSinkU64, mid); + + let pipeline = builder.build().expect("every branch wired by index"); + assert_eq!(pipeline.graph.n_steps(), 3); + } + + #[test] + fn append_step_leaves_an_unwired_branch_detectable() { + // Appending a fan-out and wiring only one of its branches must still be + // caught by `build`, exactly as with the fluent API. + let builder = PipelineBuilder::new(); + let src = builder.append_source(StubSource); + let mid = builder.append_step(StubTransform, src); + let fanout = builder.append_step(StubFanOut2, mid); + let _sink = builder.append_step(StubSinkU32, fanout); + + assert!(matches!( + builder.build(), + Err(BuildError::UnwiredOutput { step: "FanOut2", branch: "1" }) + )); + } + + #[derive(Clone)] + struct StubJoin; + impl crate::step::Step2 for StubJoin { + type InputA = u32; + type InputB = String; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "Join", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run( + &mut self, + _ctx: &mut crate::step::StepCtx2<'_, Self>, + ) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + /// `append_step2` must consume BOTH upstream branches — one into input slot + /// 0 and one into slot 1. If it wired only one, `build` would report the + /// other as unwired. + #[test] + fn append_step2_consumes_both_upstream_branches() { + let builder = PipelineBuilder::new(); + let src = builder.append_source(StubSource); + let mid = builder.append_step(StubTransform, src); + let fanout = builder.append_step(StubFanOut2, mid); + let branch_b = (fanout.0, BranchIdx(1)); + let _join = builder.append_step2(StubJoin, fanout, branch_b); + + let pipeline = builder.build().expect("both fan-out branches consumed by the join"); + assert_eq!(pipeline.graph.n_steps(), 4, "source + transform + fanout + join"); + } + + #[test] + fn dag_renders_chain_shape() { + let builder = PipelineBuilder::new(); + builder.chain(StubSource).chain(StubSinkU32).into_sink_marker(); + let pipeline = builder.build().unwrap(); + let dag = pipeline.dag(); + // Verify shape rendering — names and (sink) marker. + assert!(dag.contains("Source"), "DAG missing source name: {dag}"); + assert!(dag.contains("SinkU32"), "DAG missing sink name: {dag}"); + assert!(dag.contains("(sink)"), "DAG missing sink marker: {dag}"); + assert!(dag.contains("→ SinkU32"), "DAG missing source→sink wiring: {dag}"); + } + + #[test] + fn dag_renders_effective_collapsed_ordering_for_serial_exclusive() { + // `StubSource` is `Exclusive` and declares `BranchOrdering::ByOrdinal`, + // but `build_output_set` collapses that to `None` (no reorder stage) for + // single-input Serial/Exclusive producers. `dag()` must render the + // *effective* ordering so the diagnostic matches the transport actually + // built — it must not print the un-collapsed declared `ByOrdinal`. + let builder = PipelineBuilder::new(); + builder.chain(StubSource).chain(StubSinkU32).into_sink_marker(); + let pipeline = builder.build().unwrap(); + let dag = pipeline.dag(); + // The source's output branch line is `.0: → SinkU32`. + let source_branch_line = dag + .lines() + .find(|l| l.contains("→ SinkU32")) + .expect("DAG must have the source→sink branch line"); + assert!( + source_branch_line.contains("None"), + "DAG must render the collapsed (effective) ordering `None`: {source_branch_line}" + ); + assert!( + !source_branch_line.contains("ByOrdinal"), + "DAG must NOT render the un-collapsed declared `ByOrdinal`: {source_branch_line}" + ); + } + + // ───────────────────────────────────────────────────────────────────── + // Pipeline::run smoke tests + // ───────────────────────────────────────────────────────────────────── + + use std::sync::atomic::{AtomicU32, Ordering as AtomicOrd}; + + use crate::signal::PipelineError; + + /// Source emitting `remaining` items via a shared atomic counter; safe + /// for both single-worker and multi-worker `Parallel` execution. + /// + /// Uses an `Unbounded` output queue so the test never hits backpressure + /// (which would require the source to use the `HeldSlot>` + /// retry pattern — exercised in the bigger end-to-end smoke test below). + #[derive(Clone)] + struct SharedCountingSource { + remaining: Arc, + } + impl Step for SharedCountingSource { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "SharedSource", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::Unbounded], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> std::io::Result { + let n = self.remaining.load(AtomicOrd::Acquire); + if n == 0 { + return Ok(StepOutcome::Finished); + } + // CAS down to claim this item; only push on successful claim. + if self + .remaining + .compare_exchange(n, n - 1, AtomicOrd::AcqRel, AtomicOrd::Acquire) + .is_ok() + { + ctx.outputs.push(n).map_err(|_| { + std::io::Error::other("Unbounded queue rejected push (impossible)") + })?; + Ok(StepOutcome::Progress) + } else { + Ok(StepOutcome::NoProgress) + } + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + /// Byte-bounded variant of [`SharedCountingSource`]. Identical claim/push + /// logic but declares a `ByteBounded` output transport so it is + /// monitor-visible — required by any test that arms the deadlock monitor + /// (`deadlock_timeout_secs > 0` + stats), which requires every output edge + /// to be `ByteBounded` (see `ensure_monitor_visible_transports`) before + /// workers spawn. Using the `Unbounded` source there would fail the run with + /// `PipelineError::MonitorBlindTransport` before the worker-panic path is + /// ever reached. + #[derive(Clone)] + struct SharedCountingSourceByteBounded { + remaining: Arc, + /// A value claimed from `remaining` but not yet accepted by the output + /// (the byte-bounded push hit backpressure). Held per-worker and retried + /// on a later tick so the exact ordinal survives — rolling `remaining` + /// back instead would let a peer re-claim the count and drop/duplicate an + /// ordinal under contention. + pending: Option, + } + impl Step for SharedCountingSourceByteBounded { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "SharedSourceByteBounded", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: 1 << 20 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> std::io::Result { + // First flush any value claimed on a prior tick whose push was + // rejected by backpressure. Retrying the exact held ordinal (rather + // than rolling `remaining` back) keeps the emitted set a clean + // permutation of `1..=N` even under contention. + if let Some(n) = self.pending { + return if ctx.outputs.push(n).is_ok() { + self.pending = None; + Ok(StepOutcome::Progress) + } else { + Ok(StepOutcome::NoProgress) + }; + } + let n = self.remaining.load(AtomicOrd::Acquire); + if n == 0 { + return Ok(StepOutcome::Finished); + } + if self + .remaining + .compare_exchange(n, n - 1, AtomicOrd::AcqRel, AtomicOrd::Acquire) + .is_ok() + { + // Byte-bounded push can hit backpressure; hold the claimed + // ordinal and retry it on a later tick. Holding a claimed item + // counts as progress per the `StepOutcome` contract ("pushed or + // held an item" — see `step.rs`) and matches the production + // held-slot source (`sort::merge`), so the scheduler's deadlock + // accounting sees the claim as forward motion rather than a stall. + if ctx.outputs.push(n).is_ok() { + Ok(StepOutcome::Progress) + } else { + self.pending = Some(n); + Ok(StepOutcome::Progress) + } + } else { + Ok(StepOutcome::NoProgress) + } + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + /// Sink that pops and counts. + #[derive(Clone)] + struct ParallelCountingSink { + received: Arc, + /// Every value popped, so a test can assert record *identity* and not + /// just the count — a framework bug that duplicates one item and drops + /// another keeps the count at N. Shared across worker copies. + seen: Arc>>, + } + impl ParallelCountingSink { + fn new(received: &Arc) -> Self { + Self { + received: Arc::clone(received), + seen: Arc::new(parking_lot::Mutex::new(Vec::new())), + } + } + } + impl Step for ParallelCountingSink { + type Input = u32; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "ParallelSink", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> std::io::Result { + match ctx.input.pop() { + Some(v) => { + self.seen.lock().push(v); + self.received.fetch_add(1, AtomicOrd::Relaxed); + Ok(StepOutcome::Progress) + } + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + /// `SharedCountingSource` emits the distinct ordinals `n..=1`, so a drain + /// test can assert the exact multiset the sink received rather than only its + /// size. Sorted because worker interleaving makes arrival order a valid + /// scheduling detail; identity is the contract, order is not. + fn assert_received_every_ordinal_once(sink: &ParallelCountingSink, n: u32) { + let mut got = sink.seen.lock().clone(); + got.sort_unstable(); + let expected: Vec = (1..=n).collect(); + assert_eq!(got, expected, "every emitted ordinal 1..={n} must arrive exactly once"); + } + + #[test] + fn pipeline_run_with_threads_1_drains_chain() { + let remaining = Arc::new(AtomicU32::new(10)); + let received = Arc::new(AtomicU32::new(0)); + let sink = ParallelCountingSink::new(&received); + + let builder = PipelineBuilder::new(); + builder + .chain(SharedCountingSource { remaining: Arc::clone(&remaining) }) + .chain(sink.clone()) + .into_sink_marker(); + let pipeline = builder.build().unwrap(); + + let result = pipeline.run(PipelineConfig { threads: 1, ..Default::default() }); + assert!(result.is_ok(), "run failed: {:?}", result.err()); + assert_eq!(received.load(AtomicOrd::Relaxed), 10); + assert_received_every_ordinal_once(&sink, 10); + } + + #[test] + fn pipeline_run_with_threads_4_drains_chain() { + let remaining = Arc::new(AtomicU32::new(50)); + let received = Arc::new(AtomicU32::new(0)); + let sink = ParallelCountingSink::new(&received); + + let builder = PipelineBuilder::new(); + builder + .chain(SharedCountingSource { remaining: Arc::clone(&remaining) }) + .chain(sink.clone()) + .into_sink_marker(); + let pipeline = builder.build().unwrap(); + + let result = pipeline.run(PipelineConfig { threads: 4, ..Default::default() }); + assert!(result.is_ok(), "run failed: {:?}", result.err()); + assert_eq!(received.load(AtomicOrd::Relaxed), 50); + assert_received_every_ordinal_once(&sink, 50); + } + + /// A `u32 -> u32` pass-through step that runs on a dedicated Detached + /// thread. Pops one item per `try_run`, pushes it on (holding on + /// output-full backpressure), finishes once its input drains. + #[derive(Clone)] + struct DetachedPassThrough { + held: Option, + } + impl Step for DetachedPassThrough { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "DetachedPassThrough", + kind: StepKind::Detached, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 8 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> std::io::Result { + // A rejected push reports `NoProgress`, not `Contention`: the driver + // treats them identically, but `Contention` means "a Serial step's + // mutex was held by another worker" and feeds `contention_count`, + // which the bottleneck verdict turns into its SPIN ratio. Using it + // for ordinary output backpressure invents contention that never + // happened. + if let Some(v) = self.held.take() { + if ctx.outputs.push(v).is_err() { + self.held = Some(v); + return Ok(StepOutcome::NoProgress); + } + return Ok(StepOutcome::Progress); + } + match ctx.input.pop() { + Some(v) => match ctx.outputs.push(v) { + Ok(()) => Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held = Some(unpushed.into_item()); + Ok(StepOutcome::NoProgress) + } + }, + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } + } + + /// L2.3: a `source(Parallel) -> passthrough(Detached) -> sink(Parallel)` + /// chain runs to completion at `--threads 4` with the Detached step on its + /// own dedicated thread (off the pool). All N items flow through. + #[test] + fn detached_step_runs_and_drains_multithreaded() { + let remaining = Arc::new(AtomicU32::new(200)); + let received = Arc::new(AtomicU32::new(0)); + let sink = ParallelCountingSink::new(&received); + + let builder = PipelineBuilder::new(); + builder + .chain(SharedCountingSource { remaining: Arc::clone(&remaining) }) + .chain(DetachedPassThrough { held: None }) + .chain(sink.clone()) + .into_sink_marker(); + let pipeline = builder.build().unwrap(); + + let result = pipeline.run(PipelineConfig { threads: 4, ..Default::default() }); + assert!(result.is_ok(), "run failed: {:?}", result.err()); + assert_eq!( + received.load(AtomicOrd::Relaxed), + 200, + "all items flowed through Detached step" + ); + assert_received_every_ordinal_once(&sink, 200); + } + + /// L2.3: a chain with a Detached step and ZERO items (source finishes + /// immediately) cleanly drains the Detached thread and the run completes. + #[test] + fn detached_step_zero_items_run_completes() { + let remaining = Arc::new(AtomicU32::new(0)); + let received = Arc::new(AtomicU32::new(0)); + let sink = ParallelCountingSink::new(&received); + + let builder = PipelineBuilder::new(); + builder + .chain(SharedCountingSource { remaining: Arc::clone(&remaining) }) + .chain(DetachedPassThrough { held: None }) + .chain(sink.clone()) + .into_sink_marker(); + let pipeline = builder.build().unwrap(); + + let result = pipeline.run(PipelineConfig { threads: 4, ..Default::default() }); + assert!(result.is_ok(), "run failed: {:?}", result.err()); + assert_eq!(received.load(AtomicOrd::Relaxed), 0); + // Empty expectation: nothing was emitted, so nothing may arrive — this + // also catches a spurious item the bare count check would miss if the + // counter and the sink ever disagreed. + assert_received_every_ordinal_once(&sink, 0); + } + + /// L2.3 step 7: at `--threads 1` the fusible linear chain runs the Detached + /// step **inline** in the fused single-thread driver — no dedicated thread + /// is spawned (the fused path returns before `extract_detached_steps`), yet + /// the Detached step still drives to completion. + #[test] + fn detached_collapses_inline_at_t1() { + let remaining = Arc::new(AtomicU32::new(30)); + let received = Arc::new(AtomicU32::new(0)); + let sink = ParallelCountingSink::new(&received); + + let builder = PipelineBuilder::new(); + builder + .chain(SharedCountingSource { remaining: Arc::clone(&remaining) }) + .chain(DetachedPassThrough { held: None }) + .chain(sink.clone()) + .into_sink_marker(); + let pipeline = builder.build().unwrap(); + + let result = pipeline.run(PipelineConfig { threads: 1, ..Default::default() }); + assert!(result.is_ok(), "run failed: {:?}", result.err()); + assert_eq!(received.load(AtomicOrd::Relaxed), 30); + assert_received_every_ordinal_once(&sink, 30); + } + + /// Sink whose worker loop panics the moment it pops an item — used to drive + /// the worker-panic deferral paths in `Pipeline::run`. + #[derive(Clone)] + struct PanickingSink; + impl Step for PanickingSink { + type Input = u32; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "PanickingSink", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> std::io::Result { + match ctx.input.pop() { + Some(_) => panic!("intentional worker panic for test"), + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + /// Run `f` with the panic hook silenced so an *expected* worker panic does + /// not spew a backtrace into the test log. Safe under `nextest`, which runs + /// each test in its own process. + fn with_silenced_panic_hook(f: impl FnOnce() -> R) -> R { + let prev = std::panic::take_hook(); + std::panic::set_hook(Box::new(|_| {})); + let result = f(); + std::panic::set_hook(prev); + result + } + + /// Assert a caught panic carries `PanickingSink`'s payload. + /// + /// `with_silenced_panic_hook` suppresses the payload, so a bare + /// `is_err()` is satisfied by *any* panic — a `thread::Builder::spawn` + /// `expect`, a framework `debug_assert!`, or the monitor-visibility guard + /// tripping would all pass while the worker-panic deferral path never ran. + fn assert_is_the_intentional_worker_panic(payload: &(dyn std::any::Any + Send), ctx: &str) { + let msg = payload + .downcast_ref::<&str>() + .map(|s| (*s).to_string()) + .or_else(|| payload.downcast_ref::().cloned()) + .expect("panic payload is a string"); + assert!( + msg.contains("intentional worker panic for test"), + "{ctx}: run() must re-raise the SINK's panic, got: {msg}" + ); + } + + #[test] + fn pipeline_run_reraises_worker_panic_single_threaded() { + // A worker-loop panic on the single-threaded fast path must propagate + // out of `run` (after the common monitor/rebalancer shutdown), not be + // swallowed. The test completing at all proves the run did not hang. + let remaining = Arc::new(AtomicU32::new(10)); + let builder = PipelineBuilder::new(); + builder + .chain(SharedCountingSource { remaining: Arc::clone(&remaining) }) + .chain(PanickingSink) + .into_sink_marker(); + let pipeline = builder.build().unwrap(); + + let result = with_silenced_panic_hook(|| { + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + pipeline.run(PipelineConfig { threads: 1, ..Default::default() }) + })) + }); + let payload = result.expect_err("single-threaded worker panic must propagate out of run()"); + assert_is_the_intentional_worker_panic(payload.as_ref(), "single-threaded"); + } + + #[test] + fn pipeline_run_reraises_worker_panic_with_monitor_enabled() { + // With the deadlock monitor enabled (stats + non-zero timeout), a + // multi-worker panic must still re-raise — after the monitor is stopped + // and joined — rather than deadlocking the join loop or leaking the + // helper thread. The panicking worker signals cancellation so any wedged + // peer observes `is_done()` and exits, letting every join complete. + // + // The source MUST be byte-bounded: arming the monitor (stats + + // non-zero timeout) runs `ensure_monitor_visible_transports`, which + // requires every output edge to be `ByteBounded` in every build. An + // `Unbounded` source would fail the run with `MonitorBlindTransport` + // before any worker is spawned, so the test would "pass" on the wrong + // error and never exercise the worker-panic deferral path it covers. + let remaining = Arc::new(AtomicU32::new(1_000)); + let builder = PipelineBuilder::new(); + builder + .chain(SharedCountingSourceByteBounded { + remaining: Arc::clone(&remaining), + pending: None, + }) + .chain(PanickingSink) + .into_sink_marker(); + let pipeline = builder.build().unwrap(); + let stats = pipeline.stats(); + + let result = with_silenced_panic_hook(|| { + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + pipeline.run(PipelineConfig { + threads: 4, + stats: Some(Arc::clone(&stats)), + deadlock_timeout_secs: 5, + ..Default::default() + }) + })) + }); + let payload = result.expect_err("multi-worker panic must propagate out of run()"); + assert_is_the_intentional_worker_panic(payload.as_ref(), "multi-worker with monitor"); + } + + #[test] + fn pipeline_stats_handle_matches_chain_size() { + let builder = PipelineBuilder::new(); + builder.chain(StubSource).chain(StubSinkU32).into_sink_marker(); + let pipeline = builder.build().unwrap(); + let stats = pipeline.stats(); + assert_eq!(stats.n_steps(), 2); + assert_eq!(stats.step_name(StepIdx(0)), "Source"); + assert_eq!(stats.step_name(StepIdx(1)), "SinkU32"); + } + + #[test] + fn pipeline_run_populates_stats_when_enabled() { + let remaining = Arc::new(AtomicU32::new(20)); + let received = Arc::new(AtomicU32::new(0)); + let sink = ParallelCountingSink::new(&received); + + let builder = PipelineBuilder::new(); + builder + .chain(SharedCountingSource { remaining: Arc::clone(&remaining) }) + .chain(sink.clone()) + .into_sink_marker(); + let pipeline = builder.build().unwrap(); + let stats = pipeline.stats(); + + let cfg = + PipelineConfig { threads: 2, stats: Some(Arc::clone(&stats)), ..Default::default() }; + let result = pipeline.run(cfg); + assert!(result.is_ok(), "run failed: {:?}", result.err()); + assert_eq!(received.load(AtomicOrd::Relaxed), 20); + + let snap = stats.snapshot(); + assert_eq!(snap.steps.len(), 2); + assert_eq!(snap.steps[0].0, "SharedSource"); + assert_eq!(snap.steps[1].0, "ParallelSink"); + + // Sink saw exactly 20 items (one Progress per pop with item). + assert_eq!(snap.steps[1].1.progress_count, 20); + // Source must have made progress at least 20 times to push the items. + assert!(snap.steps[0].1.progress_count >= 20); + // Both steps accumulated wall time. + assert!(snap.steps[0].1.total_run_ns > 0); + assert!(snap.steps[1].1.total_run_ns > 0); + // No errors recorded. + assert_eq!(snap.steps[0].1.error_count, 0); + assert_eq!(snap.steps[1].1.error_count, 0); + // The source returned Finished at least once across the workers. + assert!(snap.steps[0].1.finished_count >= 1); + } + + #[test] + fn pipeline_run_without_stats_succeeds_unchanged() { + let remaining = Arc::new(AtomicU32::new(5)); + let received = Arc::new(AtomicU32::new(0)); + let sink = ParallelCountingSink::new(&received); + + let builder = PipelineBuilder::new(); + builder + .chain(SharedCountingSource { remaining: Arc::clone(&remaining) }) + .chain(sink.clone()) + .into_sink_marker(); + let pipeline = builder.build().unwrap(); + + let result = pipeline.run(PipelineConfig { threads: 2, stats: None, ..Default::default() }); + assert!(result.is_ok(), "run failed: {:?}", result.err()); + assert_eq!(received.load(AtomicOrd::Relaxed), 5); + } + + #[test] + fn pipeline_run_responds_to_cancellation() { + /// Source that emits forever (never returns `Finished`) until cancelled. + #[derive(Clone)] + struct InfiniteSource; + impl Step for InfiniteSource { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "InfiniteSource", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 8 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> std::io::Result { + let _ = ctx.outputs.push(0); + Ok(StepOutcome::Progress) + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + #[derive(Clone)] + struct DiscardSink; + impl Step for DiscardSink { + type Input = u32; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "DiscardSink", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> std::io::Result { + match ctx.input.pop() { + Some(_) => Ok(StepOutcome::Progress), + None => Ok(StepOutcome::NoProgress), + } + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + let builder = PipelineBuilder::new(); + builder.chain(InfiniteSource).chain(DiscardSink).into_sink_marker(); + let pipeline = builder.build().unwrap(); + let cancel = pipeline.cancel_handle(); + + std::thread::spawn(move || { + std::thread::sleep(std::time::Duration::from_millis(100)); + cancel.cancel(); + }); + + let result = pipeline.run(PipelineConfig { threads: 4, ..Default::default() }); + assert!(matches!(result, Err(PipelineError::Cancelled))); + } + + #[test] + fn pipeline_run_propagates_step_error() { + /// Source that emits `n` items, then returns `Err`. + #[derive(Clone)] + struct FailingSource { + remaining: Arc, + } + impl Step for FailingSource { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "FailingSource", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 8 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> std::io::Result { + let n = self.remaining.load(AtomicOrd::Acquire); + if n == 0 { + return Err(std::io::Error::other("source failed")); + } + if self + .remaining + .compare_exchange(n, n - 1, AtomicOrd::AcqRel, AtomicOrd::Acquire) + .is_ok() + { + let _ = ctx.outputs.push(n); + Ok(StepOutcome::Progress) + } else { + Ok(StepOutcome::NoProgress) + } + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + let received = Arc::new(AtomicU32::new(0)); + let sink = ParallelCountingSink::new(&received); + let remaining = Arc::new(AtomicU32::new(5)); + let builder = PipelineBuilder::new(); + builder + .chain(FailingSource { remaining: Arc::clone(&remaining) }) + .chain(sink.clone()) + .into_sink_marker(); + let pipeline = builder.build().unwrap(); + + let result = pipeline.run(PipelineConfig { threads: 4, ..Default::default() }); + match result { + Err(PipelineError::Io { step, source }) => { + assert_eq!(step, "FailingSource"); + assert_eq!(source.kind(), std::io::ErrorKind::Other); + } + other => panic!("expected Io error, got {other:?}"), + } + } + + // ───── sleep_until_stop ───── + + /// When `stop` is already set, the helper must return effectively + /// immediately, never sleeping the full duration. This is the teardown + /// fast-path: the main thread sets `stop` then `join()`s, and the helper + /// must not block for a poll interval afterward. + #[test] + fn sleep_until_stop_returns_immediately_when_already_stopped() { + let stop = StopSignal::default(); + stop.stop(); + let start = std::time::Instant::now(); + sleep_until_stop(&stop, std::time::Duration::from_secs(10)); + assert!( + start.elapsed() < std::time::Duration::from_millis(100), + "expected near-immediate return, took {:?}", + start.elapsed() + ); + } + + /// When `stop` is set partway through the sleep, the helper must wake and + /// return well before the full duration elapses (within a couple of poll + /// slices), proving it interrupts a long sleep rather than waiting it out. + #[test] + fn sleep_until_stop_wakes_when_stopped_midway() { + let stop = Arc::new(StopSignal::default()); + let stop_clone = Arc::clone(&stop); + let setter = std::thread::spawn(move || { + std::thread::sleep(std::time::Duration::from_millis(50)); + stop_clone.stop(); + }); + let start = std::time::Instant::now(); + sleep_until_stop(&stop, std::time::Duration::from_secs(10)); + let elapsed = start.elapsed(); + setter.join().unwrap(); + // Woke shortly after the 50ms flag flip — far below the 10s budget. + assert!(elapsed >= std::time::Duration::from_millis(40), "woke too early: {elapsed:?}"); + assert!(elapsed < std::time::Duration::from_millis(500), "woke too late: {elapsed:?}"); + } + + /// A duration so large that `Instant::now() + dur` is unrepresentable must + /// not panic: the helper falls back to an untimed wait and still wakes the + /// instant `stop()` fires. Regression test for the `checked_add` guard — + /// before it, `Duration::MAX` panicked here and silently disarmed the + /// monitor (whose join error `Pipeline::run` discards). + #[test] + fn sleep_until_stop_survives_an_unrepresentable_deadline() { + let stop = Arc::new(StopSignal::default()); + let stop_clone = Arc::clone(&stop); + let setter = std::thread::spawn(move || { + std::thread::sleep(std::time::Duration::from_millis(50)); + stop_clone.stop(); + }); + let start = std::time::Instant::now(); + sleep_until_stop(&stop, std::time::Duration::MAX); + let elapsed = start.elapsed(); + setter.join().unwrap(); + // Woke shortly after the 50ms flag flip, not after an (impossible) MAX wait. + assert!(elapsed >= std::time::Duration::from_millis(40), "woke too early: {elapsed:?}"); + assert!(elapsed < std::time::Duration::from_millis(500), "woke too late: {elapsed:?}"); + } + + /// When `stop` is never set, the helper sleeps for approximately the full + /// duration (it does not return early). Generous upper bound keeps it + /// non-flaky on a loaded CI host. + #[test] + fn sleep_until_stop_sleeps_full_duration_when_never_stopped() { + let stop = StopSignal::default(); + let start = std::time::Instant::now(); + sleep_until_stop(&stop, std::time::Duration::from_millis(100)); + let elapsed = start.elapsed(); + assert!(elapsed >= std::time::Duration::from_millis(95), "returned too early: {elapsed:?}"); + assert!(elapsed < std::time::Duration::from_secs(2), "ran far too long: {elapsed:?}"); + } + + #[test] + fn classify_stall_progress_is_healthy() { + // Any advance in the global progress counter clears the stall, + // regardless of how long the prior stall was or what is in flight. + assert_eq!(classify_stall(true, 999, 10, 60, 1_000_000), StallVerdict::Progressing); + } + + #[test] + fn classify_stall_empty_pipeline_is_starvation_not_deadlock() { + // No progress, but nothing is stuck anywhere: the pipeline is idle + // waiting on a slow upstream (e.g. a stdin pipe). Never fatal — mirrors + // legacy `check_deadlock_and_restore`'s starvation guard. + assert_eq!(classify_stall(false, 120, 10, 60, 0), StallVerdict::Starving); + } + + /// The `in_flight_bytes` probe is blind to queued items whose `heap_size()` + /// is 0, because `ByteBoundedQueue` accounts `T::heap_size()` only and never + /// `size_of::()`. A `ByteBounded` edge can therefore hold items and still + /// report zero bytes — which `classify_stall` reads as `Starving`, resetting + /// the stall clock on every poll so `deadlock_timeout_secs` never fires. + /// + /// This pins the accounting limitation documented on `in_flight_bytes`, so a + /// future change that starts counting `size_of::()` (closing the blind + /// spot) fails here and prompts the doc to be updated with it. + #[test] + fn byte_bounded_queue_of_zero_heap_items_reports_no_bytes_in_flight() { + use crate::queues::{ByteBoundedQueue, ItemQueue}; + + // `u32: HeapSize` reports 0 heap bytes (see `item.rs`). + let q = ByteBoundedQueue::::new(64 * 1024); + for i in 0..100u32 { + q.try_push(i).expect("64 KiB budget accepts zero-heap items"); + } + assert!(!q.is_empty(), "the items really are queued"); + assert_eq!( + q.current_bytes(), + 0, + "heap_size()-only accounting reports zero for zero-heap items" + ); + // ...and a wedge stranding exactly those items is unclassifiable. + assert_eq!( + classify_stall(false, 999, 10, 60, q.current_bytes()), + StallVerdict::Starving, + "a wedge holding only zero-heap items cannot reach Wedged" + ); + } + + #[test] + fn classify_stall_stuck_work_below_fatal_only_warns() { + // No progress with work stuck, but the stall has not persisted long + // enough to be sure it is a wedge rather than one slow dispatch. + assert_eq!(classify_stall(false, 15, 10, 60, 4096), StallVerdict::Stalled); + } + + #[test] + fn classify_stall_below_warn_threshold_keeps_watching() { + // Stalled with stuck work but not yet past the warn threshold: no-op. + assert_eq!(classify_stall(false, 5, 10, 60, 4096), StallVerdict::Watching); + } + + #[test] + fn classify_stall_stuck_work_past_fatal_is_wedged() { + // No progress with work stuck for >= the fatal threshold: a genuine + // wedge — fail fast instead of hanging forever. + assert_eq!(classify_stall(false, 60, 10, 60, 4096), StallVerdict::Wedged); + } + + /// L2.4: a Detached step legitimately blocked on EMPTY input (slow upstream) + /// must NOT trip the monitor. The Detached merge's data flows through the + /// internal `SortMergeSlot` table — invisible to `in_flight_bytes` — and its + /// framework input edge (setup events) is drained early, so while it waits + /// for decompress the byte-bounded edges feeding it are EMPTY: + /// `in_flight_bytes == 0` → `Starving` (progress-shaped), never `Wedged`, + /// no matter how long the wait. The starvation guard already encodes this; + /// this test pins the Detached interpretation against a regression. + #[test] + fn monitor_no_false_positive_on_idle_detached() { + // Long stall, well past the fatal threshold, but nothing is stuck on a + // visible byte-bounded edge (the Detached step is parked on its empty + // input). Must be Starving, not Wedged. + assert_eq!(classify_stall(false, 600, 10, 60, 0), StallVerdict::Starving); + // And the instant the slow producer feeds it, global progress advances. + assert_eq!(classify_stall(true, 600, 10, 60, 0), StallVerdict::Progressing); + } + + /// L2.4 (revision item 10): a *genuine* wedge of the Detached merge still + /// trips fatal. If the merge is truly stuck, its pool upstream (decompress) + /// can't push to the full slots either, so the byte-bounded spill→decompress + /// edge stays FULL with no progress — `in_flight_bytes > 0` past the fatal + /// window → `Wedged`. The slot-table blindness does not hide a real wedge, + /// because the upstream byte edge always reflects it. + #[test] + fn monitor_still_trips_on_genuine_detached_wedge() { + assert_eq!(classify_stall(false, 60, 10, 60, 1_048_576), StallVerdict::Wedged); + } + + /// A step with a single output branch of the given `QueueSpec`, used to + /// exercise the monitor-visibility transport check. + fn step_with_output_spec( + name: &'static str, + spec: QueueSpec, + ) -> Box { + #[derive(Clone)] + struct SpecStep { + name: &'static str, + spec: QueueSpec, + } + impl Step for SpecStep { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: self.name, + kind: StepKind::Serial, + sticky: false, + output_queues: vec![self.spec], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + Box::new(crate::erased::TypedStep::new(SpecStep { name, spec })) + } + + /// A `ChainGraph` whose per-step branch count matches each step's declared + /// `output_queues` length — the common shape where every branch has a spec. + fn graph_matching_specs( + steps: &[Box], + ) -> crate::topology::ChainGraph { + let mut g = crate::topology::ChainGraph::new(); + for step in steps { + g.register_step(step.profile().name, step.profile().output_queues.len()); + } + g + } + + #[test] + fn first_monitor_blind_transport_finds_count_and_unbounded() { + // All-ByteBounded: every transport is monitor-visible. + let visible = vec![ + step_with_output_spec("A", QueueSpec::ByteBounded { limit_bytes: 1 << 20 }), + step_with_output_spec("B", QueueSpec::ByteBounded { limit_bytes: 1 << 20 }), + ]; + assert!(first_monitor_blind_transport(&visible, &graph_matching_specs(&visible)).is_none()); + + // A CountBounded branch is flagged (it registers no byte probe). + let with_count = vec![ + step_with_output_spec("A", QueueSpec::ByteBounded { limit_bytes: 1 << 20 }), + step_with_output_spec("Blind", QueueSpec::CountBounded { capacity: 8 }), + ]; + let (name, spec) = + first_monitor_blind_transport(&with_count, &graph_matching_specs(&with_count)).unwrap(); + assert_eq!(name, "Blind"); + assert!(matches!(spec, QueueSpec::CountBounded { .. })); + + // Unbounded is flagged too. + let with_unbounded = vec![step_with_output_spec("U", QueueSpec::Unbounded)]; + assert!( + first_monitor_blind_transport(&with_unbounded, &graph_matching_specs(&with_unbounded)) + .is_some() + ); + } + + #[test] + fn first_monitor_blind_transport_flags_implicit_unbounded_branch() { + // A step whose graph branch count exceeds its declared `output_queues`: + // the extra branch resolves to `QueueSpec::Unbounded` (matching `dag()` + // and context-building), which is monitor-blind and must be flagged even + // though the step declared only one explicit, ByteBounded spec. + let steps = vec![step_with_output_spec( + "HasImplicitBranch", + QueueSpec::ByteBounded { limit_bytes: 1 << 20 }, + )]; + let mut graph = crate::topology::ChainGraph::new(); + graph.register_step("HasImplicitBranch", 2); // 2 branches, 1 explicit spec + let (name, spec) = first_monitor_blind_transport(&steps, &graph).unwrap(); + assert_eq!(name, "HasImplicitBranch"); + assert!(matches!(spec, QueueSpec::Unbounded), "implicit 2nd branch is Unbounded"); + } + + #[test] + fn ensure_monitor_visible_transports_ok_for_all_byte_bounded() { + // The armed-monitor invariant holds (returns Ok) when every output + // transport is ByteBounded. + let steps = vec![ + step_with_output_spec("A", QueueSpec::ByteBounded { limit_bytes: 1 << 20 }), + step_with_output_spec("B", QueueSpec::ByteBounded { limit_bytes: 1 << 20 }), + ]; + assert!(ensure_monitor_visible_transports(&steps, &graph_matching_specs(&steps)).is_ok()); + } + + #[test] + fn ensure_monitor_visible_transports_errs_on_count_bounded() { + // A monitor-blind transport on an armed pipeline is rejected with a + // graceful error (in every build, release included) instead of silently + // losing the wedge verdict or crashing the process. + let steps = vec![step_with_output_spec("Blind", QueueSpec::CountBounded { capacity: 8 })]; + let err = ensure_monitor_visible_transports(&steps, &graph_matching_specs(&steps)) + .expect_err("a CountBounded transport on an armed pipeline must be rejected"); + assert!( + matches!(err, PipelineError::MonitorBlindTransport { step: "Blind", .. }), + "expected MonitorBlindTransport for the blind step, got {err:?}" + ); + } + + /// Arming the monitor must enforce the monitor-visible-transport invariant + /// even with no stats handle attached. + /// + /// This goes through `Pipeline::run` rather than calling + /// `ensure_monitor_visible_transports` directly on purpose: the helper's own + /// unit tests above pass whatever condition guards the *call site*, so they + /// cannot catch a guard that skips the check. Only a run-level test can. + /// + /// The failure it pins is silent, which is what makes it worth a test: on a + /// blind edge `in_flight_bytes` reports 0, `classify_stall` reads that as + /// `Starving`, and `Starving` resets the stall clock on every poll — so an + /// armed monitor would watch a wedged pipeline forever and never fail it. + /// + /// The converse — that a blind transport is still legal on a *disarmed* run + /// — needs no test of its own: most chains in this suite pair a + /// `CountBounded` edge with the default `deadlock_timeout_secs: 0`, so an + /// over-broad guard would fail them by the hundred. + #[test] + fn pipeline_run_rejects_a_monitor_blind_transport_without_stats() { + // `StubSource` declares `CountBounded`, which the byte-accounting probe + // cannot see. The chain never terminates (`StubSinkU32` always reports + // `NoProgress`), which is fine here and load-bearing for the assertion: + // the transport check runs before any worker is spawned, so a passing + // run proves the rejection happened at startup rather than after any + // work. + let builder = PipelineBuilder::new(); + builder.chain(StubSource).chain(StubSinkU32).into_sink_marker(); + let pipeline = builder.build().unwrap(); + + let err = pipeline + .run(PipelineConfig { + threads: 2, + stats: None, + deadlock_timeout_secs: 5, + ..Default::default() + }) + .expect_err("an armed monitor on a blind transport must fail the run"); + assert!( + matches!(err, PipelineError::MonitorBlindTransport { step: "Source", .. }), + "expected MonitorBlindTransport for the blind source, got {err:?}", + ); + } + + #[test] + fn apply_stall_verdict_wedged_records_timeout_and_cancels() { + use crate::signal::PipelineError; + let stats = PipelineStats::new(vec![]); + let signal = PipelineSignal::new(); + let now = std::time::Instant::now(); + let mut mon_state = StallMonitorState { last_total: 0, stall_start: now, last_warn: None }; + + let stop = apply_stall_verdict( + StallVerdict::Wedged, + now, + 0, + 60, + 4096, + std::time::Duration::from_secs(10), + Some(&stats), + &signal, + &mut mon_state, + ); + + assert!(stop, "a wedge must stop the monitor"); + assert!(signal.is_done(), "a wedge must make workers observe is_done()"); + match signal.outcome() { + Some(PipelineError::TimedOut { stalled_secs }) => assert_eq!(*stalled_secs, 60), + other => panic!("expected TimedOut, got {other:?}"), + } + // Exercises the `TimedOut` Display arm. + assert!(signal.outcome().unwrap().to_string().contains("no progress")); + } + + #[test] + fn apply_stall_verdict_stalled_warns_once_per_window() { + let stats = PipelineStats::new(vec![]); + let signal = PipelineSignal::new(); + let now = std::time::Instant::now(); + let mut mon_state = StallMonitorState { last_total: 0, stall_start: now, last_warn: None }; + let warn = std::time::Duration::from_secs(10); + + // First stall in the window arms the throttle without failing the run. + let stop = apply_stall_verdict( + StallVerdict::Stalled, + now, + 0, + 15, + 4096, + warn, + Some(&stats), + &signal, + &mut mon_state, + ); + assert!(!stop); + assert!(mon_state.last_warn.is_some(), "a stall must arm the warn throttle"); + assert!(!signal.is_done(), "a stall must not fail the run"); + + // A second stall inside the same window must not re-arm (no re-warn). + let armed = mon_state.last_warn; + let stop2 = apply_stall_verdict( + StallVerdict::Stalled, + now, + 0, + 16, + 4096, + warn, + Some(&stats), + &signal, + &mut mon_state, + ); + assert!(!stop2); + assert_eq!(mon_state.last_warn, armed, "must not re-warn within the same window"); + } + + #[test] + fn apply_stall_verdict_progressing_and_starving_reset_the_clock() { + let stats = PipelineStats::new(vec![]); + let signal = PipelineSignal::new(); + let t0 = std::time::Instant::now(); + let warn = std::time::Duration::from_secs(10); + let mut mon_state = + StallMonitorState { last_total: 0, stall_start: t0, last_warn: Some(t0) }; + + // Progress advances the watermark and clears the warn throttle. + let later = t0 + std::time::Duration::from_secs(5); + let stop = apply_stall_verdict( + StallVerdict::Progressing, + later, + 42, + 0, + 0, + warn, + Some(&stats), + &signal, + &mut mon_state, + ); + assert!(!stop); + assert_eq!(mon_state.last_total, 42); + assert_eq!(mon_state.stall_start, later); + assert!(mon_state.last_warn.is_none()); + assert!(!signal.is_done()); + + // Starvation (nothing in flight) resets the clock without failing. + mon_state.last_warn = Some(t0); + let even_later = later + std::time::Duration::from_secs(5); + let stop2 = apply_stall_verdict( + StallVerdict::Starving, + even_later, + 99, + 0, + 0, + warn, + Some(&stats), + &signal, + &mut mon_state, + ); + assert!(!stop2); + assert_eq!(mon_state.stall_start, even_later); + assert!(mon_state.last_warn.is_none()); + assert!(!signal.is_done()); + } + + #[test] + fn in_flight_bytes_counts_transport_and_reorder_stash() { + // The wedge-vs-starvation signal must include the reorder overflow + // stash, not just transport queues — items can sit in a reorder buffer + // (waiting for a missing serial) while every transport reads empty. + use crate::item::HeapSize; + use crate::queues::{BoundedQueueHandle, ByteBoundedQueue, CountBoundedQueue, ItemQueue}; + use crate::reorder::{ReorderCapHandle, ReorderStage, Sequenced}; + use crate::runtime::contexts::{ChainContexts, RegisteredQueue}; + use crate::topology::{BranchIdx, StepIdx}; + + #[derive(Debug)] + struct Heavy(Vec); + impl HeapSize for Heavy { + fn heap_size(&self) -> usize { + self.0.len() + } + } + + // Transport queue holding 200 bytes. + let transport = Arc::new(ByteBoundedQueue::::new(10_000)); + transport.try_push(Heavy(vec![0u8; 200])).unwrap(); + assert_eq!(transport.current_bytes(), 200); + + // Reorder stage with 300 bytes stuck in its overflow stash: with + // next_serial (0) absent and the inner transport (cap 1) full, the + // second push overflows into the buffer. + let inner: Arc>> = + Arc::new(CountBoundedQueue::>::new(1)); + let reorder = Arc::new(ReorderStage::with_max_overflow_bytes(inner, 100_000)); + reorder.try_push(5, Heavy(vec![0u8; 100])).unwrap(); // -> inner transport + reorder.try_push(6, Heavy(vec![0u8; 300])).unwrap(); // -> overflow stash + assert_eq!(reorder.current_buffer_bytes(), 300); + + let handle: Arc = transport.clone(); + let cap: Arc = reorder.clone(); + let rq = RegisteredQueue { + producer_step_name: "test", + producer_step: StepIdx(0), + branch: BranchIdx(0), + handle, + reorder_cap: Some(cap), + }; + let contexts = ChainContexts { + inputs: vec![], + outputs: vec![], + bounded_queues: vec![rq], + edges: vec![], + }; + // 200 (transport) + 300 (reorder stash). + assert_eq!(in_flight_bytes(&contexts), 500); + + // An empty registry reads zero — the starvation (idle) signal. + let empty = ChainContexts { + inputs: vec![], + outputs: vec![], + bounded_queues: vec![], + edges: vec![], + }; + assert_eq!(in_flight_bytes(&empty), 0); + } +} diff --git a/crates/fgumi-pipeline-core/src/erased.rs b/crates/fgumi-pipeline-core/src/erased.rs new file mode 100644 index 000000000..29203a6e7 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/erased.rs @@ -0,0 +1,1491 @@ +//! Type erasure: `ErasedStep` trait + `TypedStep` adapter. +//! +//! The runtime holds a heterogeneous chain in `Vec>`. +//! `TypedStep` is the adapter that bridges between the type-erased +//! dispatch in the worker loop and the concrete `S::try_run` body. +//! +//! Each `ErasedStep` exposes the methods the runtime needs: +//! - `clone_boxed` — make per-worker copies for `Parallel` steps +//! - `build_output_set` — construct the producer's queue set + view from +//! `StepProfile::output_queues` + `branch_ordering` +//! - `build_input_handle` — pull the consumer's typed input handle out of +//! the producer's `OutputQueueSet` (mutable: takes ownership of the +//! branch's slot) +//! - `wrap_outputs_view` — wrap the type-erased view into a typed +//! `OutputHandles` for the worker to pass into `ctx.outputs` +//! - `mark_outputs_drained` — close all output branches (called by the +//! driver when a step returns `StepOutcome::Finished`, counter-gated for +//! `Parallel` so only the last clone closes the shared output) +//! - `is_source` — true iff `S::Input == ()` (used by chain-context +//! construction to pick the source's unit-input path) + +use std::any::{Any, TypeId}; +use std::io; +use std::marker::PhantomData; +use std::sync::Arc; + +use super::handles::{BranchInputHandle, OutputQueueSet}; +use super::outputs::StepOutputs; +use super::reorder::BranchOrdering; +use super::signal::PipelineSignal; +use super::step::{ + Affinity, DetachedGroup, OutputHandles, OutputsViewAny, Step, StepCtx, StepKind, StepOutcome, + StepProfile, +}; + +/// The branch orderings the framework actually *builds* for a single-input +/// producer of the given [`StepKind`], given its declared profile orderings. +/// +/// `Serial` / `Exclusive` producers emit items in arrival order by construction +/// (the framework's mutex serializes pushes; an Exclusive-owned step has a +/// single dispatcher), so inserting a `ReorderStage` on their output edges is +/// pure overhead. The framework collapses any declared +/// `ByOrdinal` / `ByItemOrdinal` to [`BranchOrdering::None`] for those kinds — +/// `build_queues` then constructs the direct transport with no reorder stage. +/// `Parallel` producers keep their declared orderings verbatim. +/// +/// # Precondition for `ByItemOrdinal` +/// +/// The collapse preserves *arrival* order, which equals *ordinal* order only if +/// the producer already receives its input in ordinal order. A single-input +/// `Serial` / `Exclusive` / `Detached` step fed by a +/// [`BranchOrdering::None`] edge can be handed items out of ordinal order +/// (multiple `Parallel` upstream workers push in nondeterministic order), and +/// collapsing its declared `ByItemOrdinal` output to `None` then propagates that +/// disorder to a consumer that asked for ordinal order. +/// +/// So a chain declaring `ByItemOrdinal` must keep every edge upstream of that +/// producer ordered — which the canonical BAM pattern does: the source declares +/// an ordering and every intermediate transform propagates the input's serial +/// onto its outputs. Declaring `ByItemOrdinal` downstream of an unordered +/// (`None`) edge is a chain-construction error, not something this collapse can +/// detect: the rule is a function of `(kind, declared)` only and has no view of +/// the graph. +/// +/// This is the single source of truth for the collapse rule, shared by +/// [`TypedStep::build_output_set`] (which constructs the transport) and +/// `Pipeline::dag` (which must render the *effective* — post-collapse — +/// ordering so the diagnostic matches the transport actually built). Note +/// `Step2` producers (`TypedStep2`) do **not** collapse, so this helper is only +/// for single-input steps. +#[must_use] +pub(crate) fn effective_branch_orderings( + kind: StepKind, + declared: &[BranchOrdering], +) -> Vec { + match kind { + StepKind::Parallel => declared.to_vec(), + // Single-producer kinds emit in-order from one thread, so a downstream + // reorder stage is redundant. `Detached` is single-producer too (one + // dedicated thread), so it collapses like `Serial`/`Exclusive`. + StepKind::Serial | StepKind::Exclusive | StepKind::Detached => { + declared.iter().map(|_| BranchOrdering::None).collect() + } + } +} + +/// Type-erased step interface used by the worker loop. +pub trait ErasedStep: Send + 'static { + fn profile(&self) -> StepProfile; + + /// The step's static name, returned WITHOUT building a `StepProfile`. + /// + /// `dispatch_one_step` reads the name on every dispatch (for stats / + /// error reporting), so going through `profile()` there would heap- + /// allocate the profile's two `Vec`s per dispatch (a virtual call, so + /// the optimizer cannot elide them). Adapters cache the name at + /// construction and return it here for free. + fn name(&self) -> &'static str; + + /// The step's [`StepKind`], returned WITHOUT building a `StepProfile`. + /// + /// The setup passes (`pool::assign_exclusive_owners`, + /// `pool::assign_sticky_owners`, `storage::build_worker_storage`) read the + /// kind per step; going through `profile()` there heap-allocates the + /// profile's two `Vec`s per read (a virtual call, so the optimizer cannot + /// elide them). Adapters cache the kind at construction and return it free. + fn kind(&self) -> StepKind; + + /// Whether the step is `sticky`, returned WITHOUT building a `StepProfile`. + /// Same rationale as [`Self::kind`]: the sticky-owner assignment passes read + /// it per step. Adapters cache it at construction. + fn sticky(&self) -> bool; + + /// Forward `Step::affinity` for worker-eligibility gating. Only + /// consulted for `Serial` steps; the runtime calls this once during + /// `build_worker_storage` to decide which workers get a `Shared` + /// entry vs a `Skip` placeholder. + fn affinity(&self) -> Affinity; + + /// Forward `Step::detached_group` — which dedicated driver thread a + /// `Detached` step runs on. Read once by `extract_detached_steps` to group + /// detached steps onto shared driver threads (the N+2 model). Only + /// meaningful for `Detached` kinds. + fn detached_group(&self) -> DetachedGroup; + + /// Dispatch `S::try_run` after downcasting queue handles. + /// + /// # Errors + /// + /// Forwards any I/O error from the step body. + fn try_run_erased(&mut self, ctx: &mut ErasedStepCtx<'_>) -> io::Result; + + /// Construct a fresh per-worker copy of this step. Used for Parallel + /// steps. Cheap — implementing types call `S::clone()`, where typical + /// state is unit-struct or `Arc` (one atomic increment). + fn clone_boxed(&self) -> Box; + + /// Take ownership of the consumer's input handle from the producer's + /// output queue set. Used by the framework when constructing chain + /// topology; each branch is taken exactly once. + fn build_input_handle( + &self, + producer_set: &mut OutputQueueSet, + branch_idx: usize, + ) -> Box; + + /// Input arity. Default `1` for single-input + /// [`crate::step::Step`] impls; multi-input + /// adapters (`TypedStep2`, future `StepN`) override to return their + /// arity. Used by [`crate::runtime::contexts`] + /// to decide which input-construction path to take. + fn input_arity(&self) -> usize { + 1 + } + + /// Take ownership of TWO consumer input handles from two + /// upstream output queue sets, paired into a typed + /// [`crate::handles::TwoInputHandles`] + /// per the consumer's [`crate::step::Step2`] + /// associated types. + /// + /// Default impl panics — single-input steps never have arity 2. + /// `TypedStep2` overrides to call `take_typed_input` twice + /// (once per input slot) and wrap the pair. + /// + /// # Panics + /// + /// Default impl always panics; multi-input adapters override. + fn build_two_input_handles( + &self, + _producer_sets: &mut [OutputQueueSet], + _p0_idx: usize, + _p0_branch: usize, + _p1_idx: usize, + _p1_branch: usize, + ) -> Box { + let p = self.profile(); + panic!( + "build_two_input_handles called on '{}' (kind = {:?}); \ + only Step2 adapters (`TypedStep2`) support arity-2 input \ + construction. This is a framework bug — chain-build code \ + should only dispatch arity-2 input construction to steps \ + that override `input_arity` to return 2.", + p.name, p.kind + ); + } + + /// Build this step's output queue set + outputs view from the profile's + /// per-branch queue specs and ordering directives. + fn build_output_set( + &self, + level: crate::builder::InstrumentationLevel, + ) -> (OutputQueueSet, OutputsViewAny); + + /// Like [`Self::build_output_set`] but forces every output branch to a + /// **direct** transport (no reorder stage), keeping the profile's configured + /// queue bound. + /// + /// Used only by the single-thread *fused* driver + /// ([`crate::runtime::run_fused_single_thread`]). At one worker, FIFO push + /// order is already the correct order, so the reorder stage is dead weight. + /// The count/byte bound stays: the fused driver runs a producer before its + /// consumer in each pass, so a step emitting more items per `try_run` than + /// its consumer removes would grow the edge every pass. The bound turns that + /// into ordinary backpressure — the producer holds and retries, as it must + /// under the scheduled driver anyway — instead of unbounded growth. + fn build_fused_output_set( + &self, + level: crate::builder::InstrumentationLevel, + ) -> (OutputQueueSet, OutputsViewAny); + + /// Wrap a typed `OutputsViewAny` into `Box>` + /// (type-erased as `Any`). The runtime stores this box per step so the + /// worker loop's `try_run_erased` can downcast it to the typed + /// `OutputHandles` the step expects in its `ctx.outputs`. + fn wrap_outputs_view(&self, view: OutputsViewAny) -> Box; + + /// Mark all of this step's output branches drained. Downcasts the + /// type-erased outputs handle to `OutputHandles` and calls + /// the typed `mark_all_drained`. Called by the driver when a step returns + /// `StepOutcome::Finished` (counter-gated for `Parallel`). + fn mark_outputs_drained(&self, outputs: &(dyn Any + Send + Sync)); + + /// Returns `true` iff `S::Input = ()` — i.e., this is a source step. + /// Used by chain-context construction to pick the source's unit-input + /// path (a source has no upstream queue to take an input handle from). + fn is_source(&self) -> bool; +} + +/// Identity of a type-erased handle box: the address of the value the +/// `&dyn Any` points at. +/// +/// The typed-handle caches below record this alongside each cached pointer and +/// compare it on every cached hit, so a step handed a *different* +/// `ChainContexts` is rejected instead of reading through a stale pointer. A +/// `usize` rather than a raw pointer keeps `TypedStep` / `TypedStep2` auto-`Send` +/// (raw pointers are `!Send`); the value is only ever compared, never +/// dereferenced. +fn erased_addr(handle: &(dyn Any + Send + Sync)) -> usize { + std::ptr::from_ref(handle).cast::<()>().addr() +} + +/// One slot of a step adapter's typed-handle cache: the resolved handle plus the +/// [`erased_addr`] of the `&dyn Any` box it came from. +/// +/// The `'static` lifetime is a lie — the real bound is `ChainContexts` — held up +/// by the invariants documented on [`TypedStep`], the third of which is what the +/// stored address is checked against on every cached hit. +type CachedHandle = Option<(usize, &'static H)>; + +/// Context the worker loop hands to `ErasedStep` methods. +pub struct ErasedStepCtx<'a> { + /// Boxed `BranchInputHandle`. Adapter downcasts. + pub input: &'a (dyn Any + Send + Sync), + /// `OutputHandles`. Adapter downcasts. + pub outputs: &'a (dyn Any + Send + Sync), + /// Shared signal (error/cancel). Workers consult; steps don't directly. + pub signal: &'a Arc, +} + +/// Adapter that wraps a concrete `Step` impl as an `ErasedStep`. +/// +/// ## Cached typed handles +/// +/// `try_run_erased` needs typed +/// `&BranchInputHandle` / `&OutputHandles` views +/// of the type-erased context boxes. The first call resolves them via +/// `Any::downcast_ref` (a `TypeId` compare + transmute); subsequent calls +/// reuse the cached references through the `cached_input` / +/// `cached_outputs` fields. Without the cache, a 4-thread CODEC 8M +/// run pays ≈230 samples (~1.6%) on `downcast_ref` `TypeId` compares +/// across the dispatch hot path; the cache eliminates them entirely. +/// +/// The cache is sound because: +/// +/// 1. The boxes are owned by `ChainContexts` (an `Arc` held alive +/// for the entire `Pipeline::run` call by every worker thread). +/// 2. Every `TypedStep` instance — owned (`Parallel`), shared +/// (`Serial`, behind `Mutex`), or pinned (`Exclusive`) — is +/// destroyed before `ChainContexts` goes out of scope (workers +/// exit, then the runtime drops the contexts). +/// 3. Every dispatch passes the **same** box reference for a given +/// `step_idx` (see `run_worker_loop`'s +/// `contexts.inputs[step_idx.0].as_ref()`). The cached pointer +/// always refers to that same box. +/// +/// Only (3) is invisible to the compiler, so each cache slot also records the +/// `erased_addr` of the box it was resolved from and every cached hit asserts +/// it — unconditionally, in release too. A step handed a different, +/// **still-live** `ChainContexts` panics on the spot instead of reading through +/// a pointer into the other one. +/// +/// That check is a best-effort guard, not a proof. It compares data addresses +/// only, so it cannot detect a *freed* box whose address the allocator handed +/// back out: a new `ChainContexts` whose input box lands on a dead one's address +/// passes the assert and the `transmute` reads through a dangling pointer. +/// Soundness therefore still rests entirely on invariant (2) — every +/// `TypedStep` is dropped before the `ChainContexts` it cached from. Relaxing +/// (2) is not made safe by this guard. +/// +/// `clone_boxed` calls `TypedStep::new(...)` which initializes the +/// cache to `None`, so per-worker clones (`Parallel` steps) start +/// with a fresh cache. +pub struct TypedStep { + inner: S, + /// Static step name, cached from `inner.profile().name` at + /// construction so `ErasedStep::name()` (read per dispatch) never + /// rebuilds the profile's `Vec`s. See the `ErasedStep::name` doc. + name: &'static str, + /// `StepKind`, cached at construction so the setup passes read it without + /// rebuilding the profile's `Vec`s. See the `ErasedStep::kind` doc. + kind: StepKind, + /// `sticky` flag, cached at construction for the same reason as `kind`. + sticky: bool, + /// Cached downcast of `ctx.input` after the first dispatch, paired with the + /// [`erased_addr`] of the box it was resolved from (checked on every cached + /// hit). `'static` is a lie — the actual lifetime is bounded by + /// `ChainContexts` — enforced via the `// SAFETY:` comment below. + cached_input: CachedHandle>, + /// Cached downcast of `ctx.outputs`, with its box address. Same lifetime story. + cached_outputs: CachedHandle>, + _phantom: PhantomData S>, +} + +impl TypedStep { + pub fn new(step: S) -> Self { + let profile = step.profile(); + let name = profile.name; + let kind = profile.kind; + let sticky = profile.sticky; + Self { + inner: step, + name, + kind, + sticky, + cached_input: None, + cached_outputs: None, + _phantom: PhantomData, + } + } + + /// Resolve the typed input handle from the dispatch context, caching + /// the result for subsequent dispatches. + /// + /// # Panics + /// + /// Panics on the first dispatch if the `ctx.input` box doesn't + /// downcast to `BranchInputHandle` (a chain topology + /// invariant violation; the builder should have caught this). + #[allow(unsafe_code)] + fn resolve_input<'a>(&mut self, ctx: &ErasedStepCtx<'a>) -> &'a BranchInputHandle { + if let Some((cached_addr, cached)) = self.cached_input { + // Point 3 -- every dispatch passes the same box for a given + // `step_idx` -- is the invariant the cache rests on, and the one the + // type system cannot enforce. Violating it makes `cached` point at a + // box this dispatch does not own, so reject a mismatch in EVERY + // build: a `debug_assert!` alone would leave release builds reading + // through a stale pointer. The check is one load and one compare — + // not the `downcast_ref` `TypeId` probe this cache exists to elide. + assert_eq!( + erased_addr(ctx.input), + cached_addr, + "cached input handle does not match the dispatch context — a \ + TypedStep was reused across two ChainContexts" + ); + // Same box implies the same typed pointer, so this adds nothing in + // release; keep it in debug as a check that the cached pointer was + // derived from that box correctly in the first place. + debug_assert!( + ctx.input + .downcast_ref::>() + .is_some_and(|live| std::ptr::eq(live, cached)), + "cached input handle disagrees with a fresh downcast of the same box" + ); + // SAFETY: Lifetime extension from `'static` (cache slot) back + // to `'a` (the dispatch context's lifetime). The cached + // pointer was originally a `&'a BranchInputHandle` + // pulled out of `ChainContexts.inputs[step_idx]`; that box + // outlives every `TypedStep` (point 2 in the type-level + // doc). Dispatches always pass the same box for the same + // `step_idx` (point 3), so the pointer is still pointing + // at the live box on every subsequent dispatch. + return unsafe { + std::mem::transmute::<&BranchInputHandle, &'a BranchInputHandle>( + cached, + ) + }; + } + let r: &'a BranchInputHandle = ctx + .input + .downcast_ref::>() + .expect("input handle downcast failed — chain topology invariant"); + // SAFETY: Lifetime extension from `'a` to `'static` for storage. + // The same point-2/point-3 invariants from `// SAFETY:` above + // apply: the box outlives `self`, and we only ever read the + // cache through `resolve_input`, which immediately re-extends + // back to a bounded `'a` before handing it to user code. + let cached: &'static BranchInputHandle = unsafe { + std::mem::transmute::< + &'a BranchInputHandle, + &'static BranchInputHandle, + >(r) + }; + self.cached_input = Some((erased_addr(ctx.input), cached)); + r + } + + /// Resolve the typed outputs handle from the dispatch context, + /// caching the result for subsequent dispatches. See `resolve_input` + /// for the safety argument. + #[allow(unsafe_code)] + fn resolve_outputs<'a>(&mut self, ctx: &ErasedStepCtx<'a>) -> &'a OutputHandles { + if let Some((cached_addr, cached)) = self.cached_outputs { + // See `resolve_input` for why the address check is unconditional and + // the downcast re-check is debug-only. + assert_eq!( + erased_addr(ctx.outputs), + cached_addr, + "cached outputs handle does not match the dispatch context — a \ + TypedStep was reused across two ChainContexts" + ); + debug_assert!( + ctx.outputs + .downcast_ref::>() + .is_some_and(|live| std::ptr::eq(live, cached)), + "cached outputs handle disagrees with a fresh downcast of the same box" + ); + // SAFETY: see `resolve_input`; same boxes/lifetimes story. + return unsafe { + std::mem::transmute::<&OutputHandles, &'a OutputHandles>( + cached, + ) + }; + } + let r: &'a OutputHandles = ctx + .outputs + .downcast_ref::>() + .expect("outputs handle downcast failed — chain topology invariant"); + // SAFETY: see `resolve_input`; same boxes/lifetimes story. + let cached: &'static OutputHandles = unsafe { + std::mem::transmute::<&'a OutputHandles, &'static OutputHandles>( + r, + ) + }; + self.cached_outputs = Some((erased_addr(ctx.outputs), cached)); + r + } +} + +impl ErasedStep for TypedStep +where + S: Step, +{ + fn profile(&self) -> StepProfile { + self.inner.profile() + } + + fn name(&self) -> &'static str { + self.name + } + + fn kind(&self) -> StepKind { + self.kind + } + + fn sticky(&self) -> bool { + self.sticky + } + + fn affinity(&self) -> Affinity { + self.inner.affinity() + } + + fn detached_group(&self) -> DetachedGroup { + self.inner.detached_group() + } + + fn try_run_erased(&mut self, ctx: &mut ErasedStepCtx<'_>) -> io::Result { + let input = self.resolve_input(ctx); + let outputs = self.resolve_outputs(ctx); + let mut step_ctx = StepCtx { input, outputs }; + self.inner.try_run(&mut step_ctx) + } + + fn clone_boxed(&self) -> Box { + // `TypedStep::new` initializes `cached_input` / `cached_outputs` + // to `None` so the new clone resolves them fresh on its first + // dispatch (the box references in *this* `TypedStep`'s cache + // are still valid for the new clone too — they refer to the + // same `ChainContexts` boxes — but resolving fresh is simpler + // and keeps the cache lifetime story local to each clone). + Box::new(TypedStep::new(self.inner.new_worker_copy())) + } + + fn build_input_handle( + &self, + producer_set: &mut OutputQueueSet, + branch_idx: usize, + ) -> Box { + let handle: BranchInputHandle = + producer_set.take_typed_input::(branch_idx); + Box::new(handle) + } + + fn build_output_set( + &self, + level: crate::builder::InstrumentationLevel, + ) -> (OutputQueueSet, OutputsViewAny) { + let profile = self.inner.profile(); + // Producers with `StepKind::Serial` or `StepKind::Exclusive` emit + // items in arrival order by construction (the framework's mutex + // serializes pushes; an Exclusive-owned step has a single + // dispatcher). Inserting a `ReorderStage` on those output edges + // is pure overhead — every push pays an ordinal allocation + + // `Sequenced` wrap + reorder mutex hop for items that are + // already in order. The framework collapses any + // `BranchOrdering::ByOrdinal` / `ByItemOrdinal` declaration to + // `None` here, so `build_queues` constructs the direct transport + // path with no reorder stage. Step authors keep declaring + // `ByItemOrdinal` (preserves the intent in the profile; + // documents what the consumer needs); the framework is just + // smart enough to skip the wrap when the producer can already + // satisfy that need. + let effective_orderings = + effective_branch_orderings(profile.kind, &profile.branch_ordering); + ::build_queues( + &profile.output_queues, + &effective_orderings, + level, + ) + } + + fn build_fused_output_set( + &self, + level: crate::builder::InstrumentationLevel, + ) -> (OutputQueueSet, OutputsViewAny) { + // Force a DIRECT transport on every branch — at one worker FIFO push + // order is already the correct order, so the reorder stage is pure + // overhead. The profile's queue specs are kept as-is: the fused driver + // dispatches a producer before its consumer in each pass, so a step that + // emits k > 1 items per `try_run` while its consumer takes one would + // grow the edge by k - 1 every pass. Dropping the bound turns that into + // unbounded memory growth; keeping it makes the producer hold and retry, + // exactly as it already must under the scheduled driver. + let profile = self.inner.profile(); + let orderings = vec![BranchOrdering::None; profile.output_queues.len()]; + ::build_queues(&profile.output_queues, &orderings, level) + } + + fn wrap_outputs_view(&self, view: OutputsViewAny) -> Box { + let outputs: OutputHandles = OutputHandles::new(view); + Box::new(outputs) + } + + fn mark_outputs_drained(&self, outputs: &(dyn Any + Send + Sync)) { + let typed = outputs + .downcast_ref::>() + .expect("outputs handle downcast failed in mark_outputs_drained"); + ::mark_all_drained(typed); + } + + fn is_source(&self) -> bool { + TypeId::of::() == TypeId::of::<()>() + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// TypedStep2 — adapter for `Step2` impls (two-input merge steps). +// +// Same shape as TypedStep for outputs: one `OutputHandles` +// cached after the first dispatch. Inputs differ — instead of a single +// `BranchInputHandle`, the adapter caches a single +// `&TwoInputHandles` and lends per-branch refs +// (`ctx.a`, `ctx.b`) into the wrapped `Step2::try_run`. +// +// Drain detection: a `Step2` consumer is "drained" when **both** input +// branches report drained. The step itself checks this in `try_run` +// (`ctx.a.is_drained() && ctx.b.is_drained()`) to decide when to report +// `Finished`; the typed-erased input box is a `TwoInputHandles`. +// ───────────────────────────────────────────────────────────────────────────── + +use super::handles::TwoInputHandles; +use super::step::{Step2, StepCtx2}; + +/// Adapter wrapping a [`Step2`] impl as an [`ErasedStep`]. +/// +/// Cached typed handles follow the same pattern as [`TypedStep`]: +/// the first dispatch resolves `&TwoInputHandles` +/// / `&OutputHandles` via `Any::downcast_ref`, subsequent +/// dispatches reuse the cached references through unsafe lifetime +/// extension. The safety argument is identical (the boxes are owned +/// by `ChainContexts` which outlives every adapter instance, and +/// every dispatch passes the same box reference for a given +/// `step_idx`). +pub struct TypedStep2 { + inner: S, + /// Static step name, cached at construction. See `ErasedStep::name`. + name: &'static str, + /// `StepKind`, cached at construction. See `ErasedStep::kind`. + kind: StepKind, + /// `sticky` flag, cached at construction. See `ErasedStep::sticky`. + sticky: bool, + /// Cached downcasts, each paired with the [`erased_addr`] of the box it was + /// resolved from — see [`CachedHandle`]. + cached_inputs: CachedHandle>, + cached_outputs: CachedHandle>, + _phantom: PhantomData S>, +} + +impl TypedStep2 { + pub fn new(step: S) -> Self { + let profile = step.profile(); + let name = profile.name; + let kind = profile.kind; + let sticky = profile.sticky; + Self { + inner: step, + name, + kind, + sticky, + cached_inputs: None, + cached_outputs: None, + _phantom: PhantomData, + } + } + + #[allow(unsafe_code)] + fn resolve_inputs<'a>( + &mut self, + ctx: &ErasedStepCtx<'a>, + ) -> &'a TwoInputHandles { + if let Some((cached_addr, cached)) = self.cached_inputs { + // See `TypedStep::resolve_input` for why the address check is + // unconditional and the downcast re-check is debug-only. + assert_eq!( + erased_addr(ctx.input), + cached_addr, + "cached input handles do not match the dispatch context — a \ + TypedStep2 was reused across two ChainContexts" + ); + debug_assert!( + ctx.input + .downcast_ref::>() + .is_some_and(|live| std::ptr::eq(live, cached)), + "cached input handles disagree with a fresh downcast of the same box" + ); + // SAFETY: lifetime extension from `'static` (cache slot) + // back to `'a` (the dispatch context's lifetime). The + // cached pointer was originally a + // `&'a TwoInputHandles` pulled out + // of `ChainContexts.inputs[step_idx]`; that box outlives + // every `TypedStep2` instance (same point-2/point-3 + // invariants as `TypedStep::resolve_input`). + return unsafe { + std::mem::transmute::< + &TwoInputHandles, + &'a TwoInputHandles, + >(cached) + }; + } + let r: &'a TwoInputHandles = ctx + .input + .downcast_ref::>() + .expect("input handle downcast failed — Step2 chain topology invariant"); + // SAFETY: lifetime extension from `'a` to `'static` for + // storage. The box outlives `self`; we only ever read the + // cache through `resolve_inputs`, which immediately re-extends + // back to a bounded `'a` before handing it to user code. + let cached: &'static TwoInputHandles = unsafe { + std::mem::transmute::< + &'a TwoInputHandles, + &'static TwoInputHandles, + >(r) + }; + self.cached_inputs = Some((erased_addr(ctx.input), cached)); + r + } + + #[allow(unsafe_code)] + fn resolve_outputs<'a>(&mut self, ctx: &ErasedStepCtx<'a>) -> &'a OutputHandles { + if let Some((cached_addr, cached)) = self.cached_outputs { + // See `TypedStep::resolve_input` for why the address check is + // unconditional and the downcast re-check is debug-only. + assert_eq!( + erased_addr(ctx.outputs), + cached_addr, + "cached outputs handle does not match the dispatch context — a \ + TypedStep2 was reused across two ChainContexts" + ); + debug_assert!( + ctx.outputs + .downcast_ref::>() + .is_some_and(|live| std::ptr::eq(live, cached)), + "cached outputs handle disagrees with a fresh downcast of the same box" + ); + // SAFETY: lifetime extension from `'static` (cache slot) back to + // `'a`; same boxes/lifetimes story as `resolve_inputs`. + return unsafe { + std::mem::transmute::<&OutputHandles, &'a OutputHandles>( + cached, + ) + }; + } + let r: &'a OutputHandles = ctx + .outputs + .downcast_ref::>() + .expect("outputs handle downcast failed — Step2 chain topology invariant"); + // SAFETY: lifetime extension from `'a` to `'static` for storage; same + // boxes/lifetimes story as `resolve_inputs`. + let cached: &'static OutputHandles = unsafe { + std::mem::transmute::<&'a OutputHandles, &'static OutputHandles>( + r, + ) + }; + self.cached_outputs = Some((erased_addr(ctx.outputs), cached)); + r + } +} + +impl ErasedStep for TypedStep2 { + fn profile(&self) -> StepProfile { + self.inner.profile() + } + + fn name(&self) -> &'static str { + self.name + } + + fn kind(&self) -> StepKind { + self.kind + } + + fn sticky(&self) -> bool { + self.sticky + } + + fn affinity(&self) -> Affinity { + self.inner.affinity() + } + + fn detached_group(&self) -> DetachedGroup { + self.inner.detached_group() + } + + fn try_run_erased(&mut self, ctx: &mut ErasedStepCtx<'_>) -> io::Result { + let inputs = self.resolve_inputs(ctx); + let outputs = self.resolve_outputs(ctx); + let mut typed_ctx = StepCtx2:: { a: &inputs.a, b: &inputs.b, outputs }; + self.inner.try_run(&mut typed_ctx) + } + + fn clone_boxed(&self) -> Box { + Box::new(TypedStep2::new(self.inner.new_worker_copy())) + } + + fn build_input_handle( + &self, + _producer_set: &mut OutputQueueSet, + _branch_idx: usize, + ) -> Box { + let p = self.profile(); + panic!( + "build_input_handle called on Step2 adapter '{}' (kind = {:?}); \ + multi-input steps build their inputs via build_two_input_handles. \ + This is a framework bug — chain-build code should dispatch on \ + input_arity().", + p.name, p.kind + ); + } + + fn input_arity(&self) -> usize { + 2 + } + + fn build_two_input_handles( + &self, + producer_sets: &mut [OutputQueueSet], + p0_idx: usize, + p0_branch: usize, + p1_idx: usize, + p1_branch: usize, + ) -> Box { + if p0_idx == p1_idx { + assert_ne!( + p0_branch, + p1_branch, + "Step2 inputs must consume distinct branches when they share \ + producer step {p0_idx} for step '{}'.", + self.profile().name + ); + let set = &mut producer_sets[p0_idx]; + let a: BranchInputHandle = set.take_typed_input::(p0_branch); + let b: BranchInputHandle = set.take_typed_input::(p1_branch); + return Box::new(TwoInputHandles::::new(a, b)); + } + // Borrow two disjoint elements of `producer_sets` simultaneously. + // `split_at_mut(lo+1)` puts producer_sets[lo] in the first half; + // we index into the second half for the hi side. + let (lo_idx, hi_idx, swap) = + if p0_idx < p1_idx { (p0_idx, p1_idx, false) } else { (p1_idx, p0_idx, true) }; + let (lo_half, hi_half) = producer_sets.split_at_mut(lo_idx + 1); + let lo_set: &mut OutputQueueSet = &mut lo_half[lo_idx]; + let hi_set: &mut OutputQueueSet = &mut hi_half[hi_idx - (lo_idx + 1)]; + let (a_set, b_set) = if swap { + (hi_set, lo_set) // p0 = high, p1 = low + } else { + (lo_set, hi_set) // p0 = low, p1 = high + }; + let a: BranchInputHandle = a_set.take_typed_input::(p0_branch); + let b: BranchInputHandle = b_set.take_typed_input::(p1_branch); + Box::new(TwoInputHandles::::new(a, b)) + } + + fn build_output_set( + &self, + level: crate::builder::InstrumentationLevel, + ) -> (OutputQueueSet, OutputsViewAny) { + // Unlike `TypedStep::build_output_set`, a `Step2` does NOT collapse a + // `ByOrdinal` / `ByItemOrdinal` output to `None` for `Serial` / + // `Exclusive` kinds — the ordering is passed through verbatim and the + // reorder stage is kept. + // + // The single-input collapse rests on a universal property: a Serial + // single-input step consumes one already-ordered stream, so any input + // ordinal it propagates onto an output is emitted in push order, making + // the reorder stage redundant. A two-input MERGE has no such guarantee. + // It interleaves two branches, so a `ByItemOrdinal` output that + // propagates an *input's* ordinal can be pushed out of ordinal order + // even under serial (mutex-serialized, single-dispatcher) execution — + // e.g. when branch B's next-needed ordinal hasn't arrived yet but + // branch A's later one has. For such a step the reorder stage is + // load-bearing: dropping it would deliver records out of order. + // + // Today's production `Step2`s happen not to need it (`PairRawFastq` + // declares `BranchOrdering::None`; `ZipperMergeStep` assigns fresh + // *sequential* output ordinals, so its push order already equals its + // ordinal order). But the framework cannot assume that for an arbitrary + // merge, so it conservatively keeps the reorder for every ordered Step2 + // output. Do not add the single-input collapse here: a `Step2` that + // propagates an input ordinal can emit out of ordinal order even under + // serial execution, so dropping the reorder would deliver records out of + // order. + // Bind the profile once — each `profile()` call rebuilds its two `Vec`s. + // Matches the single-input adapter and this adapter's own + // `build_fused_output_set`. + let profile = self.inner.profile(); + ::build_queues( + &profile.output_queues, + &profile.branch_ordering, + level, + ) + } + + fn build_fused_output_set( + &self, + _level: crate::builder::InstrumentationLevel, + ) -> (OutputQueueSet, OutputsViewAny) { + // `is_fusible_chain` rejects any step with `input_arity() > 1`, so a + // `Step2` never reaches the fused driver. Collapsing its ordering to + // `None` here would drop the reorder stage that `build_output_set` + // documents as load-bearing for a merge, so refuse instead of + // returning an unordered transport. + panic!( + "build_fused_output_set called on Step2 adapter '{}': fusing a \ + two-input merge would drop its load-bearing reorder stage. \ + Extend the fused driver's ordering handling before allowing this.", + self.name + ); + } + + fn wrap_outputs_view(&self, view: OutputsViewAny) -> Box { + let typed: OutputHandles = OutputHandles::::new(view); + Box::new(typed) + } + + fn mark_outputs_drained(&self, outputs: &(dyn Any + Send + Sync)) { + let typed = outputs + .downcast_ref::>() + .expect("mark_outputs_drained downcast failed — Step2 chain topology invariant"); + ::mark_all_drained(typed); + } + + fn is_source(&self) -> bool { + // Step2 consumers are never sources by definition (they have + // two non-unit input branches). + false + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + use crate::handles::BranchInputHandle; + use crate::outputs::Single; + use crate::queues::QueueSpec; + use crate::reorder::BranchOrdering; + use crate::step::{ + InputHandle, OutputHandles, Step, StepCtx, StepCtx2, StepKind, StepOutcome, StepProfile, + }; + + /// Trivial step: u32 → u32+1 (single output). + #[derive(Clone)] + struct AddOne; + + impl Step for AddOne { + type Input = u32; + type Outputs = Single; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "AddOne", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + match ctx.input.pop() { + Some(n) => { + let _ = ctx.outputs.push(n + 1); + Ok(StepOutcome::Progress) + } + None => Ok(StepOutcome::NoProgress), + } + } + + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + /// Build a `(producer_set, producer_outputs)` pair so we can simulate a + /// chain link by manually constructing the upstream side. + fn build_addone_outputs() -> (OutputQueueSet, OutputHandles>) { + let producer: Box = Box::new(TypedStep::new(AddOne)); + let (queue_set, outputs_view) = + producer.build_output_set(crate::builder::InstrumentationLevel::Off); + let outputs: OutputHandles> = OutputHandles::new(outputs_view); + (queue_set, outputs) + } + + #[rstest] + #[case(StepKind::Serial, false)] + #[case(StepKind::Serial, true)] + #[case(StepKind::Parallel, false)] + #[case(StepKind::Parallel, true)] + #[case(StepKind::Exclusive, false)] + #[case(StepKind::Exclusive, true)] + #[case(StepKind::Detached, false)] + #[case(StepKind::Detached, true)] + fn cached_kind_and_sticky_match_profile(#[case] kind: StepKind, #[case] sticky: bool) { + // The `kind()` / `sticky()` accessors return values cached at + // construction (so the setup passes don't rebuild the profile's `Vec`s). + // They must agree with `profile()` for every kind/sticky combination, + // or a stale cache could mis-route worker assignment. Each combination is + // an independent `#[case]` so a failure isolates the offending pair. + #[derive(Clone)] + struct Tagged(StepKind, bool); + impl Step for Tagged { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "Tagged", + kind: self.0, + sticky: self.1, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + let erased: Box = Box::new(TypedStep::new(Tagged(kind, sticky))); + let profile = erased.profile(); + assert_eq!(erased.kind(), profile.kind, "kind() must match profile for {kind:?}"); + assert_eq!( + erased.sticky(), + profile.sticky, + "sticky() must match profile for {kind:?}/{sticky}" + ); + } + + #[rstest] + #[case(StepKind::Serial, false)] + #[case(StepKind::Serial, true)] + #[case(StepKind::Parallel, false)] + #[case(StepKind::Parallel, true)] + #[case(StepKind::Exclusive, false)] + #[case(StepKind::Exclusive, true)] + #[case(StepKind::Detached, false)] + #[case(StepKind::Detached, true)] + fn cached_kind_and_sticky_match_profile_step2(#[case] kind: StepKind, #[case] sticky: bool) { + // Mirror `cached_kind_and_sticky_match_profile` for the `Step2` adapter. + // `TypedStep2::{new, kind, sticky}` caches the same metadata at + // construction, and runtime setup reads `ErasedStep::kind()`/`sticky()` + // for ALL erased steps — `Step` and `Step2` alike (e.g. zipper's merge + // step). A stale `TypedStep2` cache could mis-route owner assignment for + // a merge step just as a `TypedStep` cache could, so pin both paths. + #[derive(Clone)] + struct Tagged2(StepKind, bool); + impl Step2 for Tagged2 { + type InputA = u32; + type InputB = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "Tagged2", + kind: self.0, + sticky: self.1, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx2<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + let erased: Box = Box::new(TypedStep2::new(Tagged2(kind, sticky))); + let profile = erased.profile(); + assert_eq!(erased.kind(), profile.kind, "kind() must match profile for {kind:?}"); + assert_eq!( + erased.sticky(), + profile.sticky, + "sticky() must match profile for {kind:?}/{sticky}" + ); + } + + #[test] + fn typed_step_round_trips_through_erased_dispatch() { + // Producer's output set carries the input handle that the consumer + // (also an AddOne) will pull from. The consumer's own output set is + // separate. + let (mut producer_set, producer_outputs) = build_addone_outputs(); + + // Push a u32 onto the producer's output (which is the consumer's input). + producer_outputs.push(41).unwrap(); + + // Consumer takes the typed input handle out of the producer's set. + let mut consumer: Box = Box::new(TypedStep::new(AddOne)); + let input_any = consumer.build_input_handle(&mut producer_set, 0); + + // Consumer needs its own output set + view to run. + let (consumer_set, consumer_view) = + consumer.build_output_set(crate::builder::InstrumentationLevel::Off); + let consumer_outputs: OutputHandles> = OutputHandles::new(consumer_view); + + let signal = PipelineSignal::new(); + let mut ctx = ErasedStepCtx { + input: input_any.as_ref(), + outputs: &consumer_outputs as &(dyn Any + Send + Sync), + signal: &signal, + }; + let outcome = consumer.try_run_erased(&mut ctx).unwrap(); + assert_eq!(outcome, StepOutcome::Progress); + + // Consumer pushed (41 + 1) = 42 onto its own output. Pull it. + let mut consumer_set = consumer_set; + let consumer_input = consumer_set.take_typed_input::(0); + assert_eq!(consumer_input.pop(), Some(42)); + } + + /// Every other dispatch test calls `try_run_erased` exactly once, so + /// `resolve_input` / `resolve_outputs` always take the first-resolve path + /// and the cached-hit arm — which holds both `transmute`s and the + /// box-address assert guarding the dispatch-identity invariant — is never + /// executed. Dispatch twice against the same `ErasedStepCtx` so the second + /// call resolves from cache. + #[test] + fn second_dispatch_resolves_handles_from_the_cache() { + let (mut producer_set, producer_outputs) = build_addone_outputs(); + producer_outputs.push(1).unwrap(); + producer_outputs.push(2).unwrap(); + + let mut consumer: Box = Box::new(TypedStep::new(AddOne)); + let input_any = consumer.build_input_handle(&mut producer_set, 0); + let (mut consumer_set, consumer_view) = + consumer.build_output_set(crate::builder::InstrumentationLevel::Off); + let consumer_outputs: OutputHandles> = OutputHandles::new(consumer_view); + + let signal = PipelineSignal::new(); + let mut ctx = ErasedStepCtx { + input: input_any.as_ref(), + outputs: &consumer_outputs as &(dyn Any + Send + Sync), + signal: &signal, + }; + // First dispatch populates the cache; the second must hit it — and both + // its unconditional box-address assert and the debug-only downcast + // re-check, which is live in this test build. + assert_eq!(consumer.try_run_erased(&mut ctx).unwrap(), StepOutcome::Progress); + assert_eq!(consumer.try_run_erased(&mut ctx).unwrap(), StepOutcome::Progress); + + let out = consumer_set.take_typed_input::(0); + assert_eq!( + (out.pop(), out.pop()), + (Some(2), Some(3)), + "both dispatches must produce output through the cached handles" + ); + } + + /// Point 3 of the `TypedStep` safety argument — every dispatch passes the + /// same box for a given `step_idx` — is the one the compiler cannot enforce. + /// A step whose cache is populated and is then handed a *different* input box + /// must panic rather than resolve from the stale entry. The check is an + /// unconditional `assert!`, so this holds in release builds too; when it was + /// a `debug_assert!` a release build read through the pointer into the + /// previous `ChainContexts` instead. + #[test] + #[should_panic(expected = "cached input handle does not match the dispatch context")] + fn dispatch_against_a_different_context_box_panics() { + let (mut first_producer_set, first_producer_outputs) = build_addone_outputs(); + first_producer_outputs.push(1).unwrap(); + let (mut second_producer_set, second_producer_outputs) = build_addone_outputs(); + second_producer_outputs.push(2).unwrap(); + + let mut consumer: Box = Box::new(TypedStep::new(AddOne)); + let first_input = consumer.build_input_handle(&mut first_producer_set, 0); + let second_input = consumer.build_input_handle(&mut second_producer_set, 0); + let (_consumer_set, consumer_view) = + consumer.build_output_set(crate::builder::InstrumentationLevel::Off); + let consumer_outputs: OutputHandles> = OutputHandles::new(consumer_view); + let outputs_any = &consumer_outputs as &(dyn Any + Send + Sync); + + let signal = PipelineSignal::new(); + // First dispatch populates the cache from `first_input`. + let mut ctx = + ErasedStepCtx { input: first_input.as_ref(), outputs: outputs_any, signal: &signal }; + assert_eq!(consumer.try_run_erased(&mut ctx).unwrap(), StepOutcome::Progress); + + // Same step, different input box — the invariant violation. + let mut wrong_ctx = + ErasedStepCtx { input: second_input.as_ref(), outputs: outputs_any, signal: &signal }; + let _ = consumer.try_run_erased(&mut wrong_ctx); + } + + #[test] + fn typed_step_returns_noprogress_on_empty_input() { + let (mut producer_set, _producer_outputs) = build_addone_outputs(); + let mut consumer: Box = Box::new(TypedStep::new(AddOne)); + let input_any = consumer.build_input_handle(&mut producer_set, 0); + + let (_consumer_set, consumer_view) = + consumer.build_output_set(crate::builder::InstrumentationLevel::Off); + let consumer_outputs: OutputHandles> = OutputHandles::new(consumer_view); + + let signal = PipelineSignal::new(); + let mut ctx = ErasedStepCtx { + input: input_any.as_ref(), + outputs: &consumer_outputs as &(dyn Any + Send + Sync), + signal: &signal, + }; + let outcome = consumer.try_run_erased(&mut ctx).unwrap(); + assert_eq!(outcome, StepOutcome::NoProgress); + } + + #[test] + fn clone_boxed_yields_independent_step() { + let original: Box = Box::new(TypedStep::new(AddOne)); + let cloned = original.clone_boxed(); + assert_eq!(original.profile().name, "AddOne"); + assert_eq!(cloned.profile().name, "AddOne"); + } + + /// `ErasedStep::name()` (the per-dispatch, allocation-free accessor) + /// must agree with `profile().name`, including across `clone_boxed` + /// (the per-worker clone path that re-caches the name in `new`). + #[test] + fn erased_name_matches_profile_name() { + let step: Box = Box::new(TypedStep::new(AddOne)); + assert_eq!(step.name(), step.profile().name); + assert_eq!(step.name(), "AddOne"); + let cloned = step.clone_boxed(); + assert_eq!(cloned.name(), cloned.profile().name); + assert_eq!(cloned.name(), "AddOne"); + } + + #[test] + fn build_output_set_uses_profile_queues_and_ordering() { + let typed: Box = Box::new(TypedStep::new(AddOne)); + let (mut queue_set, _outputs_view) = + typed.build_output_set(crate::builder::InstrumentationLevel::Off); + assert_eq!(queue_set.n_branches(), 1); + // Verify the queue is u32-typed by taking the input handle. + let input = queue_set.take_typed_input::(0); + // Empty initially. + assert_eq!(input.pop(), None); + } + + /// The fused transport keeps the profile's queue bound; only the ordering is + /// dropped. It previously forced `QueueSpec::Unbounded` on every branch, so a + /// step emitting more items per `try_run` than its consumer removed grew the + /// edge by the difference on every pass of the fused driver — the driver runs + /// a producer before its consumer, so nothing capped it. `AddOne` declares + /// `CountBounded { capacity: 4 }`, so the fifth push must be rejected. + #[test] + fn build_fused_output_set_keeps_the_profile_queue_bound() { + let typed: Box = Box::new(TypedStep::new(AddOne)); + let (mut queue_set, outputs_view) = + typed.build_fused_output_set(crate::builder::InstrumentationLevel::Off); + assert_eq!(queue_set.n_branches(), 1); + let outputs: OutputHandles> = OutputHandles::new(outputs_view); + + for n in 0..4u32 { + outputs.push(n).expect("the capacity-4 fused transport admits four items"); + } + assert!( + outputs.push(4).is_err(), + "the fused transport must apply the profile's capacity-4 backpressure \ + rather than growing without limit" + ); + + // Backpressure, not a wedge: draining one item frees exactly one slot. + let input = queue_set.take_typed_input::(0); + assert_eq!(input.pop(), Some(0), "direct FIFO transport, no reorder stage"); + outputs.push(4).expect("one pop frees one slot"); + } + + #[test] + fn build_input_handle_downcasts_to_typed_branch_handle() { + // Producer pushes a value onto its output. + let (mut producer_set, producer_outputs) = build_addone_outputs(); + producer_outputs.push(7).unwrap(); + + // Consumer grabs the typed input handle. + let consumer: Box = Box::new(TypedStep::new(AddOne)); + let input_any = consumer.build_input_handle(&mut producer_set, 0); + let input = + input_any.downcast_ref::>().expect("input handle downcast"); + assert_eq!(input.pop(), Some(7)); + } + + #[test] + fn wrap_outputs_view_yields_typed_outputs_handles() { + let typed: Box = Box::new(TypedStep::new(AddOne)); + let (mut queue_set, view) = + typed.build_output_set(crate::builder::InstrumentationLevel::Off); + let outputs_any = typed.wrap_outputs_view(view); + + // Downcast back to OutputHandles> and push. + let outputs = outputs_any + .downcast_ref::>>() + .expect("OutputHandles downcast"); + outputs.push(99).unwrap(); + + let input = queue_set.take_typed_input::(0); + assert_eq!(input.pop(), Some(99)); + } + + #[test] + fn mark_outputs_drained_propagates_to_consumer_input() { + // Producer is an AddOne; its outputs feed a consumer's input. + let producer: Box = Box::new(TypedStep::new(AddOne)); + let (mut producer_set, producer_view) = + producer.build_output_set(crate::builder::InstrumentationLevel::Off); + let producer_outputs_any = producer.wrap_outputs_view(producer_view); + + let consumer: Box = Box::new(TypedStep::new(AddOne)); + let input_any = consumer.build_input_handle(&mut producer_set, 0); + let consumer_input = input_any + .downcast_ref::>() + .expect("consumer input handle is BranchInputHandle"); + + // Initially not drained. + assert!(!InputHandle::is_drained(consumer_input)); + + // Producer marks its outputs drained via the type-erased path. + producer.mark_outputs_drained(producer_outputs_any.as_ref()); + + // Drained signal reaches consumer. + assert!(InputHandle::is_drained(consumer_input)); + } + + /// Multi-output step: u32 → (u32, u32) split into low + high bytes. + #[derive(Clone)] + struct SplitBytes; + + impl Step for SplitBytes { + type Input = u32; + type Outputs = (u32, u32); + + fn profile(&self) -> StepProfile { + StepProfile { + name: "SplitBytes", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![ + QueueSpec::CountBounded { capacity: 2 }, + QueueSpec::CountBounded { capacity: 2 }, + ], + branch_ordering: vec![BranchOrdering::None, BranchOrdering::None], + } + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + match ctx.input.pop() { + Some(n) => { + let v = ctx.outputs.view(); + let _ = v.a.push(n & 0xFF); + let _ = v.b.push((n >> 8) & 0xFF); + Ok(StepOutcome::Progress) + } + None => Ok(StepOutcome::NoProgress), + } + } + } + + #[test] + fn split_bytes_routes_to_two_branches() { + // Producer feeding `SplitBytes`: another AddOne (Single → Single). + // We just need a way to feed u32 into SplitBytes, so reuse AddOne's outputs. + let (mut producer_set, producer_outputs) = build_addone_outputs(); + producer_outputs.push(0xABCD).unwrap(); + + // SplitBytes is the consumer; takes the typed input from the producer. + let mut splitter: Box = Box::new(TypedStep::new(SplitBytes)); + let input_any = splitter.build_input_handle(&mut producer_set, 0); + + // Splitter's own output set has two branches. + let (mut splitter_set, splitter_view) = + splitter.build_output_set(crate::builder::InstrumentationLevel::Off); + let splitter_outputs: OutputHandles<(u32, u32)> = OutputHandles::new(splitter_view); + + let signal = PipelineSignal::new(); + let mut ctx = ErasedStepCtx { + input: input_any.as_ref(), + outputs: &splitter_outputs as &(dyn Any + Send + Sync), + signal: &signal, + }; + + let outcome = splitter.try_run_erased(&mut ctx).unwrap(); + assert_eq!(outcome, StepOutcome::Progress); + + let in_a = splitter_set.take_typed_input::(0); + let in_b = splitter_set.take_typed_input::(1); + assert_eq!(in_a.pop(), Some(0xCD)); + assert_eq!(in_b.pop(), Some(0xAB)); + } + + #[test] + fn is_source_true_for_unit_input_steps() { + #[derive(Clone)] + struct UnitInputStep; + impl Step for UnitInputStep { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "UnitInput", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 1 }], + branch_ordering: vec![BranchOrdering::ByOrdinal], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::Finished) + } + } + + let source: Box = Box::new(TypedStep::new(UnitInputStep)); + let mid: Box = Box::new(TypedStep::new(AddOne)); + assert!(source.is_source()); + assert!(!mid.is_source()); + } + + /// Pins the deliberate asymmetry between `TypedStep::build_output_set` + /// (collapses ordered output → `None` for Serial/Exclusive) and + /// `TypedStep2::build_output_set` (keeps the reorder stage). A two-input + /// merge can push a `ByItemOrdinal` output out of ordinal order even under + /// serial execution, so the reorder stage is load-bearing. This test fails + /// if a future change adds the single-input collapse to the Step2 path. + #[test] + fn step2_serial_byitemordinal_output_is_reordered_not_collapsed() { + use crate::item::{HeapSize, Ordered}; + use crate::outputs::OrderedBytesSingle; + use crate::step::{Step2, StepCtx2}; + + #[derive(Debug)] + struct Ord32 { + ordinal: u64, + } + impl HeapSize for Ord32 { + fn heap_size(&self) -> usize { + 0 + } + } + impl Ordered for Ord32 { + fn ordinal(&self) -> u64 { + self.ordinal + } + } + + // A Serial two-input merge with a `ByItemOrdinal` output. We only need + // its output edge (via `build_output_set`), never run it, so `try_run` + // is trivial. + struct OutOfOrderMerge; + impl Step2 for OutOfOrderMerge { + type InputA = u32; + type InputB = u32; + type Outputs = OrderedBytesSingle; + fn profile(&self) -> StepProfile { + StepProfile { + name: "OutOfOrderMerge", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: 64 * 1024 }], + branch_ordering: vec![BranchOrdering::ByItemOrdinal], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx2<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + let step: Box = Box::new(TypedStep2::new(OutOfOrderMerge)); + let (mut queue_set, view) = + step.build_output_set(crate::builder::InstrumentationLevel::Off); + let outputs_any = step.wrap_outputs_view(view); + let outputs = outputs_any + .downcast_ref::>>() + .expect("OutputHandles downcast"); + + // Push OUT of ordinal order: 2, then 0, then 1. + outputs.push(Ord32 { ordinal: 2 }).unwrap(); + outputs.push(Ord32 { ordinal: 0 }).unwrap(); + outputs.push(Ord32 { ordinal: 1 }).unwrap(); + + // The reorder stage must deliver them in ORDINAL order (0, 1, 2). With + // the single-input collapse the consumer would see push order (2, 0, 1). + let input = queue_set.take_typed_input::(0); + let got: Vec = std::iter::from_fn(|| input.pop().map(|o| o.ordinal)).collect(); + assert_eq!(got, vec![0, 1, 2], "ordered Step2 output must be reordered by ordinal"); + } + + /// Input construction is dispatched on `input_arity()`, so a single-input + /// adapter must never be asked for two-input handles and a `Step2` adapter + /// must never be asked for one. Both defaults panic rather than silently + /// mis-wiring a chain, and both messages must name the step so a real hit + /// is debuggable. These pin the guard, not the happy path. + #[test] + #[should_panic(expected = "AddOne")] + fn build_two_input_handles_on_a_single_input_adapter_panics() { + let step: Box = Box::new(TypedStep::new(AddOne)); + let (mut producer_set, _view) = + step.build_output_set(crate::builder::InstrumentationLevel::Off); + let _ = step.build_two_input_handles(std::slice::from_mut(&mut producer_set), 0, 0, 0, 0); + } + + #[test] + #[should_panic(expected = "build_input_handle called on Step2 adapter")] + fn build_input_handle_on_a_step2_adapter_panics() { + use crate::step::{Step2, StepCtx2}; + + #[derive(Clone)] + struct JoinStub; + impl Step2 for JoinStub { + type InputA = u32; + type InputB = u32; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "JoinStub", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx2<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + let (mut producer_set, _outputs) = build_addone_outputs(); + let step: Box = Box::new(TypedStep2::new(JoinStub)); + let _ = step.build_input_handle(&mut producer_set, 0); + } +} diff --git a/crates/fgumi-pipeline-core/src/finalize.rs b/crates/fgumi-pipeline-core/src/finalize.rs new file mode 100644 index 000000000..4170d9333 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/finalize.rs @@ -0,0 +1,39 @@ +//! Post-pipeline finalize hook contract. +//! +//! [`FinalizeHook`] is the bare trait that lets a chain register +//! heterogeneous post-pipeline cleanup actions — metrics drains, summary +//! logging, gate checks, rejects-file finalization, BAI indexing — in a single +//! `Vec>` that the caller drains after `Pipeline::run` +//! returns. +//! +//! It lives in `fgumi-pipeline-core` (rather than the umbrella `fgumi` crate) +//! so that any CLI crate built directly on the pipeline core — e.g. +//! `fgumi-sort-cli` — can implement the same contract without redefining a +//! parallel trait. Keeping the trait single-sourced here is what lets the sort +//! finalize hooks be defined once and re-exported by the umbrella, instead of +//! duplicated across both crates (X1-005). +//! +//! The richer machinery built on top of this trait — `BuiltPipeline`, +//! `drain_finalize`, and the built-in stats/timing hooks — stays in the +//! umbrella crate, since it references umbrella-only types. + +use anyhow::Result; + +/// Post-pipeline cleanup. Each stage with metrics, summary logging, +/// `--min-corrected`-style gates, or rejects-file finalization registers one or +/// more hooks during chain build. The caller iterates the registered hooks +/// (typically via `try_for_each`) after `Pipeline::run` returns. +/// +/// Trait object so heterogeneous hooks (correct's metrics drain, consensus's +/// summary log, AAM's records-aligned counter, sort's BAI indexer, etc.) can +/// share a `Vec>`. +pub trait FinalizeHook: Send { + /// Run the post-pipeline action. Called exactly once, after + /// `Pipeline::run` returns and before the caller exits. + /// + /// # Errors + /// + /// Returns the underlying I/O / aggregation / gate-check error. + /// Callers typically chain via `try_for_each`. + fn finalize(self: Box) -> Result<()>; +} diff --git a/crates/fgumi-pipeline-core/src/handles.rs b/crates/fgumi-pipeline-core/src/handles.rs new file mode 100644 index 000000000..1caa59c24 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/handles.rs @@ -0,0 +1,2260 @@ +//! Framework-side input/output handles wired to the layered queue stack. +//! +//! Per-branch architecture: +//! +//! ```text +//! producer step +//! │ outputs.push(item) +//! ▼ +//! BranchOutputHandle +//! │ - allocates ordinal (if BranchOrdering::ByOrdinal) +//! │ - heap-size accounting happens inside ByteBoundedQueue +//! │ - forwards to either ReorderStage or transport directly +//! ▼ +//! ┌────────────────────────────────┐ +//! │ ReorderStage (optional) │ ← only if BranchOrdering::ByOrdinal +//! └────────────────────────────────┘ +//! ▼ +//! ┌──────────────────────────────────────┐ +//! │ Arc> │ ← Wrapped = Sequenced if ordered, else T +//! └──────────────────────────────────────┘ +//! ▼ +//! BranchInputHandle +//! │ pop() -> Option +//! ▼ +//! consumer step +//! ``` +//! +//! No code in this module surfaces serials or heap sizes to step authors — +//! those are framework-managed. Step authors see `outputs.push(item)` and +//! `input.pop() -> Option`. +//! +//! `QueueSpec::ByteBounded` requires `T: HeapSize`. Every build path — `Single`, +//! the tuple fan-outs, and the ordered-bytes shapes — bounds `T: HeapSize` and +//! routes through `build_branch_byte_aware` (or `build_branch_ordered_bytes`), so +//! `ByteBounded` is honored on every branch of every shape. Those byte-aware +//! entry points delegate every non-byte spec straight back to `build_branch`, so +//! count/unbounded branches are unaffected. +//! +//! `build_branch` itself keeps a `ByteBounded` panic as the guard on that +//! delegation: reaching it with a byte spec means a caller bypassed a byte-aware +//! entry point. The other two panics are for an item-carried ordinal without the +//! trait to read it — `ByItemOrdinal` declared without `Ordered`, and +//! `ByteBounded` + `ByItemOrdinal` on a path lacking the `Ordered` bound (which +//! must use `build_branch_ordered_bytes`). + +use std::any::Any; +use std::marker::PhantomData; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering}; + +use super::item::{HeapSize, Ordered}; +use super::outputs::Single; +use super::queues::{ByteBoundedQueue, CountBoundedQueue, ItemQueue, QueueSpec, UnboundedQueue}; +use super::reorder::{ + BranchOrdering, DEFAULT_REORDER_OVERFLOW_BYTES, ReorderCapHandle, ReorderStage, Sequenced, +}; +use super::step::{InputHandle, OutputHandles, OutputsViewAny}; + +// ───────────────────────────────────────────────────────────────────────────── +// BranchOutputHandle — one per output branch on the producer side. +// ───────────────────────────────────────────────────────────────────────────── + +/// Producer-side handle for one output branch. Step authors call `push(item)`. +/// The framework manages ordinal assignment and routes to the right destination +/// (transport or reorder stage). +pub struct BranchOutputHandle { + inner: BranchOutputInner, +} + +enum BranchOutputInner { + /// `BranchOrdering::None`: push directly to transport. + Direct(Arc>), + /// `BranchOrdering::ByOrdinal` or `ByItemOrdinal`: push through + /// `ReorderStage`. The stage internally wraps in `Sequenced` and + /// forwards to its transport. `OrdinalSource` decides where the + /// ordinal comes from on each push. + Ordered { stage: Arc>, ordinal_source: OrdinalSource }, +} + +/// Where each push's ordinal comes from on an `Ordered` branch. +enum OrdinalSource { + /// `BranchOrdering::ByOrdinal`: producer allocates a fresh ordinal via + /// this counter on every push. + Allocated(Arc), + /// `BranchOrdering::ByItemOrdinal`: items carry their own ordinal. The + /// function pointer is `|item: &T| item.ordinal()`, monomorphized at + /// branch construction time (requires `T: Ordered`). + ItemSerial(fn(&T) -> u64), +} + +impl OrdinalSource { + #[inline] + fn next(&self, item: &T) -> u64 { + match self { + Self::Allocated(c) => c.fetch_add(1, AtomicOrdering::AcqRel), + Self::ItemSerial(f) => f(item), + } + } +} + +/// An item the underlying queue rejected due to backpressure. +/// +/// `Unpushed` carries the rejected item back to the producer for retry, +/// **plus an opaque token** that — for ordered (`ByOrdinal`) branches — +/// preserves the ordinal allocated on the original push. Retrying via +/// [`BranchOutputHandle::retry`] reuses that ordinal so the consumer's +/// `ReorderStage` doesn't stall waiting for a sequence number that was +/// burned on a rejected attempt. +/// +/// Step authors store `HeldSlot>` (not `HeldSlot`) when their +/// step pushes to an ordered branch. Direct branches don't carry an ordinal, +/// so `Unpushed` is just a typed wrapper around the item there. +pub struct Unpushed { + item: T, + /// Pre-allocated ordinal from a rejected `Ordered` push. `None` for + /// `Direct` (FIFO) branches. + ordinal: Option, +} + +/// Outcome of [`OutputHandles::retry_held`] — a step's drain-time held-slot +/// retry. Lets the caller distinguish "nothing was held" from "flushed +/// successfully" (the two differ for a single-final-batch flusher: the former +/// must still build its batch, the latter is done). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum HeldRetry { + /// The held slot was empty — nothing to retry. + WasEmpty, + /// A held item was retried and accepted; the slot is now empty. + Flushed, + /// A held item was retried, still rejected by backpressure, and put back + /// in the slot. The caller must yield (return `NoProgress`/`Contention` + /// and retry on the next dispatch). + StillHeld, +} + +impl std::fmt::Debug for Unpushed { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Unpushed").field("ordinal", &self.ordinal).field("item", &"<...>").finish() + } +} + +impl Unpushed { + /// Borrow the rejected item (read-only). + pub fn item(&self) -> &T { + &self.item + } + + /// Discard the framework token and recover the item, abandoning the + /// pre-allocated ordinal (if any). After calling this, **do not** push + /// the item back into an ordered branch — the missing ordinal will + /// stall the consumer's `ReorderStage`. Use only when abandoning the + /// item entirely (e.g., during cancellation). + pub fn into_item(self) -> T { + self.item + } +} + +impl BranchOutputHandle { + /// Try to push a fresh item. + /// + /// # Errors + /// + /// Returns `Err(Unpushed)` when the underlying queue (or reorder + /// stage, for `ByOrdinal` branches) applied backpressure. The producer + /// must hand the `Unpushed` to a subsequent [`Self::retry`] call — + /// **not** call `push` with the recovered item — so any pre-allocated + /// ordinal is preserved. + pub fn push(&self, item: T) -> Result<(), Unpushed> { + match &self.inner { + BranchOutputInner::Direct(q) => { + q.try_push(item).map_err(|item| Unpushed { item, ordinal: None }) + } + BranchOutputInner::Ordered { stage, ordinal_source } => { + let ord = ordinal_source.next(&item); + stage + .try_push(ord, item) + .map_err(|(ord, item)| Unpushed { item, ordinal: Some(ord) }) + } + } + } + + /// Retry a previously-rejected push. Reuses the original ordinal for + /// ordered branches so the consumer's `ReorderStage` sees a contiguous + /// sequence. + /// + /// # Errors + /// + /// Returns `Err(Unpushed)` if the queue still applied backpressure; + /// the caller should re-hold and retry on the next worker iteration. + /// + /// # Panics + /// + /// Panics if the `Unpushed` came from a different branch kind than this + /// handle (a framework invariant violation). The `Ordered` + `ordinal = + /// None` case panics in **every** build: there is no recovery that + /// preserves the sequence (see the arm's comment). The `Direct` + + /// `ordinal = Some` case is recoverable — a `Direct` branch has no reorder + /// stage to desynchronize — so it panics only in debug builds and drops + /// the stray ordinal in release. + pub fn retry(&self, unpushed: Unpushed) -> Result<(), Unpushed> { + let Unpushed { item, ordinal } = unpushed; + match (&self.inner, ordinal) { + (BranchOutputInner::Direct(q), None) => { + q.try_push(item).map_err(|item| Unpushed { item, ordinal: None }) + } + (BranchOutputInner::Ordered { stage, .. }, Some(ord)) => stage + .try_push(ord, item) + .map_err(|(ord, item)| Unpushed { item, ordinal: Some(ord) }), + (BranchOutputInner::Direct(q), Some(_)) => { + debug_assert!(false, "Unpushed::ordinal=Some on a Direct branch"); + q.try_push(item).map_err(|item| Unpushed { item, ordinal: None }) + } + (BranchOutputInner::Ordered { .. }, None) => { + // Unrecoverable, so fail loudly in every build. `self.push` + // would allocate a FRESH ordinal and abandon the one this + // `Unpushed` was created with, leaving a permanent hole in the + // sequence — the consumer's `ReorderStage` then waits on the + // missing ordinal forever and only the deadlock monitor notices. + // A panic here names the framework bug instead. This is the + // hazard `Unpushed::into_item` documents. + panic!( + "Unpushed::ordinal=None on an Ordered branch — re-pushing would \ + abandon the allocated ordinal and stall the consumer's reorder stage" + ); + } + } + } + + /// Mark this branch drained (producer-side). Idempotent. + pub fn mark_drained(&self) { + match &self.inner { + BranchOutputInner::Direct(q) => q.mark_drained(), + BranchOutputInner::Ordered { stage, .. } => stage.mark_drained(), + } + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// BranchInputHandle — implements InputHandle. +// ───────────────────────────────────────────────────────────────────────────── + +/// Consumer-side input handle for one branch. Implements `InputHandle`. +pub struct BranchInputHandle { + inner: BranchInputInner, + /// `Some` only on an instrumented edge. The **consumer-pop** side of the + /// edge's `EdgeMetrics` is recorded here (not at the transport queue), + /// because an ordered edge pops through a `ReorderStage` that drains the + /// transport in bulk — `pop` here is the real consumer pop. Shares the same + /// `Arc` as the producer transport (push side). + metrics: Option>, + /// Whether to report each popped item's `heap_size()` to `record_pop`. + /// + /// True only on a byte-bounded edge, mirroring the push side: a + /// `CountBoundedQueue` / `UnboundedQueue` records `record_push(0)` because it + /// keeps no byte accounting, and `ordered_branch` gates its push-side byte + /// reporting on the same signal (`transport_handle.is_some()`). Recording + /// real bytes on the pop side of such an edge made `popped_bytes` nonzero + /// while `pushed_bytes` stayed 0, so `compute_edge_stats` reported a + /// `mibytes_per_s` throughput for an edge that measures no bytes at all. + record_item_bytes: bool, +} + +enum BranchInputInner { + Direct(Arc>), + Ordered(Arc>), + /// A permanently-empty, permanently-drained handle that owns no transport. + /// Used for source steps' dummy input (their input is implicitly drained + /// from t=0; the worker loop never pops from it), avoiding a per-source + /// `SegQueue` allocation just to report `is_drained() == true`. + AlwaysDrained(PhantomData T>), +} + +impl BranchInputHandle { + /// Construct a zero-state input handle that is always empty and always + /// drained, owning no backing queue. Used for source steps, whose input is + /// implicitly drained from the start and never popped. + #[must_use] + pub fn always_drained() -> Self { + Self { + inner: BranchInputInner::AlwaysDrained(PhantomData), + metrics: None, + record_item_bytes: false, + } + } +} + +impl InputHandle for BranchInputHandle { + fn pop(&self) -> Option { + // A `None` from an `Ordered` edge is *reorder-blocked* (not starved) while + // the reorder stage still holds out-of-order items waiting for an earlier + // ordinal — counting that as an empty pop inflates `pop_empties` and can + // misclassify a backlogged edge as starved. The `reorder_blocked` flag is + // reported by the SAME locked pop, so shared (`Parallel`) consumers can't + // observe a torn `(item, blocked)` pair. Direct / always-drained branches + // have no reorder buffer, so a `None` there is always a true empty pop. + let (item, reorder_blocked) = match &self.inner { + BranchInputInner::Direct(q) => (q.try_pop(), false), + BranchInputInner::Ordered(stage) => stage.try_pop_in_order_reporting_blocked(), + BranchInputInner::AlwaysDrained(_) => (None, false), + }; + if let Some(m) = &self.metrics { + if let Some(it) = &item { + // Bytes only on a byte-bounded edge — see `record_item_bytes`. + // usize→u64 is lossless on every target we build for. + #[allow(clippy::cast_possible_truncation)] + let bytes = if self.record_item_bytes { it.heap_size() as u64 } else { 0 }; + m.record_pop(bytes); + } else if !reorder_blocked { + m.record_empty(); + } + } + item + } + + fn is_drained(&self) -> bool { + match &self.inner { + BranchInputInner::Direct(q) => q.is_drained() && q.is_empty(), + BranchInputInner::Ordered(stage) => stage.is_drained(), + BranchInputInner::AlwaysDrained(_) => true, + } + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Branch construction (one branch at a time, dispatching on QueueSpec + Ordering). +// ───────────────────────────────────────────────────────────────────────────── + +/// The budget-resize handles for one byte-bounded branch: the transport +/// queue's limit setter, plus (for an ordered branch) the reorder stage's +/// overflow-cap setter. The framework's budget pass +/// (`apply_initial_queue_budget`) sets both — the transport limit and the +/// reorder cap — from the same per-edge budget, so they stay in lockstep +/// (single source of truth for `per_queue`). +pub(crate) struct BranchBudgetHandles { + pub(crate) transport: Arc, + /// `Some` for an ordered (`ByOrdinal` / `ByItemOrdinal`) byte-bounded + /// branch — its reorder overflow stash is sized from the same budget. + /// `None` for a direct (unordered) byte-bounded branch (no reorder stage). + pub(crate) reorder_cap: Option>, +} + +/// One end-to-end branch: an output handle and an input handle wired +/// to the same underlying transport (and optional reorder stage). +/// +/// `bounded_queue_handle` is `Some` iff the branch's transport is a +/// `ByteBoundedQueue`; it bundles the transport-limit and reorder-cap +/// setters (see [`BranchBudgetHandles`]). The pipeline-builder collects +/// these (across every branch in the chain) into a registry that the +/// budget pass + optional queue-memory rebalancer use to set/redistribute +/// budget. `None` for `CountBounded` / `Unbounded` branches. +pub(crate) struct Branch { + pub(crate) output: BranchOutputHandle, + pub(crate) input: BranchInputHandle, + pub(crate) bounded_queue_handle: Option, + /// `Some` on an instrumented edge — the shared `EdgeMetrics` (also held by + /// the transport queue for push counts and the input handle for pop counts). + /// Collected into the edge registry by `contexts.rs` Pass 1.5. + pub(crate) metrics: Option>, +} + +/// Mint a per-edge [`crate::runtime::metrics::EdgeMetrics`] when instrumentation is on, else `None`. +/// Called once per branch in the `build_branch*` constructors; the same handle +/// is shared by the transport (push counts) and the input handle (pop counts). +fn edge_metrics( + level: crate::builder::InstrumentationLevel, +) -> Option> { + level.is_on().then(crate::runtime::metrics::EdgeMetrics::new) +} + +/// Build one branch from a queue spec + ordering directive, where `T` does +/// not need to impl `HeapSize` or `Ordered`. +/// +/// # Panics +/// +/// Panics on `QueueSpec::ByteBounded` (requires `T: HeapSize`; use the +/// byte-aware build path) or on `BranchOrdering::ByItemOrdinal` (requires +/// `T: Ordered`; use `build_branch_ordered`). +pub(crate) fn build_branch( + spec: QueueSpec, + ordering: BranchOrdering, + level: crate::builder::InstrumentationLevel, +) -> Branch { + match (spec, ordering) { + (QueueSpec::CountBounded { capacity }, BranchOrdering::None) => { + let m = edge_metrics(level); + let q: Arc> = + Arc::new(CountBoundedQueue::::maybe_instrumented(capacity, m.clone())); + direct_branch(q, None, m) + } + (QueueSpec::CountBounded { capacity }, BranchOrdering::ByOrdinal) => { + let m = edge_metrics(level); + let transport: Arc>> = + Arc::new(CountBoundedQueue::>::maybe_instrumented( + // Ordered transport is NOT instrumented: push/reject are recorded + // at the ReorderStage boundary (a stash turns a full-transport + // `Err` into an accepted `Ok`). Pop side is on the input handle. + capacity, None, + )); + ordered_branch( + transport, + OrdinalSource::Allocated(Arc::new(AtomicU64::new(0))), + DEFAULT_REORDER_OVERFLOW_BYTES, + None, + m, + ) + } + (QueueSpec::Unbounded, BranchOrdering::None) => { + let m = edge_metrics(level); + let q: Arc> = + Arc::new(UnboundedQueue::::maybe_instrumented(m.clone())); + direct_branch(q, None, m) + } + (QueueSpec::Unbounded, BranchOrdering::ByOrdinal) => { + let m = edge_metrics(level); + let transport: Arc>> = + // Ordered transport is NOT instrumented — push/reject recorded at + // the ReorderStage boundary (see the count-bounded note above). + Arc::new(UnboundedQueue::>::maybe_instrumented(None)); + ordered_branch( + transport, + OrdinalSource::Allocated(Arc::new(AtomicU64::new(0))), + DEFAULT_REORDER_OVERFLOW_BYTES, + None, + m, + ) + } + (_, BranchOrdering::ByItemOrdinal) => { + panic!( + "BranchOrdering::ByItemOrdinal requires `T: Ordered` — \ + use `build_branch_ordered::` instead. The plain \ + `build_branch::` path doesn't have the trait bound to \ + read `item.ordinal()` at push time." + ); + } + (QueueSpec::ByteBounded { .. }, _) => { + panic!( + "QueueSpec::ByteBounded requires `T: HeapSize` — \ + use the byte-aware build path (`build_branch_byte_aware` or \ + `build_branch_ordered_bytes`)." + ); + } + } +} + +/// Build one branch where `T: Ordered` (no `HeapSize` requirement). +/// Supports `BranchOrdering::ByItemOrdinal` (uses `item.ordinal()` for the +/// reorder stage's serial). Falls through to `build_branch::` for non- +/// `ByItemOrdinal` orderings. +/// +/// # Panics +/// +/// Panics on `QueueSpec::ByteBounded` (requires `T: HeapSize`; use the +/// `_ordered_bytes` build path). +pub(crate) fn build_branch_ordered( + spec: QueueSpec, + ordering: BranchOrdering, + level: crate::builder::InstrumentationLevel, +) -> Branch { + match (spec, ordering) { + (QueueSpec::CountBounded { capacity }, BranchOrdering::ByItemOrdinal) => { + let m = edge_metrics(level); + let transport: Arc>> = + Arc::new(CountBoundedQueue::>::maybe_instrumented( + // Ordered transport is NOT instrumented: push/reject are recorded + // at the ReorderStage boundary (a stash turns a full-transport + // `Err` into an accepted `Ok`). Pop side is on the input handle. + capacity, None, + )); + ordered_branch( + transport, + OrdinalSource::ItemSerial(|item: &T| item.ordinal()), + DEFAULT_REORDER_OVERFLOW_BYTES, + None, + m, + ) + } + (QueueSpec::Unbounded, BranchOrdering::ByItemOrdinal) => { + let m = edge_metrics(level); + let transport: Arc>> = + // Ordered transport is NOT instrumented — push/reject recorded at + // the ReorderStage boundary (see the count-bounded note above). + Arc::new(UnboundedQueue::>::maybe_instrumented(None)); + ordered_branch( + transport, + OrdinalSource::ItemSerial(|item: &T| item.ordinal()), + DEFAULT_REORDER_OVERFLOW_BYTES, + None, + m, + ) + } + (QueueSpec::ByteBounded { .. }, _) => { + panic!( + "QueueSpec::ByteBounded requires `T: HeapSize` — \ + `build_branch_ordered` only bounds `T: Ordered`. Use \ + `build_branch_ordered_bytes::` for the \ + combined case." + ); + } + (other, ord) => build_branch::(other, ord, level), + } +} + +/// Build one branch where `T: HeapSize + Ordered` (the canonical BAM step +/// case). Supports all `QueueSpec` × `BranchOrdering` combinations. +pub(crate) fn build_branch_ordered_bytes( + spec: QueueSpec, + ordering: BranchOrdering, + level: crate::builder::InstrumentationLevel, +) -> Branch { + use crate::queues::BoundedQueueHandle; + + match (spec, ordering) { + (QueueSpec::ByteBounded { limit_bytes }, BranchOrdering::None) => { + let m = edge_metrics(level); + let q = Arc::new(ByteBoundedQueue::::maybe_instrumented(limit_bytes, m.clone())); + let handle: Arc = Arc::clone(&q) as Arc; + let q_dyn: Arc> = q; + direct_branch(q_dyn, Some(handle), m) + } + (QueueSpec::ByteBounded { limit_bytes }, BranchOrdering::ByOrdinal) => { + let m = edge_metrics(level); + let transport_concrete = Arc::new( + // Ordered transport is NOT instrumented — see the note on the + // count-bounded ordered transport above; push/reject are recorded + // at the ReorderStage boundary. The byte `handle` (occupancy / + // budget resize) is derived from this same queue and is unaffected. + ByteBoundedQueue::>::maybe_instrumented(limit_bytes, None), + ); + let handle: Arc = + Arc::clone(&transport_concrete) as Arc; + let transport: Arc>> = transport_concrete; + // Cap the must-accept overflow buffer so one worker grinding on a + // large ordinal can't let every later ordinal overflow unbounded + // (#29). This `DEFAULT_REORDER_OVERFLOW_BYTES` is the no-budget + // FALLBACK: when `queue_memory_total` is set (production chains), + // `apply_initial_queue_budget` re-sizes this cap thread-awarely + // from the per-edge transport budget via the registered + // `ReorderCapHandle` — so at low thread counts the stash is small + // and at high thread counts it keeps this 256 MiB ceiling. + ordered_branch( + transport, + OrdinalSource::Allocated(Arc::new(AtomicU64::new(0))), + DEFAULT_REORDER_OVERFLOW_BYTES, + Some(handle), + m, + ) + } + (QueueSpec::ByteBounded { limit_bytes }, BranchOrdering::ByItemOrdinal) => { + let m = edge_metrics(level); + let transport_concrete = Arc::new( + // Ordered transport is NOT instrumented — see the note on the + // count-bounded ordered transport above; push/reject are recorded + // at the ReorderStage boundary. The byte `handle` (occupancy / + // budget resize) is derived from this same queue and is unaffected. + ByteBoundedQueue::>::maybe_instrumented(limit_bytes, None), + ); + let handle: Arc = + Arc::clone(&transport_concrete) as Arc; + let transport: Arc>> = transport_concrete; + ordered_branch( + transport, + OrdinalSource::ItemSerial(|item: &T| item.ordinal()), + DEFAULT_REORDER_OVERFLOW_BYTES, + Some(handle), + m, + ) + } + (other, ord) => build_branch_ordered::(other, ord, level), + } +} + +/// Build one branch where `T: HeapSize`. Supports all three queue specs +/// for non-`ByItemOrdinal` orderings. `ByItemOrdinal` requires `T: Ordered` +/// — use `build_branch_ordered_bytes` for the combined case. +/// +/// The build path `build_single_queues` uses for every `Single` output +/// (`T: HeapSize`, no `Ordered` bound): it honors `ByteBounded` for +/// `BranchOrdering::None` and `ByOrdinal`, and delegates every other spec +/// straight back to `build_branch::`. Steps needing both `HeapSize` and +/// `Ordered` go through `build_branch_ordered_bytes` instead. +/// +/// # Panics +/// +/// Panics on `ByteBounded` + `ByItemOrdinal`: reading an item-carried serial +/// needs `T: Ordered`, which this signature does not bound — so the ordinal is +/// unreachable here even when the item type happens to implement `Ordered`. A +/// step needing both declares `OrderedBytesSingle`, which routes to +/// `build_branch_ordered_bytes`; reaching this arm means a step asked for +/// `ByItemOrdinal` through `Single`, whose build path cannot read one. The +/// panic is that guard rather than a supported configuration. +pub(crate) fn build_branch_byte_aware( + spec: QueueSpec, + ordering: BranchOrdering, + level: crate::builder::InstrumentationLevel, +) -> Branch { + use crate::queues::BoundedQueueHandle; + + match (spec, ordering) { + (QueueSpec::ByteBounded { limit_bytes }, BranchOrdering::None) => { + let m = edge_metrics(level); + let q = Arc::new(ByteBoundedQueue::::maybe_instrumented(limit_bytes, m.clone())); + let handle: Arc = Arc::clone(&q) as Arc; + let q_dyn: Arc> = q; + direct_branch(q_dyn, Some(handle), m) + } + (QueueSpec::ByteBounded { limit_bytes }, BranchOrdering::ByOrdinal) => { + let m = edge_metrics(level); + let transport_concrete = Arc::new( + // Ordered transport is NOT instrumented — see the note on the + // count-bounded ordered transport above; push/reject are recorded + // at the ReorderStage boundary. The byte `handle` (occupancy / + // budget resize) is derived from this same queue and is unaffected. + ByteBoundedQueue::>::maybe_instrumented(limit_bytes, None), + ); + let handle: Arc = + Arc::clone(&transport_concrete) as Arc; + let transport: Arc>> = transport_concrete; + ordered_branch( + transport, + OrdinalSource::Allocated(Arc::new(AtomicU64::new(0))), + DEFAULT_REORDER_OVERFLOW_BYTES, + Some(handle), + m, + ) + } + (QueueSpec::ByteBounded { .. }, BranchOrdering::ByItemOrdinal) => { + panic!( + "ByteBounded + ByItemOrdinal requires `T: HeapSize + Ordered` — \ + use `build_branch_ordered_bytes` instead." + ); + } + (other, ord) => build_branch::(other, ord, level), + } +} + +fn direct_branch( + q: Arc>, + transport_handle: Option>, + metrics: Option>, +) -> Branch { + // A byte-bounded direct edge carries a `transport_handle`; count/unbounded + // ones don't. Same signal `ordered_branch` uses to gate its push-side byte + // reporting — the pop side must agree, or the edge reports popped bytes it + // never counted as pushed. + let record_item_bytes = transport_handle.is_some(); + // A direct (unordered) branch has no reorder stage, so no reorder cap. + let bounded_queue_handle = + transport_handle.map(|transport| BranchBudgetHandles { transport, reorder_cap: None }); + Branch { + output: BranchOutputHandle { inner: BranchOutputInner::Direct(Arc::clone(&q)) }, + input: BranchInputHandle { + inner: BranchInputInner::Direct(q), + metrics: metrics.clone(), + record_item_bytes, + }, + bounded_queue_handle, + metrics, + } +} + +/// Build one ordered branch: a `ReorderStage` over `transport`, plus the input / +/// output handles that share it. +/// +/// `max_overflow_bytes` is mandatory rather than optional. The stage's +/// must-accept overflow stash bypasses the transport's own bound by design, so an +/// ordered branch without a cap would let memory grow with *input* size instead +/// of with config. Taking it by value also keeps `ReorderStage::new`'s unbounded +/// default — which exists for tests — unreachable from any branch builder, so +/// there is no path here that produces an uncapped stash. +fn ordered_branch( + transport: Arc>>, + ordinal_source: OrdinalSource, + max_overflow_bytes: u64, + transport_handle: Option>, + metrics: Option>, +) -> Branch { + // A byte-bounded ordered edge carries a `transport_handle`; count/unbounded + // ones don't. Record the push side's byte size only on the byte-bounded ones + // (matching the queue's own byte accounting). Push/reject are recorded at the + // ReorderStage boundary — NOT on the transport — because a must-accept stash + // turns a full-transport `Err` into an accepted `Ok` (see `push_metrics`). + let record_item_bytes = transport_handle.is_some(); + let stage = Arc::new( + ReorderStage::::with_max_overflow_bytes(transport, max_overflow_bytes) + .with_push_metrics(metrics.clone(), record_item_bytes), + ); + // Pair the transport-limit setter with this stage's overflow-cap setter so + // the budget pass sizes both from one per-edge budget. Cloned (and coerced + // to the trait object) before the stage is moved into the input handle. + let reorder_cap: Arc = stage.clone(); + let bounded_queue_handle = transport_handle + .map(|transport| BranchBudgetHandles { transport, reorder_cap: Some(reorder_cap) }); + Branch { + output: BranchOutputHandle { + inner: BranchOutputInner::Ordered { stage: Arc::clone(&stage), ordinal_source }, + }, + input: BranchInputHandle { + inner: BranchInputInner::Ordered(stage), + metrics: metrics.clone(), + record_item_bytes, + }, + bounded_queue_handle, + metrics, + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// OutputQueueSet — per-step collection of typed input handles +// ───────────────────────────────────────────────────────────────────────────── + +/// Type-erased per-branch input handles, owned by the producer step's chain +/// entry until the consumer claims them. The framework moves one entry from +/// here into each consumer's `TypedStep` at chain-build time via +/// `take_typed_input`. +pub struct OutputQueueSet { + pub(crate) branches: Vec, +} + +pub(crate) struct BranchEntry { + /// `Box` carrying a `BranchInputHandle` for the branch's `T`. + /// Set to `Box::new(())` after `take_typed_input`. + pub(crate) input_handle: Box, + /// `Some` iff this branch's transport is a `ByteBoundedQueue`; bundles + /// the transport-limit setter and (for an ordered branch) the reorder + /// overflow-cap setter (see [`BranchBudgetHandles`]). The pipeline + /// builder collects these into a registry so the budget pass + + /// optional queue-memory rebalancer can set/reallocate budget. + pub(crate) bounded_queue_handle: Option, + /// `Some` iff this edge is instrumented (`--pipeline-trace`); the shared + /// `EdgeMetrics` (also held by the transport for push counts and the input + /// handle for pop counts). Collected into the edge registry by `contexts.rs` + /// Pass 1.5. Cleared (`None`) by `take_typed_input`'s placeholder. + pub(crate) metrics: Option>, +} + +impl OutputQueueSet { + pub(crate) fn new(branches: Vec) -> Self { + Self { branches } + } + + #[must_use] + pub fn n_branches(&self) -> usize { + self.branches.len() + } + + /// Take ownership of branch `i`'s typed input handle. Called by the + /// consumer step's chain-build code once. + /// + /// # Panics + /// + /// Panics if `branch_idx >= n_branches()`, if the branch was already + /// taken, or if the requested type `T` doesn't match the producer's + /// declared branch type (build-time invariant enforced by `Chain::chain`). + pub fn take_typed_input( + &mut self, + branch_idx: usize, + ) -> BranchInputHandle { + let entry = std::mem::replace( + &mut self.branches[branch_idx], + BranchEntry { input_handle: Box::new(()), bounded_queue_handle: None, metrics: None }, + ); + let handle: Box> = + entry.input_handle.downcast::>().unwrap_or_else(|_| { + panic!("OutputQueueSet branch type mismatch at index {branch_idx}") + }); + *handle + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Two-input handle wrapper — held inside ChainContexts.inputs[step_idx] for +// every `Step2` consumer. The `TypedStep2` adapter downcasts the +// type-erased box back to `&TwoInputHandles` per +// dispatch and exposes the per-branch refs as `ctx.a` / `ctx.b`. +// +// The wrapper itself is plain owned data; `ChainContexts` builds it once +// at chain-build time by pulling each branch's `BranchInputHandle` out +// of the matching upstream's `OutputQueueSet`. +// ───────────────────────────────────────────────────────────────────────────── + +/// Two per-branch input handles for a [`crate::step::Step2`] +/// consumer, paired by branch slot (`a` = consumer's input slot 0, +/// `b` = slot 1). Boxed type-erased into `ChainContexts.inputs[step_idx]` +/// at chain-build time; the +/// [`crate::erased::TypedStep2`] adapter +/// downcasts and lends per-branch references into [`crate::step::StepCtx2`]. +pub struct TwoInputHandles +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, +{ + pub(crate) a: BranchInputHandle, + pub(crate) b: BranchInputHandle, +} + +impl TwoInputHandles +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, +{ + /// Construct from two per-branch handles. Called by the + /// chain-context builder. + pub(crate) fn new(a: BranchInputHandle, b: BranchInputHandle) -> Self { + Self { a, b } + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Per-arity typed views — held inside OutputsViewAny.inner +// ───────────────────────────────────────────────────────────────────────────── + +/// View for `Single` outputs. Internal: stored inside `OutputsViewAny` +/// and accessed via `OutputHandles>::push` / `retry`. +pub(crate) struct SingleOutputsView { + pub(crate) primary: BranchOutputHandle, +} + +impl SingleOutputsView { + pub fn mark_all_drained(&self) { + self.primary.mark_drained(); + } +} + +/// View for `OrderedBytesSingle` outputs. Same shape as `SingleOutputsView` +/// but the contained `BranchOutputHandle` is constructed via the +/// `_ordered_bytes` build path (which supports `T: HeapSize + Ordered` and +/// dispatches `BranchOrdering::ByItemOrdinal` correctly). +pub(crate) struct OrderedBytesSingleOutputsView { + pub(crate) primary: BranchOutputHandle, +} + +impl OrderedBytesSingleOutputsView { + pub fn mark_all_drained(&self) { + self.primary.mark_drained(); + } +} + +/// View for `(A, B)` tuple outputs. Internal: see `SingleOutputsView`. +pub(crate) struct Tuple2OutputsView { + pub(crate) a: BranchOutputHandle, + pub(crate) b: BranchOutputHandle, +} + +impl Tuple2OutputsView { + pub fn mark_all_drained(&self) { + self.a.mark_drained(); + self.b.mark_drained(); + } +} + +/// View for `(A, B, C)` tuple outputs. Internal: see `SingleOutputsView`. +pub(crate) struct Tuple3OutputsView +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, + C: Send + HeapSize + 'static, +{ + pub(crate) a: BranchOutputHandle, + pub(crate) b: BranchOutputHandle, + pub(crate) c: BranchOutputHandle, +} + +impl Tuple3OutputsView +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, + C: Send + HeapSize + 'static, +{ + pub fn mark_all_drained(&self) { + self.a.mark_drained(); + self.b.mark_drained(); + self.c.mark_drained(); + } +} + +/// View for `(A, B, C, D)` tuple outputs. Internal: see `SingleOutputsView`. +pub(crate) struct Tuple4OutputsView +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, + C: Send + HeapSize + 'static, + D: Send + HeapSize + 'static, +{ + pub(crate) a: BranchOutputHandle, + pub(crate) b: BranchOutputHandle, + pub(crate) c: BranchOutputHandle, + pub(crate) d: BranchOutputHandle, +} + +impl Tuple4OutputsView +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, + C: Send + HeapSize + 'static, + D: Send + HeapSize + 'static, +{ + pub fn mark_all_drained(&self) { + self.a.mark_drained(); + self.b.mark_drained(); + self.c.mark_drained(); + self.d.mark_drained(); + } +} + +/// View for `()` (sink) outputs. Internal: see `SingleOutputsView`. +/// No fields and no methods — sinks have no outputs to track or drain. +pub(crate) struct UnitOutputsView; + +// ───────────────────────────────────────────────────────────────────────────── +// Typed accessors on OutputHandles +// ───────────────────────────────────────────────────────────────────────────── + +/// The single canonical definition of the held-slot re-hold invariant, shared by +/// every single-output shape's `retry_held` (the `Single` and +/// `OrderedBytesSingle` variants differ only in which `retry` they hand in). +/// +/// If the slot holds an item, `retry` it; on rejection put it **back** in the slot +/// (so it is never dropped) and report [`HeldRetry::StillHeld`]. Used by step +/// flush-first preambles so each step doesn't re-implement the take/retry/put-back +/// dance (a copy that forgot the put-back would silently drop a final batch). +/// **Never spins** — the caller maps `StillHeld` to a yield (`NoProgress`/`Contention`) +/// and retries on the next dispatch. +#[inline] +fn retry_held_impl( + held: &mut crate::held::HeldSlot>, + retry: impl FnOnce(Unpushed) -> Result<(), Unpushed>, +) -> HeldRetry { + match held.take() { + None => HeldRetry::WasEmpty, + Some(unpushed) => match retry(unpushed) { + Ok(()) => HeldRetry::Flushed, + Err(again) => { + held.put(again); + HeldRetry::StillHeld + } + }, + } +} + +impl OutputHandles> { + /// Push a fresh item to the (single) output. + /// + /// # Errors + /// + /// Returns `Err(Unpushed)` when backpressure rejected the push. Hand + /// the rejected item to [`Self::retry`] on the next iteration so any + /// pre-allocated ordinal is preserved. + /// + /// # Panics + /// + /// Panics if the type-erased outputs view doesn't downcast to + /// `SingleOutputsView` (a framework invariant violation). + #[inline] + pub fn push(&self, item: T) -> Result<(), Unpushed> { + let view = self + .inner + .inner + .downcast_ref::>() + .expect("Single outputs view downcast failed"); + view.primary.push(item) + } + + /// Retry a previously-rejected push. + /// + /// # Errors + /// + /// Returns `Err(Unpushed)` if backpressure still rejected the push. + /// + /// # Panics + /// + /// Panics if the type-erased outputs view doesn't downcast to + /// `SingleOutputsView`. + #[inline] + pub fn retry(&self, unpushed: Unpushed) -> Result<(), Unpushed> { + let view = self + .inner + .inner + .downcast_ref::>() + .expect("Single outputs view downcast failed"); + view.primary.retry(unpushed) + } + + /// Retry a step's held output slot, re-holding on backpressure. + /// + /// Retry the held output slot for the `Single` shape, re-holding on + /// backpressure. Delegates to `retry_held_impl`, the shared canonical + /// re-hold invariant. + #[inline] + pub fn retry_held(&self, held: &mut crate::held::HeldSlot>) -> HeldRetry { + retry_held_impl(held, |unpushed| self.retry(unpushed)) + } +} + +impl OutputHandles> { + /// Push a fresh item to the heap-aware ordered output. + /// + /// # Errors + /// + /// Returns `Err(Unpushed)` when backpressure (count or byte) rejected + /// the push. Hand the rejected item to [`Self::retry`] on the next + /// iteration so the pre-allocated ordinal (if any) is preserved. + /// + /// # Panics + /// + /// Panics if the type-erased outputs view doesn't downcast to + /// `OrderedBytesSingleOutputsView` (a framework invariant violation). + #[inline] + pub fn push(&self, item: T) -> Result<(), Unpushed> { + let view = self + .inner + .inner + .downcast_ref::>() + .expect("OrderedBytesSingle outputs view downcast failed"); + view.primary.push(item) + } + + /// Retry a previously-rejected push on the heap-aware ordered output. + /// + /// # Errors + /// + /// Returns `Err(Unpushed)` if backpressure still rejected the push. + /// + /// # Panics + /// + /// Panics if the type-erased outputs view doesn't downcast to + /// `OrderedBytesSingleOutputsView`. + #[inline] + pub fn retry(&self, unpushed: Unpushed) -> Result<(), Unpushed> { + let view = self + .inner + .inner + .downcast_ref::>() + .expect("OrderedBytesSingle outputs view downcast failed"); + view.primary.retry(unpushed) + } + + /// Retry the held output slot for the `OrderedBytesSingle` shape, + /// re-holding on backpressure. Delegates to `retry_held_impl`, the shared + /// canonical re-hold invariant. + #[inline] + pub fn retry_held(&self, held: &mut crate::held::HeldSlot>) -> HeldRetry { + retry_held_impl(held, |unpushed| self.retry(unpushed)) + } +} + +impl OutputHandles<(A, B)> { + /// Borrow the typed per-branch view for a 2-tuple output. + /// + /// # Panics + /// + /// Panics if the type-erased outputs view doesn't downcast to + /// `Tuple2OutputsView` (a framework invariant violation). + #[must_use] + #[inline] + pub fn view(&self) -> Tuple2View<'_, A, B> { + let v = self + .inner + .inner + .downcast_ref::>() + .expect("Tuple-2 outputs view downcast failed"); + Tuple2View { a: &v.a, b: &v.b } + } +} + +impl OutputHandles> +where + A: Send + HeapSize + Ordered + 'static, + B: Send + HeapSize + Ordered + 'static, +{ + /// Borrow the typed per-branch view for an ordered + byte-bounded + /// 2-tuple output. Same view type as plain `(A, B)` (the view + /// only carries `BranchOutputHandle`s; the ordering is encoded in + /// the queue's transport, not the handle). + /// + /// # Panics + /// + /// Panics if the type-erased outputs view doesn't downcast to + /// `Tuple2OutputsView` (a framework invariant violation). + #[must_use] + #[inline] + pub fn view(&self) -> Tuple2View<'_, A, B> { + let v = self + .inner + .inner + .downcast_ref::>() + .expect("OrderedBytesTuple2 outputs view downcast failed"); + Tuple2View { a: &v.a, b: &v.b } + } +} + +impl OutputHandles> +where + A: Send + HeapSize + Ordered + 'static, + B: Send + HeapSize + Ordered + 'static, + C: Send + HeapSize + Ordered + 'static, +{ + /// Borrow the typed per-branch view for an ordered + byte-bounded 3-tuple + /// output. Same view type as plain `(A, B, C)` — the view carries only + /// `BranchOutputHandle`s; the ordering is encoded in the queue transport. + /// + /// # Panics + /// + /// Panics if the type-erased outputs view doesn't downcast to + /// `Tuple3OutputsView` (a framework invariant violation). + #[must_use] + #[inline] + pub fn view(&self) -> Tuple3View<'_, A, B, C> { + let v = self + .inner + .inner + .downcast_ref::>() + .expect("OrderedBytesTuple3 outputs view downcast failed"); + Tuple3View { a: &v.a, b: &v.b, c: &v.c } + } +} + +pub struct Tuple2View<'a, A: Send + HeapSize + 'static, B: Send + HeapSize + 'static> { + pub a: &'a BranchOutputHandle, + pub b: &'a BranchOutputHandle, +} + +impl OutputHandles<(A, B, C)> +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, + C: Send + HeapSize + 'static, +{ + /// Borrow the typed per-branch view for a 3-tuple output. + /// + /// # Panics + /// + /// Panics if the type-erased outputs view doesn't downcast to + /// `Tuple3OutputsView` (a framework invariant violation). + #[must_use] + #[inline] + pub fn view(&self) -> Tuple3View<'_, A, B, C> { + let v = self + .inner + .inner + .downcast_ref::>() + .expect("Tuple-3 outputs view downcast failed"); + Tuple3View { a: &v.a, b: &v.b, c: &v.c } + } +} + +pub struct Tuple3View<'a, A, B, C> +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, + C: Send + HeapSize + 'static, +{ + pub a: &'a BranchOutputHandle, + pub b: &'a BranchOutputHandle, + pub c: &'a BranchOutputHandle, +} + +impl OutputHandles<(A, B, C, D)> +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, + C: Send + HeapSize + 'static, + D: Send + HeapSize + 'static, +{ + /// Borrow the typed per-branch view for a 4-tuple output. + /// + /// # Panics + /// + /// Panics if the type-erased outputs view doesn't downcast to + /// `Tuple4OutputsView` (a framework invariant violation). + #[must_use] + #[inline] + pub fn view(&self) -> Tuple4View<'_, A, B, C, D> { + let v = self + .inner + .inner + .downcast_ref::>() + .expect("Tuple-4 outputs view downcast failed"); + Tuple4View { a: &v.a, b: &v.b, c: &v.c, d: &v.d } + } +} + +pub struct Tuple4View<'a, A, B, C, D> +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, + C: Send + HeapSize + 'static, + D: Send + HeapSize + 'static, +{ + pub a: &'a BranchOutputHandle, + pub b: &'a BranchOutputHandle, + pub c: &'a BranchOutputHandle, + pub d: &'a BranchOutputHandle, +} + +impl OutputHandles<()> { + /// Sinks have no output. Method exists for API symmetry. + pub fn noop(&self) {} +} + +// ───────────────────────────────────────────────────────────────────────────── +// Typed mark_all_drained on OutputHandles +// +// The framework's driver calls this through `TypedStep` when a step returns +// `StepOutcome::Finished` (counter-gated for `Parallel` so only the last clone +// closes the shared output). Each impl downcasts to the right per-arity view +// and forwards. +// ───────────────────────────────────────────────────────────────────────────── + +impl OutputHandles> { + /// Mark all output branches drained. + /// + /// # Panics + /// + /// Panics if the type-erased outputs view doesn't downcast to + /// `SingleOutputsView` (a framework invariant violation). + pub fn mark_all_drained(&self) { + let view = self + .inner + .inner + .downcast_ref::>() + .expect("Single outputs view downcast failed"); + view.mark_all_drained(); + } +} + +impl OutputHandles> { + /// Mark all output branches drained. + /// + /// # Panics + /// + /// Panics if the type-erased outputs view doesn't downcast to + /// `OrderedBytesSingleOutputsView` (a framework invariant violation). + pub fn mark_all_drained(&self) { + let view = self + .inner + .inner + .downcast_ref::>() + .expect("OrderedBytesSingle outputs view downcast failed"); + view.mark_all_drained(); + } +} + +impl OutputHandles<(A, B)> { + /// Mark all output branches drained. + /// + /// # Panics + /// + /// Panics if the type-erased outputs view doesn't downcast to + /// `Tuple2OutputsView` (a framework invariant violation). + pub fn mark_all_drained(&self) { + let view = self + .inner + .inner + .downcast_ref::>() + .expect("Tuple-2 outputs view downcast failed"); + view.mark_all_drained(); + } +} + +impl OutputHandles> +where + A: Send + HeapSize + Ordered + 'static, + B: Send + HeapSize + Ordered + 'static, +{ + /// Mark all output branches drained. + /// + /// # Panics + /// + /// Panics if the type-erased outputs view doesn't downcast to + /// `Tuple2OutputsView` (the view shape is shared with plain + /// `(A, B)`). + pub fn mark_all_drained(&self) { + let view = self + .inner + .inner + .downcast_ref::>() + .expect("OrderedBytesTuple2 outputs view downcast failed"); + view.mark_all_drained(); + } +} + +impl OutputHandles> +where + A: Send + HeapSize + Ordered + 'static, + B: Send + HeapSize + Ordered + 'static, + C: Send + HeapSize + Ordered + 'static, +{ + /// Mark all output branches drained. + /// + /// # Panics + /// + /// Panics if the type-erased outputs view doesn't downcast to + /// `Tuple3OutputsView` (the view shape is shared with plain + /// `(A, B, C)`). + pub fn mark_all_drained(&self) { + let view = self + .inner + .inner + .downcast_ref::>() + .expect("OrderedBytesTuple3 outputs view downcast failed"); + view.mark_all_drained(); + } +} + +impl OutputHandles<(A, B, C)> +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, + C: Send + HeapSize + 'static, +{ + /// Mark all output branches drained. + /// + /// # Panics + /// + /// Panics if the type-erased outputs view doesn't downcast to + /// `Tuple3OutputsView` (a framework invariant violation). + pub fn mark_all_drained(&self) { + let view = self + .inner + .inner + .downcast_ref::>() + .expect("Tuple-3 outputs view downcast failed"); + view.mark_all_drained(); + } +} + +impl OutputHandles<(A, B, C, D)> +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, + C: Send + HeapSize + 'static, + D: Send + HeapSize + 'static, +{ + /// Mark all output branches drained. + /// + /// # Panics + /// + /// Panics if the type-erased outputs view doesn't downcast to + /// `Tuple4OutputsView` (a framework invariant violation). + pub fn mark_all_drained(&self) { + let view = self + .inner + .inner + .downcast_ref::>() + .expect("Tuple-4 outputs view downcast failed"); + view.mark_all_drained(); + } +} + +impl OutputHandles<()> { + /// Mark all output branches drained. Sinks have no outputs; no-op. + pub fn mark_all_drained(&self) { + // No-op for unit outputs. + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// build_*_queues per arity — entry points called from outputs.rs +// ───────────────────────────────────────────────────────────────────────────── + +pub(crate) fn build_single_queues( + specs: &[QueueSpec], + ordering: &[BranchOrdering], + level: crate::builder::InstrumentationLevel, +) -> (OutputQueueSet, OutputsViewAny) { + assert_eq!(specs.len(), 1, "Single::build_queues requires 1 spec"); + assert_eq!(ordering.len(), 1, "Single::build_queues requires 1 ordering"); + + // `Single` bounds `T: HeapSize`, so use the byte-aware build path: it + // honors `QueueSpec::ByteBounded` for `BranchOrdering::None` and + // `ByOrdinal` — the two orderings a `Single` output can declare, since it + // carries no item serial — and delegates every non-byte spec straight back + // to `build_branch::`, so count/unbounded paths are unchanged. A step + // needing `ByItemOrdinal` declares `OrderedBytesSingle` instead; asking + // for it here panics (see `build_branch_byte_aware`). + let branch = build_branch_byte_aware::(specs[0], ordering[0], level); + let view = SingleOutputsView { primary: branch.output }; + let outputs_view = OutputsViewAny { inner: Box::new(view) }; + let queue_set = OutputQueueSet::new(vec![BranchEntry { + input_handle: Box::new(branch.input), + bounded_queue_handle: branch.bounded_queue_handle, + metrics: branch.metrics, + }]); + (queue_set, outputs_view) +} + +/// Build queues for `OrderedBytesSingle` outputs (`T: HeapSize + Ordered`). +/// Supports every `QueueSpec` × `BranchOrdering` combination. +pub(crate) fn build_single_queues_ordered_bytes( + specs: &[QueueSpec], + ordering: &[BranchOrdering], + level: crate::builder::InstrumentationLevel, +) -> (OutputQueueSet, OutputsViewAny) { + assert_eq!(specs.len(), 1, "OrderedBytesSingle::build_queues requires 1 spec"); + assert_eq!(ordering.len(), 1, "OrderedBytesSingle::build_queues requires 1 ordering"); + + let branch = build_branch_ordered_bytes::(specs[0], ordering[0], level); + let view = OrderedBytesSingleOutputsView { primary: branch.output }; + let outputs_view = OutputsViewAny { inner: Box::new(view) }; + let queue_set = OutputQueueSet::new(vec![BranchEntry { + input_handle: Box::new(branch.input), + bounded_queue_handle: branch.bounded_queue_handle, + metrics: branch.metrics, + }]); + (queue_set, outputs_view) +} + +/// Fan-out build path for an `(A, B)` tuple output. +/// +/// Every branch goes through `build_branch_byte_aware`, matching +/// `build_single_queues`: each type parameter is bounded `HeapSize`, so a branch +/// may declare `QueueSpec::ByteBounded`, and the plain `build_branch` panics on +/// that spec. Non-byte specs are delegated straight back to `build_branch`, so +/// count/unbounded branches are unchanged. This matters beyond fan-out ergonomics: +/// an armed deadlock monitor REQUIRES `ByteBounded` on every output transport +/// (`ensure_monitor_visible_transports`), so a fan-out chain could not be built +/// at all while these branches rejected the spec. +pub(crate) fn build_tuple2_queues( + specs: &[QueueSpec], + ordering: &[BranchOrdering], + level: crate::builder::InstrumentationLevel, +) -> (OutputQueueSet, OutputsViewAny) +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, +{ + assert_eq!(specs.len(), 2, "(A, B)::build_queues requires 2 specs"); + assert_eq!(ordering.len(), 2, "(A, B)::build_queues requires 2 orderings"); + + let ba = build_branch_byte_aware::(specs[0], ordering[0], level); + let bb = build_branch_byte_aware::(specs[1], ordering[1], level); + let view = Tuple2OutputsView { a: ba.output, b: bb.output }; + let outputs_view = OutputsViewAny { inner: Box::new(view) }; + let queue_set = OutputQueueSet::new(vec![ + BranchEntry { + input_handle: Box::new(ba.input), + bounded_queue_handle: ba.bounded_queue_handle, + metrics: ba.metrics, + }, + BranchEntry { + input_handle: Box::new(bb.input), + bounded_queue_handle: bb.bounded_queue_handle, + metrics: bb.metrics, + }, + ]); + (queue_set, outputs_view) +} + +/// Build queues for `OrderedBytesTuple2` outputs. Both branches +/// support the full `QueueSpec × BranchOrdering` cross-product because +/// `A: Ordered + HeapSize` and `B: Ordered + HeapSize`. The +/// `Tuple2OutputsView` carries only `BranchOutputHandle`s (no ordering +/// metadata), so the view type is the same as for plain `(A, B)`. +pub(crate) fn build_tuple2_queues_ordered_bytes( + specs: &[QueueSpec], + ordering: &[BranchOrdering], + level: crate::builder::InstrumentationLevel, +) -> (OutputQueueSet, OutputsViewAny) +where + A: Send + HeapSize + Ordered + 'static, + B: Send + HeapSize + Ordered + 'static, +{ + assert_eq!(specs.len(), 2, "OrderedBytesTuple2::build_queues requires 2 specs"); + assert_eq!(ordering.len(), 2, "OrderedBytesTuple2::build_queues requires 2 orderings"); + + let ba = build_branch_ordered_bytes::(specs[0], ordering[0], level); + let bb = build_branch_ordered_bytes::(specs[1], ordering[1], level); + let view = Tuple2OutputsView { a: ba.output, b: bb.output }; + let outputs_view = OutputsViewAny { inner: Box::new(view) }; + let queue_set = OutputQueueSet::new(vec![ + BranchEntry { + input_handle: Box::new(ba.input), + bounded_queue_handle: ba.bounded_queue_handle, + metrics: ba.metrics, + }, + BranchEntry { + input_handle: Box::new(bb.input), + bounded_queue_handle: bb.bounded_queue_handle, + metrics: bb.metrics, + }, + ]); + (queue_set, outputs_view) +} + +/// Fan-out build path for an `(A, B, C)` tuple output. +/// +/// Every branch goes through `build_branch_byte_aware`, matching +/// `build_single_queues`: each type parameter is bounded `HeapSize`, so a branch +/// may declare `QueueSpec::ByteBounded`, and the plain `build_branch` panics on +/// that spec. Non-byte specs are delegated straight back to `build_branch`, so +/// count/unbounded branches are unchanged. This matters beyond fan-out ergonomics: +/// an armed deadlock monitor REQUIRES `ByteBounded` on every output transport +/// (`ensure_monitor_visible_transports`), so a fan-out chain could not be built +/// at all while these branches rejected the spec. +pub(crate) fn build_tuple3_queues( + specs: &[QueueSpec], + ordering: &[BranchOrdering], + level: crate::builder::InstrumentationLevel, +) -> (OutputQueueSet, OutputsViewAny) +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, + C: Send + HeapSize + 'static, +{ + assert_eq!(specs.len(), 3, "(A, B, C)::build_queues requires 3 specs"); + assert_eq!(ordering.len(), 3, "(A, B, C)::build_queues requires 3 orderings"); + + let ba = build_branch_byte_aware::(specs[0], ordering[0], level); + let bb = build_branch_byte_aware::(specs[1], ordering[1], level); + let bc = build_branch_byte_aware::(specs[2], ordering[2], level); + let view = Tuple3OutputsView { a: ba.output, b: bb.output, c: bc.output }; + let outputs_view = OutputsViewAny { inner: Box::new(view) }; + let queue_set = OutputQueueSet::new(vec![ + BranchEntry { + input_handle: Box::new(ba.input), + bounded_queue_handle: ba.bounded_queue_handle, + metrics: ba.metrics, + }, + BranchEntry { + input_handle: Box::new(bb.input), + bounded_queue_handle: bb.bounded_queue_handle, + metrics: bb.metrics, + }, + BranchEntry { + input_handle: Box::new(bc.input), + bounded_queue_handle: bc.bounded_queue_handle, + metrics: bc.metrics, + }, + ]); + (queue_set, outputs_view) +} + +/// Build queues for `OrderedBytesTuple3` outputs. All three branches +/// support the full `QueueSpec × BranchOrdering` cross-product because each is +/// `Ordered + HeapSize`. The `Tuple3OutputsView` carries only +/// `BranchOutputHandle`s (no ordering metadata), so the view type is the same +/// as for plain `(A, B, C)`. +pub(crate) fn build_tuple3_queues_ordered_bytes( + specs: &[QueueSpec], + ordering: &[BranchOrdering], + level: crate::builder::InstrumentationLevel, +) -> (OutputQueueSet, OutputsViewAny) +where + A: Send + HeapSize + Ordered + 'static, + B: Send + HeapSize + Ordered + 'static, + C: Send + HeapSize + Ordered + 'static, +{ + assert_eq!(specs.len(), 3, "OrderedBytesTuple3::build_queues requires 3 specs"); + assert_eq!(ordering.len(), 3, "OrderedBytesTuple3::build_queues requires 3 orderings"); + + let ba = build_branch_ordered_bytes::(specs[0], ordering[0], level); + let bb = build_branch_ordered_bytes::(specs[1], ordering[1], level); + let bc = build_branch_ordered_bytes::(specs[2], ordering[2], level); + let view = Tuple3OutputsView { a: ba.output, b: bb.output, c: bc.output }; + let outputs_view = OutputsViewAny { inner: Box::new(view) }; + let queue_set = OutputQueueSet::new(vec![ + BranchEntry { + input_handle: Box::new(ba.input), + bounded_queue_handle: ba.bounded_queue_handle, + metrics: ba.metrics, + }, + BranchEntry { + input_handle: Box::new(bb.input), + bounded_queue_handle: bb.bounded_queue_handle, + metrics: bb.metrics, + }, + BranchEntry { + input_handle: Box::new(bc.input), + bounded_queue_handle: bc.bounded_queue_handle, + metrics: bc.metrics, + }, + ]); + (queue_set, outputs_view) +} + +/// Fan-out build path for an `(A, B, C, D)` tuple output. +/// +/// Every branch goes through `build_branch_byte_aware`, matching +/// `build_single_queues`: each type parameter is bounded `HeapSize`, so a branch +/// may declare `QueueSpec::ByteBounded`, and the plain `build_branch` panics on +/// that spec. Non-byte specs are delegated straight back to `build_branch`, so +/// count/unbounded branches are unchanged. This matters beyond fan-out ergonomics: +/// an armed deadlock monitor REQUIRES `ByteBounded` on every output transport +/// (`ensure_monitor_visible_transports`), so a fan-out chain could not be built +/// at all while these branches rejected the spec. +pub(crate) fn build_tuple4_queues( + specs: &[QueueSpec], + ordering: &[BranchOrdering], + level: crate::builder::InstrumentationLevel, +) -> (OutputQueueSet, OutputsViewAny) +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, + C: Send + HeapSize + 'static, + D: Send + HeapSize + 'static, +{ + assert_eq!(specs.len(), 4, "(A, B, C, D)::build_queues requires 4 specs"); + assert_eq!(ordering.len(), 4, "(A, B, C, D)::build_queues requires 4 orderings"); + + let ba = build_branch_byte_aware::(specs[0], ordering[0], level); + let bb = build_branch_byte_aware::(specs[1], ordering[1], level); + let bc = build_branch_byte_aware::(specs[2], ordering[2], level); + let bd = build_branch_byte_aware::(specs[3], ordering[3], level); + let view = Tuple4OutputsView { a: ba.output, b: bb.output, c: bc.output, d: bd.output }; + let outputs_view = OutputsViewAny { inner: Box::new(view) }; + let queue_set = OutputQueueSet::new(vec![ + BranchEntry { + input_handle: Box::new(ba.input), + bounded_queue_handle: ba.bounded_queue_handle, + metrics: ba.metrics, + }, + BranchEntry { + input_handle: Box::new(bb.input), + bounded_queue_handle: bb.bounded_queue_handle, + metrics: bb.metrics, + }, + BranchEntry { + input_handle: Box::new(bc.input), + bounded_queue_handle: bc.bounded_queue_handle, + metrics: bc.metrics, + }, + BranchEntry { + input_handle: Box::new(bd.input), + bounded_queue_handle: bd.bounded_queue_handle, + metrics: bd.metrics, + }, + ]); + (queue_set, outputs_view) +} + +pub(crate) fn build_unit_queues( + specs: &[QueueSpec], + ordering: &[BranchOrdering], + _level: crate::builder::InstrumentationLevel, +) -> (OutputQueueSet, OutputsViewAny) { + assert_eq!(specs.len(), 0, "()::build_queues requires 0 specs"); + assert_eq!(ordering.len(), 0, "()::build_queues requires 0 orderings"); + let outputs_view = OutputsViewAny { inner: Box::new(UnitOutputsView) }; + let queue_set = OutputQueueSet::new(Vec::new()); + (queue_set, outputs_view) +} + +// ───────────────────────────────────────────────────────────────────────────── +// Tests +// ───────────────────────────────────────────────────────────────────────────── + +#[cfg(test)] +mod handle_tests { + use super::*; + use rstest::rstest; + + #[test] + fn always_drained_handle_is_empty_and_drained() { + // The zero-state source input handle owns no transport: it pops nothing + // and reports drained from the start, so a source step sees its + // (implicit) input as immediately end-of-stream. + let h = BranchInputHandle::<()>::always_drained(); + assert_eq!(h.pop(), None, "always-drained handle yields no items"); + assert!(h.is_drained(), "always-drained handle reports drained"); + // Idempotent: still drained, still empty after repeated reads. + assert_eq!(h.pop(), None); + assert!(h.is_drained()); + } + + #[test] + fn build_branch_mints_metrics_iff_on() { + use crate::builder::InstrumentationLevel as L; + let on = build_branch::( + QueueSpec::CountBounded { capacity: 4 }, + BranchOrdering::None, + L::Summary, + ); + assert!(on.metrics.is_some(), "level on → edge metrics minted"); + let off = build_branch::( + QueueSpec::CountBounded { capacity: 4 }, + BranchOrdering::None, + L::Off, + ); + assert!(off.metrics.is_none(), "level off → no metrics (hot path metric-free)"); + } + + #[test] + fn metrics_shared_between_transport_push_and_input_handle_pop() { + // Producer-push (transport) and consumer-pop (input handle) share one + // EdgeMetrics: push lands via the queue, pop/empty via the input handle. + use crate::builder::InstrumentationLevel as L; + let b = build_branch::( + QueueSpec::CountBounded { capacity: 4 }, + BranchOrdering::None, + L::Summary, + ); + let m = b.metrics.clone().expect("metrics present"); + b.output.push(7).unwrap(); + assert_eq!(b.input.pop(), Some(7)); + assert_eq!(b.input.pop(), None); // empty + let s = m.snapshot(); + assert_eq!(s.pushed_items, 1, "producer push counted at the transport"); + assert_eq!(s.popped_items, 1, "consumer pop counted at the input handle"); + assert_eq!(s.pop_empties, 1, "empty pop counted at the input handle"); + } + + #[test] + fn ordered_reorder_blocked_pop_is_not_counted_empty() { + // Regression: an ordered edge's `try_pop_in_order` returns `None` both + // when the edge is genuinely starved AND when it is only reorder-blocked + // (later ordinals are buffered while it waits for the next in-order + // ordinal). Counting the reorder-blocked case as an empty pop inflates + // `pop_empties` and can misclassify a backlogged edge as starved, so + // the pop path must skip `record_empty()` while the reorder buffer holds + // out-of-order items. + use crate::runtime::metrics::EdgeMetrics; + + let transport: Arc>> = + Arc::new(CountBoundedQueue::>::new(8)); + let stage = Arc::new(ReorderStage::new(transport)); + // Buffer ordinal 1 while ordinal 0 is still absent: the next in-order + // pop is blocked, not starved. + stage.try_push(1, 100).unwrap(); + + let m = EdgeMetrics::new(); + let handle = BranchInputHandle { + inner: BranchInputInner::Ordered(stage.clone()), + metrics: Some(m.clone()), + // Count-bounded transport, so no byte accounting on either side. + record_item_bytes: false, + }; + + // Reorder-blocked: yields nothing yet, but there is buffered work. + assert_eq!(handle.pop(), None, "next ordinal (0) absent → no in-order item"); + assert_eq!( + m.snapshot().pop_empties, + 0, + "a reorder-blocked pop is backlog, not starvation — must not count as empty" + ); + + // Now the missing ordinal arrives and both items drain in order; still + // no empty pops recorded. + stage.try_push(0, 50).unwrap(); + assert_eq!(handle.pop(), Some(50)); + assert_eq!(handle.pop(), Some(100)); + assert_eq!(m.snapshot().pop_empties, 0, "in-order drains record no empties"); + + // Genuinely drained now: this pop IS a starved/empty pop. + assert_eq!(handle.pop(), None); + assert_eq!( + m.snapshot().pop_empties, + 1, + "an empty pop with no buffered work is a true empty pop" + ); + } + + #[test] + fn count_bounded_fifo_round_trip() { + let b = build_branch::( + QueueSpec::CountBounded { capacity: 4 }, + BranchOrdering::None, + crate::builder::InstrumentationLevel::Off, + ); + b.output.push(1).unwrap(); + b.output.push(2).unwrap(); + assert_eq!(b.input.pop(), Some(1)); + assert_eq!(b.input.pop(), Some(2)); + assert_eq!(b.input.pop(), None); + } + + #[test] + fn count_bounded_ordered_emits_in_order() { + let b = build_branch::( + QueueSpec::CountBounded { capacity: 8 }, + BranchOrdering::ByOrdinal, + crate::builder::InstrumentationLevel::Off, + ); + // Single producer pushes ordinals 0,1,2 in arrival order. + b.output.push(100).unwrap(); + b.output.push(200).unwrap(); + b.output.push(300).unwrap(); + assert_eq!(b.input.pop(), Some(100)); + assert_eq!(b.input.pop(), Some(200)); + assert_eq!(b.input.pop(), Some(300)); + } + + #[test] + fn drained_signal_propagates() { + let b = build_branch::( + QueueSpec::CountBounded { capacity: 4 }, + BranchOrdering::None, + crate::builder::InstrumentationLevel::Off, + ); + b.output.push(1).unwrap(); + b.output.mark_drained(); + // Marker set but item still buffered: not drained. + assert!(!b.input.is_drained()); + assert_eq!(b.input.pop(), Some(1)); + assert!(b.input.is_drained()); + } + + #[test] + fn unbounded_branch_works() { + let b = build_branch::( + QueueSpec::Unbounded, + BranchOrdering::None, + crate::builder::InstrumentationLevel::Off, + ); + for i in 0..1024 { + b.output.push(i).unwrap(); + } + for i in 0..1024 { + assert_eq!(b.input.pop(), Some(i)); + } + } + + #[derive(Debug)] + struct Bytes(Vec); + impl HeapSize for Bytes { + fn heap_size(&self) -> usize { + self.0.len() + } + } + + #[test] + fn byte_bounded_branch_works() { + let b = build_branch_byte_aware::( + QueueSpec::ByteBounded { limit_bytes: 1000 }, + BranchOrdering::None, + crate::builder::InstrumentationLevel::Off, + ); + b.output.push(Bytes(vec![0; 500])).unwrap(); + let popped = b.input.pop().unwrap(); + assert_eq!(popped.0.len(), 500); + } + + /// Every fan-out shape must accept `QueueSpec::ByteBounded` on every branch. + /// The tuple builders used the non-byte-aware `build_branch`, which panics on + /// that spec, so a fan-out step declaring it could not be built — and an armed + /// deadlock monitor *requires* `ByteBounded` on every output transport + /// (`ensure_monitor_visible_transports`), which made a monitored fan-out chain + /// unbuildable. Asserting the registered byte handle (not just "it didn't + /// panic") also pins that the branch is genuinely byte-bounded and therefore + /// visible to the monitor and the queue-memory budget. + #[rstest] + #[case::tuple2(2)] + #[case::tuple3(3)] + #[case::tuple4(4)] + fn byte_bounded_accepted_on_every_fan_out_branch(#[case] n_branches: usize) { + let specs = vec![QueueSpec::ByteBounded { limit_bytes: 1000 }; n_branches]; + let orderings = vec![BranchOrdering::None; n_branches]; + let level = crate::builder::InstrumentationLevel::Off; + let (queue_set, _view) = match n_branches { + 2 => build_tuple2_queues::(&specs, &orderings, level), + 3 => build_tuple3_queues::(&specs, &orderings, level), + 4 => build_tuple4_queues::(&specs, &orderings, level), + other => panic!("unhandled branch count {other}"), + }; + assert_eq!(queue_set.n_branches(), n_branches); + for branch in 0..n_branches { + let handles = queue_set.branches[branch] + .bounded_queue_handle + .as_ref() + .unwrap_or_else(|| panic!("branch {branch} must register a byte-bounded handle")); + assert_eq!( + handles.transport.limit_bytes(), + 1000, + "branch {branch} must carry the declared byte bound" + ); + } + } + + /// `popped_bytes` must agree with `pushed_bytes` about whether an edge counts + /// bytes at all. A count/unbounded transport records `record_push(0)` by + /// design (no byte accounting on its hot path), so the pop side must report 0 + /// too — otherwise `compute_edge_stats` divides real popped bytes by wall time + /// and reports a `mibytes_per_s` throughput for an edge whose pushed bytes are + /// always 0. A byte-bounded edge reports real sizes on both sides. + #[rstest] + #[case::count_bounded(QueueSpec::CountBounded { capacity: 4 }, 0)] + #[case::unbounded(QueueSpec::Unbounded, 0)] + #[case::byte_bounded(QueueSpec::ByteBounded { limit_bytes: 4096 }, 700)] + fn popped_bytes_recorded_only_on_byte_bounded_edges( + #[case] spec: QueueSpec, + #[case] expected_popped_bytes: u64, + ) { + // `Summary` mints the edge metrics; `Off` would leave `metrics: None`. + let b = build_branch_byte_aware::( + spec, + BranchOrdering::None, + crate::builder::InstrumentationLevel::Summary, + ); + b.output.push(Bytes(vec![0; 700])).expect("first push fits every spec here"); + assert!(b.input.pop().is_some(), "the pushed item comes back"); + + let ms = b.metrics.as_ref().expect("Summary mints edge metrics").snapshot(); + assert_eq!( + ms.popped_bytes, expected_popped_bytes, + "popped_bytes must mirror the push side's byte accounting" + ); + assert_eq!( + ms.pushed_bytes, expected_popped_bytes, + "pushed and popped byte accounting must agree for the same edge" + ); + assert_eq!(ms.popped_items, 1, "the item itself is always counted"); + } + + #[test] + #[should_panic(expected = "ByteBounded requires `T: HeapSize`")] + fn byte_bounded_panics_in_non_heap_aware_builder() { + let _ = build_branch::( + QueueSpec::ByteBounded { limit_bytes: 1000 }, + BranchOrdering::None, + crate::builder::InstrumentationLevel::Off, + ); + } + + #[test] + fn byte_bounded_single_builds_via_user_facing_build_queues() { + // Regression: a user step declaring `QueueSpec::ByteBounded` in its + // `StepProfile::output_queues` with a `Single` output reaches + // `StepOutputs::build_queues` → `build_single_queues`. `Single` + // bounds `T: HeapSize`, so this builds a byte-bounded queue (the + // documented "byte-bounded without item-carried serials" shape) + // rather than panicking. + use crate::outputs::{Single, StepOutputs}; + use crate::step::OutputHandles; + let (mut queue_set, outputs_view) = as StepOutputs>::build_queues( + &[QueueSpec::ByteBounded { limit_bytes: 1000 }], + &[BranchOrdering::None], + crate::builder::InstrumentationLevel::Off, + ); + let outputs: OutputHandles> = OutputHandles::new(outputs_view); + outputs.push(7).unwrap(); + let input = queue_set.take_typed_input::(0); + assert_eq!(input.pop(), Some(7)); + } + + #[test] + fn single_retry_held_drains_then_reholds() { + // Exercises `OutputHandles>::retry_held`: empty → WasEmpty, + // a rejected push re-held → StillHeld while the queue is full → Flushed + // once a slot frees, with the held item never dropped. + use crate::held::HeldSlot; + use crate::outputs::{Single, StepOutputs}; + use crate::step::OutputHandles; + let (mut queue_set, outputs_view) = as StepOutputs>::build_queues( + &[QueueSpec::CountBounded { capacity: 1 }], + &[BranchOrdering::None], + crate::builder::InstrumentationLevel::Off, + ); + let outputs: OutputHandles> = OutputHandles::new(outputs_view); + let mut held: HeldSlot> = HeldSlot::new(); + + // Empty slot → WasEmpty (queue untouched). + assert!(matches!(outputs.retry_held(&mut held), HeldRetry::WasEmpty)); + + // Fill the capacity-1 queue; the next push is rejected and re-held. + outputs.push(1).unwrap(); + let rejected = outputs.push(2).expect_err("capacity-1 queue must reject the 2nd push"); + held.put(rejected); + + // Queue still full → StillHeld; the item is put back, not dropped. + assert!(matches!(outputs.retry_held(&mut held), HeldRetry::StillHeld)); + assert!(held.is_held()); + + // Drain one item; the held push now flushes. + let input = queue_set.take_typed_input::(0); + assert_eq!(input.pop(), Some(1)); + assert!(matches!(outputs.retry_held(&mut held), HeldRetry::Flushed)); + assert!(!held.is_held()); + assert_eq!(input.pop(), Some(2)); + } + + #[test] + fn byte_bounded_plus_byordinal_works_after_phase3_amendment() { + // Phase 3 amendment 2 lifted the deferred ByteBounded + ByOrdinal + // path: `Sequenced: HeapSize` is now impl'd, so a byte-bounded + // transport can wrap `Sequenced` and the `Allocated` ordinal + // source provides per-branch ordinals. + let b = build_branch_byte_aware::( + QueueSpec::ByteBounded { limit_bytes: 1000 }, + BranchOrdering::ByOrdinal, + crate::builder::InstrumentationLevel::Off, + ); + b.output.push(Bytes(vec![0; 100])).unwrap(); + let popped = b.input.pop().unwrap(); + assert_eq!(popped.0.len(), 100); + } + + #[test] + #[should_panic(expected = "BranchOrdering::ByItemOrdinal requires `T: Ordered`")] + fn by_item_ordinal_panics_in_non_ordered_builder() { + let _ = build_branch::( + QueueSpec::CountBounded { capacity: 4 }, + BranchOrdering::ByItemOrdinal, + crate::builder::InstrumentationLevel::Off, + ); + } + + /// Trivial `Ordered`-impl test fixture for the by-item-ordinal builders. + #[derive(Debug, PartialEq, Eq)] + struct OrdItem { + ord: u64, + v: u32, + } + impl crate::item::Ordered for OrdItem { + fn ordinal(&self) -> u64 { + self.ord + } + } + + #[test] + fn by_item_ordinal_uses_item_serial() { + let b = build_branch_ordered::( + QueueSpec::CountBounded { capacity: 8 }, + BranchOrdering::ByItemOrdinal, + crate::builder::InstrumentationLevel::Off, + ); + // Push out of order: ordinals 2, 0, 1. + b.output.push(OrdItem { ord: 2, v: 200 }).unwrap(); + b.output.push(OrdItem { ord: 0, v: 0 }).unwrap(); + b.output.push(OrdItem { ord: 1, v: 100 }).unwrap(); + // Consumer sees them in item-ordinal order. + assert_eq!(b.input.pop().unwrap().v, 0); + assert_eq!(b.input.pop().unwrap().v, 100); + assert_eq!(b.input.pop().unwrap().v, 200); + } + + impl HeapSize for OrdItem { + fn heap_size(&self) -> usize { + std::mem::size_of::() + } + } + + #[test] + fn ordered_bytes_supports_byte_bounded_with_item_serial() { + let b = build_branch_ordered_bytes::( + QueueSpec::ByteBounded { limit_bytes: 4096 }, + BranchOrdering::ByItemOrdinal, + crate::builder::InstrumentationLevel::Off, + ); + b.output.push(OrdItem { ord: 1, v: 1 }).unwrap(); + b.output.push(OrdItem { ord: 0, v: 0 }).unwrap(); + assert_eq!(b.input.pop().unwrap().v, 0); + assert_eq!(b.input.pop().unwrap().v, 1); + } + + #[test] + fn ordered_branch_preserves_ordinal_across_retry() { + // Regression for C1 (ordinal-burn): with capacity 2, push three items + // through an ordered branch, draining one at a time. Without the + // Unpushed retry, the third push would burn an ordinal on + // backpressure and stall the reorder stage. + let b = build_branch::( + QueueSpec::CountBounded { capacity: 2 }, + BranchOrdering::ByOrdinal, + crate::builder::InstrumentationLevel::Off, + ); + + // Pump 5 items through the branch, draining sequentially. Capacity 2 + // means every push after the second hits backpressure once. + let mut received = Vec::new(); + for n in 0..5_u32 { + // Try to push. Retry on backpressure until accepted, draining the + // input handle in between to make progress. + let mut held: Option> = None; + let mut fresh: Option = Some(n); + + loop { + if let Some(unpushed) = held.take() { + match b.output.retry(unpushed) { + Ok(()) => {} + Err(again) => { + held = Some(again); + } + } + } + if held.is_none() + && let Some(item) = fresh.take() + { + match b.output.push(item) { + Ok(()) => {} + Err(unpushed) => { + held = Some(unpushed); + } + } + } + if held.is_none() && fresh.is_none() { + break; + } + // Drain any items the consumer has ready, freeing transport space. + while let Some(v) = b.input.pop() { + received.push(v); + } + } + } + // Final drain — drain everything still buffered. + while let Some(v) = b.input.pop() { + received.push(v); + } + assert_eq!(received, vec![0, 1, 2, 3, 4], "ordinals preserved across retries"); + } + + /// An `Unpushed` carrying no ordinal handed to an `Ordered` branch's + /// `retry` must panic in every build, not silently re-push. Re-pushing + /// allocates a fresh ordinal and abandons the original, so the consumer's + /// `ReorderStage` waits on the missing ordinal forever — a stall the + /// deadlock monitor reports with no hint at the cause. + #[test] + #[should_panic(expected = "Unpushed::ordinal=None on an Ordered branch")] + fn ordered_retry_without_ordinal_panics() { + let b = build_branch::( + QueueSpec::CountBounded { capacity: 2 }, + BranchOrdering::ByOrdinal, + crate::builder::InstrumentationLevel::Off, + ); + let _ = b.output.retry(Unpushed { item: 7, ordinal: None }); + } + + #[test] + fn output_queue_set_take_typed_input() { + let (mut set, _view) = build_single_queues::( + &[QueueSpec::CountBounded { capacity: 4 }], + &[BranchOrdering::None], + crate::builder::InstrumentationLevel::Off, + ); + let _input: BranchInputHandle = set.take_typed_input::(0); + } +} + +#[cfg(test)] +mod build_queues_tests { + use super::*; + use crate::outputs::{Single, StepOutputs}; + use crate::step::OutputHandles; + + fn count_specs(arity: usize, capacity: usize) -> (Vec, Vec) { + (vec![QueueSpec::CountBounded { capacity }; arity], vec![BranchOrdering::None; arity]) + } + + #[test] + fn single_round_trip() { + let (specs, ordering) = count_specs(1, 4); + let (mut queue_set, outputs_view) = as StepOutputs>::build_queues( + &specs, + &ordering, + crate::builder::InstrumentationLevel::Off, + ); + let outputs: OutputHandles> = OutputHandles::new(outputs_view); + outputs.push(7).unwrap(); + + let input = queue_set.take_typed_input::(0); + assert_eq!(input.pop(), Some(7)); + } + + #[test] + fn tuple_2_round_trip() { + let (specs, ordering) = count_specs(2, 2); + let (mut queue_set, outputs_view) = <(u32, String) as StepOutputs>::build_queues( + &specs, + &ordering, + crate::builder::InstrumentationLevel::Off, + ); + let outputs: OutputHandles<(u32, String)> = OutputHandles::new(outputs_view); + + let v = outputs.view(); + v.a.push(10).unwrap(); + v.b.push("hello".to_string()).unwrap(); + + let in_a = queue_set.take_typed_input::(0); + let in_b = queue_set.take_typed_input::(1); + assert_eq!(in_a.pop(), Some(10)); + assert_eq!(in_b.pop(), Some("hello".to_string())); + } + + #[test] + fn tuple_3_round_trip() { + let (specs, ordering) = count_specs(3, 2); + let (mut queue_set, outputs_view) = <(u32, u64, String) as StepOutputs>::build_queues( + &specs, + &ordering, + crate::builder::InstrumentationLevel::Off, + ); + let outputs: OutputHandles<(u32, u64, String)> = OutputHandles::new(outputs_view); + + let v = outputs.view(); + v.a.push(1).unwrap(); + v.b.push(2).unwrap(); + v.c.push("three".to_string()).unwrap(); + + assert_eq!(queue_set.take_typed_input::(0).pop(), Some(1)); + assert_eq!(queue_set.take_typed_input::(1).pop(), Some(2)); + assert_eq!(queue_set.take_typed_input::(2).pop(), Some("three".to_string())); + } + + #[test] + fn tuple_4_round_trip() { + let (specs, ordering) = count_specs(4, 2); + let (mut queue_set, outputs_view) = + <(u32, u64, String, Vec) as StepOutputs>::build_queues( + &specs, + &ordering, + crate::builder::InstrumentationLevel::Off, + ); + let outputs: OutputHandles<(u32, u64, String, Vec)> = OutputHandles::new(outputs_view); + + let v = outputs.view(); + v.a.push(1).unwrap(); + v.b.push(2).unwrap(); + v.c.push("three".to_string()).unwrap(); + v.d.push(vec![4u8, 5, 6]).unwrap(); + + assert_eq!(queue_set.take_typed_input::(0).pop(), Some(1)); + assert_eq!(queue_set.take_typed_input::(1).pop(), Some(2)); + assert_eq!(queue_set.take_typed_input::(2).pop(), Some("three".to_string())); + assert_eq!(queue_set.take_typed_input::>(3).pop(), Some(vec![4u8, 5, 6])); + } + + #[test] + fn unit_build_queues_yields_empty_set() { + let (queue_set, outputs_view) = + <() as StepOutputs>::build_queues(&[], &[], crate::builder::InstrumentationLevel::Off); + let outputs: OutputHandles<()> = OutputHandles::new(outputs_view); + outputs.noop(); + assert_eq!(queue_set.n_branches(), 0); + } + + #[test] + fn mark_all_drained_propagates_to_input_handle() { + let (specs, ordering) = count_specs(1, 4); + let (mut queue_set, outputs_view) = as StepOutputs>::build_queues( + &specs, + &ordering, + crate::builder::InstrumentationLevel::Off, + ); + let outputs: OutputHandles> = OutputHandles::new(outputs_view); + + let input = queue_set.take_typed_input::(0); + assert!(!input.is_drained()); + outputs.mark_all_drained(); + assert!(input.is_drained()); + } + + #[test] + fn ordered_branch_preserves_emission_order_through_typed_path() { + let specs = vec![QueueSpec::CountBounded { capacity: 8 }]; + let ordering = vec![BranchOrdering::ByOrdinal]; + let (mut queue_set, outputs_view) = as StepOutputs>::build_queues( + &specs, + &ordering, + crate::builder::InstrumentationLevel::Off, + ); + let outputs: OutputHandles> = OutputHandles::new(outputs_view); + + outputs.push(10).unwrap(); + outputs.push(20).unwrap(); + outputs.push(30).unwrap(); + + let input = queue_set.take_typed_input::(0); + assert_eq!(input.pop(), Some(10)); + assert_eq!(input.pop(), Some(20)); + assert_eq!(input.pop(), Some(30)); + } +} diff --git a/crates/fgumi-pipeline-core/src/header.rs b/crates/fgumi-pipeline-core/src/header.rs new file mode 100644 index 000000000..7109d33b7 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/header.rs @@ -0,0 +1,495 @@ +//! One-shot lazy SAM/BAM header handle. +//! +//! `HeaderHandle` lets a downstream sink open its output file before the +//! full output header is known, and consume the header at first record +//! time once an upstream step has produced it. The primary motivating +//! consumer is the writer downstream of `AlignAndMergeStep`: the +//! aligner's `@PG` (and any `@RG`/`@CO` lines it adds) are runtime +//! contributions that aren't available at `Pipeline::build` time. +//! +//! ### Contract +//! +//! - `set` and `poison` are one-shot — calling either after the handle +//! has been resolved returns `AlreadySetError` without altering state. +//! - `try_get` is non-blocking and dispatcher-friendly: a step that +//! needs the header but observes `None` returns `StepOutcome::NoProgress` +//! and is rescheduled. +//! - Cloning shares state (via `Arc>`), so multiple +//! readers see the same set-or-poison outcome. The cross-thread +//! happens-before guarantee comes from `OnceLock`'s internal +//! synchronization (`set` is Release, `get` is Acquire); the +//! `Arc>` wrapper exists only for shareability. +//! +//! ### Producer-side invariant +//! +//! Every consumer that polls `try_get` assumes the handle will +//! eventually resolve via `set` or `poison`. Producer steps owning a +//! `HeaderHandle` must poison their handle in their `Drop` impl as a +//! backstop against panics (or any exit path that bypasses normal +//! completion). The framework does not enforce this; it's a step-level +//! convention. A future revision may split the handle into typed +//! setter / reader halves so an orphaned setter is statically detectable. +//! +//! An orphaned setter — a producer that exits without resolving — does +//! **not** hang the run, provided the consumer follows the shape below. +//! Two framework backstops catch it, both pinned by tests in this module +//! (`orphaned_setter_*`): +//! +//! 1. **Producer left no items.** When a step reports `Finished` the +//! framework closes its output edges, so the consumer's +//! `is_input_drained()` arm fires and it can fail with a diagnostic. +//! 2. **Producer left items queued.** A header-blocked consumer pops +//! nothing, so `is_drained()` (which is `drained && empty`) stays +//! false and backstop 1 is unreachable — but those stranded items are +//! in flight, so the deadlock monitor sees non-zero `in_flight_bytes`, +//! classifies the run `Wedged` past the fatal timeout, and fails it +//! with `PipelineError::TimedOut`. +//! +//! Backstop 2 depends on the stranded items carrying real heap bytes: +//! queue accounting is `HeapSize::heap_size()`-only, so a zero-heap item +//! type leaves `in_flight_bytes` at 0, which classifies as `Starving` and +//! resets the stall clock forever. That gap is general to the monitor +//! rather than specific to headers; see the `in_flight_bytes` docs in +//! `builder.rs`. +//! +//! ### Consumer-side contract +//! +//! A consumer must never poll `try_get` in isolation. Pair the probe with +//! `is_input_drained()` and treat "input drained before the header +//! resolved" as an error — that arm is backstop 1, and without it a +//! consumer whose producer left no items polls `None` forever. + +use std::fmt; +use std::io; +use std::sync::{Arc, OnceLock}; + +use noodles::sam::Header; + +/// Shared, one-shot header slot. +/// +/// See module docs for the full contract. +#[derive(Clone, Default, Debug)] +pub struct HeaderHandle { + inner: Arc>>>, +} + +impl HeaderHandle { + /// Construct an empty handle. `try_get` returns `None` until `set` + /// or `poison` is called. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Construct a handle that already carries `header`. `try_get` + /// returns `Some(Ok(&header))` immediately. Convenience for call + /// sites that have an eager header today and want to opt into the + /// handle-typed sink API without changing behavior. + /// + /// # Panics + /// The internal `set` call is on a freshly-constructed cell and + /// is therefore unreachable as a failure — the `expect` is a + /// belt-and-braces assertion documenting the invariant. + #[must_use] + #[allow(clippy::missing_panics_doc)] // documented above; clippy + // doesn't see the # Panics section across this constructor's + // delegation to `set`. + pub fn from_header(header: Header) -> Self { + let handle = Self::new(); + handle.set(header).expect("fresh HeaderHandle accepts its first set"); + handle + } + + /// Resolve the handle to `header`. Returns `AlreadySetError` if the + /// handle was previously set or poisoned. + /// + /// # Errors + /// Returns `AlreadySetError` on the second and subsequent calls. + pub fn set(&self, header: Header) -> Result<(), AlreadySetError> { + self.inner.set(Ok(header)).map_err(|_| AlreadySetError) + } + + /// Resolve the handle to a failure. Subsequent `try_get` calls + /// surface the error. + /// + /// # Errors + /// Returns `AlreadySetError` on the second and subsequent calls. + pub fn poison(&self, error: io::Error) -> Result<(), AlreadySetError> { + self.inner.set(Err(Arc::new(error))).map_err(|_| AlreadySetError) + } + + /// Non-blocking handle probe. + /// + /// Returns: + /// - `None` if neither `set` nor `poison` has been called yet — the + /// caller should yield (e.g. return `StepOutcome::NoProgress`). + /// - `Some(Ok(&header))` if `set(header)` was called. + /// - `Some(Err(e))` if `poison(e)` was called. A fresh `io::Error` + /// is constructed on every call from the stored kind + message + /// (because `io::Error` is not `Clone`). **Note:** the + /// original error's `source()` chain and any structured + /// payload are not preserved — only kind + display string + /// round-trip. Producers that need to surface structured + /// diagnostic context should encode it into the display + /// string before poisoning the handle. + #[must_use] + pub fn try_get(&self) -> Option> { + self.inner.get().map(|stored| match stored { + Ok(header) => Ok(header), + Err(err) => Err(io::Error::new(err.kind(), err.to_string())), + }) + } + + /// `true` once `set` or `poison` has resolved this handle. + #[must_use] + pub fn is_set(&self) -> bool { + self.inner.get().is_some() + } +} + +/// Returned by `set` / `poison` when the handle was previously +/// resolved. +#[derive(Debug, Clone, Copy)] +pub struct AlreadySetError; + +impl fmt::Display for AlreadySetError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "HeaderHandle has already been set or poisoned") + } +} + +impl std::error::Error for AlreadySetError {} + +#[cfg(test)] +mod tests { + use super::*; + use std::thread; + + #[test] + fn new_handle_has_no_value() { + let h = HeaderHandle::new(); + assert!(!h.is_set()); + assert!(h.try_get().is_none()); + } + + #[test] + fn from_header_resolves_immediately() { + let h = HeaderHandle::from_header(Header::default()); + assert!(h.is_set()); + let got = h.try_get().expect("set").expect("ok"); + assert_eq!(got, &Header::default()); + } + + #[test] + fn set_then_try_get_returns_ok() { + let h = HeaderHandle::new(); + h.set(Header::default()).expect("first set"); + assert!(h.is_set()); + let got = h.try_get().expect("set").expect("ok"); + assert_eq!(got, &Header::default()); + } + + #[test] + fn poison_then_try_get_returns_err() { + let h = HeaderHandle::new(); + h.poison(io::Error::other("aligner crashed")).expect("first poison"); + assert!(h.is_set()); + let err = h.try_get().expect("set").expect_err("poisoned"); + assert_eq!(err.kind(), io::ErrorKind::Other); + assert_eq!(err.to_string(), "aligner crashed"); + } + + #[test] + fn second_set_returns_already_set() { + let h = HeaderHandle::new(); + h.set(Header::default()).expect("first"); + let err = h.set(Header::default()).expect_err("second set"); + let _ = err; // type asserted by the binding + } + + #[test] + fn set_then_poison_returns_already_set() { + let h = HeaderHandle::new(); + h.set(Header::default()).expect("first"); + let _err: AlreadySetError = + h.poison(io::Error::other("late")).expect_err("poison after set"); + } + + #[test] + fn poison_then_set_returns_already_set() { + let h = HeaderHandle::new(); + h.poison(io::Error::other("first")).expect("first"); + let _err: AlreadySetError = h.set(Header::default()).expect_err("set after poison"); + } + + #[test] + fn clones_share_state() { + let a = HeaderHandle::new(); + let b = a.clone(); + assert!(!a.is_set() && !b.is_set()); + a.set(Header::default()).expect("first"); + assert!(a.is_set() && b.is_set()); + assert!(b.try_get().expect("set").is_ok()); + } + + #[test] + fn poison_preserves_error_kind_and_message_across_calls() { + let h = HeaderHandle::new(); + h.poison(io::Error::new(io::ErrorKind::BrokenPipe, "stderr ring: foo")).unwrap(); + for _ in 0..3 { + let e = h.try_get().unwrap().unwrap_err(); + assert_eq!(e.kind(), io::ErrorKind::BrokenPipe); + assert_eq!(e.to_string(), "stderr ring: foo"); + } + } + + #[test] + fn set_from_another_thread_is_observable() { + let h = HeaderHandle::new(); + let h2 = h.clone(); + let join = thread::spawn(move || { + h2.set(Header::default()).expect("first set"); + }); + join.join().expect("thread join"); + assert!(h.is_set()); + assert!(h.try_get().unwrap().is_ok()); + } + + #[test] + fn from_header_then_set_returns_already_set() { + let h = HeaderHandle::from_header(Header::default()); + let _err: AlreadySetError = + h.set(Header::default()).expect_err("from_header consumes the slot"); + } + + #[test] + fn from_header_then_poison_returns_already_set() { + let h = HeaderHandle::from_header(Header::default()); + let _err: AlreadySetError = + h.poison(io::Error::other("late")).expect_err("from_header consumes the slot"); + } + + // ───────────────────────────────────────────────────────────────────── + // Orphaned-setter backstops. + // + // These pin the two escapes described in the module's "Producer-side + // invariant" section, so the claim there is tested rather than asserted. + // ───────────────────────────────────────────────────────────────────── + + /// Shared fixtures for the orphaned-setter tests: a producer that finishes + /// without ever resolving its handle, and a consumer shaped like the real + /// BGZF writer (holds records until the header resolves; treats "input + /// drained before the header arrived" as a hard error). + mod orphan { + use super::HeaderHandle; + use std::io; + + use crate::item::HeapSize; + use crate::outputs::Single; + use crate::queues::QueueSpec; + use crate::reorder::BranchOrdering; + use crate::step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}; + + /// Item with a real heap payload, so a `ByteBounded` transport holding + /// one reports non-zero `current_bytes` (the queue's accounting is + /// `heap_size()`-only — `size_of::()` is not counted). + pub struct Block { + pub payload: Vec, + } + impl HeapSize for Block { + fn heap_size(&self) -> usize { + self.payload.capacity() + } + } + + /// Producer that emits `remaining` blocks and then finishes, never + /// calling `set`/`poison` on the handle it owns. + #[derive(Clone)] + pub struct OrphanedSetter { + pub remaining: u32, + pub block_bytes: usize, + pub spec: QueueSpec, + pub _handle: HeaderHandle, + } + impl Step for OrphanedSetter { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "OrphanedSetter", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![self.spec], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if self.remaining == 0 { + // Finishes WITHOUT resolving `_handle` — the orphan case. + return Ok(StepOutcome::Finished); + } + self.remaining -= 1; + // Asserted, not discarded: backstop 2's whole premise is that + // these items are stranded IN the queue. A silently dropped push + // would leave `in_flight_bytes` at 0 and the test would pass for + // the wrong reason. Budgets below always admit every block. + ctx.outputs + .push(Block { payload: vec![0u8; self.block_bytes] }) + .map_err(|_| io::Error::other("orphan test budget must admit every block"))?; + Ok(StepOutcome::Progress) + } + } + + /// Consumer mirroring the real `WriteBgzfFile` sink. + #[derive(Clone)] + pub struct HeaderBlockedSink { + pub handle: HeaderHandle, + } + impl Step for HeaderBlockedSink { + type Input = Block; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "HeaderBlockedSink", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + // The header gates consumption: nothing may be written before + // it resolves, so a blocked sink pops nothing. + let header_ready = self.handle.try_get().transpose()?.is_some(); + if header_ready && ctx.input.pop().is_some() { + return Ok(StepOutcome::Progress); + } + if ctx.input.is_drained() { + if !header_ready { + return Err(io::Error::other( + "HeaderBlockedSink: input drained before HeaderHandle was resolved", + )); + } + return Ok(StepOutcome::Finished); + } + Ok(StepOutcome::NoProgress) + } + } + } + + /// Backstop 1 — orphaned setter that emitted **nothing**: the consumer's + /// `is_drained()` arm fires and the run fails fast with a diagnostic. + /// + /// `is_drained()` is `queue.is_drained() && queue.is_empty()`, so this arm + /// is only reachable when the orphaned producer left no items behind. That + /// precondition is what backstop 2 covers. + #[test] + fn orphaned_setter_with_empty_output_fails_the_run_fast() { + use std::sync::mpsc; + use std::time::Duration; + + use crate::queues::QueueSpec; + use crate::{Pipeline, PipelineConfig}; + + let handle = HeaderHandle::new(); + let for_run = handle.clone(); + let (tx, rx) = mpsc::channel(); + let join = thread::spawn(move || { + let builder = Pipeline::builder(); + builder + .chain(orphan::OrphanedSetter { + remaining: 0, + block_bytes: 0, + spec: QueueSpec::CountBounded { capacity: 8 }, + _handle: for_run.clone(), + }) + .chain(orphan::HeaderBlockedSink { handle: for_run }) + .into_sink_marker(); + let pipeline = builder.build().expect("pipeline build"); + let _ = tx.send(pipeline.run(PipelineConfig { threads: 2, ..Default::default() })); + }); + + let result = rx + .recv_timeout(Duration::from_secs(10)) + .expect("an orphaned setter with an empty output must not hang the run"); + join.join().expect("run thread panicked"); + + let err = result.expect_err("an unresolved HeaderHandle must fail the run"); + assert!( + format!("{err}").contains("input drained before HeaderHandle was resolved"), + "run must fail with the sink's drained-without-header error, got: {err}" + ); + assert!(!handle.is_set(), "nobody ever resolved the handle"); + } + + /// Backstop 2 — orphaned setter that emitted items: those items sit unpopped + /// (the blocked sink never consumes), so `is_drained()` stays false forever + /// and the drained arm above can never fire. The deadlock monitor is what + /// catches this: the stranded items are in flight on a `ByteBounded` edge, + /// so `in_flight_bytes` is non-zero, `classify_stall` returns `Wedged` past + /// the fatal timeout, and the run fails with `PipelineError::TimedOut`. + /// + /// The item must carry real heap bytes for that to hold — queue accounting + /// is `heap_size()`-only, so a zero-heap item type would leave + /// `in_flight_bytes` at 0, yielding `Starving`, which resets the stall clock + /// on every poll and leaves the wedge uncatchable. + #[test] + fn orphaned_setter_with_inflight_bytes_is_caught_by_the_deadlock_monitor() { + use std::sync::mpsc; + use std::time::Duration; + + use crate::queues::QueueSpec; + use crate::signal::PipelineError; + use crate::{Pipeline, PipelineConfig}; + + let handle = HeaderHandle::new(); + let for_run = handle.clone(); + let (tx, rx) = mpsc::channel(); + let join = thread::spawn(move || { + let builder = Pipeline::builder(); + builder + .chain(orphan::OrphanedSetter { + remaining: 4, + block_bytes: 4096, + spec: QueueSpec::ByteBounded { limit_bytes: 1 << 20 }, + _handle: for_run.clone(), + }) + .chain(orphan::HeaderBlockedSink { handle: for_run }) + .into_sink_marker(); + let pipeline = builder.build().expect("pipeline build"); + let stats = pipeline.stats(); + let _ = tx.send(pipeline.run(PipelineConfig { + threads: 2, + stats: Some(stats), + // warn at 1s, fatal at 1s * DEADLOCK_FATAL_MULTIPLE. + deadlock_timeout_secs: 1, + ..Default::default() + })); + }); + + let result = rx + .recv_timeout(Duration::from_secs(60)) + .expect("the deadlock monitor must fail a header-wedged run, not let it hang"); + join.join().expect("run thread panicked"); + + let err = result.expect_err("a header-wedged run must fail"); + assert!( + matches!(err, PipelineError::TimedOut { .. }), + "the wedge must surface as TimedOut, got: {err:?}" + ); + assert!(!handle.is_set(), "nobody ever resolved the handle"); + } + + #[test] + fn poison_visible_to_all_clones() { + let a = HeaderHandle::new(); + let b = a.clone(); + let c = a.clone(); + a.poison(io::Error::new(io::ErrorKind::BrokenPipe, "boom")).unwrap(); + for clone in [&a, &b, &c] { + let err = clone.try_get().expect("set").expect_err("poisoned"); + assert_eq!(err.kind(), io::ErrorKind::BrokenPipe); + assert_eq!(err.to_string(), "boom"); + } + } +} diff --git a/crates/fgumi-pipeline-core/src/held.rs b/crates/fgumi-pipeline-core/src/held.rs new file mode 100644 index 000000000..f4a2d3af0 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/held.rs @@ -0,0 +1,89 @@ +//! Generic held-item slot for non-blocking back-pressure. +//! +//! When a step can't push to a downstream queue (it's full), it stashes the +//! item in a `HeldSlot` and returns `StepOutcome::Progress`. The next +//! `try_run` call drains the held item before doing new work. +//! +//! Draining the held item before new work is a *step-author convention*, not +//! something `HeldSlot` enforces: the type only provides single-slot put/take +//! (with a double-put panic). The step is responsible for checking and draining +//! the slot first on each `try_run`. + +pub struct HeldSlot { + inner: Option, +} + +impl Default for HeldSlot { + fn default() -> Self { + Self { inner: None } + } +} + +impl HeldSlot { + #[must_use] + pub fn new() -> Self { + Self::default() + } + + #[must_use] + pub fn is_held(&self) -> bool { + self.inner.is_some() + } + + /// Stash an item. + /// + /// # Panics + /// + /// Panics if the slot is already occupied (a contract violation: callers + /// must drain via `take` before calling `put` again). This is a hard + /// `assert!` rather than `debug_assert!` because silently overwriting would + /// **drop the previously held item** — a record lost from the pipeline — + /// which must not pass unnoticed in release builds. The check is a single + /// predictable branch on the back-pressure path, so the cost is negligible. + pub fn put(&mut self, item: T) { + assert!(self.inner.is_none(), "HeldSlot already occupied"); + self.inner = Some(item); + } + + pub fn take(&mut self) -> Option { + self.inner.take() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn empty_slot_is_empty() { + let s: HeldSlot = HeldSlot::new(); + assert!(!s.is_held()); + } + + #[test] + fn put_then_take_round_trips() { + let mut s = HeldSlot::new(); + s.put(42_u32); + assert!(s.is_held()); + assert_eq!(s.take(), Some(42)); + assert!(!s.is_held()); + } + + #[test] + fn take_from_empty_returns_none() { + let mut s: HeldSlot = HeldSlot::new(); + assert_eq!(s.take(), None); + } + + #[test] + #[should_panic(expected = "HeldSlot already occupied")] + fn double_put_panics() { + // Pin the message, not just the unwind: `catch_unwind` accepts ANY panic, + // so a `put` that started failing for an unrelated reason would still pass + // while the double-put guard itself had silently gone. Must fire in debug + // AND release — overwriting a held item is silent record loss. + let mut s = HeldSlot::new(); + s.put(1_u32); + s.put(2_u32); + } +} diff --git a/crates/fgumi-pipeline-core/src/item.rs b/crates/fgumi-pipeline-core/src/item.rs new file mode 100644 index 000000000..438361c2d --- /dev/null +++ b/crates/fgumi-pipeline-core/src/item.rs @@ -0,0 +1,162 @@ +//! Item-level orthogonal traits used by the queue layer. +//! +//! These traits are independent of each other: +//! - [`HeapSize`] is required only for items in a `ByteBoundedQueue`. +//! - [`Ordered`] is required only for items routed through a `ReorderStage`. +//! +//! An item type may implement both, neither, or just one, depending on which +//! queues and operators it flows through. The framework checks bounds at +//! queue construction time (see [`crate::queues`] and [`crate::reorder`]). + +/// Approximate heap footprint, in bytes, of an item. +/// +/// Used by `ByteBoundedQueue` to enforce a memory budget rather than an +/// item-count budget. Items that hold variable-size buffers (BAM batches, +/// FASTQ batches, decompressed BGZF blocks) implement this manually. +/// +/// **Default impl.** The default method body returns `0`, so an empty +/// `impl HeapSize for MyType {}` opts in but reports zero bytes — every +/// item then fits in any byte-bounded queue regardless of the limit. Step +/// authors must override `heap_size` for any type whose actual footprint +/// matters. +/// +/// **No blanket impl.** We do *not* provide `impl HeapSize for T` so that +/// pushing a non-impl type into a `ByteBoundedQueue` is a compile error +/// rather than a silent zero-count. +pub trait HeapSize { + fn heap_size(&self) -> usize { + 0 + } +} + +/// Source-step input type. Carries no heap allocation. +impl HeapSize for () {} + +// Primitive impls for the integer/string types that flow through tests +// and pipeline plumbing. Production BAM/FASTQ types impl `HeapSize` +// explicitly with their actual heap footprint; primitives report zero +// (their `Vec` capacity, when applicable, is reported by the wrapping +// type's own impl). +impl HeapSize for u8 {} +impl HeapSize for u16 {} +impl HeapSize for u32 {} +impl HeapSize for u64 {} +impl HeapSize for usize {} +impl HeapSize for i8 {} +impl HeapSize for i16 {} +impl HeapSize for i32 {} +impl HeapSize for i64 {} +impl HeapSize for isize {} +impl HeapSize for bool {} +impl HeapSize for f32 {} +impl HeapSize for f64 {} +impl HeapSize for char {} + +impl HeapSize for String { + fn heap_size(&self) -> usize { + self.capacity() + } +} + +impl HeapSize for Vec { + fn heap_size(&self) -> usize { + self.capacity() * std::mem::size_of::() + + self.iter().map(HeapSize::heap_size).sum::() + } +} + +impl HeapSize for Option { + fn heap_size(&self) -> usize { + self.as_ref().map_or(0, HeapSize::heap_size) + } +} + +/// Producer-emitted serial ordinal. +/// +/// Items routed through a `ReorderStage` carry a monotonically-increasing +/// ordinal assigned by the producer step at push time (the framework +/// allocates from a per-branch `AtomicU64`). The reorder stage uses this +/// ordinal to deliver items to the consumer in producer-emitted order +/// regardless of inter-thread arrival skew. +/// +/// Steps don't implement this directly. The framework wraps the user's +/// pushed item in an internal `Sequenced` newtype that carries the +/// ordinal alongside the item; `Sequenced` impls `Ordered`. See the +/// reorder stage for the wrapper type. +pub trait Ordered { + fn ordinal(&self) -> u64; +} + +#[cfg(test)] +mod tests { + use super::*; + + struct WithHeap(Vec); + impl HeapSize for WithHeap { + fn heap_size(&self) -> usize { + self.0.len() + } + } + + struct WithoutHeap; + impl HeapSize for WithoutHeap {} // explicit empty impl: opts in, default returns 0 + + struct WithOrdinal { + o: u64, + } + impl Ordered for WithOrdinal { + fn ordinal(&self) -> u64 { + self.o + } + } + + #[test] + fn heap_size_default_is_zero() { + assert_eq!(WithoutHeap.heap_size(), 0); + } + + #[test] + fn heap_size_override_works() { + assert_eq!(WithHeap(vec![0; 1024]).heap_size(), 1024); + } + + #[test] + fn ordered_returns_ordinal() { + assert_eq!(WithOrdinal { o: 42 }.ordinal(), 42); + } + + /// A `ByteBoundedQueue` charges each item its `heap_size()`, so the blanket + /// impls decide how much budget a queue of owned data actually accounts + /// for. They report *capacity*, not length — an over-allocated buffer costs + /// what it reserved, because that is what the process is holding. + #[test] + fn string_heap_size_is_its_capacity() { + let mut s = String::with_capacity(64); + s.push_str("abc"); + assert_eq!(s.heap_size(), 64, "capacity, not the 3 bytes in use"); + assert_eq!(String::new().heap_size(), 0, "an unallocated String costs nothing"); + } + + #[test] + fn vec_heap_size_counts_its_buffer_and_its_elements() { + // Flat elements: just the buffer. + let mut flat: Vec = Vec::with_capacity(8); + flat.push(1); + assert_eq!(flat.heap_size(), 8 * std::mem::size_of::()); + + // Nested elements add their own heap. Without the per-element sum a + // queue of `Vec>` would be charged only for the outer spine and + // could blow well past its byte budget. + let nested: Vec = vec![WithHeap(vec![0; 100]), WithHeap(vec![0; 200])]; + let spine = nested.capacity() * std::mem::size_of::(); + assert_eq!(nested.heap_size(), spine + 300); + + assert_eq!(Vec::::new().heap_size(), 0, "an unallocated Vec costs nothing"); + } + + #[test] + fn option_heap_size_delegates_to_the_payload_and_none_is_free() { + assert_eq!(Some(WithHeap(vec![0; 77])).heap_size(), 77); + assert_eq!(None::.heap_size(), 0); + } +} diff --git a/crates/fgumi-pipeline-core/src/lib.rs b/crates/fgumi-pipeline-core/src/lib.rs new file mode 100644 index 000000000..df0e62461 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/lib.rs @@ -0,0 +1,62 @@ +//! Core types and traits for the typed-step pipeline framework. +//! +//! A pipeline is a graph of typed [`Step`]s joined by bounded [`queues`] and run +//! by a work-stealing worker pool. This crate carries only the framework +//! primitives — the step traits, the queues, the [`reorder`] stage that restores +//! input order across parallel branches, and the [`runtime`] that schedules and +//! drives them. Nothing here reads or writes sequencing data, so the dependency +//! graph stays light (`ahash` / `crossbeam-queue` / `parking_lot` / `log`, plus +//! `anyhow` for the one [`FinalizeHook::finalize`] return type and one +//! `noodles::sam` type for the shared header handle) and the crate compiles +//! fast in isolation. +//! +//! The concrete steps that do the I/O and the computation live elsewhere and +//! plug in by implementing [`Step`] (one input) or [`Step2`] (two inputs); +//! [`PipelineBuilder`] wires them into a [`Pipeline`]. + +#![deny(unsafe_code)] + +pub mod builder; +pub mod erased; +pub mod finalize; +pub mod handles; +pub mod header; +pub mod held; +pub mod item; +pub mod liveness; +pub mod outputs; +pub mod queues; +pub mod reorder; +pub mod runtime; +pub mod signal; +pub mod step; +pub mod topology; + +#[cfg(test)] +mod tests; + +pub use builder::{ + BuildError, Chain, DEFAULT_DEADLOCK_TIMEOUT_SECS, InstrumentationLevel, MultiChain2, + MultiChain2Ordered, MultiChain3, MultiChain4, Pipeline, PipelineBuilder, PipelineConfig, +}; +pub use erased::{ErasedStep, ErasedStepCtx, TypedStep, TypedStep2}; +pub use finalize::FinalizeHook; +pub use handles::{ + BranchInputHandle, HeldRetry, OutputQueueSet, Tuple2View, Tuple3View, Tuple4View, + TwoInputHandles, Unpushed, +}; +pub use header::{AlreadySetError, HeaderHandle}; +pub use held::HeldSlot; +pub use item::{HeapSize, Ordered}; +pub use liveness::LivenessCounter; +pub use outputs::{ + MAX_ARITY, OrderedBytesSingle, OrderedBytesTuple2, OrderedBytesTuple3, Single, StepOutputs, +}; +pub use queues::{ByteBoundedQueue, CountBoundedQueue, ItemQueue, QueueSpec, UnboundedQueue}; +pub use reorder::{BranchOrdering, ReorderStage, Sequenced}; +pub use signal::{CancelHandle, PipelineError, PipelineSignal}; +pub use step::{ + Affinity, DetachedGroup, InputHandle, OutputHandles, OutputsViewAny, Step, Step2, StepCtx, + StepCtx2, StepKind, StepOutcome, StepProfile, +}; +pub use topology::{BranchIdx, ChainGraph, StepIdx}; diff --git a/crates/fgumi-pipeline-core/src/liveness.rs b/crates/fgumi-pipeline-core/src/liveness.rs new file mode 100644 index 000000000..11bf598a2 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/liveness.rs @@ -0,0 +1,210 @@ +//! A cheap, always-on liveness signal for the deadlock monitor. +//! +//! # Why this exists separately from [`PipelineStats`](crate::runtime::stats::PipelineStats) +//! +//! The deadlock monitor needs to answer exactly one question: *has anything in +//! the pipeline made progress since I last looked?* It compares one total +//! against the previous total and cares about nothing else. +//! +//! It used to get that total from `PipelineStats`, which coupled the monitor to +//! full instrumentation — and full instrumentation is deliberately not free. +//! `dispatch_one_step` times each dispatch with `Instant::now()` (~20–50 ns on +//! Apple Silicon, ~50–100 ns on `x86_64`) and gates that on `stats.is_some()` +//! precisely to keep the uninstrumented path zero-cost. So the monitor could +//! only be armed by paying a per-dispatch timing cost on every run, which is why +//! it shipped disarmed by default and every wedge hung silently instead. +//! +//! This type breaks that coupling: liveness is a counter, profiling is a +//! separate opt-in. The monitor can therefore run with no `PipelineStats` +//! attached at all — a stats handle now only adds the per-step snapshot to the +//! stall report, and its absence costs the diagnostic, not the detection. +//! +//! That removes the *cost* argument for shipping disarmed; it does not by +//! itself arm anything. `PipelineConfig::default()` still sets +//! `deadlock_timeout_secs: 0`, because an armed monitor additionally requires +//! every output transport to be `ByteBounded` (see +//! [`PipelineError::MonitorBlindTransport`](crate::signal::PipelineError)) and +//! chains such as `Process2` use `CountBounded` today. +//! +//! Both of those are properties of the **scheduled** path. A run that fuses to +//! a single thread returns before the transport check and before the monitor +//! spawns, so it neither gains the monitor nor is rejected for a byte-blind +//! edge; it is bounded by the fused path's own stall budget instead. +//! +//! # Why the counter is sharded +//! +//! A single shared `AtomicU64` bumped by every worker on every productive +//! dispatch is a false-sharing hotspot: each increment is a read-modify-write on +//! one cache line, so N workers ping-pong that line between cores and the cost +//! grows with thread count — exactly the wrong shape, since more threads is when +//! the monitor matters most. +//! +//! Each worker therefore owns a slot padded to its own cache line, so a bump is +//! normally an uncontended increment on a line no other worker touches. (The one +//! exception is a dedicated driver thread, which reuses `worker_slot` 0 and so +//! shares slot 0 with pool worker 0; the bump is an atomic `fetch_add`, so a +//! coincident increment is counted correctly rather than lost — it only costs the +//! two threads a shared cache line, not accuracy.) The monitor sums +//! the slots, which it does once per poll interval (seconds), so the read side's +//! cost is irrelevant. +//! +//! The sum is not a synchronized snapshot — slots are read one at a time under +//! `Relaxed`, so the total may mix values from slightly different instants. That +//! is fine for the only question asked of it: a torn total still *changes* when +//! any worker makes progress, and monotonicity per slot means it can never +//! spuriously appear frozen while work is happening. + +use std::sync::atomic::{AtomicU64, Ordering}; + +/// Cache-line padding. 128 rather than 64 because Apple Silicon and some `x86_64` +/// prefetchers pair adjacent 64-byte lines, so a 64-byte stride can still share +/// a coherence unit. +const CACHE_LINE_BYTES: usize = 128; + +#[repr(align(128))] +struct PaddedCounter(AtomicU64); + +const _: () = assert!(std::mem::size_of::() == CACHE_LINE_BYTES); + +/// Per-worker progress counters, summed by the deadlock monitor. +pub struct LivenessCounter { + slots: Box<[PaddedCounter]>, + /// `slots.len() - 1`, valid because the slot count is always a power of two. + /// The bump path masks with this instead of `%`: a runtime modulo is an + /// integer division (tens of cycles on both `aarch64` and `x86_64`) on a path + /// that runs once per productive dispatch, which measured as a double-digit + /// percentage of dispatch cost. + mask: usize, +} + +impl LivenessCounter { + /// One slot per worker. `n_workers` must cover every index passed to + /// [`Self::bump`]; indices are taken modulo the slot count so an unexpected + /// worker id degrades to sharing a slot rather than panicking on the hot + /// path. + #[must_use] + pub fn new(n_workers: usize) -> Self { + // Round up to a power of two so the bump path can mask instead of + // dividing. The slack is a few unused cache lines — irrelevant next to + // removing a division from a per-dispatch path. + let n = n_workers.max(1).next_power_of_two(); + Self { slots: (0..n).map(|_| PaddedCounter(AtomicU64::new(0))).collect(), mask: n - 1 } + } + + /// Record one unit of progress for `worker`. + /// + /// `Relaxed` is sufficient: the value only has to change, and the monitor + /// re-reads it on a multi-second cadence, so no ordering relative to other + /// memory is required. + #[inline] + pub fn bump(&self, worker: usize) { + // Mask, not modulo — see `mask`. The slot count is a power of two, so + // this is exact for in-range ids and wraps harmlessly for anything else. + let slot = &self.slots[worker & self.mask]; + slot.0.fetch_add(1, Ordering::Relaxed); + } + + /// Sum every worker's progress. Called once per monitor poll. + #[must_use] + pub fn total(&self) -> u64 { + self.slots.iter().map(|s| s.0.load(Ordering::Relaxed)).sum() + } + + /// Number of slots. This is the requested worker count rounded up to a + /// power of two, not the requested count itself. + #[must_use] + pub fn len(&self) -> usize { + self.slots.len() + } + + /// Always false — `new` clamps to at least one slot. Present because clippy + /// requires it alongside `len`. + #[must_use] + pub fn is_empty(&self) -> bool { + false + } +} + +impl std::fmt::Debug for LivenessCounter { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // Hand-written rather than derived: the per-slot atomics are an + // implementation detail, and the only useful facts are how many slots + // there are and what they currently sum to. + f.debug_struct("LivenessCounter") + .field("slots", &self.slots.len()) + .field("mask", &self.mask) + .field("total", &self.total()) + .finish() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + + #[test] + fn total_sums_every_slot() { + let counter = LivenessCounter::new(4); + counter.bump(0); + counter.bump(1); + counter.bump(1); + counter.bump(3); + assert_eq!(counter.total(), 4); + } + + #[test] + fn an_out_of_range_worker_wraps_rather_than_panicking() { + // The hot path must never panic on an unexpected worker id; sharing a + // slot only costs contention, and the total stays correct. + let counter = LivenessCounter::new(2); + counter.bump(7); + assert_eq!(counter.total(), 1); + } + + #[test] + fn a_single_worker_pipeline_still_gets_a_slot() { + let counter = LivenessCounter::new(0); + assert_eq!(counter.len(), 1, "the slot count is clamped to at least one"); + counter.bump(0); + assert_eq!(counter.total(), 1); + } + + /// Every increment must be observed — the point of the type is that a frozen + /// total means a frozen pipeline, so a lost bump would be a false wedge + /// report. + #[test] + fn concurrent_bumps_are_all_counted() { + const WORKERS: usize = 8; + const PER_WORKER: u64 = 10_000; + + let counter = Arc::new(LivenessCounter::new(WORKERS)); + let handles: Vec<_> = (0..WORKERS) + .map(|w| { + let counter = Arc::clone(&counter); + std::thread::spawn(move || { + for _ in 0..PER_WORKER { + counter.bump(w); + } + }) + }) + .collect(); + for h in handles { + h.join().expect("worker thread joins"); + } + assert_eq!(counter.total(), WORKERS as u64 * PER_WORKER); + } + + /// Each slot must sit on its own cache line — the whole reason for sharding. + #[test] + fn slots_are_cache_line_separated() { + let counter = LivenessCounter::new(4); + let a = std::ptr::from_ref(&counter.slots[0]) as usize; + let b = std::ptr::from_ref(&counter.slots[1]) as usize; + assert_eq!( + b - a, + CACHE_LINE_BYTES, + "adjacent slots must be a full cache line apart or they false-share", + ); + } +} diff --git a/crates/fgumi-pipeline-core/src/outputs.rs b/crates/fgumi-pipeline-core/src/outputs.rs new file mode 100644 index 000000000..f84edbc45 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/outputs.rs @@ -0,0 +1,525 @@ +//! `StepOutputs`: type-level description of a step's outputs. +//! +//! Single-output steps declare `type Outputs = Single;`. +//! Multi-output steps declare `type Outputs = (A, B, C);` (positional access). +//! Sinks declare `type Outputs = ();`. +//! +//! **Maximum arity: 4.** [`MAX_ARITY`] is the largest tuple shape +//! that has a `StepOutputs` impl. A user step that declares +//! `type Outputs = (A, B, C, D, E);` (arity 5) will fail to compile with a +//! "trait `StepOutputs` is not implemented" error. Larger arities require +//! adding the additional impls in this file plus the matching +//! `build_tupleN_queues` constructor in `handles.rs`. +//! +//! Each `StepOutputs` impl provides: +//! - `arity()` — number of independent output channels. +//! - `build_queues(specs, ordering, level) -> (OutputQueueSet, OutputsViewAny)` — +//! constructs the typed queues + reorder operators (where applicable) +//! and the type-erased view the framework stores. Implemented in +//! `handles.rs` (one impl per arity). +//! +//! The `specs` and `ordering` slices both have length `arity()`, sourced +//! from `StepProfile::output_queues` and `StepProfile::branch_ordering`. + +use std::marker::PhantomData; + +use super::handles::OutputQueueSet; +use super::item::{HeapSize, Ordered}; +use super::queues::QueueSpec; +use super::reorder::BranchOrdering; +use super::step::OutputsViewAny; + +/// The largest tuple `Outputs` shape that has a `StepOutputs` impl. +/// `Single` and `()` are also valid output shapes. +pub const MAX_ARITY: usize = 4; + +/// Marker trait for a step's `Outputs` associated type. +pub trait StepOutputs: Send + 'static { + /// Number of independent output channels. + fn arity() -> usize; + + /// Construct the typed queues + type-erased view for this Outputs shape. + /// + /// `specs[i]` and `ordering[i]` together describe branch `i`. Panics if + /// `specs.len() != Self::arity()` or `ordering.len() != Self::arity()` + /// (the builder ensures this invariant before calling). + fn build_queues( + specs: &[QueueSpec], + ordering: &[BranchOrdering], + level: crate::builder::InstrumentationLevel, + ) -> (OutputQueueSet, OutputsViewAny); + + /// Mark all output branches drained, dispatching through the typed + /// `OutputHandles::mark_all_drained` method. Implemented for + /// each per-arity variant in this file; called by `TypedStep` from + /// `mark_outputs_drained` in the worker loop's drain-propagation path. + fn mark_all_drained(handles: &super::step::OutputHandles) + where + Self: Sized; +} + +/// Wrapper for single-output steps. `type Outputs = Single;` +pub struct Single(PhantomData T>); + +impl StepOutputs for Single { + #[inline] + fn arity() -> usize { + 1 + } + + fn build_queues( + specs: &[QueueSpec], + ordering: &[BranchOrdering], + level: crate::builder::InstrumentationLevel, + ) -> (OutputQueueSet, OutputsViewAny) { + super::handles::build_single_queues::(specs, ordering, level) + } + + fn mark_all_drained(handles: &super::step::OutputHandles) { + handles.mark_all_drained(); + } +} + +/// Wrapper for single-output steps where the item type is heap-aware AND +/// carries its own ordinal. The canonical Phase 3 BAM step output shape: +/// every BAM step's output type impls both `HeapSize` (for byte-bounded +/// queues) and `Ordered` (so a `batch_serial: u64` field carries record- +/// read order through every Parallel transform). +/// +/// Steps that need byte-bounded outputs but don't have item-carried serials +/// use `Single` with `BranchOrdering::None` or `ByOrdinal`. Steps that +/// need item-ordinal ordering but not byte-bounded queues are uncommon +/// (and currently not supported as a separate shape) — they can use this +/// shape with `QueueSpec::CountBounded` since `T: HeapSize` is required by +/// the bound but not consulted by count-bounded queues. +pub struct OrderedBytesSingle(PhantomData T>); + +impl StepOutputs for OrderedBytesSingle { + #[inline] + fn arity() -> usize { + 1 + } + + fn build_queues( + specs: &[QueueSpec], + ordering: &[BranchOrdering], + level: crate::builder::InstrumentationLevel, + ) -> (OutputQueueSet, OutputsViewAny) { + super::handles::build_single_queues_ordered_bytes::(specs, ordering, level) + } + + fn mark_all_drained(handles: &super::step::OutputHandles) { + handles.mark_all_drained(); + } +} + +impl StepOutputs for (A, B) { + #[inline] + fn arity() -> usize { + 2 + } + + fn build_queues( + specs: &[QueueSpec], + ordering: &[BranchOrdering], + level: crate::builder::InstrumentationLevel, + ) -> (OutputQueueSet, OutputsViewAny) { + super::handles::build_tuple2_queues::(specs, ordering, level) + } + + fn mark_all_drained(handles: &super::step::OutputHandles) { + handles.mark_all_drained(); + } +} + +/// Ordered + byte-bounded tuple-2 outputs. Use when a step fans out to +/// two branches that both need `BranchOrdering::ByItemOrdinal` + +/// byte-bounded queues — e.g., `filter` emitting kept and rejected +/// records as parallel ordered streams that downstream +/// `BgzfCompress` / `WriteBgzfFile` sinks can consume. +/// +/// The plain `(A, B)` `StepOutputs` impl only bounds `A: HeapSize` +/// (suitable for `Process2`'s `BranchOrdering::None` case). To use +/// `ByItemOrdinal` on either branch, both branches must satisfy +/// `Ordered + HeapSize` — that's what this shape encodes at the type +/// level. +pub struct OrderedBytesTuple2(PhantomData (A, B)>) +where + A: Send + HeapSize + Ordered + 'static, + B: Send + HeapSize + Ordered + 'static; + +impl StepOutputs for OrderedBytesTuple2 +where + A: Send + HeapSize + Ordered + 'static, + B: Send + HeapSize + Ordered + 'static, +{ + #[inline] + fn arity() -> usize { + 2 + } + + fn build_queues( + specs: &[QueueSpec], + ordering: &[BranchOrdering], + level: crate::builder::InstrumentationLevel, + ) -> (OutputQueueSet, OutputsViewAny) { + super::handles::build_tuple2_queues_ordered_bytes::(specs, ordering, level) + } + + fn mark_all_drained(handles: &super::step::OutputHandles) { + handles.mark_all_drained(); + } +} + +/// Ordered + byte-bounded tuple-3 outputs. The 3-branch analog of +/// [`OrderedBytesTuple2`]: use when a step fans out to three branches that all +/// need `BranchOrdering::ByItemOrdinal` + byte-bounded queues — e.g. paired +/// FASTQ output splitting one record batch into R1 / R2 / other byte streams, +/// each feeding an ordered `WriteRawFile` sink. +/// +/// The plain `(A, B, C)` `StepOutputs` impl only bounds each branch `HeapSize`. +/// To use `ByItemOrdinal` on any branch, all three must satisfy +/// `Ordered + HeapSize` — that's what this shape encodes at the type level. +#[allow(clippy::type_complexity)] +pub struct OrderedBytesTuple3(PhantomData (A, B, C)>) +where + A: Send + HeapSize + Ordered + 'static, + B: Send + HeapSize + Ordered + 'static, + C: Send + HeapSize + Ordered + 'static; + +impl StepOutputs for OrderedBytesTuple3 +where + A: Send + HeapSize + Ordered + 'static, + B: Send + HeapSize + Ordered + 'static, + C: Send + HeapSize + Ordered + 'static, +{ + #[inline] + fn arity() -> usize { + 3 + } + + fn build_queues( + specs: &[QueueSpec], + ordering: &[BranchOrdering], + level: crate::builder::InstrumentationLevel, + ) -> (OutputQueueSet, OutputsViewAny) { + super::handles::build_tuple3_queues_ordered_bytes::(specs, ordering, level) + } + + fn mark_all_drained(handles: &super::step::OutputHandles) { + handles.mark_all_drained(); + } +} + +impl StepOutputs for (A, B, C) +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, + C: Send + HeapSize + 'static, +{ + #[inline] + fn arity() -> usize { + 3 + } + + fn build_queues( + specs: &[QueueSpec], + ordering: &[BranchOrdering], + level: crate::builder::InstrumentationLevel, + ) -> (OutputQueueSet, OutputsViewAny) { + super::handles::build_tuple3_queues::(specs, ordering, level) + } + + fn mark_all_drained(handles: &super::step::OutputHandles) { + handles.mark_all_drained(); + } +} + +impl StepOutputs for (A, B, C, D) +where + A: Send + HeapSize + 'static, + B: Send + HeapSize + 'static, + C: Send + HeapSize + 'static, + D: Send + HeapSize + 'static, +{ + #[inline] + fn arity() -> usize { + 4 + } + + fn build_queues( + specs: &[QueueSpec], + ordering: &[BranchOrdering], + level: crate::builder::InstrumentationLevel, + ) -> (OutputQueueSet, OutputsViewAny) { + super::handles::build_tuple4_queues::(specs, ordering, level) + } + + fn mark_all_drained(handles: &super::step::OutputHandles) { + handles.mark_all_drained(); + } +} + +impl StepOutputs for () { + #[inline] + fn arity() -> usize { + 0 + } + + fn build_queues( + specs: &[QueueSpec], + ordering: &[BranchOrdering], + level: crate::builder::InstrumentationLevel, + ) -> (OutputQueueSet, OutputsViewAny) { + super::handles::build_unit_queues(specs, ordering, level) + } + + fn mark_all_drained(handles: &super::step::OutputHandles) { + handles.mark_all_drained(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::step::InputHandle; + + #[test] + fn single_arity_is_one() { + assert_eq!( as StepOutputs>::arity(), 1); + } + + #[test] + fn tuple_2_arity_is_two() { + assert_eq!(<(u32, u64) as StepOutputs>::arity(), 2); + } + + #[test] + fn tuple_3_arity_is_three() { + assert_eq!(<(u32, u64, String) as StepOutputs>::arity(), 3); + } + + #[derive(Clone, Copy)] + struct OrdU64(u64); + impl crate::item::HeapSize for OrdU64 {} + impl crate::item::Ordered for OrdU64 { + fn ordinal(&self) -> u64 { + self.0 + } + } + + #[test] + fn ordered_bytes_tuple_3_arity_is_three() { + // OrderedBytesTuple3 is the ordered + byte-bounded 3-way fan-out shape. + // Its three branches carry Ordered + HeapSize items (here the u64/u32 + // stand-ins just need to satisfy the bounds at the type level). + fn assert_arity() -> usize { + O::arity() + } + assert_eq!(assert_arity::>(), 3); + } + + #[test] + fn tuple_4_arity_is_four() { + assert_eq!(<(u32, u64, String, Vec) as StepOutputs>::arity(), 4); + } + + #[test] + fn unit_arity_is_zero() { + assert_eq!(<() as StepOutputs>::arity(), 0); + } + + #[test] + fn ordered_bytes_tuple_2_arity_is_two() { + fn assert_arity() -> usize { + O::arity() + } + assert_eq!(assert_arity::>(), 2); + } + + // ───────────────────────────────────────────────────────────────────────── + // `build_queues` / `mark_all_drained` round-trips, one test per output shape. + // + // Every branch a shape declares must come back from `build_queues` as its + // own live edge: `n_branches()` matches `arity()`, each branch's input handle + // transports the item type its position declares, each starts OPEN, and + // `mark_all_drained` closes all of them. The open-before / closed-after pair + // is what makes these discriminating — asserting only the closed-after state + // would pass just as well if `build_queues` returned branches already closed + // and `mark_all_drained` did nothing. + // + // These stay separate tests rather than an `#[rstest]` case table: each shape + // is a distinct type with distinct per-branch item types, so the cases cannot + // share a function signature. Every shape with a `StepOutputs` impl has one — + // `Single`, `OrderedBytesSingle`, the 2-/3-/4-tuples, `OrderedBytesTuple2`, + // `OrderedBytesTuple3`, and `()`. Adding a shape without adding its + // round-trip leaves its `build_queues` / `mark_all_drained` pair unexercised. + // ───────────────────────────────────────────────────────────────────────── + + fn specs(n: usize) -> Vec { + vec![QueueSpec::CountBounded { capacity: 4 }; n] + } + + fn orderings(n: usize) -> Vec { + vec![BranchOrdering::None; n] + } + + /// Build a shape's queues and hand back the branch set plus a typed + /// `OutputHandles` view, the pair the runtime hands to a step. + fn build(n: usize) -> (OutputQueueSet, crate::step::OutputHandles) { + let (queues, view) = + O::build_queues(&specs(n), &orderings(n), crate::builder::InstrumentationLevel::Off); + (queues, crate::step::OutputHandles::::new(view)) + } + + #[test] + fn single_builds_one_branch_and_drains() { + type Shape = Single; + let (mut queues, handles) = build::(1); + assert_eq!(queues.n_branches(), 1, "one branch per declared output"); + + let a = queues.take_typed_input::(0); + assert!(!a.is_drained(), "branch starts open"); + + ::mark_all_drained(&handles); + assert!(a.is_drained(), "mark_all_drained closes branch 0"); + } + + /// The only shape whose `build_queues` routes to + /// `build_single_queues_ordered_bytes`, so without this its drain path has no + /// coverage here at all. + #[test] + fn ordered_bytes_single_builds_one_branch_and_drains() { + type Shape = OrderedBytesSingle; + let (mut queues, handles) = build::(1); + assert_eq!(queues.n_branches(), 1, "one branch per declared output"); + + let a = queues.take_typed_input::(0); + assert!(!a.is_drained(), "branch starts open"); + + ::mark_all_drained(&handles); + assert!(a.is_drained(), "mark_all_drained closes branch 0"); + } + + #[test] + fn tuple_2_builds_two_independent_branches() { + type Shape = (u32, u64); + let (mut queues, handles) = build::(2); + assert_eq!(queues.n_branches(), 2, "one branch per declared output"); + + let a = queues.take_typed_input::(0); + let b = queues.take_typed_input::(1); + assert!(!a.is_drained(), "branches start open"); + assert!(!b.is_drained()); + + ::mark_all_drained(&handles); + assert!(a.is_drained(), "mark_all_drained closes branch 0"); + assert!(b.is_drained(), "mark_all_drained closes branch 1"); + } + + #[test] + fn tuple_3_builds_three_independent_branches() { + let (mut queues, handles) = build::<(u32, u64, String)>(3); + assert_eq!(queues.n_branches(), 3, "one branch per declared output"); + + // Each branch's input handle must downcast to that position's type — + // a mis-wired builder would panic here or hand back the wrong branch. + let a = queues.take_typed_input::(0); + let b = queues.take_typed_input::(1); + let c = queues.take_typed_input::(2); + assert!(!a.is_drained(), "branches start open"); + assert!(!b.is_drained()); + assert!(!c.is_drained()); + + <(u32, u64, String) as StepOutputs>::mark_all_drained(&handles); + assert!(a.is_drained(), "mark_all_drained closes branch 0"); + assert!(b.is_drained(), "mark_all_drained closes branch 1"); + assert!(c.is_drained(), "mark_all_drained closes branch 2"); + } + + #[test] + fn tuple_4_builds_four_independent_branches() { + let (mut queues, handles) = build::<(u32, u64, String, Vec)>(4); + assert_eq!(queues.n_branches(), 4, "one branch per declared output"); + + let a = queues.take_typed_input::(0); + let b = queues.take_typed_input::(1); + let c = queues.take_typed_input::(2); + let d = queues.take_typed_input::>(3); + assert!(!a.is_drained(), "branches start open"); + assert!(!b.is_drained()); + assert!(!c.is_drained()); + assert!(!d.is_drained()); + + <(u32, u64, String, Vec) as StepOutputs>::mark_all_drained(&handles); + assert!(a.is_drained(), "mark_all_drained closes branch 0"); + assert!(b.is_drained(), "mark_all_drained closes branch 1"); + assert!(c.is_drained(), "mark_all_drained closes branch 2"); + assert!(d.is_drained(), "mark_all_drained closes branch 3"); + } + + #[test] + fn ordered_bytes_tuple_2_builds_two_independent_branches() { + type Shape = OrderedBytesTuple2; + let (mut queues, handles) = build::(2); + assert_eq!(queues.n_branches(), 2, "one branch per declared output"); + + let a = queues.take_typed_input::(0); + let b = queues.take_typed_input::(1); + assert!(!a.is_drained(), "branches start open"); + assert!(!b.is_drained()); + + // Identity, not just type. Both branches carry the SAME item type, so the + // per-position downcast cannot tell them apart, and `mark_all_drained` + // closes every branch — a builder that aliased branch 0's queue onto both + // positions would satisfy every other assertion here. Route a distinct + // value through each and read it back off its own edge. + let view = handles.view(); + view.a.push(OrdU64(10)).expect("branch 0 accepts one item"); + view.b.push(OrdU64(11)).expect("branch 1 accepts one item"); + assert_eq!(a.pop().map(|v| v.0), Some(10), "branch 0 is its own edge"); + assert_eq!(b.pop().map(|v| v.0), Some(11), "branch 1 is its own edge"); + + ::mark_all_drained(&handles); + assert!(a.is_drained(), "mark_all_drained closes branch 0"); + assert!(b.is_drained(), "mark_all_drained closes branch 1"); + } + + #[test] + fn ordered_bytes_tuple_3_builds_three_independent_branches() { + type Shape = OrderedBytesTuple3; + let (mut queues, handles) = build::(3); + assert_eq!(queues.n_branches(), 3, "one branch per declared output"); + + let a = queues.take_typed_input::(0); + let b = queues.take_typed_input::(1); + let c = queues.take_typed_input::(2); + assert!(!a.is_drained(), "branches start open"); + assert!(!b.is_drained()); + assert!(!c.is_drained()); + + // Identity, not just type — see the tuple-2 test above for why the + // same-item-type shapes need this. + let view = handles.view(); + view.a.push(OrdU64(10)).expect("branch 0 accepts one item"); + view.b.push(OrdU64(11)).expect("branch 1 accepts one item"); + view.c.push(OrdU64(12)).expect("branch 2 accepts one item"); + assert_eq!(a.pop().map(|v| v.0), Some(10), "branch 0 is its own edge"); + assert_eq!(b.pop().map(|v| v.0), Some(11), "branch 1 is its own edge"); + assert_eq!(c.pop().map(|v| v.0), Some(12), "branch 2 is its own edge"); + + ::mark_all_drained(&handles); + assert!(a.is_drained(), "mark_all_drained closes branch 0"); + assert!(b.is_drained(), "mark_all_drained closes branch 1"); + assert!(c.is_drained(), "mark_all_drained closes branch 2"); + } + + /// A sink declares no outputs, so its queue set is empty and + /// `mark_all_drained` is a no-op rather than a panic. + #[test] + fn unit_builds_no_branches_and_drains_without_panicking() { + let (queues, handles) = build::<()>(0); + assert_eq!(queues.n_branches(), 0, "a sink owns no output branches"); + <() as StepOutputs>::mark_all_drained(&handles); + } +} diff --git a/crates/fgumi-pipeline-core/src/queues.rs b/crates/fgumi-pipeline-core/src/queues.rs new file mode 100644 index 000000000..faa5f0931 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/queues.rs @@ -0,0 +1,815 @@ +//! Transport-layer queue trait + three concrete impls. +//! +//! Concerns: pure transport (push/pop, drained signal). **Not** ordering — +//! see [`crate::reorder`] for the `ReorderStage` operator that adds +//! ordinal-based reordering on top of any `ItemQueue`. **Not** memory +//! bookkeeping at the trait level — `ByteBoundedQueue` is a +//! concrete impl that knows about heap size, but the trait surface is +//! type-uniform. +//! +//! Backpressure is expressed as `try_push -> Result<(), T>`: `Err(item)` +//! returns the rejected item back to the producer (which holds it in a +//! `HeldSlot` and re-pushes on the next worker iteration). No blocking, +//! no awaiting — pure non-blocking surface. +//! +//! Drained-signal protocol: +//! - Producer (output side) calls `mark_drained()` exactly once when the +//! producing step returns `StepOutcome::Finished` (counter-gated for +//! `Parallel` so only the last clone closes the shared queue). Subsequent +//! `try_push` calls panic — in every build, not just debug (a contract +//! violation: producer pushed after declaring done, and the item would be +//! silently lost). See `assert_not_drained`. +//! - Consumer (input side) checks `is_drained() && is_empty()` to detect +//! end-of-stream. Once both are true, no further items will arrive. + +use crossbeam_queue::{ArrayQueue, SegQueue}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; + +use super::item::HeapSize; +use super::runtime::metrics::EdgeMetrics; + +/// Transport-layer queue trait. Type-uniform across queue impls: the +/// `try_push` surface accepts any `T` regardless of whether the impl uses +/// item-count or memory bookkeeping internally. +/// +/// `Send + Sync`: queues are shared between worker threads via `Arc`. +pub trait ItemQueue: Send + Sync { + /// Non-blocking push. `Err(item)` returns the rejected item to the + /// caller; the framework holds it in a `HeldSlot` and retries. + /// + /// # Errors + /// + /// Returns `Err(item)` when the queue is at its backpressure limit + /// (item-count or byte-budget, depending on the impl). + fn try_push(&self, item: T) -> Result<(), T>; + + /// Non-blocking pop. `None` means the queue is currently empty (which + /// is *not* the same as drained — combine with `is_drained()`). + fn try_pop(&self) -> Option; + + /// True when no items are currently buffered. May race with concurrent + /// pushes; consumers that need a quiescent check combine with + /// `is_drained()`. + fn is_empty(&self) -> bool; + + /// Mark the queue drained (producer-side: "I'm done pushing"). Idempotent. + /// A `try_push` after `mark_drained` panics — see `assert_not_drained`. + fn mark_drained(&self); + + /// True if `mark_drained` has been called. + fn is_drained(&self) -> bool; +} + +/// Panic if `try_push` is called after `mark_drained`. +/// +/// The consumer treats a drained queue as closed, so an item pushed afterwards +/// may never be popped: the item is silently lost and the loss surfaces (if at +/// all) as a short output far from its cause. That makes this a framework +/// contract violation rather than a recoverable condition, so — like +/// `BranchOutputHandle::retry`'s `Ordered` + `ordinal = None` arm — it fails +/// loudly in **every** build. It was `debug_assert!`-only, which left release +/// builds performing exactly the silent push the message warns about. +/// +/// `Relaxed` is sufficient here and is the cheaper load on a per-item path. +/// `drained` is monotonic — its only write anywhere is `store(true, Release)` in +/// `mark_drained` — so a `Relaxed` load can return a stale `false` (a missed +/// detection when the producer races the close on another thread) but never a +/// spurious `true`. It cannot panic a correct program. +#[inline] +fn assert_not_drained(drained: &AtomicBool, queue_kind: &'static str) { + assert!( + !drained.load(Ordering::Relaxed), + "{queue_kind}::try_push after mark_drained — producer contract violation" + ); +} + +/// One entry per output branch in `StepProfile::output_queues`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum QueueSpec { + /// Item-count bounded. `try_push` rejects when `len() >= capacity`. + /// Best for fixed-size items (parsed records, compressed BGZF blocks + /// of known size, etc). + CountBounded { capacity: usize }, + /// Memory-bounded. `try_push` rejects when adding the item would push + /// the running byte counter past `limit_bytes`. Requires `T: HeapSize`. + /// Best for variable-size BAM batches and FASTQ batches. + ByteBounded { limit_bytes: u64 }, + /// No backpressure. `try_push` always succeeds. Use only when the + /// branch is naturally rate-limited upstream (e.g., a header-once + /// emit on pipeline start). + Unbounded, +} + +// ───────────────────────────────────────────────────────────────────────────── +// CountBoundedQueue +// ───────────────────────────────────────────────────────────────────────────── + +/// Item-count bounded transport. Backed by `crossbeam_queue::ArrayQueue`. +pub struct CountBoundedQueue { + inner: ArrayQueue, + drained: AtomicBool, + /// `Some` only on an instrumented edge (`--pipeline-trace`); `None` keeps the + /// hot path metric-free. Producer-push counts are recorded here; consumer-pop + /// counts are recorded at the `BranchInputHandle` (see `handles.rs`). + metrics: Option>, +} + +impl CountBoundedQueue { + /// Construct a count-bounded transport with the given capacity. + /// + /// # Panics + /// + /// Panics if `capacity == 0` (a zero-capacity queue would always reject). + #[must_use] + pub fn new(capacity: usize) -> Self { + Self::build(capacity, None) + } + + /// Like [`new`](Self::new) but recording producer-push metrics into `metrics` + /// (an instrumented edge). The non-blocking `try_*` surface is unchanged. + /// + /// # Panics + /// + /// Panics if `capacity == 0`. + #[must_use] + pub fn new_instrumented(capacity: usize, metrics: Arc) -> Self { + Self::build(capacity, Some(metrics)) + } + + /// [`new`](Self::new) when `metrics` is `None`, [`new_instrumented`](Self::new_instrumented) + /// when `Some`. Lets branch builders thread an optional metrics handle uniformly. + /// + /// # Panics + /// + /// Panics if `capacity == 0`. + #[must_use] + pub fn maybe_instrumented(capacity: usize, metrics: Option>) -> Self { + Self::build(capacity, metrics) + } + + fn build(capacity: usize, metrics: Option>) -> Self { + assert!(capacity > 0, "CountBoundedQueue capacity must be > 0"); + Self { inner: ArrayQueue::new(capacity), drained: AtomicBool::new(false), metrics } + } +} + +impl ItemQueue for CountBoundedQueue { + fn try_push(&self, item: T) -> Result<(), T> { + assert_not_drained(&self.drained, "CountBoundedQueue"); + if let Err(item) = self.inner.push(item) { + if let Some(m) = &self.metrics { + m.record_reject(); + } + return Err(item); + } + if let Some(m) = &self.metrics { + m.record_push(0); // count-bounded: items only, no byte size + } + Ok(()) + } + + fn try_pop(&self) -> Option { + let item = self.inner.pop()?; + Some(item) + } + + fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + fn mark_drained(&self) { + self.drained.store(true, Ordering::Release); + } + + fn is_drained(&self) -> bool { + self.drained.load(Ordering::Acquire) + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// ByteBoundedQueue +// ───────────────────────────────────────────────────────────────────────────── + +/// Backing slot capacity for `ByteBoundedQueue`. Since the queue's real +/// gate is the byte budget, this just needs to be large enough that the +/// count never matters for any sane workload. 1024 slots is well past +/// the working set of any single pipeline edge — even for the smallest +/// items the byte cap (default 4 MiB) imposes a tighter bound. +/// +/// Sized in pages of `crossbeam_queue::ArrayQueue` storage (one +/// pre-allocated slot array, no per-push allocation). Mirrors the +/// `ArrayQueue::new(queue_capacity)` strategy the legacy pipeline used — it +/// also used a fixed-capacity `ArrayQueue` everywhere for the same +/// reason: `SegQueue` allocates segments on demand under load, and +/// the resulting allocator churn shows up as `mi_*` overhead in +/// profiles (≈260 samples vs legacy on CODEC 8M). +const BYTE_BOUNDED_QUEUE_SLOT_CAPACITY: usize = 1024; + +/// Memory-bounded transport. Backed by +/// `crossbeam_queue::ArrayQueue<(T, u64)>` plus an atomic byte counter. +/// `try_push` rejects when the running byte counter has already reached +/// `limit_bytes`. Requires `T: HeapSize`. +/// +/// ## Concurrency / ordering +/// +/// The check-then-add is two atomics, so two concurrent pushes can both +/// observe `cur < limit` and both succeed, yielding a small overshoot. +/// The next push will see the overshoot and reject; the budget is +/// enforced as "approximate within one item's worth per producer." This +/// trade-off avoids a CAS loop and is fine for backpressure semantics. +/// +/// `current_bytes` is **only** a backpressure heuristic — it's not used +/// to synchronize handoff of the items themselves. The handoff is the +/// `ArrayQueue`'s job; `ArrayQueue`'s internal atomics provide the +/// happens-before relationship between `inner.push` and `inner.pop`. We +/// therefore use `Relaxed` ordering on every `current_bytes` access: +/// the worst case is a slightly stale reading of the budget, never an +/// observability violation on the items. +/// +/// ## Cached size at push +/// +/// The size is stored alongside the item in the inner queue +/// (`ArrayQueue<(T, u64)>`) so `try_pop` doesn't need to recompute +/// `T::heap_size()` for the budget update. For types whose +/// `heap_size()` is O(items inside) (e.g. `BatchedRawPositionGroups`, +/// `OrderedRawPositionGroup`) this avoids recomputing a O(group) +/// walk on every pop. Mirrors the legacy `ReorderBuffer`'s +/// cached-size storage strategy (`(T, usize)` there; `(T, u64)` here, +/// matching `inner`'s `ArrayQueue<(T, u64)>` above). +pub struct ByteBoundedQueue { + inner: ArrayQueue<(T, u64)>, + current_bytes: AtomicU64, + /// Mutable byte-budget cap. The rebalancer (when enabled via + /// `PipelineConfig::queue_memory_total`) updates this atomic at + /// runtime to shift budget across queues based on observed + /// fullness. Producers read it on every `try_push`; the + /// `Relaxed` ordering matches `current_bytes` (this is a + /// best-effort backpressure heuristic, not a correctness gate). + limit_bytes: AtomicU64, + drained: AtomicBool, + /// Per-instance one-shot guard so the "slot cap hit before byte budget" + /// warning (see `try_push`) is emitted at most once *per queue*, not once + /// per process. A process-global flag would silence the warning for every + /// later queue (e.g. a second `runall` stage, or many pipelines in one + /// long-lived host / test harness) after the first occurrence. The hot-path + /// cost is a single relaxed swap after the first hit. + slot_cap_warned: AtomicBool, + /// `Some` only on an instrumented edge; producer-push (items + bytes) and + /// rejections are recorded here. Consumer-pop is recorded at the + /// `BranchInputHandle` (see `handles.rs`). + metrics: Option>, +} + +impl ByteBoundedQueue { + /// Construct a byte-bounded transport with the given memory limit. + /// + /// # Panics + /// + /// Panics if `limit_bytes == 0` (a zero-budget queue would always reject). + #[must_use] + pub fn new(limit_bytes: u64) -> Self { + Self::build(limit_bytes, None) + } + + /// Like [`new`](Self::new) but recording producer-push metrics (items + bytes + /// + rejections) into `metrics`. Byte-budget semantics unchanged. + /// + /// # Panics + /// + /// Panics if `limit_bytes == 0`. + #[must_use] + pub fn new_instrumented(limit_bytes: u64, metrics: Arc) -> Self { + Self::build(limit_bytes, Some(metrics)) + } + + /// [`new`](Self::new) when `metrics` is `None`, [`new_instrumented`](Self::new_instrumented) + /// when `Some`. + /// + /// # Panics + /// + /// Panics if `limit_bytes == 0`. + #[must_use] + pub fn maybe_instrumented(limit_bytes: u64, metrics: Option>) -> Self { + Self::build(limit_bytes, metrics) + } + + fn build(limit_bytes: u64, metrics: Option>) -> Self { + assert!(limit_bytes > 0, "ByteBoundedQueue limit_bytes must be > 0"); + Self { + inner: ArrayQueue::new(BYTE_BOUNDED_QUEUE_SLOT_CAPACITY), + current_bytes: AtomicU64::new(0), + limit_bytes: AtomicU64::new(limit_bytes), + drained: AtomicBool::new(false), + slot_cap_warned: AtomicBool::new(false), + metrics, + } + } + + /// Best-effort, stale-tolerant `Relaxed` read of the running byte + /// counter. Used by the rebalancer as a budget heuristic, not as a + /// correctness gate — it may lag a concurrent `try_push`/`try_pop`. + #[must_use] + pub fn current_bytes(&self) -> u64 { + self.current_bytes.load(Ordering::Relaxed) + } + + /// Best-effort, stale-tolerant `Relaxed` read of the byte-budget cap. + /// A concurrent `set_limit_bytes` (rebalancer) may not yet be visible; + /// callers use this as a heuristic, never as a correctness gate. + #[must_use] + pub fn limit_bytes(&self) -> u64 { + self.limit_bytes.load(Ordering::Relaxed) + } + + /// Update the byte-budget cap. Called by the rebalancer when + /// reallocating budget across queues. Concurrent `try_push`es + /// see the new cap on their next read; transient overshoot + /// (pushes already in flight that read the old cap) is + /// self-correcting. + pub fn set_limit_bytes(&self, new_limit: u64) { + // Floor at 1. `try_push` rejects when `current_bytes >= limit_bytes`, so a + // limit of 0 rejects unconditionally — even on an empty edge — and wedges + // the producer permanently. `new` asserts `limit_bytes > 0` for exactly + // this reason; without a floor here that invariant could be undone after + // construction, which is the one case the constructor cannot guard. + // + // Clamped rather than asserted: this runs on a live pipeline (the budget + // pass and the rebalancer), where degrading to a 1-byte limit still makes + // progress — `try_push` admits an item whenever `current_bytes` is under + // the limit, regardless of item size — while a panic would take down a + // running pipeline over a recoverable arithmetic slip. Every current caller + // already applies its own positive per-queue floor. + self.limit_bytes.store(new_limit.max(1), Ordering::Relaxed); + } +} + +/// Type-erased handle for a byte-bounded queue. The pipeline +/// rebalancer iterates over registered handles to read fullness +/// (`current_bytes / limit_bytes`) and reallocate budget across +/// queues by calling `set_limit_bytes`. The trait deliberately +/// does not surface the queue's item type or its `ItemQueue` +/// methods — rebalancing only needs the byte counters. +pub trait BoundedQueueHandle: Send + Sync { + /// Bytes currently held in the queue. + fn current_bytes(&self) -> u64; + /// Current byte-budget cap. May change between calls if a + /// rebalancer is active. + fn limit_bytes(&self) -> u64; + /// Update the byte-budget cap. Concurrent producers see the + /// new value on their next push. + fn set_limit_bytes(&self, new_limit: u64); +} + +impl BoundedQueueHandle for ByteBoundedQueue { + fn current_bytes(&self) -> u64 { + self.current_bytes() + } + fn limit_bytes(&self) -> u64 { + self.limit_bytes() + } + fn set_limit_bytes(&self, new_limit: u64) { + self.set_limit_bytes(new_limit); + } +} + +impl ItemQueue for ByteBoundedQueue { + fn try_push(&self, item: T) -> Result<(), T> { + assert_not_drained(&self.drained, "ByteBoundedQueue"); + // Like the legacy `ReorderBufferState::can_proceed`, this + // gates on `heap_bytes < limit` — accept if currently *under* + // budget, regardless of incoming item size. Per-item-larger-than + // -limit is a real case (busy-locus position-group batches can + // be tens of MB while the queue limit is 4 MiB), so a strict + // `cur + size <= limit` would deadlock the producer. + // + // Once `cur` reaches `limit_bytes`, subsequent pushes reject + // until a consumer drains. Transient overshoot under concurrent + // pushes is self-correcting on the next round. + let cur = self.current_bytes.load(Ordering::Relaxed); + let limit = self.limit_bytes.load(Ordering::Relaxed); + if cur >= limit { + if let Some(m) = &self.metrics { + m.record_reject(); + } + return Err(item); + } + let size = item.heap_size() as u64; + // Reserve bytes before pushing so a concurrent consumer cannot pop and + // decrement the counter before we add our share, which would cause the + // counter to underflow and create permanent false backpressure. + self.current_bytes.fetch_add(size, Ordering::Relaxed); + // ArrayQueue::push returns Err((item, size)) on full; roll back the + // reservation and return the item to the caller for retry. (In practice + // the slot cap should never be hit before the byte budget triggers a + // reject above, but defend against it anyway.) + match self.inner.push((item, size)) { + Ok(()) => { + if let Some(m) = &self.metrics { + m.record_push(size); + } + Ok(()) + } + Err((item, _size)) => { + // Roll back the byte reservation — the item never entered the queue. + self.current_bytes.fetch_sub(size, Ordering::Relaxed); + if let Some(m) = &self.metrics { + m.record_reject(); + } + // The fixed 1024-slot backing was hit before the byte budget. + // This degrades byte-backpressure into a hard count cap for + // small items (heap_size ≲ limit/1024) — correctness is + // preserved (the producer retries) but throughput silently + // suffers. Surface it once so it is observable rather than a + // silent foot-gun; near-zero cost after the first hit. + if !self.slot_cap_warned.swap(true, Ordering::Relaxed) { + log::warn!( + "ByteBoundedQueue hit its {BYTE_BOUNDED_QUEUE_SLOT_CAPACITY}-slot count \ + cap before the byte budget; small items are degrading byte-backpressure \ + into a count cap (throughput, not correctness, is affected)." + ); + } + Err(item) + } + } + } + + fn try_pop(&self) -> Option { + let (item, size) = self.inner.pop()?; + self.current_bytes.fetch_sub(size, Ordering::Relaxed); + Some(item) + } + + fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + fn mark_drained(&self) { + self.drained.store(true, Ordering::Release); + } + + fn is_drained(&self) -> bool { + self.drained.load(Ordering::Acquire) + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// UnboundedQueue +// ───────────────────────────────────────────────────────────────────────────── + +/// Unbounded transport. `try_push` always succeeds. Backed by `SegQueue`. +pub struct UnboundedQueue { + inner: SegQueue, + drained: AtomicBool, + /// `Some` only on an instrumented edge; producer-push items are recorded + /// here (unbounded → never rejects, no byte tracking). Consumer-pop is at the + /// `BranchInputHandle`. + metrics: Option>, +} + +impl UnboundedQueue { + #[must_use] + pub fn new() -> Self { + Self { inner: SegQueue::new(), drained: AtomicBool::new(false), metrics: None } + } + + /// Like [`new`](Self::new) but recording producer-push item counts into + /// `metrics`. Unbounded edges have no byte budget and never reject; depth is + /// reported as raw length only. + #[must_use] + pub fn new_instrumented(metrics: Arc) -> Self { + Self { inner: SegQueue::new(), drained: AtomicBool::new(false), metrics: Some(metrics) } + } + + /// [`new`](Self::new) when `metrics` is `None`, [`new_instrumented`](Self::new_instrumented) + /// when `Some`. + #[must_use] + pub fn maybe_instrumented(metrics: Option>) -> Self { + Self { inner: SegQueue::new(), drained: AtomicBool::new(false), metrics } + } +} + +impl Default for UnboundedQueue { + fn default() -> Self { + Self::new() + } +} + +impl ItemQueue for UnboundedQueue { + fn try_push(&self, item: T) -> Result<(), T> { + assert_not_drained(&self.drained, "UnboundedQueue"); + self.inner.push(item); + if let Some(m) = &self.metrics { + m.record_push(0); + } + Ok(()) + } + + fn try_pop(&self) -> Option { + self.inner.pop() + } + + fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + fn mark_drained(&self) { + self.drained.store(true, Ordering::Release); + } + + fn is_drained(&self) -> bool { + self.drained.load(Ordering::Acquire) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + use std::sync::Arc; + + /// A `try_push` after `mark_drained` must panic on every transport impl, in + /// every build. It was `debug_assert!`-only, so a release build pushed the + /// item into a queue the consumer had already closed — a silent loss that + /// surfaces only as a short output. `#[values]` covers all three impls so a + /// new transport that forgets the guard is caught by the same table. + #[rstest] + #[case::count_bounded(Arc::new(CountBoundedQueue::new(2)) as Arc>)] + #[case::byte_bounded(Arc::new(ByteBoundedQueue::new(1024)) as Arc>)] + #[case::unbounded(Arc::new(UnboundedQueue::new()) as Arc>)] + #[should_panic(expected = "try_push after mark_drained — producer contract violation")] + fn try_push_after_mark_drained_panics(#[case] q: Arc>) { + q.mark_drained(); + let _ = q.try_push(1); + } + + /// The guard must not fire before `mark_drained` — a plain push on a fresh + /// queue still succeeds on every impl. + #[rstest] + #[case::count_bounded(Arc::new(CountBoundedQueue::new(2)) as Arc>)] + #[case::byte_bounded(Arc::new(ByteBoundedQueue::new(1024)) as Arc>)] + #[case::unbounded(Arc::new(UnboundedQueue::new()) as Arc>)] + fn try_push_before_mark_drained_succeeds(#[case] q: Arc>) { + assert!(q.try_push(1).is_ok(), "an undrained queue must still accept a push"); + assert_eq!(q.try_pop(), Some(1)); + } + + #[test] + fn count_bounded_round_trip() { + let q: Arc> = Arc::new(CountBoundedQueue::new(2)); + assert!(q.try_push(1).is_ok()); + assert!(q.try_push(2).is_ok()); + assert_eq!(q.try_push(3), Err(3)); + assert_eq!(q.try_pop(), Some(1)); + assert_eq!(q.try_pop(), Some(2)); + assert_eq!(q.try_pop(), None); + } + + #[test] + fn count_bounded_drain_signal() { + let q = CountBoundedQueue::::new(4); + assert!(!q.is_drained()); + q.mark_drained(); + assert!(q.is_drained()); + } + + #[derive(Debug)] + struct Heavy(Vec); + impl HeapSize for Heavy { + fn heap_size(&self) -> usize { + self.0.len() + } + } + + #[test] + fn byte_bounded_slot_cap_reject_is_observable() { + // Tiny (0-byte heap) items with a huge byte limit: the byte budget is + // never reached, so the fixed slot backing becomes the binding cap. + // The first SLOT_CAPACITY pushes succeed; the next rejects on the slot + // cap even though current_bytes is far below the limit. Regression for + // F02 — this path silently degraded byte-backpressure into a count cap; + // it is now warn-once observable, and this pins the reject behaviour. + let q = ByteBoundedQueue::::new(1_000_000); + for i in 0..BYTE_BOUNDED_QUEUE_SLOT_CAPACITY { + assert!(q.try_push(Heavy(Vec::new())).is_ok(), "push {i} within slot cap"); + } + assert_eq!(q.current_bytes(), 0, "0-byte items leave the byte budget unused"); + assert!( + q.try_push(Heavy(Vec::new())).is_err(), + "push #{} must reject on the slot cap, not the byte budget", + BYTE_BOUNDED_QUEUE_SLOT_CAPACITY + 1 + ); + } + + /// The slot-cap reject path reserves `size` bytes *before* pushing and rolls + /// the reservation back when `ArrayQueue::push` reports full. The sibling test + /// above fills with 0-byte items, so that `fetch_sub` runs with `size == 0` + /// and a leak or double-subtract is invisible. Fill with nonzero items under a + /// limit large enough that the slot cap still binds, and pin the byte counter + /// across the rejected push. + #[test] + fn byte_bounded_slot_cap_reject_rolls_back_reserved_bytes() { + const ITEM_BYTES: usize = 8; + // Large enough that 1024 * 8 bytes never reaches it, so the reject below + // is the slot cap and not the byte budget. + let q = ByteBoundedQueue::::new(1_000_000); + for i in 0..BYTE_BOUNDED_QUEUE_SLOT_CAPACITY { + assert!(q.try_push(Heavy(vec![0; ITEM_BYTES])).is_ok(), "push {i} within slot cap"); + } + let before = q.current_bytes(); + assert_eq!( + before, + (BYTE_BOUNDED_QUEUE_SLOT_CAPACITY * ITEM_BYTES) as u64, + "every admitted item's bytes are accounted" + ); + assert!(before < 1_000_000, "the byte budget must not be the binding cap here"); + assert!( + q.try_push(Heavy(vec![0; ITEM_BYTES])).is_err(), + "push #{} must reject on the slot cap", + BYTE_BOUNDED_QUEUE_SLOT_CAPACITY + 1 + ); + assert_eq!( + q.current_bytes(), + before, + "a slot-cap reject must roll its reservation back exactly — leaking bytes here \ + would create permanent false backpressure" + ); + } + + /// A 0 limit makes `try_push` reject unconditionally (`current_bytes >= 0` is + /// always true), wedging the producer forever. `new` asserts against it, so + /// the setter must not be able to reintroduce it after construction. Clamping + /// to 1 keeps the edge alive: `try_push` admits an item whenever + /// `current_bytes` is *under* the limit, whatever the item's size. + #[test] + fn set_limit_bytes_clamps_zero_to_one_so_the_edge_still_admits() { + let q = ByteBoundedQueue::::new(4096); + q.set_limit_bytes(0); + assert_eq!(q.limit_bytes(), 1, "a 0 limit is floored to 1, never stored as 0"); + assert!( + q.try_push(Heavy(vec![0; 64])).is_ok(), + "an empty edge must still admit one item — a 0 limit would reject forever" + ); + // Now over the 1-byte limit, so the next push rejects: still a real bound, + // not a silent promotion to unbounded. + assert!(q.try_push(Heavy(vec![0; 64])).is_err(), "the clamped limit still applies"); + } + + #[test] + fn slot_cap_warn_flag_is_per_instance_not_process_global() { + // The "slot cap hit before byte budget" warn-once guard lives on the + // queue instance, so a second queue (e.g. a later runall stage, or a new + // pipeline in a long-lived host) still warns on its own first hit — the + // signal is not silenced process-wide by an earlier queue. + let fill_to_slot_cap = |q: &ByteBoundedQueue| { + for _ in 0..BYTE_BOUNDED_QUEUE_SLOT_CAPACITY { + q.try_push(Heavy(Vec::new())).expect("push within slot cap"); + } + // This push trips the slot cap and (first time) sets the flag. + assert!(q.try_push(Heavy(Vec::new())).is_err(), "push must reject on slot cap"); + }; + + let q1 = ByteBoundedQueue::::new(1_000_000); + assert!(!q1.slot_cap_warned.load(Ordering::Relaxed)); + fill_to_slot_cap(&q1); + assert!(q1.slot_cap_warned.load(Ordering::Relaxed), "first queue must warn on its hit"); + + // A fresh queue starts un-warned even though q1 already warned, so it + // will warn on its own first hit (per-queue, not process-global). + let q2 = ByteBoundedQueue::::new(1_000_000); + assert!( + !q2.slot_cap_warned.load(Ordering::Relaxed), + "a second queue must NOT inherit the first queue's warned state" + ); + fill_to_slot_cap(&q2); + assert!( + q2.slot_cap_warned.load(Ordering::Relaxed), + "second queue must warn on its own hit" + ); + } + + #[test] + fn byte_bounded_respects_limit() { + let q = ByteBoundedQueue::::new(100); + // Empty queue accepts even an oversized item (legacy semantics: + // gate on `cur < limit`, not `cur + size <= limit`). This is the + // fix for the per-item-larger-than-limit deadlock. + assert!(q.try_push(Heavy(vec![0; 200])).is_ok()); + assert_eq!(q.current_bytes(), 200); + // Now `cur >= limit`, all subsequent pushes reject regardless + // of size. + let rejected = q.try_push(Heavy(vec![0; 1])); + assert!(rejected.is_err(), "queue at/over budget should reject"); + assert_eq!(q.current_bytes(), 200); + // After a pop drops `cur` below limit, pushes succeed again. + let _ = q.try_pop().unwrap(); + assert_eq!(q.current_bytes(), 0); + assert!(q.try_push(Heavy(vec![0; 50])).is_ok()); + assert_eq!(q.current_bytes(), 50); + } + + #[test] + fn byte_bounded_oversized_first_push_succeeds() { + // Regression: previously a single push larger than `limit_bytes` + // would always reject (`0 + size > limit`), deadlocking + // producers that emit oversized batches (e.g. busy-locus + // position-group batches). With the legacy `cur < limit` + // semantics, the oversized push goes through. + let q = ByteBoundedQueue::::new(100); + assert!(q.try_push(Heavy(vec![0; 1024])).is_ok()); + } + + #[test] + fn byte_bounded_decrements_on_pop() { + let q = ByteBoundedQueue::::new(1000); + q.try_push(Heavy(vec![0; 200])).unwrap(); + assert_eq!(q.current_bytes(), 200); + let _ = q.try_pop().unwrap(); + assert_eq!(q.current_bytes(), 0); + } + + #[test] + fn unbounded_never_rejects() { + let q = UnboundedQueue::::new(); + for i in 0..1024 { + assert!(q.try_push(i).is_ok()); + } + } + + // ── Per-edge metrics (L2-instrumentation Task 2) ───────────────────────── + + #[test] + fn instrumented_queue_counts_push_and_reject() { + let m = EdgeMetrics::new(); + let q = CountBoundedQueue::::new_instrumented(1, Arc::clone(&m)); + assert!(q.try_push(1).is_ok()); + assert_eq!(q.try_push(2), Err(2)); // full (cap 1) → reject + let s = m.snapshot(); + assert_eq!(s.pushed_items, 1, "one successful push"); + assert_eq!(s.push_rejections, 1, "one rejection"); + // Producer-push only at this layer; pop is counted at the input handle. + assert_eq!(s.popped_items, 0); + } + + #[test] + fn byte_bounded_instrumented_push_bytes_and_depth() { + let m = EdgeMetrics::new(); + let q = ByteBoundedQueue::::new_instrumented(1000, Arc::clone(&m)); + q.try_push(Heavy(vec![0; 200])).unwrap(); + let s = m.snapshot(); + assert_eq!(s.pushed_items, 1); + assert_eq!(s.pushed_bytes, 200); + } + + #[test] + fn byte_bounded_instrumented_counts_reject() { + // The byte-budget reject path increments push_rejections (distinct from + // the CountBounded slot reject above). Fill to budget, then a push rejects. + let m = EdgeMetrics::new(); + let q = ByteBoundedQueue::::new_instrumented(100, Arc::clone(&m)); + q.try_push(Heavy(vec![0; 200])).unwrap(); // accepted (cur::new(4).metrics.is_none()); + assert!(ByteBoundedQueue::::new(100).metrics.is_none()); + assert!(UnboundedQueue::::new().metrics.is_none()); + // And instrumented constructors do attach metrics. + assert!( + CountBoundedQueue::::new_instrumented(4, EdgeMetrics::new()).metrics.is_some() + ); + assert!(UnboundedQueue::::new_instrumented(EdgeMetrics::new()).metrics.is_some()); + assert!( + ByteBoundedQueue::::new_instrumented(100, EdgeMetrics::new()).metrics.is_some() + ); + + // `maybe_instrumented` is the constructor the branch builders actually + // call, and it was the only one with no coverage: one that dropped a + // `Some(metrics)` would silently produce an edge that reports nothing + // under `--pipeline-trace`, with every other test still green. Both + // directions, all three impls. + assert!(CountBoundedQueue::::maybe_instrumented(4, None).metrics.is_none()); + assert!(ByteBoundedQueue::::maybe_instrumented(100, None).metrics.is_none()); + assert!(UnboundedQueue::::maybe_instrumented(None).metrics.is_none()); + assert!( + CountBoundedQueue::::maybe_instrumented(4, Some(EdgeMetrics::new())) + .metrics + .is_some() + ); + assert!( + ByteBoundedQueue::::maybe_instrumented(100, Some(EdgeMetrics::new())) + .metrics + .is_some() + ); + assert!( + UnboundedQueue::::maybe_instrumented(Some(EdgeMetrics::new())).metrics.is_some() + ); + } +} diff --git a/crates/fgumi-pipeline-core/src/reorder.rs b/crates/fgumi-pipeline-core/src/reorder.rs new file mode 100644 index 000000000..c0f1c06b3 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/reorder.rs @@ -0,0 +1,1036 @@ +//! Reorder operator on top of an `ItemQueue>`. +//! +//! `ReorderStage` is layered between a producer's transport queue and +//! the consumer's `InputHandle`. The producer pushes items wrapped in +//! a framework-managed `Sequenced { ordinal, item }`. The reorder +//! stage buffers items until their ordinal equals `next_serial`, then +//! releases them in order. +//! +//! Smart backpressure (deadlock avoidance): +//! - Until `next_serial` is *observable in the reorder buffer*, the stage +//! MUST accept everything (refusing could deadlock — the producer of +//! `next_serial` may be one of the backpressured producers). +//! - Once `next_serial` is in the buffer, the consumer can make progress +//! and producers MAY be rejected by the underlying transport's normal +//! backpressure. +//! +//! "Observable in the reorder buffer" is the conservative test: items +//! already in transport but not yet pulled into the buffer count as +//! *not* observable. This means producers over-accept when `next_serial` is +//! sitting in transport waiting to be pulled — but it can never deadlock, +//! and the over-acceptance window closes the next time a consumer calls +//! `try_pop_in_order` (which drains transport into the buffer). +//! +//! Storage layering: +//! - "In flight" items live in the underlying `ItemQueue>` +//! transport (`CountBounded` / `Unbounded` — see PR 1 caveat below). +//! - "Stashed" items (must-accept overflow when transport rejected, or +//! items pulled but not yet at their turn) live in `state.buffer: +//! AHashMap` (ahash — the ordinals are trivial monotonic `u64` +//! keys, so the default `SipHash` buys nothing on this per-item path). +//! - The transport's backpressure budget covers only in-flight items. +//! The overflow stash has its own byte cap (`max_overflow_bytes`). The +//! framework sizes that cap thread-awarely from the per-edge transport +//! budget (see `apply_initial_queue_budget` / `set_max_overflow_bytes`), +//! clamped to a fixed ceiling — so at low thread counts the stash stays +//! small (a streaming footprint) and at high thread counts it keeps the +//! prior lookahead headroom. `next_serial` is always exempt from the cap, +//! so the stash bound is purely a memory/throughput knob, never a +//! liveness constraint (any cap ≥ 0 is deadlock-free). +//! +//! `Sequenced` impls `HeapSize` (see below), so `BranchOrdering::ByOrdinal` +//! / `ByItemOrdinal` compose with `QueueSpec::ByteBounded` — the canonical BAM +//! step output shape (`build_branch_ordered_bytes`). + +use ahash::AHashMap; +use parking_lot::Mutex; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; + +use super::item::{HeapSize, Ordered}; +use super::queues::ItemQueue; + +/// Default cap on a `ReorderStage`'s must-accept overflow buffer, in +/// bytes. Mirrors legacy `BACKPRESSURE_THRESHOLD_BYTES / 2` (`base.rs:704`): +/// the legacy pipeline gates non-`next_seq` reorder pushes at half the +/// 512 MB threshold (= 256 MB). We use the same value per branch so a +/// multi-stage pipeline with four ordered edges peaks around 1 GB of +/// reorder overflow under heterogeneous load — far below the unbounded +/// growth Task #29 hit (23+ GB on a 53M-record group workload). Items +/// at `next_serial` are exempt to preserve liveness. +pub const DEFAULT_REORDER_OVERFLOW_BYTES: u64 = 256 * 1024 * 1024; + +/// Per-branch ordering directive in `StepProfile::branch_ordering`. +/// +/// Three modes; pick based on what the consumer needs: +/// +/// - **`None`** — FIFO by arrival. Cheapest. Consumer sees items in the +/// order workers happened to push them, which under multi-producer +/// concurrency is non-deterministic. +/// +/// - **`ByOrdinal`** — producer-allocated ordinals via a per-branch +/// `AtomicU64` counter. The framework wraps each pushed item in a +/// `Sequenced` and inserts a `ReorderStage` in front of the consumer. +/// Imposes a total order at the producer's emission point but **does +/// not preserve any pre-existing global ordering** — under multi- +/// producer Parallel concurrency, the ordinal a worker gets bears no +/// relation to the order of the input it processed. Useful for +/// single-producer steps (sources) and for cases where any deterministic +/// total order suffices. +/// +/// - **`ByItemOrdinal`** — items carry their own serial via the [`Ordered`] +/// trait. The framework reads `item.ordinal()` instead of allocating +/// one. Preserves global ordering across multi-step Parallel transforms +/// when each step propagates the input's serial onto its outputs (the +/// canonical pattern for BAM pipelines: every batch carries +/// `batch_serial: u64` from its read order, and every transform +/// preserves that serial on its output items). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BranchOrdering { + /// FIFO by arrival. + None, + /// Producer-allocated ordinal via a per-branch `AtomicU64` counter. + ByOrdinal, + /// Items carry their own ordinal via the `Ordered` trait. Requires + /// `T: Ordered` at branch construction time. + ByItemOrdinal, +} + +/// Framework-internal wrapper carrying a producer-assigned ordinal. +/// +/// Step authors never see this type; it is `pub` only because it appears in the +/// public `ItemQueue>` trait bounds that ordered queues (and their +/// tests) instantiate, so callers may construct it directly when wiring queues. +pub struct Sequenced { + /// Producer-assigned monotonic ordinal used to restore emission order. + pub ordinal: u64, + /// The wrapped payload carried alongside its `ordinal`. + pub item: T, +} + +impl Ordered for Sequenced { + fn ordinal(&self) -> u64 { + self.ordinal + } +} + +/// `Sequenced` forwards `HeapSize` to the inner item, allowing byte- +/// bounded queues to wrap ordered items. The `ordinal` field is stack- +/// only and contributes nothing to the heap budget. +impl HeapSize for Sequenced { + fn heap_size(&self) -> usize { + self.item.heap_size() + } +} + +/// Reorder operator. Wraps an `Arc>>` and +/// presents a `try_pop_in_order() -> Option` surface plus an +/// ordinal-tagged `try_push(ordinal, item)`. +pub struct ReorderStage { + transport: Arc>>, + state: Mutex>, + /// Cached "is `next_serial` currently in the buffer?" snapshot, + /// updated under the state lock by writers and read **without** the + /// lock by `try_push`'s fast path. Mirrors legacy + /// `OrderedQueue::has_next` (`queue.rs:54`). + /// + /// Semantics: + /// - `false` → `next_serial` NOT in buffer → producers MUST overflow + /// into the buffer if transport rejects (the slow path takes the + /// state lock and re-checks under it; the I1 fix from Phase 1). + /// - `true` → `next_serial` IS in buffer → consumer can drain → + /// producers CAN apply backpressure. Fast path: a lock-free + /// `transport.try_push` whose rejection surfaces as `Err`. + /// + /// Stale-`true` is harmless (the slow path is always correct). + /// Stale-`false` is also harmless (just a missed fast-path + /// opportunity; producer takes the slow path which still does the + /// right thing). + next_serial_buffered: AtomicBool, + /// Sticky drained-observed flag for `is_drained()` callers that want + /// to short-circuit on subsequent calls. Set lazily by `is_drained` + /// when it observes the drained-and-empty condition. + drained_observed: AtomicBool, + /// Byte cap on the must-accept overflow stash: once the buffered bytes + /// reach the cap, must-accept rejects new pushes for ordinals other than + /// `next_serial` (the producer then holds via its `held_slot`). + /// `next_serial` is always exempt — its producer cannot be backpressured + /// by this cap (see Task #29) — so any cap value is deadlock-free. + /// `u64::MAX` (`REORDER_OVERFLOW_UNBOUNDED`) means unbounded. + /// + /// Set once at construction (`with_max_overflow_bytes`, the no-budget + /// fallback) and re-set once before workers spawn by the budget pass + /// (`set_max_overflow_bytes` via `apply_initial_queue_budget`) to a + /// thread-aware value from the per-edge transport budget. Read only on the + /// must-accept slow path (under the `state` lock), so the atomic costs + /// nothing on the lock-free fast path. Mirrors legacy + /// `ReorderBufferState::can_proceed` (`base.rs:766-781`). + max_overflow_bytes: AtomicU64, + /// Producer-side push metrics for an **ordered** edge, recorded at THIS + /// boundary rather than on the transport queue. `try_push` can turn a + /// full-transport `Err` into a must-accept stash `Ok`, so recording on the + /// transport would miscount a stashed (accepted) item as a rejection. Here we + /// record `record_push` on every accepted push (transport OR stash) and + /// `record_reject` only on a true `Err`, keeping `pushed_items` / + /// `push_rejections` accurate under producer skew. `None` when the edge is + /// not instrumented; the pop side is recorded separately by the input handle. + push_metrics: Option>, + /// Whether `push_metrics` records each push's `heap_size` (byte-bounded edge) + /// or `0` (count/unbounded edge) — mirroring the underlying queue's own byte + /// accounting. Set from the branch's queue kind at construction. + record_item_bytes: bool, +} + +/// Unbounded sentinel for [`ReorderStage::max_overflow_bytes`]. +const REORDER_OVERFLOW_UNBOUNDED: u64 = u64::MAX; + +/// Type-erased setter for a [`ReorderStage`]'s overflow cap, so the runtime +/// can size it without naming the branch's item type `T`. The framework +/// collects one per ordered byte-bounded branch and sets the cap in +/// `apply_initial_queue_budget` (before any worker spawns). Kept separate +/// from `BoundedQueueHandle` (the transport-resize handle) so existing +/// transport-handle impls are untouched. +pub trait ReorderCapHandle: Send + Sync { + /// Set the overflow byte cap (`u64::MAX` = unbounded). + fn set_max_overflow_bytes(&self, bytes: u64); + + /// Bytes currently held in the must-accept overflow stash. The deadlock + /// monitor sums this across branches (with the transport queues) to tell a + /// real wedge (work stuck) from upstream starvation (everything empty); the + /// cap-enforcement tests use it to observe stash growth. + fn current_buffer_bytes(&self) -> u64; +} + +impl ReorderCapHandle for ReorderStage { + fn set_max_overflow_bytes(&self, bytes: u64) { + // Set once before workers spawn; `Relaxed` is sufficient because the + // worker pool's spawn establishes the happens-before edge, and the + // value is only read on the must-accept slow path under `state.lock()`. + self.max_overflow_bytes.store(bytes, Ordering::Relaxed); + } + + fn current_buffer_bytes(&self) -> u64 { + self.state.lock().buffer_bytes + } +} + +#[cfg(test)] +impl ReorderStage { + /// Read the current overflow cap (`u64::MAX` = unbounded). Test-only — + /// lets the budget-wiring test assert `apply_initial_queue_budget` set it. + pub(crate) fn current_max_overflow_bytes(&self) -> u64 { + self.max_overflow_bytes.load(Ordering::Relaxed) + } +} + +struct ReorderState { + /// Items pulled from transport but not yet at their turn, indexed by + /// ordinal. Each entry caches the item's `heap_size()` measured at + /// insert time so we don't pay an O(item) walk on every pop or + /// transport→buffer drain. Mirrors legacy `ReorderBuffer` + /// (`fgumi-bam-io/src/reorder.rs:50`) which stores `(T, usize)`. + /// The size is consumed by `buffer_bytes` accounting; for items + /// whose `heap_size()` is O(records) (e.g. position groups) this + /// caching avoids 3× the `heap_size` cost per item flowing through. + buffer: AHashMap, + /// Tracked heap bytes of items currently in `buffer`. Updated on + /// insert + on `try_pop_in_order` drain. Used to enforce + /// `max_overflow_bytes`. Mirrors legacy + /// `ReorderBufferState::heap_bytes` (`base.rs:746`). + buffer_bytes: u64, + /// Next ordinal we'll release. + next_serial: u64, +} + +impl ReorderStage { + #[must_use] + pub fn new(transport: Arc>>) -> Self { + Self { + transport, + state: Mutex::new(ReorderState { + buffer: AHashMap::new(), + buffer_bytes: 0, + next_serial: 0, + }), + next_serial_buffered: AtomicBool::new(false), + drained_observed: AtomicBool::new(false), + max_overflow_bytes: AtomicU64::new(REORDER_OVERFLOW_UNBOUNDED), + push_metrics: None, + record_item_bytes: false, + } + } + + /// Variant that caps the must-accept overflow buffer to `max_bytes` + /// of accumulated heap (computed via `T::heap_size()` per item). + /// `next_serial` is exempt (always accepted to preserve liveness); + /// other ordinals are rejected back to the producer when the buffer + /// is at the byte cap, so producers hold via their `held_slot` and + /// retry. Without a cap, heterogeneous-size workloads (e.g. one + /// large position group from a busy locus while others crank + /// through small groups) can OOM via must-accept overflow. + /// + /// Mirrors legacy `ReorderBufferState::can_proceed` semantics + /// (`base.rs:766-781`): the legacy gates non-`next_seq` pushes on + /// `heap_bytes < memory_limit / 2`. We collapse the halving into + /// the caller-supplied cap so the hot-path check is one load + one + /// compare. + /// + /// In production the framework constructs the stage with a fallback cap + /// here and then RE-SIZES it thread-awarely via [`set_max_overflow_bytes`] + /// in `apply_initial_queue_budget` (from the same per-edge budget as the + /// transport queue). So this constructor's value is the no-budget fallback; + /// the live cap tracks the transport budget. + /// + /// [`set_max_overflow_bytes`]: ReorderCapHandle::set_max_overflow_bytes + #[must_use] + pub fn with_max_overflow_bytes( + transport: Arc>>, + max_bytes: u64, + ) -> Self { + Self { + transport, + state: Mutex::new(ReorderState { + buffer: AHashMap::new(), + buffer_bytes: 0, + next_serial: 0, + }), + next_serial_buffered: AtomicBool::new(false), + drained_observed: AtomicBool::new(false), + max_overflow_bytes: AtomicU64::new(max_bytes), + push_metrics: None, + record_item_bytes: false, + } + } + + /// Attach producer-side push metrics recorded at the `ReorderStage` boundary + /// (see `push_metrics`). `record_item_bytes` is `true` + /// for a byte-bounded edge (record each push's `heap_size`) and `false` for a + /// count/unbounded edge (record `0`), matching the underlying queue's byte + /// accounting. Called once by the ordered-branch builder. + #[must_use] + pub fn with_push_metrics( + mut self, + push_metrics: Option>, + record_item_bytes: bool, + ) -> Self { + self.push_metrics = push_metrics; + self.record_item_bytes = record_item_bytes; + self + } + + /// Producer-side push. The framework allocates `ordinal` from a + /// per-branch `AtomicU64` counter (see `handles.rs`). + /// + /// Backpressure semantics: + /// - If `next_serial` is not yet in the buffer, MUST accept (overflows + /// into the buffer if transport rejects). This prevents the deadlock + /// where the producer of `next_serial` is itself backpressured. + /// - If `next_serial` is in the buffer, the transport's normal + /// `try_push` rules apply; rejection surfaces as `Err((ordinal, item))`. + /// + /// Performance: when `next_serial` is observed in the buffer (the + /// steady-state case once the consumer is keeping up), the push is + /// **lock-free** — only the transport's atomic `try_push` runs. The + /// state mutex is acquired only when overflow into the buffer might + /// be needed. + /// + /// # Errors + /// + /// Returns `Err((ordinal, item))` when the transport rejected and we + /// were not in must-accept mode (i.e., `next_serial` is already + /// observable, so the consumer can drain). + pub fn try_push(&self, ordinal: u64, item: T) -> Result<(), (u64, T)> { + // Record push-side metrics at THIS boundary (see `push_metrics`). The + // must-accept path turns a full-transport `Err` into a stash `Ok`, so the + // transport queue can't tell a stashed (accepted) push from a reject — + // only the final `Result` here can. `heap_size()` is read before `item` + // moves into the inner push. Every `Ok` (transport OR stash) is a push; + // every `Err` (backpressure or stash-cap held) is a reject. + // + // Gate the `heap_size()` call on metrics being present: on the default + // instrumentation-off path (`push_metrics == None`) the byte figure is + // never recorded, so computing it would be pure hot-path overhead. + let bytes = if self.push_metrics.is_some() && self.record_item_bytes { + item.heap_size() as u64 + } else { + 0 + }; + let result = self.try_push_inner(ordinal, item); + if let Some(m) = &self.push_metrics { + match &result { + Ok(()) => m.record_push(bytes), + Err(_) => m.record_reject(), + } + } + result + } + + /// Inner push: the transport / must-accept-stash decision, without metrics. + /// See [`try_push`](Self::try_push) for the public contract; metrics are + /// recorded there so a stashed push is not miscounted as a rejection. + fn try_push_inner(&self, ordinal: u64, item: T) -> Result<(), (u64, T)> { + // Lock-free fast path. If next_serial is in the buffer (consumer + // can drain), the producer is in pure-backpressure mode: a + // transport push that succeeds returns Ok; a rejection returns + // Err. No overflow into the buffer is possible, so we don't need + // the state lock at all. + // + // Liveness: returning `Err` here (instead of falling back to the + // slow path) is safe because `try_pop_in_order` re-derives + // `next_serial_buffered` under the state lock, and the round-robin + // worker driver guarantees a `try_pop_in_order` runs between + // producer retries — so a stale-`true` flag here is corrected on + // the next consumer poll and the producer is re-dispatched. + if self.next_serial_buffered.load(Ordering::Acquire) { + let seq = Sequenced { ordinal, item }; + return match self.transport.try_push(seq) { + Ok(()) => Ok(()), + Err(seq) => Err((seq.ordinal, seq.item)), + }; + } + + // Slow path: must_accept may apply. Hold the state lock across + // the must_accept check AND the transport push so a concurrent + // consumer can't drain `next_serial` between the two — which + // would let producers keep over-accepting into the overflow + // buffer indefinitely (I1 from the Phase 1 review). + // + // Transport pushes are non-blocking, so holding the lock briefly + // is safe (we don't risk priority inversion against blocking I/O). + let mut state = self.state.lock(); + let must_accept = !state.buffer.contains_key(&state.next_serial); + let seq = Sequenced { ordinal, item }; + + if must_accept { + // Apply the byte-aware overflow cap *before* attempting the + // transport push, not only when transport is full. The consumer's + // `try_pop_in_order` drain loop relocates the entire transport into + // the stash while `next_serial` is absent, so a cap that only fires + // on transport-full never binds — the consumer keeps transport + // non-full — and the stash grows without bound (#330 zipper OOM: + // 7.7 GB of reorder stash vs 466 MB of transport). Gating the push + // on `buffer_bytes >= cap` regardless of transport room bounds the + // stash to roughly `cap + one transport-worth`. + // + // Liveness preserved: `next_serial` is exempt (always accepted), + // so the producer of `next_serial` can never be backpressured by + // this cap and the consumer can always make progress. Any cap value + // is deadlock-free (see `concurrent_tiny_cap_drains_all_in_order`, + // which proves this with a 1-byte cap). + let landed_next = ordinal == state.next_serial; + let cap = self.max_overflow_bytes.load(Ordering::Relaxed); + if !landed_next && cap != REORDER_OVERFLOW_UNBOUNDED && state.buffer_bytes >= cap { + let Sequenced { ordinal, item } = seq; + return Err((ordinal, item)); + } + match self.transport.try_push(seq) { + Ok(()) => Ok(()), + Err(seq) => { + // Transport full: overflow into the stash. The cap was + // already checked above; `next_serial` is exempt either + // way, so the must-accept liveness guarantee holds. + // `usize → u64` is a lossless widen on every supported + // (≤64-bit) target, so this cast never truncates. + let item_bytes = seq.item.heap_size() as u64; + // A duplicate ordinal would silently drop the buffered item and + // leak its bytes into `buffer_bytes`. `ByOrdinal` serials are + // unique by construction; a hit here means a `ByItemOrdinal` + // upstream emitted two items with the same serial (a step bug). + // Fail loud in release too (like the `next_serial` overflow guard + // below): silently dropping a buffered record is a data-integrity bug. + assert!( + !state.buffer.contains_key(&seq.ordinal), + "duplicate reorder ordinal {} — ByItemOrdinal upstream serials must be unique", + seq.ordinal + ); + state.buffer.insert(seq.ordinal, (seq.item, item_bytes)); + state.buffer_bytes = state.buffer_bytes.saturating_add(item_bytes); + if landed_next { + self.next_serial_buffered.store(true, Ordering::Release); + } + Ok(()) + } + } + } else { + // Cache was stale (false) but next_serial actually IS in the + // buffer. Update the cache so subsequent producers take the + // fast path, then apply backpressure. + self.next_serial_buffered.store(true, Ordering::Release); + match self.transport.try_push(seq) { + Ok(()) => Ok(()), + Err(seq) => Err((seq.ordinal, seq.item)), + } + } + } + + /// Consumer-side pop. Returns `Some(T)` if `next_serial` is available, + /// `None` if still waiting for it. + /// + /// # Panics + /// + /// Panics if the in-order ordinal counter would overflow `u64` (i.e. + /// `next_serial == u64::MAX`). Branch ordinals start at 0 and increment, so + /// this is unreachable on any real workload (~1.8e19 items on one edge); the + /// guard exists only to fail loudly rather than silently wrap and wait + /// forever for ordinal 0. + pub fn try_pop_in_order(&self) -> Option { + self.try_pop_in_order_reporting_blocked().0 + } + + /// Like [`try_pop_in_order`](Self::try_pop_in_order), but also reports whether + /// a `None` result was *reorder-blocked* — later ordinals are buffered while + /// the stage waits for an earlier one — rather than genuinely drained. The + /// pop and the reorder-blocked check are computed under the SAME state lock, + /// so a shared (`Parallel`) consumer cannot observe a torn `(item, blocked)` + /// pair that would skew the `pop_empties` starvation metric. The flag is + /// always `false` when an item is returned. + /// + /// # Panics + /// + /// Panics if the in-order ordinal counter would overflow `u64` (i.e. + /// `next_serial == u64::MAX`). Branch ordinals start at 0 and increment, so + /// this is unreachable on any real workload (~1.8e19 items on one edge); the + /// guard exists only to fail loudly rather than silently wrap and wait + /// forever for ordinal 0. + pub fn try_pop_in_order_reporting_blocked(&self) -> (Option, bool) { + let mut state = self.state.lock(); + let next = state.next_serial; + + // Drain transport into buffer until we see next_serial or transport + // is empty. This is the only place transport → buffer movement + // happens; it's the visibility synchronizer between producer pushes + // and consumer reads. Each move costs one `heap_size()` call, + // cached alongside the item so subsequent pop / cap accounting + // doesn't pay it again. + if !state.buffer.contains_key(&next) { + while let Some(seq) = self.transport.try_pop() { + // `usize → u64` is a lossless widen on every supported + // (≤64-bit) target, so this cast never truncates. + let bytes = seq.item.heap_size() as u64; + // See the drain-path insert above: a duplicate ordinal here would + // silently drop the buffered item and corrupt `buffer_bytes`. + // Fail loud in release too (like the `next_serial` overflow guard): + // silently dropping a buffered record is a data-integrity bug. + assert!( + !state.buffer.contains_key(&seq.ordinal), + "duplicate reorder ordinal {} — ByItemOrdinal upstream serials must be unique", + seq.ordinal + ); + state.buffer.insert(seq.ordinal, (seq.item, bytes)); + state.buffer_bytes = state.buffer_bytes.saturating_add(bytes); + if state.buffer.contains_key(&next) { + break; + } + } + } + + if let Some((item, item_bytes)) = state.buffer.remove(&next) { + state.buffer_bytes = state.buffer_bytes.saturating_sub(item_bytes); + // Advance the in-order cursor. `ByOrdinal` ordinals come from an + // `AtomicU64` allocator and `ByItemOrdinal` from upstream item + // ordinals; both start at 0 and increment, so `u64::MAX` is + // unreachable on any real workload (~1.8e19 items on one edge). + // `checked_add` makes that invariant explicit: a wrap here would + // silently wait forever for ordinal 0, so we fail loudly instead. + let new_next = next.checked_add(1).expect("reorder ordinal overflow (next_serial)"); + state.next_serial = new_next; + // Update the cache: is the NEW next_serial in the buffer? + // Producers reading post-update see the right state. + let new_buffered = state.buffer.contains_key(&new_next); + self.next_serial_buffered.store(new_buffered, Ordering::Release); + (Some(item), false) + } else { + // No in-order item. Reorder-blocked (backlog, not starvation) iff the + // buffer still holds out-of-order items awaiting an earlier ordinal — + // `next` was just confirmed absent, so any remaining entry is a later + // ordinal. Computed here under the same lock as the pop above. + let reorder_blocked = !state.buffer.is_empty(); + (None, reorder_blocked) + } + } + + /// True iff transport is drained, transport is empty, and the reorder + /// buffer is empty. Sticky: once observed true, stays true. + pub fn is_drained(&self) -> bool { + if self.drained_observed.load(Ordering::Acquire) { + return true; + } + if !self.transport.is_drained() { + return false; + } + if !self.transport.is_empty() { + return false; + } + let buf_empty = self.state.lock().buffer.is_empty(); + if buf_empty { + self.drained_observed.store(true, Ordering::Release); + } + buf_empty + } + + /// Producer-side: mark transport drained. Once drained-and-empty is + /// observed by a consumer, `is_drained()` returns true. + pub fn mark_drained(&self) { + self.transport.mark_drained(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::queues::CountBoundedQueue; + + fn make_stage(transport_capacity: usize) -> ReorderStage { + let q: Arc>> = + Arc::new(CountBoundedQueue::>::new(transport_capacity)); + ReorderStage::new(q) + } + + #[test] + fn pops_in_serial_order_regardless_of_push_order() { + let s = make_stage(8); + s.try_push(2, 200).unwrap(); + s.try_push(0, 100).unwrap(); + s.try_push(1, 150).unwrap(); + assert_eq!(s.try_pop_in_order(), Some(100)); + assert_eq!(s.try_pop_in_order(), Some(150)); + assert_eq!(s.try_pop_in_order(), Some(200)); + assert_eq!(s.try_pop_in_order(), None); + } + + /// Q1 (audit D4): a `ByItemOrdinal` upstream supplies each item's serial, so + /// a buggy step could emit two items with the same ordinal. Inserting the + /// second would silently drop the first and leak its bytes into + /// `buffer_bytes`; the always-on `assert!` must turn that into a loud failure + /// (in release builds too). + #[test] + #[should_panic(expected = "duplicate reorder ordinal")] + fn duplicate_ordinal_trips_assert() { + let s = make_stage(8); + s.try_push(1, 100).unwrap(); + s.try_push(1, 200).unwrap(); // duplicate ordinal (simulated upstream bug) + s.try_push(0, 0).unwrap(); + // Draining moves both ordinal-1 items from the transport into the + // in-order buffer; the second insert hits the duplicate guard. + let _ = s.try_pop_in_order(); + } + + /// Sibling of `duplicate_ordinal_trips_assert`, covering the OTHER always-on + /// duplicate-ordinal guard: the must-accept *stash-insert* path in + /// `try_push_inner` (transport full → overflow into `state.buffer`), not the + /// transport-drain path in `try_pop_in_order_reporting_blocked`. + /// + /// With transport capacity 1 and `next_serial` (0) still absent, every push + /// is must-accept: the first fills the transport, and each later push finds + /// the transport full and overflows into the stash. The second and third + /// pushes both carry ordinal 1, so the second stashes it and the third trips + /// the stash-insert `assert!` — no `try_pop_in_order` runs, so the drain-path + /// guard is never reached. This test fails if that stash-insert guard is + /// removed. + #[test] + #[should_panic(expected = "duplicate reorder ordinal")] + fn duplicate_ordinal_trips_stash_assert() { + let s = make_stage(1); // capacity 1 so the transport fills and later items stash + s.try_push(1, 100).unwrap(); // ordinal 1 -> transport (fills capacity-1 transport) + s.try_push(1, 200).unwrap(); // transport full -> ordinal 1 stashed into buffer + // Transport still full and ordinal 1 already in the stash: the overflow + // insert hits the duplicate guard on the push path. + let _ = s.try_push(1, 300); + } + + #[test] + fn waits_for_next_serial() { + let s = make_stage(8); + s.try_push(1, 100).unwrap(); + s.try_push(2, 200).unwrap(); + assert_eq!(s.try_pop_in_order(), None); + s.try_push(0, 0).unwrap(); + assert_eq!(s.try_pop_in_order(), Some(0)); + assert_eq!(s.try_pop_in_order(), Some(100)); + assert_eq!(s.try_pop_in_order(), Some(200)); + } + + #[test] + fn must_accept_overflows_when_transport_full() { + // Transport capacity 1. Push 4 items in arrival order; all must be + // accepted because we're waiting for serial 0. + let s = make_stage(1); + s.try_push(3, 30).unwrap(); + s.try_push(2, 20).unwrap(); + s.try_push(1, 10).unwrap(); + s.try_push(0, 0).unwrap(); + // Now drain in order. + assert_eq!(s.try_pop_in_order(), Some(0)); + assert_eq!(s.try_pop_in_order(), Some(10)); + assert_eq!(s.try_pop_in_order(), Some(20)); + assert_eq!(s.try_pop_in_order(), Some(30)); + } + + #[test] + fn reporting_pop_flags_reorder_blocked_vs_drained() { + // The reporting pop returns the reorder-blocked flag from the same locked + // path as the pop itself, so a shared consumer sees a consistent pair. + let s = make_stage(8); + // Buffer a later ordinal while ordinal 0 is absent → reorder-blocked. + s.try_push(1, 100).unwrap(); + assert_eq!(s.try_pop_in_order_reporting_blocked(), (None, true), "blocked, not drained"); + // Ordinal 0 arrives; the pop yields it and is not blocked. + s.try_push(0, 0).unwrap(); + assert_eq!(s.try_pop_in_order_reporting_blocked(), (Some(0), false)); + assert_eq!(s.try_pop_in_order_reporting_blocked(), (Some(100), false)); + // Genuinely drained now → None and NOT reorder-blocked (a true empty pop). + assert_eq!(s.try_pop_in_order_reporting_blocked(), (None, false), "drained, not blocked"); + } + + #[test] + fn push_metrics_count_stashed_item_as_push_not_reject() { + use crate::runtime::metrics::EdgeMetrics; + // Transport capacity 1, next_serial (0) absent → must-accept. The first + // push fills the transport; the second must-accept-overflows into the + // stash, returning Ok. Recording at the ReorderStage boundary must count + // BOTH as pushes and NEITHER as a rejection — the transport's internal + // `Err` on the stashed push is not a real backpressure event. (Recording + // on the transport, as before, miscounted the stashed item as a reject.) + let m = EdgeMetrics::new(); + let q: Arc>> = Arc::new(CountBoundedQueue::new(1)); + let stage = ReorderStage::new(q).with_push_metrics(Some(Arc::clone(&m)), false); + assert!(stage.try_push(1, 10).is_ok(), "ordinal 1 into transport"); + assert!(stage.try_push(2, 20).is_ok(), "ordinal 2 stashed (transport full)"); + let s = m.snapshot(); + assert_eq!(s.pushed_items, 2, "both accepted pushes counted"); + assert_eq!(s.push_rejections, 0, "a stashed push is not a rejection"); + assert_eq!(s.pushed_bytes, 0, "count-bounded edge records 0 push bytes"); + } + + #[test] + fn push_metrics_record_item_bytes_when_byte_bounded() { + use crate::runtime::metrics::EdgeMetrics; + #[derive(Debug)] + struct Heavy; + impl HeapSize for Heavy { + fn heap_size(&self) -> usize { + 100 + } + } + // With `record_item_bytes = true` (byte-bounded edge), each accepted push + // records its `heap_size` — so `pushed_bytes` tracks the bare `T`, the + // same size the pop side records. + let m = EdgeMetrics::new(); + let q: Arc>> = Arc::new(CountBoundedQueue::new(8)); + let stage = ReorderStage::new(q).with_push_metrics(Some(Arc::clone(&m)), true); + assert!(stage.try_push(0, Heavy).is_ok()); + assert!(stage.try_push(1, Heavy).is_ok()); + let s = m.snapshot(); + assert_eq!(s.pushed_items, 2); + assert_eq!(s.pushed_bytes, 200, "byte-bounded edge records heap_size per push"); + } + + #[test] + fn push_without_metrics_skips_heap_size() { + use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; + // On the default instrumentation-off path (`push_metrics == None`) the + // push wrapper must NOT compute `heap_size` — that byte figure is only + // needed to record a push, so computing it would be pure hot-path cost. + #[derive(Debug)] + struct Counted(Arc); + impl HeapSize for Counted { + fn heap_size(&self) -> usize { + self.0.fetch_add(1, AtomicOrdering::Relaxed); + 0 + } + } + let calls = Arc::new(AtomicUsize::new(0)); + let q: Arc>> = Arc::new(CountBoundedQueue::new(8)); + // Byte-bounded edge (`record_item_bytes = true`) but no push metrics. + let stage = ReorderStage::new(q).with_push_metrics(None, true); + // ordinal 0 == next_serial → transport accepts (no stash), so the only + // `heap_size` call would be the wrapper's — which the gate must skip. + stage.try_push(0, Counted(Arc::clone(&calls))).unwrap(); + assert_eq!( + calls.load(AtomicOrdering::Relaxed), + 0, + "heap_size must not be computed when push metrics are disabled" + ); + } + + #[test] + fn max_overflow_bytes_caps_must_accept_buffer() { + // Use `Heavy(u32)` carrying a pretend heap size of 100 bytes per + // item. Transport capacity 1, byte cap 200. Buffer accepts items + // until `buffer_bytes >= 200`, then rejects (except for + // `next_serial`, which always gets in). + #[derive(Debug, PartialEq)] + struct Heavy(u32); + impl HeapSize for Heavy { + fn heap_size(&self) -> usize { + 100 + } + } + let q: Arc>> = + Arc::new(CountBoundedQueue::>::new(1)); + let s = ReorderStage::with_max_overflow_bytes(q, 200); + + // Fill transport (capacity 1). + s.try_push(3, Heavy(30)).unwrap(); + // Overflow into buffer: ordinal 4 (100 B). buffer_bytes = 100, + // < cap=200, accepted. + s.try_push(4, Heavy(40)).unwrap(); + // Ordinal 5 lands in buffer too (buffer_bytes = 100 at the cap + // check, before insert; insert lifts it to 200). + s.try_push(5, Heavy(50)).unwrap(); + // Pushing ordinal 6: buffer_bytes = 200 >= cap. Reject. Producer's + // job to retry via held_slot. + assert_eq!(s.try_push(6, Heavy(60)), Err((6, Heavy(60)))); + // But the next_serial=0 ALWAYS gets in (liveness exemption). + s.try_push(0, Heavy(0)).unwrap(); + // Drain to advance next_serial. + assert_eq!(s.try_pop_in_order(), Some(Heavy(0))); + // next_serial is now 1; ordinal 1 IS the new next_serial → must + // always be accepted regardless of byte cap. + s.try_push(1, Heavy(10)).unwrap(); + assert_eq!(s.try_pop_in_order(), Some(Heavy(10))); + s.try_push(2, Heavy(20)).unwrap(); + assert_eq!(s.try_pop_in_order(), Some(Heavy(20))); + assert_eq!(s.try_pop_in_order(), Some(Heavy(30))); + assert_eq!(s.try_pop_in_order(), Some(Heavy(40))); + assert_eq!(s.try_pop_in_order(), Some(Heavy(50))); + // After buffer drains, we can finally push the rejected 6. + s.try_push(6, Heavy(60)).unwrap(); + assert_eq!(s.try_pop_in_order(), Some(Heavy(60))); + } + + #[test] + fn drain_into_stash_respects_cap_on_success_path() { + // Regression for the zipper reorder-stash OOM (issue #330). When + // `next_serial` is withheld by a lagging producer, the consumer's + // `try_pop_in_order` drain loop relocates the *entire* transport into + // the stash hunting for the absent serial. If the must-accept push + // only consults the byte cap on the transport-FULL path, the cap never + // fires — the consumer keeps transport non-full — and the stash grows + // without bound (measured: 7.7 GB on a 60M-read zipper run). The cap + // must gate non-`next_serial` pushes regardless of transport room. + #[derive(Debug, PartialEq)] + struct Heavy(u32); + impl HeapSize for Heavy { + fn heap_size(&self) -> usize { + 100 + } + } + + // Transport capacity 1 item (≤100 B in flight); stash byte cap 200 B. + let q: Arc>> = + Arc::new(CountBoundedQueue::>::new(1)); + let s = ReorderStage::with_max_overflow_bytes(q, 200); + + // Serial 0 is never pushed (the lagging worker). Push ordinals 1..=50 + // in order, polling the consumer after each push — the poll returns + // None (serial 0 absent) but drains transport into the stash as a side + // effect, which is the relocation that bypassed the cap. Track the peak + // stash size after each drain. + let mut rejected = 0u32; + let mut peak_stash = 0u64; + for ord in 1u32..=50 { + if s.try_push(u64::from(ord), Heavy(ord)).is_err() { + rejected += 1; + } + assert_eq!(s.try_pop_in_order(), None, "serial 0 absent → no item releases"); + peak_stash = peak_stash.max(s.current_buffer_bytes()); + } + + // The cap must fire (without the fix, the success path bypasses it + // entirely and `rejected` stays 0). + assert!(rejected > 0, "stash cap never fired: drain-into-stash bypassed the byte cap"); + // Peak stash is bounded by the cap (200 B) plus at most one + // transport-worth (100 B) of overshoot. + assert!(peak_stash <= 300, "stash grew past cap + one transport-worth: {peak_stash} B"); + } + + #[test] + fn setter_tightens_cap_and_tiny_cap_preserves_liveness() { + #[derive(Debug, PartialEq)] + struct Heavy(u32); + impl HeapSize for Heavy { + fn heap_size(&self) -> usize { + 100 + } + } + + // (a) The `ReorderCapHandle` setter tightens an initially-unbounded + // stage's cap (this is how `apply_initial_queue_budget` sizes it). + let q: Arc>> = + Arc::new(CountBoundedQueue::>::new(1)); + let s = ReorderStage::new(q); + s.set_max_overflow_bytes(150); + s.try_push(3, Heavy(30)).unwrap(); // -> transport (cap 1) + s.try_push(4, Heavy(40)).unwrap(); // overflow buffer: 0 < 150 + s.try_push(5, Heavy(50)).unwrap(); // overflow buffer: 100 < 150 + // buffer_bytes is now 200 >= 150 → a non-next push is rejected, + // proving the setter's value took effect. + assert_eq!(s.try_push(6, Heavy(60)), Err((6, Heavy(60))), "tightened cap rejects"); + + // (b) Liveness for any cap (§10.3): a near-zero (1-byte) cap, far below + // any item, still drains every item in serial order with none lost. + // The `next_serial` exemption guarantees progress; the round bound + // catches a livelock blow-up (not just a hard deadlock). + let q2: Arc>> = + Arc::new(CountBoundedQueue::>::new(1)); + let s2 = ReorderStage::new(q2); + s2.set_max_overflow_bytes(1); + let n = 16u32; + let mut pending: Vec<(u64, Heavy)> = + (0..n).rev().map(|i| (u64::from(i), Heavy(i))).collect(); + let mut out: Vec = Vec::new(); + let mut rounds = 0u32; + while out.len() < n as usize { + rounds += 1; + assert!(rounds <= 4 * n, "tiny cap must converge (no livelock blow-up)"); + // Retry every pending push (held-slot semantics), keeping rejects. + let mut still = Vec::new(); + for (ord, item) in pending.drain(..) { + if let Err(rej) = s2.try_push(ord, item) { + still.push(rej); + } + } + pending = still; + while let Some(Heavy(v)) = s2.try_pop_in_order() { + out.push(v); + } + } + assert_eq!(out, (0..n).collect::>(), "all drain in serial order, none lost"); + } + + #[test] + fn concurrent_tiny_cap_drains_all_in_order() { + // The load-bearing liveness test (§10.3/§10.5): at t>1, deadlock-freedom + // is "not proven by the next_serial exemption alone" — the worry is all + // producers wedged on rejected later-ordinal pushes while next_serial is + // unpushed. Here N producer threads concurrently push a disjoint, + // interleaved set of ordinals into one stage with a 1-byte cap (far + // below any item), retrying rejects (held-slot semantics); a consumer + // drains. All items must emerge in serial order with none lost. Progress + // is guaranteed because the producer owning `next_serial` always has it + // as its current (ascending) push, and `next_serial` is cap-exempt. + use std::thread; + + // Heavy carries its ordinal as `usize` (identity) with a fixed heap + // size, so the byte cap binds without any narrowing casts in the test. + #[derive(Debug, PartialEq)] + struct Heavy(usize); + impl HeapSize for Heavy { + fn heap_size(&self) -> usize { + 100 + } + } + + let n_threads = 4usize; + let per = 64usize; + let total = n_threads * per; + let q: Arc>> = + Arc::new(CountBoundedQueue::>::new(4)); + let stage = Arc::new(ReorderStage::::new(q)); + stage.set_max_overflow_bytes(1); // 1 byte ≪ any item + + let producers: Vec<_> = (0..n_threads) + .map(|t| { + let s = Arc::clone(&stage); + thread::spawn(move || { + // This thread owns ordinals {t, t+N, t+2N, ...}, pushed in + // ascending order (each retried until accepted). + let mine: Vec = (0..per).map(|k| t + k * n_threads).collect(); + let mut i = 0; + while i < mine.len() { + let ord = mine[i]; + match s.try_push(ord as u64, Heavy(ord)) { + Ok(()) => i += 1, + Err(_) => thread::yield_now(), // held-slot retry + } + } + }) + }) + .collect(); + + // Deadline guard: this test proves deadlock-freedom, so a liveness + // regression must *fail fast* rather than hang the whole `nextest` run + // until a global harness timeout (if any) fires. 30s is generous — + // 256 tiny items drain in milliseconds when healthy — so it only trips + // on a genuine wedge, not on a slow CI host. Mirrors the + // bounded-convergence guard the single-threaded sibling uses. + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(30); + let mut out: Vec = Vec::with_capacity(total); + while out.len() < total { + assert!( + std::time::Instant::now() < deadline, + "concurrent reorder drain did not complete within 30s ({} of {total} drained) — \ + likely a deadlock/livelock regression", + out.len(), + ); + match stage.try_pop_in_order() { + Some(Heavy(v)) => out.push(v), + None => thread::yield_now(), + } + } + // Bound producer completion with the same deadline. The drain loop above + // only proves the *consumer* made progress; if a regression let + // `out.len()` reach `total` (e.g. a double-emit) while a producer is + // still stuck in its `try_push` retry loop, an unbounded `join()` would + // hang the run forever. Poll `is_finished()` against the deadline so a + // stuck producer fails fast instead. + for p in producers { + while !p.is_finished() { + assert!( + std::time::Instant::now() < deadline, + "a producer thread did not finish within 30s — likely a stuck try_push \ + retry loop (liveness regression)", + ); + thread::yield_now(); + } + p.join().unwrap(); + } + assert_eq!(out, (0..total).collect::>(), "all drain in serial order, none lost"); + } + + #[test] + fn drained_propagates() { + let s = make_stage(4); + s.try_push(0, 0).unwrap(); + s.try_push(1, 1).unwrap(); + s.mark_drained(); + // Not drained while items remain (need at least one pop to hydrate + // the buffer/transport visibility test). + assert!(!s.is_drained()); + let _ = s.try_pop_in_order().unwrap(); + let _ = s.try_pop_in_order().unwrap(); + assert!(s.is_drained()); + } + + #[test] + fn backpressure_applies_after_next_serial_buffered() { + // Drive the stage into a state where the reorder buffer holds + // `next_serial`, then verify producers see backpressure rejection. + // + // Step 1: push ordinal 1 with next_serial=0 still pending. + // transport: [(1,1)], buf: {}, next_serial=0. + // Step 2: a consumer call can't return anything (still waiting for 0), + // but it drains transport into buf as a side effect. + // transport: [], buf: {1}, next_serial=0. + // Step 3: push ordinal 0 — must_accept; goes to transport. + // transport: [(0,0)], buf: {1}, next_serial=0. + // Step 4: pop returns 0; advances next_serial=1; buf still has 1. + // transport: [], buf: {1}, next_serial=1. + // Step 5: now buf has next_serial=1, so backpressure path is active. + // Pushes 2, 3 fill transport (capacity 2). Push 4 rejects. + let s = make_stage(2); + s.try_push(1, 100).unwrap(); + assert_eq!(s.try_pop_in_order(), None); // drains transport into buf as side effect + s.try_push(0, 0).unwrap(); + assert_eq!(s.try_pop_in_order(), Some(0)); + // buf now contains next_serial=1; backpressure path active. + s.try_push(2, 200).unwrap(); + s.try_push(3, 300).unwrap(); + // Transport full, buf has next_serial=1: producer 4 must be rejected. + let result = s.try_push(4, 400); + assert!(result.is_err(), "expected backpressure rejection"); + let (ord, item) = result.unwrap_err(); + assert_eq!((ord, item), (4, 400)); + } +} diff --git a/crates/fgumi-pipeline-core/src/runtime/contexts.rs b/crates/fgumi-pipeline-core/src/runtime/contexts.rs new file mode 100644 index 000000000..1265b85ad --- /dev/null +++ b/crates/fgumi-pipeline-core/src/runtime/contexts.rs @@ -0,0 +1,644 @@ +//! `ChainContexts`: per-step typed input + outputs handles, constructed +//! at `Pipeline::run` start by walking `ChainGraph` and calling each +//! step's `build_input_handle` / `build_output_set`. +//! +//! Layout: +//! - `inputs[step_idx]` — `Box` carrying `BranchInputHandle` +//! (or a dummy unit handle for sources). The worker loop hands this to +//! `ErasedStepCtx.input`. +//! - `outputs[step_idx]` — `Box` carrying `OutputHandles`. +//! The worker loop hands this to `ErasedStepCtx.outputs` AND uses it +//! for the typed `mark_outputs_drained` dispatch via the step's +//! `ErasedStep::mark_outputs_drained` method. +//! +//! Since `OutputsViewAny` no longer carries per-branch drained-flag Arcs +//! (option (c) pivot — drain marking is per-branch via +//! `BranchOutputHandle::mark_drained`, dispatched through the typed view), +//! the framework's drain-propagation path goes through +//! `ErasedStep::mark_outputs_drained(outputs_any.as_ref())`. + +use std::any::Any; +use std::sync::Arc; + +use crate::erased::ErasedStep; +use crate::handles::{BranchInputHandle, OutputQueueSet}; +use crate::item::HeapSize; +use crate::topology::{BranchIdx, ChainGraph, StepIdx}; + +/// Per-step input + outputs handles. The worker loop hands these to +/// `ErasedStepCtx` for each `try_run_erased` dispatch and to +/// `ErasedStep::mark_outputs_drained` for the drain-propagation path. +pub struct ChainContexts { + /// `inputs[step_idx]` = boxed `BranchInputHandle` (or a dummy + /// drained unit handle for sources). + pub inputs: Vec>, + /// `outputs[step_idx]` = boxed `OutputHandles` (the typed + /// surface step authors push into; also dispatchable via + /// `ErasedStep::mark_outputs_drained`). + pub outputs: Vec>, + /// Registry of every byte-bounded queue in the chain. Populated + /// during `build_chain_contexts` from each `Branch`'s + /// `bounded_queue_handle`. Empty when no byte-bounded queues + /// exist (e.g. a chain composed entirely of `CountBounded` / + /// `Unbounded` steps). + pub bounded_queues: Vec, + /// Registry of every instrumented edge (`--pipeline-trace`), with its + /// `EdgeMetrics` + producer/consumer + (for byte-bounded edges) a depth + /// source. Empty when instrumentation is `Off`. Read by the occupancy + /// sampler and the end-of-run edge report. + pub edges: Vec, +} + +/// One instrumented edge: its shared [`EdgeMetrics`](crate::runtime::metrics::EdgeMetrics) +/// (push counts from the transport, pop counts from the input handle), its +/// producer + consumer steps, and a depth source for occupancy sampling. +pub struct RegisteredEdge { + pub producer_step: StepIdx, + pub producer_name: &'static str, + /// `None` for a terminal branch with no consumer (e.g. a `--rejects` tail). + pub consumer_step: Option, + pub consumer_name: Option<&'static str>, + pub branch: BranchIdx, + /// On an **ordered** edge, the push-side counters (`pushed_items` / + /// `pushed_bytes` / `push_rejections`) are recorded at the `ReorderStage` + /// boundary, not on the transport queue: the reorder stage's must-accept path + /// turns a full-transport reject into an accepted (stashed) push, so recording + /// on the transport would miscount stashed items as backpressure. Recorded at + /// the boundary, `pushed_bytes` counts the bare `T` (`heap_size`) and so + /// matches `popped_bytes` (the ordinal tag adds no heap). + pub metrics: std::sync::Arc, + /// `Some` for byte-bounded edges — the transport queue's `current_bytes` + /// (plus, for an ordered edge, the `reorder_depth` stash bytes) over + /// `limit_bytes` gives the occupancy fraction the sampler records. `None` + /// for count/unbounded edges (counters still apply; occupancy histogram is + /// skipped). + pub depth_source: Option>, + /// `Some` for an **ordered** byte-bounded edge — the `ReorderStage`'s + /// overflow-stash handle. Its buffered bytes are added to the transport's + /// `current_bytes` when sampling depth, so an ordered edge reflects total + /// buffered bytes (transport + reorder stash) rather than the transport + /// queue alone (which under-reports occupancy when items pile in the reorder + /// buffer under producer skew). `None` for direct/count/unbounded edges. + pub reorder_depth: Option>, +} + +/// One byte-bounded queue's location in the chain plus the handles for +/// updating its budget: the transport-limit setter (`handle`) and, for an +/// ordered branch, the reorder stage's overflow-cap setter (`reorder_cap`). +/// Used by the budget pass (which sets both from one per-edge budget) and +/// the queue-memory rebalancer (which shifts transport budget at runtime). +pub struct RegisteredQueue { + pub producer_step_name: &'static str, + pub producer_step: StepIdx, + pub branch: BranchIdx, + pub handle: std::sync::Arc, + /// `Some` for an ordered byte-bounded branch — its reorder overflow stash + /// is sized from the same per-edge budget as `handle`. `None` for a direct + /// (unordered) byte-bounded branch (no reorder stage). + pub reorder_cap: Option>, +} + +/// Build the chain's per-step contexts. +/// +/// Walks `graph` to discover each step's producer and pulls the typed input +/// handle out of the producer's `OutputQueueSet`. Sources get a dummy +/// pre-drained `BranchInputHandle<()>` (their input is implicitly drained +/// from the start; framework never pops from it). +/// +/// # Panics +/// +/// Panics if `graph.n_steps() != steps.len()` or if any non-source step +/// has no producer (the builder's all-wired check should prevent this). +#[must_use] +pub fn build_chain_contexts( + steps: &[Box], + graph: &ChainGraph, + level: crate::builder::InstrumentationLevel, +) -> ChainContexts { + build_chain_contexts_inner(steps, graph, false, level) +} + +/// Build the chain's per-step contexts with **direct** inter-step transports +/// (no reorder stages, profile queue bounds retained) — the wiring the +/// single-thread *fused* driver runs over. +/// +/// Identical to [`build_chain_contexts`] except every producer's output set is +/// built via [`ErasedStep::build_fused_output_set`] instead of +/// [`ErasedStep::build_output_set`]. At one worker FIFO push order is already +/// the correct order, so dropping the reorder stage is sound; the count/byte +/// bound stays, so a producer that outruns its consumer within a pass gets +/// backpressure rather than growing the edge without limit. Only call this for a +/// linear (single source → single sink, single-input) chain at `--threads 1`; +/// see [`super::run_fused_single_thread`]. +#[must_use] +pub fn build_chain_contexts_fused( + steps: &[Box], + graph: &ChainGraph, +) -> ChainContexts { + // Fusion is only chosen when instrumentation is off (see + // `should_fuse_single_thread`), and the fused path registers no edges, so + // force `Off` regardless of the run's level. + build_chain_contexts_inner(steps, graph, true, crate::builder::InstrumentationLevel::Off) +} + +/// Build the producer→consumer edge map (keyed by `(producer_step, branch)`) +/// used to label registered edges with their consumer. Empty when +/// instrumentation is off (no edges are registered, so the map is unused). +fn build_consumer_map( + steps: &[Box], + graph: &ChainGraph, + level: crate::builder::InstrumentationLevel, +) -> std::collections::HashMap<(usize, usize), usize> { + let mut m = std::collections::HashMap::new(); + if !level.is_on() { + return m; + } + for (consumer_idx, step) in steps.iter().enumerate() { + if step.is_source() { + continue; + } + match step.input_arity() { + 1 => { + let (p, b) = find_producer(graph, StepIdx(consumer_idx)); + m.insert((p.0, b.0), consumer_idx); + } + 2 => { + for (p, b) in find_all_producers(graph, StepIdx(consumer_idx)) { + m.insert((p.0, b.0), consumer_idx); + } + } + _ => {} + } + } + m +} + +fn build_chain_contexts_inner( + steps: &[Box], + graph: &ChainGraph, + direct: bool, + level: crate::builder::InstrumentationLevel, +) -> ChainContexts { + assert_eq!(steps.len(), graph.n_steps(), "chain length mismatch"); + + let n_steps = steps.len(); + let mut output_sets: Vec = Vec::with_capacity(n_steps); + let mut outputs: Vec> = Vec::with_capacity(n_steps); + let mut inputs: Vec> = Vec::with_capacity(n_steps); + + // Pass 1: build each step's output set + boxed `OutputHandles`. The + // fused driver uses direct transports with the profile's queue bounds + // (`build_fused_output_set`, which only drops the ordering); the scheduled + // path honours each step's profiled queue spec + ordering. + for step in steps { + let (set, view) = + if direct { step.build_fused_output_set(level) } else { step.build_output_set(level) }; + let outputs_box = step.wrap_outputs_view(view); + output_sets.push(set); + outputs.push(outputs_box); + } + + // Producer→consumer map for edge registration (inverse of `find_producer`, + // which Pass 2 uses consumer→producer). Empty when instrumentation is off. + let consumer_of = build_consumer_map(steps, graph, level); + + // Pass 1.5: collect byte-bounded queue handles into the registry. + // Must run before Pass 2 because `take_typed_input` (called via + // `build_input_handle`) replaces the `BranchEntry` with a fresh + // one whose `bounded_queue_handle` is `None` — by then the + // handles have been moved out of the chain. + let mut bounded_queues: Vec = Vec::new(); + let mut edges: Vec = Vec::new(); + for (step_idx_usize, set) in output_sets.iter().enumerate() { + for (branch_idx_usize, entry) in set.branches.iter().enumerate() { + if let Some(handles) = entry.bounded_queue_handle.as_ref() { + bounded_queues.push(RegisteredQueue { + producer_step_name: steps[step_idx_usize].profile().name, + producer_step: StepIdx(step_idx_usize), + branch: BranchIdx(branch_idx_usize), + handle: std::sync::Arc::clone(&handles.transport), + reorder_cap: handles.reorder_cap.clone(), + }); + } + // Instrumented edge: register its shared metrics + producer/consumer + // + (for byte-bounded edges) a depth source for occupancy sampling. + if let Some(metrics) = entry.metrics.as_ref() { + let consumer_step = + consumer_of.get(&(step_idx_usize, branch_idx_usize)).map(|c| StepIdx(*c)); + edges.push(RegisteredEdge { + producer_step: StepIdx(step_idx_usize), + producer_name: steps[step_idx_usize].profile().name, + consumer_step, + consumer_name: consumer_step.map(|c| steps[c.0].profile().name), + branch: BranchIdx(branch_idx_usize), + metrics: std::sync::Arc::clone(metrics), + depth_source: entry + .bounded_queue_handle + .as_ref() + .map(|h| std::sync::Arc::clone(&h.transport)), + reorder_depth: entry + .bounded_queue_handle + .as_ref() + .and_then(|h| h.reorder_cap.clone()), + }); + } + } + } + + // Pass 2: input handles. Sources get a dummy drained unit handle; + // single-input mid-steps and sinks pull one handle from their + // producer's `OutputQueueSet`; two-input merge steps (`Step2` + // adapters, identified by `input_arity() == 2`) pull TWO handles + // — one per consumer-input-slot — and wrap them in a + // `TwoInputHandles`. + for (consumer_idx, step) in steps.iter().enumerate() { + let input_box = if step.is_source() { + // Source step (Input = ()). Build a permanently-drained + // dummy unit handle — the worker loop never pops from it. + // Chains with multiple sources (e.g. zipper's mapped + + // unmapped subchains converging at a Step2 merger) all + // share this code path; the runtime walks every source + // independently to Finished. + dummy_unit_input_handle() + } else { + match step.input_arity() { + 1 => { + let (producer_idx, branch_idx) = find_producer(graph, StepIdx(consumer_idx)); + step.build_input_handle(&mut output_sets[producer_idx.0], branch_idx.0) + } + 2 => { + let edges = find_all_producers(graph, StepIdx(consumer_idx)); + assert_eq!( + edges.len(), + 2, + "Step2 consumer {:?} expects 2 input edges, found {}", + StepIdx(consumer_idx), + edges.len() + ); + let (p0, p0_branch) = edges[0]; + let (p1, p1_branch) = edges[1]; + step.build_two_input_handles( + &mut output_sets, + p0.0, + p0_branch.0, + p1.0, + p1_branch.0, + ) + } + n => panic!("unsupported input_arity {n} for step {:?}", StepIdx(consumer_idx)), + } + }; + inputs.push(input_box); + } + + // `output_sets` is consumed implicitly here — every branch was taken + // exactly once via `build_input_handle`, leaving placeholder slots. + drop(output_sets); + + ChainContexts { inputs, outputs, bounded_queues, edges } +} + +/// Construct a `BranchInputHandle<()>` that's already drained — for source +/// steps whose input is implicitly empty + drained from t=0. Uses the zero-state +/// `always_drained` handle (no backing queue), so building a source costs no +/// `SegQueue` allocation for a handle whose only job is to report +/// `is_drained() == true` (the worker loop never pops from a source's input). +fn dummy_unit_input_handle() -> Box { + Box::new(BranchInputHandle::<()>::always_drained()) +} + +/// Find the single (producer, branch) that produces the given single-input +/// consumer step. +/// +/// # Panics +/// +/// Panics if no producer exists, or if more than one producer branch targets +/// the consumer (the builder's all-wired check should have prevented either on +/// a built pipeline). +fn find_producer(graph: &ChainGraph, consumer: StepIdx) -> (StepIdx, BranchIdx) { + let mut found: Option<(StepIdx, BranchIdx)> = None; + for producer_usize in 0..graph.n_steps() { + let producer = StepIdx(producer_usize); + let n_branches = graph.branch_count(producer); + for branch_usize in 0..n_branches { + let branch = BranchIdx(branch_usize); + if graph.consumer(producer, branch) == Some(consumer) { + // A single-input consumer has exactly one incoming edge. Two + // producer branches wired to it would leave the second one + // unpopped: the caller takes an input handle for the first edge + // only, so the extra branch's queue is never drained and its + // producer wedges on backpressure. Reject the graph instead of + // silently keeping whichever edge was scanned first. Mirrors the + // per-slot uniqueness check in `find_all_producers`. + if let Some((prev_producer, prev_branch)) = found { + panic!( + "single-input consumer {consumer:?} is wired by more than one \ + producer branch ({prev_producer:?} branch {prev_branch:?} and \ + {producer:?} branch {branch:?}); the extra edge would never be \ + popped" + ); + } + found = Some((producer, branch)); + } + } + } + found.unwrap_or_else(|| { + panic!( + "step {consumer:?} has no producer in chain graph; \ + the builder's all-wired check should have caught this" + ) + }) +} + +/// Find all `(producer, producer_branch)` edges feeding the given +/// multi-input consumer, sorted by the consumer's input-slot index +/// (`returned[i]` feeds the consumer's input slot `i`). Used by +/// [`build_chain_contexts`] to assemble per-branch input handles +/// for `Step2` and future `StepN` consumers. +/// +/// # Panics +/// +/// Panics if any of the consumer's input slots is unwired (the +/// builder's all-wired check should have caught this on a built +/// pipeline). +fn find_all_producers(graph: &ChainGraph, consumer: StepIdx) -> Vec<(StepIdx, BranchIdx)> { + let arity = graph.input_arity(consumer); + let mut edges: Vec> = vec![None; arity]; + for producer_usize in 0..graph.n_steps() { + let producer = StepIdx(producer_usize); + let n_branches = graph.branch_count(producer); + for branch_usize in 0..n_branches { + let branch = BranchIdx(branch_usize); + if graph.consumer(producer, branch) == Some(consumer) { + let slot = graph + .consumer_input_slot(producer, branch) + .expect("consumer_input_slot must be set when consumer is wired"); + assert!( + slot < arity, + "consumer {consumer:?} edge has input slot {slot} but arity is {arity}" + ); + assert!( + edges[slot].is_none(), + "consumer {consumer:?} input slot {slot} wired more than once", + ); + edges[slot] = Some((producer, branch)); + } + } + } + edges + .into_iter() + .enumerate() + .map(|(slot, e)| { + e.unwrap_or_else(|| panic!("consumer {consumer:?} input slot {slot} has no producer")) + }) + .collect() +} + +/// Convenience: borrow the typed `BranchInputHandle` for a given step. +/// +/// # Panics +/// +/// Panics if the step's input handle doesn't downcast to `T` (a framework +/// invariant violation). +#[must_use] +pub fn input_as( + contexts: &Arc, + step: StepIdx, +) -> &BranchInputHandle { + contexts.inputs[step.0] + .downcast_ref::>() + .expect("BranchInputHandle downcast failed in input_as") +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io; + + use proptest::prelude::*; + use rstest::rstest; + + use crate::erased::TypedStep; + use crate::outputs::Single; + use crate::queues::QueueSpec; + use crate::reorder::BranchOrdering; + use crate::step::{ + InputHandle, OutputHandles, Step, StepCtx, StepKind, StepOutcome, StepProfile, + }; + + #[derive(Clone)] + struct StubSource; + impl Step for StubSource { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "Source", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::ByOrdinal], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::Finished) + } + } + + #[derive(Clone)] + struct StubSink; + impl Step for StubSink { + type Input = u32; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "Sink", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + /// A `u32 → u32` pass-through step used to grow a linear chain to an + /// arbitrary length between the source and sink. + #[derive(Clone)] + struct MiddleStep; + impl Step for MiddleStep { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "Middle", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + /// Build an `n`-step linear chain `Source → Middle×(n - 2) → Sink`, + /// returning the erased step boxes alongside the fully wired graph. + /// + /// # Panics + /// + /// Panics if `n < 2` (a linear chain needs at least a source and a sink). + fn linear_chain(n: usize) -> (Vec>, ChainGraph) { + assert!(n >= 2, "linear chain needs at least a source and a sink"); + + let mut graph = ChainGraph::new(); + let mut steps: Vec> = Vec::with_capacity(n); + let mut indices: Vec = Vec::with_capacity(n); + + indices.push(graph.register_step("Source", 1)); + steps.push(Box::new(TypedStep::new(StubSource))); + for _ in 1..n - 1 { + indices.push(graph.register_step("Middle", 1)); + steps.push(Box::new(TypedStep::new(MiddleStep))); + } + indices.push(graph.register_step("Sink", 0)); + steps.push(Box::new(TypedStep::new(StubSink))); + + for pair in indices.windows(2) { + graph.wire(pair[0], BranchIdx(0), pair[1]); + } + + (steps, graph) + } + + /// Assert the invariants `build_chain_contexts` must uphold for an + /// `n`-step linear chain: one input/output slot per step, no byte-bounded + /// queues (the stubs only use `CountBounded` transport), the source's + /// input is a dummy pre-drained `BranchInputHandle<()>`, and every + /// downstream step receives a real `BranchInputHandle` wired from its + /// producer. + fn assert_linear_chain_invariants(ctx: &ChainContexts, n: usize) { + assert_eq!(ctx.inputs.len(), n); + assert_eq!(ctx.outputs.len(), n); + assert!(ctx.bounded_queues.is_empty()); + + let source = ctx.inputs[0] + .downcast_ref::>() + .expect("source must have a dummy unit input handle"); + // Pin the `always_drained()` invariant: a source's input is implicitly + // drained from t=0, so a regression that swapped it back to a real + // (never-draining) queue would be caught here, not just the downcast. + assert!( + source.is_drained(), + "source's dummy unit input handle must report drained (always_drained invariant)" + ); + // Identity, not just type: every `MiddleStep` and the sink share the + // type `BranchInputHandle`, so a `find_producer` regression that + // handed several consumers the SAME producer's branch handle would still + // downcast and still pass. Push a per-producer distinct value and assert + // it arrives on that producer's consumer. + for i in 1..n { + let input = ctx.inputs[i] + .downcast_ref::>() + .unwrap_or_else(|| panic!("step {i} must have a BranchInputHandle")); + let tag = 1000 + u32::try_from(i).expect("chain length fits in u32"); + let producer = ctx.outputs[i - 1] + .downcast_ref::>>() + .unwrap_or_else(|| panic!("step {} must have OutputHandles>", i - 1)); + producer.push(tag).expect("producer transport accepts one item"); + assert_eq!( + input.pop(), + Some(tag), + "step {i}'s input must be wired to step {}'s output branch 0", + i - 1 + ); + } + } + + /// `build_chain_contexts` wires linear chains of varying length: the + /// three-step case also exercises `find_all_producers` for a middle step. + #[rstest] + #[case(2)] + #[case(3)] + #[case(4)] + fn build_chain_contexts_linear(#[case] n: usize) { + let (steps, graph) = linear_chain(n); + let ctx = build_chain_contexts(&steps, &graph, crate::builder::InstrumentationLevel::Off); + assert_linear_chain_invariants(&ctx, n); + } + + /// `ChainGraph::wire` only rejects re-wiring the same *producer* branch, so + /// nothing stops two producers from targeting one single-input consumer. + /// `find_producer` must reject that graph rather than return the + /// first-scanned edge: the caller takes an input handle for that edge only, + /// so the second producer's branch is never popped and it wedges on + /// backpressure once its transport fills. + #[test] + #[should_panic(expected = "is wired by more than one producer branch")] + fn find_producer_rejects_duplicate_incoming_edges() { + let mut graph = ChainGraph::new(); + let source_a = graph.register_step("SourceA", 1); + let source_b = graph.register_step("SourceB", 1); + let sink = graph.register_step("Sink", 0); + graph.wire(source_a, BranchIdx(0), sink); + graph.wire(source_b, BranchIdx(0), sink); + let _ = find_producer(&graph, sink); + } + + /// The single-producer case still resolves to that one edge — the + /// uniqueness check must not have turned the happy path into a panic. + #[test] + fn find_producer_resolves_the_single_incoming_edge() { + let (_steps, graph) = linear_chain(3); + assert_eq!(find_producer(&graph, StepIdx(2)), (StepIdx(1), BranchIdx(0))); + } + + #[test] + fn registry_covers_edges_with_producer_and_consumer() { + use crate::builder::InstrumentationLevel as L; + // Source → Middle → Sink: two producing edges, each with a consumer. + let (steps, graph) = linear_chain(3); + let ctx = build_chain_contexts(&steps, &graph, L::Summary); + assert_eq!(ctx.edges.len(), 2, "every bounded producing edge is registered"); + // Identity, not just presence. Labelling each edge with its own consumer is + // the whole job of `build_consumer_map`, and a count-plus-`is_some()` check + // passes just as well if the two consumers were swapped or if both edges + // were labelled with the same one. Pin the pairs. + let mut pairs: Vec<(&str, Option<&str>)> = + ctx.edges.iter().map(|e| (e.producer_name, e.consumer_name)).collect(); + pairs.sort_unstable(); + assert_eq!( + pairs, + vec![("Middle", Some("Sink")), ("Source", Some("Middle"))], + "each edge is labelled with its own consumer" + ); + for e in &ctx.edges { + assert!( + e.consumer_step.is_some(), + "edge from {} has a resolved consumer", + e.producer_name + ); + assert!(e.consumer_name.is_some()); + // CountBounded edges expose no byte depth source (occupancy via len only). + assert!(e.depth_source.is_none(), "count-bounded edge has no byte depth source"); + } + // Off → no edges registered (zero-overhead path). + let (steps, graph) = linear_chain(3); + let off = build_chain_contexts(&steps, &graph, L::Off); + assert!(off.edges.is_empty(), "level Off registers no edges"); + } + + proptest! { + /// The typed-handle and bounded-queue invariants hold for linear + /// chains of any length, not just the hand-picked rstest cases. + #[test] + fn build_chain_contexts_linear_invariants(n in 2usize..=8) { + let (steps, graph) = linear_chain(n); + let ctx = build_chain_contexts(&steps, &graph, crate::builder::InstrumentationLevel::Off); + assert_linear_chain_invariants(&ctx, n); + } + } +} diff --git a/crates/fgumi-pipeline-core/src/runtime/detached.rs b/crates/fgumi-pipeline-core/src/runtime/detached.rs new file mode 100644 index 000000000..176c32040 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/runtime/detached.rs @@ -0,0 +1,799 @@ +//! Detached-step extraction and the dedicated **driver thread** that runs them +//! off the work-stealing pool. +//! +//! A [`StepKind::Detached`] step is excluded +//! from the pool ([`build_worker_storage`](crate::runtime::build_worker_storage) +//! gives every pool worker a `Skip` entry for it) and instead runs on a dedicated +//! OS thread spawned at run start alongside the deadlock-monitor / +//! queue-rebalancer and joined after the workers. This is the legacy sort's +//! "N + 2" threading: N pool workers do the parallel (compression-bound) work +//! while off-pool driver threads do the serial coordination + I/O, so neither +//! steals a pool worker slot. +//! +//! ## The driver IS a 1-thread pool ([`run_detached_driver`]) +//! +//! A driver thread does **not** have a bespoke drive loop. It runs the *same* +//! [`run_worker_loop`] the pool uses — over a +//! purpose-built storage row where its group's steps are `Owned` and every other +//! step is `Skip` ([`build_driver_storage`]) — with a +//! [`WorkerCore::driver`](crate::runtime::WorkerCore) (Park backoff, off-pool +//! stats attribution) and the [`DrainFirstScheduler`]. So a driver is literally a +//! 1-thread (or, for a group, still 1-thread over several `Owned` steps) instance +//! of the pool loop. Several detached steps sharing a +//! [`DetachedGroup::Shared`](crate::step::DetachedGroup) label are driven by ONE +//! thread that round-robins them; [`DetachedGroup::PerStep`] (the default) keeps +//! one thread per step. +//! +//! Because the group's steps and their phases are temporally disjoint (e.g. the +//! sort's phase-1 admit/sort/frame finish and leave the live set before the +//! phase-2 merge runs), one driver thread covers a whole phase's coordination +//! without oversubscription — exactly like main's single main thread. +//! +//! ## Lock-ordering acyclicity (non-negotiable, the deadlock proof) +//! +//! A driver thread must NEVER hold one queue's internal lock while parking on +//! another. `run_worker_loop` upholds that by construction: +//! +//! 1. Each `try_run_erased` is a single non-blocking call. The step body pops +//! from its input transport (a lock-free `crossbeam ArrayQueue` — `try_pop`, +//! no lock held across the call) and pushes to its outputs (`try_push`, +//! likewise); a full output / empty input is reported back as `NoProgress` / +//! `Contention`. It never *blocks* inside `try_run`. +//! 2. The only blocking a driver does is the loop's `WorkerCore::sleep_backoff` +//! (`park_timeout` under the Park policy), which holds NO queue lock. +//! 3. The loop tries EVERY live step in a pass before it parks (round-robin, +//! park only after a full no-progress pass). So when one grouped step is +//! blocked, a sibling on the same driver still runs — a park-on-first-idle +//! loop would wedge (see `driver_round_robins_all_live_before_parking`). +//! +//! So there is no cycle: a driver parks only *between* `try_run` calls, never +//! while holding a transport lock, so a two-sided step (consuming from the pool +//! AND producing to it, both bounded) cannot deadlock. The +//! `detached_two_sided_no_deadlock` test pins this with a wall-clock watchdog. + +use std::any::Any; +use std::collections::HashMap; +use std::sync::Arc; + +use crate::erased::{ErasedStep, ErasedStepCtx}; +use crate::runtime::contexts::ChainContexts; +use crate::runtime::drain::StepDrainCounter; +use crate::runtime::driver::run_worker_loop; +use crate::runtime::scheduler::DrainFirstScheduler; +use crate::runtime::stats::PipelineStats; +use crate::runtime::storage::WorkerStepEntry; +use crate::runtime::worker_core::WorkerCore; +use crate::signal::PipelineSignal; +use crate::step::{Affinity, DetachedGroup, OutputsViewAny, StepKind, StepOutcome, StepProfile}; +use crate::topology::StepIdx; + +/// Sentinel left in the chain's `steps` vec in place of a `Detached` step after +/// its real instance has been extracted for its dedicated thread (see +/// [`extract_detached_steps`]). Keeping a same-position placeholder preserves +/// the `step_idx`-aligned indexing that `build_worker_storage` and +/// `ChainContexts` rely on. `build_worker_storage` only reads `kind()` (matches +/// `Detached` → every worker gets `Skip`) and then drops the box, so the +/// placeholder never has any other method invoked; they panic to catch a +/// framework bug if one ever is. +struct DetachedPlaceholder { + name: &'static str, + /// The real step's group, carried through the swap. + /// + /// Not consulted on the run path — `extract_detached_steps` reads the real + /// step's group *before* installing this placeholder. Preserved because + /// `extract_detached_steps` is `pub` and leaves these placeholders in the + /// caller's slice: an external caller regrouping from that slice would read + /// `PerStep` for a step that declared `Shared(..)` and split one shared + /// group across separate driver threads. `profile()` panics for the same + /// class of reason — returning fabricated metadata misinforms that caller. + group: DetachedGroup, +} + +impl ErasedStep for DetachedPlaceholder { + fn profile(&self) -> StepProfile { + // Panics like every other placeholder method rather than returning empty + // `output_queues` / `branch_ordering`. `Pipeline::run` reads + // `step.profile().kind` for the drain counters BEFORE + // `extract_detached_steps` swaps these in, and everything after reads the + // cached `kind()` / `name()` accessors — so nothing on the run path calls + // this. But `extract_detached_steps` is `pub` and leaves placeholders in + // the caller's slice: silently reporting "no outputs" for a step that + // declares them would misinform any later caller. + panic!( + "DetachedPlaceholder::profile invoked for {:?} — the real instance was extracted; \ + read the cached kind()/name() accessors instead", + self.name + ); + } + fn name(&self) -> &'static str { + self.name + } + fn kind(&self) -> StepKind { + StepKind::Detached + } + fn sticky(&self) -> bool { + false + } + fn affinity(&self) -> Affinity { + Affinity::None + } + fn detached_group(&self) -> DetachedGroup { + // See the field doc: off the run path, but must not report `PerStep` for a + // step that declared `Shared(..)`. + self.group + } + fn try_run_erased(&mut self, _ctx: &mut ErasedStepCtx<'_>) -> std::io::Result { + panic!("DetachedPlaceholder::try_run_erased invoked — the real instance was extracted"); + } + fn clone_boxed(&self) -> Box { + panic!("DetachedPlaceholder::clone_boxed invoked — placeholder is never cloned"); + } + fn build_input_handle( + &self, + _producer_set: &mut crate::handles::OutputQueueSet, + _branch_idx: usize, + ) -> Box { + panic!("DetachedPlaceholder::build_input_handle invoked"); + } + fn build_output_set( + &self, + _level: crate::builder::InstrumentationLevel, + ) -> (crate::handles::OutputQueueSet, OutputsViewAny) { + panic!("DetachedPlaceholder::build_output_set invoked"); + } + fn build_fused_output_set( + &self, + _level: crate::builder::InstrumentationLevel, + ) -> (crate::handles::OutputQueueSet, OutputsViewAny) { + panic!("DetachedPlaceholder::build_fused_output_set invoked"); + } + fn wrap_outputs_view(&self, _view: OutputsViewAny) -> Box { + panic!("DetachedPlaceholder::wrap_outputs_view invoked"); + } + fn mark_outputs_drained(&self, _outputs: &(dyn Any + Send + Sync)) { + panic!("DetachedPlaceholder::mark_outputs_drained invoked"); + } + fn is_source(&self) -> bool { + false + } +} + +/// One dedicated driver thread's worth of extracted detached steps, in chain +/// (`StepIdx`) order. The caller spawns one OS thread per group and drives it +/// with [`run_detached_driver`]. +pub struct DetachedDriverGroup { + /// The steps this one driver thread runs, in chain order. Always non-empty + /// and all [`StepKind::Detached`]. Private so those invariants — enforced by + /// [`Self::new`] — cannot be bypassed by an external caller building the + /// struct directly (which could otherwise trigger `steps[0]` panics in + /// `primary_step` / `label`, or run non-detached work on an off-pool driver). + steps: Vec<(StepIdx, Box)>, +} + +impl DetachedDriverGroup { + /// Wrap a driver thread's extracted steps, enforcing the invariants every + /// consumer relies on: the group is **non-empty** (`primary_step` / `label` + /// index `steps[0]`) and **every step is [`StepKind::Detached`]** (a + /// non-detached step must not run off-pool on a dedicated driver thread, and + /// a `Parallel` step would hang its shared output — see + /// [`build_driver_storage`]). + /// + /// # Panics + /// + /// Panics if `steps` is empty or contains a non-`Detached` step. + #[must_use] + fn new(steps: Vec<(StepIdx, Box)>) -> Self { + assert!(!steps.is_empty(), "detached driver group must be non-empty"); + assert!( + steps.iter().all(|(_, step)| step.kind() == StepKind::Detached), + "detached driver group may only contain Detached steps" + ); + Self { steps } + } + + /// The group's representative step (first in chain order). Used as the + /// driver thread's off-pool stats key (`WorkerCore::driver`) and as a stable + /// label. Non-empty by construction. + #[must_use] + pub fn primary_step(&self) -> StepIdx { + self.steps[0].0 + } + + /// The name of the group's representative step (first in chain order), used + /// to label a `PerStep` driver thread. Non-empty by construction. + #[must_use] + pub fn primary_name(&self) -> &'static str { + self.steps[0].1.name() + } + + /// The group label — derived from the steps (every step in the group reports + /// the same [`DetachedGroup`]), used to name the driver thread. Non-empty by + /// construction. + #[must_use] + pub fn label(&self) -> DetachedGroup { + self.steps[0].1.detached_group() + } + + /// Consume the group, yielding its steps for [`build_driver_storage`]. + #[must_use] + fn into_steps(self) -> Vec<(StepIdx, Box)> { + self.steps + } +} + +/// Remove every [`StepKind::Detached`] step's +/// real instance from `steps`, replacing each in place with a +/// `DetachedPlaceholder` so the surviving slots keep their `step_idx` +/// positions (which `build_worker_storage` and `ChainContexts` index by). +/// +/// Groups the extracted steps by their [`DetachedGroup`]: every +/// [`DetachedGroup::Shared`] label collects onto ONE group (one driver thread); +/// each [`DetachedGroup::PerStep`] step becomes its own singleton group (the +/// legacy one-thread-per-step behavior — the default, so non-sort chains are +/// unchanged). Within a group and across groups, order follows chain order +/// (first appearance). Called by `Pipeline::run` **before** `build_worker_storage` +/// consumes `steps`, while the (read-only) `ChainContexts` have already been +/// built from `&steps`. +#[must_use] +pub fn extract_detached_steps(steps: &mut [Box]) -> Vec { + // Accumulate each driver thread's steps as a raw vec, then wrap through + // `DetachedDriverGroup::new` so the non-empty / all-Detached invariants are + // enforced in one place rather than trusting each construction site. + let mut group_steps: Vec)>> = Vec::new(); + // Shared(label) -> index into `group_steps`, for O(1) append. PerStep steps + // never share, so they are not indexed (each starts its own group). + let mut shared_index: HashMap<&'static str, usize> = HashMap::new(); + for (idx, slot) in steps.iter_mut().enumerate() { + if slot.kind() != StepKind::Detached { + continue; + } + let group = slot.detached_group(); + let placeholder: Box = + Box::new(DetachedPlaceholder { name: slot.name(), group }); + let real = std::mem::replace(slot, placeholder); + let entry = (StepIdx(idx), real); + match group { + DetachedGroup::PerStep => group_steps.push(vec![entry]), + DetachedGroup::Shared(label) => { + if let Some(&gi) = shared_index.get(label) { + group_steps[gi].push(entry); + } else { + shared_index.insert(label, group_steps.len()); + group_steps.push(vec![entry]); + } + } + } + } + group_steps.into_iter().map(DetachedDriverGroup::new).collect() +} + +/// Build a driver thread's storage row: a full-length `Vec` +/// (length `n_total_steps`, indexed by global `step_idx` like every other row) +/// where each of `group_steps` is `Owned` and every other slot is `Skip`. The +/// driver thread runs [`run_worker_loop`] over this row exactly as a pool worker +/// runs over its own row. +/// +/// # Panics +/// +/// - if a group step's kind is `Parallel` — a 1-thread driver's +/// [`StepDrainCounter`] is init 1 (single finisher), which would never close a +/// `Parallel` step's shared output (that needs init N, all clones finishing), +/// hanging the downstream consumer; +/// - if two group steps map to the same `step_idx` (a double registration), or a +/// step index is out of range. +#[must_use] +pub fn build_driver_storage( + group_steps: Vec<(StepIdx, Box)>, + n_total_steps: usize, +) -> Vec { + let mut row: Vec = (0..n_total_steps).map(|_| WorkerStepEntry::Skip).collect(); + for (idx, step) in group_steps { + assert!( + idx.0 < n_total_steps, + "driver group step index {} out of range (chain has {n_total_steps} steps)", + idx.0 + ); + assert_ne!( + step.kind(), + StepKind::Parallel, + "driver group step `{}` is Parallel; a 1-thread driver's StepDrainCounter (init 1) \ + would never close its shared output — group only single-runner steps", + step.name() + ); + assert!( + matches!(row[idx.0], WorkerStepEntry::Skip), + "driver group step index {} registered twice (dual registration)", + idx.0 + ); + row[idx.0] = WorkerStepEntry::Owned { step }; + } + row +} + +/// Drive one [`DetachedDriverGroup`] to completion on the calling (dedicated) +/// thread — the unified "1-thread pool". Builds the group's `Owned`/`Skip` +/// storage row ([`build_driver_storage`]) and runs the *same* +/// [`run_worker_loop`] the pool uses, with a [`WorkerCore::driver`] (Park +/// backoff, off-pool stats) and the [`DrainFirstScheduler`] (drain/seal +/// downstream before producing more — frees the sort's capacity-1 arena fastest). +/// +/// `drain_counters` is the full per-step slice (init 1 for each detached step, so +/// the single finisher closes its output edges — the downstream consumer's +/// end-of-stream signal). On a cancel before `Finished`, outputs are NOT closed +/// (the run is tearing down; the recorded error/cancel is what propagates) — +/// `run_worker_loop`'s top-of-loop `is_done` break upholds this. +pub fn run_detached_driver( + group: DetachedDriverGroup, + contexts: &Arc, + drain_counters: &[Arc], + signal: &Arc, + stats: Option<&Arc>, + liveness: &crate::liveness::LivenessCounter, +) { + let primary = group.primary_step(); + let mut row = build_driver_storage(group.into_steps(), contexts.inputs.len()); + let mut worker = WorkerCore::driver(primary); + run_worker_loop( + &mut worker, + &mut row, + contexts, + drain_counters, + signal, + stats, + liveness, + &DrainFirstScheduler, + ); +} + +#[cfg(test)] +mod tests { + use std::io; + use std::sync::atomic::{AtomicU32, Ordering}; + use std::time::Duration; + + use super::*; + use crate::builder::InstrumentationLevel; + use crate::erased::TypedStep; + use crate::outputs::Single; + use crate::queues::QueueSpec; + use crate::reorder::BranchOrdering; + use crate::step::{InputHandle, OutputHandles, Step, StepCtx, StepKind, StepProfile}; + + /// `() -> u32` source stub used only so `build_output_set` constructs the + /// transport that becomes the Detached step's input edge. Never run. + #[derive(Clone)] + struct SrcStub { + capacity: usize, + } + impl Step for SrcStub { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "Src", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: self.capacity }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + /// `u32 -> u32` pass-through Detached step: pop one item per `try_run`, + /// push it on (holding it on output-full backpressure), report `Finished` + /// once input drains and nothing is held. + #[derive(Clone)] + struct PassThroughDetached { + held: Option, + } + impl Step for PassThroughDetached { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "PassThroughDetached", + kind: StepKind::Detached, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + // A rejected push reports `NoProgress`, not `Contention`: the driver + // treats them identically, but `Contention` means "a Serial step's + // mutex was held by another worker" and feeds `contention_count`, + // which the bottleneck verdict turns into its SPIN ratio. Using it + // for ordinary output backpressure invents contention that never + // happened. + if let Some(v) = self.held.take() { + if ctx.outputs.push(v).is_err() { + self.held = Some(v); + return Ok(StepOutcome::NoProgress); + } + return Ok(StepOutcome::Progress); + } + match ctx.input.pop() { + Some(v) => match ctx.outputs.push(v) { + Ok(()) => Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held = Some(unpushed.into_item()); + Ok(StepOutcome::NoProgress) + } + }, + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } + } + + /// Assemble a `Source -> PassThroughDetached -> Sink` shaped context by + /// hand: build the producer's output set (= the Detached input edge), wire + /// the Detached step's input from it, build the Detached step's own output + /// set (= the downstream consumer's input edge), and return the pieces. + #[allow(clippy::type_complexity)] + fn build_one_detached( + src_capacity: usize, + ) -> ( + Box, // the detached step + Arc, // contexts for step_idx 1 + Arc>, // producer outputs (push side) + crate::handles::BranchInputHandle, // downstream consumer (pop side) + ) { + let producer: Box = + Box::new(TypedStep::new(SrcStub { capacity: src_capacity })); + let (mut producer_set, producer_view) = + producer.build_output_set(InstrumentationLevel::Off); + let producer_outputs_any = producer.wrap_outputs_view(producer_view); + + let det: Box = Box::new(TypedStep::new(PassThroughDetached { held: None })); + let det_input = det.build_input_handle(&mut producer_set, 0); + let (mut det_set, det_view) = det.build_output_set(InstrumentationLevel::Off); + let det_outputs_any = det.wrap_outputs_view(det_view); + let det_output_consumer = det_set.take_typed_input::(0); + + let contexts = Arc::new(ChainContexts { + inputs: vec![Box::new(()), det_input, Box::new(())], + outputs: vec![Box::new(()), det_outputs_any, Box::new(())], + bounded_queues: vec![], + edges: vec![], + }); + (det, contexts, Arc::new(producer_outputs_any), det_output_consumer) + } + + /// Drive a single detached `step` at `step_idx` through the unified driver — + /// a `PerStep` group of one — on the calling thread, with a full-length + /// `drain_counters` slice (init 1 each). Mirrors what `builder.rs` step 4d + /// does for a one-step group. + fn drive_single( + step: Box, + step_idx: StepIdx, + contexts: &Arc, + signal: &Arc, + ) { + let drain_counters: Vec> = + (0..contexts.inputs.len()).map(|_| StepDrainCounter::new(1)).collect(); + let group = DetachedDriverGroup::new(vec![(step_idx, step)]); + run_detached_driver( + group, + contexts, + &drain_counters, + signal, + None, + &crate::liveness::LivenessCounter::new(1), + ); + } + + /// Items pushed onto a Detached step's input flow through to its output; + /// the step finishes once its input is drained and closes the output edge. + #[test] + fn detached_step_flows_items_and_finishes() { + let (det, contexts, producer_outputs_any, consumer) = build_one_detached(64); + let producer_outputs = + producer_outputs_any.downcast_ref::>>().unwrap(); + producer_outputs.push(10).unwrap(); + producer_outputs.push(20).unwrap(); + producer_outputs.push(30).unwrap(); + producer_outputs.mark_all_drained(); + + let signal = PipelineSignal::new(); + drive_single(det, StepIdx(1), &contexts, &signal); + + let mut got = Vec::new(); + while let Some(v) = consumer.pop() { + got.push(v); + } + assert_eq!(got, vec![10, 20, 30]); + assert!(InputHandle::is_drained(&consumer), "output closed on Finished"); + assert!(!signal.is_done(), "clean completion, no error"); + } + + /// Zero items in (input drained from the start): the Detached step cleanly + /// drains its output and returns. + #[test] + fn detached_step_zero_items_clean_drain() { + let (det, contexts, producer_outputs_any, consumer) = build_one_detached(64); + let producer_outputs = + producer_outputs_any.downcast_ref::>>().unwrap(); + producer_outputs.mark_all_drained(); // no items + + let signal = PipelineSignal::new(); + drive_single(det, StepIdx(1), &contexts, &signal); + + assert!(consumer.pop().is_none(), "no items produced"); + assert!(InputHandle::is_drained(&consumer), "output closed on clean empty drain"); + } + + /// A two-sided Detached step — consumer of a bounded pool-fed input AND + /// producer to a bounded pool-drained output, both tiny — completes without + /// deadlock. A wall-clock watchdog aborts (fails the test) if it wedges. + /// This is the sort merge's topology. + #[test] + fn detached_two_sided_no_deadlock() { + const N: u32 = 5_000; + + // cap-2 input AND cap-4 output both force interleaved backpressure. + let (det, contexts, producer_outputs_any, consumer) = build_one_detached(2); + let signal = PipelineSignal::new(); + + // Watchdog: a deadlock parks forever; abort so the test FAILS loudly. + let done = Arc::new(std::sync::atomic::AtomicBool::new(false)); + { + let done = Arc::clone(&done); + std::thread::spawn(move || { + for _ in 0..200 { + std::thread::sleep(Duration::from_millis(50)); + if done.load(Ordering::SeqCst) { + return; + } + } + eprintln!("detached_two_sided_no_deadlock: WEDGED (deadlock)"); + std::process::abort(); + }); + } + + // Producer thread: push N items into the cap-2 input (backpressure), + // then close it so the Detached step's input drains. + let pushed = Arc::new(AtomicU32::new(0)); + let producer_handle = { + let producer_outputs_any = Arc::clone(&producer_outputs_any); + let pushed = Arc::clone(&pushed); + std::thread::spawn(move || { + let outputs = + producer_outputs_any.downcast_ref::>>().unwrap(); + let mut held: Option = None; + let mut next = 0u32; + loop { + if let Some(v) = held.take() { + match outputs.push(v) { + Ok(()) => { + pushed.fetch_add(1, Ordering::Relaxed); + } + Err(unpushed) => { + held = Some(unpushed.into_item()); + std::thread::yield_now(); + } + } + continue; + } + if next >= N { + break; + } + match outputs.push(next) { + Ok(()) => { + pushed.fetch_add(1, Ordering::Relaxed); + next += 1; + } + Err(unpushed) => { + // The value at `next` is now held for retry; advance + // `next` so the fresh-push branch doesn't re-emit it + // after `held` flushes (which would double-count). + held = Some(unpushed.into_item()); + next += 1; + std::thread::yield_now(); + } + } + } + outputs.mark_all_drained(); + }) + }; + + // Consumer thread: pop everything the Detached step produces. Keep the + // VALUES, not just a count — a count-only assertion passes for a step that + // emits N copies of one item, or that duplicates the held value while + // dropping a popped one, which is exactly the loss this test exists to + // catch. + let received = Arc::new(parking_lot::Mutex::new(Vec::::new())); + let consumer_handle = { + let received = Arc::clone(&received); + std::thread::spawn(move || { + loop { + if let Some(v) = consumer.pop() { + received.lock().push(v); + } else if InputHandle::is_drained(&consumer) { + break; + } else { + std::thread::yield_now(); + } + } + }) + }; + + // Drive the Detached step (as a one-step group) on this thread to + // completion via the unified driver. + drive_single(det, StepIdx(1), &contexts, &signal); + + producer_handle.join().unwrap(); + consumer_handle.join().unwrap(); + done.store(true, Ordering::SeqCst); + + assert_eq!(pushed.load(Ordering::Relaxed), N, "all items pushed"); + // Sorted multiset, not the sequence: the producer's backpressure branch + // holds `next` and advances, so a retried item can arrive after a later + // one. Order is not the invariant here; every distinct item arriving + // exactly once is. + let mut got = received.lock().clone(); + got.sort_unstable(); + assert_eq!( + got, + (0..N).collect::>(), + "every distinct item must flow through the two-sided Detached step exactly once" + ); + } + + /// A `Detached` step declaring a `Shared` group label. Used only to exercise + /// `extract_detached_steps` grouping — never actually run. + #[derive(Clone)] + struct SharedDetached { + label: &'static str, + } + impl Step for SharedDetached { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "SharedDetached", + kind: StepKind::Detached, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn detached_group(&self) -> crate::step::DetachedGroup { + crate::step::DetachedGroup::Shared(self.label) + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::Finished) + } + } + + /// A `Parallel` step — must never be placed on a 1-thread driver. + #[derive(Clone)] + struct ParallelStub; + impl Step for ParallelStub { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "ParallelStub", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + /// `extract_detached_steps` collects every `Shared(label)` onto one group, + /// keeps each `PerStep` (default) step as its own singleton, preserves chain + /// order within and across groups, and leaves non-detached steps in place. + #[test] + fn extract_groups_shared_together_and_perstep_alone() { + use crate::step::DetachedGroup; + // idx0 non-detached; idx1/idx3 Shared("coord"); idx2 PerStep; idx4 Shared("io"). + let mut steps: Vec> = vec![ + Box::new(TypedStep::new(SrcStub { capacity: 4 })), + Box::new(TypedStep::new(SharedDetached { label: "coord" })), + Box::new(TypedStep::new(PassThroughDetached { held: None })), + Box::new(TypedStep::new(SharedDetached { label: "coord" })), + Box::new(TypedStep::new(SharedDetached { label: "io" })), + ]; + let groups = extract_detached_steps(&mut steps); + + assert_eq!(groups.len(), 3, "coord{{1,3}}, perstep{{2}}, io{{4}}"); + // Order is first-appearance in chain order. + assert_eq!(groups[0].label(), DetachedGroup::Shared("coord")); + assert_eq!( + groups[0].steps.iter().map(|(i, _)| i.0).collect::>(), + vec![1, 3], + "shared group keeps both steps in chain order" + ); + assert_eq!(groups[0].primary_step(), StepIdx(1)); + assert_eq!(groups[1].label(), DetachedGroup::PerStep); + assert_eq!(groups[1].steps.iter().map(|(i, _)| i.0).collect::>(), vec![2]); + assert_eq!(groups[2].label(), DetachedGroup::Shared("io")); + + // Non-detached step survives; detached slots became placeholders that + // still report `Detached` (so `build_worker_storage` Skips them on pool). + assert_eq!(steps[0].kind(), StepKind::Exclusive); + assert_eq!(steps[1].kind(), StepKind::Detached); + assert_eq!(steps[2].kind(), StepKind::Detached); + + // Each placeholder reports the group its real step declared, not a blanket + // `PerStep`. `extract_detached_steps` is `pub` and hands this slice back, + // so a caller regrouping from it would otherwise split the `coord` group + // across separate driver threads. + assert_eq!( + steps[1].detached_group(), + DetachedGroup::Shared("coord"), + "placeholder must not downgrade a Shared group to PerStep" + ); + assert_eq!(steps[2].detached_group(), DetachedGroup::PerStep); + } + + #[test] + #[should_panic(expected = "must be non-empty")] + fn driver_group_rejects_empty() { + // An empty group would panic later in `primary_step`/`label` (`steps[0]`); + // the constructor rejects it up front. + let _ = DetachedDriverGroup::new(vec![]); + } + + #[test] + #[should_panic(expected = "only contain Detached steps")] + fn driver_group_rejects_non_detached_step() { + // `SrcStub` is Exclusive, not Detached — running it off-pool on a driver + // thread is a bug, so the constructor rejects the group. + let step: Box = Box::new(TypedStep::new(SrcStub { capacity: 1 })); + let _ = DetachedDriverGroup::new(vec![(StepIdx(0), step)]); + } + + /// `build_driver_storage` makes the group's steps `Owned` and every other + /// slot `Skip`, at the correct global indices. + #[test] + fn build_driver_storage_owns_group_skips_rest() { + let det: Box = Box::new(TypedStep::new(PassThroughDetached { held: None })); + let row = build_driver_storage(vec![(StepIdx(2), det)], 5); + assert_eq!(row.len(), 5); + assert!(matches!(row[2], WorkerStepEntry::Owned { .. }), "group step is Owned"); + for i in [0usize, 1, 3, 4] { + assert!(matches!(row[i], WorkerStepEntry::Skip), "non-group slot {i} is Skip"); + } + } + + /// G3: a `Parallel` step must never be grouped onto a 1-thread driver (its + /// init-1 counter would never close the shared output). + #[test] + #[should_panic(expected = "is Parallel")] + fn build_driver_storage_rejects_parallel_group_step() { + let par: Box = Box::new(TypedStep::new(ParallelStub)); + let _ = build_driver_storage(vec![(StepIdx(0), par)], 2); + } + + /// G3: two group steps at the same index is a dual registration — rejected. + #[test] + #[should_panic(expected = "registered twice")] + fn build_driver_storage_rejects_dual_registration() { + let a: Box = Box::new(TypedStep::new(PassThroughDetached { held: None })); + let b: Box = Box::new(TypedStep::new(PassThroughDetached { held: None })); + let _ = build_driver_storage(vec![(StepIdx(1), a), (StepIdx(1), b)], 3); + } + + /// G3: a group step index past `n_total_steps` is the documented out-of-range + /// panic, matching the guard the sibling `build_worker_storage` applies — + /// without it the bare `row[idx.0]` index would panic with an opaque message. + #[test] + #[should_panic(expected = "out of range")] + fn build_driver_storage_rejects_out_of_range_index() { + let det: Box = Box::new(TypedStep::new(PassThroughDetached { held: None })); + let _ = build_driver_storage(vec![(StepIdx(5), det)], 3); + } +} diff --git a/crates/fgumi-pipeline-core/src/runtime/drain.rs b/crates/fgumi-pipeline-core/src/runtime/drain.rs new file mode 100644 index 000000000..fe60db784 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/runtime/drain.rs @@ -0,0 +1,99 @@ +//! `StepDrainCounter`: coordinates last-worker-wins for closing a step's +//! shared output queue when it reports `StepOutcome::Finished`. +//! +//! For `Parallel` steps, init to N (= worker count). Each clone returns +//! `Finished` independently when the shared input edge drains; each calls +//! `observe_drain`, and the clone that takes the counter to 0 is the "last +//! worker" — only it calls `mark_outputs_drained` (closing the shared output). +//! Otherwise a clone could close the output while a sibling is still pushing. +//! +//! For `Serial` / `Exclusive` steps, init to 1. The single finisher wins on +//! its first call. + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; + +#[derive(Debug)] +pub struct StepDrainCounter { + remaining: AtomicUsize, +} + +impl StepDrainCounter { + /// Construct a counter with the given initial decrement budget. Returns + /// an `Arc` for sharing across worker threads. + #[must_use] + pub fn new(initial: usize) -> Arc { + Arc::new(Self { remaining: AtomicUsize::new(initial) }) + } + + /// Called by a worker when it observes drain on this step. Returns + /// `true` if this is the last worker (counter went from 1 to 0). + /// Subsequent calls (counter already 0) return `false`. + pub fn observe_drain(&self) -> bool { + // CAS-decrement loop: atomically decrement only if `> 0`. Avoids + // underflow under concurrent over-calls. + let mut prev = self.remaining.load(AtomicOrdering::Acquire); + loop { + if prev == 0 { + return false; + } + match self.remaining.compare_exchange_weak( + prev, + prev - 1, + AtomicOrdering::AcqRel, + AtomicOrdering::Acquire, + ) { + Ok(_) => return prev == 1, + Err(actual) => prev = actual, + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn first_worker_with_init_one_wins() { + let counter = StepDrainCounter::new(1); + assert!(counter.observe_drain()); + assert!(!counter.observe_drain()); + } + + #[test] + fn last_worker_with_init_n_wins() { + let counter = StepDrainCounter::new(4); + assert!(!counter.observe_drain()); + assert!(!counter.observe_drain()); + assert!(!counter.observe_drain()); + assert!(counter.observe_drain()); + } + + #[test] + fn extra_calls_after_zero_return_false() { + let counter = StepDrainCounter::new(2); + counter.observe_drain(); + counter.observe_drain(); + assert!(!counter.observe_drain()); + } + + #[test] + fn concurrent_decrements_have_exactly_one_winner() { + use std::thread; + let counter = StepDrainCounter::new(8); + // Collect the handles BEFORE joining: iterator adapters are lazy, so a + // chained `.map(spawn).map(join)` pulls one element at a time and joins + // each thread before spawning the next — making this the same sequential + // path as `last_worker_with_init_n_wins`, and passing even against a + // non-atomic load/store implementation. + let handles: Vec<_> = (0..8) + .map(|_| { + let c = Arc::clone(&counter); + thread::spawn(move || c.observe_drain()) + }) + .collect(); + let winners: Vec = handles.into_iter().map(|h| h.join().unwrap()).collect(); + assert_eq!(winners.iter().filter(|&&w| w).count(), 1, "exactly one winner"); + } +} diff --git a/crates/fgumi-pipeline-core/src/runtime/driver.rs b/crates/fgumi-pipeline-core/src/runtime/driver.rs new file mode 100644 index 000000000..bee9ed3dc --- /dev/null +++ b/crates/fgumi-pipeline-core/src/runtime/driver.rs @@ -0,0 +1,1347 @@ +//! Worker loop body. Each worker thread runs `run_worker_loop` until +//! `signal.is_done()` or all steps are drained. +//! +//! Loop structure (per iteration): +//! 1. Check `signal.is_done()`; bail if true. +//! 2. **Sticky re-entry**: if this worker owns an `Exclusive` step that's +//! sticky, drive it to a stop (`Progress`→loop; `Finished`→done; +//! `NoProgress` or `Contention`→exit sticky), for at most +//! `STICKY_BURST_LIMIT` consecutive calls. Sticky avoids context- +//! switch overhead for source-style steps that emit in tight bursts. +//! 3. **Round-robin dispatch**: try each step in chain order. On +//! `Progress`, restart from step 0 (priority) — except for this worker's +//! sticky owner, which already had its burst in step 2. On `Finished`, the +//! step is complete — `mark_outputs_drained` (counter-gated for `Parallel`) +//! and remove it from the worklist. `NoProgress`/`Contention` are idle ticks. +//! 4. If no work happened this iteration, exponential-backoff sleep. +//! +//! Completion: every step — source, mid, or sink — terminates by returning +//! `StepOutcome::Finished` from `try_run` once its input edges are drained and +//! it holds no buffered output. The framework then closes its output edges and +//! drops it from the per-worker worklist. For a `Parallel` step the per-step +//! [`StepDrainCounter`] gates `mark_outputs_drained` so only the last clone to +//! finish closes the shared output queue (see `dispatch_one_step`). + +use std::any::Any; +use std::sync::Arc; +use std::time::Instant; + +use crate::erased::ErasedStepCtx; +use crate::liveness::LivenessCounter; +use crate::runtime::contexts::ChainContexts; +use crate::runtime::drain::StepDrainCounter; +use crate::runtime::live::LiveSteps; +use crate::runtime::scheduler::{Scheduler, WalkDirection}; +use crate::runtime::stats::PipelineStats; +use crate::runtime::storage::WorkerStepEntry; +use crate::runtime::worker_core::{WorkerCore, WorkerRole}; +use crate::signal::{PipelineError, PipelineSignal}; +use crate::step::StepOutcome; +use crate::topology::StepIdx; + +/// Maximum consecutive sticky re-entries before the worker drops back to a +/// round-robin pass. +/// +/// `StepOutcome::Progress` means the step "pushed **or held** an item" (see the +/// [`crate::step`] completion contract), so a sticky step whose output is full +/// reports `Progress` on every call while moving nothing. Re-entering without a +/// bound would spin on that step forever and never dispatch the downstream step +/// that drains the output — at one worker, a hang; at several, a burnt core. +/// +/// The bound is high enough that the fast path keeps its point (a source +/// emitting a tight burst skips the outer loop's bookkeeping per item) and low +/// enough that a step stuck on a full output yields promptly. It caps starvation, +/// it is not a tuning knob: forward progress comes from this bound *together +/// with* round-robin declining to restart the walk on the sticky owner's +/// `Progress`. +const STICKY_BURST_LIMIT: usize = 1024; + +/// Run the worker loop for one worker thread. +/// +/// `entries[step_idx] = WorkerStepEntry` — this worker's storage. +/// `contexts` — shared per-step input/output handles. +/// `drain_counters[step_idx]` — the per-step `StepDrainCounter` that gates the +/// output close on `Finished` (init N for Parallel so only the last clone +/// closes the shared output; init 1 for Serial/Exclusive). +/// `signal` — error/cancel broadcast. +#[allow(clippy::too_many_arguments)] // per-step shared state plus the liveness shard; a struct would only rename it +pub fn run_worker_loop( + worker: &mut WorkerCore, + entries: &mut [WorkerStepEntry], + contexts: &Arc, + drain_counters: &[Arc], + signal: &Arc, + stats: Option<&Arc>, + liveness: &LivenessCounter, + scheduler: &dyn Scheduler, +) { + // Per-worker worklist of still-dispatchable steps, in chain order. A step + // is removed when it returns `StepOutcome::Finished`; the worker exits once + // the list is empty. Build-time `Skip` placeholders (Exclusive steps owned + // by other workers, Serial steps gated out by affinity) never enter the + // worklist. + let mut live = LiveSteps::from_entries(entries); + + // Cache whether this worker's sticky owner (if any) is still dispatchable, + // so the hot sticky fast-path does not run a linear `live.contains` scan on + // every outer-loop iteration (the sticky path exists precisely to shave + // per-iteration overhead). `sticky_owner` is fixed for the worker's + // lifetime; the step leaves `live` exactly when it returns `Finished`, + // either via the sticky block below or via round-robin dispatch — both + // sites flip this flag false. A `None` sticky owner is permanently "not + // live" so the fast path is skipped entirely. + let mut sticky_live = worker.sticky_owner.is_some_and(|idx| live.contains(idx)); + + // A driver thread attributes busy time PER grouped step on the off-pool + // detached line (so a multi-step `Shared` group shows each step's real + // busy, not the whole thread's time under one name), recorded inside + // `dispatch_one_step`; a pool worker records aggregate busy by `thread_id` + // below. Idle/park stay thread-level (keyed to the group's primary step). + let is_driver = matches!(worker.role(), WorkerRole::Driver { .. }); + + loop { + if signal.is_done() { + break; + } + // Exit when this worker has nothing left to dispatch. + if live.is_empty() { + break; + } + + let mut did_work = false; + + // Time the dispatch (busy) section vs the backoff sleep (idle) below, + // per worker, only when stats are on (`Instant::now()` is otherwise not + // called — the no-stats loop stays zero-cost). + let work_start = stats.map(|_| Instant::now()); + + // 1. Sticky re-entry for sticky-owned steps (either an Exclusive + // sticky step this worker owns, or a Serial+sticky step whose + // Affinity targets this worker). The Pipeline::run path only + // sets `sticky_owner` when the step is actually sticky, so no + // per-iteration profile peek is needed. Re-enter while the + // step makes Progress; exit on Finished / NoProgress / Contention + // / Err. Remove from the worklist on Finished (source drain) or + // observed input drain (mid-step drain). Gated on the step still + // being live — once removed, the sticky fast-path is disabled. + // + // Bounded at `STICKY_BURST_LIMIT` calls: `Progress` also covers "held an + // item", so a step whose output is full reports it indefinitely without + // moving anything. An unbounded re-entry would then never reach the + // round-robin pass that dispatches the downstream step draining that + // output. + if let Some(owned_idx) = worker.sticky_owner.filter(|_| sticky_live) { + let mut mark_skip = false; + for _ in 0..STICKY_BURST_LIMIT { + if signal.is_done() { + break; + } + let entry = &mut entries[owned_idx.0]; + let Some(info) = dispatch_one_step( + entry, + owned_idx, + contexts, + &drain_counters[owned_idx.0], + signal, + stats, + liveness, + worker.thread_id, + is_driver, + ) else { + break; // Skip + }; + match info.result { + Ok(StepOutcome::Progress) => { + did_work = true; + // Continue sticky. + } + Ok(StepOutcome::Finished) => { + // Outputs were marked drained under the dispatch guard. + did_work = true; + mark_skip = true; + break; + } + // Nothing to do this call — yield out of the sticky loop + // back to round-robin. The step terminates via `Finished`, + // not a drain protocol. + Ok(StepOutcome::NoProgress | StepOutcome::Contention) => break, + Err(io_err) => { + signal.record_error(PipelineError::Io { step: info.name, source: io_err }); + break; + } + } + } + if mark_skip { + live.remove(owned_idx); + // The sticky owner finished here; disable the fast path. + sticky_live = false; + } + } + + // 2. Round-robin priority dispatch over all live steps. + if !signal.is_done() { + let outcome = round_robin_dispatch( + entries, + &mut live, + worker.sticky_owner, + contexts, + drain_counters, + signal, + stats, + liveness, + worker.thread_id, + scheduler.walk(), + is_driver, + ); + did_work |= outcome.did_work; + if outcome.removed_sticky_owner { + // The sticky owner finished during round-robin; disable the + // fast path so subsequent iterations skip the sticky block. + sticky_live = false; + } + } + + // Attribute the dispatch section's wall time to this thread's busy total. + // Pool workers sum the whole pass into the N-worker utilisation line (by + // thread_id). Driver threads instead record each grouped step's own busy + // inside `dispatch_one_step` (on the off-pool detached line, by step) so a + // multi-step `Shared` group isn't collapsed onto one name — so nothing to + // record here for a driver. + if let (Some(stats), Some(ws)) = (stats, work_start) + && let WorkerRole::Pool = worker.role() + { + let ns = u64::try_from(ws.elapsed().as_nanos()).unwrap_or(u64::MAX); + stats.record_worker_busy(worker.thread_id, ns); + } + + // 3. Exponential-backoff sleep on no-progress; reset on progress. The + // sleep is the worker's idle/blocked time — attribute it per worker so + // pool under-utilisation (cores parked while one worker drives a Serial + // step) is visible in `--pipeline-stats`. + if did_work { + worker.reset_backoff(); + } else if signal.is_done() { + break; + } else { + let sleep_start = stats.map(|_| Instant::now()); + worker.sleep_backoff(); + worker.increase_backoff(); + if let (Some(stats), Some(ss)) = (stats, sleep_start) { + let ns = u64::try_from(ss.elapsed().as_nanos()).unwrap_or(u64::MAX); + match worker.role() { + WorkerRole::Pool => stats.record_worker_idle(worker.thread_id, ns), + WorkerRole::Driver { primary_step } => { + stats.record_detached_idle(primary_step, ns); + stats.record_detached_park(primary_step); + } + } + } + } + } +} + +/// Result of one round-robin pass: whether any step did useful work (caller +/// resets backoff), and whether the worker's sticky owner finished during the +/// pass (caller clears its `sticky_live` cache so the sticky fast-path is not +/// re-attempted on a removed step). +struct RoundRobinOutcome { + did_work: bool, + removed_sticky_owner: bool, +} + +/// One pass of the round-robin dispatch over this worker's live steps, in +/// chain order. Finished steps are removed from `live` at end-of-pass (deferred +/// so the in-progress walk over `live.order()` is not mutated underneath it). +/// `sticky_owner` (if any) has two roles here: it is reported back via +/// [`RoundRobinOutcome::removed_sticky_owner`] when it finishes in this pass, so +/// the caller can disable the per-iteration sticky fast-path without a linear +/// `live.contains` scan; and its `Progress` does **not** trigger the priority +/// restart, so the walk continues to the steps downstream of it (see the +/// `Progress` arm). +#[allow(clippy::too_many_arguments)] // shared per-step state + the walk policy; a struct would not clarify +fn round_robin_dispatch( + entries: &mut [WorkerStepEntry], + live: &mut LiveSteps, + sticky_owner: Option, + contexts: &Arc, + drain_counters: &[Arc], + signal: &Arc, + stats: Option<&Arc>, + liveness: &LivenessCounter, + worker_slot: usize, + walk: WalkDirection, + is_driver: bool, +) -> RoundRobinOutcome { + let mut did_work = false; + // Steps that finished this pass, removed from `live` after the walk. A + // step is visited at most once per pass (the cursor only advances; we + // `break` on `Progress`/error, never revisit), so deferring removal is + // safe and avoids reorder-under-iteration. + let mut finished: Vec = Vec::new(); + let n = live.len(); + for i in 0..n { + if signal.is_done() { + break; + } + // The Scheduler selects the walk DIRECTION over this worker's live + // steps: `Forward` = chain order (upstream-first, favour production); + // `Reverse` = downstream-first (favour draining buffered work before + // producing more). Everything else — skip-on-contention, the sticky + // source/sink fast-path, Finished handling — is direction-agnostic. + let pos = match walk { + WalkDirection::Forward => i, + WalkDirection::Reverse => n - 1 - i, + }; + let step_idx = live.order()[pos]; + let entry = &mut entries[step_idx.0]; + let mut mark_skip = false; + let mut restart_priority = false; + let Some(info) = dispatch_one_step( + entry, + step_idx, + contexts, + &drain_counters[step_idx.0], + signal, + stats, + liveness, + worker_slot, + is_driver, + ) else { + continue; // Skip (build-time placeholder; should not appear in `live`) + }; + match info.result { + Ok(StepOutcome::Progress) => { + did_work = true; + // Restart the walk at the top so the highest-priority step runs + // again — EXCEPT for this worker's sticky owner. That step just + // had its dedicated burst in phase 1 of the worker loop, so a + // priority restart here only re-runs it. Worse, `Progress` also + // means "held an item", so a sticky step sitting on a full output + // reports it forever: breaking the pass on that outcome means the + // downstream step that would drain the output is never reached + // and the pair livelocks at one worker. Walking past the sticky + // owner is what turns the bounded burst into real forward + // progress. + restart_priority = sticky_owner != Some(step_idx); + } + // Nothing to do this call. The step terminates by returning + // `Finished` (handled below); `NoProgress`/`Contention` are idle + // ticks — there is no separate drain protocol. + Ok(StepOutcome::NoProgress | StepOutcome::Contention) => {} + Ok(StepOutcome::Finished) => { + // Any step (source, mid, or sink) may report `Finished` once + // all its inputs are drained and it holds no buffered output. + // Outputs were marked drained under the dispatch guard (and, for + // a Serial step, the shared `finished` latch was set so the + // other workers stop re-dispatching it — see `dispatch_one_step`). + did_work = true; + mark_skip = true; + } + Err(io_err) => { + signal.record_error(PipelineError::Io { step: info.name, source: io_err }); + break; + } + } + if mark_skip { + finished.push(step_idx); + } + if restart_priority { + break; + } + } + let removed_sticky_owner = sticky_owner.is_some_and(|owner| finished.contains(&owner)); + for step_idx in finished { + live.remove(step_idx); + } + RoundRobinOutcome { did_work, removed_sticky_owner } +} + +/// Outcome of dispatching one step, plus the `name` captured *during* the +/// dispatch — under the same `Shared`-mutex guard as the run itself — for +/// error reporting. +struct DispatchInfo { + result: std::io::Result, + name: &'static str, +} + +/// Dispatch one step's `try_run_erased`. Returns: +/// - `Some(DispatchInfo)` — dispatched (the `result` carries the outcome or +/// the step's `Err`); on `Finished`, outputs are already marked drained. +/// - `None` — entry is `Skip` (caller continues to next step). +/// +/// For a `Serial` (`Shared`) step the shared `finished` latch on its +/// `DrainGate` is consulted *before* acquiring the step mutex: once any worker +/// has finished the step (returned `Finished`, or completed its cooperative +/// drain), the latch is set and every other worker short-circuits to a synthetic +/// `Finished` here rather than re-`try_lock`-ing and re-running an already-done +/// step. The winning worker sets the latch under the dispatch guard before +/// `mark_outputs_drained`, so a non-idempotent flusher can never be re-entered. +#[allow(clippy::too_many_arguments)] // per-step shared state plus the liveness shard; a struct would only rename it +fn dispatch_one_step( + entry: &mut WorkerStepEntry, + step_idx: StepIdx, + contexts: &ChainContexts, + counter: &StepDrainCounter, + signal: &Arc, + stats: Option<&Arc>, + // Always-on liveness signal for the deadlock monitor, sharded per worker so + // the bump is normally an uncontended increment (a dedicated driver reuses + // `worker_slot` 0 and so shares slot 0 with pool worker 0 — the atomic + // `fetch_add` counts a coincident bump correctly; see `crate::liveness`). + // Separate from `stats` on purpose: + // liveness must be free enough to leave armed, while `stats` pays for + // per-dispatch timing and stays opt-in. See `crate::liveness`. + liveness: &LivenessCounter, + worker_slot: usize, + // When true (a dedicated driver thread), attribute this dispatch's wall time + // to the off-pool detached line keyed by `step_idx` — so each grouped step + // reports its own busy. Pool workers pass `false` and record aggregate busy + // by `thread_id` in the loop instead. + is_driver: bool, +) -> Option { + let outputs_any: &(dyn Any + Send + Sync) = contexts.outputs[step_idx.0].as_ref(); + let mut ctx = + ErasedStepCtx { input: contexts.inputs[step_idx.0].as_ref(), outputs: outputs_any, signal }; + + // Time the dispatch only when stats collection is on. `Instant::now()` + // is ~20-50ns on Apple Silicon, ~50-100ns on x86_64; gating on + // `stats.is_some()` keeps the no-stats path zero-cost. + let start = stats.map(|_| Instant::now()); + + // Run the step and capture `name` *while still holding the `Shared` guard* + // (or with direct `&mut` for Owned/Exclusive). On `Finished`, mark outputs + // drained here too — under the same guard — so the caller never re-acquires + // the lock for any post-dispatch inspection. + // + // `mark_outputs_drained` is gated behind `counter.observe_drain()` (the + // per-step `StepDrainCounter`): for a `Parallel` step (counter init N) + // every clone returns `Finished` independently when the shared input edge + // is drained, but only the LAST clone to finish (the one that takes the + // counter to 0) closes the shared output queue — otherwise a clone could + // `mark_drained` while a sibling is still pushing (`try_push`-after-drained + // panic). For `Serial`/`Exclusive` (counter init 1) the single finisher + // wins on its first call, unchanged. + // + // INVARIANT: for a `Parallel` step, `counter` init == clone count == the + // worker count, and a clone leaves its worklist ONLY by returning + // `Finished`, so the counter reaches 0 exactly when every clone has + // finished. Any future scheduler change that removes a Parallel clone for + // another reason (work-stealing, per-worker early exit) — or makes a source + // `Parallel` — would leave the counter stuck above 0 and never close the + // shared output, hanging the downstream consumer. Keep the init (builder.rs) + // and this gate in lockstep. + let info: Option = match entry { + WorkerStepEntry::Owned { step } | WorkerStepEntry::Exclusive { step } => { + let result = step.try_run_erased(&mut ctx); + // `name()` returns the cached static name — no per-dispatch + // `StepProfile` (and its two `Vec`s) is built. + let name = step.name(); + if matches!(result, Ok(StepOutcome::Finished)) && counter.observe_drain() { + step.mark_outputs_drained(outputs_any); + } + Some(DispatchInfo { result, name }) + } + WorkerStepEntry::Shared { step, drain } => { + if drain.is_finished() { + // Another worker already finished this Serial step. Don't + // re-`try_lock`/re-run it — report a synthetic `Finished` so the + // caller drops it from this worker's live set. Outputs were + // already marked drained by the finishing worker. + Some(DispatchInfo { + result: Ok(StepOutcome::Finished), + name: "", + }) + } else { + match step.try_lock() { + None => Some(DispatchInfo { + result: Ok(StepOutcome::Contention), + name: "", + }), + Some(mut guard) => { + let result = guard.try_run_erased(&mut ctx); + let name = guard.name(); + if matches!(result, Ok(StepOutcome::Finished)) { + // Set the shared finished latch under the guard, + // before marking outputs drained, so a concurrent + // worker that observes the latch never re-runs the + // step nor re-marks its outputs. + drain.mark_finished(); + if counter.observe_drain() { + guard.mark_outputs_drained(outputs_any); + } + } + Some(DispatchInfo { result, name }) + } + } + } + } + WorkerStepEntry::Skip => None, + }; + + // Liveness first, and unconditionally: this is what the deadlock monitor + // samples, so it must not depend on `stats` being attached. Only productive + // outcomes count — a wedged pipeline still spins through `NoProgress` and + // `Contention` dispatches forever, so counting those would make a wedge look + // alive and defeat the whole detector. + if let Some(i) = info.as_ref() + && matches!(i.result, Ok(StepOutcome::Progress | StepOutcome::Finished)) + { + liveness.bump(worker_slot); + } + + if let (Some(stats), Some(start)) = (stats, start) { + let elapsed_ns = u64::try_from(start.elapsed().as_nanos()).unwrap_or(u64::MAX); + // Wall ns at dispatch start, relative to pipeline start. + let start_ns = stats.elapsed_ns().saturating_sub(elapsed_ns); + // `None` (Skip) attempted no work, so there is nothing to record. + if let Some(i) = info.as_ref() { + match &i.result { + Ok(outcome) => stats.record(step_idx, *outcome, start_ns, elapsed_ns), + Err(_) => stats.record_error(step_idx, start_ns, elapsed_ns), + } + // On a driver thread, this step's try_run wall is its own off-pool + // busy (excluded from the pool%); each grouped step accrues its own. + if is_driver { + stats.record_detached_busy(step_idx, elapsed_ns); + } + } + } + + info +} + +#[cfg(test)] +mod tests { + use std::io; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + + use super::*; + use crate::erased::{ErasedStep, TypedStep}; + use crate::handles::BranchInputHandle; + use crate::outputs::Single; + use crate::queues::QueueSpec; + use crate::reorder::BranchOrdering; + use crate::runtime::contexts::build_chain_contexts; + use crate::runtime::storage::DrainGate; + use crate::step::{InputHandle, Step, StepCtx, StepKind, StepOutcome, StepProfile}; + use crate::topology::{BranchIdx, ChainGraph}; + use parking_lot::Mutex; + + #[test] + fn run_worker_loop_exits_on_signal_done() { + let signal = PipelineSignal::new(); + let mut entries: Vec = vec![]; + let contexts = Arc::new(ChainContexts { + inputs: vec![], + outputs: vec![], + bounded_queues: vec![], + edges: vec![], + }); + let drain_counters: Vec> = vec![]; + let _ = ChainGraph::new(); + let mut worker = WorkerCore::new(0, None, None); + + signal.cancel(); + run_worker_loop( + &mut worker, + &mut entries, + &contexts, + &drain_counters, + &signal, + None, + &crate::liveness::LivenessCounter::new(1), + &crate::runtime::scheduler::ChainOrderScheduler, + ); + // If we reach this line, the loop exited cleanly. + } + + // ── Test steps for dispatch-level coverage ────────────────────────────── + + /// `() → u32` source that returns `Finished` immediately. + #[derive(Clone)] + struct SrcFinished; + impl Step for SrcFinished { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "Src", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::Finished) + } + } + + /// `() → u32` source that returns `NoProgress` on its first `try_run` (so + /// the sticky fast-path yields back to round-robin without removing it) and + /// `Finished` on every later call (so it is removed during the round-robin + /// pass, exercising `RoundRobinOutcome::removed_sticky_owner`). + #[derive(Clone)] + struct SrcIdleThenFinish { + calls: Arc, + } + impl Step for SrcIdleThenFinish { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "SrcIdleThenFinish", + kind: StepKind::Exclusive, + sticky: true, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + let n = self.calls.fetch_add(1, Ordering::Relaxed); + if n == 0 { Ok(StepOutcome::NoProgress) } else { Ok(StepOutcome::Finished) } + } + } + + /// `u32 → u32` step that always returns `Finished`. Used both as a + /// `Parallel` body (counter-gated output close) and a `Serial` body + /// (`DrainGate` short-circuit). The `runs` counter records every `try_run` + /// so the short-circuit test can prove a worker did NOT re-run the step. + #[derive(Clone)] + struct FinishStep { + kind: StepKind, + runs: Arc, + } + impl Step for FinishStep { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "Finish", + kind: self.kind, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + self.runs.fetch_add(1, Ordering::Relaxed); + Ok(StepOutcome::Finished) + } + fn new_worker_copy(&self) -> Self { + // Clones share the `runs` counter so the test can total runs across + // every Parallel clone. + self.clone() + } + } + + #[derive(Clone)] + struct SinkStep; + impl Step for SinkStep { + type Input = u32; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "Sink", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + // Drain any inputs, then finish once the upstream edge is drained so + // the worker loop can terminate (a real sink finishes on drain). + while ctx.input.pop().is_some() {} + if ctx.input.is_drained() { + Ok(StepOutcome::Finished) + } else { + Ok(StepOutcome::NoProgress) + } + } + } + + /// A non-sticky `Exclusive` sink that deliberately stays live for one extra + /// round-robin pass: it ignores its input-drain status and finishes purely + /// on an internal tick counter — `NoProgress` on the first `try_run`, + /// `Finished` after. Keeping a second step alive for one more outer + /// iteration *after* the sticky source is removed is what makes + /// `sticky_owner_removed_via_round_robin_and_loop_exits` branch-specific: a + /// plain `SinkStep` finishes in the same round-robin pass as the source + /// (its input is already drained), emptying `live` so the loop exits via + /// `live.is_empty()` even if the `removed_sticky_owner` branch had failed to + /// clear `sticky_live`. Lingering forces the extra iteration on which a + /// stale `sticky_live` would re-enter the sticky fast-path and re-invoke the + /// already-removed source. + struct LingerThenFinish { + ticks: usize, + } + impl Step for LingerThenFinish { + type Input = u32; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "Linger", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + while ctx.input.pop().is_some() {} + self.ticks += 1; + // Stay live for exactly one extra round-robin pass before finishing, + // regardless of input-drain status. + if self.ticks >= 2 { Ok(StepOutcome::Finished) } else { Ok(StepOutcome::NoProgress) } + } + } + + /// Build `Src → Finish → Sink` (Finish having the given kind) and return the + /// erased steps + the wired graph. The `Finish` step's `try_run` counter is + /// returned so tests can assert how many times it actually ran. + fn three_step_chain( + finish_kind: StepKind, + ) -> (Vec>, ChainGraph, Arc) { + let runs = Arc::new(AtomicUsize::new(0)); + let mut graph = ChainGraph::new(); + let src = graph.register_step("Src", 1); + let mid = graph.register_step("Finish", 1); + let sink = graph.register_step("Sink", 0); + graph.wire(src, BranchIdx(0), mid); + graph.wire(mid, BranchIdx(0), sink); + let steps: Vec> = vec![ + Box::new(TypedStep::new(SrcFinished)), + Box::new(TypedStep::new(FinishStep { kind: finish_kind, runs: Arc::clone(&runs) })), + Box::new(TypedStep::new(SinkStep)), + ]; + (steps, graph, runs) + } + + /// A `Parallel` step's shared output queue must be closed exactly once — by + /// the LAST clone to finish (the one that takes the `StepDrainCounter` to + /// 0). Earlier finishers must leave the downstream input un-drained so a + /// sibling could still push. + #[test] + fn parallel_last_finisher_closes_shared_output_exactly_once() { + const N: usize = 4; + let (steps, graph, _runs) = three_step_chain(StepKind::Parallel); + let contexts = Arc::new(build_chain_contexts( + &steps, + &graph, + crate::builder::InstrumentationLevel::Off, + )); + let mid = StepIdx(1); + let counter = StepDrainCounter::new(N); + let signal = PipelineSignal::new(); + + // One `Owned` clone per worker, each sharing `contexts.outputs[mid]`. + let mut clones: Vec = + (0..N).map(|_| WorkerStepEntry::Owned { step: steps[mid.0].clone_boxed() }).collect(); + + let sink_input = contexts.inputs[2].downcast_ref::>().unwrap(); + + for (i, clone) in clones.iter_mut().enumerate() { + assert!( + !InputHandle::is_drained(sink_input), + "downstream input drained before the last clone finished (after {i} of {N})" + ); + let info = dispatch_one_step( + clone, + mid, + &contexts, + &counter, + &signal, + None, + &LivenessCounter::new(1), + 0, + false, + ) + .unwrap(); + assert!(matches!(info.result, Ok(StepOutcome::Finished))); + } + assert!( + InputHandle::is_drained(sink_input), + "downstream input must be drained once the last Parallel clone finished" + ); + } + + /// Once one worker finishes a `Serial` step (setting the shared `DrainGate`), + /// a second `dispatch_one_step` short-circuits to a synthetic `Finished` + /// WITHOUT re-acquiring the mutex or re-running the step. + #[test] + fn serial_drain_gate_short_circuits_second_worker() { + let (steps, graph, runs) = three_step_chain(StepKind::Serial); + let contexts = Arc::new(build_chain_contexts( + &steps, + &graph, + crate::builder::InstrumentationLevel::Off, + )); + let mid = StepIdx(1); + let counter = StepDrainCounter::new(1); + let signal = PipelineSignal::new(); + + let shared = Arc::new(Mutex::new(steps.into_iter().nth(1).unwrap())); + let drain = Arc::new(DrainGate::default()); + + // Worker 1 finishes the step: runs once, sets the latch. + let mut entry1 = + WorkerStepEntry::Shared { step: Arc::clone(&shared), drain: Arc::clone(&drain) }; + let info1 = dispatch_one_step( + &mut entry1, + mid, + &contexts, + &counter, + &signal, + None, + &LivenessCounter::new(1), + 0, + false, + ) + .unwrap(); + assert!(matches!(info1.result, Ok(StepOutcome::Finished))); + assert_eq!(runs.load(Ordering::Relaxed), 1, "step ran exactly once on the first worker"); + assert!(drain.is_finished(), "first finisher must set the DrainGate latch"); + + // Worker 2 dispatches the same step: short-circuit, no re-run. + let mut entry2 = + WorkerStepEntry::Shared { step: Arc::clone(&shared), drain: Arc::clone(&drain) }; + let info2 = dispatch_one_step( + &mut entry2, + mid, + &contexts, + &counter, + &signal, + None, + &LivenessCounter::new(1), + 0, + false, + ) + .unwrap(); + assert!(matches!(info2.result, Ok(StepOutcome::Finished))); + assert_eq!( + info2.name, "", + "second worker takes the latch short-circuit" + ); + assert_eq!( + runs.load(Ordering::Relaxed), + 1, + "the Serial step must NOT be re-run after the DrainGate latch is set" + ); + } + + /// A sticky-owned source driven through `run_worker_loop` completes and the + /// loop exits even though the cached `sticky_live` flag (not a per-iteration + /// `live.contains` scan) gates the fast path. Exercises the S1b-006 cache: + /// the sticky step is removed once it returns `Finished`, after which the + /// fast path must be disabled and the loop must terminate. + #[test] + fn sticky_owner_completes_and_loop_exits() { + let mut graph = ChainGraph::new(); + let src = graph.register_step("Src", 1); + let sink = graph.register_step("Sink", 0); + graph.wire(src, BranchIdx(0), sink); + let steps: Vec> = + vec![Box::new(TypedStep::new(SrcFinished)), Box::new(TypedStep::new(SinkStep))]; + let contexts = Arc::new(build_chain_contexts( + &steps, + &graph, + crate::builder::InstrumentationLevel::Off, + )); + + // Single worker; the source (idx 0) is its Exclusive sticky owner, the + // sink (idx 1) is Exclusive owned by the same worker for this 1-worker + // run. + let mut entries: Vec = vec![ + WorkerStepEntry::Exclusive { step: steps.into_iter().next().unwrap() }, + WorkerStepEntry::Exclusive { step: Box::new(TypedStep::new(SinkStep)) }, + ]; + let drain_counters = vec![StepDrainCounter::new(1), StepDrainCounter::new(1)]; + let signal = PipelineSignal::new(); + let mut worker = WorkerCore::new(0, Some(src), Some(src)); + + // The source finishes immediately; the sink then sees its input drained + // and finishes too. `run_worker_loop` must return (no hang). + run_worker_loop( + &mut worker, + &mut entries, + &contexts, + &drain_counters, + &signal, + None, + &crate::liveness::LivenessCounter::new(1), + &crate::runtime::scheduler::ChainOrderScheduler, + ); + } + + /// A sticky owner that returns `NoProgress` on its first call (yielding out + /// of the sticky fast-path back to round-robin) and `Finished` later must be + /// removed via the round-robin path (`RoundRobinOutcome::removed_sticky_owner`), + /// after which the next outer iteration skips the sticky re-entry. This pins + /// the round-robin removal branch (lines around `outcome.removed_sticky_owner`), + /// not just the sticky fast-path removal exercised by + /// `sticky_owner_completes_and_loop_exits`. + /// + /// A second `LingerThenFinish` step is kept alive for one extra round-robin + /// pass *after* the source is removed, so the worker loop must run one more + /// outer iteration. That iteration is where a stale `sticky_live` would + /// wrongly re-enter the sticky fast-path and re-invoke the + /// already-removed-from-`live` source — making the `== 2` source-call + /// assertion below uniquely diagnostic of the `removed_sticky_owner` branch. + /// (Without the linger, a plain sink would finish in the same pass as the + /// source, emptying `live` so the loop exits via `live.is_empty()` whether + /// or not `sticky_live` was cleared — and `== 2` would not be branch-specific.) + #[test] + fn sticky_owner_removed_via_round_robin_and_loop_exits() { + let mut graph = ChainGraph::new(); + let src = graph.register_step("SrcIdleThenFinish", 1); + let linger = graph.register_step("Linger", 0); + graph.wire(src, BranchIdx(0), linger); + + let calls = Arc::new(AtomicUsize::new(0)); + let steps: Vec> = vec![ + Box::new(TypedStep::new(SrcIdleThenFinish { calls: Arc::clone(&calls) })), + Box::new(TypedStep::new(LingerThenFinish { ticks: 0 })), + ]; + let contexts = Arc::new(build_chain_contexts( + &steps, + &graph, + crate::builder::InstrumentationLevel::Off, + )); + + let mut entries: Vec = vec![ + WorkerStepEntry::Exclusive { step: steps.into_iter().next().unwrap() }, + WorkerStepEntry::Exclusive { + step: Box::new(TypedStep::new(LingerThenFinish { ticks: 0 })), + }, + ]; + let drain_counters = vec![StepDrainCounter::new(1), StepDrainCounter::new(1)]; + let signal = PipelineSignal::new(); + let mut worker = WorkerCore::new(0, Some(src), Some(src)); + + // First sticky call → NoProgress (yield to round-robin); the source then + // returns Finished during a round-robin pass, which must remove it and + // disable the sticky fast-path so the loop terminates rather than hangs. + // The `Linger` step stays alive for one more iteration, forcing the + // post-removal outer iteration that exercises the cleared fast path. + run_worker_loop( + &mut worker, + &mut entries, + &contexts, + &drain_counters, + &signal, + None, + &crate::liveness::LivenessCounter::new(1), + &crate::runtime::scheduler::ChainOrderScheduler, + ); + + // The source must have been called EXACTLY twice: call 1 = idle in the + // sticky fast-path (`NoProgress`, which does NOT remove it there — that + // block only reaps a `Finished`), call 2 = `Finished` during the + // round-robin pass. The lingering second step guarantees one more outer + // iteration after that removal, so `== 2` is the branch-specific signal + // for the `removed_sticky_owner` path: if that branch had failed to clear + // `sticky_live`, the extra iteration's sticky fast-path would re-invoke + // the (already-removed-from-`live`) source — the sticky block dispatches + // `entries[owned_idx]` directly, not gated on `live` membership — + // producing a third call. `== 2` therefore proves removal happened via + // round-robin AND that it correctly disabled the fast path. + assert_eq!( + calls.load(Ordering::Relaxed), + 2, + "source must be called exactly twice (sticky idle, then round-robin finish); \ + a different count means the removed_sticky_owner branch did not gate the fast path", + ); + } + + // ── G1: the load-bearing driver invariant ─────────────────────────────── + // + // A driver thread (`WorkerCore::driver`) is just `run_worker_loop` over a + // few `Owned` steps. Its no-deadlock property rests ENTIRELY on the loop + // trying EVERY live step in a pass before it parks. A naive "drive the first + // live step until it yields, then park" loop would wedge the coordination + // driver: it parks on a step whose input isn't ready yet while a *sibling* + // step on the same driver holds the work that would unblock it (e.g. park on + // `FindBoundariesAndSort` while `SpillGather` holds the chunk that frees the + // capacity-1 arena). These two steps pin that the whole-pass discipline holds. + + /// `Owned` step wedged on `NoProgress` until `gate` is flipped by a sibling, + /// then `Finished`. Placed FIRST in the walk so a park-on-first-NoProgress + /// loop would never let the sibling run — the gate never flips — and hang. + #[derive(Clone)] + struct WedgedUntilGate { + gate: Arc, + } + impl Step for WedgedUntilGate { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "WedgedUntilGate", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + if self.gate.load(Ordering::Acquire) { + Ok(StepOutcome::Finished) + } else { + Ok(StepOutcome::NoProgress) + } + } + } + + /// `Owned` sibling that flips `gate` and finishes on its first dispatch — + /// reached only if the loop tries all live steps in a pass rather than + /// parking on the wedged step's `NoProgress`. + #[derive(Clone)] + struct GateOpener { + gate: Arc, + } + impl Step for GateOpener { + type Input = u32; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "GateOpener", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + self.gate.store(true, Ordering::Release); + Ok(StepOutcome::Finished) + } + } + + /// A driver (`WorkerCore::driver`, Park backoff) driving `[Wedged, Opener]` + /// as two `Owned` steps must drive the sibling that unblocks the wedged step + /// and terminate. A watchdog aborts the process on a wedge so the failure is + /// loud rather than a silent hang. + #[test] + fn driver_round_robins_all_live_before_parking() { + let mut graph = ChainGraph::new(); + let wedged = graph.register_step("WedgedUntilGate", 1); + let opener = graph.register_step("GateOpener", 0); + graph.wire(wedged, BranchIdx(0), opener); + + let gate = Arc::new(AtomicBool::new(false)); + let steps: Vec> = vec![ + Box::new(TypedStep::new(WedgedUntilGate { gate: Arc::clone(&gate) })), + Box::new(TypedStep::new(GateOpener { gate: Arc::clone(&gate) })), + ]; + let contexts = Arc::new(build_chain_contexts( + &steps, + &graph, + crate::builder::InstrumentationLevel::Off, + )); + + // Hand-built driver row: both steps Owned on the one driver thread. + let mut entries: Vec = + steps.into_iter().map(|step| WorkerStepEntry::Owned { step }).collect(); + let drain_counters = vec![StepDrainCounter::new(1), StepDrainCounter::new(1)]; + let signal = PipelineSignal::new(); + + // Watchdog: a wedge parks forever; abort so the test FAILS loudly. + let done = Arc::new(AtomicBool::new(false)); + { + let done = Arc::clone(&done); + std::thread::spawn(move || { + for _ in 0..200 { + std::thread::sleep(std::time::Duration::from_millis(25)); + if done.load(Ordering::SeqCst) { + return; + } + } + eprintln!("driver_round_robins_all_live_before_parking: WEDGED"); + std::process::abort(); + }); + } + + // Forward walk (ChainOrderScheduler) tries `wedged` (idx 0) first: it + // yields NoProgress, and the loop MUST proceed to `opener` in the same + // pass, flip the gate, then finish `wedged` on the next pass. + let mut worker = WorkerCore::driver(wedged); + run_worker_loop( + &mut worker, + &mut entries, + &contexts, + &drain_counters, + &signal, + None, + &crate::liveness::LivenessCounter::new(1), + &crate::runtime::scheduler::ChainOrderScheduler, + ); + done.store(true, Ordering::SeqCst); + + assert!(gate.load(Ordering::Acquire), "the sibling opener must have run"); + assert!(!signal.is_done(), "clean completion, no error"); + } + + // ── Sticky forward-progress ───────────────────────────────────────────── + // + // `StepOutcome::Progress` means "pushed OR HELD an item", so a sticky step + // whose output is full keeps reporting `Progress` while moving nothing. Two + // things must hold for the pair below to make progress at ONE worker: the + // sticky burst is bounded, and round-robin does not restart the walk on the + // sticky owner's `Progress` (which would break the pass before the sink is + // ever reached). Removing either one hangs `sticky_holding_source_yields_to_ + // its_draining_consumer`. + + /// Sticky source over a capacity-1 output. Emits `remaining` items; when the + /// transport rejects a push it *holds* the item and reports `Progress` — the + /// contract's "pushed or held" case, and the outcome that makes an unbounded + /// sticky loop spin forever. + struct StickyHoldingSource { + remaining: u32, + held: Option, + calls: Arc, + } + impl Step for StickyHoldingSource { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "StickyHoldingSource", + kind: StepKind::Exclusive, + sticky: true, + // Capacity 1 so the second item in a burst always backs up. + output_queues: vec![QueueSpec::CountBounded { capacity: 1 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + self.calls.fetch_add(1, Ordering::Relaxed); + // Flush-first: retry the held item before producing a new one. + if let Some(item) = self.held.take() { + return match ctx.outputs.push(item) { + Ok(()) => Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held = Some(unpushed.into_item()); + Ok(StepOutcome::Progress) + } + }; + } + if self.remaining == 0 { + return Ok(StepOutcome::Finished); + } + let item = self.remaining; + self.remaining -= 1; + match ctx.outputs.push(item) { + Ok(()) => Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held = Some(unpushed.into_item()); + Ok(StepOutcome::Progress) + } + } + } + } + + /// Sink that pops one item per dispatch — the only thing that frees a slot in + /// the source's capacity-1 output. + struct DrainingSink { + received: Arc>>, + } + impl Step for DrainingSink { + type Input = u32; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "DrainingSink", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + match ctx.input.pop() { + Some(item) => { + self.received.lock().push(item); + Ok(StepOutcome::Progress) + } + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } + } + + /// One worker owning a sticky source (capacity-1 output) and the sink that + /// drains it must deliver every item and terminate. A watchdog aborts on a + /// wedge so the failure is loud rather than a silent hang. + #[test] + fn sticky_holding_source_yields_to_its_draining_consumer() { + const N_ITEMS: u32 = 4; + + let mut graph = ChainGraph::new(); + let src = graph.register_step("StickyHoldingSource", 1); + let sink = graph.register_step("DrainingSink", 0); + graph.wire(src, BranchIdx(0), sink); + + let calls = Arc::new(AtomicUsize::new(0)); + let received = Arc::new(Mutex::new(Vec::new())); + let steps: Vec> = vec![ + Box::new(TypedStep::new(StickyHoldingSource { + remaining: N_ITEMS, + held: None, + calls: Arc::clone(&calls), + })), + Box::new(TypedStep::new(DrainingSink { received: Arc::clone(&received) })), + ]; + let contexts = Arc::new(build_chain_contexts( + &steps, + &graph, + crate::builder::InstrumentationLevel::Off, + )); + + let mut entries: Vec = + steps.into_iter().map(|step| WorkerStepEntry::Exclusive { step }).collect(); + let drain_counters = vec![StepDrainCounter::new(1), StepDrainCounter::new(1)]; + let signal = PipelineSignal::new(); + + // Watchdog: an unbounded sticky loop never returns, so abort loudly. + let done = Arc::new(AtomicBool::new(false)); + { + let done = Arc::clone(&done); + std::thread::spawn(move || { + for _ in 0..400 { + std::thread::sleep(std::time::Duration::from_millis(25)); + if done.load(Ordering::SeqCst) { + return; + } + } + eprintln!( + "sticky_holding_source_yields_to_its_draining_consumer: WEDGED — the \ + sticky fast-path is starving its downstream consumer" + ); + std::process::abort(); + }); + } + + // Forward walk: the source is idx 0 and is this worker's sticky owner. + let mut worker = WorkerCore::new(0, Some(src), Some(src)); + run_worker_loop( + &mut worker, + &mut entries, + &contexts, + &drain_counters, + &signal, + None, + &crate::liveness::LivenessCounter::new(1), + &crate::runtime::scheduler::ChainOrderScheduler, + ); + done.store(true, Ordering::SeqCst); + + assert_eq!( + *received.lock(), + (1..=N_ITEMS).rev().collect::>(), + "every item must reach the sink, in emission order" + ); + assert!(!signal.is_done(), "clean completion, no error"); + // Each item costs at most one full sticky burst plus a round-robin + // dispatch, so the source cannot have been called an unbounded number of + // times. Loose on purpose — it pins "bounded", not an exact schedule. + let n_calls = calls.load(Ordering::Relaxed); + let ceiling = (usize::try_from(N_ITEMS).unwrap() + 2) * (STICKY_BURST_LIMIT + 2); + assert!( + n_calls <= ceiling, + "source dispatches must stay bounded by the sticky burst limit: \ + {n_calls} calls > {ceiling}" + ); + } + + /// A step that spends a measurable, nonzero span inside `try_run` so its + /// dispatch busy-time rounds above 0 ns. + #[derive(Clone)] + struct SlowFinish; + impl Step for SlowFinish { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "SlowFinish", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + std::thread::sleep(std::time::Duration::from_micros(200)); + Ok(StepOutcome::Finished) + } + } + + /// A driver dispatch (`is_driver = true`) records the step's busy on the + /// off-pool detached line keyed by THAT step's own index — so a multi-step + /// `Shared` group attributes each member's real time, not the whole thread's + /// under one name. A pool dispatch (`is_driver = false`) records nothing there. + #[test] + fn driver_dispatch_records_detached_busy_by_own_step() { + let mut graph = ChainGraph::new(); + let a = graph.register_step("SlowFinish", 1); + let sink = graph.register_step("Sink", 0); + graph.wire(a, BranchIdx(0), sink); + let steps: Vec> = + vec![Box::new(TypedStep::new(SlowFinish)), Box::new(TypedStep::new(SinkStep))]; + let contexts = Arc::new(build_chain_contexts( + &steps, + &graph, + crate::builder::InstrumentationLevel::Off, + )); + let mid = StepIdx(0); + let counter = StepDrainCounter::new(1); + let signal = PipelineSignal::new(); + + // Driver dispatch of step 0 → its busy lands on the detached line keyed + // to step 0 (not some group primary). + let stats = Arc::new(PipelineStats::new(vec!["SlowFinish", "Sink"])); + let mut entry = WorkerStepEntry::Owned { step: Box::new(TypedStep::new(SlowFinish)) }; + let _ = dispatch_one_step( + &mut entry, + mid, + &contexts, + &counter, + &signal, + Some(&stats), + &LivenessCounter::new(1), + 0, + true, + ); + let snap = stats.snapshot(); + assert!( + snap.detached.iter().any(|&(step, name, busy, ..)| { + step == mid.0 && name == "SlowFinish" && busy > 0 + }), + "driver dispatch must record detached busy for the dispatched step itself" + ); + + // Pool dispatch (is_driver=false) records nothing on the detached line. + // A FRESH counter: the driver dispatch above consumed the first one, so + // reusing it would make `observe_drain()` return false here and skip + // `mark_outputs_drained` — the second dispatch would silently stop + // exercising the same output-close path as the first. + let counter_pool = StepDrainCounter::new(1); + let stats_pool = Arc::new(PipelineStats::new(vec!["SlowFinish", "Sink"])); + let mut entry_pool = WorkerStepEntry::Owned { step: Box::new(TypedStep::new(SlowFinish)) }; + let _ = dispatch_one_step( + &mut entry_pool, + mid, + &contexts, + &counter_pool, + &signal, + Some(&stats_pool), + &LivenessCounter::new(1), + 0, + false, + ); + assert!( + stats_pool.snapshot().detached.is_empty(), + "pool dispatch must not record on the off-pool detached line" + ); + } +} diff --git a/crates/fgumi-pipeline-core/src/runtime/fused.rs b/crates/fgumi-pipeline-core/src/runtime/fused.rs new file mode 100644 index 000000000..17620fdf7 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/runtime/fused.rs @@ -0,0 +1,1078 @@ +//! Single-thread *fused* execution mode (issue #330). +//! +//! At `--threads 1` a forward-wired `source → … → sink` chain gains nothing from +//! the scheduled worker pool: there is one worker, so the inter-step bounded +//! queues, round-robin polling, held-slot retries, and reorder bookkeeping are +//! pure overhead (profiling showed ~2/3 of `try_run` calls do no useful work). +//! +//! This module is the structural fix. Fusion is **not** a per-command rewrite +//! — it is an execution mode of the existing pipeline. [`is_fusible_chain`] +//! detects a fusible chain (forward-wired, fan-out allowed); +//! [`run_fused_single_thread`] then drives the chain's +//! own type-erased steps inline, in topological order, over **direct** buffers +//! (built by [`build_chain_contexts_fused`]). FIFO push order is already the +//! correct order at one worker, so the reorder stage is dropped; the profile's +//! count/byte bound is kept, because the driver runs a producer before its +//! consumer and a step emitting more per `try_run` than its consumer removes +//! would otherwise grow the edge on every pass. The runtime's +//! generic wiring (`build_chain_contexts`) and dispatch (`try_run_erased`) are +//! reused verbatim — no step logic is duplicated. +//! +//! Non-linear chains (two-input `Step2` merges like zipper, multi-output splits +//! like `correct --rejects`) and `--threads ≥ 2` are not eligible and fall back +//! to the scheduled [`run_worker_loop`](super::driver::run_worker_loop). + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use super::contexts::build_chain_contexts_fused; +use super::stats::PipelineStats; +use crate::builder::InstrumentationLevel; +use crate::erased::{ErasedStep, ErasedStepCtx}; +use crate::signal::{PipelineError, PipelineSignal}; +use crate::step::StepOutcome; +use crate::topology::{BranchIdx, ChainGraph, StepIdx}; + +/// Returns `true` iff `steps` (in chain-construction order) form a chain the +/// fused driver can run on one worker: at least two steps, exactly one source +/// (at index 0), no two-input (`Step2`) merge, and every output branch wired +/// **forward** to a later step. +/// +/// "Forward-wired" allows fan-out — a step may have more than one output branch +/// (e.g. the kept/rejects split of `--rejects`) as long as each branch feeds a +/// later step. Build order equals a topological order (`append_source` then +/// `append_step`/`append_step2`, producers always before consumers), so the +/// driver can walk `steps` by index in a single pass and every producer runs +/// before its consumer. A chain may therefore have **multiple sinks** (one per +/// fan-out leaf), each a zero-branch step at some index; the only structural +/// requirement is that the last step has no unwired-forward branch, which forces +/// at least one terminal sink. +/// +/// Excluded: single-step chains (`n < 2` — nothing to fuse; the scheduled +/// single-thread path is already optimal), and any chain with a `Step2` merge +/// (two input streams — a single-worker inline drive assumes one source). Those +/// fall back to the scheduled worker pool. +#[must_use] +pub fn is_fusible_chain(steps: &[Box], graph: &ChainGraph) -> bool { + let n = steps.len(); + if n < 2 { + return false; + } + // Exactly one source, and it must be the first step. + if !steps[0].is_source() || steps[1..].iter().any(|s| s.is_source()) { + return false; + } + for i in 0..n { + let idx = StepIdx(i); + // No two-input (`Step2`) consumers: a single-worker inline drive follows + // one input stream. Sources register arity 0 (their input is implicit), + // single-input steps 1, and `Step2` 2 — so anything above 1 is a merge. + if graph.input_arity(idx) > 1 { + return false; + } + // Every output branch (one for a linear step, ≥2 for a fan-out like the + // `--rejects` split) must be wired to a strictly later step. A sink has + // zero branches, so its loop body is skipped; the last step necessarily + // has no forward target, so it must be a sink. + for b in 0..graph.branch_count(idx) { + match graph.consumer(idx, BranchIdx(b)) { + Some(StepIdx(j)) if j > i => {} + _ => return false, + } + } + } + true +} + +/// Whether `build_chain`'s fused single-thread fast path may be taken for this +/// chain. +/// +/// The fast path drives a fusible chain inline over direct buffers, skipping the +/// scheduler — and with it the per-edge instrumentation the scheduled path sets +/// up (edge [`EdgeMetrics`](super::metrics::EdgeMetrics), the occupancy sampler, +/// and the `snapshot_with_edges` bottleneck verdict). An instrumented run +/// (`InstrumentationLevel != Off`) must therefore NOT fuse, or its +/// `--pipeline-stats` / `--pipeline-trace` output would silently omit all edge / +/// occupancy data. When instrumentation is `Off` (the default), fusion is taken +/// whenever the chain is single-thread and fusible (zero-overhead fast path +/// preserved). +#[must_use] +pub fn should_fuse_single_thread( + n_threads: usize, + instrumentation: InstrumentationLevel, + steps: &[Box], + graph: &ChainGraph, +) -> bool { + n_threads == 1 && !instrumentation.is_on() && is_fusible_chain(steps, graph) +} + +/// Drive a fusible chain to completion on the calling thread, fused. +/// +/// Builds the chain's per-step contexts with direct inter-step transports (no +/// reorder stage, profile queue bounds retained — see +/// [`ErasedStep::build_fused_output_set`]), then repeatedly walks the steps in +/// topological order — popping +/// from each step's input, pushing to its output(s) — until **every** step has +/// reported [`StepOutcome::Finished`]. On a step's `Finished` the driver marks +/// all its output branches drained so downstream steps see their inputs closed +/// (the same drain propagation the scheduled driver performs, minus the +/// `Parallel` counter gate — there is exactly one instance of each step at one +/// worker). Waiting for *all* steps (not just the last) is what lets a fan-out +/// chain (e.g. the `--rejects` split) finish both its sink subchains. +/// +/// Callers must have confirmed [`is_fusible_chain`] first. +/// +/// `queue_memory_total` is the run's `--queue-memory-total`, if set. Because the +/// fused contexts are built here rather than by the caller, this function is the +/// only place that budget can be applied to the fused transports. +/// +/// `deadlock_timeout_secs` is the run's stall patience. A pass in which no step +/// progresses is retried with a backoff until this budget is exhausted, and only +/// then reported as a stall; `0` selects a built-in default rather than +/// "unbounded" (see the constant's comment for why). +/// +/// # Errors +/// +/// Returns [`PipelineError::Io`] if any step's `try_run` returned `Err` (the +/// first such error wins, carrying the originating step's name), or +/// [`PipelineError::Cancelled`] if the run was cancelled via the pipeline's +/// [`CancelHandle`](crate::signal::CancelHandle) — matching +/// [`crate::builder::Pipeline::run`]'s contract. +pub fn run_fused_single_thread( + steps: Vec>, + graph: &ChainGraph, + signal: &Arc, + stats: Option<&Arc>, + queue_memory_total: Option, + deadlock_timeout_secs: u64, +) -> Result<(), PipelineError> { + // A single no-progress pass is NOT a stall. `NoProgress` is the transient + // "input momentarily empty but not drained" outcome — a source waiting on a + // background reader thread returns it legitimately — so failing on the first + // idle pass aborts a healthy run and truncates its output. Tolerate idling + // until this wall-clock budget is exhausted, which is what the scheduled path + // does via `WorkerCore::sleep_backoff`. That path has no stall limit at all + // because the deadlock monitor catches wedges for it; the fused path is not + // monitored, so it needs its own bound rather than hanging forever. + // + // Budgeted by TIME, not by a pass count: `thread::sleep` granularity varies by + // platform (a requested 50µs can round up to ~1ms), so a fixed number of + // passes would mean wildly different real budgets across hosts. + // + // The budget comes from `PipelineConfig::deadlock_timeout_secs`, so a chain + // with a genuinely slow source (a network-backed reader, say) can raise it + // instead of having a library-chosen default abort a legitimate run. + // + // `deadlock_timeout_secs == 0` — the config default, meaning "monitor + // disarmed" on the scheduled path — falls back to `DEFAULT_STALL_BUDGET` here + // rather than meaning "unbounded". That asymmetry is deliberate and is the + // whole reason this bound exists: the scheduled path has a deadlock monitor to + // arm, and this one does not, so disabling the bound would leave a fused wedge + // with nothing at all to detect it. A caller wanting more patience raises the + // number; there is deliberately no way to remove the bound. + // Deliberately generous. This bound exists to catch a PERMANENT wedge, not to + // police a slow source, and it is the value most runs get (`deadlock_timeout_secs` + // defaults to 0). A source whose background reader blocks on a cold page cache or + // network-backed input can legitimately idle for tens of seconds, and the idle + // timer only resets on progress — so a tight default trades a real risk of failing + // a healthy run for a few seconds off the report of a wedge that has already hung. + // The costs are asymmetric; err long. + const DEFAULT_STALL_BUDGET: Duration = Duration::from_secs(60); + // Long enough to stop pegging the calling thread at 100%, short enough that it + // adds no meaningful latency to a source that is about to produce. + const IDLE_BACKOFF: Duration = Duration::from_micros(50); + + let stall_budget = if deadlock_timeout_secs > 0 { + Duration::from_secs(deadlock_timeout_secs) + } else { + DEFAULT_STALL_BUDGET + }; + + let n = steps.len(); + let contexts = build_chain_contexts_fused(&steps, graph); + // Re-bind `steps` as a local declared *after* `contexts` so reverse-declaration + // drop order runs `steps` before `contexts` on EVERY exit — the normal return, + // and a panic unwinding out of `try_run_erased` or the stall `debug_assert!`. + // The typed-handle cache in `TypedStep`/`TypedStep2` stores handles borrowed + // from `contexts` as `'static` (see the "Approved typed-handle cache" note in + // CLAUDE.md), which is sound only while every step is dropped before the + // `ChainContexts` it cached from. A by-value parameter drops *after* locals, so + // without this shadow an unwind would free `contexts` while the cached refs in + // `steps` still point into it. `mut` because the drive loop dispatches steps. + let mut steps = steps; + // Honour `--queue-memory-total` here as the scheduled path does. The fused + // transports keep each step's profiled byte bound (see + // `ErasedStep::build_fused_output_set`), so without this the user's budget + // would be silently ignored in favour of the per-step defaults. These + // contexts are local to this call, so the budget cannot be applied by + // `Pipeline::run` on its behalf. + if let Some(total) = queue_memory_total { + crate::builder::apply_initial_queue_budget(&contexts.bounded_queues, total); + } + let mut finished = vec![false; n]; + // `Some(t)` while the driver has been idle since `t`; cleared by any pass that + // makes progress. + let mut idle_since: Option = None; + + 'drive: loop { + // Bail promptly on an external cancel or a prior-pass error. + if signal.is_done() { + break; + } + let mut progressed = false; + for i in 0..n { + if finished[i] { + continue; + } + let outputs_any = contexts.outputs[i].as_ref(); + let mut ctx = + ErasedStepCtx { input: contexts.inputs[i].as_ref(), outputs: outputs_any, signal }; + + // Time the dispatch only when stats collection is on (mirrors + // `dispatch_one_step`): `Instant::now()` is non-trivial on the hot + // path, so gate it on `stats.is_some()`. + let start = stats.map(|_| Instant::now()); + let result = steps[i].try_run_erased(&mut ctx); + if let (Some(stats), Some(start)) = (stats, start) { + let elapsed_ns = u64::try_from(start.elapsed().as_nanos()).unwrap_or(u64::MAX); + let start_ns = stats.elapsed_ns().saturating_sub(elapsed_ns); + match &result { + Ok(outcome) => stats.record(StepIdx(i), *outcome, start_ns, elapsed_ns), + Err(_) => stats.record_error(StepIdx(i), start_ns, elapsed_ns), + } + } + + match result { + Ok(StepOutcome::Progress) => progressed = true, + Ok(StepOutcome::Finished) => { + // Close this step's output branches so downstream drains. + // No `StepDrainCounter` gate: one instance per step here. + steps[i].mark_outputs_drained(outputs_any); + finished[i] = true; + progressed = true; + } + Ok(StepOutcome::NoProgress | StepOutcome::Contention) => {} + Err(io_err) => { + signal + .record_error(PipelineError::Io { step: steps[i].name(), source: io_err }); + break 'drive; + } + } + } + // Done when every step has finished. For a fan-out chain (`--rejects`) + // that means BOTH sink subchains have drained — checking only the last + // step would break before an earlier-indexed reject sink had flushed. + if finished.iter().all(|&f| f) { + break; + } + if progressed { + idle_since = None; + } else { + // Idle pass. Back off and retry the same state until the budget runs + // out — only a sustained run of idle passes is evidence of a wedge. + let idle_start = *idle_since.get_or_insert_with(Instant::now); + if idle_start.elapsed() < stall_budget { + std::thread::sleep(IDLE_BACKOFF); + continue 'drive; + } + // Budget exhausted. For a well-formed fusible chain this is + // unreachable: the source either progresses or finishes, and + // `Finished` cascades drain downstream so some step always advances + // until every sink closes. Guard against an infinite spin rather than + // trusting that invariant blindly. + debug_assert!( + false, + "fused single-thread driver stalled: no step progressed for {stall_budget:?} \ + and not all steps have finished" + ); + // In release builds the `debug_assert!` is compiled out, so record + // the stall as an error before breaking. Otherwise the loop would + // exit and map to `Ok`, silently truncating the output instead of + // surfacing the broken invariant. Name the still-unfinished steps so + // the error points at the wedged step(s) rather than just the + // synthetic "fused-driver" name. + let stalled_steps: Vec<&'static str> = finished + .iter() + .enumerate() + .filter(|(_, f)| !**f) + .map(|(i, _)| steps[i].name()) + .collect(); + signal.record_error(PipelineError::Io { + step: "fused-driver", + source: std::io::Error::other(format!( + "fused single-thread driver stalled: no step progressed for {stall_budget:?} \ + and not all steps have finished; unfinished step(s): {}", + stalled_steps.join(", "), + )), + }); + break 'drive; + } + } + + // `steps` (re-bound as a local after `contexts` above) drops before `contexts` + // here by reverse-declaration order, keeping the typed-handle cache invariant. + + // Map the recorded outcome to the run result (same shape as `Pipeline::run`). + // `to_result` reconstructs the non-`Clone` `PipelineError` and synthesizes + // `Cancelled` from the state when an external cancel published the terminal + // state but its `OnceLock` payload is not yet visible to this thread. + signal.to_result() +} + +#[cfg(test)] +mod tests { + use std::io; + use std::sync::{Arc, Mutex}; + + use rstest::rstest; + + use super::*; + use crate::erased::TypedStep; + use crate::outputs::Single; + use crate::queues::QueueSpec; + use crate::reorder::BranchOrdering; + use crate::step::{Step, StepCtx, StepKind, StepProfile}; + + // ── Stub steps: source → +100 mid → collecting sink ────────────────── + + /// Source emitting `0, 1, …, count-1` then `Finished`. + struct CountSource { + next: u32, + count: u32, + } + impl Step for CountSource { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "CountSource", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::Unbounded], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if self.next >= self.count { + return Ok(StepOutcome::Finished); + } + let _ = ctx.outputs.push(self.next); + self.next += 1; + Ok(StepOutcome::Progress) + } + } + + /// Mid: pops a `u32`, pushes `+100`; `Finished` once its input is drained. + #[derive(Clone)] + struct AddHundred; + impl Step for AddHundred { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "AddHundred", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::Unbounded], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + match ctx.input.pop() { + Some(v) => { + let _ = ctx.outputs.push(v + 100); + Ok(StepOutcome::Progress) + } + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } + } + + /// Sink: collects popped values into a shared `Vec`; `Finished` on drain. + struct CollectSink { + out: Arc>>, + } + impl Step for CollectSink { + type Input = u32; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "CollectSink", + 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(v) => { + self.out.lock().unwrap().push(v); + Ok(StepOutcome::Progress) + } + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } + } + + /// Build a linear `source → mid → sink` chain (graph + boxed steps). + fn linear_chain( + count: u32, + out: &Arc>>, + ) -> (Vec>, ChainGraph) { + let mut graph = ChainGraph::new(); + let s = graph.register_step("CountSource", 1); + let m = graph.register_step("AddHundred", 1); + let k = graph.register_step("CollectSink", 0); + graph.wire(s, BranchIdx(0), m); + graph.wire(m, BranchIdx(0), k); + let steps: Vec> = vec![ + Box::new(TypedStep::new(CountSource { next: 0, count })), + Box::new(TypedStep::new(AddHundred)), + Box::new(TypedStep::new(CollectSink { out: Arc::clone(out) })), + ]; + (steps, graph) + } + + /// Fan-out split: routes even values to branch 0, odd to branch 1 — the + /// shape of a `--rejects` kept/rejects split. + struct EvenOddSplit; + impl Step for EvenOddSplit { + type Input = u32; + type Outputs = (u32, u32); + fn profile(&self) -> StepProfile { + StepProfile { + name: "EvenOddSplit", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::Unbounded, QueueSpec::Unbounded], + branch_ordering: vec![BranchOrdering::None, BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + match ctx.input.pop() { + Some(v) => { + let view = ctx.outputs.view(); + if v % 2 == 0 { + let _ = view.a.push(v); + } else { + let _ = view.b.push(v); + } + Ok(StepOutcome::Progress) + } + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } + } + + /// Build a fan-out `source → split → (even sink, odd sink)` chain. + fn fan_out_chain( + count: u32, + even: &Arc>>, + odd: &Arc>>, + ) -> (Vec>, ChainGraph) { + let mut graph = ChainGraph::new(); + let s = graph.register_step("CountSource", 1); + let m = graph.register_step("EvenOddSplit", 2); + let k0 = graph.register_step("CollectSink", 0); + let k1 = graph.register_step("CollectSink", 0); + graph.wire(s, BranchIdx(0), m); + graph.wire(m, BranchIdx(0), k0); + graph.wire(m, BranchIdx(1), k1); + let steps: Vec> = vec![ + Box::new(TypedStep::new(CountSource { next: 0, count })), + Box::new(TypedStep::new(EvenOddSplit)), + Box::new(TypedStep::new(CollectSink { out: Arc::clone(even) })), + Box::new(TypedStep::new(CollectSink { out: Arc::clone(odd) })), + ]; + (steps, graph) + } + + #[test] + fn is_fusible_detects_source_mid_sink() { + let out = Arc::new(Mutex::new(Vec::new())); + let (steps, graph) = linear_chain(3, &out); + assert!(is_fusible_chain(&steps, &graph)); + } + + #[test] + fn is_fusible_rejects_single_step() { + // A self-contained source+sink (Input=(), no output branches) is one + // step — nothing to fuse, so not eligible. + let mut graph = ChainGraph::new(); + graph.register_step("CountSource", 0); + let steps: Vec> = + vec![Box::new(TypedStep::new(CountSource { next: 0, count: 0 }))]; + assert!(!is_fusible_chain(&steps, &graph)); + } + + #[test] + fn is_fusible_rejects_empty() { + let steps: Vec> = vec![]; + let graph = ChainGraph::new(); + assert!(!is_fusible_chain(&steps, &graph)); + } + + #[test] + fn is_fusible_accepts_fan_out() { + // A fan-out (the `--rejects` shape: one step with two output branches, + // each wired forward to its own sink) IS fusible — both branches feed + // strictly later steps. + let even = Arc::new(Mutex::new(Vec::new())); + let odd = Arc::new(Mutex::new(Vec::new())); + let (steps, graph) = fan_out_chain(0, &even, &odd); + assert!(is_fusible_chain(&steps, &graph)); + } + + /// The `input_arity(idx) > 1` guard: a chain containing a `Step2` merge + /// (input arity 2) is NOT fusible — the single-worker inline drive follows + /// one input stream per step. Every other condition holds (one leading + /// source; all branches wired strictly forward), so the merge arity is the + /// sole reason fusion is refused. + #[test] + fn is_fusible_rejects_a_two_input_merge() { + let out = Arc::new(Mutex::new(Vec::new())); + let mut graph = ChainGraph::new(); + let s = graph.register_step("CountSource", 1); + let m = graph.register_step_with_input_arity("AddHundred", 1, 2); + let k = graph.register_step("CollectSink", 0); + graph.wire(s, BranchIdx(0), m); + graph.wire(m, BranchIdx(0), k); + let steps: Vec> = vec![ + Box::new(TypedStep::new(CountSource { next: 0, count: 0 })), + Box::new(TypedStep::new(AddHundred)), + Box::new(TypedStep::new(CollectSink { out: Arc::clone(&out) })), + ]; + assert!( + !is_fusible_chain(&steps, &graph), + "a Step2 merge (input_arity 2) must block fusion" + ); + } + + /// The forward-wiring guard: an output branch wired to an earlier-or-equal + /// step index is NOT fusible — the inline drive is a single topological pass, + /// so a back-edge would revisit an already-drained step. Here the mid step's + /// branch loops back to the source (consumer index 0 ≤ producer index 1); + /// every other condition holds, so the back-edge is the sole reason. + #[test] + fn is_fusible_rejects_a_backward_wired_branch() { + let out = Arc::new(Mutex::new(Vec::new())); + let mut graph = ChainGraph::new(); + let s = graph.register_step("CountSource", 1); + let m = graph.register_step("AddHundred", 1); + let _k = graph.register_step("CollectSink", 0); + graph.wire(s, BranchIdx(0), m); + graph.wire(m, BranchIdx(0), s); // back-edge: consumer index 0 ≤ producer index 1 + let steps: Vec> = vec![ + Box::new(TypedStep::new(CountSource { next: 0, count: 0 })), + Box::new(TypedStep::new(AddHundred)), + Box::new(TypedStep::new(CollectSink { out: Arc::clone(&out) })), + ]; + assert!( + !is_fusible_chain(&steps, &graph), + "a branch wired to an earlier-or-equal step must block fusion" + ); + } + + #[test] + fn is_fusible_rejects_a_second_source() { + // The single-source guard: the inline drive assumes exactly one + // implicit-input source, at index 0. A second source at index ≥ 1 must + // block fusion — otherwise the driver would dispatch a step whose input is + // implicit as if it were wired. The other three structural guards + // (`n < 2`, `input_arity > 1`, back-edge) each have a test; this pins the + // fourth. Register the second source with input arity 1 so the wire is + // accepted — the boxed step is still a `CountSource`, so + // `ErasedStep::is_source()` reports true and it is the source guard, not + // the arity/wiring checks, that rejects the chain. + let out = Arc::new(Mutex::new(Vec::new())); + let mut graph = ChainGraph::new(); + let s0 = graph.register_step("CountSource", 1); + let s1 = graph.register_step_with_input_arity("CountSource", 1, 1); + let k = graph.register_step("CollectSink", 0); + graph.wire(s0, BranchIdx(0), s1); + graph.wire(s1, BranchIdx(0), k); + let steps: Vec> = vec![ + Box::new(TypedStep::new(CountSource { next: 0, count: 0 })), + Box::new(TypedStep::new(CountSource { next: 0, count: 0 })), + Box::new(TypedStep::new(CollectSink { out })), + ]; + assert!( + !is_fusible_chain(&steps, &graph), + "a second source at index >= 1 must block fusion" + ); + } + + // A fusible chain fuses ONLY when it is single-thread AND uninstrumented: the + // fused fast path skips the scheduled path's edge metrics / occupancy sampler + // / bottleneck verdict, so any instrumentation level (or ≥2 threads) must fall + // through to the scheduled path instead of silently dropping that output. + #[rstest] + #[case::off_single_thread(1, InstrumentationLevel::Off, true)] + #[case::summary_single_thread(1, InstrumentationLevel::Summary, false)] + #[case::timeline_single_thread(1, InstrumentationLevel::Timeline, false)] + #[case::deep_single_thread(1, InstrumentationLevel::Deep, false)] + #[case::off_multi_thread(2, InstrumentationLevel::Off, false)] + fn should_fuse_only_when_uninstrumented_single_thread( + #[case] n_threads: usize, + #[case] level: InstrumentationLevel, + #[case] expected: bool, + ) { + let out = Arc::new(Mutex::new(Vec::new())); + let (steps, graph) = linear_chain(4, &out); + assert!(is_fusible_chain(&steps, &graph), "linear chain is fusible"); + assert_eq!(should_fuse_single_thread(n_threads, level, &steps, &graph), expected); + } + + #[test] + fn drive_runs_chain_to_completion_in_order() { + let out = Arc::new(Mutex::new(Vec::new())); + let (steps, graph) = linear_chain(5, &out); + let signal = PipelineSignal::new(); + run_fused_single_thread(steps, &graph, &signal, None, None, 0).expect("clean run"); + // Source emits 0..5, mid adds 100, sink collects in FIFO order. + assert_eq!(*out.lock().unwrap(), vec![100, 101, 102, 103, 104]); + } + + /// Source with a *byte-bounded* output, so a fused chain built from it + /// registers a real bounded transport. Emits `0, 1, …, count-1`, holding and + /// retrying whatever the transport rejects — the contract every step already + /// owes the scheduled path. + struct ByteBoundedSource { + next: u32, + count: u32, + held: Option, + } + impl Step for ByteBoundedSource { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "ByteBoundedSource", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: PROFILE_LIMIT_BYTES }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if let Some(item) = self.held.take() { + if let Err(unpushed) = ctx.outputs.push(item) { + self.held = Some(unpushed.into_item()); + } + return Ok(StepOutcome::Progress); + } + if self.next >= self.count { + return Ok(StepOutcome::Finished); + } + let item = self.next; + self.next += 1; + if let Err(unpushed) = ctx.outputs.push(item) { + self.held = Some(unpushed.into_item()); + } + Ok(StepOutcome::Progress) + } + } + + /// Deliberately unusual so the assertions below cannot pass by matching a + /// default or a budget-derived value. + const PROFILE_LIMIT_BYTES: u64 = 7 * 1024 * 1024; + + fn byte_bounded_chain( + count: u32, + out: &Arc>>, + ) -> (Vec>, ChainGraph) { + let mut graph = ChainGraph::new(); + let src = graph.register_step("ByteBoundedSource", 1); + let sink = graph.register_step("CollectSink", 0); + graph.wire(src, BranchIdx(0), sink); + let steps: Vec> = vec![ + Box::new(TypedStep::new(ByteBoundedSource { next: 0, count, held: None })), + Box::new(TypedStep::new(CollectSink { out: Arc::clone(out) })), + ]; + (steps, graph) + } + + /// The fused path must honour `--queue-memory-total`, as the scheduled path + /// does. Two facts are load-bearing and both were false before fused + /// transports kept their profile bounds: the fused contexts register a + /// bounded queue at all, and the budget resizes it off the profile default. + #[test] + fn fused_contexts_register_bounded_queues_and_take_the_budget() { + let out = Arc::new(Mutex::new(Vec::new())); + let (steps, graph) = byte_bounded_chain(4, &out); + let contexts = build_chain_contexts_fused(&steps, &graph); + + assert_eq!( + contexts.bounded_queues.len(), + 1, + "a byte-bounded fused edge must be registered, else there is nothing \ + for the queue-memory budget to apply to" + ); + assert_eq!( + contexts.bounded_queues[0].handle.limit_bytes(), + PROFILE_LIMIT_BYTES, + "the fused transport starts at the profile's declared bound" + ); + + // 64 MiB over one queue → 64 MiB per queue, well above the 1 MiB floor, + // and distinct from the profile default so the assert cannot pass by + // accident. + let total = 64 * 1024 * 1024; + crate::builder::apply_initial_queue_budget(&contexts.bounded_queues, total); + assert_eq!( + contexts.bounded_queues[0].handle.limit_bytes(), + total, + "the user's budget must override the per-step default on the fused path" + ); + } + + /// End-to-end: a byte-bounded fused chain run with a budget still delivers + /// every item in order. Pins that threading the budget through does not wedge + /// the driver — a bounded fused edge relies on the producer's hold-and-retry. + #[rstest] + #[case::no_budget(None)] + #[case::with_budget(Some(2 * 1024 * 1024))] + fn fused_byte_bounded_chain_completes(#[case] queue_memory_total: Option) { + let out = Arc::new(Mutex::new(Vec::new())); + let (steps, graph) = byte_bounded_chain(5, &out); + let signal = PipelineSignal::new(); + run_fused_single_thread(steps, &graph, &signal, None, queue_memory_total, 0) + .expect("clean run"); + assert_eq!(*out.lock().unwrap(), vec![0, 1, 2, 3, 4]); + } + + /// Source that reports `NoProgress` for its first few dispatches before it + /// starts emitting — the shape of a source waiting on a background reader. + struct IdleThenEmitSource { + idle_left: u32, + remaining: u32, + dispatches: Arc>, + } + impl Step for IdleThenEmitSource { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "IdleThenEmitSource", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + *self.dispatches.lock().unwrap() += 1; + if self.idle_left > 0 { + self.idle_left -= 1; + // Legitimate transient: nothing to hand over *yet*. + return Ok(StepOutcome::NoProgress); + } + if self.remaining == 0 { + return Ok(StepOutcome::Finished); + } + let item = self.remaining; + self.remaining -= 1; + let _ = ctx.outputs.push(item); + Ok(StepOutcome::Progress) + } + } + + /// A pass in which no step progressed is not a stall — `NoProgress` is the + /// transient "input momentarily empty but not drained" outcome. The driver + /// used to fail the run on the FIRST such pass, recording `PipelineError::Io` + /// and truncating the output (and panicking through the `debug_assert!` in + /// test builds). It must back off and retry instead, so a source that idles + /// before producing still completes. + #[test] + fn drive_tolerates_transient_no_progress_passes() { + const IDLE_PASSES: u32 = 5; + const N_ITEMS: u32 = 3; + + let mut graph = ChainGraph::new(); + let src = graph.register_step("IdleThenEmitSource", 1); + let sink = graph.register_step("CollectSink", 0); + graph.wire(src, BranchIdx(0), sink); + + let out = Arc::new(Mutex::new(Vec::new())); + let dispatches = Arc::new(Mutex::new(0)); + let steps: Vec> = vec![ + Box::new(TypedStep::new(IdleThenEmitSource { + idle_left: IDLE_PASSES, + remaining: N_ITEMS, + dispatches: Arc::clone(&dispatches), + })), + Box::new(TypedStep::new(CollectSink { out: Arc::clone(&out) })), + ]; + let signal = PipelineSignal::new(); + run_fused_single_thread(steps, &graph, &signal, None, None, 0) + .expect("transient NoProgress must not fail the run"); + + assert_eq!( + *out.lock().unwrap(), + vec![3, 2, 1], + "every item must still be delivered after the idle passes" + ); + assert!(!signal.is_done(), "no error recorded for a transient idle pass"); + assert!( + *dispatches.lock().unwrap() > IDLE_PASSES, + "the driver must have retried past the idle passes, not given up on the first" + ); + } + + /// Source that never progresses and never finishes — a permanent wedge, the + /// case the stall bound exists for. + struct WedgedSource; + impl Step for WedgedSource { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "WedgedSource", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + /// The stall budget is `deadlock_timeout_secs`, and `0` selects the built-in + /// default rather than "unbounded" — a fused wedge has no deadlock monitor to + /// fall back on, so the bound must not be disableable. A 1-second budget is + /// used here to keep the test fast; the assertion is that the run fails inside + /// its own budget and names the wedged step. + /// + /// Gated `#[should_panic]` rather than an unconditional one: the + /// `debug_assert!` on the stall path fires first in a debug test build, which + /// is itself the contract (a genuine wedge must be loud in tests). In a + /// release test build that `debug_assert!` is compiled out, so the driver + /// records the stall and returns `Err` instead — the in-body assertions cover + /// that path, and the `should_panic` expectation must not apply there. + #[test] + #[cfg_attr(debug_assertions, should_panic(expected = "fused single-thread driver stalled"))] + fn drive_reports_a_real_stall_within_the_configured_budget() { + let mut graph = ChainGraph::new(); + let src = graph.register_step("WedgedSource", 1); + let sink = graph.register_step("CollectSink", 0); + graph.wire(src, BranchIdx(0), sink); + + let out = Arc::new(Mutex::new(Vec::new())); + let steps: Vec> = vec![ + Box::new(TypedStep::new(WedgedSource)), + Box::new(TypedStep::new(CollectSink { out })), + ]; + let signal = PipelineSignal::new(); + let started = Instant::now(); + let result = run_fused_single_thread(steps, &graph, &signal, None, None, 1); + // Only reached in a release test build, where the `debug_assert!` is gone. + assert!(result.is_err(), "a permanent wedge must not report success"); + assert!( + started.elapsed() < Duration::from_secs(30), + "the configured 1s budget must bound the wait, not the default" + ); + } + + #[test] + fn drive_runs_fan_out_to_completion() { + // Both sink subchains of a fan-out must drain — the driver waits for ALL + // steps to finish, not just the last-indexed one. + let even = Arc::new(Mutex::new(Vec::new())); + let odd = Arc::new(Mutex::new(Vec::new())); + let (steps, graph) = fan_out_chain(6, &even, &odd); + let signal = PipelineSignal::new(); + run_fused_single_thread(steps, &graph, &signal, None, None, 0).expect("clean run"); + // Source emits 0..6; evens route to branch 0, odds to branch 1. + assert_eq!(*even.lock().unwrap(), vec![0, 2, 4]); + assert_eq!(*odd.lock().unwrap(), vec![1, 3, 5]); + } + + #[test] + fn drive_propagates_step_error() { + /// Mid that errors on the first item. + struct Boom; + impl Step for Boom { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "Boom", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::Unbounded], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if ctx.input.pop().is_some() { + return Err(io::Error::other("boom")); + } + if ctx.input.is_drained() { + Ok(StepOutcome::Finished) + } else { + Ok(StepOutcome::NoProgress) + } + } + } + + let out = Arc::new(Mutex::new(Vec::new())); + let mut graph = ChainGraph::new(); + let s = graph.register_step("CountSource", 1); + let m = graph.register_step("Boom", 1); + let k = graph.register_step("CollectSink", 0); + graph.wire(s, BranchIdx(0), m); + graph.wire(m, BranchIdx(0), k); + let steps: Vec> = vec![ + Box::new(TypedStep::new(CountSource { next: 0, count: 3 })), + Box::new(TypedStep::new(Boom)), + Box::new(TypedStep::new(CollectSink { out })), + ]; + let signal = PipelineSignal::new(); + let err = + run_fused_single_thread(steps, &graph, &signal, None, None, 0).expect_err("must error"); + assert!(matches!(err, PipelineError::Io { step: "Boom", .. })); + } + + #[test] + fn drive_records_stats_for_success_and_error_dispatches() { + // `PipelineBuilder::run` can hand the fused driver `Some(&Arc)` + // even with instrumentation Off; the driver's stats block then records + // every dispatch. A single erroring chain exercises BOTH arms: the source + // dispatches `Progress` (→ `record`), then the mid errors (→ `record_error`) + // before the loop breaks. Every other fused test runs with `None` stats. + struct Boom; + impl Step for Boom { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "Boom", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::Unbounded], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if ctx.input.pop().is_some() { + return Err(io::Error::other("boom")); + } + if ctx.input.is_drained() { + Ok(StepOutcome::Finished) + } else { + Ok(StepOutcome::NoProgress) + } + } + } + + let out = Arc::new(Mutex::new(Vec::new())); + let mut graph = ChainGraph::new(); + let s = graph.register_step("CountSource", 1); + let m = graph.register_step("Boom", 1); + let k = graph.register_step("CollectSink", 0); + graph.wire(s, BranchIdx(0), m); + graph.wire(m, BranchIdx(0), k); + let steps: Vec> = vec![ + Box::new(TypedStep::new(CountSource { next: 0, count: 3 })), + Box::new(TypedStep::new(Boom)), + Box::new(TypedStep::new(CollectSink { out })), + ]; + let stats = Arc::new(PipelineStats::new(vec!["CountSource", "Boom", "CollectSink"])); + let signal = PipelineSignal::new(); + run_fused_single_thread(steps, &graph, &signal, Some(&stats), None, 0) + .expect_err("must error"); + + let snap = stats.snapshot(); + // `record` logged the source's Progress dispatch. + assert_eq!(snap.steps[0].0, "CountSource"); + assert!( + snap.steps[0].1.progress_count > 0, + "record must log the source's progress under Some(stats)" + ); + // `record_error` logged Boom's failure (error_count, not an outcome bucket). + assert_eq!(snap.steps[1].0, "Boom"); + assert!( + snap.steps[1].1.error_count > 0, + "record_error must log the failing step under Some(stats)" + ); + assert!(snap.steps[1].1.try_run_total > 0, "the erroring dispatch is still counted"); + } + + /// The `# Errors` contract promises `PipelineError::Cancelled` when the run + /// is cancelled through the `CancelHandle`, and the driver's only delivery + /// of it is the top-of-loop `signal.is_done()` break mapped by + /// `to_result()`. Every other test drives a clean run, a budget, a transient + /// idle, a stall, or a step error — none a cancel — so a regression that + /// moved or dropped this check (for example, breaking only once every step + /// finished) would keep the whole suite green while a cancelled fused run + /// reported `Ok`. + #[test] + fn drive_maps_an_external_cancel_to_cancelled() { + /// Source that cancels the shared signal on its first dispatch, then + /// keeps emitting. Models an external `CancelHandle::cancel` landing + /// mid-run: the typed `StepCtx` deliberately does not expose the signal, + /// so the step holds its own clone, exactly as an external canceller + /// (which owns a `CancelHandle` over the same signal) does. + struct CancelOnFirstDispatch { + signal: Arc, + next: u32, + count: u32, + } + impl Step for CancelOnFirstDispatch { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "CancelOnFirstDispatch", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::Unbounded], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if self.next == 0 { + self.signal.cancel(); + } + if self.next >= self.count { + return Ok(StepOutcome::Finished); + } + let _ = ctx.outputs.push(self.next); + self.next += 1; + Ok(StepOutcome::Progress) + } + } + + const COUNT: u32 = 8; + let out = Arc::new(Mutex::new(Vec::new())); + let signal = PipelineSignal::new(); + let mut graph = ChainGraph::new(); + let s = graph.register_step("CancelOnFirstDispatch", 1); + let k = graph.register_step("CollectSink", 0); + graph.wire(s, BranchIdx(0), k); + let steps: Vec> = vec![ + Box::new(TypedStep::new(CancelOnFirstDispatch { + signal: Arc::clone(&signal), + next: 0, + count: COUNT, + })), + Box::new(TypedStep::new(CollectSink { out: Arc::clone(&out) })), + ]; + let err = run_fused_single_thread(steps, &graph, &signal, None, None, 0) + .expect_err("a cancelled run must not report success"); + assert!( + matches!(err, PipelineError::Cancelled), + "an external cancel must map to PipelineError::Cancelled, got {err:?}" + ); + assert!( + out.lock().unwrap().len() < COUNT as usize, + "the driver must break on the cancel before draining every item" + ); + } +} diff --git a/crates/fgumi-pipeline-core/src/runtime/live.rs b/crates/fgumi-pipeline-core/src/runtime/live.rs new file mode 100644 index 000000000..4600e3b8b --- /dev/null +++ b/crates/fgumi-pipeline-core/src/runtime/live.rs @@ -0,0 +1,154 @@ +//! `LiveSteps`: per-worker worklist of still-dispatchable steps. +//! +//! Each worker dispatches steps by walking a `Vec` of the steps it +//! still has work for, in chain order. A step is **removed** from this list +//! when it returns `StepOutcome::Finished`. The worker exits once the list is +//! empty. +//! +//! This replaces the previous scheme of mutating each finished step's +//! `WorkerStepEntry` to `WorkerStepEntry::Skip` in place and re-scanning the +//! full `entries` vec every pass (`entries.iter().all(|e| !e.is_dispatchable())` +//! for the exit check, plus a per-pass `continue` over every inert `Skip` +//! slot). The worklist: +//! +//! - never re-visits a finished or build-time-excluded step, +//! - makes "done" an honest removal rather than an in-place inert variant, and +//! - keeps `StepIdx` stable as the canonical identity into the parallel +//! `entries` / `contexts.inputs` / `contexts.outputs` / `drain_counters` +//! arrays (the worklist holds indices *into* that stable storage; it does +//! not renumber anything). +//! +//! Removal is **stable** (`Vec::remove`, not `swap_remove`): the chain-order +//! invariant the round-robin dispatch relies on (attempt upstream steps before +//! downstream, restart from the front on `Progress`) must survive a removal. At +//! the handful of steps in a chain the linear `remove`/`position` cost is +//! irrelevant. + +use crate::runtime::storage::WorkerStepEntry; +use crate::topology::StepIdx; + +/// One per worker: the steps this worker can still dispatch, in chain order. +pub struct LiveSteps { + order: Vec, +} + +impl LiveSteps { + /// Build from this worker's storage. Every dispatchable entry (i.e. not a + /// build-time `WorkerStepEntry::Skip` placeholder for an Exclusive step + /// owned by another worker, or a Serial step this worker's affinity gates + /// out) enters the worklist, in `StepIdx` (chain) order. + #[must_use] + pub fn from_entries(entries: &[WorkerStepEntry]) -> Self { + let order = entries + .iter() + .enumerate() + .filter(|(_, entry)| entry.is_dispatchable()) + .map(|(idx, _)| StepIdx(idx)) + .collect(); + Self { order } + } + + /// The steps this worker can still dispatch, in chain order. + #[must_use] + pub fn order(&self) -> &[StepIdx] { + &self.order + } + + /// Number of still-dispatchable steps. + #[must_use] + pub fn len(&self) -> usize { + self.order.len() + } + + /// `true` once this worker has nothing left to dispatch (loop exit). + #[must_use] + pub fn is_empty(&self) -> bool { + self.order.is_empty() + } + + /// `true` if `step_idx` is still dispatchable by this worker. + #[must_use] + pub fn contains(&self, step_idx: StepIdx) -> bool { + self.order.contains(&step_idx) + } + + /// Stably remove a finished step. No-op if already absent (idempotent). + /// Preserves the relative order of the remaining steps. + pub fn remove(&mut self, step_idx: StepIdx) { + if let Some(pos) = self.order.iter().position(|&s| s == step_idx) { + self.order.remove(pos); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + + use parking_lot::Mutex; + + use crate::erased::{ErasedStep, TypedStep}; + use crate::outputs::Single; + use crate::queues::QueueSpec; + use crate::reorder::BranchOrdering; + use crate::runtime::storage::DrainGate; + use crate::step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}; + + #[derive(Clone)] + struct Nop; + impl Step for Nop { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "Nop", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 1 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> std::io::Result { + Ok(StepOutcome::NoProgress) + } + } + + fn shared_entry() -> WorkerStepEntry { + let step: Box = Box::new(TypedStep::new(Nop)); + WorkerStepEntry::Shared { + step: Arc::new(Mutex::new(step)), + drain: Arc::new(DrainGate::default()), + } + } + + #[test] + fn from_entries_excludes_build_time_skips() { + let entries = + vec![shared_entry(), WorkerStepEntry::Skip, shared_entry(), WorkerStepEntry::Skip]; + let live = LiveSteps::from_entries(&entries); + assert_eq!(live.order(), &[StepIdx(0), StepIdx(2)]); + assert!(!live.is_empty()); + assert_eq!(live.len(), 2); + } + + #[test] + fn remove_is_stable_and_idempotent() { + let entries = vec![shared_entry(), shared_entry(), shared_entry()]; + let mut live = LiveSteps::from_entries(&entries); + assert_eq!(live.order(), &[StepIdx(0), StepIdx(1), StepIdx(2)]); + + // Remove the middle; the survivors keep their relative order. + live.remove(StepIdx(1)); + assert_eq!(live.order(), &[StepIdx(0), StepIdx(2)]); + assert!(!live.contains(StepIdx(1))); + + // Removing an already-absent step is a no-op. + live.remove(StepIdx(1)); + assert_eq!(live.order(), &[StepIdx(0), StepIdx(2)]); + + live.remove(StepIdx(0)); + live.remove(StepIdx(2)); + assert!(live.is_empty()); + } +} diff --git a/crates/fgumi-pipeline-core/src/runtime/metrics.rs b/crates/fgumi-pipeline-core/src/runtime/metrics.rs new file mode 100644 index 000000000..7d4981380 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/runtime/metrics.rs @@ -0,0 +1,308 @@ +//! Per-edge instrumentation counters + occupancy histogram. +//! +//! An [`EdgeMetrics`] is the edge-side complement of the step-side +//! [`PipelineStats`](super::stats::PipelineStats): one per instrumented queue +//! edge, shared (`Arc`) between the producer's transport (push/reject counts) +//! and the consumer's input handle (pop/empty counts), and sampled periodically +//! for occupancy. All counters are `Relaxed` atomics — these are statistics, not +//! synchronization (staleness is fine), mirroring `ByteBoundedQueue::current_bytes`. +//! +//! What the histogram alone can classify is [`RawOccupancy`]; the richer +//! `Empty`-vs-`Starved` / `Full`-vs-`Backpressured` split needs the reject/empty +//! *rates* and is done by the renderer, not here. + +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +/// Number of occupancy histogram buckets (depth fraction `0.0..=1.0` split into +/// `OCCUPANCY_BUCKETS` equal bands; the top band captures exactly-full). +pub const OCCUPANCY_BUCKETS: usize = 8; + +/// Per-edge counters + occupancy histogram. Construct with [`EdgeMetrics::new`] +/// (returns an `Arc` so the producer transport and consumer input handle share +/// one instance). All methods are lock-free `Relaxed` atomic updates. +#[derive(Debug)] +pub struct EdgeMetrics { + /// Items the producer successfully pushed. + pushed_items: AtomicU64, + /// Bytes pushed (only meaningful for byte-bounded edges; `0` otherwise). + pushed_bytes: AtomicU64, + /// Items the consumer popped. + popped_items: AtomicU64, + /// Bytes popped (byte-bounded edges only). + popped_bytes: AtomicU64, + /// `try_push` rejections — backpressure events (producer wanted to push, the + /// edge was full). + push_rejections: AtomicU64, + /// Consumer `pop` on an empty edge — starvation events. + pop_empties: AtomicU64, + /// Number of occupancy samples taken (sampler ticks). + depth_samples: AtomicU64, + /// Histogram of occupancy fraction at sample time. + occupancy_buckets: [AtomicU64; OCCUPANCY_BUCKETS], + /// Σ of `depth_fraction * 1000` over all samples, for the mean. + occupancy_sum_milli: AtomicU64, + /// Σ of raw occupied **bytes** over all samples. Kept alongside the fraction + /// sum so the mean occupancy can be reported in absolute bytes — which, + /// unlike `fraction × final_limit`, stays correct when the byte limit is + /// rebalanced at runtime (`queue_memory_total`). For an ordered edge the + /// sampled bytes include the `ReorderStage` overflow stash. + occupancy_bytes_sum: AtomicU64, +} + +impl EdgeMetrics { + /// Construct a fresh metrics instance, shared via `Arc`. + #[must_use] + pub fn new() -> Arc { + Arc::new(Self { + pushed_items: AtomicU64::new(0), + pushed_bytes: AtomicU64::new(0), + popped_items: AtomicU64::new(0), + popped_bytes: AtomicU64::new(0), + push_rejections: AtomicU64::new(0), + pop_empties: AtomicU64::new(0), + depth_samples: AtomicU64::new(0), + occupancy_buckets: std::array::from_fn(|_| AtomicU64::new(0)), + occupancy_sum_milli: AtomicU64::new(0), + occupancy_bytes_sum: AtomicU64::new(0), + }) + } + + /// Record a successful producer push of `bytes` (pass `0` for count-bounded + /// / unbounded edges that don't track bytes). + pub fn record_push(&self, bytes: u64) { + self.pushed_items.fetch_add(1, Ordering::Relaxed); + self.pushed_bytes.fetch_add(bytes, Ordering::Relaxed); + } + + /// Record a successful consumer pop of `bytes`. + pub fn record_pop(&self, bytes: u64) { + self.popped_items.fetch_add(1, Ordering::Relaxed); + self.popped_bytes.fetch_add(bytes, Ordering::Relaxed); + } + + /// Record a `try_push` rejection (backpressure). + pub fn record_reject(&self) { + self.push_rejections.fetch_add(1, Ordering::Relaxed); + } + + /// Record a consumer pop that found the edge empty (starvation). + pub fn record_empty(&self) { + self.pop_empties.fetch_add(1, Ordering::Relaxed); + } + + /// Record one occupancy sample from the edge's current `occupied_bytes` and + /// its byte `limit_bytes`. The histogram bucket and fraction use + /// `occupied_bytes / limit_bytes` clamped to `0.0..=1.0`; the **raw** + /// `occupied_bytes` (un-clamped — so it stays truthful when an ordered + /// edge's reorder stash pushes total buffered bytes past the transport + /// limit) is summed for the byte-accurate mean the latency estimate uses. A + /// `limit_bytes` of `0` is treated as fraction `0.0` defensively (the + /// sampler already skips count/unbounded edges). + // Casts are bounded statistics: `f ∈ [0,1]` so `f * BUCKETS` ∈ [0,8] and + // `f * 1000` ∈ [0,1000] — no meaningful truncation/sign loss/precision loss. + #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss, clippy::cast_precision_loss)] + pub fn record_depth(&self, occupied_bytes: u64, limit_bytes: u64) { + let f = if limit_bytes == 0 { + 0.0 + } else { + (occupied_bytes as f32 / limit_bytes as f32).clamp(0.0, 1.0) + }; + let bucket = ((f * OCCUPANCY_BUCKETS as f32) as usize).min(OCCUPANCY_BUCKETS - 1); + self.occupancy_buckets[bucket].fetch_add(1, Ordering::Relaxed); + self.occupancy_sum_milli.fetch_add((f * 1000.0) as u64, Ordering::Relaxed); + self.occupancy_bytes_sum.fetch_add(occupied_bytes, Ordering::Relaxed); + self.depth_samples.fetch_add(1, Ordering::Relaxed); + } + + /// Take a consistent-enough point-in-time snapshot of the counters. + // Mean is a display statistic; u64→f32 precision loss past 2^23 samples is + // irrelevant to a 0..1 occupancy mean. + #[allow(clippy::cast_precision_loss)] + #[must_use] + pub fn snapshot(&self) -> EdgeMetricsSnapshot { + let buckets: [u64; OCCUPANCY_BUCKETS] = + std::array::from_fn(|i| self.occupancy_buckets[i].load(Ordering::Relaxed)); + let depth_samples = self.depth_samples.load(Ordering::Relaxed); + let mean_occupancy = if depth_samples == 0 { + 0.0 + } else { + (self.occupancy_sum_milli.load(Ordering::Relaxed) as f32 / 1000.0) + / depth_samples as f32 + }; + let mean_occupancy_bytes = if depth_samples == 0 { + 0.0 + } else { + self.occupancy_bytes_sum.load(Ordering::Relaxed) as f64 / depth_samples as f64 + }; + EdgeMetricsSnapshot { + pushed_items: self.pushed_items.load(Ordering::Relaxed), + pushed_bytes: self.pushed_bytes.load(Ordering::Relaxed), + popped_items: self.popped_items.load(Ordering::Relaxed), + popped_bytes: self.popped_bytes.load(Ordering::Relaxed), + push_rejections: self.push_rejections.load(Ordering::Relaxed), + pop_empties: self.pop_empties.load(Ordering::Relaxed), + depth_samples, + raw_occupancy: RawOccupancy::from_buckets(&buckets, depth_samples), + mean_occupancy, + mean_occupancy_bytes, + } + } +} + +/// A point-in-time read of an [`EdgeMetrics`]. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct EdgeMetricsSnapshot { + /// Items successfully pushed onto the edge by the producer. + pub pushed_items: u64, + /// Bytes successfully pushed onto the edge (0 for count/unbounded edges). + pub pushed_bytes: u64, + /// Items successfully popped off the edge by the consumer. + pub popped_items: u64, + /// Bytes successfully popped off the edge (0 for count/unbounded edges). + pub popped_bytes: u64, + /// Push attempts rejected by backpressure (edge at its limit). + pub push_rejections: u64, + /// Pop attempts that found the edge empty (starvation signal). + pub pop_empties: u64, + /// Number of occupancy samples the background sampler took on this edge. + pub depth_samples: u64, + /// What the occupancy histogram alone says (no rate-based refinement). + pub raw_occupancy: RawOccupancy, + /// Mean occupancy fraction `0.0..=1.0` over all samples (`0.0` if none). + pub mean_occupancy: f32, + /// Mean occupied **bytes** over all samples (`0.0` if none). Unlike + /// `mean_occupancy` (a fraction that must be multiplied by a limit to + /// recover bytes), this is measured directly at sample time, so the + /// Little's-Law latency estimate stays correct even when `queue_memory_total` + /// rebalances the byte limit mid-run. For an ordered edge it includes the + /// `ReorderStage` overflow stash. + pub mean_occupancy_bytes: f64, +} + +/// Occupancy classification derivable from the histogram **alone** (no +/// reject/empty rates). The renderer refines `MostlyEmpty`→`Empty`/`Starved` +/// and `MostlyFull`→`Full`/`Backpressured` using the rates. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RawOccupancy { + /// Insufficient samples to classify. + Unknown, + /// ≥80% of samples in the bottom bucket. + MostlyEmpty, + /// ≥80% of samples in the top bucket. + MostlyFull, + /// Bottom and top buckets each hold >25% of samples (bursty coupling). + Bimodal, + /// Spread across the middle — neither side bound. + Healthy, +} + +impl RawOccupancy { + /// Classify from the bucket histogram and total sample count. + // Ratios are display statistics; u64→f64 precision loss is irrelevant to the + // 0.25/0.80 thresholds. + #[allow(clippy::cast_precision_loss)] + #[must_use] + pub fn from_buckets(buckets: &[u64; OCCUPANCY_BUCKETS], samples: u64) -> Self { + if samples == 0 { + return Self::Unknown; + } + let s = samples as f64; + let bottom = buckets[0] as f64 / s; + let top = buckets[OCCUPANCY_BUCKETS - 1] as f64 / s; + if bottom >= 0.80 { + Self::MostlyEmpty + } else if top >= 0.80 { + Self::MostlyFull + } else if bottom > 0.25 && top > 0.25 { + Self::Bimodal + } else { + Self::Healthy + } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[test] + fn edge_metrics_counts_are_exact() { + let m = EdgeMetrics::new(); + for _ in 0..3 { + m.record_push(100); + } + m.record_reject(); + for _ in 0..2 { + m.record_pop(100); + } + m.record_empty(); + let s = m.snapshot(); + assert_eq!(s.pushed_items, 3); + assert_eq!(s.pushed_bytes, 300); + assert_eq!(s.push_rejections, 1); + assert_eq!(s.popped_items, 2); + assert_eq!(s.popped_bytes, 200); + assert_eq!(s.pop_empties, 1); + } + + // Each case records a depth pattern, then asserts the classification (and, + // where the pattern pins one, the mean occupancy within a tolerance). The + // `mostly_empty` case reuses the `(mean, tol)` form as `(0.0, 0.02)`, i.e. + // `mean < 0.02`, since occupancy is non-negative. `fill` is a non-capturing + // closure coerced to a `fn` pointer so it can ride as an rstest case value. + #[rstest] + // Samples are `(occupied_bytes, limit_bytes)`; a 1000-byte limit makes the + // occupied byte count read directly as the depth fraction ×1000. + #[case::mostly_full(|m: &EdgeMetrics| for _ in 0..100 { m.record_depth(1000, 1000); }, RawOccupancy::MostlyFull, Some((1.0, 0.02)))] + #[case::mostly_empty(|m: &EdgeMetrics| for _ in 0..100 { m.record_depth(0, 1000); }, RawOccupancy::MostlyEmpty, Some((0.0, 0.02)))] + #[case::bimodal(|m: &EdgeMetrics| for _ in 0..50 { m.record_depth(0, 1000); m.record_depth(1000, 1000); }, RawOccupancy::Bimodal, None)] + // Spread uniformly across the middle buckets → mean ≈ 0.45. + #[case::healthy(|m: &EdgeMetrics| for i in 0..100u32 { m.record_depth(300 + u64::from(i % 4) * 100, 1000); }, RawOccupancy::Healthy, Some((0.45, 0.05)))] + #[case::unknown(|_m: &EdgeMetrics| {}, RawOccupancy::Unknown, None)] + fn occupancy_classification( + #[case] fill: fn(&EdgeMetrics), + #[case] expected: RawOccupancy, + #[case] mean_within: Option<(f32, f32)>, + ) { + let m = EdgeMetrics::new(); + fill(&m); + let s = m.snapshot(); + assert_eq!(s.raw_occupancy, expected); + if let Some((mean, tol)) = mean_within { + assert!( + (s.mean_occupancy - mean).abs() < tol, + "mean_occupancy {} not within {tol} of {mean}", + s.mean_occupancy + ); + } + } + + #[test] + fn mean_occupancy_bytes_tracks_absolute_bytes_independent_of_limit() { + // Sampled occupied bytes are averaged in absolute terms, not as a + // fraction of the limit — so a later limit change can't distort the mean + // (the property the derived-latency estimate relies on). Two 600-byte + // samples average to 600 bytes regardless of the 1000-byte limit used. + let m = EdgeMetrics::new(); + m.record_depth(600, 1000); + m.record_depth(600, 1000); + let s = m.snapshot(); + assert!((s.mean_occupancy_bytes - 600.0).abs() < 1e-9, "{}", s.mean_occupancy_bytes); + assert!((s.mean_occupancy - 0.6).abs() < 0.01, "{}", s.mean_occupancy); + } + + #[test] + fn mean_occupancy_bytes_keeps_raw_bytes_when_stash_exceeds_limit() { + // An ordered edge's reorder stash can push total buffered bytes past the + // transport limit; the fraction saturates at 1.0 but the byte mean stays + // truthful (1500 bytes), so latency is not under-counted. + let m = EdgeMetrics::new(); + m.record_depth(1500, 1000); // 150% of limit + let s = m.snapshot(); + assert!((s.mean_occupancy - 1.0).abs() < 1e-6, "fraction clamps to 1.0"); + assert!((s.mean_occupancy_bytes - 1500.0).abs() < 1e-9, "raw bytes preserved"); + } +} diff --git a/crates/fgumi-pipeline-core/src/runtime/mod.rs b/crates/fgumi-pipeline-core/src/runtime/mod.rs new file mode 100644 index 000000000..39143a3ac --- /dev/null +++ b/crates/fgumi-pipeline-core/src/runtime/mod.rs @@ -0,0 +1,30 @@ +//! Runtime: per-worker step storage, chain contexts, drain coordination, +//! worker pool, worker loop body. Built on top of Phase 1's trait surface. + +pub mod contexts; +pub mod detached; +pub mod drain; +pub mod driver; +pub mod fused; +pub mod live; +pub mod metrics; +pub mod pool; +pub mod sampler; +pub mod scheduler; +pub mod stats; +pub mod storage; +pub mod worker_core; + +pub use contexts::{ChainContexts, build_chain_contexts, build_chain_contexts_fused}; +pub use detached::{ + DetachedDriverGroup, build_driver_storage, extract_detached_steps, run_detached_driver, +}; +pub use drain::StepDrainCounter; +pub use driver::run_worker_loop; +pub use fused::{is_fusible_chain, run_fused_single_thread, should_fuse_single_thread}; +pub use live::LiveSteps; +pub use pool::{assign_exclusive_owners, assign_sticky_owners}; +pub use scheduler::{ChainOrderScheduler, DrainFirstScheduler, Scheduler, WalkDirection}; +pub use stats::{PipelineStats, StatsSnapshot, StepStatsSnapshot}; +pub use storage::{WorkerStepEntry, build_worker_storage}; +pub use worker_core::{BackoffPolicy, WorkerCore, WorkerRole}; diff --git a/crates/fgumi-pipeline-core/src/runtime/pool.rs b/crates/fgumi-pipeline-core/src/runtime/pool.rs new file mode 100644 index 000000000..51737618b --- /dev/null +++ b/crates/fgumi-pipeline-core/src/runtime/pool.rs @@ -0,0 +1,296 @@ +//! Exclusive owner assignment. + +use crate::erased::ErasedStep; +use crate::signal::PipelineError; +use crate::step::StepKind; +use crate::topology::StepIdx; + +/// Assign Exclusive steps to specific worker threads in chain declaration +/// order. Returns `Ok(owners)` where `owners[step_idx] == Some(worker_id)` +/// for Exclusive steps and `None` otherwise. +/// +/// # Errors +/// +/// Returns `PipelineError::NotEnoughThreads` if more Exclusive steps exist +/// than worker threads available. +pub fn assign_exclusive_owners( + steps: &[Box], + n_threads: usize, +) -> Result>, PipelineError> { + let total_exclusive = steps.iter().filter(|s| s.kind() == StepKind::Exclusive).count(); + if total_exclusive > n_threads { + return Err(PipelineError::NotEnoughThreads { + required: total_exclusive, + available: n_threads, + }); + } + + let mut owners: Vec> = vec![None; steps.len()]; + let mut next_owner = 0usize; + for (idx, step) in steps.iter().enumerate() { + if step.kind() == StepKind::Exclusive { + owners[idx] = Some(next_owner); + next_owner += 1; + } + } + Ok(owners) +} + +/// Compute each worker's sticky-driven step (the step its `WorkerCore::sticky_owner` +/// will hold). A step contributes to a worker's `sticky_owner` iff it's flagged +/// `sticky=true` in its profile AND that worker is its sole eligible dispatcher: +/// - `Exclusive sticky` step → sticky owner is the worker assigned by +/// `assign_exclusive_owners` (read from `exclusive_owners[step_idx]`). +/// - `Serial sticky` + `Affinity::Reader` → sticky owner is worker 0. +/// - `Serial sticky` + `Affinity::Writer` → sticky owner is worker `N-1`. +/// - `Serial sticky` + `Affinity::Worker(idx)` → sticky owner is worker `idx`. +/// - `Serial sticky` + `Affinity::None` → no sticky owner (every worker +/// is eligible, so no single worker can drive sticky without starving +/// the others; we silently drop the sticky hint here). +/// - `Parallel` steps are never sticky-driven (each worker has its own +/// clone — sticky drive on one would not gate others). +/// +/// If two steps' sticky-ownership rules collide on the same worker, the first +/// writer to that worker's slot wins — the framework only models one +/// sticky-owned step per worker today. The Exclusive pass runs before the Serial +/// pass and each pass only fills empty slots, so an Exclusive owner beats a +/// later Serial-sticky-Affinity target on the same worker; the same first-wins +/// rule resolves Exclusive-vs-Exclusive and Serial-vs-Serial collisions too. +/// Returns `None` for workers without any sticky-owned step. +#[must_use] +pub fn assign_sticky_owners( + steps: &[Box], + exclusive_owners: &[Option], + n_workers: usize, +) -> Vec> { + debug_assert_eq!(steps.len(), exclusive_owners.len()); + let mut sticky: Vec> = vec![None; n_workers]; + + // First pass: Exclusive-sticky owners (highest priority). + for (step_usize, step) in steps.iter().enumerate() { + if step.kind() == StepKind::Exclusive + && step.sticky() + && let Some(owner) = exclusive_owners[step_usize] + && owner < n_workers + && sticky[owner].is_none() + { + sticky[owner] = Some(StepIdx(step_usize)); + } + } + + // Second pass: Serial-sticky-Affinity owners (only fill empty slots). + for (step_usize, step) in steps.iter().enumerate() { + // Resolve through the shared `Affinity::target_worker` helper rather + // than re-matching the variants here: `Affinity::eligible` gates which + // worker may dispatch the step, and a local copy of that mapping could + // drift and hand a sticky owner to a worker that `Skip`s the step. + if step.kind() == StepKind::Serial + && step.sticky() + && let Some(target) = step.affinity().target_worker(n_workers) + && target < n_workers + && sticky[target].is_none() + { + sticky[target] = Some(StepIdx(step_usize)); + } + } + + sticky +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io; + + use rstest::rstest; + + use crate::erased::TypedStep; + use crate::outputs::Single; + use crate::queues::QueueSpec; + use crate::reorder::BranchOrdering; + use crate::step::{Affinity, Step, StepCtx, StepOutcome, StepProfile}; + + fn stub_step(kind: StepKind) -> Box { + #[derive(Clone)] + struct StubStep(StepKind); + impl Step for StubStep { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "Stub", + kind: self.0, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + Box::new(TypedStep::new(StubStep(kind))) + } + + #[test] + fn no_exclusives_returns_all_none() { + let steps = vec![stub_step(StepKind::Parallel), stub_step(StepKind::Serial)]; + let owners = assign_exclusive_owners(&steps, 4).unwrap(); + assert_eq!(owners, vec![None, None]); + } + + #[test] + fn one_exclusive_assigned_to_thread_zero() { + let steps = vec![stub_step(StepKind::Exclusive), stub_step(StepKind::Parallel)]; + let owners = assign_exclusive_owners(&steps, 4).unwrap(); + assert_eq!(owners, vec![Some(0), None]); + } + + #[test] + fn two_exclusives_assigned_zero_and_one() { + let steps = vec![ + stub_step(StepKind::Exclusive), + stub_step(StepKind::Parallel), + stub_step(StepKind::Exclusive), + ]; + let owners = assign_exclusive_owners(&steps, 4).unwrap(); + assert_eq!(owners, vec![Some(0), None, Some(1)]); + } + + #[test] + fn too_many_exclusives_returns_not_enough_threads() { + let steps = vec![ + stub_step(StepKind::Exclusive), + stub_step(StepKind::Exclusive), + stub_step(StepKind::Exclusive), + ]; + let result = assign_exclusive_owners(&steps, 2); + assert!(matches!( + result, + Err(PipelineError::NotEnoughThreads { required: 3, available: 2 }) + )); + } + + fn sticky_exclusive_step() -> Box { + #[derive(Clone)] + struct StickyExclusive; + impl Step for StickyExclusive { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "StickyExclusive", + kind: StepKind::Exclusive, + sticky: true, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + Box::new(TypedStep::new(StickyExclusive)) + } + + fn sticky_serial_step(affinity: crate::step::Affinity) -> Box { + struct StickySerial(crate::step::Affinity); + impl Step for StickySerial { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "StickySerial", + kind: StepKind::Serial, + sticky: true, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn affinity(&self) -> crate::step::Affinity { + self.0 + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + Box::new(TypedStep::new(StickySerial(affinity))) + } + + /// Assert that `sticky` holds `Some(StepIdx(0))` in exactly `expected_slot` + /// (when `Some`) and `None` everywhere else across `n_workers` slots. + fn assert_only_slot( + sticky: &[Option], + n_workers: usize, + expected_slot: Option, + ) { + for (slot, &got) in sticky.iter().enumerate().take(n_workers) { + let want = if Some(slot) == expected_slot { Some(StepIdx(0)) } else { None }; + assert_eq!(got, want, "slot {slot}; expected owner slot {expected_slot:?}"); + } + } + + #[rstest] + #[case::in_range(2, Some(2))] + #[case::out_of_range(10, None)] + fn sticky_exclusive_owner_maps_to_slot( + #[case] owner: usize, + #[case] expected_slot: Option, + ) { + // A sticky-exclusive step's slot is its (in-range) owner worker; an + // out-of-range owner is skipped, leaving every slot empty. + let steps = vec![sticky_exclusive_step(), stub_step(StepKind::Parallel)]; + let exclusive_owners = vec![Some(owner), None]; + let sticky = assign_sticky_owners(&steps, &exclusive_owners, 4); + assert_only_slot(&sticky, 4, expected_slot); + } + + #[test] + fn sticky_exclusive_occupied_slot_not_overwritten() { + // Two sticky-exclusive steps competing for the same worker slot. + let steps = vec![sticky_exclusive_step(), sticky_exclusive_step()]; + let exclusive_owners = vec![Some(0_usize), Some(0_usize)]; + let sticky = assign_sticky_owners(&steps, &exclusive_owners, 4); + // First step wins; second is skipped because slot[0] is already occupied. + // The loser must not land in any other slot either. + assert_only_slot(&sticky, 4, Some(0)); + } + + #[rstest] + #[case::reader(Affinity::Reader, Some(0))] + #[case::writer(Affinity::Writer, Some(3))] + #[case::worker_in_range(Affinity::Worker(2), Some(2))] + #[case::worker_out_of_range(Affinity::Worker(10), None)] + #[case::none(Affinity::None, None)] + fn sticky_serial_affinity_maps_to_slot( + #[case] affinity: Affinity, + #[case] expected_slot: Option, + ) { + // Serial affinity resolves to a single worker slot: Reader→0, + // Writer→last, Worker(i)→i; an out-of-range Worker index and None are + // skipped, leaving every slot empty. + let steps = vec![sticky_serial_step(affinity)]; + let exclusive_owners = vec![None]; + let sticky = assign_sticky_owners(&steps, &exclusive_owners, 4); + assert_only_slot(&sticky, 4, expected_slot); + } + + #[test] + fn sticky_exclusive_beats_sticky_serial_on_same_slot() { + // Exclusive pass runs first; serial pass only fills empty slots. + let exc = sticky_exclusive_step(); + let ser = sticky_serial_step(crate::step::Affinity::Reader); // also targets slot 0 + let steps: Vec> = vec![exc, ser]; + let exclusive_owners = vec![Some(0_usize), None]; + let sticky = assign_sticky_owners(&steps, &exclusive_owners, 4); + // Exclusive (step 0) wins slot 0; Serial (step 1) is blocked — and the + // blocked serial step must not land in any other slot either. + assert_only_slot(&sticky, 4, Some(0)); + } +} diff --git a/crates/fgumi-pipeline-core/src/runtime/sampler.rs b/crates/fgumi-pipeline-core/src/runtime/sampler.rs new file mode 100644 index 000000000..63c080c23 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/runtime/sampler.rs @@ -0,0 +1,441 @@ +//! Background occupancy sampler for `--pipeline-trace`. +//! +//! When instrumentation is on, `Pipeline::run` spawns one +//! [`run_occupancy_sampler`] thread (same lifecycle slot as the deadlock +//! monitor / queue rebalancer) that periodically reads each byte-bounded edge's +//! depth (`current_bytes / limit_bytes`) and feeds it to the edge's +//! [`EdgeMetrics`](super::metrics::EdgeMetrics) occupancy histogram. Only edges +//! with a `depth_source` (byte-bounded) are sampled — count/unbounded edges still +//! get their push/pop counters, just no occupancy histogram. +//! +//! Reads are cheap but not free. A direct byte-bounded edge costs one `Relaxed` +//! load per tick. An **ordered** edge additionally reads its `ReorderStage` +//! overflow stash, and `ReorderCapHandle::current_buffer_bytes` takes the stage's +//! `state` mutex — the same one every must-accept push and every +//! `try_pop_in_order` holds. So an ordered edge does briefly touch a worker-hot +//! lock once per tick (at the default interval, negligible against per-item +//! traffic, but not zero). +//! +//! Each tick therefore reads every edge's depth **once**, via `read_depths`, +//! and hands the result to both consumers (`record_depths` and the timeline +//! writer). Letting each consumer read for itself would take that mutex twice per +//! ordered edge per tick and could record two different depths for one tick. + +use std::fmt::Write as _; +use std::io::{BufWriter, Write}; +use std::path::PathBuf; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::{Duration, Instant}; + +use super::contexts::RegisteredEdge; + +/// Default sampling interval. Two milliseconds is cheap (one atomic load per +/// edge) yet fine-grained enough to resolve a chain's phase structure over a +/// multi-second run. +pub const DEFAULT_SAMPLE_INTERVAL: Duration = Duration::from_millis(2); + +/// Poll each byte-bounded edge's occupancy into its histogram until `stop` is +/// set. Edges without a `depth_source` (count/unbounded) are skipped. When +/// `trace_path` is `Some` (the `Timeline` level), also append one TSV row per +/// tick — `t_ms` plus, per edge, its depth fraction and cumulative +/// pushed/popped item counts — so the run's phase structure can be plotted. +pub fn run_occupancy_sampler( + stop: &AtomicBool, + edges: &[RegisteredEdge], + interval: Duration, + trace_path: Option, +) { + let mut trace = trace_path.and_then(|p| TraceWriter::open(&p, edges)); + let start = Instant::now(); + let mut sampled_in_loop = false; + while !stop.load(Ordering::Relaxed) { + // One read per tick, shared by the histogram and the timeline row. + let depths = read_depths(edges); + record_depths(edges, &depths); + if let Some(t) = trace.as_mut() { + t.write_row(edges, &depths, start.elapsed()); + } + sampled_in_loop = true; + std::thread::sleep(interval); + } + // Guard the final sample: only take it when the loop never sampled (a run + // so short `stop` was already set before the first iteration). Sampling + // unconditionally here would add an extra occupancy point + timeline row + // taken AFTER the pipeline already drained, biasing the mean toward the + // empty final state. + if !sampled_in_loop { + let depths = read_depths(edges); + record_depths(edges, &depths); + if let Some(t) = trace.as_mut() { + t.write_row(edges, &depths, start.elapsed()); + } + } + if let Some(mut t) = trace { + t.flush(); + } +} + +/// Bytes buffered in an ordered edge's `ReorderStage` overflow stash (0 for a +/// direct/count/unbounded edge). Added to the transport queue's `current_bytes` +/// when sampling depth so an ordered edge reflects total buffered bytes rather +/// than reading empty while items pile in the reorder buffer awaiting an earlier +/// ordinal. +fn reorder_stash_bytes(edge: &RegisteredEdge) -> u64 { + edge.reorder_depth.as_ref().map_or(0, |r| r.current_buffer_bytes()) +} + +/// One edge's depth for a single tick: `(occupied_bytes, limit_bytes)`, or `None` +/// for a count/unbounded edge (no `depth_source`) or one whose limit reads 0. +type EdgeDepth = Option<(u64, u64)>; + +/// Read every edge's depth once, for one tick. +/// +/// Called once per tick and shared by [`record_depths`] and the timeline writer +/// so that (a) an ordered edge's `ReorderStage` mutex is taken once per tick +/// rather than once per consumer, and (b) the histogram sample and the timeline +/// row for a given tick always report the same number. +fn read_depths(edges: &[RegisteredEdge]) -> Vec { + edges + .iter() + .map(|edge| { + let src = edge.depth_source.as_ref()?; + let limit = src.limit_bytes(); + if limit == 0 { + return None; + } + Some((src.current_bytes().saturating_add(reorder_stash_bytes(edge)), limit)) + }) + .collect() +} + +/// Feed one tick's depths into each edge's occupancy histogram. +fn record_depths(edges: &[RegisteredEdge], depths: &[EdgeDepth]) { + for (edge, depth) in edges.iter().zip(depths) { + if let Some((occupied, limit)) = *depth { + edge.metrics.record_depth(occupied, limit); + } + } +} + +/// Column-name prefix for one edge's timeline columns. Includes the producer +/// step index and output branch so fan-out edges (one producer, several +/// branches) and repeated step names produce distinct, collision-free headers — +/// a bare `producer__consumer` prefix duplicates columns whenever two edges +/// share both names. `(producer_step, branch)` uniquely identifies an edge. +fn edge_column_prefix(e: &RegisteredEdge) -> String { + format!( + "{}__{}#{}b{}", + e.producer_name, + e.consumer_name.unwrap_or("sink"), + e.producer_step.0, + e.branch.0, + ) +} + +/// Per-tick TSV writer for the `Timeline` level. Best-effort: a write error is +/// logged once and further rows are dropped (instrumentation never aborts a run). +struct TraceWriter { + writer: BufWriter, + failed: bool, +} + +impl TraceWriter { + /// Open `path` and write the header (`t_ms` + three columns per edge). + /// Returns `None` (with a warning) if the file can't be created. + fn open(path: &std::path::Path, edges: &[RegisteredEdge]) -> Option { + match std::fs::File::create(path) { + Ok(file) => { + let mut writer = BufWriter::new(file); + let mut header = String::from("t_ms"); + for e in edges { + let edge = edge_column_prefix(e); + let _ = write!(header, "\t{edge}.depth\t{edge}.pushed\t{edge}.popped"); + } + if writeln!(writer, "{header}").is_err() { + log::warn!( + "pipeline-trace: failed to write timeline header to {}", + path.display() + ); + return None; + } + Some(Self { writer, failed: false }) + } + Err(e) => { + log::warn!("pipeline-trace: cannot create timeline file {}: {e}", path.display()); + None + } + } + } + + /// `depths` is this tick's depths from [`read_depths`], shared with + /// [`record_depths`] so the row and the histogram agree and each ordered + /// edge's reorder mutex is taken once per tick. + #[allow(clippy::cast_precision_loss)] + fn write_row(&mut self, edges: &[RegisteredEdge], depths: &[EdgeDepth], elapsed: Duration) { + if self.failed { + return; + } + let mut row = format!("{}", elapsed.as_millis()); + for (e, depth) in edges.iter().zip(depths) { + // Count/unbounded edges (no `depth_source`) are unsampled — emit `NA` + // rather than `0.000`, which would misread as "empty" instead of + // "not measured". Byte-bounded edges report total buffered depth + // (transport + reorder stash) as a fraction of the limit. + let depth = depth.map_or_else( + || "NA".to_string(), + |(occupied, limit)| format!("{:.3}", occupied as f32 / limit as f32), + ); + let ms = e.metrics.snapshot(); + let _ = write!(row, "\t{depth}\t{}\t{}", ms.pushed_items, ms.popped_items); + } + if writeln!(self.writer, "{row}").is_err() { + log::warn!("pipeline-trace: timeline write failed; dropping further rows"); + self.failed = true; + } + } + + fn flush(&mut self) { + if self.failed { + return; + } + // A dropped flush error can silently lose buffered rows after every + // write appeared to succeed — warn and mark the writer failed, matching + // `write_row`'s best-effort error handling. + if self.writer.flush().is_err() { + log::warn!("pipeline-trace: timeline flush failed; buffered rows may be lost"); + self.failed = true; + } + } +} + +/// One sampling sweep over all edges. Extracted so tests can drive a single +/// deterministic tick without the sleep loop. +/// +/// The sampler loop does not call this — it uses `read_depths` once per tick +/// and shares the result with the timeline writer, so both record the same +/// numbers from a single read (see the module doc). +pub fn sample_once(edges: &[RegisteredEdge]) { + record_depths(edges, &read_depths(edges)); +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + use std::sync::atomic::AtomicBool; + + use crate::item::HeapSize; + use crate::queues::{BoundedQueueHandle, ByteBoundedQueue, ItemQueue}; + use crate::runtime::metrics::EdgeMetrics; + use crate::topology::{BranchIdx, StepIdx}; + + #[derive(Debug)] + struct Heavy(Vec); + impl HeapSize for Heavy { + fn heap_size(&self) -> usize { + self.0.len() + } + } + + fn edge_over( + metrics: Arc, + depth_source: Option>, + ) -> RegisteredEdge { + RegisteredEdge { + producer_step: StepIdx(0), + producer_name: "producer", + consumer_step: Some(StepIdx(1)), + consumer_name: Some("consumer"), + branch: BranchIdx(0), + metrics, + depth_source, + reorder_depth: None, + } + } + + #[test] + fn sample_once_records_occupancy_from_depth_source() { + let m = EdgeMetrics::new(); + let q = Arc::new(ByteBoundedQueue::::new(1000)); + q.try_push(Heavy(vec![0; 500])).unwrap(); // 50% of the 1000-byte budget + let edge = edge_over(Arc::clone(&m), Some(Arc::clone(&q) as Arc)); + for _ in 0..10 { + sample_once(std::slice::from_ref(&edge)); + } + let s = m.snapshot(); + assert_eq!(s.depth_samples, 10); + assert!((s.mean_occupancy - 0.5).abs() < 0.05, "mean ≈ 0.5, got {}", s.mean_occupancy); + } + + #[test] + fn count_edge_without_depth_source_is_skipped() { + // An edge with no depth_source (count/unbounded) records no occupancy. + let m = EdgeMetrics::new(); + let edge = edge_over(Arc::clone(&m), None); + for _ in 0..5 { + sample_once(std::slice::from_ref(&edge)); + } + assert_eq!(m.snapshot().depth_samples, 0, "no depth source → no occupancy samples"); + } + + #[test] + fn run_occupancy_sampler_stops_and_records() { + let m = EdgeMetrics::new(); + let q = Arc::new(ByteBoundedQueue::::new(1000)); + q.try_push(Heavy(vec![0; 800])).unwrap(); + let edges = + vec![edge_over(Arc::clone(&m), Some(Arc::clone(&q) as Arc))]; + let stop = Arc::new(AtomicBool::new(false)); + let stop_c = Arc::clone(&stop); + let handle = std::thread::spawn(move || { + run_occupancy_sampler(&stop_c, &edges, Duration::from_millis(1), None); + }); + std::thread::sleep(Duration::from_millis(30)); + stop.store(true, Ordering::Relaxed); + handle.join().unwrap(); + let s = m.snapshot(); + assert!(s.depth_samples > 0, "sampler recorded at least one tick"); + assert!((s.mean_occupancy - 0.8).abs() < 0.1, "mean ≈ 0.8, got {}", s.mean_occupancy); + } + + #[test] + fn edge_columns_are_unique_when_step_names_collide() { + // Regression: two edges that share producer AND consumer names (fan-out, + // or duplicate step names) must still produce distinct TSV columns — a + // bare `producer__consumer` prefix would emit duplicate column headers. + let m0 = EdgeMetrics::new(); + let m1 = EdgeMetrics::new(); + let e0 = RegisteredEdge { + producer_step: StepIdx(0), + producer_name: "dup", + consumer_step: Some(StepIdx(1)), + consumer_name: Some("sink"), + branch: BranchIdx(0), + metrics: m0, + depth_source: None, + reorder_depth: None, + }; + // Same names, different (producer_step, branch): a fan-out sibling. + let e1 = RegisteredEdge { + producer_step: StepIdx(0), + producer_name: "dup", + consumer_step: Some(StepIdx(2)), + consumer_name: Some("sink"), + branch: BranchIdx(1), + metrics: m1, + depth_source: None, + reorder_depth: None, + }; + let p0 = edge_column_prefix(&e0); + let p1 = edge_column_prefix(&e1); + assert_ne!(p0, p1, "colliding names must yield distinct column prefixes"); + assert_eq!(p0, "dup__sink#0b0"); + assert_eq!(p1, "dup__sink#0b1"); + } + + #[test] + fn ordered_edge_occupancy_includes_reorder_stash() { + use crate::queues::CountBoundedQueue; + use crate::reorder::{ReorderCapHandle, ReorderStage, Sequenced}; + // The transport (occupancy depth source) is empty, but the reorder stash + // holds 400 buffered bytes waiting for an earlier ordinal. Sampled + // occupancy must reflect the stash (400/1000 = 0.4), not read empty — + // otherwise a producer-skewed ordered edge looks idle while backed up. + let m = EdgeMetrics::new(); + let transport = Arc::new(ByteBoundedQueue::::new(1000)); + + let reorder_transport: Arc>> = + Arc::new(CountBoundedQueue::>::new(8)); + let stage = Arc::new(ReorderStage::new(reorder_transport)); + stage.try_push(1, Heavy(vec![0; 400])).unwrap(); // out-of-order → stashed + assert!(stage.try_pop_in_order().is_none(), "ordinal 0 absent → nothing pops"); + assert!(stage.current_buffer_bytes() >= 400, "stash holds the buffered bytes"); + + let edge = RegisteredEdge { + producer_step: StepIdx(0), + producer_name: "p", + consumer_step: Some(StepIdx(1)), + consumer_name: Some("c"), + branch: BranchIdx(0), + metrics: Arc::clone(&m), + depth_source: Some(Arc::clone(&transport) as Arc), + reorder_depth: Some(Arc::clone(&stage) as Arc), + }; + for _ in 0..10 { + sample_once(std::slice::from_ref(&edge)); + } + let s = m.snapshot(); + assert_eq!(s.depth_samples, 10); + assert!( + (s.mean_occupancy - 0.4).abs() < 0.05, + "occupancy reflects the reorder stash, got {}", + s.mean_occupancy + ); + assert!( + (s.mean_occupancy_bytes - 400.0).abs() < 1.0, + "byte mean equals the stashed bytes, got {}", + s.mean_occupancy_bytes + ); + } + + #[test] + fn timeline_tsv_has_header_and_rows() { + let m = EdgeMetrics::new(); + let q = Arc::new(ByteBoundedQueue::::new(1000)); + q.try_push(Heavy(vec![0; 400])).unwrap(); + let edges = + vec![edge_over(Arc::clone(&m), Some(Arc::clone(&q) as Arc))]; + let dir = std::env::temp_dir(); + let path = dir.join(format!("fgumi-trace-test-{}.tsv", std::process::id())); + let stop = Arc::new(AtomicBool::new(false)); + let stop_c = Arc::clone(&stop); + let path_c = path.clone(); + let handle = std::thread::spawn(move || { + run_occupancy_sampler(&stop_c, &edges, Duration::from_millis(2), Some(path_c)); + }); + std::thread::sleep(Duration::from_millis(30)); + stop.store(true, Ordering::Relaxed); + handle.join().unwrap(); + + let content = std::fs::read_to_string(&path).expect("trace file written"); + let _ = std::fs::remove_file(&path); + let mut lines = content.lines(); + let header = lines.next().expect("header row"); + assert!(header.starts_with("t_ms"), "header begins with t_ms"); + assert!(header.contains("producer__consumer#0b0.depth"), "per-edge depth column"); + let rows: Vec<&str> = lines.collect(); + assert!(!rows.is_empty(), "at least one data row"); + // First field of a data row is a monotonic t_ms integer. + let first_t: u128 = rows[0].split('\t').next().unwrap().parse().expect("t_ms is an int"); + let last_t: u128 = rows.last().unwrap().split('\t').next().unwrap().parse().unwrap(); + assert!(last_t >= first_t, "t_ms is monotonic"); + } + + #[test] + fn timeline_tsv_marks_unsampled_edge_na() { + // A count/unbounded edge (no depth_source) is unsampled: its depth column + // must read `NA`, not `0.000` (which would misread as an empty byte edge). + let m = EdgeMetrics::new(); + let edges = vec![edge_over(Arc::clone(&m), None)]; + let dir = std::env::temp_dir(); + let path = dir.join(format!("fgumi-trace-na-{}.tsv", std::process::id())); + let stop = Arc::new(AtomicBool::new(false)); + let stop_c = Arc::clone(&stop); + let path_c = path.clone(); + let handle = std::thread::spawn(move || { + run_occupancy_sampler(&stop_c, &edges, Duration::from_millis(2), Some(path_c)); + }); + std::thread::sleep(Duration::from_millis(30)); + stop.store(true, Ordering::Relaxed); + handle.join().unwrap(); + + let content = std::fs::read_to_string(&path).expect("trace file written"); + let _ = std::fs::remove_file(&path); + let mut lines = content.lines(); + let _header = lines.next().expect("header row"); + let row = lines.next().expect("at least one data row"); + // Columns: t_ms, .depth, .pushed, .popped. + let depth = row.split('\t').nth(1).expect("depth column"); + assert_eq!(depth, "NA", "unsampled edge's depth column is NA, row: {row}"); + } +} diff --git a/crates/fgumi-pipeline-core/src/runtime/scheduler.rs b/crates/fgumi-pipeline-core/src/runtime/scheduler.rs new file mode 100644 index 000000000..0fde6eb00 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/runtime/scheduler.rs @@ -0,0 +1,106 @@ +//! Pluggable per-worker dispatch-order policy for the round-robin pool driver. +//! +//! The worker loop ([`run_worker_loop`](crate::runtime::run_worker_loop)) walks +//! each worker's *live* steps once per pass and runs the first that makes +//! progress. A [`Scheduler`] decides the ORDER of that walk — the only thing it +//! controls; it never changes which steps exist, the sticky source/sink +//! fast-path, or the Serial/Exclusive contention rules. +//! +//! Two policies ship: +//! +//! - [`ChainOrderScheduler`] (the default) — walk **upstream-first** (chain +//! order). A worker attempts the earliest-in-chain step with work, favouring +//! production. This is the historical behaviour; every command keeps it unless +//! it opts out, so existing pipelines are byte-for-byte unaffected. +//! - [`DrainFirstScheduler`] — walk **downstream-first** (reverse chain order). +//! A worker attempts the deepest step with work first, favouring *draining* +//! buffered work before producing more. Combined with skip-on-Serial-contention +//! this self-balances: for a Serial step fed by an N-way Parallel producer, one +//! worker grabs the Serial drain (mutex) while the rest find it contended, skip, +//! and fall through to the producer — so the drain overlaps production instead +//! of starving behind it on the shared pool. A sticky step is **not** exempt +//! from the walk: it takes its bounded burst on the sticky fast-path first and +//! is then still visited in the walk itself (under `Reverse`, last rather than +//! first). Only its `Progress` priority restart is suppressed, so the walk +//! continues past it to the steps that drain its output — see +//! `super::driver::round_robin_dispatch`. +//! +//! This mirrors main's `Scheduler`-trait design (`BalancedChaseDrainScheduler` +//! among others), but generically over an arbitrary step list rather than a +//! fixed set of named BAM stages: the only lever exposed here is walk direction, +//! which is all the generic driver needs to express drain-first scheduling. + +/// The order in which a worker attempts its live steps in one round-robin pass. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WalkDirection { + /// Chain order (upstream → downstream): favour production. + Forward, + /// Reverse chain order (downstream → upstream): favour draining. + Reverse, +} + +/// A per-worker dispatch-order policy. Selected per pipeline via +/// [`PipelineConfig::with_scheduler`](crate::builder::PipelineConfig::with_scheduler) +/// and shared across all workers (the shipped policies are stateless). +pub trait Scheduler: Send + Sync + std::fmt::Debug { + /// Direction to walk this worker's live steps this pass. Called once per + /// round-robin pass, so it must be cheap. + fn walk(&self) -> WalkDirection; + + /// Human-readable name for diagnostics / `--pipeline-stats`. + fn name(&self) -> &'static str; +} + +/// Default upstream-first (chain-order) walk. Preserves the historical dispatch +/// behaviour for every pipeline that does not opt into a different policy. +#[derive(Debug, Default, Clone, Copy)] +pub struct ChainOrderScheduler; + +impl Scheduler for ChainOrderScheduler { + #[inline] + fn walk(&self) -> WalkDirection { + WalkDirection::Forward + } + fn name(&self) -> &'static str { + "chain-order" + } +} + +/// Downstream-first (reverse chain-order) walk — drain buffered work before +/// producing more. Opt-in per pipeline (e.g. the sort chain, to overlap the +/// serial boundary/key scan with the parallel inflate instead of starving it on +/// the shared pool). +#[derive(Debug, Default, Clone, Copy)] +pub struct DrainFirstScheduler; + +impl Scheduler for DrainFirstScheduler { + #[inline] + fn walk(&self) -> WalkDirection { + WalkDirection::Reverse + } + fn name(&self) -> &'static str { + "drain-first" + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn default_is_forward() { + assert_eq!(ChainOrderScheduler.walk(), WalkDirection::Forward); + assert_eq!(ChainOrderScheduler.name(), "chain-order"); + } + + #[test] + fn drain_first_is_reverse() { + assert_eq!(DrainFirstScheduler.walk(), WalkDirection::Reverse); + assert_eq!(DrainFirstScheduler.name(), "drain-first"); + } + + // The forward/reverse position→step arithmetic these policies drive is + // exercised end-to-end against the real dispatcher in driver.rs + // (`driver_round_robins_all_live_before_parking`); a standalone test here + // would only re-derive the formula, not the driver's actual mapping. +} diff --git a/crates/fgumi-pipeline-core/src/runtime/stats.rs b/crates/fgumi-pipeline-core/src/runtime/stats.rs new file mode 100644 index 000000000..b04d57890 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/runtime/stats.rs @@ -0,0 +1,1392 @@ +//! Per-step pipeline statistics: dispatch counts, outcome breakdown, and +//! cumulative `try_run` time. Indexed by `StepIdx` so it works for any +//! chain shape. +//! +//! Stats are opt-in. When `PipelineConfig::stats` is `Some(arc)`, the +//! worker loop times each `dispatch_one_step` call and records the +//! outcome. When `None`, the loop pays no extra cost. +//! +//! ## What's recorded +//! +//! Per step: +//! - `try_run_total` — total `try_run_erased` dispatches that returned +//! a result (excluding `Skip` entries). +//! - `progress_count` — `StepOutcome::Progress`. +//! - `no_progress_count` — `StepOutcome::NoProgress`. +//! - `contention_count` — `StepOutcome::Contention` (Serial step mutex +//! held by another worker; or skipped via `try_lock`). Always 0 under the +//! fused single-thread driver, which holds no mutex and never contends. +//! - `finished_count` — `StepOutcome::Finished` (any step on end-of-stream: +//! source, mid, or sink all record `Finished` once their inputs drain). +//! - `error_count` — `try_run_erased` returned `Err`. +//! - `total_run_ns` — cumulative wall time across all dispatches. +//! +//! Per edge, when instrumentation is on: +//! - Queue-depth samples. [`crate::runtime::sampler::sample_once`] reads each +//! registered edge's occupancy and feeds +//! [`EdgeMetrics::record_depth`](crate::runtime::metrics::EdgeMetrics::record_depth); +//! [`PipelineStats::snapshot_with_edges`] turns the histogram into the +//! per-edge occupancy on [`EdgeStatsSnapshot`]. Plain +//! [`snapshot`](PipelineStats::snapshot) leaves `edges` empty. +//! +//! ## What's *not* recorded yet +//! +//! - Per-thread step counts. Legacy carries per-`(thread, step)` counters +//! for bottleneck attribution. Useful for the rebalancer (#17) port, +//! not for first-cut observability. Deferred until the rebalancer +//! work lands. + +use std::fmt; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Instant; + +use crate::step::StepOutcome; +use crate::topology::StepIdx; + +/// Atomic counters for a single step. +#[derive(Debug)] +pub struct StepStats { + /// Total `try_run` dispatches (every outcome, including errors). + pub try_run_total: AtomicU64, + /// Dispatches that returned `StepOutcome::Progress`. + pub progress_count: AtomicU64, + /// Dispatches that returned `StepOutcome::NoProgress`. + pub no_progress_count: AtomicU64, + /// Dispatches that returned `StepOutcome::Contention`. + pub contention_count: AtomicU64, + /// Dispatches that returned `StepOutcome::Finished`. + pub finished_count: AtomicU64, + /// Dispatches that returned an error from `try_run_erased`. + pub error_count: AtomicU64, + /// Cumulative wall-ns spent inside `try_run` across all dispatches. + pub total_run_ns: AtomicU64, + /// Wall ns from pipeline start to the start of this step's FIRST + /// Progress dispatch. `u64::MAX` until set. Lets us see when each + /// step actually first did useful work. + pub first_progress_ns: AtomicU64, + /// Wall ns from pipeline start to the END of this step's LAST + /// Progress dispatch. Updated on every Progress (monotonic max). + /// `0` if the step never made progress. + pub last_progress_ns: AtomicU64, +} + +impl Default for StepStats { + fn default() -> Self { + Self { + try_run_total: AtomicU64::new(0), + progress_count: AtomicU64::new(0), + no_progress_count: AtomicU64::new(0), + contention_count: AtomicU64::new(0), + finished_count: AtomicU64::new(0), + error_count: AtomicU64::new(0), + total_run_ns: AtomicU64::new(0), + first_progress_ns: AtomicU64::new(u64::MAX), + last_progress_ns: AtomicU64::new(0), + } + } +} + +impl StepStats { + fn snapshot(&self) -> StepStatsSnapshot { + StepStatsSnapshot { + try_run_total: self.try_run_total.load(Ordering::Relaxed), + progress_count: self.progress_count.load(Ordering::Relaxed), + no_progress_count: self.no_progress_count.load(Ordering::Relaxed), + contention_count: self.contention_count.load(Ordering::Relaxed), + finished_count: self.finished_count.load(Ordering::Relaxed), + error_count: self.error_count.load(Ordering::Relaxed), + total_run_ns: self.total_run_ns.load(Ordering::Relaxed), + first_progress_ns: self.first_progress_ns.load(Ordering::Relaxed), + last_progress_ns: self.last_progress_ns.load(Ordering::Relaxed), + } + } +} + +/// Upper bound on per-worker utilization slots. Worker `thread_id`s index the +/// `worker_busy_ns` / `worker_idle_ns` arrays; ids at or above this are not +/// tracked (a no-op, never a panic). Sized well above any realistic +/// `--threads`, so a fixed array avoids threading `num_workers` through every +/// `PipelineStats::new` call site. +const MAX_TRACKED_WORKERS: usize = 512; + +/// Per-step counter container. Sized to match the pipeline's chain length; +/// callers obtain one via `Pipeline::stats()`. +#[derive(Debug)] +pub struct PipelineStats { + steps: Box<[StepStats]>, + step_names: Box<[&'static str]>, + /// Per-worker wall-ns spent dispatching (the sticky + round-robin work + /// section of `run_worker_loop`), indexed by `WorkerCore::thread_id`. + worker_busy_ns: Box<[AtomicU64]>, + /// Per-worker wall-ns spent in the no-progress backoff sleep (idle/blocked + /// waiting for upstream work), indexed by `thread_id`. Together with + /// `worker_busy_ns` this answers "are workers utilised, or blocked". + worker_idle_ns: Box<[AtomicU64]>, + /// Per-step wall-ns a `StepKind::Detached` step's dedicated thread spent + /// inside `try_run` (busy) and parked on its backoff (idle), indexed by + /// `step_idx`. Detached threads have no `WorkerCore::thread_id`, so they are + /// tracked here separately and are intentionally EXCLUDED from the pool + /// utilisation line (which only sums `worker_busy_ns` / `worker_idle_ns`). + /// Reported on their own line so the legacy "N + 2" split is visible without + /// diluting the N-worker pool%. + detached_busy_ns: Box<[AtomicU64]>, + detached_idle_ns: Box<[AtomicU64]>, + /// Per-step count of backoff-park events on a `StepKind::Detached` step's + /// dedicated thread (one per [`backoff_park`](crate::runtime::detached) call), + /// indexed by `step_idx`. With `detached_idle_ns` this gives the average park + /// duration and the park-to-progress ratio — the signal for whether the + /// backoff (vs a precise per-slot condvar) adds latency on the merge's path. + detached_park_events: Box<[AtomicU64]>, + /// Anchor for first/last-progress timestamps. Set at `PipelineStats` + /// construction; all `first_progress_ns` / `last_progress_ns` values + /// are wall-ns elapsed from this `Instant`. + pipeline_start: Instant, +} + +impl PipelineStats { + /// Construct a stats container sized to `step_names.len()`. Each step's + /// counters start at zero. Wrap in `Arc` to share across worker threads. + #[must_use] + pub fn new(step_names: Vec<&'static str>) -> Self { + let steps = (0..step_names.len()).map(|_| StepStats::default()).collect::>(); + let worker_busy_ns = + (0..MAX_TRACKED_WORKERS).map(|_| AtomicU64::new(0)).collect::>(); + let worker_idle_ns = + (0..MAX_TRACKED_WORKERS).map(|_| AtomicU64::new(0)).collect::>(); + let n_steps = steps.len(); + let detached_busy_ns = (0..n_steps).map(|_| AtomicU64::new(0)).collect::>(); + let detached_idle_ns = (0..n_steps).map(|_| AtomicU64::new(0)).collect::>(); + let detached_park_events = (0..n_steps).map(|_| AtomicU64::new(0)).collect::>(); + Self { + steps: steps.into_boxed_slice(), + step_names: step_names.into_boxed_slice(), + worker_busy_ns: worker_busy_ns.into_boxed_slice(), + worker_idle_ns: worker_idle_ns.into_boxed_slice(), + detached_busy_ns: detached_busy_ns.into_boxed_slice(), + detached_idle_ns: detached_idle_ns.into_boxed_slice(), + detached_park_events: detached_park_events.into_boxed_slice(), + pipeline_start: Instant::now(), + } + } + + /// Accumulate `ns` of `try_run` (busy) time for a `StepKind::Detached` + /// step's dedicated thread, keyed by `step_idx`. Excluded from pool + /// utilisation; reported on the Detached line. Out-of-range steps are a + /// silent no-op. + #[inline] + pub fn record_detached_busy(&self, step: StepIdx, ns: u64) { + if let Some(c) = self.detached_busy_ns.get(step.0) { + c.fetch_add(ns, Ordering::Relaxed); + } + } + + /// Accumulate `ns` of backoff-park (idle) time for a `StepKind::Detached` + /// step's dedicated thread, keyed by `step_idx`. Excluded from pool + /// utilisation; reported on the Detached line. + #[inline] + pub fn record_detached_idle(&self, step: StepIdx, ns: u64) { + if let Some(c) = self.detached_idle_ns.get(step.0) { + c.fetch_add(ns, Ordering::Relaxed); + } + } + + /// Increment the backoff-park event count for a `StepKind::Detached` step's + /// dedicated thread (called once per park). Out-of-range steps are a silent + /// no-op. + #[inline] + pub fn record_detached_park(&self, step: StepIdx) { + if let Some(c) = self.detached_park_events.get(step.0) { + c.fetch_add(1, Ordering::Relaxed); + } + } + + /// Accumulate `ns` of dispatch (busy) time for `worker` (its `thread_id`). + /// Ids `>= MAX_TRACKED_WORKERS` are silently dropped. + #[inline] + pub fn record_worker_busy(&self, worker: usize, ns: u64) { + if let Some(c) = self.worker_busy_ns.get(worker) { + c.fetch_add(ns, Ordering::Relaxed); + } + } + + /// Accumulate `ns` of backoff-sleep (idle) time for `worker`. + #[inline] + pub fn record_worker_idle(&self, worker: usize, ns: u64) { + if let Some(c) = self.worker_idle_ns.get(worker) { + c.fetch_add(ns, Ordering::Relaxed); + } + } + + /// Wall ns elapsed since pipeline start. Used by the driver to stamp + /// per-step first/last progress timestamps. + #[must_use] + pub fn elapsed_ns(&self) -> u64 { + u64::try_from(self.pipeline_start.elapsed().as_nanos()).unwrap_or(u64::MAX) + } + + #[must_use] + pub fn n_steps(&self) -> usize { + self.steps.len() + } + + #[must_use] + pub fn step_name(&self, step: StepIdx) -> &'static str { + self.step_names[step.0] + } + + /// Record a successful `try_run_erased` outcome for the given step. + /// Hot path: relaxed atomics, no allocation, no locking. `start_ns` + /// is wall-ns at dispatch start (relative to `pipeline_start`); + /// `elapsed_ns` is the dispatch duration. On `Progress` we stamp + /// the step's first/last active timestamps. + /// + /// # Panics + /// + /// Panics if `step` is outside `self.steps` — it indexes directly rather than + /// using `.get()`. That is deliberate and asymmetric with the detached + /// recorders (`record_detached_busy` and friends), which silently no-op on an + /// out-of-range index: an out-of-range `StepIdx` reaching *this* path is a + /// driver bug that would otherwise silently lose every sample for the step, so + /// it should be loud. Do not assume the detached recorders' no-op contract + /// here. + #[inline] + pub fn record(&self, step: StepIdx, outcome: StepOutcome, start_ns: u64, elapsed_ns: u64) { + let s = &self.steps[step.0]; + s.try_run_total.fetch_add(1, Ordering::Relaxed); + s.total_run_ns.fetch_add(elapsed_ns, Ordering::Relaxed); + match outcome { + StepOutcome::Progress => { + s.progress_count.fetch_add(1, Ordering::Relaxed); + // CAS-min first_progress_ns (initially u64::MAX). + let mut cur = s.first_progress_ns.load(Ordering::Relaxed); + while start_ns < cur { + match s.first_progress_ns.compare_exchange_weak( + cur, + start_ns, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => break, + Err(seen) => cur = seen, + } + } + // last_progress_ns = max(last, start + elapsed). + let end_ns = start_ns.saturating_add(elapsed_ns); + let mut cur = s.last_progress_ns.load(Ordering::Relaxed); + while end_ns > cur { + match s.last_progress_ns.compare_exchange_weak( + cur, + end_ns, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => break, + Err(seen) => cur = seen, + } + } + } + StepOutcome::NoProgress => { + s.no_progress_count.fetch_add(1, Ordering::Relaxed); + } + StepOutcome::Contention => { + s.contention_count.fetch_add(1, Ordering::Relaxed); + } + StepOutcome::Finished => { + s.finished_count.fetch_add(1, Ordering::Relaxed); + } + } + } + + /// Record an error returned by `try_run_erased`. Counted toward + /// `try_run_total` and `total_run_ns`; outcome buckets are not bumped. + /// + /// # Panics + /// + /// Panics if `step` is outside `self.steps` — it indexes directly rather than + /// using `.get()`. That is deliberate and asymmetric with the detached + /// recorders (`record_detached_busy` and friends), which silently no-op on an + /// out-of-range index: an out-of-range `StepIdx` reaching *this* path is a + /// driver bug that would otherwise silently lose every sample for the step, so + /// it should be loud. Do not assume the detached recorders' no-op contract + /// here. + #[inline] + pub fn record_error(&self, step: StepIdx, _start_ns: u64, elapsed_ns: u64) { + let s = &self.steps[step.0]; + s.try_run_total.fetch_add(1, Ordering::Relaxed); + s.error_count.fetch_add(1, Ordering::Relaxed); + s.total_run_ns.fetch_add(elapsed_ns, Ordering::Relaxed); + } + + /// Snapshot all per-step counters into an owned, lock-free struct + /// suitable for printing or further analysis. + #[must_use] + pub fn snapshot(&self) -> StatsSnapshot { + let steps = self + .steps + .iter() + .zip(self.step_names.iter()) + .map(|(stats, &name)| (name, stats.snapshot())) + .collect(); + // Only workers that recorded any activity (busy or idle) — the array is + // sized to MAX_TRACKED_WORKERS but typically few slots are live. + let workers = (0..self.worker_busy_ns.len()) + .map(|w| { + ( + w, + self.worker_busy_ns[w].load(Ordering::Relaxed), + self.worker_idle_ns[w].load(Ordering::Relaxed), + ) + }) + .filter(|(_, busy, idle)| *busy != 0 || *idle != 0) + .collect(); + StatsSnapshot { steps, workers, detached: self.detached_snapshot(), edges: Vec::new() } + } + + /// Collect `(step, step_name, busy_ns, idle_ns, park_events)` for every step + /// whose Detached thread recorded any activity. Shared by both snapshot + /// builders so the Detached line renders identically with or without + /// `--pipeline-trace`. The leading `StepIdx` lets consumers match entries by + /// identity instead of by (possibly duplicated) name. + fn detached_snapshot(&self) -> Vec<(usize, &'static str, u64, u64, u64)> { + (0..self.detached_busy_ns.len()) + .map(|s| { + ( + s, + self.step_names[s], + self.detached_busy_ns[s].load(Ordering::Relaxed), + self.detached_idle_ns[s].load(Ordering::Relaxed), + self.detached_park_events[s].load(Ordering::Relaxed), + ) + }) + .filter(|(_, _, busy, idle, _)| *busy != 0 || *idle != 0) + .collect() + } + + /// Like [`snapshot`](Self::snapshot) but also derives per-edge throughput / + /// occupancy / latency from the chain's instrumented `edges` over a + /// `wall_ns` run. Used by the `--pipeline-trace` end-of-run report. + #[must_use] + pub fn snapshot_with_edges( + &self, + edges: &[crate::runtime::contexts::RegisteredEdge], + wall_ns: u64, + ) -> StatsSnapshot { + let mut snap = self.snapshot(); + snap.edges = edges + .iter() + .map(|e| { + let ms = e.metrics.snapshot(); + let limit_bytes = e.depth_source.as_ref().map(|s| s.limit_bytes()); + compute_edge_stats( + e.producer_name, + e.consumer_name, + e.producer_step.0, + e.consumer_step.map(|s| s.0), + &ms, + limit_bytes, + wall_ns, + ) + }) + .collect(); + snap + } +} + +/// Plain (non-atomic) snapshot of `PipelineStats` at a moment in time. +#[derive(Debug, Clone)] +pub struct StatsSnapshot { + pub steps: Vec<(&'static str, StepStatsSnapshot)>, + /// `(thread_id, busy_ns, idle_ns)` for each worker that did anything. + pub workers: Vec<(usize, u64, u64)>, + /// `(step, step_name, busy_ns, idle_ns, park_events)` for each + /// `StepKind::Detached` step's dedicated thread that did anything. Tracked + /// separately from `workers` and EXCLUDED from the pool utilisation line + /// (legacy "N + 2" — the merge / writer threads are not pool workers); + /// rendered on their own line. `park_events` is the backoff-park count + /// (latency signal). + /// + /// `step` is the entry's `StepIdx`, carried so consumers can match a + /// Detached step by identity rather than by name — two steps may share a + /// name, and `bottleneck_verdict`'s SPIN exemption must not leak from one to + /// the other. Entries are filtered, so the indices are not contiguous. + pub detached: Vec<(usize, &'static str, u64, u64, u64)>, + /// Per-edge throughput / occupancy / latency. Empty unless the snapshot was + /// built via [`PipelineStats::snapshot_with_edges`] (i.e. instrumentation on). + pub edges: Vec, +} + +/// Occupancy classification refined from the histogram's [`RawOccupancy`](crate::runtime::metrics::RawOccupancy) plus +/// the reject/empty rates: `MostlyEmpty` + high empty-rate → `Starved`, +/// `MostlyFull` + high reject-rate → `Backpressured`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OccupancyClass { + /// No occupancy samples (count/unbounded edge, or run too short). + Unknown, + /// Mostly empty, low empty-rate — idle, not starved. + Empty, + /// Mostly empty AND consumer frequently found it empty — producer-starved. + Starved, + /// Spread across the middle — neither side bound. + Healthy, + /// Mostly full AND producer frequently rejected — consumer-backpressured. + Backpressured, + /// Mostly full, low reject-rate — full but not actively rejecting. + Full, + /// Oscillates empty↔full (bursty / batch coupling). + Bimodal, +} + +/// Rate above which a `MostlyEmpty`/`MostlyFull` edge is reclassified as +/// `Starved`/`Backpressured`. 10% of attempts hitting the wall is a clear signal. +const RATE_REFINE_THRESHOLD: f64 = 0.10; + +/// Per-edge derived statistics for the end-of-run report. +#[derive(Debug, Clone)] +pub struct EdgeStatsSnapshot { + pub producer: &'static str, + pub consumer: Option<&'static str>, + /// Producer step index. The bottleneck verdict attributes edges to steps by + /// this identity (not by `producer`/`consumer` name), so fan-out, fan-in, or + /// duplicate step names are not misattributed to whichever edge matches a + /// name first. + pub producer_step: usize, + /// Consumer step index, `None` for a terminal edge with no consumer. + pub consumer_step: Option, + pub items_per_s: f64, + /// Consumer throughput in MiB/s (popped bytes / wall time / 2^20). + pub mibytes_per_s: f64, + pub class: OccupancyClass, + pub mean_occupancy: f32, + /// Fraction of push attempts rejected (backpressure signal). + pub reject_rate: f64, + /// Fraction of pop attempts that found the edge empty (starvation signal). + pub empty_rate: f64, + /// Derived residence time (Little's Law `L/X`), byte-bounded edges only; + /// `None` for count/unbounded edges (no occupancy) or zero throughput. + pub derived_latency_ms: Option, +} + +/// Refine the histogram's raw occupancy class using the reject/empty rates. +#[must_use] +fn refine_occupancy( + raw: crate::runtime::metrics::RawOccupancy, + reject_rate: f64, + empty_rate: f64, +) -> OccupancyClass { + use crate::runtime::metrics::RawOccupancy; + match raw { + RawOccupancy::Unknown => OccupancyClass::Unknown, + RawOccupancy::Healthy => OccupancyClass::Healthy, + RawOccupancy::Bimodal => OccupancyClass::Bimodal, + RawOccupancy::MostlyEmpty if empty_rate >= RATE_REFINE_THRESHOLD => OccupancyClass::Starved, + RawOccupancy::MostlyEmpty => OccupancyClass::Empty, + RawOccupancy::MostlyFull if reject_rate >= RATE_REFINE_THRESHOLD => { + OccupancyClass::Backpressured + } + RawOccupancy::MostlyFull => OccupancyClass::Full, + } +} + +/// Little's-Law residence time (`L/X`), byte-bounded edges only. Mean occupancy +/// in ITEMS = `mean_occupancy_bytes / mean_item_bytes`. +/// +/// `mean_occupancy_bytes` is sampled directly (absolute bytes at each tick), NOT +/// reconstructed as `fraction × limit`. That matters because `queue_memory_total` +/// can rebalance an edge's byte limit at runtime: the limit at teardown may +/// differ from the limits in force during sampling, so `fraction × final_limit` +/// would yield a materially wrong byte figure — and hence a wrong latency. +/// `limit_bytes` is retained only as the byte-bounded gate (`None` for +/// count/unbounded edges, which have no occupancy and thus no derived latency). +#[must_use] +#[allow(clippy::cast_precision_loss)] +fn derived_latency_ms( + ms: &crate::runtime::metrics::EdgeMetricsSnapshot, + limit_bytes: Option, + items_per_s: f64, +) -> Option { + if limit_bytes.is_none() || ms.pushed_items == 0 || items_per_s <= 0.0 { + return None; + } + let mean_item_bytes = ms.pushed_bytes as f64 / ms.pushed_items as f64; + if mean_item_bytes <= 0.0 { + return None; + } + let mean_items = ms.mean_occupancy_bytes / mean_item_bytes; + Some(1000.0 * mean_items / items_per_s) +} + +/// Compute one edge's derived stats from its metrics snapshot + (for a +/// byte-bounded edge) its byte budget, over a `wall_ns` run. Pure — unit-tested +/// directly. `limit_bytes` is `None` for count/unbounded edges (no occupancy → +/// no derived latency). +#[must_use] +#[allow(clippy::cast_precision_loss)] +pub(crate) fn compute_edge_stats( + producer: &'static str, + consumer: Option<&'static str>, + producer_step: usize, + consumer_step: Option, + ms: &crate::runtime::metrics::EdgeMetricsSnapshot, + limit_bytes: Option, + wall_ns: u64, +) -> EdgeStatsSnapshot { + let wall_secs = (wall_ns as f64) / 1e9; + let items_per_s = if wall_secs > 0.0 { ms.popped_items as f64 / wall_secs } else { 0.0 }; + // Divisor is 1 MiB (2^20), so the field/column is MiB/s, not MB/s. + let mibytes_per_s = + if wall_secs > 0.0 { (ms.popped_bytes as f64 / wall_secs) / 1_048_576.0 } else { 0.0 }; + let push_attempts = ms.pushed_items + ms.push_rejections; + let reject_rate = + if push_attempts > 0 { ms.push_rejections as f64 / push_attempts as f64 } else { 0.0 }; + let pop_attempts = ms.popped_items + ms.pop_empties; + let empty_rate = + if pop_attempts > 0 { ms.pop_empties as f64 / pop_attempts as f64 } else { 0.0 }; + + EdgeStatsSnapshot { + producer, + consumer, + producer_step, + consumer_step, + items_per_s, + mibytes_per_s, + class: refine_occupancy(ms.raw_occupancy, reject_rate, empty_rate), + mean_occupancy: ms.mean_occupancy, + reject_rate, + empty_rate, + derived_latency_ms: derived_latency_ms(ms, limit_bytes, items_per_s), + } +} + +#[derive(Debug, Clone, Copy)] +pub struct StepStatsSnapshot { + /// Total `try_run` dispatches (every outcome, including errors). + pub try_run_total: u64, + /// Dispatches that returned `StepOutcome::Progress`. + pub progress_count: u64, + /// Dispatches that returned `StepOutcome::NoProgress`. + pub no_progress_count: u64, + /// Dispatches that returned `StepOutcome::Contention`. + pub contention_count: u64, + /// Dispatches that returned `StepOutcome::Finished`. + pub finished_count: u64, + /// Dispatches that returned an error from `try_run_erased`. + pub error_count: u64, + /// Cumulative wall-ns spent inside `try_run` across all dispatches. + pub total_run_ns: u64, + /// Wall ns (from pipeline start) of this step's first Progress + /// dispatch start. `u64::MAX` if no Progress was ever recorded. + pub first_progress_ns: u64, + /// Wall ns (from pipeline start) of this step's last Progress + /// dispatch end. `0` if no Progress was ever recorded. + pub last_progress_ns: u64, +} + +impl StepStatsSnapshot { + #[must_use] + pub fn avg_run_ns(&self) -> Option { + if self.try_run_total == 0 { None } else { Some(self.total_run_ns / self.try_run_total) } + } +} + +/// Severity of a [`Finding`] from the bottleneck verdict. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Severity { + /// The single rate-limiting step (full input edge + empty output edge). + Primary, + /// A contributing issue (spin, starvation) that is not the prime mover. + Secondary, + /// Whole-chain observation (e.g. latency-bound, not throughput-bound). + Info, +} + +/// One mechanically-derived diagnosis from [`bottleneck_verdict`]. +#[derive(Debug, Clone)] +pub(crate) struct Finding { + /// Triage classification. Set by [`bottleneck_verdict`] and asserted on by + /// the crate's tests; the human-readable report renders only `message`, so + /// the lib build never reads this field. + #[allow(dead_code)] + pub severity: Severity, + pub message: String, +} + +/// Rate above which a Serial step's `contention/tries` ratio is flagged as spin. +const SPIN_THRESHOLD: f64 = 0.20; + +/// Mechanically locate the chain's bottleneck and contributing issues from a +/// snapshot's step + edge stats. Rules (all from already-collected numbers): +/// - **Primary**: a step whose input edge is `Full`/`Backpressured` AND output +/// edge is `Empty`/`Starved` — work piles into it, it can't fill downstream. +/// Tagged CPU-bound if it dominates `total_run_ns`, else coordination. +/// - **Secondary (spin)**: a step with a high `contention/tries` ratio. +/// - **Secondary (starvation)**: an edge whose consumer frequently finds it +/// empty — the producer (upstream) can't keep up. +/// - **Info**: if no primary and no full/empty edge, the chain is +/// latency/coordination-bound, not throughput-bound. +#[must_use] +#[allow(clippy::cast_precision_loss)] +pub(crate) fn bottleneck_verdict(snap: &StatsSnapshot) -> Vec { + use OccupancyClass::{Backpressured, Empty, Full, Starved}; + let mut findings = Vec::new(); + let total_cpu_ns: u64 = snap.steps.iter().map(|(_, s)| s.total_run_ns).sum(); + // Attribute edges to steps by step IDENTITY (index), not by name. `snap.steps` + // is in `StepIdx` order, so a step's index is its id. A step's input is "full" + // if ANY input edge (whose consumer is this step) is Full/Backpressured, and + // its output is "empty" if ANY output edge (whose producer is this step) is + // Empty/Starved — aggregating over all matching edges rather than the first + // name match, so fan-out / fan-in / duplicate step names aren't misattributed. + let input_full = |step: usize| { + snap.edges + .iter() + .filter(|e| e.consumer_step == Some(step)) + .any(|e| matches!(e.class, Full | Backpressured)) + }; + let output_empty = |step: usize| { + snap.edges + .iter() + .filter(|e| e.producer_step == step) + .any(|e| matches!(e.class, Empty | Starved)) + }; + + // Primary bottleneck: full input edge + empty output edge. + for (step, (name, s)) in snap.steps.iter().enumerate() { + if input_full(step) && output_empty(step) { + let cpu_share = + if total_cpu_ns > 0 { s.total_run_ns as f64 / total_cpu_ns as f64 } else { 0.0 }; + let cause = if cpu_share >= 0.30 { + format!("CPU-bound ({:.0}% of pipeline try_run time)", cpu_share * 100.0) + } else { + "coordination-bound (low CPU share — likely a serialization stall)".to_string() + }; + findings.push(Finding { + severity: Severity::Primary, + message: format!( + "BOTTLENECK: step `{name}` (input edge full, output edge empty) — {cause}" + ), + }); + } + } + + // Secondary: Serial-step spin (contention thrash). Skip Detached steps — + // their dedicated-thread backoff loop records a NoProgress/Contention on + // every idle poll, which inflates the contention ratio, and "Detach + // candidate" is meaningless for a step that is already Detached. + // Match by `StepIdx`, not by name — as the edge attribution above does. Two + // steps may share a name, and a name match would let one Detached step's + // exemption silence a genuinely thrashing pool step that happens to share it. + let detached_steps: std::collections::HashSet = + snap.detached.iter().map(|&(step, ..)| step).collect(); + for (step, (name, s)) in snap.steps.iter().enumerate() { + if detached_steps.contains(&step) { + continue; + } + if s.try_run_total > 0 { + let spin = s.contention_count as f64 / s.try_run_total as f64; + if spin >= SPIN_THRESHOLD { + findings.push(Finding { + severity: Severity::Secondary, + message: format!( + "SPIN: step `{name}` contended on {:.0}% of dispatches — affinity / fuse / Detach candidate", + spin * 100.0 + ), + }); + } + } + } + + // Secondary: starvation (consumer of an edge frequently finds it empty). + for e in &snap.edges { + if matches!(e.class, Starved) { + findings.push(Finding { + severity: Severity::Secondary, + message: format!( + "STARVATION: `{}` is starved by upstream `{}` (empty {:.0}% of pops) — look upstream", + e.consumer.unwrap_or("(sink)"), + e.producer, + e.empty_rate * 100.0 + ), + }); + } + } + + // Info: no clear bottleneck and no full/empty edge → latency-bound. + let any_extreme = + snap.edges.iter().any(|e| matches!(e.class, Full | Backpressured | Empty | Starved)); + if findings.is_empty() && !snap.edges.is_empty() && !any_extreme { + findings.push(Finding { + severity: Severity::Info, + message: "No throughput bottleneck — edges have headroom; the chain is \ + latency/coordination-bound (look at residence time, not service time)." + .to_string(), + }); + } + findings +} + +impl fmt::Display for StatsSnapshot { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + writeln!(f, "Pipeline stats ({} step{}):", self.steps.len(), pluralize(self.steps.len()))?; + // Total dispatch time across steps — the denominator for each step's + // `cpu%` share (the Amdahl ranking: which step to optimize first) and + // for the headroom line. + let total_cpu_ns: u64 = self.steps.iter().map(|(_, s)| s.total_run_ns).sum(); + self.write_steps(f, total_cpu_ns)?; + self.write_utilization(f, total_cpu_ns)?; + // Per-edge throughput / occupancy / latency + the bottleneck verdict + // (both only when instrumented). + self.write_edges(f)?; + self.write_verdict(f) + } +} + +impl StatsSnapshot { + /// Render the per-step counter table, including each step's `cpu%` share of + /// total dispatch time (the per-step cost ranking). `total_cpu_ns` is the + /// sum of all steps' `total_run_ns`, passed in to avoid recomputing. + fn write_steps(&self, f: &mut fmt::Formatter<'_>, total_cpu_ns: u64) -> fmt::Result { + writeln!( + f, + " {:<28} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>14} {:>7} {:>12} {:>12}", + "step", + "tries", + "progress", + "noprog", + "content", + "fin", + "err", + "total_ms", + "cpu%", + "first_ms", + "last_ms", + )?; + for (name, s) in &self.steps { + // ns -> ms with three decimal places. The `as f64` cast can lose + // precision above 2^52 ns (~52 days of cumulative run time per + // step), well beyond any real pipeline run. + #[allow(clippy::cast_precision_loss)] + let total_ms = (s.total_run_ns as f64) / 1_000_000.0; + let first_str = if s.first_progress_ns == u64::MAX { + "-".to_string() + } else { + #[allow(clippy::cast_precision_loss)] + let ms = (s.first_progress_ns as f64) / 1_000_000.0; + format!("{ms:.1}") + }; + let last_str = if s.last_progress_ns == 0 { + "-".to_string() + } else { + #[allow(clippy::cast_precision_loss)] + let ms = (s.last_progress_ns as f64) / 1_000_000.0; + format!("{ms:.1}") + }; + // Share of total dispatch time — the per-step cost ranking. The + // chain-order rows stay readable; this column gives the magnitude. + #[allow(clippy::cast_precision_loss)] + let cpu_pct = if total_cpu_ns == 0 { + 0.0 + } else { + s.total_run_ns as f64 / total_cpu_ns as f64 * 100.0 + }; + writeln!( + f, + " {:<28} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>14.3} {:>6.1}% {:>12} {:>12}", + name, + s.try_run_total, + s.progress_count, + s.no_progress_count, + s.contention_count, + s.finished_count, + s.error_count, + total_ms, + cpu_pct, + first_str, + last_str, + )?; + } + Ok(()) + } + + /// Render per-worker utilisation, the pool-level utilisation, and the + /// squeeze-headroom synthesis. No-op when no worker recorded activity. + /// + /// Per-worker: busy = dispatch time, idle = backoff-sleep time. A high + /// idle% (cores parked while one worker drives a Serial step) is the + /// signature of pool under-utilisation. The headroom line reads straight off + /// these numbers: the pool idle% is recoverable via better step overlap; + /// once the pool saturates (no idle left) the only remaining wins are fewer + /// cycles/item or more cores, and the hottest step is the Amdahl target. + fn write_utilization(&self, f: &mut fmt::Formatter<'_>, total_cpu_ns: u64) -> fmt::Result { + if self.workers.is_empty() { + // No pool workers recorded activity, but a Detached thread may still + // have run (e.g. a degenerate chain). Render its line if so. + self.write_detached(f)?; + return Ok(()); + } + #[allow(clippy::cast_precision_loss)] + let ms = |ns: u64| (ns as f64) / 1_000_000.0; + let (mut sum_busy, mut sum_idle) = (0u64, 0u64); + writeln!(f, " {:<8} {:>14} {:>14} {:>8}", "worker", "busy_ms", "idle_ms", "busy%")?; + for &(id, busy, idle) in &self.workers { + sum_busy += busy; + sum_idle += idle; + let pct = if busy + idle == 0 { + 0.0 + } else { + #[allow(clippy::cast_precision_loss)] + let p = busy as f64 / (busy + idle) as f64 * 100.0; + p + }; + writeln!(f, " {id:<8} {:>14.3} {:>14.3} {pct:>7.1}%", ms(busy), ms(idle))?; + } + let pool_pct = if sum_busy + sum_idle == 0 { + 0.0 + } else { + #[allow(clippy::cast_precision_loss)] + let p = sum_busy as f64 / (sum_busy + sum_idle) as f64 * 100.0; + p + }; + writeln!( + f, + " pool utilisation: {pool_pct:.1}% busy across {} worker{} ({:.3} ms busy / {:.3} ms idle)", + self.workers.len(), + pluralize(self.workers.len()), + ms(sum_busy), + ms(sum_idle), + )?; + if let Some((hot_name, hot)) = self.steps.iter().max_by_key(|(_, s)| s.total_run_ns) { + #[allow(clippy::cast_precision_loss)] + let hot_pct = if total_cpu_ns == 0 { + 0.0 + } else { + hot.total_run_ns as f64 / total_cpu_ns as f64 * 100.0 + }; + writeln!( + f, + " headroom: {:.1}% pool idle (recoverable via overlap); hottest step `{hot_name}` = {hot_pct:.1}% of dispatch time (Amdahl target once pool saturates)", + 100.0 - pool_pct, + )?; + } + self.write_detached(f) + } + + /// Render the `StepKind::Detached` threads' busy/idle on their own line(s). + /// These threads are NOT pool workers (legacy "N + 2"), so they are reported + /// separately and never folded into the pool utilisation %. No-op when no + /// Detached thread recorded activity (every non-sort chain). + fn write_detached(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + if self.detached.is_empty() { + return Ok(()); + } + #[allow(clippy::cast_precision_loss)] + let ms = |ns: u64| (ns as f64) / 1_000_000.0; + for &(_step, name, busy, idle, parks) in &self.detached { + let pct = if busy + idle == 0 { + 0.0 + } else { + #[allow(clippy::cast_precision_loss)] + let p = busy as f64 / (busy + idle) as f64 * 100.0; + p + }; + // avg park = idle / parks: a few long parks (idle at end) is benign; + // many short parks during productive phases is the backoff-latency + // signal (the deferred per-slot condvar would eliminate them). + #[allow(clippy::cast_precision_loss)] + let avg_park_us = if parks == 0 { 0.0 } else { (idle as f64 / parks as f64) / 1000.0 }; + writeln!( + f, + " detached `{name}`: {:.3} ms busy / {:.3} ms idle ({pct:.1}% busy, off pool); {parks} parks (avg {avg_park_us:.1}µs)", + ms(busy), + ms(idle), + )?; + } + Ok(()) + } + + /// Render the mechanically-derived bottleneck verdict. No-op when there are + /// no instrumented edges (nothing to diagnose). + fn write_verdict(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + if self.edges.is_empty() { + return Ok(()); + } + let findings = bottleneck_verdict(self); + if findings.is_empty() { + return Ok(()); + } + writeln!(f, "Bottleneck verdict:")?; + for finding in findings { + writeln!(f, " {}", finding.message)?; + } + Ok(()) + } + + /// Render the per-edge table (throughput / occupancy class / rates / + /// latency). No-op when there are no instrumented edges. + fn write_edges(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + if self.edges.is_empty() { + return Ok(()); + } + writeln!(f, "Pipeline edges ({}):", self.edges.len())?; + writeln!( + f, + " {:<40} {:>12} {:>10} {:>14} {:>6} {:>6} {:>10}", + "producer→consumer", "items/s", "MiB/s", "class", "rej%", "empt%", "lat_ms", + )?; + for e in &self.edges { + let edge = format!("{}→{}", e.producer, e.consumer.unwrap_or("(none)")); + let lat = e.derived_latency_ms.map_or_else(|| "-".to_string(), |l| format!("{l:.2}")); + writeln!( + f, + " {:<40} {:>12.0} {:>10.1} {:>14} {:>5.1}% {:>5.1}% {:>10}", + edge, + e.items_per_s, + e.mibytes_per_s, + format!("{:?}", e.class), + e.reject_rate * 100.0, + e.empty_rate * 100.0, + lat, + )?; + } + writeln!( + f, + " (note: throughput is depressed under tracing — confirm wall/RSS with --pipeline-trace off)" + ) + } +} + +fn pluralize(n: usize) -> &'static str { + if n == 1 { "" } else { "s" } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[test] + fn record_increments_progress_bucket() { + let stats = PipelineStats::new(vec!["A", "B"]); + stats.record(StepIdx(0), StepOutcome::Progress, 0, 100); + stats.record(StepIdx(0), StepOutcome::Progress, 200, 50); + stats.record(StepIdx(1), StepOutcome::NoProgress, 0, 25); + + let snap = stats.snapshot(); + assert_eq!(snap.steps.len(), 2); + + let (name_a, a) = &snap.steps[0]; + assert_eq!(*name_a, "A"); + assert_eq!(a.try_run_total, 2); + assert_eq!(a.progress_count, 2); + assert_eq!(a.no_progress_count, 0); + assert_eq!(a.total_run_ns, 150); + assert_eq!(a.first_progress_ns, 0); + assert_eq!(a.last_progress_ns, 250); + + let (name_b, b) = &snap.steps[1]; + assert_eq!(*name_b, "B"); + assert_eq!(b.try_run_total, 1); + assert_eq!(b.no_progress_count, 1); + assert_eq!(b.total_run_ns, 25); + assert_eq!(b.first_progress_ns, u64::MAX); + assert_eq!(b.last_progress_ns, 0); + } + + #[test] + fn record_buckets_each_outcome() { + let stats = PipelineStats::new(vec!["S"]); + stats.record(StepIdx(0), StepOutcome::Progress, 0, 1); + stats.record(StepIdx(0), StepOutcome::NoProgress, 10, 2); + stats.record(StepIdx(0), StepOutcome::Contention, 20, 3); + stats.record(StepIdx(0), StepOutcome::Finished, 30, 4); + stats.record_error(StepIdx(0), 40, 5); + + let snap = stats.snapshot(); + let (_name, s) = &snap.steps[0]; + assert_eq!(s.try_run_total, 5); + assert_eq!(s.progress_count, 1); + assert_eq!(s.no_progress_count, 1); + assert_eq!(s.contention_count, 1); + assert_eq!(s.finished_count, 1); + assert_eq!(s.error_count, 1); + assert_eq!(s.total_run_ns, 1 + 2 + 3 + 4 + 5); + } + + #[test] + fn avg_run_ns_handles_empty_step() { + let snap = StepStatsSnapshot { + try_run_total: 0, + progress_count: 0, + no_progress_count: 0, + contention_count: 0, + finished_count: 0, + error_count: 0, + total_run_ns: 0, + first_progress_ns: u64::MAX, + last_progress_ns: 0, + }; + assert_eq!(snap.avg_run_ns(), None); + + let snap2 = StepStatsSnapshot { try_run_total: 4, total_run_ns: 1000, ..snap }; + assert_eq!(snap2.avg_run_ns(), Some(250)); + } + + #[allow(clippy::too_many_arguments)] // test builder: one arg per snapshot field + fn edge_ms( + pushed_items: u64, + pushed_bytes: u64, + popped_items: u64, + pop_empties: u64, + push_rejections: u64, + raw: crate::runtime::metrics::RawOccupancy, + mean_occupancy: f32, + mean_occupancy_bytes: f64, + ) -> crate::runtime::metrics::EdgeMetricsSnapshot { + crate::runtime::metrics::EdgeMetricsSnapshot { + pushed_items, + pushed_bytes, + popped_items, + popped_bytes: pushed_bytes, + push_rejections, + pop_empties, + depth_samples: 100, + raw_occupancy: raw, + mean_occupancy, + mean_occupancy_bytes, + } + } + + #[test] + fn edge_stats_byte_bounded_latency_is_dimensionally_correct() { + use crate::runtime::metrics::RawOccupancy; + // 1000 items/s; mean_item_bytes = 100_000/1000 = 100; sampled mean + // occupancy 1000 bytes → mean_items 10 → latency = 1000*10/1000 = 10ms. + // (The byte figure is sampled directly, so the limit only gates that the + // edge is byte-bounded — its value no longer feeds the latency.) + let ms = edge_ms(1000, 100_000, 1000, 0, 0, RawOccupancy::Healthy, 0.5, 1000.0); + let e = compute_edge_stats("p", Some("c"), 0, Some(1), &ms, Some(2000), 1_000_000_000); + assert!((e.items_per_s - 1000.0).abs() < 1.0, "items/s ≈ 1000, got {}", e.items_per_s); + let lat = e.derived_latency_ms.expect("byte edge has derived latency"); + assert!((lat - 10.0).abs() < 0.5, "latency ≈ 10ms, got {lat}"); + } + + #[test] + fn edge_stats_count_edge_has_no_latency() { + use crate::runtime::metrics::RawOccupancy; + let ms = edge_ms(500, 0, 500, 0, 0, RawOccupancy::Unknown, 0.0, 0.0); + let e = compute_edge_stats("p", Some("c"), 0, Some(1), &ms, None, 1_000_000_000); + assert!(e.derived_latency_ms.is_none(), "count/unbounded edge → no derived latency"); + assert!((e.items_per_s - 500.0).abs() < 1.0); + } + + // The raw occupancy class is refined by the empty-rate and reject-rate the + // edge actually observed; each case pins one (raw class, rates, byte bound) + // combination to its refined `OccupancyClass`. + #[rstest] + // MostlyEmpty + high empty-rate (90/100 pops empty) → Starved. + #[case::starved( + edge_ms(10, 0, 10, 90, 0, crate::runtime::metrics::RawOccupancy::MostlyEmpty, 0.0, 0.0), + None, + OccupancyClass::Starved + )] + // MostlyEmpty + low empty-rate → Empty (idle, not starved). + #[case::empty( + edge_ms(100, 0, 100, 0, 0, crate::runtime::metrics::RawOccupancy::MostlyEmpty, 0.0, 0.0), + None, + OccupancyClass::Empty + )] + // MostlyFull + high reject-rate → Backpressured. + #[case::backpressured( + edge_ms(10, 0, 10, 0, 90, crate::runtime::metrics::RawOccupancy::MostlyFull, 1.0, 1000.0), + Some(1000), + OccupancyClass::Backpressured + )] + fn occupancy_class_refined_by_rates( + #[case] ms: crate::runtime::metrics::EdgeMetricsSnapshot, + #[case] limit_bytes: Option, + #[case] expected: OccupancyClass, + ) { + assert_eq!( + compute_edge_stats("p", None, 0, None, &ms, limit_bytes, 1_000_000_000).class, + expected + ); + } + + fn step_stat( + name: &'static str, + total_run_ns: u64, + contention: u64, + tries: u64, + ) -> (&'static str, StepStatsSnapshot) { + ( + name, + StepStatsSnapshot { + try_run_total: tries, + progress_count: tries, + no_progress_count: 0, + contention_count: contention, + finished_count: 1, + error_count: 0, + total_run_ns, + first_progress_ns: 0, + last_progress_ns: total_run_ns, + }, + ) + } + + fn edge_stat( + producer: &'static str, + consumer: Option<&'static str>, + producer_step: usize, + consumer_step: Option, + class: OccupancyClass, + empty_rate: f64, + ) -> EdgeStatsSnapshot { + EdgeStatsSnapshot { + producer, + consumer, + producer_step, + consumer_step, + items_per_s: 1000.0, + mibytes_per_s: 1.0, + class, + mean_occupancy: 0.5, + reject_rate: 0.0, + empty_rate, + derived_latency_ms: Some(1.0), + } + } + + #[test] + fn verdict_locates_cpu_bound_bottleneck() { + // `Slow`'s input edge is Full and output edge is Empty, and it dominates + // CPU → Primary, CPU-bound. + let snap = StatsSnapshot { + steps: vec![step_stat("Slow", 1000, 0, 10)], + workers: vec![], + detached: vec![], + // `Slow` is step 0. Its input edge (consumer_step 0) is Full; its + // output edge (producer_step 0) is Empty. `Up`/`Down` are non-step + // ids (1/1) so only `Slow` matches by identity. + edges: vec![ + edge_stat("Up", Some("Slow"), 1, Some(0), OccupancyClass::Full, 0.0), + edge_stat("Slow", Some("Down"), 0, Some(1), OccupancyClass::Empty, 0.0), + ], + }; + let v = bottleneck_verdict(&snap); + let primary: Vec<_> = v.iter().filter(|f| f.severity == Severity::Primary).collect(); + assert_eq!(primary.len(), 1, "exactly one primary bottleneck"); + assert!(primary[0].message.contains("Slow")); + assert!(primary[0].message.contains("CPU-bound")); + } + + #[test] + fn verdict_matches_edges_by_step_identity_across_fan_in() { + // Regression for name-based misattribution: `Merge` (step 2) has TWO + // input edges (fan-in from steps 0 and 1). The first by iteration order + // is Healthy; the second is Full. Matching the FIRST edge by consumer + // name would see only the Healthy input and report no bottleneck; + // matching by step identity aggregates both inputs and correctly flags + // `Merge` (input full + output empty). + let snap = StatsSnapshot { + steps: vec![ + step_stat("A", 100, 0, 10), + step_stat("B", 100, 0, 10), + step_stat("Merge", 1000, 0, 10), + ], + workers: vec![], + detached: vec![], + edges: vec![ + edge_stat("A", Some("Merge"), 0, Some(2), OccupancyClass::Healthy, 0.0), + edge_stat("B", Some("Merge"), 1, Some(2), OccupancyClass::Full, 0.0), + edge_stat("Merge", Some("Sink"), 2, Some(3), OccupancyClass::Empty, 0.0), + ], + }; + let v = bottleneck_verdict(&snap); + let primary: Vec<_> = v.iter().filter(|f| f.severity == Severity::Primary).collect(); + assert_eq!(primary.len(), 1, "fan-in bottleneck detected via identity: {v:?}"); + assert!(primary[0].message.contains("Merge")); + } + + #[test] + fn verdict_flags_spin_and_starvation() { + let snap = StatsSnapshot { + // 5/10 dispatches contended → spin. + steps: vec![step_stat("Serializer", 100, 5, 10)], + workers: vec![], + detached: vec![], + // A starved edge (steps 1→2, not the `Serializer` step 0): consumer + // frequently finds it empty. + edges: vec![edge_stat( + "Producer", + Some("Consumer"), + 1, + Some(2), + OccupancyClass::Starved, + 0.7, + )], + }; + let v = bottleneck_verdict(&snap); + assert!( + v.iter().any(|f| f.severity == Severity::Secondary && f.message.contains("SPIN")), + "spin finding present" + ); + assert!( + v.iter().any(|f| f.message.contains("STARVATION") && f.message.contains("Consumer")), + "starvation finding present" + ); + } + + #[test] + fn verdict_skips_spin_for_detached_steps() { + // A Detached step's backoff loop records Contention on every idle poll, + // so its contention/tries ratio is high — but "SPIN: … Detach candidate" + // is nonsensical for an already-Detached step, so it must be skipped. + let snap = StatsSnapshot { + // 8/10 "contended" — would trip SPIN if it were a pool step. + steps: vec![step_stat("SortMerge", 100, 8, 10)], + workers: vec![], + detached: vec![(0, "SortMerge", 900_000_000, 100_000_000, 4_200)], + edges: vec![], + }; + let v = bottleneck_verdict(&snap); + assert!( + !v.iter().any(|f| f.message.contains("SPIN")), + "no SPIN finding for a Detached step: {v:?}" + ); + } + + /// The exemption is per step IDENTITY, not per name. Two steps sharing a + /// name — one Detached, one a thrashing pool step — must be judged + /// independently: exempting by name silenced the pool step's SPIN finding. + #[test] + fn verdict_reports_spin_for_a_pool_step_sharing_a_detached_step_name() { + let snap = StatsSnapshot { + // Step 0 is the Detached merge; step 1 is a pool step that happens to + // carry the same name and is genuinely thrashing (8/10 contended). + steps: vec![step_stat("SortMerge", 100, 8, 10), step_stat("SortMerge", 100, 8, 10)], + workers: vec![], + detached: vec![(0, "SortMerge", 900_000_000, 100_000_000, 4_200)], + edges: vec![], + }; + let v = bottleneck_verdict(&snap); + assert_eq!( + v.iter().filter(|f| f.message.contains("SPIN")).count(), + 1, + "exactly the pool step (index 1) reports SPIN; the Detached step (index 0) is \ + exempt: {v:?}" + ); + } + + #[test] + fn verdict_reports_latency_bound_when_no_extremes() { + let snap = StatsSnapshot { + steps: vec![step_stat("A", 100, 0, 10)], + workers: vec![], + detached: vec![], + edges: vec![edge_stat("A", Some("B"), 0, Some(1), OccupancyClass::Healthy, 0.0)], + }; + let v = bottleneck_verdict(&snap); + assert_eq!(v.len(), 1); + assert_eq!(v[0].severity, Severity::Info); + assert!(v[0].message.contains("latency")); + } + + #[test] + fn snapshot_display_renders_without_edges() { + // Back-compat: the edge-less snapshot() still renders (empty edges). + let stats = PipelineStats::new(vec!["A"]); + let out = format!("{}", stats.snapshot()); + assert!(out.contains("Pipeline stats")); + assert!(!out.contains("Pipeline edges"), "no edge section when edges empty"); + } + + #[test] + fn display_renders_cpu_share_and_headroom() { + // Hot owns 750/1000 = 75% of dispatch time; Cool owns the rest. The pool + // is 80% busy (800 ms / 1000 ms) → 20% idle headroom. + let snap = StatsSnapshot { + steps: vec![step_stat("Hot", 750, 0, 10), step_stat("Cool", 250, 0, 10)], + workers: vec![(0, 800, 200)], + detached: vec![], + edges: vec![], + }; + let out = format!("{snap}"); + assert!(out.contains("cpu%"), "per-step table has a cpu% column: {out}"); + assert!(out.contains("75.0%"), "Hot step shows its 75% CPU share: {out}"); + assert!(out.contains("pool utilisation: 80.0%"), "pool utilisation line present: {out}"); + assert!( + out.contains("headroom:") && out.contains("20.0% pool idle") && out.contains("`Hot`"), + "headroom line names the hottest step and the idle %: {out}" + ); + } + + #[test] + fn utilization_section_absent_without_workers() { + // No worker activity → no pool/headroom lines (fused or stats-off runs). + let snap = StatsSnapshot { + steps: vec![step_stat("Solo", 100, 0, 10)], + workers: vec![], + detached: vec![], + edges: vec![], + }; + let out = format!("{snap}"); + assert!(!out.contains("pool utilisation"), "no pool line without workers: {out}"); + assert!(!out.contains("headroom:"), "no headroom line without workers: {out}"); + } + + #[test] + fn display_renders_header_and_rows() { + let stats = PipelineStats::new(vec!["StepA", "StepB"]); + stats.record(StepIdx(0), StepOutcome::Progress, 0, 10); + let s = format!("{}", stats.snapshot()); + assert!(s.contains("Pipeline stats (2 steps):")); + assert!(s.contains("StepA")); + assert!(s.contains("StepB")); + assert!(s.contains("step")); + } + + /// L2.4: a Detached thread's busy/idle is recorded separately, renders on + /// its own line, and is EXCLUDED from the pool utilisation % (legacy + /// "N + 2"). The snapshot must render without panicking. + #[test] + fn detached_busy_excluded_from_pool_and_on_own_line() { + let stats = PipelineStats::new(vec!["Source", "SortMerge"]); + // One pool worker: 800 ms busy / 200 ms idle → pool% = 80%. + stats.record_worker_busy(0, 800_000_000); + stats.record_worker_idle(0, 200_000_000); + // The Detached merge: 900 ms busy / 100 ms idle (off pool). + stats.record_detached_busy(StepIdx(1), 900_000_000); + stats.record_detached_idle(StepIdx(1), 100_000_000); + + let snap = stats.snapshot(); + assert_eq!(snap.detached.len(), 1, "one Detached thread recorded"); + assert_eq!( + (snap.detached[0].0, snap.detached[0].1), + (1, "SortMerge"), + "the entry carries its StepIdx alongside the name" + ); + + let out = format!("{snap}"); + // Pool% reflects ONLY the worker (80%), NOT the Detached thread (90%). + assert!( + out.contains("pool utilisation: 80.0%"), + "pool% must exclude the Detached thread: {out}" + ); + // The Detached thread is reported on its own line. + assert!( + out.contains("detached `SortMerge`") && out.contains("90.0% busy, off pool"), + "Detached line present with its own busy%: {out}" + ); + } + + /// L2.4: the Detached line renders even when no pool worker recorded + /// activity (the `write_utilization` early-return path), without panicking. + #[test] + fn detached_line_renders_with_no_pool_workers() { + let stats = PipelineStats::new(vec!["OnlyDetached"]); + stats.record_detached_busy(StepIdx(0), 5_000_000); + let out = format!("{}", stats.snapshot()); + assert!(out.contains("detached `OnlyDetached`"), "Detached line present: {out}"); + assert!(!out.contains("pool utilisation"), "no pool line without workers: {out}"); + } +} diff --git a/crates/fgumi-pipeline-core/src/runtime/storage.rs b/crates/fgumi-pipeline-core/src/runtime/storage.rs new file mode 100644 index 000000000..3ab10d59d --- /dev/null +++ b/crates/fgumi-pipeline-core/src/runtime/storage.rs @@ -0,0 +1,431 @@ +//! `WorkerStepEntry`: per-worker per-step storage shape. Determined by +//! `StepKind` at run start; immutable thereafter. + +use parking_lot::Mutex; +use std::sync::{Arc, OnceLock}; + +use crate::erased::ErasedStep; +use crate::step::{Affinity, StepKind}; + +/// Shared "this step has finished" latch for a `Serial` step. +/// +/// A `Serial` step is one shared instance behind a `Mutex` across all workers, +/// but each worker tracks its own worklist. Without a shared latch, after one +/// worker runs the step to `Finished` the others still hold it in their +/// worklist and would re-`try_lock`/re-run an already-finished step. The +/// finishing worker sets this latch under the dispatch guard (before +/// `mark_outputs_drained`); every other worker observes it and short-circuits +/// to a synthetic `Finished` instead of re-running the step. +#[derive(Default)] +pub struct DrainGate { + finished: OnceLock<()>, +} + +impl DrainGate { + /// True once the step has run to `StepOutcome::Finished` on some worker. + pub fn is_finished(&self) -> bool { + self.finished.get().is_some() + } + + /// Latch the step as finished (idempotent). Set under the dispatch guard, + /// before `mark_outputs_drained`, so a concurrent worker that observes the + /// latch never re-runs the step. + pub fn mark_finished(&self) { + let _ = self.finished.set(()); + } +} + +/// One per (worker, step) cell. +pub enum WorkerStepEntry { + /// `Parallel` step: this worker owns its private clone via `clone_boxed`. + /// Direct `&mut` access; no locking on `try_run` dispatch. Completion is + /// coordinated by the per-step `StepDrainCounter` (init N) in the driver: + /// every clone returns `Finished` when the input drains, but only the last + /// to finish closes the shared output. + Owned { step: Box }, + /// `Serial` step: shared instance, mutex-protected. Any worker can acquire. + /// The shared `DrainGate` finished-latch lets a worker that finishes the + /// step stop the others from re-running it. + Shared { step: Arc>>, drain: Arc }, + /// `Exclusive` step: only this worker (the owner) ever runs it. + /// Stored locally on the owner; other workers have `Skip`. + Exclusive { step: Box }, + /// This worker doesn't run this step (it's an `Exclusive` step owned by + /// another worker). The worker loop skips it in dispatch. + Skip, +} + +impl WorkerStepEntry { + /// True if this worker should attempt to dispatch this step. + #[must_use] + pub fn is_dispatchable(&self) -> bool { + !matches!(self, Self::Skip) + } +} + +/// Assemble per-worker step storage. +/// +/// `steps`: the chain in order (consumed). +/// `exclusive_owners[step_idx] == Some(worker_id)` for each Exclusive step. +/// `n_workers`: number of worker threads. +/// +/// Returns `entries[worker_id][step_idx] = WorkerStepEntry`. +/// +/// # Panics +/// +/// Panics if `exclusive_owners.len() != steps.len()` or if an Exclusive +/// step has no owner assignment (`assign_exclusive_owners` must run first). +#[must_use] +pub fn build_worker_storage( + steps: Vec>, + exclusive_owners: &[Option], + n_workers: usize, +) -> Vec> { + assert_eq!( + steps.len(), + exclusive_owners.len(), + "exclusive_owners length must match step count" + ); + assert!(n_workers > 0, "build_worker_storage requires at least one worker"); + + let mut entries: Vec> = + (0..n_workers).map(|_| Vec::with_capacity(steps.len())).collect(); + + for (step_idx, step) in steps.into_iter().enumerate() { + let kind = step.kind(); + match kind { + StepKind::Parallel => { + // Each worker gets its own clone; the original goes to worker N-1. + for entries_for_worker in entries.iter_mut().take(n_workers - 1) { + entries_for_worker.push(WorkerStepEntry::Owned { step: step.clone_boxed() }); + } + entries[n_workers - 1].push(WorkerStepEntry::Owned { step }); + } + StepKind::Serial => { + // Snapshot the affinity hint before moving `step` into + // the shared `Arc>` — affinity is part of the + // static `Step` description, so calling it here is a + // single virtual dispatch. + let affinity = step.affinity(); + // Always-on (not `debug_assert!`): an out-of-range affinity + // gates the Serial step out of every worker, so it is never + // dispatched and the pipeline deadlocks. The check is a cheap + // pure predicate, so it stays enabled in release builds too. + assert!( + affinity_in_range(affinity, n_workers), + "Serial step affinity {affinity:?} requests a worker out of range \ + (n_workers = {n_workers}); pipeline would deadlock" + ); + let drain = Arc::new(DrainGate::default()); + let shared = Arc::new(Mutex::new(step)); + for (worker_id, entries_for_worker) in entries.iter_mut().enumerate() { + if affinity.eligible(worker_id, n_workers) { + entries_for_worker.push(WorkerStepEntry::Shared { + step: Arc::clone(&shared), + drain: Arc::clone(&drain), + }); + } else { + // This worker is gated out by the Serial step's + // affinity hint. It will never `try_lock` the step's + // mutex, eliminating the `Contention` thrash that + // pure `Serial` exhibits at high thread counts. + entries_for_worker.push(WorkerStepEntry::Skip); + } + } + } + StepKind::Exclusive => { + let owner = exclusive_owners[step_idx].expect( + "Exclusive step has no owner assignment; \ + assign_exclusive_owners must run before build_worker_storage", + ); + // Always-on, matching the Serial arm's affinity check above: an + // out-of-range owner would otherwise surface as an opaque slice + // index panic on `entries[owner]`. `assign_sticky_owners` + // explicitly tolerates an out-of-range owner (and its + // `out_of_range` test case treats one as reachable input), so the + // two functions must not disagree about whether it can happen. + assert!( + owner < n_workers, + "Exclusive step owner {owner} is out of range (n_workers = {n_workers}); \ + the step would never be dispatched and the pipeline would deadlock" + ); + // Push Skip placeholders for all workers, then replace owner's slot. + for entries_for_worker in &mut entries { + entries_for_worker.push(WorkerStepEntry::Skip); + } + entries[owner][step_idx] = WorkerStepEntry::Exclusive { step }; + } + StepKind::Detached => { + // A Detached step never runs on the pool. Every real Detached + // instance is extracted by `extract_detached_steps` BEFORE this + // function runs (it drives the step on its own dedicated + // thread) and is replaced in `steps` by a `DetachedPlaceholder` + // whose `kind()` still reports `StepKind::Detached`. So on every + // real sort run this arm IS reached — once per placeholder — and + // its job is exactly this: give every pool worker a `Skip` entry + // and drop the placeholder (it holds no state to run). + for entries_for_worker in &mut entries { + entries_for_worker.push(WorkerStepEntry::Skip); + } + drop(step); + } + } + } + + entries +} + +/// Validate that a Serial step's affinity refers to a worker that exists. +/// `Affinity::None` / `Reader` / `Writer` always resolve to a real worker +/// when `n_workers >= 1` (which the runtime's `assert!(n_workers > 0)` +/// already guarantees). `Worker(idx)` is only valid when `idx < n_workers`. +/// +/// Routed through [`Affinity::target_worker`] rather than re-matching the +/// variants here: that method is the single source of truth for the +/// affinity→worker mapping, and the dispatch gate resolves the target the same +/// way. A future variant therefore cannot make this range check and the gate +/// disagree about which worker a step lands on. +#[inline] +fn affinity_in_range(affinity: Affinity, n_workers: usize) -> bool { + affinity.target_worker(n_workers).is_none_or(|target| target < n_workers) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io; + + use rstest::rstest; + + use crate::erased::TypedStep; + use crate::outputs::Single; + use crate::queues::QueueSpec; + use crate::reorder::BranchOrdering; + use crate::step::{Affinity, Step, StepCtx, StepOutcome, StepProfile}; + + fn profile_for(name: &'static str, kind: StepKind, sticky: bool) -> StepProfile { + StepProfile { + name, + kind, + sticky, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + + #[derive(Clone)] + struct ParallelStep; + impl Step for ParallelStep { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + profile_for("Par", StepKind::Parallel, false) + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + #[derive(Clone)] + struct SerialStep; + impl Step for SerialStep { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + profile_for("Ser", StepKind::Serial, false) + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + #[derive(Clone)] + struct ExclusiveStep; + impl Step for ExclusiveStep { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + profile_for("Excl", StepKind::Exclusive, false) + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + #[derive(Clone)] + struct DetachedStep; + impl Step for DetachedStep { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + profile_for("Det", StepKind::Detached, false) + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + #[test] + fn skip_is_not_dispatchable() { + assert!(!WorkerStepEntry::Skip.is_dispatchable()); + } + + /// L2.1: a `Detached` step is excluded from every pool worker's dispatch + /// list (all workers get `Skip`). The dedicated-thread extraction is L2.3; + /// here we only prove the pool never sees it. + #[test] + fn detached_step_is_skipped_on_all_workers() { + let steps: Vec> = vec![Box::new(TypedStep::new(DetachedStep))]; + let owners = vec![None]; + let entries = build_worker_storage(steps, &owners, 4); + assert_eq!(entries.len(), 4); + for w in &entries { + assert_eq!(w.len(), 1); + assert!( + matches!(w[0], WorkerStepEntry::Skip), + "Detached step must be Skip on every pool worker" + ); + } + } + + #[test] + fn parallel_step_yields_owned_per_worker() { + let steps: Vec> = vec![Box::new(TypedStep::new(ParallelStep))]; + let owners = vec![None]; + let entries = build_worker_storage(steps, &owners, 4); + assert_eq!(entries.len(), 4); + for w in &entries { + assert_eq!(w.len(), 1); + assert!(matches!(w[0], WorkerStepEntry::Owned { .. })); + } + } + + #[test] + fn serial_step_yields_shared_arc_for_all_workers() { + let steps: Vec> = vec![Box::new(TypedStep::new(SerialStep))]; + let owners = vec![None]; + let entries = build_worker_storage(steps, &owners, 3); + for w in &entries { + assert!(matches!(w[0], WorkerStepEntry::Shared { .. })); + } + let arc0 = if let WorkerStepEntry::Shared { step, .. } = &entries[0][0] { + Arc::clone(step) + } else { + panic!() + }; + let arc1 = if let WorkerStepEntry::Shared { step, .. } = &entries[1][0] { + Arc::clone(step) + } else { + panic!() + }; + assert!(Arc::ptr_eq(&arc0, &arc1)); + } + + #[test] + fn exclusive_step_owner_gets_exclusive_others_skip() { + let steps: Vec> = vec![Box::new(TypedStep::new(ExclusiveStep))]; + let owners = vec![Some(2)]; + let entries = build_worker_storage(steps, &owners, 4); + assert!(matches!(entries[0][0], WorkerStepEntry::Skip)); + assert!(matches!(entries[1][0], WorkerStepEntry::Skip)); + assert!(matches!(entries[2][0], WorkerStepEntry::Exclusive { .. })); + assert!(matches!(entries[3][0], WorkerStepEntry::Skip)); + } + + struct AffinitySerialStep(crate::step::Affinity); + impl Step for AffinitySerialStep { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + profile_for("AffinitySer", StepKind::Serial, false) + } + fn affinity(&self) -> crate::step::Affinity { + self.0 + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + #[rstest] + #[case::reader(Affinity::Reader, 0)] + #[case::writer(Affinity::Writer, 2)] // last of 3 workers + #[case::worker_idx(Affinity::Worker(1), 1)] + fn serial_affinity_eligible_only_for_target_worker( + #[case] affinity: Affinity, + #[case] eligible: usize, + ) { + // A Serial step's affinity makes exactly one worker eligible (`Shared`); + // every other worker `Skip`s it. Reader→0, Writer→last, Worker(i)→i. + let steps: Vec> = + vec![Box::new(TypedStep::new(AffinitySerialStep(affinity)))]; + let owners = vec![None]; + let entries = build_worker_storage(steps, &owners, 3); + for (worker, entry) in entries.iter().enumerate() { + if worker == eligible { + assert!( + matches!(entry[0], WorkerStepEntry::Shared { .. }), + "worker {worker} should be eligible (Shared) for {affinity:?}" + ); + } else { + assert!( + matches!(entry[0], WorkerStepEntry::Skip), + "worker {worker} should Skip for {affinity:?}" + ); + } + } + } + + #[test] + #[should_panic(expected = "Serial step affinity")] + fn serial_out_of_range_worker_panics_in_storage() { + let steps: Vec> = + vec![Box::new(TypedStep::new(AffinitySerialStep(Affinity::Worker(99))))]; + let owners = vec![None]; + // build_worker_storage checks affinity in range via assert!; should panic. + let _ = build_worker_storage(steps, &owners, 3); + } + + /// The Exclusive arm's owner-range check is an always-on `assert!` (not a + /// `debug_assert!`) because an out-of-range owner would otherwise surface as + /// an opaque `entries[owner]` slice-index panic, and `assign_sticky_owners` + /// explicitly tolerates one. Without this case only the Serial arm was + /// covered, so a refactor could weaken this one to `debug_assert!` and CI + /// would stay green. + #[test] + #[should_panic(expected = "Exclusive step owner")] + fn exclusive_out_of_range_owner_panics_in_storage() { + let steps: Vec> = vec![Box::new(TypedStep::new(ExclusiveStep))]; + let owners = vec![Some(99)]; + let _ = build_worker_storage(steps, &owners, 3); + } + + #[test] + fn mixed_chain_assigns_correctly() { + let steps: Vec> = vec![ + Box::new(TypedStep::new(ExclusiveStep)), + Box::new(TypedStep::new(ParallelStep)), + Box::new(TypedStep::new(SerialStep)), + Box::new(TypedStep::new(ExclusiveStep)), + ]; + let owners = vec![Some(0), None, None, Some(1)]; + let entries = build_worker_storage(steps, &owners, 4); + + assert!(matches!(entries[0][0], WorkerStepEntry::Exclusive { .. })); + assert!(matches!(entries[1][0], WorkerStepEntry::Skip)); + + for w in &entries { + assert!(matches!(w[1], WorkerStepEntry::Owned { .. })); + } + + for w in &entries { + assert!(matches!(w[2], WorkerStepEntry::Shared { .. })); + } + + assert!(matches!(entries[1][3], WorkerStepEntry::Exclusive { .. })); + assert!(matches!(entries[0][3], WorkerStepEntry::Skip)); + } +} diff --git a/crates/fgumi-pipeline-core/src/runtime/worker_core.rs b/crates/fgumi-pipeline-core/src/runtime/worker_core.rs new file mode 100644 index 000000000..2274419a0 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/runtime/worker_core.rs @@ -0,0 +1,247 @@ +//! `WorkerCore`: per-thread state carried by the worker loop. + +use std::time::Duration; + +use crate::topology::StepIdx; + +// Pool worker idle bounds: a pool worker is never unparked, so it sleeps and can +// ramp to a coarse cap without hurting wake latency. +const SLEEP_INITIAL_US: u64 = 1; +const SLEEP_MAX_US: u64 = 50_000; // 50 milliseconds + +// Dedicated-driver idle bounds: a driver drives a small step subset off the pool +// and must stay responsive to its peers (the pool filling/draining its edges), so +// it parks with a tight cap. `park_timeout` also lets a producer holding the +// thread handle unpark it early; the cap is the bounded fallback either way. +const PARK_INITIAL_US: u64 = 10; +const PARK_MAX_US: u64 = 500; + +/// How a worker idles on a no-progress tick. Selected per thread at construction: +/// the N-worker pool uses [`Sleep`](BackoffPolicy::Sleep); each dedicated driver +/// thread (the unified "1-thread pool") uses [`Park`](BackoffPolicy::Park). This +/// is the *only* behavioral difference between a pool worker and a driver — both +/// run the same [`run_worker_loop`](crate::runtime::run_worker_loop). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BackoffPolicy { + /// `thread::sleep`, ramp 1µs→50ms. For pool workers (never unparked). + Sleep, + /// `thread::park_timeout`, ramp 10µs→500µs. For dedicated driver threads: a + /// producer *may* unpark early, and the tight cap bounds wake latency when it + /// does not. With no unparker this behaves as a bounded sleep. + Park, +} + +impl BackoffPolicy { + #[inline] + fn initial_us(self) -> u64 { + match self { + Self::Sleep => SLEEP_INITIAL_US, + Self::Park => PARK_INITIAL_US, + } + } + + #[inline] + fn max_us(self) -> u64 { + match self { + Self::Sleep => SLEEP_MAX_US, + Self::Park => PARK_MAX_US, + } + } +} + +/// Whether a `run_worker_loop` thread is an N-pool worker or a dedicated driver +/// (the unified "1-thread pool" for a set of off-pool steps). Controls only how +/// the thread's busy/idle time is attributed in `--pipeline-stats`, always +/// excluded from the pool% so the "N + 2" split stays visible: +/// - pool threads sum their whole-pass busy/idle into the N-worker utilisation +/// line, by `thread_id`; +/// - driver threads record each grouped step's own busy on the off-pool +/// "detached" line (by that step's index — so a multi-step `Shared` group +/// shows each member's real time, not the whole thread's under one name), and +/// attribute the thread-level idle/park to the group's `primary_step`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WorkerRole { + /// N-worker pool thread; busy/idle keyed by `thread_id`. + Pool, + /// Dedicated driver thread. Per-step busy is recorded by each step's own + /// index in `dispatch_one_step`; thread-level idle/park is keyed to + /// `primary_step` (the group's representative). + Driver { primary_step: StepIdx }, +} + +pub struct WorkerCore { + /// `0..n_workers` for pool threads; unused (0) for driver threads. + pub thread_id: usize, + /// If this worker owns an `Exclusive` step, that step's index. + /// Used by the runtime to skip `WorkerStepEntry::Skip` placeholders + /// for Exclusive steps owned by other workers. + pub exclusive_owner: Option, + /// If this worker is the sole eligible dispatcher for a `sticky` step + /// (either an `Exclusive sticky` step it owns, or a `Serial + sticky` + /// step whose `Affinity` targets this worker), the step's index. + /// The driver drives this step in a tight inner loop until it returns + /// `NoProgress` / `Contention` / `Finished`, then yields to round- + /// robin. Mirrors the legacy pipeline's sticky read. + pub sticky_owner: Option, + /// Pool worker vs dedicated driver. Determines both the idle backoff policy + /// (`Pool → Sleep`, `Driver → Park`, via [`Self::policy`]) and how the + /// thread's aggregate busy/idle time is attributed in `--pipeline-stats`. + role: WorkerRole, + /// Backoff duration in microseconds. Doubled on no-progress; reset on progress. + /// Bounds come from `role`'s policy. + backoff_us: u64, +} + +impl WorkerCore { + /// A pool worker (`Sleep` backoff). Existing callers are unchanged. + #[must_use] + pub fn new( + thread_id: usize, + exclusive_owner: Option, + sticky_owner: Option, + ) -> Self { + Self { + thread_id, + exclusive_owner, + sticky_owner, + role: WorkerRole::Pool, + backoff_us: BackoffPolicy::Sleep.initial_us(), + } + } + + /// A dedicated driver thread (the unified "1-thread pool"): `Driver` role + + /// `Park` backoff (derived from the role). Drives a set of `Owned` steps off + /// the pool; its aggregate busy/idle is attributed to `primary_step` on the + /// off-pool detached line. `thread_id` is unused for drivers; it never owns + /// Exclusive or sticky steps. + #[must_use] + pub fn driver(primary_step: StepIdx) -> Self { + let mut worker = Self::new(0, None, None); + worker.role = WorkerRole::Driver { primary_step }; + worker.backoff_us = worker.policy().initial_us(); + worker + } + + /// This thread's pool/driver role (drives stats attribution in the loop). + #[must_use] + pub fn role(&self) -> WorkerRole { + self.role + } + + /// The idle backoff policy implied by this thread's role: pool workers + /// `Sleep` (never unparked), driver threads `Park`. + #[must_use] + pub fn policy(&self) -> BackoffPolicy { + match self.role { + WorkerRole::Pool => BackoffPolicy::Sleep, + WorkerRole::Driver { .. } => BackoffPolicy::Park, + } + } + + pub fn reset_backoff(&mut self) { + self.backoff_us = self.policy().initial_us(); + } + + pub fn sleep_backoff(&self) { + let dur = Duration::from_micros(self.backoff_us); + match self.policy() { + BackoffPolicy::Sleep => std::thread::sleep(dur), + BackoffPolicy::Park => std::thread::park_timeout(dur), + } + } + + pub fn increase_backoff(&mut self) { + self.backoff_us = self.backoff_us.saturating_mul(2).min(self.policy().max_us()); + } + + #[cfg(test)] + fn current_backoff_us(&self) -> u64 { + self.backoff_us + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[test] + fn fresh_backoff_is_initial() { + let w = WorkerCore::new(0, None, None); + assert_eq!(w.current_backoff_us(), SLEEP_INITIAL_US); + } + + // One doubling-then-cap sequence, exercised for both policies: the pool's + // `Sleep` bounds and a driver's tight `Park` bounds. Asserting the whole + // sequence (not just the final cap) would fail a buggy `increase_backoff` + // that jumped straight to the cap or incremented by a constant; the closing + // reset confirms it returns to *this* policy's initial, not the other's. + #[rstest] + #[case::sleep( + WorkerCore::new(0, None, None), + SLEEP_INITIAL_US, + SLEEP_MAX_US, + BackoffPolicy::Sleep + )] + #[case::park(WorkerCore::driver(StepIdx(0)), PARK_INITIAL_US, PARK_MAX_US, BackoffPolicy::Park)] + fn backoff_doubles_then_caps( + #[case] mut w: WorkerCore, + #[case] initial: u64, + #[case] max: u64, + #[case] policy: BackoffPolicy, + ) { + assert_eq!(w.policy(), policy); + assert_eq!(w.current_backoff_us(), initial); + let mut expected = initial; + for _ in 0..20 { + w.increase_backoff(); + expected = (expected * 2).min(max); + assert_eq!( + w.current_backoff_us(), + expected, + "backoff must double each step until it saturates at the cap" + ); + } + assert_eq!(w.current_backoff_us(), max); + w.reset_backoff(); + assert_eq!(w.current_backoff_us(), initial); + } + + #[test] + fn reset_after_progress() { + let mut w = WorkerCore::new(0, None, None); + for _ in 0..5 { + w.increase_backoff(); + } + w.reset_backoff(); + assert_eq!(w.current_backoff_us(), SLEEP_INITIAL_US); + } + + #[test] + fn sleep_and_park_have_distinct_caps() { + // Guard against the two policies drifting to the same bound. + assert_ne!(SLEEP_MAX_US, PARK_MAX_US); + assert_ne!(SLEEP_INITIAL_US, PARK_INITIAL_US); + // Policy is derived from role: pool worker → Sleep, driver → Park. + assert_eq!(WorkerCore::new(0, None, None).policy(), BackoffPolicy::Sleep); + assert_eq!(WorkerCore::driver(StepIdx(0)).policy(), BackoffPolicy::Park); + } + + #[test] + fn worker_with_exclusive_role() { + let w = WorkerCore::new(2, Some(StepIdx(7)), Some(StepIdx(7))); + assert_eq!(w.thread_id, 2); + assert_eq!(w.exclusive_owner, Some(StepIdx(7))); + assert_eq!(w.sticky_owner, Some(StepIdx(7))); + } + + #[test] + fn worker_with_sticky_serial_owner_only() { + // E.g., the Serial+Affinity::Reader source — sticky_owner set, + // exclusive_owner None. + let w = WorkerCore::new(0, None, Some(StepIdx(0))); + assert_eq!(w.exclusive_owner, None); + assert_eq!(w.sticky_owner, Some(StepIdx(0))); + } +} diff --git a/crates/fgumi-pipeline-core/src/signal.rs b/crates/fgumi-pipeline-core/src/signal.rs new file mode 100644 index 000000000..eda14a32e --- /dev/null +++ b/crates/fgumi-pipeline-core/src/signal.rs @@ -0,0 +1,429 @@ +//! Pipeline-wide error broadcast and cancellation, folded into a single +//! shared atomic so the worker loop pays one relaxed load per iteration. + +use std::io; +use std::sync::Arc; +use std::sync::OnceLock; +use std::sync::atomic::{AtomicU8, Ordering as AtomicOrdering}; + +const STATE_OK: u8 = 0; +const STATE_CANCELLED: u8 = 1; +const STATE_ERROR: u8 = 2; + +/// Errors a pipeline run can return to its caller. +#[derive(Debug)] +pub enum PipelineError { + /// A step's `try_run` returned `Err`. Carries the originating step's + /// name and the underlying I/O error. + Io { step: &'static str, source: io::Error }, + /// `cancel_handle().cancel()` was called from outside the pipeline. + Cancelled, + /// The chain has more `Exclusive` steps than the configured thread count. + NotEnoughThreads { required: usize, available: usize }, + /// The deadlock monitor observed no global progress for `stalled_secs` + /// while work was still stuck in queues/reorder buffers — a wedge. The + /// pipeline is failed fast rather than left to hang forever. + TimedOut { stalled_secs: u64 }, + /// The pipeline was built with the deadlock monitor armed + /// (`deadlock_timeout_secs > 0`), but `step` declares a non-`ByteBounded` + /// output transport (`spec`, e.g. `CountBounded`/`Unbounded`). The + /// `in_flight_bytes` probe cannot see a wedge on such an edge, so fail-fast + /// would be silently disabled there. The pipeline is rejected at startup + /// rather than allowed to hang. This is a chain-construction error, not a + /// runtime condition: production chains wire only `ByteBounded` transports. + MonitorBlindTransport { step: &'static str, spec: String }, +} + +impl PipelineError { + /// Reconstruct an owned copy of this error. + /// + /// `PipelineError` is not `Clone` because its `Io` variant holds an + /// [`io::Error`], which is not `Clone`; this rebuilds that source from its + /// `kind()` + display string. Used to turn the borrowed `signal.outcome()` + /// into the owned error returned from the run drivers (`Pipeline::run` and + /// the fused single-thread driver), which previously open-coded the same + /// per-variant rematch in two places. + pub(crate) fn reconstruct(&self) -> PipelineError { + match self { + Self::Cancelled => Self::Cancelled, + Self::Io { step, source } => { + Self::Io { step, source: io::Error::new(source.kind(), format!("{source}")) } + } + Self::NotEnoughThreads { required, available } => { + Self::NotEnoughThreads { required: *required, available: *available } + } + Self::TimedOut { stalled_secs } => Self::TimedOut { stalled_secs: *stalled_secs }, + Self::MonitorBlindTransport { step, spec } => { + Self::MonitorBlindTransport { step, spec: spec.clone() } + } + } + } +} + +impl std::fmt::Display for PipelineError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Io { step, source } => write!(f, "step {step:?} failed: {source}"), + Self::Cancelled => write!(f, "pipeline cancelled"), + Self::NotEnoughThreads { required, available } => write!( + f, + "pipeline requires {required} threads for Exclusive steps, only {available} available" + ), + Self::TimedOut { stalled_secs } => write!( + f, + "pipeline deadlock detected: no progress for {stalled_secs}s with work still in flight" + ), + Self::MonitorBlindTransport { step, spec } => write!( + f, + "deadlock monitor is armed but step {step:?} declares a {spec} output transport, \ + which is invisible to the in_flight_bytes probe — a wedge on that edge would \ + silently disable fail-fast. Production transports must be QueueSpec::ByteBounded." + ), + } + } +} + +impl std::error::Error for PipelineError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Self::Io { source, .. } => Some(source), + _ => None, + } + } +} + +/// Single shared per-pipeline state covering both error broadcast and cancel. +/// Workers check `is_done()` once per iteration — one relaxed atomic load. +#[derive(Default)] +pub struct PipelineSignal { + state: AtomicU8, + payload: OnceLock, +} + +impl PipelineSignal { + #[must_use] + pub fn new() -> Arc { + Arc::new(Self::default()) + } + + /// Hot-path: workers call this once per loop iteration. + /// + /// Uses `Relaxed` because the only thing a worker does on `true` is stop + /// polling — workers never read the payload. The payload is read only by + /// the run drivers via `to_result`, after every internal writer + /// (workers, deadlock monitor) has been joined, so the join supplies the + /// happens-before that makes a recorded error's payload visible. + /// + /// The one writer not covered by a join is an external + /// [`CancelHandle::cancel`], which races the driver's terminal read with no + /// synchronizing edge; `to_result` derives `Cancelled` from the + /// (coherence-visible) state rather than the payload to close that window. + #[inline] + #[must_use] + pub fn is_done(&self) -> bool { + self.state.load(AtomicOrdering::Relaxed) != STATE_OK + } + + #[inline] + #[must_use] + pub fn is_cancelled(&self) -> bool { + self.state.load(AtomicOrdering::Relaxed) == STATE_CANCELLED + } + + /// First writer wins; later writers are silently dropped (the + /// `compare_exchange` rejects the state transition and the + /// `OnceLock::set` rejects the payload write). + pub fn record_error(&self, err: PipelineError) { + // Set the payload only if we won the state CAS, so the error surfaced by + // `outcome()` always matches the writer that transitioned the state. + // Otherwise writer A could win the state CAS while writer B wins the + // `OnceLock::set`, leaving `state == ERROR` reporting B's error — + // violating the documented "first writer wins" contract. + if self + .state + .compare_exchange( + STATE_OK, + STATE_ERROR, + AtomicOrdering::Release, + AtomicOrdering::Relaxed, + ) + .is_ok() + { + let _ = self.payload.set(err); + } + } + + /// First writer wins; later writers (including a `record_error` after + /// `cancel`) are silently dropped. + pub fn cancel(&self) { + // Only publish the payload if THIS call won the state transition. + // Setting it unconditionally races a concurrent `record_error`: if + // that call wins the CAS (state → ERROR) but has not yet set its + // payload, an unconditional `payload.set(Cancelled)` here can win the + // `OnceLock` and leave `state == ERROR` while `outcome()` reports + // `Cancelled` — an inconsistent state/payload pair. Guarding on the CAS + // (as `record_error` does) keeps the two consistent: whoever wins the + // state transition is the one that sets the payload. + // (See the loom NOTE at the end of the tests module for why this + // race-guard is not exercised by a unit test.) + if self + .state + .compare_exchange( + STATE_OK, + STATE_CANCELLED, + AtomicOrdering::Release, + AtomicOrdering::Relaxed, + ) + .is_ok() + { + let _ = self.payload.set(PipelineError::Cancelled); + } + } + + /// Read the recorded error payload, if any. + /// + /// The `OnceLock::get` acquire fence pairs with the `OnceLock::set` release + /// on the writer side, so a reader synchronized with the writer (via the + /// worker/monitor join in the run drivers) sees a recorded error's payload. + /// + /// Beware the gap this does NOT cover: `is_done() == true` does not imply + /// `outcome().is_some()`. `cancel()`/`record_error()` publish the terminal + /// `state` (CAS) before `payload.set()`, so a reader can observe the + /// terminal state while the payload `OnceLock` is still empty — most + /// reachably for an external `CancelHandle::cancel`, whose writer is never + /// joined against the driver's read. Map an outcome to a run result through + /// `to_result`, which handles that window, not by branching on + /// `outcome()` directly. + #[must_use] + pub fn outcome(&self) -> Option<&PipelineError> { + self.payload.get() + } + + /// Map the recorded outcome to a run `Result`, as returned by the run + /// drivers ([`crate::builder::Pipeline::run`] and the fused single-thread + /// driver). + /// + /// A recorded error's payload is always visible here: every `record_error` + /// writer (workers, the deadlock monitor, the fused driver itself) is + /// joined before this read, so the join orders its `payload.set()` ahead of + /// the read. + /// + /// An external [`CancelHandle::cancel`] is the exception — its writer is + /// never joined against this read, so there is no happens-before edge to + /// make its `OnceLock` payload visible, and it can leave `state == + /// CANCELLED` with `outcome() == None`. A naive `match outcome()` would then + /// map a genuinely cancelled run (workers observed `is_done()` and stopped + /// early) to `Ok(())`. `Cancelled` carries no payload data, so synthesize it + /// from the coherence-visible `state` via [`Self::is_cancelled`] instead. + pub(crate) fn to_result(&self) -> Result<(), PipelineError> { + match self.outcome() { + Some(err) => Err(err.reconstruct()), + None if self.is_cancelled() => Err(PipelineError::Cancelled), + // Terminal state with no payload published yet. `record_error` wins + // the state CAS *before* it runs `payload.set`, so there is a window + // in which `state == STATE_ERROR` and `outcome()` is still `None` — + // the same publish window `is_cancelled` rescues above, on the error + // path instead of the cancel path. Falling through to `Ok(())` here + // would report a failed run as successful. + // + // The argument that every `record_error` writer is joined before this + // read is a property of the drivers (pool workers AND the detached + // driver threads, on every early-return path), not of this type. Derive + // the answer from the terminal state instead, so the guarantee does not + // depend on that staying true. A successful run never reaches this arm: + // `is_done()` is `state != STATE_OK`, so it is false unless something + // already transitioned the state away from OK. + None if self.is_done() => Err(PipelineError::Io { + step: "unknown", + source: io::Error::other( + "pipeline failed but its error payload was never published; the \ + recording thread was still in the publish window when the run \ + result was read", + ), + }), + None => Ok(()), + } + } +} + +/// External handle for cancelling an in-flight pipeline. +#[derive(Clone)] +pub struct CancelHandle { + signal: Arc, +} + +impl CancelHandle { + /// Construct a `CancelHandle` from a shared signal. Used by + /// `Pipeline::cancel_handle` (see `builder.rs`) to hand a handle to the + /// caller of a built pipeline, before `Pipeline::run` kicks off worker + /// threads. + pub(crate) fn from_signal(signal: Arc) -> Self { + Self { signal } + } + + pub fn cancel(&self) { + self.signal.cancel(); + } + + #[must_use] + pub fn is_cancelled(&self) -> bool { + self.signal.is_cancelled() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn fresh_signal_is_not_done() { + let signal = PipelineSignal::new(); + assert!(!signal.is_done()); + assert!(!signal.is_cancelled()); + } + + #[test] + fn record_error_marks_done() { + let signal = PipelineSignal::new(); + signal.record_error(PipelineError::Io { + step: "test_step", + source: io::Error::other("boom"), + }); + assert!(signal.is_done()); + assert!(!signal.is_cancelled()); + assert!(matches!(signal.outcome(), Some(PipelineError::Io { step: "test_step", .. }))); + } + + #[test] + fn cancel_marks_done_and_cancelled() { + let signal = PipelineSignal::new(); + signal.cancel(); + assert!(signal.is_done()); + assert!(signal.is_cancelled()); + assert!(matches!(signal.outcome(), Some(PipelineError::Cancelled))); + } + + #[test] + fn first_recorded_error_wins() { + let signal = PipelineSignal::new(); + signal.record_error(PipelineError::Io { step: "first", source: io::Error::other("first") }); + signal + .record_error(PipelineError::Io { step: "second", source: io::Error::other("second") }); + match signal.outcome().unwrap() { + PipelineError::Io { step, .. } => assert_eq!(*step, "first"), + other => panic!("expected Io, got {other:?}"), + } + } + + #[test] + fn cancel_handle_propagates() { + let signal = PipelineSignal::new(); + let handle = CancelHandle::from_signal(Arc::clone(&signal)); + assert!(!handle.is_cancelled()); + handle.cancel(); + assert!(handle.is_cancelled()); + assert!(signal.is_cancelled()); + } + + #[test] + fn reconstruct_preserves_every_variant() { + // `reconstruct` rebuilds an owned error from a borrow (PipelineError + // isn't Clone). Each variant must round-trip; the `Io` source is + // rebuilt from kind + display, the rest are copied verbatim. + let io = PipelineError::Io { step: "s", source: io::Error::other("boom") }.reconstruct(); + match io { + PipelineError::Io { step, source } => { + assert_eq!(step, "s"); + assert_eq!(source.kind(), io::ErrorKind::Other); + assert_eq!(source.to_string(), "boom"); + } + other => panic!("expected Io, got {other:?}"), + } + + assert!(matches!(PipelineError::Cancelled.reconstruct(), PipelineError::Cancelled)); + + let net = PipelineError::NotEnoughThreads { required: 4, available: 2 }.reconstruct(); + assert!(matches!(net, PipelineError::NotEnoughThreads { required: 4, available: 2 })); + + let to = PipelineError::TimedOut { stalled_secs: 60 }.reconstruct(); + assert!(matches!(to, PipelineError::TimedOut { stalled_secs: 60 })); + + let blind = PipelineError::MonitorBlindTransport { + step: "s", + spec: "QueueSpec::Unbounded".to_string(), + } + .reconstruct(); + assert!( + matches!(blind, PipelineError::MonitorBlindTransport { step: "s", spec } if spec == "QueueSpec::Unbounded") + ); + } + + #[test] + fn to_result_maps_clean_error_and_cancel() { + let clean = PipelineSignal::new(); + assert!(clean.to_result().is_ok()); + + let errored = PipelineSignal::new(); + errored.record_error(PipelineError::Io { step: "s", source: io::Error::other("boom") }); + assert!(matches!(errored.to_result(), Err(PipelineError::Io { step: "s", .. }))); + + let cancelled = PipelineSignal::new(); + cancelled.cancel(); + assert!(matches!(cancelled.to_result(), Err(PipelineError::Cancelled))); + } + + #[test] + fn to_result_reports_cancel_when_payload_not_yet_published() { + // Reproduce the external-cancel window: a `CancelHandle::cancel` on an + // un-joined thread has published the terminal `state` (the CAS) but has + // not yet run `payload.set(Cancelled)`. A driver reading the outcome in + // that window sees `is_done() == true` but `outcome() == None`. Branching + // on `outcome()` alone would map this genuinely cancelled run to `Ok(())`; + // `to_result` must instead synthesize `Cancelled` from the state. + let signal = PipelineSignal::new(); + signal.state.store(STATE_CANCELLED, AtomicOrdering::Release); + assert!(signal.is_done(), "terminal state must read as done"); + assert!(signal.is_cancelled()); + assert!(signal.outcome().is_none(), "payload not set in this window"); + assert!( + matches!(signal.to_result(), Err(PipelineError::Cancelled)), + "a cancel observed before its payload is published must not map to Ok" + ); + } + + /// The sibling of the test above, on the error path. `record_error` wins the + /// state CAS before it runs `payload.set`, so `state == STATE_ERROR` with + /// `outcome() == None` is reachable. Branching on `outcome()` and + /// `is_cancelled()` alone mapped that to `Ok(())` — reporting a failed run as + /// successful, the worst possible direction for this to be wrong in. + #[test] + fn to_result_reports_failure_when_error_payload_not_yet_published() { + let signal = PipelineSignal::new(); + signal.state.store(STATE_ERROR, AtomicOrdering::Release); + assert!(signal.is_done(), "terminal state must read as done"); + assert!(!signal.is_cancelled(), "this is the error path, not the cancel path"); + assert!(signal.outcome().is_none(), "payload not set in this window"); + assert!( + matches!(signal.to_result(), Err(PipelineError::Io { step: "unknown", .. })), + "an error observed before its payload is published must not map to Ok" + ); + } + + /// The guard above must not fire for a run that never failed: `is_done()` is + /// `state != STATE_OK`, so an untouched signal still reports success. + #[test] + fn to_result_is_ok_for_an_untouched_signal() { + assert!(PipelineSignal::new().to_result().is_ok()); + } + + // NOTE: the `cancel` vs `record_error` state/payload-consistency fix (guarding + // `payload.set` behind the CAS) is deliberately NOT covered by a unit test. + // The bug is a true concurrent interleaving — `cancel` losing the CAS in the + // window between `record_error`'s CAS-win and its own `payload.set` — and + // `OnceLock` masks the inconsistency in every *sequential* ordering, so no + // single-threaded test can distinguish fixed from buggy. A spawned-thread + // stress loop does not reliably hit the two-instruction window either (a + // 2000-trial loop passed against the buggy code). Deterministically forcing + // the interleaving would require `loom`; until that dependency is justified, + // the fix rests on the inline argument at the `cancel` call site. +} diff --git a/crates/fgumi-pipeline-core/src/step.rs b/crates/fgumi-pipeline-core/src/step.rs new file mode 100644 index 000000000..bf8577e10 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/step.rs @@ -0,0 +1,616 @@ +//! Step trait, profile types, and context. +//! +//! Completion contract (full detail on [`StepOutcome::Finished`]): +//! +//! - Every step — source, mid, or sink — returns `StepOutcome::Finished` from +//! `try_run` once its input edges are drained (empty + upstream closed) and +//! it holds no buffered output. The framework then closes the step's output +//! edges (`mark_outputs_drained`) and drops it from the worklist. +//! - `try_run` returns `Progress` when it pushed or held an item, `NoProgress` +//! when there's nothing useful to do this call (input empty but not drained, +//! no held items), and `Contention` when a Serial-step mutex is held by +//! another worker; the scheduler reroutes. +//! +//! **Last-worker barrier for `Parallel` steps.** A `Parallel` step has N +//! per-worker `Clone`s sharing one output queue and a single drained input +//! edge. When the input drains, every clone returns `Finished`, but only the +//! LAST clone to finish (the one that takes the per-step `StepDrainCounter` to +//! 0) closes the shared output queue. Otherwise a clone could `mark_drained` +//! while a sibling is still pushing — a `try_push`-after-`mark_drained` +//! violation, which `ItemQueue::try_push` rejects with a panic in every build. +//! The gate lives in the driver's `dispatch_one_step`. + +/// Concurrency profile. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StepKind { + /// Any number of workers may be inside `try_run` concurrently. Each + /// worker holds its own `Clone` of the step. + Parallel, + /// At most one worker at a time. Framework holds a per-step mutex; any + /// worker can acquire it; one instance shared. + Serial, + /// Exactly one worker (the owner) ever runs this step. Framework + /// assigns owners at run start in chain declaration order. Other + /// workers skip the step entirely. + Exclusive, + /// Runs on a **dedicated OS thread**, spawned at run start alongside the + /// deadlock-monitor / queue-rebalancer — never dispatched by the + /// work-stealing pool. The dedicated thread drives the step with the + /// *same* `run_worker_loop` the pool uses (via + /// `runtime::detached::run_detached_driver`): a `WorkerCore::driver` (Park + /// backoff) over an `Owned`-for-its-group row. The step body still uses + /// non-blocking `try_pop`/`try_push` and never blocks inside `try_run`, so + /// it consumes no pool worker slot. A driver is literally a 1-thread pool. + /// + /// This mirrors the legacy sort's "N + 2" threading: N pool workers do + /// the parallel (compression-bound) work while off-pool driver threads do + /// the serial coordination + I/O, keeping the pool saturated. Several + /// detached steps sharing a [`DetachedGroup::Shared`] label are driven by + /// ONE thread; [`DetachedGroup::PerStep`] (the default) keeps one thread + /// per step. + /// + /// **Opt-in.** No existing step declares `Detached` unless the sort chain + /// wires it. A single shared instance (like `Serial`/`Exclusive`, never + /// `new_worker_copy`'d) runs on the driver thread; every pool worker gets a + /// `Skip` entry for it. + Detached, +} + +/// Optional scheduling hint for `Serial` steps. Restricts which worker(s) +/// are eligible to attempt the step's mutex on each round-robin pass — +/// non-eligible workers `Skip` the step entirely, eliminating the +/// `try_lock` thrash that pure `Serial` exhibits at high thread counts. +/// +/// The framework's per-step mutex is still in place (so a step author +/// can rely on single-thread-at-a-time semantics), but with a non-`None` +/// affinity only the hinted worker(s) ever acquire it. +/// +/// Mirrors the legacy framework's per-thread `exclusive_step_owned` +/// mapping, where T0 is the reader, T(N-1) is the writer, and interior +/// exclusive steps fan out from both ends. +/// +/// Default `None` keeps the existing pure-mutex-shared `Serial` behavior. +/// Ignored for `Parallel`, `Exclusive`, and `Detached` kinds: +/// `build_worker_storage` reads affinity only in its `StepKind::Serial` arm, so +/// a `Detached` step declaring `Affinity::Worker(99)` is silently ignored and — +/// unlike the `Serial` case — never range-checked at run start. A `Detached` +/// step runs on its own driver thread, so no pool worker gates on it. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Affinity { + /// Any worker may attempt this Serial step. (Default; current behavior.) + None, + /// Restrict attempts to worker 0. Other workers `Skip` this step in + /// dispatch — no `try_lock` thrash. Use for I/O sources where keeping + /// reads on a single thread improves the kernel readahead pattern. + Reader, + /// Restrict attempts to worker `N - 1` (the last worker). Other + /// workers `Skip`. Use for I/O sinks where the writer's `BufWriter` + /// benefits from thread locality. + Writer, + /// Restrict attempts to a specific worker index. Out-of-range values + /// (`>= n_threads`) trigger an assertion failure (panic) at run start + /// via an always-on `assert!` in `build_worker_storage`. + Worker(usize), +} + +impl Affinity { + /// The single worker this affinity pins a step to, or `None` when the step + /// is not pinned ([`Affinity::None`] — every worker is eligible). + /// + /// The sole source of truth for the affinity → worker mapping, shared by + /// [`Self::eligible`] (which gates dispatch) and + /// `pool::assign_sticky_owners` (which picks each worker's sticky-driven + /// step). Keeping one mapping means a future variant cannot make those two + /// disagree and assign a sticky owner to a worker that would `Skip` the step. + /// + /// The returned index may be `>= n_workers` for an out-of-range + /// [`Affinity::Worker`]; callers must range-check before indexing. + #[must_use] + pub fn target_worker(self, n_workers: usize) -> Option { + match self { + Self::None => None, + Self::Reader => Some(0), + Self::Writer => Some(n_workers.saturating_sub(1)), + Self::Worker(idx) => Some(idx), + } + } + + /// Returns `true` if `worker_id` is eligible to attempt this Serial + /// step under this affinity hint. + #[must_use] + pub fn eligible(self, worker_id: usize, n_workers: usize) -> bool { + match self.target_worker(n_workers) { + // Unpinned: every worker may attempt the step. + None => true, + // `Writer` maps to `n_workers - 1`, so a saturating target of 0 with + // `n_workers == 0` has no eligible worker — and there are none to ask. + Some(target) => worker_id == target, + } + } +} + +/// Which dedicated driver thread a [`StepKind::Detached`] step runs on. +/// +/// Detached steps run off the work-stealing pool on dedicated OS threads. By +/// default each detached step gets its own thread ([`PerStep`](Self::PerStep)) +/// — the legacy "one dedicated thread per detached step" behavior. Steps that +/// declare the same [`Shared`](Self::Shared) label instead share ONE dedicated +/// driver thread that round-robins them through the same `run_worker_loop` the +/// pool uses (the unified "1-thread pool"). This realizes the true N+2 model: +/// the sort chain groups its serial coordination steps onto one driver and its +/// writers onto another, keeping the N pool workers on pure (de)compression. +/// +/// Ignored for non-`Detached` kinds. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DetachedGroup { + /// One dedicated thread for this step alone (default; legacy behavior). + PerStep, + /// Share one dedicated driver thread with every other detached step that + /// declares the same label. + Shared(&'static str), +} + +/// Static description of how a step is scheduled, plus per-output queue +/// configuration. +/// +/// `output_queues[i]` selects the transport-layer queue type for branch `i` +/// (count-bounded, byte-bounded, or unbounded). `branch_ordering[i]` selects +/// whether the framework inserts a `ReorderStage` in front of the consumer's +/// input handle (so consumers see items in producer-emitted ordinal order). +/// +/// Both vectors must have length equal to `S::Outputs::arity()`. +#[derive(Debug, Clone)] +pub struct StepProfile { + pub name: &'static str, + pub kind: StepKind, + pub sticky: bool, + pub output_queues: Vec, + pub branch_ordering: Vec, +} + +/// Outcome of a single `try_run` call. +/// +/// **`Finished` contract.** Any step — source, mid, or sink — returns +/// `Finished` once it will never push to its output again: all its input edges +/// are drained (empty + upstream closed) and it holds no buffered output. The +/// framework then closes the step's output edges (`mark_outputs_drained`) and +/// drops it from the worklist. The step must already have flushed every item +/// (via `try_run`'s flush-first path) before returning `Finished` — returning +/// it with work still buffered loses data. For a `Parallel` step the per-step +/// `StepDrainCounter` gates the output close so only the last clone to finish +/// closes the shared queue (a sibling could still be pushing). Sources are the +/// degenerate case: their (unit) input edge is drained from birth, so they +/// return `Finished` on EOF. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StepOutcome { + /// Step did useful work (pushed an item or held one for later). + Progress, + /// Step had nothing to do this call (input empty, no held work). + NoProgress, + /// Step's Serial-step mutex was contended; scheduler reroutes. + Contention, + /// The step has drained all its input and holds no buffered output — it + /// will never push again. The framework marks its output queues drained + /// (counter-gated for `Parallel`) and drops it. See the `Finished` + /// contract in the enum docs above. + Finished, +} + +/// Type-erased holder for the per-step outputs view, downcast to a typed +/// view (`SingleOutputsView`, `Tuple{2,3,4}OutputsView`, `UnitOutputsView`) +/// inside the `TypedStep` adapter. Defined here so the `Step` trait +/// declaration can reference it; concrete views live in `handles.rs`. +/// +/// Drain marking is per-branch via `BranchOutputHandle::mark_drained`, +/// exposed by typed `mark_all_drained()` methods on each `*OutputsView`. +/// `TypedStep` (Phase 1 Task 15) downcasts `inner` to the right typed +/// view and calls `mark_all_drained` from the worker loop's drain path. +pub struct OutputsViewAny { + pub(crate) inner: Box, +} + +use std::io; + +use super::item::HeapSize; +use super::outputs::StepOutputs; + +/// Handle to this step's input queue. +/// +/// Implemented by `OrderedQueueInputHandle` in `handles.rs`. Steps see +/// only the trait surface; the concrete handle is constructed by the +/// framework at chain build time. +pub trait InputHandle: Send + Sync { + /// Pop the next item, or `None` if the queue is empty right now. + /// Non-blocking. + fn pop(&self) -> Option; + + /// Returns `true` if upstream has marked the queue drained AND the + /// queue is currently empty. Once `is_drained()` returns true, no + /// further items will arrive. Used by mid-steps and sinks to detect + /// when to stop pulling. + fn is_drained(&self) -> bool; +} + +/// Handle to this step's output queues, shaped by `S::Outputs`. +/// +/// Concrete shape is the per-arity view: `SingleOutputsView` for +/// `Single` outputs, `Tuple2OutputsView` for `(A, B)`, etc. +/// Per-arity extension impls in `handles.rs` provide typed `push` methods +/// (e.g., `OutputHandles>::push(&self, item: T) -> Result<(), T>`). +pub struct OutputHandles { + pub(crate) inner: OutputsViewAny, + /// `PhantomData O>` (not `PhantomData`) so `OutputHandles` + /// is always `Send + Sync` regardless of `O` — required so the runtime + /// can box it as `Box` without forcing a `Sync` + /// bound on item types. + pub(crate) _phantom: std::marker::PhantomData O>, +} + +impl OutputHandles { + /// Wrap a type-erased outputs view into a typed `OutputHandles`. + /// + /// Called from the `ErasedStep::wrap_outputs_view` impls in `erased.rs` — + /// `TypedStep` for one-input steps and `TypedStep2` for the two-input + /// (zipper) shape — while `build_chain_contexts` assembles one context per + /// step. That happens on the thread calling `Pipeline::run`, *before* any + /// worker is spawned; workers never construct one, they borrow the box built + /// here (which is what `TypedStep::resolve_outputs`' cache relies on). + pub(crate) fn new(inner: OutputsViewAny) -> Self { + Self { inner, _phantom: std::marker::PhantomData } + } +} + +/// Step trait — every chain link implements this. +/// +/// `Send + 'static`: +/// - `Send`: instances move across threads (workers run on dedicated threads). +/// - `'static`: no borrowed references in step state. +/// +/// `Clone` is **not** a super-trait. Per-worker copies for `Parallel` steps +/// go through [`Step::new_worker_copy`] instead, which only `Parallel` +/// authors need to implement; `Serial`/`Exclusive` step authors inherit the +/// default panic and pay nothing. +/// +/// # Per-worker copy patterns by step kind +/// +/// **`Parallel` steps** must override [`Step::new_worker_copy`]. Each worker +/// thread calls it once during `build_worker_storage` to materialize its +/// private instance. The typical impl forwards to a regular `Clone` impl +/// (most parallel steps either `#[derive(Clone)]` or hand-roll a `Clone` +/// that resets per-worker scratch state). For example: +/// +/// ```ignore +/// #[derive(Clone)] +/// pub struct ParseBamRecords; +/// +/// impl Step for ParseBamRecords { +/// // ... +/// fn new_worker_copy(&self) -> Self { self.clone() } +/// } +/// ``` +/// +/// Closure-driven `Parallel` steps (`process(fn)`, `serialize(fn)`) wrap +/// their closure in `Arc` so the per-worker copy is one atomic +/// increment. +/// +/// **`Serial` and `Exclusive` steps** do not override +/// [`Step::new_worker_copy`]. The framework holds a single instance (behind +/// a `Mutex` for `Serial`, pinned to one worker for `Exclusive`) and never +/// asks for additional copies. The default impl panics with a descriptive +/// message — if it ever fires, that's a framework bug. +pub trait Step: Send + Sized + 'static { + type Input: Send + HeapSize + 'static; + type Outputs: StepOutputs; + + /// Static description of how this step is scheduled. + fn profile(&self) -> StepProfile; + + /// Optional scheduling hint for `Serial` kinds. Defaults to + /// `Affinity::None` (any worker may attempt the step). Override for + /// I/O sources/sinks that benefit from thread locality — see the + /// [`Affinity`] enum docs. + /// + /// Ignored for `Parallel` and `Exclusive` kinds (those have their + /// own per-worker dispatch model). + fn affinity(&self) -> Affinity { + Affinity::None + } + + /// Which dedicated driver thread this step runs on, for [`StepKind::Detached`] + /// steps. Defaults to [`DetachedGroup::PerStep`] (its own thread — legacy + /// behavior). Override to [`DetachedGroup::Shared`] to co-locate several + /// detached steps on one driver thread (the true N+2 model). Ignored for + /// non-`Detached` kinds. + fn detached_group(&self) -> DetachedGroup { + DetachedGroup::PerStep + } + + /// Step body. Pop from `ctx.input`, push to `ctx.outputs`. Returns + /// `Progress` / `NoProgress` / `Contention` / `Finished`. Errors propagate via `Err`. + /// + /// # Errors + /// + /// Returns the underlying I/O error from any failed read, write, or + /// codec operation inside the step body. The framework records the + /// first error via `PipelineSignal::record_error` and broadcasts to + /// other workers. + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result; + + /// Construct a fresh per-worker copy of this step. Only `Parallel` + /// steps need to override this — the framework calls it during + /// `build_worker_storage` to materialize one instance per worker. + /// `Serial`/`Exclusive` steps inherit the default panic; the framework + /// guarantees it is never invoked on them (one shared instance behind + /// a `Mutex` for `Serial`, pinned to a single owner worker for + /// `Exclusive`). + /// + /// # Panics + /// + /// Default impl panics with the step name + kind. Hitting it indicates + /// a framework bug (the runtime tried to clone a non-`Parallel` step). + #[must_use] + fn new_worker_copy(&self) -> Self { + let p = self.profile(); + panic!( + "Step::new_worker_copy invoked on '{}' (kind = {:?}); \ + only Parallel steps need to override this. The framework \ + never clones Serial/Exclusive steps — file a bug.", + p.name, p.kind + ); + } +} + +/// Context passed to `try_run`. +pub struct StepCtx<'a, S: Step> { + pub input: &'a dyn InputHandle, + pub outputs: &'a OutputHandles, +} + +/// Two-input variant of [`Step`]. Used by merge steps (zipper, +/// AAM-aligner-output + original-record-stream) that need to pop +/// from two upstream queues independently. +/// +/// Sits **alongside** [`Step`]; single-input steps are unaffected. +/// The framework's [`TypedStep2`](crate::erased::TypedStep2) adapter +/// (`crate::erased::TypedStep2`) is the +/// dual of [`TypedStep`](crate::erased::TypedStep), implementing the same [`ErasedStep`](crate::erased::ErasedStep) +/// contract so the runtime sees no difference between single- and +/// two-input steps at dispatch time. +/// +/// The two input branches can be the same type (`InputA == InputB` +/// — what zipper uses, both `RawRecord`) or different types +/// (heterogeneous — also supported, no special case in the +/// framework). +pub trait Step2: Send + Sized + 'static { + type InputA: Send + HeapSize + 'static; + type InputB: Send + HeapSize + 'static; + type Outputs: StepOutputs; + + fn profile(&self) -> StepProfile; + + /// Same semantics as [`Step::affinity`]. Defaults to + /// `Affinity::None`. + fn affinity(&self) -> Affinity { + Affinity::None + } + + /// Same semantics as [`Step::detached_group`]. Defaults to + /// [`DetachedGroup::PerStep`]. + fn detached_group(&self) -> DetachedGroup { + DetachedGroup::PerStep + } + + /// Step body. Pop from `ctx.a` (input branch 0) and `ctx.b` + /// (input branch 1), push to `ctx.outputs`. + /// + /// # Errors + /// + /// Same handling as [`Step::try_run`]. + fn try_run(&mut self, ctx: &mut StepCtx2<'_, Self>) -> io::Result; + + /// Same semantics as [`Step::new_worker_copy`]. Default panics; + /// only `Parallel` Step2 impls need to override. + /// + /// # Panics + /// + /// Default impl panics with the step name + kind. Hitting it + /// indicates a framework bug (the runtime tried to clone a + /// non-`Parallel` step). + #[must_use] + fn new_worker_copy(&self) -> Self { + let p = self.profile(); + panic!( + "Step2::new_worker_copy invoked on '{}' (kind = {:?}); \ + only Parallel steps need to override this. The framework \ + never clones Serial/Exclusive steps — file a bug.", + p.name, p.kind + ); + } +} + +/// Context passed to [`Step2::try_run`]. Holds independent input handles +/// for each branch. +pub struct StepCtx2<'a, S: Step2> { + pub a: &'a dyn InputHandle, + pub b: &'a dyn InputHandle, + pub outputs: &'a OutputHandles, +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::queues::QueueSpec; + use crate::reorder::BranchOrdering; + + #[test] + fn step_profile_constructs() { + let p = StepProfile { + name: "Test", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 64 }], + branch_ordering: vec![BranchOrdering::None], + }; + assert_eq!(p.name, "Test"); + assert_eq!(p.kind, StepKind::Parallel); + assert!(!p.sticky); + assert_eq!(p.output_queues.len(), 1); + assert!(matches!(p.output_queues[0], QueueSpec::CountBounded { capacity: 64 })); + assert_eq!(p.branch_ordering, vec![BranchOrdering::None]); + } + + #[test] + fn step_kind_is_comparable() { + assert_eq!(StepKind::Parallel, StepKind::Parallel); + assert_ne!(StepKind::Parallel, StepKind::Serial); + assert_ne!(StepKind::Serial, StepKind::Exclusive); + } + + #[test] + fn step_outcome_includes_finished() { + let o = StepOutcome::Finished; + assert_eq!(o, StepOutcome::Finished); + assert_ne!(StepOutcome::Finished, StepOutcome::Progress); + assert_ne!(StepOutcome::Finished, StepOutcome::NoProgress); + assert_ne!(StepOutcome::Finished, StepOutcome::Contention); + } +} + +#[cfg(test)] +mod step_trait_compile_tests { + use super::*; + use crate::outputs::Single; + use crate::queues::QueueSpec; + use crate::reorder::BranchOrdering; + + /// Stub step that does nothing — exercises the trait declaration. + #[derive(Clone)] + struct NopStep; + + impl Step for NopStep { + type Input = u32; + type Outputs = Single; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "Nop", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 1 }], + branch_ordering: vec![BranchOrdering::None], + } + } + + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + #[test] + fn nop_step_advertises_profile() { + let s = NopStep; + let p = s.profile(); + assert_eq!(p.name, "Nop"); + assert_eq!(p.kind, StepKind::Parallel); + } + + /// Stub multi-output step — exercises tuple Outputs. + #[derive(Clone)] + struct FanOutStep; + + impl Step for FanOutStep { + type Input = u32; + type Outputs = (u32, String); + + fn profile(&self) -> StepProfile { + StepProfile { + name: "FanOut", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![ + QueueSpec::CountBounded { capacity: 1 }, + QueueSpec::CountBounded { capacity: 1 }, + ], + branch_ordering: vec![BranchOrdering::None, BranchOrdering::None], + } + } + + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + #[test] + fn fan_out_step_advertises_profile() { + let s = FanOutStep; + assert_eq!(s.profile().name, "FanOut"); + } + + /// Default `new_worker_copy` panics on Serial steps. The framework + /// guarantees it never invokes this on Serial/Exclusive in practice; + /// this test pins the safety-net behavior in case a future refactor + /// accidentally calls it. The panic message must include the step + /// name so a real-world hit is debuggable. + #[test] + #[should_panic(expected = "SerialStub")] + fn new_worker_copy_default_panics_for_serial() { + struct SerialStub; + impl Step for SerialStub { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "SerialStub", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 1 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + let _ = SerialStub.new_worker_copy(); + } + + /// A two-input step that overrides nothing beyond the required items, so + /// the trait defaults are what the framework will read. + struct Step2Stub; + impl Step2 for Step2Stub { + type InputA = u32; + type InputB = u64; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "Step2Stub", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 1 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx2<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + /// `Step2` mirrors `Step`'s defaults: any worker may attempt a Serial step, + /// and a step that opts into `Detached` gets its own driver thread unless it + /// asks to share one. Pinning these keeps the two traits from drifting. + #[test] + fn step2_defaults_match_step_defaults() { + let s = Step2Stub; + assert!(matches!(s.affinity(), Affinity::None)); + assert!(matches!(s.detached_group(), DetachedGroup::PerStep)); + } + + /// Same safety net as `new_worker_copy_default_panics_for_serial`, for the + /// two-input trait: the framework never clones a non-`Parallel` step, so + /// reaching this default is a framework bug and must name the step. + #[test] + #[should_panic(expected = "Step2Stub")] + fn step2_new_worker_copy_default_panics_for_serial() { + let _ = Step2Stub.new_worker_copy(); + } +} diff --git a/crates/fgumi-pipeline-core/src/tests.rs b/crates/fgumi-pipeline-core/src/tests.rs new file mode 100644 index 000000000..b76f59fd5 --- /dev/null +++ b/crates/fgumi-pipeline-core/src/tests.rs @@ -0,0 +1,1241 @@ +//! Phase 1 cross-module tests. Validates that the type vocabulary, queue +//! handles, and `TypedStep` adapter compose into a working chain when +//! wired manually. +//! +//! The runtime that wires these automatically (worker pool, worker loop) +//! lives in Phase 2. These tests confirm the Phase 1 pieces support the +//! unified report-`Finished` completion contract end-to-end. + +use std::collections::VecDeque; +use std::io; +use std::sync::Arc; + +use super::erased::{ErasedStep, ErasedStepCtx}; +use super::outputs::Single; +use super::queues::QueueSpec; +use super::reorder::BranchOrdering; +use super::signal::PipelineSignal; +use super::step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}; + +// ───────────────────────────────────────────────────────────────────────────── +// F1' — Step2 (multi-input merge) tests. +// +// Validates that: +// * `TypedStep2` adapter dispatches correctly through `ErasedStep`. +// * `TwoInputHandles` round-trips per-branch handles via the +// ChainContexts construction path. +// * Both branches' drain signals propagate so `is_input_drained` +// reports drained only when BOTH branches are drained. +// ───────────────────────────────────────────────────────────────────────────── + +use super::erased::TypedStep2; +use super::step::{InputHandle as _, Step2, StepCtx2}; + +/// Test `Step2`: sum of `(a, b)` → `u64`. Emits `a + b as u64` once both +/// branches have yielded an item, `NoProgress` otherwise. Mirrors what zipper +/// does at the level of "pop from two queues in lockstep, combine". +/// +/// Buffers a popped item whose sibling branch was empty rather than popping both +/// eagerly: `(ctx.a.pop(), ctx.b.pop())` discards any `Some` that lands in the +/// non-pair arm, which is silent data loss. This test's producers happen to be +/// pre-loaded with equal counts so the hazard never fires here — which is +/// exactly why it must not be modelled this way: a step author copying the +/// pattern into real code loses records, and a future change to the test's input +/// counts would turn the bug into a silently passing test. Same shape as +/// `PairSummer` below. +#[derive(Default)] +struct SumPairStep { + pending_a: Option, + pending_b: Option, +} + +impl Step2 for SumPairStep { + type InputA = u32; + type InputB = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "SumPair", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 8 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx2<'_, Self>) -> io::Result { + if self.pending_a.is_none() { + self.pending_a = ctx.a.pop(); + } + if self.pending_b.is_none() { + self.pending_b = ctx.b.pop(); + } + match (self.pending_a.take(), self.pending_b.take()) { + (Some(a), Some(b)) => { + ctx.outputs + .push(u64::from(a) + u64::from(b)) + .expect("capacity-8 output must accept the pair"); + Ok(StepOutcome::Progress) + } + // Hold whichever branch did yield until its sibling catches up. + (a, b) => { + self.pending_a = a; + self.pending_b = b; + Ok(StepOutcome::NoProgress) + } + } + } +} + +#[test] +fn step2_typed_dispatch_pairs_both_branches() { + // Build two single-output sub-chains feeding into a SumPair Step2. + // Producer A and producer B each output u32 on branch 0. We manually + // wire the TypedStep2's TwoInputHandles via + // build_two_input_handles, then drive it through ErasedStep. + + let sum_step: Box = Box::new(TypedStep2::new(SumPairStep::default())); + + // Producer A's output set + handles. + let (a_set, a_view) = as super::outputs::StepOutputs>::build_queues( + &[QueueSpec::CountBounded { capacity: 4 }], + &[BranchOrdering::None], + crate::builder::InstrumentationLevel::Off, + ); + let a_outputs: super::step::OutputHandles> = + super::step::OutputHandles::new(a_view); + + // Producer B's output set + handles. + let (b_set, b_view) = as super::outputs::StepOutputs>::build_queues( + &[QueueSpec::CountBounded { capacity: 4 }], + &[BranchOrdering::None], + crate::builder::InstrumentationLevel::Off, + ); + let b_outputs: super::step::OutputHandles> = + super::step::OutputHandles::new(b_view); + + // Build the merge step's TwoInputHandles by handing the + // ErasedStep adapter both producer sets. + let mut producer_sets = vec![a_set, b_set]; + let merge_input_any = sum_step.build_two_input_handles(&mut producer_sets, 0, 0, 1, 0); + let _ = producer_sets; // both branches taken by build_two_input_handles + + // Sanity: input_arity reports 2. + assert_eq!(sum_step.input_arity(), 2); + + // Merge step's output set. + let (mut merge_outset, merge_view) = + sum_step.build_output_set(crate::builder::InstrumentationLevel::Off); + let merge_outputs_any = sum_step.wrap_outputs_view(merge_view); + let merge_consumer_input = merge_outset.take_typed_input::(0); + + // Push a couple of items on each producer. + a_outputs.push(10).unwrap(); + a_outputs.push(20).unwrap(); + b_outputs.push(1).unwrap(); + b_outputs.push(2).unwrap(); + + let signal = PipelineSignal::new(); + // Drive the merge step twice — once per pair. + let mut sum_step = sum_step; + for _ in 0..2 { + let mut ctx = ErasedStepCtx { + input: merge_input_any.as_ref(), + outputs: merge_outputs_any.as_ref(), + signal: &signal, + }; + let outcome = sum_step.try_run_erased(&mut ctx).unwrap(); + assert_eq!(outcome, StepOutcome::Progress); + } + // No more pairs available — both branches non-drained but empty. + { + let mut ctx = ErasedStepCtx { + input: merge_input_any.as_ref(), + outputs: merge_outputs_any.as_ref(), + signal: &signal, + }; + let outcome = sum_step.try_run_erased(&mut ctx).unwrap(); + assert_eq!(outcome, StepOutcome::NoProgress); + } + + // Verify the merge step pushed 11, 22 (10+1, 20+2) in order. + assert_eq!(merge_consumer_input.pop(), Some(11)); + assert_eq!(merge_consumer_input.pop(), Some(22)); + assert_eq!(merge_consumer_input.pop(), None); +} + +#[test] +fn step2_build_two_input_handles_rejects_same_producer_and_branch() { + // When both `Step2` inputs share ONE producer step, + // `build_two_input_handles` asserts the two branches differ — it does NOT + // require distinct producer indices. Taking the same branch twice is the + // illegal case; pin the panic so a refactor that drops the branch check + // surfaces immediately. (The legal same-producer/distinct-branch case is + // covered by the test below.) + let sum_step: Box = Box::new(TypedStep2::new(SumPairStep::default())); + let (a_set, _a_view) = as super::outputs::StepOutputs>::build_queues( + &[QueueSpec::CountBounded { capacity: 4 }], + &[BranchOrdering::None], + crate::builder::InstrumentationLevel::Off, + ); + let mut producer_sets = vec![a_set]; + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + sum_step.build_two_input_handles(&mut producer_sets, 0, 0, 0, 0) + })); + assert!(result.is_err(), "expected panic on same producer AND same branch"); +} + +/// Like `SumPairStep` but with a **non-commutative** combiner, so a test can tell +/// input A from input B. `SumPairStep` cannot: `a + b` makes `10 + 1` and `1 + 10` +/// both 11, so a swapped branch-to-input wiring produces an identical result and +/// is invisible. `a * 1000 + b` distinguishes them (10001 vs 1010). +#[derive(Default)] +struct AsymmetricPairStep { + pending_a: Option, + pending_b: Option, +} + +impl Step2 for AsymmetricPairStep { + type InputA = u32; + type InputB = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "AsymmetricPair", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 8 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx2<'_, Self>) -> io::Result { + if self.pending_a.is_none() { + self.pending_a = ctx.a.pop(); + } + if self.pending_b.is_none() { + self.pending_b = ctx.b.pop(); + } + match (self.pending_a, self.pending_b) { + (Some(a), Some(b)) => { + self.pending_a = None; + self.pending_b = None; + ctx.outputs + .push(u64::from(a) * 1000 + u64::from(b)) + .expect("capacity-8 output must accept the pair"); + Ok(StepOutcome::Progress) + } + // Hold whichever branch did yield until its sibling catches up. + _ => Ok(StepOutcome::NoProgress), + } + } +} + +/// A fan-out producer feeding BOTH of a `Step2`'s inputs from its two distinct +/// branches is legal, and had no coverage — the duplicate-producer test above +/// passes only because both its branches are `0`, so a regression that took the +/// same branch twice would not be caught by it. +#[test] +fn step2_build_two_input_handles_accepts_one_producer_with_distinct_branches() { + // Asymmetric combiner on purpose: `SumPairStep`'s `a + b` is commutative, so a + // builder that wired branch 1 to input A and branch 0 to input B would produce + // the same 11 and this test would still pass while asserting the mapping. + let sum_step: Box = Box::new(TypedStep2::new(AsymmetricPairStep::default())); + + // One producer, two branches: `(A, B)` is the two-branch output shape. + let (fanout_set, fanout_view) = <(u32, u32) as super::outputs::StepOutputs>::build_queues( + &[QueueSpec::CountBounded { capacity: 4 }, QueueSpec::CountBounded { capacity: 4 }], + &[BranchOrdering::None, BranchOrdering::None], + crate::builder::InstrumentationLevel::Off, + ); + let fanout_outputs: super::step::OutputHandles<(u32, u32)> = + super::step::OutputHandles::new(fanout_view); + + // Same producer index (0) for both inputs, distinct branches 0 and 1. + let mut producer_sets = vec![fanout_set]; + let merge_input_any = sum_step.build_two_input_handles(&mut producer_sets, 0, 0, 0, 1); + + let (mut merge_outset, merge_view) = + sum_step.build_output_set(crate::builder::InstrumentationLevel::Off); + let merge_outputs_any = sum_step.wrap_outputs_view(merge_view); + let merge_consumer_input = merge_outset.take_typed_input::(0); + + // Branch 0 feeds input A, branch 1 feeds input B. + let view = fanout_outputs.view(); + view.a.push(10).unwrap(); + view.b.push(1).unwrap(); + + let signal = PipelineSignal::new(); + let mut sum_step = sum_step; + let mut ctx = ErasedStepCtx { + input: merge_input_any.as_ref(), + outputs: merge_outputs_any.as_ref(), + signal: &signal, + }; + assert_eq!(sum_step.try_run_erased(&mut ctx).unwrap(), StepOutcome::Progress); + assert_eq!( + merge_consumer_input.pop(), + Some(10 * 1000 + 1), + "branch 0 must feed input A and branch 1 input B — a swapped mapping yields 1010" + ); +} + +// ───────────────────────────────────────────────────────────────────────────── +// F1' integration test — end-to-end Pipeline::run with TWO sources joined +// at a Step2, demonstrating that the entire chain-builder + runtime path +// (PipelineBuilder.chain ×2 → MultiChain2::from_chains → MultiChain2::join +// → Chain::chain → PipelineBuilder::build → Pipeline::run) works through +// the real worker loop. +// ───────────────────────────────────────────────────────────────────────────── + +#[test] +#[allow(clippy::too_many_lines)] +fn step2_end_to_end_pipeline_pairs_two_sources_through_runtime() { + use super::builder::{MultiChain2, Pipeline, PipelineConfig}; + + /// Source A: emits 1..=5 then Finished. + #[derive(Clone)] + struct SourceA { + remaining: u32, + } + impl Step for SourceA { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "SourceA", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 8 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if self.remaining == 0 { + return Ok(StepOutcome::Finished); + } + let n = self.remaining; + self.remaining -= 1; + ctx.outputs.push(n).expect("capacity-8 output must accept all 5 items"); + Ok(StepOutcome::Progress) + } + } + + /// Source B: emits 10..=50 (step 10) then Finished — same count + /// as Source A so every paired emit consumes one from each. + #[derive(Clone)] + struct SourceB { + remaining: u32, + } + impl Step for SourceB { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "SourceB", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 8 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if self.remaining == 0 { + return Ok(StepOutcome::Finished); + } + let n = self.remaining * 10; + self.remaining -= 1; + ctx.outputs.push(n).expect("capacity-8 output must accept all 5 items"); + Ok(StepOutcome::Progress) + } + } + + /// Step2 merger: sum each pair (a, b) → u64. Buffers a popped + /// item from one branch if the sibling branch is empty — + /// otherwise the popped item would be dropped when the sibling + /// pop returns None, and we'd lose data. + struct PairSummer { + pending_a: Option, + pending_b: Option, + } + impl Step2 for PairSummer { + type InputA = u32; + type InputB = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "PairSummer", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 8 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx2<'_, Self>) -> io::Result { + if self.pending_a.is_none() { + self.pending_a = ctx.a.pop(); + } + if self.pending_b.is_none() { + self.pending_b = ctx.b.pop(); + } + match (self.pending_a, self.pending_b) { + (Some(a), Some(b)) => { + self.pending_a = None; + self.pending_b = None; + ctx.outputs + .push(u64::from(a) + u64::from(b)) + .expect("capacity-8 output must accept all 5 pairs"); + Ok(StepOutcome::Progress) + } + // Both inputs drained AND nothing buffered: genuinely done. + _ if ctx.a.is_drained() + && ctx.b.is_drained() + && self.pending_a.is_none() + && self.pending_b.is_none() => + { + Ok(StepOutcome::Finished) + } + // Both inputs drained but one branch item is still buffered — + // the two branches emitted different counts. Reporting + // `Finished` here would silently drop that item, exactly the + // data loss the `PairSummer` doc warns step authors about, so + // fail loudly instead of modelling the bug. + _ if ctx.a.is_drained() && ctx.b.is_drained() => { + panic!("PairSummer: unpaired item left buffered after both branches drained") + } + _ => Ok(StepOutcome::NoProgress), + } + } + } + + /// Sink: records received items so the test can assert pairing. + #[derive(Clone)] + struct CollectSink { + received: Arc>>, + } + impl Step for CollectSink { + type Input = u64; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "CollectSink", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + match ctx.input.pop() { + Some(n) => { + self.received.lock().push(n); + Ok(StepOutcome::Progress) + } + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } + } + + let received: Arc>> = Arc::new(parking_lot::Mutex::new(Vec::new())); + let received_for_run = Arc::clone(&received); + + run_with_deadlock_timeout("Step2 PairSummer join at --threads 4", move || { + let builder = Pipeline::builder(); + let chain_a = builder.chain(SourceA { remaining: 5 }); + let chain_b = builder.chain(SourceB { remaining: 5 }); + MultiChain2::::from_chains(chain_a, chain_b) + .join(PairSummer { pending_a: None, pending_b: None }) + .chain(CollectSink { received: received_for_run }) + .into_sink_marker(); + + let pipeline = builder.build().expect("pipeline build"); + // 4 threads: 2 Exclusive sources own 2 workers, the Exclusive sink + // owns a third, leaving 1 free worker for the Serial PairSummer. + pipeline.run(PipelineConfig { threads: 4, ..Default::default() }).expect("pipeline run"); + }); + + // Both sources emit 5 items in step (a: 5,4,3,2,1; b: 50,40,30,20,10), so + // pair sums are 55, 44, 33, 22, 11 — and that ORDER is guaranteed, not just + // the multiset. Each branch queue is FIFO, `PairSummer` pops at most one + // item per branch per tick into `pending_a`/`pending_b`, so A's k-th item + // always pairs with B's k-th; and the step is `Serial`, so pairs enter the + // output queue in emit order and the Exclusive sink pops them in that order. + // Sorting here would discard an ordering guarantee the transport does + // provide, letting a reordering regression in Serial dispatch pass. + let collected = received.lock().clone(); + assert_eq!( + collected, + vec![55, 44, 33, 22, 11], + "Serial dispatch must preserve FIFO pairing and output order" + ); +} + +/// Sink: records every value it receives so tests can assert no records were +/// dropped, duplicated, or corrupted — not just that the count matched. +#[derive(Clone)] +struct DrainReproSink { + received: Arc>>, +} +impl Step for DrainReproSink { + type Input = u32; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "DrainReproSink", + 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(n) => { + self.received.lock().push(n); + Ok(StepOutcome::Progress) + } + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Unified model (issue #330): a mid-chain step (not just sources) reports +// `StepOutcome::Finished` once its input is drained and it holds no buffered +// output, folding its final flush into `try_run`'s flush-first path. A +// `NoProgress` tick mid-flush (full output queue) is just an idle yield — the +// step keeps flushing across re-dispatches until it reports `Finished`. At >1 +// thread the shared `finished` latch must stop the other workers from +// re-running the finished Serial step. +// ───────────────────────────────────────────────────────────────────────────── + +/// Source emitting `0..count` then `Finished`, with an output queue wide enough +/// to hold every item so `push` never hits backpressure: a Serial source +/// hammered by N workers must not silently drop on a full queue, which is why +/// the push below is asserted rather than discarded. Capacity 256 covers the +/// test's 64 items. +#[derive(Clone)] +struct WideQueueSource { + remaining: u32, +} +impl Step for WideQueueSource { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "WideQueueSource", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 256 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if self.remaining == 0 { + return Ok(StepOutcome::Finished); + } + self.remaining -= 1; + ctx.outputs.push(self.remaining).expect("wide source queue must never reject"); + Ok(StepOutcome::Progress) + } +} + +/// Mid (Serial): buffers every input item, then flushes the whole buffer in +/// `try_run`'s flush-first path through a capacity-1 output queue, and reports +/// `Finished` once input is drained and the buffer is empty — the unified +/// completion contract. +#[derive(Clone)] +struct ReportsFinishedBuffer { + // A `VecDeque` so the flush-first drain pops the FIFO front in O(1). Real + // steps process millions of records: a `Vec` with `remove(0)` would be O(n) + // per flushed item (quadratic drain). Step authors copying this worked + // example must keep the front-drain O(1) — never `Vec::remove(0)`. + buffered: VecDeque, +} +impl Step for ReportsFinishedBuffer { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "ReportsFinishedBuffer", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 1 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + // 1. Flush-first: push held items until the queue rejects one. + let mut pushed = false; + while let Some(&item) = self.buffered.front() { + match ctx.outputs.push(item) { + Ok(()) => { + self.buffered.pop_front(); + pushed = true; + } + Err(_unpushed) => break, // queue full — yield, retry next pass + } + } + if pushed { + return Ok(StepOutcome::Progress); + } + // 2. Consume input. + if let Some(n) = ctx.input.pop() { + self.buffered.push_back(n); + return Ok(StepOutcome::Progress); + } + // 3. Completion: input drained AND nothing held — never push again. + if ctx.input.is_drained() && self.buffered.is_empty() { + return Ok(StepOutcome::Finished); + } + // Input drained but buffer still non-empty (queue was full this call): + // NoProgress is just an idle tick — the round-robin moves on to the + // sink, which drains the queue, and we flush more on the next pass, + // eventually emptying the buffer and reporting Finished above. There is + // no drain protocol to misfire here. + Ok(StepOutcome::NoProgress) + } + // No `new_worker_copy` override: this step is `Serial`, so the framework + // holds ONE shared instance behind a mutex and never clones it. An override + // here would be dead code implying `buffered` could be split per worker. +} + +/// Run `build_and_run` on its own thread and fail the test if it has not +/// returned within ten seconds. +/// +/// Every end-to-end test in this module must go through here. These pipelines' +/// primary failure modes are deadlock and stall, and calling `Pipeline::run` +/// directly on the test thread turns either one into a hung harness with no +/// message — the run never returns, so no assertion is ever reached. `context` +/// names the pipeline in the timeout message. +/// +/// A panic inside `build_and_run` is re-raised as itself, so an assertion +/// failure inside the run reports its own message rather than a spurious +/// deadlock. +fn run_with_deadlock_timeout(context: &str, build_and_run: impl FnOnce() + Send + 'static) { + let (tx, rx) = std::sync::mpsc::channel::<()>(); + let handle = std::thread::spawn(move || { + build_and_run(); + let _ = tx.send(()); + }); + match rx.recv_timeout(std::time::Duration::from_secs(10)) { + Ok(()) => handle.join().expect("worker thread panicked"), + // The sender dropped without sending, so `build_and_run` panicked — + // typically a failed assertion inside the run. Re-raise that panic + // instead of blaming a deadlock: reporting "DEADLOCKED or stalled" here + // would invert the diagnosis for every assertion failure in a pipeline. + Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => { + let panic = handle.join().expect_err("a disconnect implies the closure panicked"); + std::panic::resume_unwind(panic); + } + // Nothing sent and the sender is still alive: a genuine stall. + Err(std::sync::mpsc::RecvTimeoutError::Timeout) => { + panic!("{context} DEADLOCKED or stalled") + } + } +} + +fn run_reports_finished_pipeline(threads: usize, n_items: u32) { + let received: Arc>> = Arc::new(parking_lot::Mutex::new(Vec::new())); + let received_for_run = Arc::clone(&received); + + run_with_deadlock_timeout( + &format!("reports_finished mid-step pipeline at --threads {threads}"), + move || { + use crate::{Pipeline, PipelineConfig}; + let builder = Pipeline::builder(); + builder + .chain(WideQueueSource { remaining: n_items }) + .chain(ReportsFinishedBuffer { buffered: VecDeque::new() }) + .chain(DrainReproSink { received: received_for_run }) + .into_sink_marker(); + let pipeline = builder.build().expect("pipeline build"); + pipeline.run(PipelineConfig { threads, ..Default::default() }).expect("pipeline run"); + }, + ); + // WideQueueSource emits exactly 0..n_items. Assert the multiset (not just the + // count) so a dropped, duplicated, or value-corrupted record is caught, not + // only an off-by-N total. + let mut got = received.lock().clone(); + if threads == 1 { + // At one worker the sequence is fully determined, so sorting would throw + // away a checkable guarantee: `WideQueueSource` emits `n_items-1..0`, + // every transport is FIFO, and the single `ReportsFinishedBuffer` drains + // its `VecDeque` front-first — so a reordering regression in the + // flush-first path is only visible here. Assert identity, not multiset. + let expected_order: Vec = (0..n_items).rev().collect(); + assert_eq!( + got, expected_order, + "single worker must deliver items in emit order (flush-first is FIFO)" + ); + return; + } + // Above one worker, order is not asserted: the Serial step's re-dispatch + // interleaving across workers is a valid scheduling detail. + got.sort_unstable(); + let expected: Vec = (0..n_items).collect(); + assert_eq!( + got, expected, + "sink must receive every item 0..{n_items} exactly once (no drop/dup/corruption) \ + at --threads {threads}" + ); +} + +#[test] +fn reports_finished_mid_step_single_thread_no_premature_drain() { + // Lone worker: returning NoProgress mid-flush must not lose buffered output; + // the round-robin must reach the sink so the full queue drains (no spin), + // and the step eventually reports Finished. (The old BUG #5 drain-spin shape.) + run_reports_finished_pipeline(1, 5); +} + +#[test] +fn reports_finished_mid_step_multi_thread_shared_latch_stops_redispatch() { + // 4 workers: the mid step is Serial + Affinity::None, so it sits in every + // worker's live set. When one worker finishes it, the shared `finished` + // latch must stop the others from re-running it (a non-idempotent flusher + // re-entered would push into a drained queue → panic). A larger item count + // widens the window for a concurrent re-dispatch. + run_reports_finished_pipeline(4, 64); +} + +// ───────────────────────────────────────────────────────────────────────────── +// `MultiChain2Ordered::{from_chains, join}` smoke test. Mirrors +// `step2_end_to_end_pipeline_pairs_two_sources_through_runtime` but uses +// `OrderedBytesSingle<_>` source outputs — the chain wrapper real BAM/FASTQ +// source subchains produce — so the converge method is exercised on its +// native input shape. +// ───────────────────────────────────────────────────────────────────────────── + +#[test] +#[allow(clippy::too_many_lines)] +fn multi_chain2_ordered_pairs_two_byte_bounded_sources() { + use super::builder::{MultiChain2Ordered, Pipeline, PipelineConfig}; + use super::item::{HeapSize, Ordered}; + use super::outputs::OrderedBytesSingle; + use super::step::Step2; + use super::step::StepCtx2; + + /// Bytes of heap payload each test item carries. Non-zero so the + /// `ByteBounded` transports below actually account for these items: + /// `ByteBoundedQueue` sums `HeapSize::heap_size()` only (it never counts + /// `size_of::()`), so a zero-heap item would make every push free, the + /// byte cap unreachable, and this test blind to a regression that let a + /// `ByteBounded` queue grow without bound — the memory-bound failure mode it + /// exists to cover. + const PAYLOAD_BYTES: usize = 1024; + + /// Byte budget per output edge. Deliberately smaller than + /// `5 * PAYLOAD_BYTES` so the sources hit the cap partway through and must + /// hold the rejected item and retry — exercising backpressure rather than + /// just fitting everything in one go. + const EDGE_LIMIT_BYTES: usize = 2 * PAYLOAD_BYTES; + + /// Minimal ordered, byte-bounded item type for the test. + #[derive(Debug, Clone)] + struct OrderedU32 { + ordinal: u64, + value: u32, + payload: Vec, + } + + impl OrderedU32 { + fn new(ordinal: u64, value: u32) -> Self { + Self { ordinal, value, payload: vec![0u8; PAYLOAD_BYTES] } + } + } + + impl HeapSize for OrderedU32 { + fn heap_size(&self) -> usize { + self.payload.capacity() + } + } + + impl Ordered for OrderedU32 { + fn ordinal(&self) -> u64 { + self.ordinal + } + } + + /// Ordered source: emits items with monotonic `ordinal` carrying + /// values N..=1 then `Finished`. + /// + /// `branch_ordering` is `BranchOrdering::None` deliberately, despite the + /// `OrderedBytesSingle` output shape and the `Ordered` items. The shape is + /// chosen for its `HeapSize + Ordered` bounds and its byte-bounded transport, + /// not to engage the reorder stage: `build_branch_ordered_bytes` maps + /// `(ByteBounded, None)` to a **direct** branch, so these steps exercise + /// byte-backpressure over a plain FIFO edge, which is exactly what the + /// ordered-output assertion below reasons from (FIFO branch queues plus one + /// pop per branch per tick). Declaring `ByOrdinal`/`ByItemOrdinal` here would + /// change what the test covers, not strengthen it. + /// + /// Reorder-stage restoration is covered where it belongs — see + /// `erased::tests::step2_serial_byitemordinal_output_is_reordered_not_collapsed`, + /// `handles::handle_tests::ordered_branch_preserves_ordinal_across_retry`, and + /// the `reorder` module's own tests. + /// + /// Holds an item the byte-bounded output + /// rejected and retries it on a later tick — dropping it would punch a hole + /// in the ordinal sequence while still reporting `Progress`, so the sink's + /// count assertion would fail with no indication of the cause. + #[derive(Clone)] + struct OrderedSource { + next_ordinal: u64, + remaining: u32, + held: Option, + } + + impl Step for OrderedSource { + type Input = (); + type Outputs = OrderedBytesSingle; + fn profile(&self) -> StepProfile { + StepProfile { + name: "OrderedSource", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { + limit_bytes: EDGE_LIMIT_BYTES as u64, + }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + // Flush-first: retry the item a previous tick's full queue rejected. + // A rejected push reports `NoProgress`, not `Contention`: the driver + // treats the two identically, but `Contention` means "a Serial step's + // mutex was held by another worker" and feeds `contention_count`, from + // which the bottleneck verdict derives its SPIN ratio. Reporting it + // for ordinary output backpressure would invent mutex contention that + // never happened and can trip a bogus SPIN finding. + if let Some(item) = self.held.take() { + return match ctx.outputs.push(item) { + Ok(()) => Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held = Some(unpushed.into_item()); + Ok(StepOutcome::NoProgress) + } + }; + } + if self.remaining == 0 { + return Ok(StepOutcome::Finished); + } + let value = self.remaining; + self.remaining -= 1; + let ordinal = self.next_ordinal; + self.next_ordinal += 1; + match ctx.outputs.push(OrderedU32::new(ordinal, value)) { + Ok(()) => Ok(StepOutcome::Progress), + // Hold, never drop: the ordinal was already consumed above, so + // discarding the item would leave a hole in the sequence. + Err(unpushed) => { + self.held = Some(unpushed.into_item()); + Ok(StepOutcome::NoProgress) + } + } + } + } + + /// Step2 merger over two ordered inputs. Sums per-pair, emits + /// `OrderedU64`. Branch queues are FIFO and this step pops at most one item + /// per branch per tick, so pair order is deterministic — the test asserts + /// the exact output sequence, not just the multiset. Backpressure changes + /// only *when* each item moves, never the order. + #[derive(Debug, Clone)] + struct OrderedU64 { + ordinal: u64, + value: u64, + payload: Vec, + } + impl OrderedU64 { + fn new(ordinal: u64, value: u64) -> Self { + Self { ordinal, value, payload: vec![0u8; PAYLOAD_BYTES] } + } + } + impl HeapSize for OrderedU64 { + fn heap_size(&self) -> usize { + self.payload.capacity() + } + } + impl Ordered for OrderedU64 { + fn ordinal(&self) -> u64 { + self.ordinal + } + } + + struct OrderedPairSummer { + pending_a: Option, + pending_b: Option, + /// A summed pair the byte-bounded output rejected, retried on a later + /// tick. As in `OrderedSource`, dropping it would consume an out-ordinal + /// and lose a record while still reporting `Progress`. + held: Option, + next_out_ordinal: u64, + } + impl Step2 for OrderedPairSummer { + type InputA = OrderedU32; + type InputB = OrderedU32; + type Outputs = OrderedBytesSingle; + fn profile(&self) -> StepProfile { + StepProfile { + name: "OrderedPairSummer", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { + limit_bytes: EDGE_LIMIT_BYTES as u64, + }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, ctx: &mut StepCtx2<'_, Self>) -> io::Result { + // Flush-first: retry a pair the full output rejected earlier. + // `NoProgress` rather than `Contention` on a rejected push, for the + // reason given in `OrderedSource::try_run`. + if let Some(out) = self.held.take() { + return match ctx.outputs.push(out) { + Ok(()) => Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held = Some(unpushed.into_item()); + Ok(StepOutcome::NoProgress) + } + }; + } + if self.pending_a.is_none() { + self.pending_a = ctx.a.pop(); + } + if self.pending_b.is_none() { + self.pending_b = ctx.b.pop(); + } + match (self.pending_a.as_ref(), self.pending_b.as_ref()) { + (Some(_), Some(_)) => { + let a = self.pending_a.take().unwrap(); + let b = self.pending_b.take().unwrap(); + let out = OrderedU64::new( + self.next_out_ordinal, + u64::from(a.value) + u64::from(b.value), + ); + self.next_out_ordinal += 1; + match ctx.outputs.push(out) { + Ok(()) => Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held = Some(unpushed.into_item()); + Ok(StepOutcome::NoProgress) + } + } + } + // Same completion guard as `PairSummer`: both inputs drained and + // neither branch item buffered. (`held` is always `None` here — + // the flush-first block above returns when it is `Some`.) + _ if ctx.a.is_drained() + && ctx.b.is_drained() + && self.pending_a.is_none() + && self.pending_b.is_none() => + { + Ok(StepOutcome::Finished) + } + _ if ctx.a.is_drained() && ctx.b.is_drained() => { + panic!( + "OrderedPairSummer: unpaired item left buffered after both branches drained" + ) + } + _ => Ok(StepOutcome::NoProgress), + } + } + } + + /// Records `(ordinal, value)`, not just `value`. + /// + /// The out-ordinal is the thing `OrderedPairSummer`'s hold-and-retry path + /// exists to protect: it assigns `next_out_ordinal` *before* the push, so a + /// rejected push must retry the SAME ordinal. Recording only `value` left that + /// invariant unasserted — a regression that reassigned or skipped an ordinal on + /// the retry path leaves the value sequence intact and the test still passes. + #[derive(Clone)] + struct OrderedSink { + received: Arc>>, + } + impl Step for OrderedSink { + type Input = OrderedU64; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "OrderedSink", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + match ctx.input.pop() { + Some(item) => { + self.received.lock().push((item.ordinal, item.value)); + Ok(StepOutcome::Progress) + } + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } + } + + // Pin the byte budget's teeth deterministically, independent of how the + // scheduler happens to interleave the run below: an edge must admit at least + // one item (or the pipeline would wedge) and must reject before all five fit + // (or the cap would never bind and the test would prove nothing about byte + // bounding). Both follow from PAYLOAD_BYTES vs EDGE_LIMIT_BYTES, so a future + // edit to either constant that silently removes backpressure fails here. + { + use crate::queues::{ByteBoundedQueue, ItemQueue}; + let q = ByteBoundedQueue::::new(EDGE_LIMIT_BYTES as u64); + let admitted = (0..5u32) + .take_while(|i| q.try_push(OrderedU32::new(u64::from(*i), *i)).is_ok()) + .count(); + assert!(admitted >= 1, "a byte-bounded edge must always admit one item, else it wedges"); + assert!( + admitted < 5, + "the byte cap must bind before all 5 items fit — otherwise this test \ + exercises no byte bounding at all (admitted {admitted})" + ); + } + + let received: Arc>> = + Arc::new(parking_lot::Mutex::new(Vec::new())); + let received_for_run = Arc::clone(&received); + + run_with_deadlock_timeout("MultiChain2Ordered byte-bounded join at --threads 4", move || { + let builder = Pipeline::builder(); + let chain_a = builder.chain(OrderedSource { next_ordinal: 0, remaining: 5, held: None }); + let chain_b = builder.chain(OrderedSource { next_ordinal: 0, remaining: 5, held: None }); + MultiChain2Ordered::::from_chains(chain_a, chain_b) + .join(OrderedPairSummer { + pending_a: None, + pending_b: None, + held: None, + next_out_ordinal: 0, + }) + .chain(OrderedSink { received: received_for_run }) + .into_sink_marker(); + + let pipeline = builder.build().expect("pipeline build"); + pipeline.run(PipelineConfig { threads: 4, ..Default::default() }).expect("pipeline run"); + }); + + let collected = received.lock().clone(); + // Both sources emit values 5,4,3,2,1 in step → pair sums 10,8,6,4,2, in that + // order for the same reason as the `PairSummer` test above: FIFO branch + // queues, one pop per branch per tick, and a `Serial` merge feeding an + // `Exclusive` sink. Backpressure (the byte cap is smaller than the total + // payload) changes only WHEN each item moves, never the order. + assert_eq!( + collected, + vec![(0, 10), (1, 8), (2, 6), (3, 4), (4, 2)], + "Serial dispatch must preserve FIFO pairing, output order, AND the out-ordinal \ + sequence under backpressure — the retry path must reuse an ordinal, not skip it" + ); +} + +// ───────────────────────────────────────────────────────────────────────────── +// Flattened public API. +// +// `lib.rs` re-exports the step-author surface at the crate root so an +// implementor writes `fgumi_pipeline_core::StepProfile`, not +// `::step::StepProfile`. A type reachable only through its module path is an +// asymmetry a step author trips over, and nothing else catches a dropped +// `pub use` — everything in-crate refers to these by module path anyway, and +// the crate has no in-tree consumer to break. +// ───────────────────────────────────────────────────────────────────────────── + +/// Every type a step author needs must resolve at the crate root. This is a +/// compile-time assertion: naming them in type position fails to build if a +/// `pub use` is dropped or a new peer shape is added without one. +/// +/// `DetachedGroup`, the ordered tuple shapes and `InstrumentationLevel` are the +/// ones this caught — `Step::detached_group` returns `DetachedGroup` while `Step` +/// itself was flattened, `OrderedBytesTuple2`/`3` sat beside an exported `Single` +/// and `OrderedBytesSingle`, and `InstrumentationLevel` appears in +/// `StepOutputs::build_queues`' signature while the trait itself was flattened. +/// +/// `InstrumentationLevel` also turned up a whole class the earlier passes had +/// only sampled: it is the *parameter* of `build_queues`, and sweeping every +/// `pub` type named by a fully-public signature found the **return** side +/// unflattened too — `OutputQueueSet`, the `Tuple*View`s from +/// `OutputHandles::view()`, and `MultiChain2Ordered` from `Chain::into_multi()`. +/// Hence the sweep, not just the one symbol, is what this test now pins. +/// +/// Deliberately NOT asserted: `queues::BoundedQueueHandle` and +/// `reorder::ReorderCapHandle`. They are runtime budget plumbing, named only via +/// `runtime::contexts::RegisteredQueue`, which is not flattened either — and +/// `ByteBoundedQueue` exposes inherent `limit_bytes` / `set_limit_bytes` / +/// `current_bytes`, so a direct user never needs the trait in scope. Promoting +/// them would make the root surface *less* coherent, not more. +/// +/// Also deliberately NOT asserted: `handles::BranchOutputHandle`. No public +/// signature anywhere in the crate mentions it — a step author reaches its +/// methods through `OutputHandles` / the `Tuple*View`s and never names the type — +/// so there is no signature it has to be nameable in. +/// +/// The rule this test encodes is therefore "nameable because a public signature +/// names it", not "`pub`, therefore flattened". `BranchInputHandle` is asserted +/// under exactly that rule: `OutputQueueSet::take_typed_input` returns it, so +/// flattening `OutputQueueSet` here brought it into reach of a root-only import +/// and it had to follow. +#[test] +fn crate_root_reexports_the_step_author_surface() { + use crate as api; + use std::marker::PhantomData; + + fn arity_of() -> usize { + O::arity() + } + + // Step / Step2 vocabulary. + let _: fn() -> api::StepKind = || api::StepKind::Serial; + let _: fn() -> api::StepOutcome = || api::StepOutcome::Progress; + let _: fn() -> api::Affinity = || api::Affinity::None; + let _: fn() -> api::DetachedGroup = || api::DetachedGroup::PerStep; + let _: fn() -> api::QueueSpec = || api::QueueSpec::Unbounded; + let _: fn() -> api::BranchOrdering = || api::BranchOrdering::None; + // `StepOutputs::build_queues` names this in its public signature, so anyone + // implementing that trait by hand has to be able to name it too. + let _: fn() -> api::InstrumentationLevel = || api::InstrumentationLevel::Off; + + // Return types of public methods, which a step author has to name to store + // one or to write a helper that takes it. `PhantomData` rather than a value: + // naming the type in a position that checks its bounds is the whole + // assertion, and none of these are constructible from outside the runtime. + // `build_queues` returns `(OutputQueueSet, OutputsViewAny)` — the second half + // was already flattened, the first was not. + let _: PhantomData = PhantomData; + // `OutputHandles::view()`, one per fan-out arity. + let _: PhantomData> = PhantomData; + let _: PhantomData> = PhantomData; + let _: PhantomData> = PhantomData; + // `Chain::into_multi()`, the ordered sibling of the exported `MultiChain2`. + let _: PhantomData, api::Sequenced>> = + PhantomData; + // `OutputQueueSet::take_typed_input()` — flattening `OutputQueueSet` is what + // put this one within reach of a root-only import. + let _: PhantomData> = PhantomData; + // Naming it with its type is the whole assertion — a comparison against a + // literal would be a constant expression, not a check. + let _: usize = api::MAX_ARITY; + + // Output shapes: one arity per declared shape, so a new shape added without + // a re-export shows up here. + assert_eq!(arity_of::>(), 1); + assert_eq!(arity_of::>>(), 1); + assert_eq!(arity_of::<(u32, u32)>(), 2); + assert_eq!( + arity_of::, api::Sequenced>>(), + 2, + "OrderedBytesTuple2 must be nameable at the crate root" + ); + assert_eq!( + arity_of::< + api::OrderedBytesTuple3, api::Sequenced, api::Sequenced>, + >(), + 3, + "OrderedBytesTuple3 must be nameable at the crate root" + ); +} + +/// The crate docs claim the dependency graph "stays light" and then enumerate +/// it. That list is a promise about the whole graph, so it has to be exhaustive +/// — a reader weighing this crate as a dependency reads the list, not the +/// manifest. Nothing else notices when the two drift: adding a dependency +/// compiles fine, and the prose keeps asserting the old, shorter graph. +/// +/// `anyhow` is the one this caught. It backs `FinalizeHook::finalize`'s return +/// type and went undocumented, so the list understated the graph by one crate. +/// +/// Dev-dependencies are deliberately out of scope: the claim is about what a +/// consumer links, and `proptest` / `rstest` / `trybuild` are not that. +/// Platform-gated tables are in scope, because a consumer on that platform does +/// link them — `fgumi-sort` already carries `[target.'cfg(unix)'.dependencies]`, +/// so reading only the plain `[dependencies]` table would fail *open* the day +/// this crate grows one. +#[test] +fn crate_docs_enumerate_every_runtime_dependency() { + // `include_str!` resolves against this file's directory, so both paths are + // the real files the claim is made in and about — not a copy that can drift. + let manifest = include_str!("../Cargo.toml"); + let crate_docs = include_str!("lib.rs"); + + // `[dependencies]` plus every `[target.'cfg(..)'.dependencies]`. + // `[dev-dependencies]` and `[build-dependencies]` end in `-dependencies`, + // so neither form matches them, at top level or under a `target` table. + let is_runtime_table = + |header: &str| header == "dependencies" || header.ends_with(".dependencies"); + + let mut dependencies: Vec<&str> = Vec::new(); + let mut in_runtime_dependencies = false; + for line in manifest.lines() { + if let Some(header) = line.trim().strip_prefix('[').and_then(|h| h.strip_suffix(']')) { + in_runtime_dependencies = is_runtime_table(header); + // A crate may instead declare itself in the header, as + // `[dependencies.log]` or `[target.'cfg(unix)'.dependencies.libc]`, + // with only its own keys in the body. Recognizing the sub-table by + // its parent is what keeps both spellings in the check — matching + // the `dependencies.` prefix alone would take the first and let the + // target-scoped one fall through undetected. The trim handles the + // quoted-key spelling, `[dependencies."log"]`, which is also legal. + if !in_runtime_dependencies + && let Some((parent, name)) = header.rsplit_once('.') + && is_runtime_table(parent) + { + dependencies.push(name.trim_matches('"')); + } + continue; + } + // A dependency key sits at column 0. Skipping indented and commented + // lines is what keeps a multi-line inline table's `features = [..]` + // continuation, or a commented-out `# tokio = ..`, from being read as a + // crate name and failing this test under a name that is not a crate. + if !in_runtime_dependencies || line.starts_with([' ', '\t', '#']) { + continue; + } + if let Some((name, _)) = line.split_once('=') { + // Cargo allows a dotted dependency key (`log.workspace = true`), whose + // crate name is the segment before the dot. Take that, so the code-span + // check below matches the crate name rather than `log.workspace` and + // never reports a false undocumented dependency. + let name = name.trim(); + let name = name.split_once('.').map_or(name, |(crate_name, _)| crate_name); + if !name.is_empty() { + dependencies.push(name); + } + } + } + // A renamed section or a reordered manifest would otherwise leave this test + // asserting over an empty list and passing vacuously. + assert!( + !dependencies.is_empty(), + "parsed no runtime dependencies from Cargo.toml; the section header or layout moved \ + and this test would silently stop checking anything" + ); + + // Odd-index pieces of a backtick split are the code spans. Matching spans + // rather than raw substrings is what keeps `log` from being "documented" by + // the word `logging`, and what lets `` `noodles::sam` `` document `noodles`. + let code_spans: Vec<&str> = crate_docs + .lines() + .filter(|line| line.starts_with("//!")) + .flat_map(|line| line.split('`').skip(1).step_by(2)) + .collect(); + let undocumented: Vec<&str> = dependencies + .iter() + .filter(|dependency| { + !code_spans.iter().any(|span| { + span == *dependency + || span.strip_prefix(*dependency).is_some_and(|rest| rest.starts_with("::")) + }) + }) + .copied() + .collect(); + + assert!( + undocumented.is_empty(), + "crate docs in lib.rs enumerate the dependency graph but omit {undocumented:?}; \ + add them to the list or drop the claim" + ); +} diff --git a/crates/fgumi-pipeline-core/src/topology.rs b/crates/fgumi-pipeline-core/src/topology.rs new file mode 100644 index 000000000..dd2a8787c --- /dev/null +++ b/crates/fgumi-pipeline-core/src/topology.rs @@ -0,0 +1,380 @@ +//! `ChainGraph`: per-step queue + branch-consumer tracking. Used by +//! `PipelineBuilder::build()` to assert all-outputs-wired and by the +//! runtime to construct queue topology. + +/// Static branch-index display names covering every branch up to [`crate::outputs::MAX_ARITY`]. +/// Output branch counts are bounded by `MAX_ARITY` at registration, so a valid +/// branch index always maps to a name and the build-error message never prints a +/// placeholder. Out-of-range indices fall back to `"?"` (never hit in practice, +/// but keeps the function total). +fn branch_name(branch: usize) -> &'static str { + const BRANCH_NAMES: [&str; crate::outputs::MAX_ARITY] = ["0", "1", "2", "3"]; + BRANCH_NAMES.get(branch).copied().unwrap_or("?") +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct StepIdx(pub usize); + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct BranchIdx(pub usize); + +#[derive(Debug, Default)] +pub struct ChainGraph { + /// `consumers[(producer, branch)] = Some(consumer)` if wired. + consumers: Vec>, + /// `consumer_input_slots[(producer, branch)] = Some(slot)` — + /// which input branch of the consumer this edge feeds. + /// Defaults to `0` for single-input consumers; multi-input + /// consumers (`Step2` / future `StepN`) record their per-input + /// slot explicitly so [`crate::runtime::contexts`] + /// can build the right [`crate::handles::TwoInputHandles`] wrapper. + consumer_input_slots: Vec>, + /// `branch_count[step]` — number of output branches for that step. + branch_counts: Vec, + /// `branch_offsets[step]` — running prefix sum of `branch_counts` *before* + /// `step` (i.e. `sum(branch_counts[..step])`). This is the base index into + /// the flat `consumers` / `consumer_input_slots` arrays for `step`'s first + /// output branch, so `consumer_slot_index` is O(1) instead of re-summing the + /// prefix on every call. Pushed once per step in + /// [`Self::register_step_with_input_arity`]; static thereafter. + branch_offsets: Vec, + /// `input_arities[step]` — number of input branches for that + /// step (1 for single-input [`Step`](crate::step::Step) impls, 2 for [`Step2`](crate::step::Step2), + /// N for future `StepN`). Defaults to 1 via [`Self::register_step`]; + /// multi-input steps register via + /// [`Self::register_step_with_input_arity`]. + input_arities: Vec, + /// Step names for error messages. + step_names: Vec<&'static str>, +} + +impl ChainGraph { + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Register a step with a single input branch (`input_arity = 1`). + /// Convenience shorthand for [`Self::register_step_with_input_arity`]. + pub fn register_step(&mut self, name: &'static str, branch_count: usize) -> StepIdx { + self.register_step_with_input_arity(name, branch_count, 1) + } + + /// Register a step with an explicit input arity. Multi-input + /// steps (`Step2` and future `StepN`) pass `input_arity > 1`; + /// sources pass `input_arity = 0` (their input is implicit). + pub fn register_step_with_input_arity( + &mut self, + name: &'static str, + branch_count: usize, + input_arity: usize, + ) -> StepIdx { + let idx = StepIdx(self.branch_counts.len()); + // The offset for this step is the total branch count of all prior steps, + // i.e. the current length of the flat `consumers` array before we grow it. + self.branch_offsets.push(self.consumers.len()); + self.branch_counts.push(branch_count); + self.input_arities.push(input_arity); + self.step_names.push(name); + self.consumers.resize(self.consumers.len() + branch_count, None); + self.consumer_input_slots.resize(self.consumer_input_slots.len() + branch_count, None); + idx + } + + /// Wire a (producer, branch) → consumer link, into consumer's + /// input branch 0. Convenience shorthand for + /// [`Self::wire_to_slot`] used by single-input consumers. + /// + /// # Panics + /// + /// Panics if the (producer, branch) is already wired (the type system + + /// `Chain` move semantics ensure single-consumer in practice; this is a + /// defensive check). + pub fn wire(&mut self, producer: StepIdx, branch: BranchIdx, consumer: StepIdx) { + self.wire_to_slot(producer, branch, consumer, 0); + } + + /// Wire a (producer, branch) → (consumer, consumer's input slot) + /// link. Multi-input consumers (`Step2` / future `StepN`) record + /// the consumer's input-slot index so + /// [`crate::runtime::contexts`] can + /// route each edge to the right per-branch input handle. + /// + /// # Panics + /// + /// Panics if the (producer, branch) is already wired (defensive), or if + /// `producer`, `branch`, `consumer`, or `consumer_input_slot` are out of + /// range. + pub fn wire_to_slot( + &mut self, + producer: StepIdx, + branch: BranchIdx, + consumer: StepIdx, + consumer_input_slot: usize, + ) { + assert!( + consumer.0 < self.input_arities.len(), + "consumer StepIdx({}) out of range (graph has {} steps)", + consumer.0, + self.input_arities.len() + ); + let consumer_arity = self.input_arities[consumer.0]; + assert!( + consumer_input_slot < consumer_arity, + "consumer_input_slot {consumer_input_slot} out of range for step '{}' \ + with input_arity {consumer_arity}", + self.step_names[consumer.0] + ); + let slot = self.consumer_slot_index(producer, branch); + assert!( + self.consumers[slot].is_none(), + "branch already wired: {:?} branch {:?} → {:?}", + producer, + branch, + self.consumers[slot] + ); + self.consumers[slot] = Some(consumer); + self.consumer_input_slots[slot] = Some(consumer_input_slot); + } + + /// Returns the consumer-input-slot this (producer, branch) edge + /// feeds, if wired. Single-input consumers always return + /// `Some(0)`; multi-input consumers return `Some(0)` or + /// `Some(1)` depending on which input branch the edge feeds. + #[must_use] + pub fn consumer_input_slot(&self, producer: StepIdx, branch: BranchIdx) -> Option { + let slot = self.consumer_slot_index(producer, branch); + self.consumer_input_slots[slot] + } + + /// Returns the input arity of a step (1 for single-input + /// [`Step`](crate::step::Step) impls, 2 for [`Step2`](crate::step::Step2), etc.). Sources have arity 0. + #[must_use] + pub fn input_arity(&self, step: StepIdx) -> usize { + self.input_arities[step.0] + } + + /// Returns the first unwired branch as `(producer, branch_idx, branch_name)`, + /// or `None` if every output branch is wired. + #[must_use] + pub fn first_unwired(&self) -> Option<(StepIdx, BranchIdx, &'static str)> { + for (producer_usize, &branch_count) in self.branch_counts.iter().enumerate() { + for branch in 0..branch_count { + let producer = StepIdx(producer_usize); + let slot = self.consumer_slot_index(producer, BranchIdx(branch)); + if self.consumers[slot].is_none() { + return Some((producer, BranchIdx(branch), branch_name(branch))); + } + } + } + None + } + + #[must_use] + pub fn consumer(&self, producer: StepIdx, branch: BranchIdx) -> Option { + let slot = self.consumer_slot_index(producer, branch); + self.consumers[slot] + } + + /// Returns the first step wired INTO `consumer`, or `None` if nothing feeds + /// it. Used by `PipelineBuilder::build` to reject an edge into a source, + /// whose implicit input means the edge would never be consumed. + #[must_use] + pub fn first_producer_into(&self, consumer: StepIdx) -> Option { + for (producer_usize, &branch_count) in self.branch_counts.iter().enumerate() { + let producer = StepIdx(producer_usize); + for branch in 0..branch_count { + if self.consumer(producer, BranchIdx(branch)) == Some(consumer) { + return Some(producer); + } + } + } + None + } + + #[must_use] + pub fn step_name(&self, step: StepIdx) -> &'static str { + self.step_names[step.0] + } + + #[must_use] + pub fn n_steps(&self) -> usize { + self.branch_counts.len() + } + + #[must_use] + pub fn branch_count(&self, step: StepIdx) -> usize { + self.branch_counts[step.0] + } + + fn consumer_slot_index(&self, producer: StepIdx, branch: BranchIdx) -> usize { + assert!( + producer.0 < self.branch_counts.len(), + "producer StepIdx({}) out of range (graph has {} steps)", + producer.0, + self.branch_counts.len() + ); + let branch_count = self.branch_counts[producer.0]; + assert!( + branch.0 < branch_count, + "branch BranchIdx({}) out of range for producer '{}' with {} branches", + branch.0, + self.step_names[producer.0], + branch_count + ); + // `branch_offsets[producer]` is the precomputed prefix sum of all prior + // steps' branch counts (maintained in `register_step_with_input_arity`), + // so this is O(1) rather than re-summing `branch_counts[..producer.0]`. + self.branch_offsets[producer.0] + branch.0 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn fresh_graph_is_empty() { + let g = ChainGraph::new(); + assert_eq!(g.n_steps(), 0); + assert!(g.first_unwired().is_none()); + } + + #[test] + fn sink_with_no_branches_is_wired() { + let mut g = ChainGraph::new(); + g.register_step("Sink", 0); + assert!(g.first_unwired().is_none()); + } + + #[test] + fn unwired_single_output_step_is_detected() { + let mut g = ChainGraph::new(); + g.register_step("Source", 1); + let (producer, branch, name) = g.first_unwired().unwrap(); + assert_eq!(producer, StepIdx(0)); + assert_eq!(branch, BranchIdx(0)); + assert_eq!(name, "0"); + } + + #[test] + fn wired_chain_is_clean() { + let mut g = ChainGraph::new(); + let src = g.register_step("Source", 1); + let sink = g.register_step("Sink", 0); + g.wire(src, BranchIdx(0), sink); + assert!(g.first_unwired().is_none()); + assert_eq!(g.consumer(src, BranchIdx(0)), Some(sink)); + } + + #[test] + fn unwired_fanout_branch_is_detected() { + let mut g = ChainGraph::new(); + let src = g.register_step("Source", 1); + let mid = g.register_step("FanOut", 2); + let sink = g.register_step("Sink", 0); + g.wire(src, BranchIdx(0), mid); + g.wire(mid, BranchIdx(0), sink); + let (producer, branch, _) = g.first_unwired().unwrap(); + assert_eq!(producer, StepIdx(1)); + assert_eq!(branch, BranchIdx(1)); + } + + #[test] + fn branch_offsets_route_multi_branch_consumers_correctly() { + // Several multi-branch producers in a row: the precomputed branch + // offsets must keep each (producer, branch) edge mapped to a distinct + // flat slot so `consumer`/`consumer_input_slot` return the right links. + let mut g = ChainGraph::new(); + let src = g.register_step("Source", 1); // offset 0, slot 0 + let fan3 = g.register_step("Fan3", 3); // offset 1, slots 1..4 + let fan2 = g.register_step("Fan2", 2); // offset 4, slots 4..6 + let s0 = g.register_step("S0", 0); + let s1 = g.register_step("S1", 0); + let s2 = g.register_step("S2", 0); + let s3 = g.register_step("S3", 0); + + g.wire(src, BranchIdx(0), fan3); + g.wire(fan3, BranchIdx(0), s0); + g.wire(fan3, BranchIdx(1), fan2); + g.wire(fan3, BranchIdx(2), s1); + g.wire(fan2, BranchIdx(0), s2); + g.wire(fan2, BranchIdx(1), s3); + + assert!(g.first_unwired().is_none()); + assert_eq!(g.consumer(src, BranchIdx(0)), Some(fan3)); + assert_eq!(g.consumer(fan3, BranchIdx(0)), Some(s0)); + assert_eq!(g.consumer(fan3, BranchIdx(1)), Some(fan2)); + assert_eq!(g.consumer(fan3, BranchIdx(2)), Some(s1)); + assert_eq!(g.consumer(fan2, BranchIdx(0)), Some(s2)); + assert_eq!(g.consumer(fan2, BranchIdx(1)), Some(s3)); + } + + #[test] + #[should_panic(expected = "already wired")] + fn double_wire_panics() { + let mut g = ChainGraph::new(); + let src = g.register_step("Source", 1); + let sink1 = g.register_step("Sink1", 0); + let sink2 = g.register_step("Sink2", 0); + g.wire(src, BranchIdx(0), sink1); + g.wire(src, BranchIdx(0), sink2); // panics + } + + /// `wire_to_slot` must record which of a multi-input consumer's slots each + /// edge feeds. `contexts::find_all_producers` orders the consumer's input + /// handles by that slot, so a mix-up silently swaps `StepCtx2::a` and + /// `StepCtx2::b` and mis-wires a merge step — with no type error, since both + /// inputs are commonly the same type. + #[test] + fn wire_to_slot_records_each_consumer_input_slot() { + let mut g = ChainGraph::new(); + let a = g.register_step("A", 1); + let b = g.register_step("B", 1); + let join = g.register_step_with_input_arity("Join", 0, 2); + g.wire_to_slot(a, BranchIdx(0), join, 0); + g.wire_to_slot(b, BranchIdx(0), join, 1); + assert_eq!(g.consumer_input_slot(a, BranchIdx(0)), Some(0)); + assert_eq!(g.consumer_input_slot(b, BranchIdx(0)), Some(1)); + assert_eq!(g.input_arity(join), 2); + } + + /// `first_producer_into` backs `PipelineBuilder::build`'s wired-into-a-source + /// check, which is the only guard covering a source appended through + /// `Chain::chain` (registered with a consumer's arity 1, so `wire_to_slot` + /// accepts the edge). + #[test] + fn first_producer_into_finds_the_incoming_edge() { + let mut g = ChainGraph::new(); + let src = g.register_step("Source", 1); + let mid = g.register_step("Mid", 1); + let sink = g.register_step("Sink", 0); + g.wire(src, BranchIdx(0), mid); + g.wire(mid, BranchIdx(0), sink); + + assert_eq!(g.first_producer_into(src), None, "nothing feeds the source"); + assert_eq!(g.first_producer_into(mid), Some(src)); + assert_eq!(g.first_producer_into(sink), Some(mid)); + } + + #[test] + #[should_panic(expected = "out of range")] + fn wire_into_zero_arity_source_panics() { + // Sources register with `input_arity = 0` (their input is implicit), so + // even slot 0 must be rejected — there is no valid input branch to wire. + let mut g = ChainGraph::new(); + let producer = g.register_step("Producer", 1); + let source = g.register_step_with_input_arity("Source", 1, 0); + g.wire(producer, BranchIdx(0), source); // slot 0, but arity 0 → panics + } + + #[test] + #[should_panic(expected = "consumer StepIdx(5) out of range")] + fn wire_to_out_of_range_consumer_panics() { + // A consumer index past the registered steps must produce a + // deterministic range error rather than an opaque out-of-bounds panic. + let mut g = ChainGraph::new(); + let producer = g.register_step("Producer", 1); + g.wire(producer, BranchIdx(0), StepIdx(5)); // no step 5 registered + } +} diff --git a/crates/fgumi-pipeline-core/tests/compile-fail/chain_input_type_mismatch.rs b/crates/fgumi-pipeline-core/tests/compile-fail/chain_input_type_mismatch.rs new file mode 100644 index 000000000..0e693e648 --- /dev/null +++ b/crates/fgumi-pipeline-core/tests/compile-fail/chain_input_type_mismatch.rs @@ -0,0 +1,54 @@ +//! `Chain>::chain` requires `S::Input = u32`. Chaining a +//! step whose `Input = u64` must fail to compile, not at runtime. + +use std::io; + +use fgumi_pipeline_core::PipelineBuilder; +use fgumi_pipeline_core::outputs::Single; +use fgumi_pipeline_core::queues::QueueSpec; +use fgumi_pipeline_core::reorder::BranchOrdering; +use fgumi_pipeline_core::step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}; + +#[derive(Clone)] +struct U32Source; +impl Step for U32Source { + type Input = (); + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "U32Source", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 1 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::Finished) + } +} + +#[derive(Clone)] +struct U64Sink; +impl Step for U64Sink { + type Input = u64; // mismatch with U32Source's Single output + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "U64Sink", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } +} + +fn main() { + let builder = PipelineBuilder::new(); + // Type error: cannot chain U64Sink (Input = u64) onto Chain>. + let _ = builder.chain(U32Source).chain(U64Sink); +} diff --git a/crates/fgumi-pipeline-core/tests/compile-fail/chain_input_type_mismatch.stderr b/crates/fgumi-pipeline-core/tests/compile-fail/chain_input_type_mismatch.stderr new file mode 100644 index 000000000..e96b7ff33 --- /dev/null +++ b/crates/fgumi-pipeline-core/tests/compile-fail/chain_input_type_mismatch.stderr @@ -0,0 +1,21 @@ +error[E0271]: type mismatch resolving `::Input == u32` + --> tests/compile-fail/chain_input_type_mismatch.rs:53:44 + | +53 | let _ = builder.chain(U32Source).chain(U64Sink); + | ----- ^^^^^^^ type mismatch resolving `::Input == u32` + | | + | required by a bound introduced by this call + | +note: expected this to be `u32` + --> tests/compile-fail/chain_input_type_mismatch.rs:34:18 + | +34 | type Input = u64; // mismatch with U32Source's Single output + | ^^^ +note: required by a bound in `fgumi_pipeline_core::Chain::<'b, Single>::chain` + --> src/builder.rs + | + | pub fn chain(self, step: S) -> Chain<'b, S::Outputs> + | ----- required by a bound in this associated function + | where + | S: Step, + | ^^^^^^^^^ required by this bound in `Chain::<'b, Single>::chain` diff --git a/crates/fgumi-pipeline-core/tests/compile-fail/ordered_bytes_single_requires_heapsize.rs b/crates/fgumi-pipeline-core/tests/compile-fail/ordered_bytes_single_requires_heapsize.rs new file mode 100644 index 000000000..d09b209d8 --- /dev/null +++ b/crates/fgumi-pipeline-core/tests/compile-fail/ordered_bytes_single_requires_heapsize.rs @@ -0,0 +1,22 @@ +//! `OrderedBytesSingle` requires `T: HeapSize + Ordered`. A type that +//! impls `Ordered` but NOT `HeapSize` must fail to compile. + +use fgumi_pipeline_core::item::Ordered; +use fgumi_pipeline_core::outputs::OrderedBytesSingle; + +struct OrderedNoHeap { + serial: u64, +} + +impl Ordered for OrderedNoHeap { + fn ordinal(&self) -> u64 { + self.serial + } +} + +fn _instantiate() -> std::marker::PhantomData> { + // Bound `T: HeapSize` is unsatisfied for `OrderedNoHeap`. + std::marker::PhantomData +} + +fn main() {} diff --git a/crates/fgumi-pipeline-core/tests/compile-fail/ordered_bytes_single_requires_heapsize.stderr b/crates/fgumi-pipeline-core/tests/compile-fail/ordered_bytes_single_requires_heapsize.stderr new file mode 100644 index 000000000..7b97cc491 --- /dev/null +++ b/crates/fgumi-pipeline-core/tests/compile-fail/ordered_bytes_single_requires_heapsize.stderr @@ -0,0 +1,53 @@ +error[E0277]: the trait bound `OrderedNoHeap: HeapSize` is not satisfied + --> tests/compile-fail/ordered_bytes_single_requires_heapsize.rs:17:47 + | +17 | fn _instantiate() -> std::marker::PhantomData> { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ unsatisfied trait bound + | +help: the trait `HeapSize` is not implemented for `OrderedNoHeap` + --> tests/compile-fail/ordered_bytes_single_requires_heapsize.rs:7:1 + | + 7 | struct OrderedNoHeap { + | ^^^^^^^^^^^^^^^^^^^^ + = help: the following other types implement trait `HeapSize`: + () + Option + Sequenced + String + Vec + bool + char + f32 + and $N others +note: required by a bound in `OrderedBytesSingle` + --> src/outputs.rs + | + | pub struct OrderedBytesSingle(PhantomData T>); + | ^^^^^^^^ required by this bound in `OrderedBytesSingle` + +error[E0277]: the trait bound `OrderedNoHeap: HeapSize` is not satisfied + --> tests/compile-fail/ordered_bytes_single_requires_heapsize.rs:19:5 + | +19 | std::marker::PhantomData + | ^^^^^^^^^^^^^^^^^^^^^^^^ unsatisfied trait bound + | +help: the trait `HeapSize` is not implemented for `OrderedNoHeap` + --> tests/compile-fail/ordered_bytes_single_requires_heapsize.rs:7:1 + | + 7 | struct OrderedNoHeap { + | ^^^^^^^^^^^^^^^^^^^^ + = help: the following other types implement trait `HeapSize`: + () + Option + Sequenced + String + Vec + bool + char + f32 + and $N others +note: required by a bound in `OrderedBytesSingle` + --> src/outputs.rs + | + | pub struct OrderedBytesSingle(PhantomData T>); + | ^^^^^^^^ required by this bound in `OrderedBytesSingle` diff --git a/crates/fgumi-pipeline-core/tests/compile-fail/ordered_bytes_single_requires_ordered.rs b/crates/fgumi-pipeline-core/tests/compile-fail/ordered_bytes_single_requires_ordered.rs new file mode 100644 index 000000000..2c1d6dfc6 --- /dev/null +++ b/crates/fgumi-pipeline-core/tests/compile-fail/ordered_bytes_single_requires_ordered.rs @@ -0,0 +1,22 @@ +//! `OrderedBytesSingle` requires `T: HeapSize + Ordered`. A type that +//! impls `HeapSize` but NOT `Ordered` must fail to compile. + +use fgumi_pipeline_core::item::HeapSize; +use fgumi_pipeline_core::outputs::OrderedBytesSingle; + +struct HeapNoOrdered { + bytes: Vec, +} + +impl HeapSize for HeapNoOrdered { + fn heap_size(&self) -> usize { + self.bytes.len() + } +} + +fn _instantiate() -> std::marker::PhantomData> { + // Bound `T: Ordered` is unsatisfied for `HeapNoOrdered`. + std::marker::PhantomData +} + +fn main() {} diff --git a/crates/fgumi-pipeline-core/tests/compile-fail/ordered_bytes_single_requires_ordered.stderr b/crates/fgumi-pipeline-core/tests/compile-fail/ordered_bytes_single_requires_ordered.stderr new file mode 100644 index 000000000..6da30361f --- /dev/null +++ b/crates/fgumi-pipeline-core/tests/compile-fail/ordered_bytes_single_requires_ordered.stderr @@ -0,0 +1,43 @@ +error[E0277]: the trait bound `HeapNoOrdered: Ordered` is not satisfied + --> tests/compile-fail/ordered_bytes_single_requires_ordered.rs:17:47 + | +17 | fn _instantiate() -> std::marker::PhantomData> { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ unsatisfied trait bound + | +help: the trait `Ordered` is not implemented for `HeapNoOrdered` + --> tests/compile-fail/ordered_bytes_single_requires_ordered.rs:7:1 + | + 7 | struct HeapNoOrdered { + | ^^^^^^^^^^^^^^^^^^^^ +help: the trait `Ordered` is implemented for `Sequenced` + --> src/reorder.rs + | + | impl Ordered for Sequenced { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +note: required by a bound in `OrderedBytesSingle` + --> src/outputs.rs + | + | pub struct OrderedBytesSingle(PhantomData T>); + | ^^^^^^^ required by this bound in `OrderedBytesSingle` + +error[E0277]: the trait bound `HeapNoOrdered: Ordered` is not satisfied + --> tests/compile-fail/ordered_bytes_single_requires_ordered.rs:19:5 + | +19 | std::marker::PhantomData + | ^^^^^^^^^^^^^^^^^^^^^^^^ unsatisfied trait bound + | +help: the trait `Ordered` is not implemented for `HeapNoOrdered` + --> tests/compile-fail/ordered_bytes_single_requires_ordered.rs:7:1 + | + 7 | struct HeapNoOrdered { + | ^^^^^^^^^^^^^^^^^^^^ +help: the trait `Ordered` is implemented for `Sequenced` + --> src/reorder.rs + | + | impl Ordered for Sequenced { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +note: required by a bound in `OrderedBytesSingle` + --> src/outputs.rs + | + | pub struct OrderedBytesSingle(PhantomData T>); + | ^^^^^^^ required by this bound in `OrderedBytesSingle` diff --git a/crates/fgumi-pipeline-core/tests/compile_fail.rs b/crates/fgumi-pipeline-core/tests/compile_fail.rs new file mode 100644 index 000000000..615eec272 --- /dev/null +++ b/crates/fgumi-pipeline-core/tests/compile_fail.rs @@ -0,0 +1,56 @@ +//! Compile-fail tests for the unified pipeline core. Each `.rs` file under +//! `tests/compile-fail/` is expected to fail to compile with the trait +//! bounds and type checks the framework promises. +//! +//! trybuild covers the type-mismatch compile-fail invariants: that a chain +//! link's `Step::Input` must match its producer's item type +//! (`tests/compile-fail/chain_input_type_mismatch.rs`), and that +//! `OrderedBytesSingle` requires both `HeapSize` and `Ordered` +//! (`ordered_bytes_single_requires_{heapsize,ordered}.rs`). +//! +//! ## The `.stderr` fixtures are whole-diagnostic, not assertions +//! +//! trybuild compares each fixture to rustc's normalized stderr for **equality**; +//! it has no wildcard or subset matching, so a `.stderr` file records the entire +//! diagnostic — including rustc's `the following other types implement trait +//! HeapSize` / `the trait Ordered is implemented for Sequenced` help blocks. +//! Those blocks are an inventory of the crate's impls, not part of the bound +//! under test, so adding a `HeapSize` or `Ordered` impl anywhere rewrites them +//! and fails this test for a reason unrelated to `OrderedBytesSingle`. +//! +//! They cannot simply be deleted from the fixtures — that is a mismatch like any +//! other and fails immediately. So when a diff here is confined to those help +//! blocks (or to rustc's wording), it is a **re-bless**, not a contract +//! regression: check that the `error[E0277]` lines and the `required by a bound +//! in OrderedBytesSingle` notes still say what they should, then regenerate with +//! `TRYBUILD=overwrite cargo nextest run -p fgumi-pipeline-core --test compile_fail`. +//! A diff that touches an `E0277` line or a `required by a bound` note is the +//! real signal and must not be blessed away. +//! +//! That re-bless runs through a bare `cargo nextest run`, not the repo's +//! `cargo ci-test` alias: the alias hardcodes `--workspace`, which wins over the +//! `-p` filter and also selects `fgumi-cli-macros`' own `compile_fail` binary, so +//! blessing through it would rewrite that crate's fixtures in the same pass. + +/// The number of fixtures `tests/compile-fail/` is expected to hold. Asserted +/// before handing the glob to trybuild because `compile_fail` **passes when the +/// glob matches nothing** — so renaming or moving the fixture directory would +/// silently retire every compile-time contract above while CI stayed green. +/// Raise this when adding a fixture. +const EXPECTED_FIXTURES: usize = 3; + +#[test] +fn pipeline_core_compile_fail() { + let fixtures = std::fs::read_dir("tests/compile-fail") + .expect("tests/compile-fail must exist") + .filter_map(Result::ok) + .filter(|e| e.path().extension().is_some_and(|x| x == "rs")) + .count(); + assert!( + fixtures >= EXPECTED_FIXTURES, + "expected at least {EXPECTED_FIXTURES} compile-fail fixtures, found {fixtures}; \ + a zero-match glob makes `compile_fail` a silent no-op" + ); + let t = trybuild::TestCases::new(); + t.compile_fail("tests/compile-fail/*.rs"); +} diff --git a/crates/fgumi-pipeline-io/Cargo.toml b/crates/fgumi-pipeline-io/Cargo.toml new file mode 100644 index 000000000..55176de9a --- /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 = { workspace = true } +noodles = { workspace = true, 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, fgumi-bgzf, and tempfile are omitted here: each is already a +# normal dependency with identical configuration, which tests and benches +# inherit, so a dev-dependency entry would be redundant. +fgumi-raw-bam = { workspace = true, features = ["noodles", "test-utils"] } +fgumi-sort = { workspace = true, features = ["test-utils"] } +noodles = { workspace = true, features = ["bam", "sam"] } +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..95e43db3b --- /dev/null +++ b/crates/fgumi-pipeline-io/src/boundaries.rs @@ -0,0 +1,668 @@ +//! `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; + +/// Upper bound on a single BAM record's `block_size`. The `block_size` prefix is +/// read straight from the input bytes, so a corrupt or hostile value (e.g. +/// `0xFFFF_FFFF`) would otherwise make the scanner buffer the whole unconsumed +/// tail into `leftover` — up to ~4 GiB from a 4-byte corruption — and only fail +/// at [`BoundaryState::finish`]. A record larger than this is corruption, not +/// data; reject it at the point the prefix is read. +const MAX_RECORD_BLOCK_SIZE: usize = 64 * 1024 * 1024; + +/// 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; + if block_size > MAX_RECORD_BLOCK_SIZE { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "BAM record block_size {block_size} exceeds the maximum \ + {MAX_RECORD_BLOCK_SIZE} (corrupt or truncated input)" + ), + )); + } + 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; + + if block_size > MAX_RECORD_BLOCK_SIZE { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "BAM record block_size {block_size} exceeds the maximum \ + {MAX_RECORD_BLOCK_SIZE} (corrupt or truncated input)" + ), + )); + } + 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..9e6cad0e2 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sink/write_bgzf.rs @@ -0,0 +1,516 @@ +//! `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"); + } + + /// Positive coverage for the drained-finish branch of `try_run`, driven + /// through a real 2-step pipeline (block source → sink): the sink must write + /// exactly one trailing `BGZF_EOF` and retire its state. The negative branch + /// (`drop_before_finish_does_not_append_eof_marker`) and the by-hand EOF + /// (`header_only_round_trip`) left this positive path — the one that decides + /// whether an output BAM is a valid, EOF-terminated stream — otherwise unpinned. + #[test] + fn try_run_drained_finish_writes_exactly_one_eof() { + use fgumi_pipeline_core::{ + Unpushed, + builder::{Pipeline, PipelineConfig}, + held::HeldSlot, + outputs::OrderedBytesSingle, + queues::QueueSpec, + reorder::BranchOrdering, + }; + + /// Exclusive source draining a `Vec`, one block per `try_run`. + struct BlockSource { + blocks: Vec, + held: HeldSlot>, + } + impl Step for BlockSource { + type Input = (); + type Outputs = OrderedBytesSingle; + fn profile(&self) -> StepProfile { + StepProfile { + name: "BlockSource", + kind: StepKind::Exclusive, + sticky: true, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: 1 << 20 }], + branch_ordering: vec![BranchOrdering::ByItemOrdinal], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if let Some(unpushed) = self.held.take() + && let Err(again) = ctx.outputs.retry(unpushed) + { + self.held.put(again); + return Ok(StepOutcome::Progress); + } + let Some(block) = self.blocks.pop() else { + return Ok(StepOutcome::Finished); + }; + if let Err(unpushed) = ctx.outputs.push(block) { + self.held.put(unpushed); + } + Ok(StepOutcome::Progress) + } + } + + let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); + let sink = WriteBgzfFile::new(&path, &empty_header(), 1).unwrap(); + + // Two payload blocks with dense ordinals (0, 1). `pop()` drains the tail + // first, so ordinal 0 is emitted before ordinal 1 for the `ByItemOrdinal` + // reorder stage. The bytes are arbitrary non-EOF markers. + let blocks = vec![ + BgzfBlock { batch_serial: 1, bytes: vec![0xAB; 16], uncompressed_size: 0 }, + BgzfBlock { batch_serial: 0, bytes: vec![0xCD; 16], uncompressed_size: 0 }, + ]; + let source = BlockSource { blocks, held: HeldSlot::new() }; + + let builder = Pipeline::builder(); + builder.chain(source).chain(sink).into_sink_marker(); + let pipeline = builder.build().expect("pipeline builds"); + pipeline + .run(PipelineConfig { threads: 2, ..Default::default() }) + .expect("pipeline runs to completion"); + + let bytes = std::fs::read(&path).unwrap(); + assert!(bytes.len() >= 28, "header + payload + EOF present"); + assert_eq!( + &bytes[bytes.len() - 28..], + &BGZF_EOF, + "the drained-finish path terminates the stream with a BGZF EOF marker" + ); + // Exactly one trailing EOF: stripping it must not reveal a second, which + // is what a stray `Drop`-side append would produce. + let without_eof = &bytes[..bytes.len() - 28]; + assert!( + without_eof.len() < 28 || without_eof[without_eof.len() - 28..] != BGZF_EOF, + "only the drained-finish path may append EOF; Drop must not add a second" + ); + // Both payload blocks reached disk. + assert!(bytes.windows(16).any(|w| w == [0xAB; 16]), "first block written"); + assert!(bytes.windows(16).any(|w| w == [0xCD; 16]), "second block written"); + } + + #[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..d89e53845 --- /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, + runs_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, + runs_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..889144e9e --- /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 advanced = + frame_one_block(&chunk, next, BGZF_MAX_BLOCK_SIZE, &mut out).expect("template framing"); + assert!(advanced > next, "framing must make progress on every call"); + next = advanced; + 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..9516cb6e3 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/source/read_bam.rs @@ -0,0 +1,454 @@ +//! `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; + +/// `Serial + 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)?; + // Honor `async_reader` for the reopened raw block stream too — not just the + // temporary header reader above — so a regular file prefetches like the + // stdin path (`read_bam_stdin`) does. `verify_crc` does not apply here: + // `ReadBgzfBlocks` forwards compressed bytes without decoding them. + let reader: Box = if opts.async_reader { + log::info!("async read enabled: spawning fgumi-prefetch thread for {}", path.display()); + Box::new(fgumi_bam_io::prefetch_reader::PrefetchReader::from_file(file)) + } else { + 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..bb8ab245a --- /dev/null +++ b/crates/fgumi-pipeline-io/src/types.rs @@ -0,0 +1,315 @@ +//! 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. + /// + /// Each range must satisfy `start <= end <= backing.len()`; callers own that + /// invariant (the boundary scans that produce these ranges already validate + /// it). Violating it panics later in + /// [`iter_record_bytes`](Self::iter_record_bytes) when the range is sliced, + /// far from the site that produced it — the `debug_assert!` pins the failure + /// at the constructor boundary instead. Note that `total_bytes` reports + /// `backing.len()`, which differs from the sum of the record bodies if the + /// ranges do not cover the whole buffer. + #[must_use] + pub fn from_parsed(batch_serial: u64, backing: Vec, ranges: Vec<(u32, u32)>) -> Self { + debug_assert!( + ranges.iter().all(|&(s, e)| s <= e && e as usize <= backing.len()), + "RecordBatch::from_parsed: range outside backing buffer" + ); + 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-raw-bam/src/lib.rs b/crates/fgumi-raw-bam/src/lib.rs index 44c5045fb..9e460d001 100644 --- a/crates/fgumi-raw-bam/src/lib.rs +++ b/crates/fgumi-raw-bam/src/lib.rs @@ -127,9 +127,10 @@ pub use tags::{ array_tag_element_u16, array_tag_to_vec_u16, extract_aux_string_tags, extract_int_value, extract_template_aux_tags, find_array_tag, find_float_tag, find_int_tag, find_mc_tag_in_record, find_string_tag, find_string_tag_in_record, find_string_tag_position, find_tag_type, - find_uint8_tag, normalize_int_tag_to_smallest_signed, read_tc_template_coordinate, remove_tag, - reverse_array_tag_in_place, reverse_complement_string_tag_in_place, - reverse_string_tag_in_place, update_int_tag, update_string_tag, + find_two_string_tags_in_record, find_uint8_tag, normalize_int_tag_to_smallest_signed, + read_tc_template_coordinate, remove_tag, reverse_array_tag_in_place, + reverse_complement_string_tag_in_place, reverse_string_tag_in_place, update_int_tag, + update_string_tag, }; #[cfg(any(test, feature = "test-utils"))] diff --git a/crates/fgumi-raw-bam/src/tags.rs b/crates/fgumi-raw-bam/src/tags.rs index 1d68a5c37..5a708ff05 100644 --- a/crates/fgumi-raw-bam/src/tags.rs +++ b/crates/fgumi-raw-bam/src/tags.rs @@ -95,6 +95,73 @@ pub fn find_string_tag_in_record(bam: &[u8], tag: impl AsTagBytes) -> Option<&[u find_string_tag(aux, tag) } +/// Find two string (Z-type) tags in a complete BAM record in a **single** aux-data +/// walk, returning each value's bytes (without the NUL terminator) or `None` if that +/// tag is absent or not Z-type. +/// +/// Equivalent to calling [`find_string_tag_in_record`] twice, but it walks the aux +/// block only once — relevant on hot paths that need both an MI tag and a cell-barcode +/// tag per record. The returned slices borrow the record's aux data. +#[inline] +#[must_use] +pub fn find_two_string_tags_in_record( + bam: &[u8], + first: impl AsTagBytes, + second: impl AsTagBytes, +) -> (Option<&[u8]>, Option<&[u8]>) { + let first = u16::from_le_bytes(*first.as_tag_bytes()); + let second = u16::from_le_bytes(*second.as_tag_bytes()); + let aux = aux_data_slice(bam); + + let mut first_done = false; + let mut second_done = false; + let mut first_val: Option<&[u8]> = None; + let mut second_val: Option<&[u8]> = None; + let mut p = 0; + while p + 3 <= aux.len() { + let entry_u16 = u16::from_le_bytes([aux[p], aux[p + 1]]); + let val_type = aux[p + 2]; + + // Each tag resolves on its FIRST matching entry, exactly as + // `find_string_tag_in_record` does (it returns the first tag match and + // yields a value only if that match is a NUL-terminated `Z` entry). A + // non-`Z` (or unterminated) first match resolves the tag to `None`; we + // must NOT skip it and pick up a later `Z` duplicate, or the two-tag + // walk would disagree with two single-tag lookups on malformed aux. + let matches_first = entry_u16 == first && !first_done; + let matches_second = entry_u16 == second && !second_done; + if matches_first || matches_second { + let value = if val_type == b'Z' { + let start = p + 3; + aux[start..].iter().position(|&b| b == 0).map(|len| &aux[start..start + len]) + } else { + None + }; + // Independent (not `else`) so a single entry resolves both slots + // when `first == second`, preserving the "call twice" contract. + if matches_first { + first_done = true; + first_val = value; + } + if matches_second { + second_done = true; + second_val = value; + } + } + + // Stop once both tags have been resolved (each to a value or `None`). + if first_done && second_done { + break; + } + + match tag_value_size(val_type, &aux[p + 3..]) { + Some(size) => p += 3 + size, + None => break, + } + } + (first_val, second_val) +} + /// Find the byte range `[start, end)` of an entire tag entry (tag+type+value) in aux data. /// /// Returns offsets relative to the start of `aux_data`. @@ -4526,4 +4593,134 @@ mod tests { assert_eq!(find_string_tag(&record, SamTag::CB), Some(b"ACGT".as_ref())); assert_eq!(&record[8..], b"CBZACGT\0", "verbatim tag+type+value bytes"); } + + #[test] + fn test_find_two_string_tags_both_present_either_order() { + // MI before CB. + let aux = b"MIZ7\x00CBZACGT\x00"; + let rec = make_bam_bytes(0, 0, 0, b"rea", &[], 0, -1, -1, aux); + let (mi, cb) = find_two_string_tags_in_record(&rec, SamTag::MI, SamTag::CB); + assert_eq!(mi, Some(b"7".as_ref())); + assert_eq!(cb, Some(b"ACGT".as_ref())); + + // CB before MI: result is independent of aux ordering and of arg order. + let aux = b"CBZACGT\x00MIZ7\x00"; + let rec = make_bam_bytes(0, 0, 0, b"rea", &[], 0, -1, -1, aux); + let (mi, cb) = find_two_string_tags_in_record(&rec, SamTag::MI, SamTag::CB); + assert_eq!(mi, Some(b"7".as_ref())); + assert_eq!(cb, Some(b"ACGT".as_ref())); + } + + #[test] + fn test_find_two_string_tags_agrees_with_single_lookups() { + // Interleave a non-Z tag to make the walk non-trivial. + let mut aux = Vec::new(); + aux.extend_from_slice(b"NMC"); + aux.push(3); + aux.extend_from_slice(b"MIZ42\x00"); + aux.extend_from_slice(b"CBZTGCA\x00"); + let rec = make_bam_bytes(0, 0, 0, b"rea", &[], 0, -1, -1, &aux); + let (mi, cb) = find_two_string_tags_in_record(&rec, SamTag::MI, SamTag::CB); + assert_eq!(mi, find_string_tag_in_record(&rec, SamTag::MI)); + assert_eq!(cb, find_string_tag_in_record(&rec, SamTag::CB)); + } + + #[test] + fn test_find_two_string_tags_first_match_wins_on_malformed_dup() { + // A non-Z MI entry precedes a Z MI duplicate. `find_string_tag_in_record` + // resolves the FIRST MI match (non-Z → None) and never reaches the later + // Z duplicate; the two-tag walk must match that rather than skipping the + // non-Z entry and returning the later Z value. + let mut aux = Vec::new(); + aux.extend_from_slice(b"MIC"); // MI, type 'C' (uint8) ... + aux.push(9); // ... 1-byte value (non-Z → MI resolves None) + aux.extend_from_slice(b"MIZ77\x00"); // later Z duplicate — must be ignored + aux.extend_from_slice(b"CBZACGT\x00"); + let rec = make_bam_bytes(0, 0, 0, b"rea", &[], 0, -1, -1, &aux); + + let (mi, cb) = find_two_string_tags_in_record(&rec, SamTag::MI, SamTag::CB); + assert_eq!(mi, None, "first (non-Z) MI match resolves the tag to None"); + assert_eq!(cb, Some(b"ACGT".as_ref())); + // Must agree with two independent single-tag lookups. + assert_eq!(mi, find_string_tag_in_record(&rec, SamTag::MI)); + assert_eq!(cb, find_string_tag_in_record(&rec, SamTag::CB)); + } + + /// A `Z` entry with no terminating NUL, in the two shapes that behave + /// differently. + /// + /// The doc contract says a first match that is non-`Z` *or unterminated* + /// resolves the tag to `None`. Only the non-`Z` half was pinned, and the + /// unterminated half is subtler than it reads: the value scan looks for the + /// next NUL anywhere in the remaining aux block, not just within the entry. + /// So "unterminated" resolves to `None` only when there is no NUL left at + /// all — if a later entry supplies one, the value runs *through* that entry + /// and swallows it. + /// + /// Both shapes are pinned because they fail differently: the first stops the + /// walk via `tag_value_size` returning `None`, the second consumes the + /// remaining bytes as one oversized value. Either way the two-tag walk must + /// agree with two independent single-tag lookups, which is the contract that + /// makes `find_two_string_tags_in_record` a safe substitution for calling + /// `find_string_tag_in_record` twice. + #[rstest] + // No NUL anywhere after the tag: the value scan finds nothing, so MI is + // `None`, and the walk stops rather than advancing past a sizeless entry. + #[case::no_nul_anywhere(b"MIZ77".as_ref(), None, None)] + // A later entry supplies the NUL, so MI's value runs through `CBZACGT` and + // consumes it — CB is never seen as an entry of its own. + #[case::nul_supplied_by_a_later_entry( + b"MIZ77CBZACGT\x00".as_ref(), + Some(b"77CBZACGT".as_ref()), + None + )] + fn test_find_two_string_tags_unterminated_z_first_match( + #[case] aux: &[u8], + #[case] expected_mi: Option<&[u8]>, + #[case] expected_cb: Option<&[u8]>, + ) { + let rec = make_bam_bytes(0, 0, 0, b"rea", &[], 0, -1, -1, aux); + + let (mi, cb) = find_two_string_tags_in_record(&rec, SamTag::MI, SamTag::CB); + assert_eq!(mi, expected_mi); + assert_eq!(cb, expected_cb); + // The parity contract: identical to two independent single-tag lookups. + assert_eq!( + mi, + find_string_tag_in_record(&rec, SamTag::MI), + "MI must match a single lookup" + ); + assert_eq!( + cb, + find_string_tag_in_record(&rec, SamTag::CB), + "CB must match a single lookup" + ); + } + + #[test] + fn test_find_two_string_tags_same_tag_fills_both() { + // When the same tag is requested for both first and second, a single + // matching aux entry must populate both outputs (not just one). + let aux = b"MIZ7\x00CBZACGT\x00"; + let rec = make_bam_bytes(0, 0, 0, b"rea", &[], 0, -1, -1, aux); + let (a, b) = find_two_string_tags_in_record(&rec, SamTag::MI, SamTag::MI); + assert_eq!(a, Some(b"7".as_ref())); + assert_eq!(b, Some(b"7".as_ref())); + // Both outputs must agree with an independent single-tag lookup, proving + // the same-tag case behaves like two separate `find_string_tag_in_record` + // calls (the contract the `mi_group.rs` production path relies on). + let mi = find_string_tag_in_record(&rec, SamTag::MI); + assert_eq!(a, mi, "first output must match an independent MI lookup"); + assert_eq!(b, mi, "second output must match an independent MI lookup"); + + // Same-tag MISS path: when the (identical) requested tag is absent, both + // outputs must be None — the zero-match branch of the same-tag case, + // matching two independent single-tag lookups that each miss. + let rec_without_mi = make_bam_bytes(0, 0, 0, b"rea", &[], 0, -1, -1, b"CBZACGT\x00"); + let (a_missing, b_missing) = + find_two_string_tags_in_record(&rec_without_mi, SamTag::MI, SamTag::MI); + assert_eq!(a_missing, None); + assert_eq!(b_missing, None); + assert_eq!(a_missing, find_string_tag_in_record(&rec_without_mi, SamTag::MI)); + } } diff --git a/crates/fgumi-sort/Cargo.toml b/crates/fgumi-sort/Cargo.toml index 018280508..160f1c355 100644 --- a/crates/fgumi-sort/Cargo.toml +++ b/crates/fgumi-sort/Cargo.toml @@ -13,6 +13,12 @@ keywords = ["bioinformatics", "bam", "sort", "sequencing", "ngs"] [lints.clippy] pedantic = { level = "deny", priority = -1 } +[lints.rust] +# `loom` is a config flag set only for the model-checking test target +# (`tests/loom_merge_slots.rs`, run via `RUSTFLAGS="--cfg loom"`). Declare it so +# the unexpected-cfg lint stays quiet under the normal `-D warnings` build. +unexpected_cfgs = { level = "warn", check-cfg = ['cfg(loom)'] } + [dependencies] fgumi-raw-bam = { workspace = true, features = ["noodles"] } fgumi-bgzf = { workspace = true } @@ -35,7 +41,12 @@ bytes = { workspace = true } bytesize = { workspace = true } fs4 = { workspace = true } log = { workspace = true } +# `smallvec` is crate-local, not workspace-inherited: `keys.rs` is its only user +# in the workspace, and the root manifest reserves `[workspace.dependencies]` for +# crates shared by two or more members. +smallvec = { version = "1.15", features = ["const_generics"] } zstd = "0.13" +voracious_radix_sort = { version = "1", features = ["voracious_multithread"] } [target.'cfg(unix)'.dependencies] libc = "0.2" @@ -51,6 +62,7 @@ nix = { workspace = true } [features] default = [] memory-debug = [] +test-utils = [] [dev-dependencies] fgumi-bam-io = { workspace = true } @@ -61,6 +73,19 @@ proptest = { workspace = true } criterion = { workspace = true } tempfile = { workspace = true } +# loom is only compiled under `--cfg loom`. It is a *regular* (not dev) +# dependency so `merge_slots.rs` itself can swap `std::sync` -> `loom::sync` +# under `#[cfg(loom)]` and have its REAL atomics/locks model-checked by +# `tests/loom_merge_slots.rs` (which drives the real `SortMergeSlot`). The test +# is `#![cfg(loom)]` (empty under a normal build). +# Run with: RUSTFLAGS="--cfg loom" cargo test -p fgumi-sort --test loom_merge_slots --release +[target.'cfg(loom)'.dependencies] +loom = "0.7" + +[[bench]] +name = "chunk_sorter" +harness = false + [[bench]] -name = "template_key" +name = "queryname_keys" harness = false diff --git a/crates/fgumi-sort/benches/chunk_sorter.rs b/crates/fgumi-sort/benches/chunk_sorter.rs new file mode 100644 index 000000000..ed020b5ae --- /dev/null +++ b/crates/fgumi-sort/benches/chunk_sorter.rs @@ -0,0 +1,76 @@ +//! Microbenchmark for the P6 `SortBuffer` chunk sorters' sort + materialize +//! step (`take_sorted_chunk`). +//! +//! This isolates the cost the A/B smoke campaign flagged: `take_sorted_chunk` +//! materialized a full owned `Vec<(K, RawRecord)>` (one heap copy + malloc per +//! record) while the source `RecordBuffer` arena was still alive — doubling +//! peak memory and paying a full per-record copy. The fix routes coordinate / +//! template through the zero-copy `InMemoryChunk` (arena-move). Run before and +//! after to confirm the materialize step gets cheaper: +//! +//! cargo bench -p fgumi-sort --bench chunk_sorter + +use criterion::{BatchSize, Criterion, Throughput, criterion_group, criterion_main}; +use fgumi_raw_bam::testutil::make_bam_bytes; +use fgumi_sort::{RawExternalSorter, SortOrder}; +use noodles::sam::Header; + +/// Build `n` aligned records at scattered positions on tid 0 so the sort does +/// real reordering work; ~100 bp reads ≈ a realistic per-record size. +fn synth_records(n: usize) -> Vec> { + (0..n) + .map(|i| { + let pos = (i as u64).wrapping_mul(2_654_435_761) % 5_000_000; + let name = format!("r{i:08}"); + make_bam_bytes(0, pos as i32, 0, name.as_bytes(), &[], 100, -1, -1, &[]) + }) + .collect() +} + +fn bench_materialize(c: &mut Criterion) { + const N: usize = 500_000; + let records = synth_records(N); + + let mut group = c.benchmark_group("chunk_sorter_materialize"); + group.throughput(Throughput::Elements(N as u64)); + group.sample_size(10); + + group.bench_function("coordinate_500k", |b| { + b.iter_batched( + || { + let mut sorter = RawExternalSorter::new(SortOrder::Coordinate) + .threads(4) + .into_coordinate_chunk_sorter(&Header::default()) + .expect("build coordinate chunk sorter"); + for r in &records { + sorter.push(r).expect("push"); + } + sorter + }, + |mut sorter| std::hint::black_box(sorter.take_sorted_chunk()), + BatchSize::PerIteration, + ); + }); + + group.bench_function("template_coordinate_500k", |b| { + b.iter_batched( + || { + let mut sorter = RawExternalSorter::new(SortOrder::TemplateCoordinate) + .threads(4) + .into_template_chunk_sorter(&Header::default()) + .expect("build template chunk sorter"); + for r in &records { + sorter.push(r).expect("push"); + } + sorter + }, + |mut sorter| std::hint::black_box(sorter.take_sorted_chunk()), + BatchSize::PerIteration, + ); + }); + + group.finish(); +} + +criterion_group!(benches, bench_materialize); +criterion_main!(benches); diff --git a/crates/fgumi-sort/benches/queryname_keys.rs b/crates/fgumi-sort/benches/queryname_keys.rs new file mode 100644 index 000000000..7dd42d3da --- /dev/null +++ b/crates/fgumi-sort/benches/queryname_keys.rs @@ -0,0 +1,197 @@ +//! Microbenchmarks isolating the queryname sort's per-record key cost, to +//! decide two things before touching the production queryname path: +//! +//! 1. **Prefix key (solution C)** — does packing the first 8 name bytes into a +//! `u64` for a cheap first-pass compare actually help? The answer depends +//! entirely on the name distribution: Illumina read names share a long +//! common prefix (`instrument:run:flowcell:lane:` is identical for every +//! read in a lane) and only vary in the `tile:x:y` suffix — so a *front* +//! prefix key may discriminate nothing and add pure overhead. We therefore +//! bench with REALISTIC Illumina-style names, not random ones. +//! +//! 2. **Null-key ceilings** — the wall-time floor if the key were free: +//! - `*_ingest`: owned-name extraction (alloc + memcpy per record) vs a +//! cheap `u64` extraction (no alloc). The delta is the ceiling on what +//! solution A (borrow the name from the arena, drop the per-record Vec) +//! can recover on the ingest side. +//! - `*_sort`: full-name comparator vs a cheap `u64` comparator. The delta +//! is the ceiling on what a cheaper comparator (C, or radix on a packed +//! key) can recover on the sort side. +//! +//! cargo bench -p fgumi-sort --bench queryname_keys + +use std::cmp::Ordering; + +use criterion::{BatchSize, Criterion, Throughput, criterion_group, criterion_main}; +use fgumi_raw_bam::testutil::make_bam_bytes; +use fgumi_sort::{RawQuerynameKey, RawQuerynameLexKey, RawSortKey}; + +/// Number of records per bench iteration. +const N: usize = 1_000_000; + +/// Build `n` realistic Illumina-style read names: a fixed +/// `instrument:run:flowcell:lane:` prefix shared by every read (22 bytes), then +/// a varying `tile:x:y` suffix. This is the adversarial case for a front-prefix +/// key — the first ~22 bytes are identical across all records. +fn illumina_names(n: usize) -> Vec> { + (0..n) + .map(|i| { + // Vary tile (1101..1112), x (0..) and y (0..) like a real lane. + let tile = 1101 + (i % 12); + let x = (i.wrapping_mul(2_654_435_761)) % 30_000; + let y = (i.wrapping_mul(40_503)) % 30_000; + format!("A00123:45:HGVWXDSXY:1:{tile}:{x}:{y}").into_bytes() + }) + .collect() +} + +/// Pack the first 8 bytes of `name` into a big-endian `u64` so numeric ordering +/// of the u64 matches lexicographic ordering of the first 8 bytes (short names +/// zero-pad on the right). +#[inline] +fn prefix8(name: &[u8]) -> u64 { + let mut buf = [0u8; 8]; + let take = name.len().min(8); + buf[..take].copy_from_slice(&name[..take]); + u64::from_be_bytes(buf) +} + +/// Build BAM records carrying `names` (paired R1 flag), so key extraction walks +/// a real record body exactly as production does. +fn bam_records(names: &[Vec]) -> Vec> { + names.iter().map(|name| make_bam_bytes(0, 0, 0, name, &[], 100, -1, -1, &[])).collect() +} + +fn bench_ingest(c: &mut Criterion) { + let names = illumina_names(N); + let records = bam_records(&names); + + let mut group = c.benchmark_group("queryname_ingest_1M"); + group.throughput(Throughput::Elements(N as u64)); + group.sample_size(10); + + // Owned natural key: alloc + memcpy of the name per record (current path). + group.bench_function("natural_owned_extract", |b| { + b.iter(|| { + let keys: Vec = + records.iter().map(|r| RawQuerynameKey::extract_from_record(r)).collect(); + std::hint::black_box(keys) + }); + }); + + // Owned lex key: same alloc + memcpy. + group.bench_function("lex_owned_extract", |b| { + b.iter(|| { + let keys: Vec = + records.iter().map(|r| RawQuerynameLexKey::extract_from_record(r)).collect(); + std::hint::black_box(keys) + }); + }); + + // Null key: a cheap u64 prefix, NO allocation. The delta vs the owned + // extracts is the ceiling on solution A (borrow, drop the per-record Vec). + group.bench_function("null_u64_extract", |b| { + b.iter(|| { + let keys: Vec = records + .iter() + .map(|r| { + // True no-alloc floor: walk the raw name bytes directly + // (`l_read_name` at byte 8, name at offset 32) with no + // `RawQuerynameLexKey`/`NameBuf` allocation, then take the u64 + // prefix. `fields::read_name` yields the same name bytes + // `RawQuerynameLexKey::name()` does, so the only delta vs + // `lex_owned_extract` is the alloc + memcpy this path skips — + // a valid lower bound for solution A. + prefix8(fgumi_raw_bam::fields::read_name(r)) + }) + .collect(); + std::hint::black_box(keys) + }); + }); + + group.finish(); +} + +fn bench_sort(c: &mut Criterion) { + let names = illumina_names(N); + let records = bam_records(&names); + + // Pre-build the three key representations once, outside the timed loop. + let nat_keys: Vec = + records.iter().map(|r| RawQuerynameKey::extract_from_record(r)).collect(); + let lex_keys: Vec = + records.iter().map(|r| RawQuerynameLexKey::extract_from_record(r)).collect(); + // Prefix + full: (u64 prefix, full lex key). Prefix compared first, full + // name only on prefix ties. + let prefix_keys: Vec<(u64, RawQuerynameLexKey)> = records + .iter() + .map(|r| { + let k = RawQuerynameLexKey::extract_from_record(r); + (prefix8(k.name()), k) + }) + .collect(); + let u64_keys: Vec = prefix_keys.iter().map(|(p, _)| *p).collect(); + + let mut group = c.benchmark_group("queryname_sort_1M"); + group.throughput(Throughput::Elements(N as u64)); + group.sample_size(10); + + // Full natural comparator (samtools strnum_cmp) — current path. + group.bench_function("natural_full_cmp", |b| { + b.iter_batched( + || nat_keys.clone(), + |mut keys| { + keys.sort_unstable(); + std::hint::black_box(keys) + }, + BatchSize::PerIteration, + ); + }); + + // Full lex comparator (Vec::cmp) — current lex path. + group.bench_function("lex_full_cmp", |b| { + b.iter_batched( + || lex_keys.clone(), + |mut keys| { + keys.sort_unstable(); + std::hint::black_box(keys) + }, + BatchSize::PerIteration, + ); + }); + + // Prefix-then-full: cheap u64 compare first, full name only on prefix ties. + // On Illumina names the shared front prefix means this almost always falls + // through to the full compare — the bench proves whether it helps or hurts. + group.bench_function("lex_prefix_then_full_cmp", |b| { + b.iter_batched( + || prefix_keys.clone(), + |mut keys| { + keys.sort_unstable_by(|a, b| match a.0.cmp(&b.0) { + Ordering::Equal => a.1.cmp(&b.1), + other => other, + }); + std::hint::black_box(keys) + }, + BatchSize::PerIteration, + ); + }); + + // Null key: sort the u64 prefix alone. The comparator ceiling — no full + // name ever compared. + group.bench_function("null_u64_cmp", |b| { + b.iter_batched( + || u64_keys.clone(), + |mut keys| { + keys.sort_unstable(); + std::hint::black_box(keys) + }, + BatchSize::PerIteration, + ); + }); + + group.finish(); +} + +criterion_group!(benches, bench_ingest, bench_sort); +criterion_main!(benches); diff --git a/crates/fgumi-sort/benches/template_key.rs b/crates/fgumi-sort/benches/template_key.rs deleted file mode 100644 index 817ca90b1..000000000 --- a/crates/fgumi-sort/benches/template_key.rs +++ /dev/null @@ -1,134 +0,0 @@ -//! Per-record cost of the template-coordinate ingest path. -//! -//! Phase 1 of a spill-heavy sort is bound by its serial main thread: on -//! `1kg-wgs-HG00096` at 16 threads that thread is busy 219.5s of a 240.7s -//! phase, and `perf` puts 39.5% of it in `extract_template_key_inline` with a -//! further 12.3% in the slice-iterator `next()` the aux-tag scan walks bytes -//! with. This bench sizes that function directly, so a change to it is measured -//! before a seven-minute whole-genome run is spent on it. -//! -//! The records are built to the aux layout the measured sample actually carries -//! (`PG:Z AS:i XS:i MD:Z NM:i RG:Z MQ:i MC:Z`, 118 bytes) rather than a minimal -//! one: the scan's cost is per aux byte, so a record with two tags would report -//! a number that has nothing to do with the workload. `with_xa` adds the long -//! `XA:Z` alt-hit tag that a fraction of records carry, since it roughly doubles -//! the aux data and is where the scan's tail lives. - -use bstr::BString; -use criterion::{Criterion, criterion_group, criterion_main}; -use fgumi_raw_bam::SamTag; -use noodles::sam::Header; -use noodles::sam::header::record::value::Map; -use noodles::sam::header::record::value::map::ReadGroup; -use noodles::sam::header::record::value::map::read_group::tag as rg_tag; -use std::hint::black_box; - -const READ_NAME: &[u8] = b"A00132:53:HFHJKDSXX:1:1646:26467:33332\0"; -const RG_ID: &str = "HG00096_CGGACAAC-TCCGGATT_HFHJKDSXX_L001"; -const SEQ_LEN: usize = 151; - -/// Append a `Z`-typed aux tag (two tag bytes, `Z`, value, NUL). -fn push_z(aux: &mut Vec, tag: SamTag, value: &[u8]) { - aux.extend_from_slice(&*tag); - aux.push(b'Z'); - aux.extend_from_slice(value); - aux.push(0); -} - -/// Append an `i`-typed (32-bit signed) aux tag. -fn push_i(aux: &mut Vec, tag: SamTag, value: i32) { - aux.extend_from_slice(&*tag); - aux.push(b'i'); - aux.extend_from_slice(&value.to_le_bytes()); -} - -/// One BAM record body (from `ref_id`, i.e. without the `block_size` prefix), -/// carrying the aux layout of the measured 1000 Genomes WGS sample. -fn record(pos: i32, mate_pos: i32, with_xa: bool) -> Vec { - let mut aux = Vec::with_capacity(256); - push_z(&mut aux, SamTag::PG, b"MarkDuplicates"); - push_i(&mut aux, SamTag::AS, 64); - push_i(&mut aux, SamTag::XS, 61); - push_z(&mut aux, SamTag::MD, b"0N0N0N0N0N2A61"); - push_i(&mut aux, SamTag::NM, 6); - push_z(&mut aux, SamTag::RG, RG_ID.as_bytes()); - push_i(&mut aux, SamTag::MQ, 0); - push_z(&mut aux, SamTag::MC, b"81S69M"); - if with_xa { - push_z( - &mut aux, - SamTag::new(b'X', b'A'), - b"chr3,+198173832,34M116S,0;chr12,-108091,117S33M,0;chr1,-180805,118S32M,0;\ - chr1,-10052,118S32M,0;chr3,-10519,118S32M,0;", - ); - } - - let n_cigar_op: u16 = 2; - let mut rec = Vec::with_capacity(32 + READ_NAME.len() + 8 + SEQ_LEN * 2 + aux.len()); - rec.extend_from_slice(&0i32.to_le_bytes()); // ref_id - rec.extend_from_slice(&pos.to_le_bytes()); // pos - rec.push(u8::try_from(READ_NAME.len()).expect("read name fits a u8")); - rec.push(60); // mapq - rec.extend_from_slice(&0u16.to_le_bytes()); // bin - rec.extend_from_slice(&n_cigar_op.to_le_bytes()); - // PAIRED | PROPER_PAIR | REVERSE | FIRST_SEGMENT. The mate is forward, so - // the mate lane resolves through `unclipped_other_start` (leading clips). - rec.extend_from_slice(&0x0053u16.to_le_bytes()); - rec.extend_from_slice(&i32::try_from(SEQ_LEN).expect("seq len fits").to_le_bytes()); - rec.extend_from_slice(&0i32.to_le_bytes()); // next_ref_id - rec.extend_from_slice(&mate_pos.to_le_bytes()); - rec.extend_from_slice(&0i32.to_le_bytes()); // tlen - rec.extend_from_slice(READ_NAME); - // CIGAR 81S69M: (len << 4) | op, S = 4, M = 0. - rec.extend_from_slice(&((81u32 << 4) | 4).to_le_bytes()); - rec.extend_from_slice(&(69u32 << 4).to_le_bytes()); - rec.resize(rec.len() + SEQ_LEN.div_ceil(2), 0x11); // packed seq - rec.resize(rec.len() + SEQ_LEN, 30); // qual - rec.extend_from_slice(&aux); - rec -} - -/// A header declaring the sample's single read group, so the RG lookup resolves -/// rather than missing (a miss and a hit take different paths through the map). -fn header() -> Header { - let read_group = Map::::builder() - .insert(rg_tag::LIBRARY, String::from("lib1")) - .build() - .expect("a read group with only LB is valid"); - Header::builder().add_read_group(BString::from(RG_ID), read_group).build() -} - -fn bench_template_key(c: &mut Criterion) { - let lookup = fgumi_sort::LibraryLookup::from_header(&header()); - let hasher = fgumi_sort::cb_hasher(); - - for (label, with_xa) in [("aux-122b", false), ("aux-with-xa", true)] { - // A batch rather than one record: the scan is memory-bound over the aux - // bytes, and re-reading one cache-hot record would flatter it. - let records: Vec> = - (0..1024).map(|i| record(10_000 + i, 133_000_000 + i, with_xa)).collect(); - let aux_bytes: usize = - records.iter().map(|r| fgumi_raw_bam::aux_data_slice(r).len()).sum::() - / records.len(); - - let mut group = c.benchmark_group("extract_template_key_inline"); - group.throughput(criterion::Throughput::Elements(records.len() as u64)); - group.bench_function(format!("{label}-mean-aux-{aux_bytes}"), |b| { - b.iter(|| { - for rec in &records { - let key = fgumi_sort::extract_template_key_inline( - black_box(rec.as_slice()), - &lookup, - None, - &hasher, - ); - black_box(&key); - } - }); - }); - group.finish(); - } -} - -criterion_group!(benches, bench_template_key); -criterion_main!(benches); diff --git a/crates/fgumi-sort/examples/arena_bench.rs b/crates/fgumi-sort/examples/arena_bench.rs new file mode 100644 index 000000000..ba8ef8f74 --- /dev/null +++ b/crates/fgumi-sort/examples/arena_bench.rs @@ -0,0 +1,255 @@ +//! Arena-lifetime microbenchmark — settles N_arena (in-flight sort arenas) and +//! the gather/compress release point, isolated from the typed-step pipeline. +// Throwaway diagnostic bench (not production): relax pedantic lints. +#![allow( + clippy::doc_markdown, + clippy::too_many_lines, + clippy::map_unwrap_or, + clippy::cast_possible_truncation +)] +//! +//! Models the proposed Phase-1 dataflow with a **reused arena pool** (no +//! mem-take-and-mint): a producer fills + par-sorts one `RecordBuffer` to a +//! memory limit, hands ownership to a consumer that gathers (frames sorted +//! records into ≤64 KiB blocks) + compresses, then returns the buffer to the +//! pool. The pool depth is `N_arena`; the release point is `split` +//! (return the buffer right after gather, before compress) or `fused` +//! (return only after compress). +//! +//! Uses the REAL kernels: `RecordBuffer` (fill/par_sort/refs/get_record), +//! `frame_keyed_record_into`, `SpillBlockCompressor::compress_block`. +//! +//! Run one config per process (clean peak RSS): +//! cargo run --release -p fgumi-sort --example arena_bench -- +//! e.g. cargo run --release -p fgumi-sort --example arena_bench -- split 1 512 6 4 +//! +//! Output (stdout, one TSV line): mode N_arena chunk_mib n_chunks workers wall_s fill_s sort_s gather_s compress_s +//! Peak RSS is sampled externally by the runner wrapper (this process prints its PID first). + +use std::sync::mpsc::sync_channel; +use std::time::{Duration, Instant}; + +use fgumi_sort::{ + RawCoordinateKey, RawSortKey, RecordBuffer, SpillBlockCompressor, SpillCodec, + frame_keyed_record_into, +}; + +const BLOCK_SIZE: usize = 65280; // BGZF_MAX_BLOCK_SIZE +const RECORD_BODY: usize = 200; // ~realistic BAM record body bytes + +/// Per-record framing overhead `frame_keyed_record_into` adds: a 4-byte `u32` +/// record-length prefix, plus the serialized key ONLY when the key is not +/// already embedded in the record. Derived from `RawCoordinateKey`'s own trait +/// constants so it tracks the framing (and key size) instead of hardcoding `12`, +/// which would silently mis-size blocks if the key encoding ever changed. +const FRAME_OVERHEAD: usize = 4 // u32 record-length prefix + + if RawCoordinateKey::EMBEDDED_IN_RECORD { + 0 + } else { + // The coordinate key is fixed-size; a variable-length key (`None`) can't + // be statically bounded, but this bench only frames `RawCoordinateKey`. + match RawCoordinateKey::SERIALIZED_SIZE { + Some(n) => n, + None => 0, + } + }; + +/// Synthesize one ~200-byte BAM-ish record into `out`: tid (bytes 0-3), pos +/// (4-7), then a semi-compressible payload so zstd does realistic work. `n` is +/// the record ordinal (varies tid/pos so keys differ). +fn synth_record(out: &mut Vec, n: u64) { + out.clear(); + let tid: i32 = (n % 25) as i32; // 25 contigs + let pos: i32 = ((n.wrapping_mul(2_654_435_761) >> 8) % 250_000_000) as i32; + out.extend_from_slice(&tid.to_le_bytes()); + out.extend_from_slice(&pos.to_le_bytes()); + // bytes 8..16 — remaining fixed BAM header fields (arbitrary) + out.extend_from_slice(&[0u8; 8]); + // payload: semi-compressible (a few repeating motifs keyed off n) + let motif = [(n & 0xFF) as u8, b'A', b'C', b'G', b'T', ((n >> 3) & 0xFF) as u8]; + while out.len() < RECORD_BODY { + out.extend_from_slice(&motif); + } + out.truncate(RECORD_BODY); +} + +/// Frame the sorted records of `buf` into ≤64 KiB raw blocks, cut-before-overflow. +/// Reads `buf.refs()` (sorted) + `buf.get_record()` — the real gather access pattern. +fn gather_range(buf: &RecordBuffer, start: usize, end: usize) -> Vec> { + let refs = buf.refs(); + let mut blocks = Vec::new(); + let mut cur = Vec::with_capacity(BLOCK_SIZE + 1024); + for r in &refs[start..end] { + let key = RawCoordinateKey { sort_key: r.sort_key }; + let body = buf.get_record(r); + // would this record overflow a non-empty block? cut first. + if !cur.is_empty() && cur.len() + FRAME_OVERHEAD + body.len() > BLOCK_SIZE { + blocks.push(std::mem::replace(&mut cur, Vec::with_capacity(BLOCK_SIZE + 1024))); + } + frame_keyed_record_into(&mut cur, &key, body).expect("frame"); + } + if !cur.is_empty() { + blocks.push(cur); + } + blocks +} + +fn main() { + let args: Vec = std::env::args().collect(); + // Validate rather than fall through to a default. This writes a TSV row per + // run into a results table, so a typo that silently measured a different + // configuration than the label claims would corrupt the table rather than + // fail — the worst outcome for a benchmark. + let mode = args.get(1).map(String::as_str).unwrap_or("split").to_string(); + assert!( + mode == "split" || mode == "fused", + "mode must be `split` or `fused`, got `{mode}` — any other string would run fused \ + and be recorded under the name you typed", + ); + let n_arena: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(1); + // `sync_channel(0)` is a rendezvous channel and the seeding loop below sends + // nothing, so the producer would block on `free_rx.recv()` forever. + assert!(n_arena > 0, "N_arena must be at least 1 (0 deadlocks the producer)"); + let chunk_mib: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(512); + let n_chunks: usize = args.get(4).and_then(|s| s.parse().ok()).unwrap_or(6); + let workers: usize = args.get(5).and_then(|s| s.parse().ok()).unwrap_or(4); + // With `workers == 0` the per-worker `ranges` below are empty, so the + // benchmark skips gather and compression but still records a result row. + assert!(workers > 0, "workers must be at least 1"); + let codec = SpillCodec::Zstd; + let level = 1u32; + let mem_limit = chunk_mib * 1024 * 1024; + + eprintln!("PID {}", std::process::id()); + eprintln!( + "config: mode={mode} N_arena={n_arena} chunk_mib={chunk_mib} n_chunks={n_chunks} workers={workers} codec=zstd{level}" + ); + + // Per-phase timing accumulators (nanos). + let fill_ns = std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0)); + let sort_ns = std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0)); + let gather_ns = std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0)); + let compress_ns = std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0)); + + // Arena pool: a free-list channel seeded with N_arena empty buffers, and a + // work channel of filled+sorted buffers. Bounded so the producer blocks when + // no arena is free (this IS the N_arena bound). + let (free_tx, free_rx) = sync_channel::(n_arena); + let (work_tx, work_rx) = sync_channel::(n_arena); + // Estimate records/arena for with_capacity pre-sizing. + let est_records = mem_limit / (RECORD_BODY + 8 + 16); + for _ in 0..n_arena { + free_tx.send(RecordBuffer::with_capacity(est_records, mem_limit, 25)).expect("seed pool"); + } + + let wall = Instant::now(); + + // Producer: fill + par_sort each chunk, hand the owned buffer to the consumer. + let prod_fill = fill_ns.clone(); + let prod_sort = sort_ns.clone(); + let producer = std::thread::spawn(move || { + let mut scratch = Vec::with_capacity(RECORD_BODY + 8); + let mut ordinal: u64 = 0; + for _ in 0..n_chunks { + let mut buf = free_rx.recv().expect("acquire arena"); + let t = Instant::now(); + while buf.memory_usage() < mem_limit { + synth_record(&mut scratch, ordinal); + ordinal += 1; + buf.push_coordinate(&scratch).expect("push"); + } + prod_fill + .fetch_add(t.elapsed().as_nanos() as u64, std::sync::atomic::Ordering::Relaxed); + let t = Instant::now(); + buf.par_sort(); + prod_sort + .fetch_add(t.elapsed().as_nanos() as u64, std::sync::atomic::Ordering::Relaxed); + work_tx.send(buf).expect("hand to consumer"); + } + // drop work_tx → consumer loop ends + }); + + // Consumer: gather + compress each chunk; release the arena per `mode`. + let split = mode == "split"; + while let Ok(mut buf) = work_rx.recv() { + let n = buf.len(); + // contiguous index ranges, one per worker (`workers > 0` asserted above) + let per = n.div_ceil(workers); + let ranges: Vec<(usize, usize)> = (0..workers) + .map(|w| (w * per, ((w + 1) * per).min(n))) + .filter(|(s, e)| s < e) + .collect(); + + if split { + // GATHER (parallel) → owned raw blocks; then RELEASE arena; then COMPRESS (parallel). + let tg = Instant::now(); + let raw: Vec> = std::thread::scope(|sc| { + let bref = &buf; + let handles: Vec<_> = ranges + .iter() + .map(|&(s, e)| sc.spawn(move || gather_range(bref, s, e))) + .collect(); + handles.into_iter().flat_map(|h| h.join().expect("gather join")).collect() + }); + gather_ns + .fetch_add(tg.elapsed().as_nanos() as u64, std::sync::atomic::Ordering::Relaxed); + // RELEASE the arena BEFORE compress (the whole point of "split"). + buf.clear(); + let _ = free_tx.send(buf); + // COMPRESS the owned raw blocks in parallel (arena already free). + let tc = Instant::now(); + let chunks: Vec<&[Vec]> = + raw.chunks(raw.len().div_ceil(workers.max(1)).max(1)).collect(); + std::thread::scope(|sc| { + for ch in &chunks { + sc.spawn(move || { + let mut comp = SpillBlockCompressor::new(codec, level).expect("compressor"); + for blk in *ch { + let out = comp.compress_block(blk).expect("compress"); + std::hint::black_box(&out); + } + }); + } + }); + compress_ns + .fetch_add(tc.elapsed().as_nanos() as u64, std::sync::atomic::Ordering::Relaxed); + } else { + // FUSED: each worker gathers+compresses its range from the arena; + // the arena is released only after ALL workers finish. + let tf = Instant::now(); + std::thread::scope(|sc| { + for &(s, e) in &ranges { + let bref = &buf; + sc.spawn(move || { + let mut comp = SpillBlockCompressor::new(codec, level).expect("compressor"); + for blk in gather_range(bref, s, e) { + let out = comp.compress_block(&blk).expect("compress"); + std::hint::black_box(&out); + } + }); + } + }); + // fused gather+compress are interleaved; attribute the whole span to compress. + compress_ns + .fetch_add(tf.elapsed().as_nanos() as u64, std::sync::atomic::Ordering::Relaxed); + buf.clear(); + let _ = free_tx.send(buf); + } + let _ = n; + } + producer.join().expect("producer join"); + let wall_s = wall.elapsed().as_secs_f64(); + + let g = |a: &std::sync::atomic::AtomicU64| { + Duration::from_nanos(a.load(std::sync::atomic::Ordering::Relaxed)).as_secs_f64() + }; + // headline TSV + println!( + "{mode}\t{n_arena}\t{chunk_mib}\t{n_chunks}\t{workers}\t{:.2}\t{:.2}\t{:.2}\t{:.2}\t{:.2}", + wall_s, + g(&fill_ns), + g(&sort_ns), + g(&gather_ns), + g(&compress_ns) + ); +} diff --git a/crates/fgumi-sort/examples/read_ladder.rs b/crates/fgumi-sort/examples/read_ladder.rs deleted file mode 100644 index 80cb3339a..000000000 --- a/crates/fgumi-sort/examples/read_ladder.rs +++ /dev/null @@ -1,664 +0,0 @@ -//! An additive ladder over Phase 1's input path, one stage at a time. -//! -//! Phase 1's in-process instrumentation says what each stage costs *in situ*, but -//! two things it structurally cannot say: -//! -//! 1. **What the disk can actually deliver to this process.** The 605 MB/s -//! single-stream / 1110 MB/s four-stream figures the campaign quotes were taken -//! with direct I/O in a different session. `raw` re-measures them buffered, from -//! the same binary and the same boot, which is also the *actionable* ceiling — -//! an `O_DIRECT` reader would need aligned buffers and `unsafe`. -//! 2. **What our reader does with nothing downstream.** `record_step` times the -//! `ReadInputBlocks` step only when it *runs*; a reader sitting ineligible -//! because `raw_input_blocks` is full (cap = `num_workers * 8` = 128) or because -//! it is holding blocks is invisible to it. So the measured 358 MB/s is -//! consistent with both "one buffered stream can do no better" and "the reader -//! is fine and backpressure is throttling it" — which imply opposite fixes. -//! `blocks` separates them: it is the same framing code with no consumer at all. -//! -//! # Reading the result — the ladder is serial, the pipeline is not -//! -//! Each rung adds one stage to the one below it, so `rung(n) - rung(n-1)` is that -//! stage's cost **on one thread, with nothing else running**. The real sort does -//! not work that way: reading is one exclusive thread, decompression is spread over -//! all 16 workers, and framing/key/push are the serial main thread. The pipeline's -//! floor is the **max over its serial resources**, not the sum of these rungs. -//! -//! Concretely: decompression is the largest single block of CPU in Phase 1 -//! (231.4s of worker busy) and would dominate this ladder, but across 16 workers it -//! is ~14.5s of effective wall and nowhere near binding. A rung being expensive -//! here does **not** mean it binds in production. The `resource` column names where -//! each stage actually lands, and that is the column that decides whether a rung's -//! cost matters. -//! -//! What the ladder is good for: per-stage *ceilings*, and cross-validating the -//! in-process partition by a completely independent method. -//! -//! # Usage -//! -//! One rung per process, so no rung inherits a warm page cache from the one below -//! it — drop caches between invocations (see `scripts/read-ladder.sh`). -//! -//! ```text -//! read_ladder [--buf BYTES] [--streams N] [--limit-gb G] -//! -//! rungs, each adding one stage to the previous: -//! raw read() into a reused buffer [reader thread] -//! blocks + BGZF framing (read_raw_blocks) [reader thread] -//! decompress + inflate each block [16 workers] -//! records + walk the BAM records in the stream [main thread] -//! key + extract the template-coordinate sort key [main thread] -//! push + copy the record into an arena and keep a ref [main thread] -//! ``` - -#![deny(unsafe_code)] - -use std::fs::File; -use std::io::BufReader; -use std::os::unix::fs::FileExt; -use std::path::{Path, PathBuf}; -use std::sync::Arc; -use std::time::Instant; - -use anyhow::{Context, Result, bail}; -use fgumi_bgzf::reader::{decompress_block_into, read_raw_blocks}; -use fgumi_raw_bam::BAM_MAGIC; -use fgumi_sort::{LibraryLookup, cb_hasher, extract_template_key_inline}; -use libdeflater::Decompressor; - -/// The sort pipeline's own input buffer (`SORT_INPUT_BUFFER_SIZE`). -const DEFAULT_BUF_BYTES: usize = 2 * 1024 * 1024; - -/// The sort pipeline's own batch size (`INPUT_READ_BATCH_SIZE`). -const BLOCKS_PER_BATCH: usize = 16; - -/// Which resource a stage lands on in the real pipeline. Printed with every rung -/// because a serial ladder over a concurrent pipeline is misread by default. -fn resource_for(rung: &str) -> &'static str { - match rung { - "raw" | "blocks" => "reader thread (exclusive, worker 0)", - "decompress" => "16-worker pool", - _ => "ingest thread (serial, main)", - } -} - -/// How the `raw` rung's streams divide the file. -#[derive(Clone, Copy, PartialEq, Eq)] -enum Pattern { - /// Each stream owns one large contiguous span. Every stream is purely - /// sequential, which is the friendliest case for kernel read-ahead -- and is - /// *not* what an in-order reader can do, since stream 3 would be delivering - /// bytes from 30 GB in while the framer still needs byte 0. - Contiguous, - /// Streams take turns on fixed-size chunks (stream k takes chunk k, k+N, ...). - /// This is the access pattern an in-order parallel reader actually generates, - /// so it is the one worth believing. - Interleaved, -} - -/// What the reader tells the kernel about the access pattern it is about to make. -/// -/// The question these exist to answer: a blocking `read()` keeps roughly one -/// read-ahead window outstanding at the device, and on this volume that is ~3x -/// too little to saturate it (358 MB/s single-stream against a ~1050 MB/s cap). -/// Four threads fix that by carrying four windows. These advise calls try to fix -/// it *without* extra threads, by asking the kernel to keep the queue primed on -/// a single ordered stream -- which, unlike parallel streams, cannot deliver -/// blocks out of order. -#[derive(Clone, Copy, PartialEq, Eq)] -enum Advice { - /// No hint: the default read-ahead window and nothing else. - None, - /// `POSIX_FADV_SEQUENTIAL` once per stream. Linux doubles the read-ahead - /// window for the file; a 2x lever at most. - Sequential, - /// `POSIX_FADV_WILLNEED` for a window ahead of the cursor, reissued as the - /// cursor advances. Not bounded by the read-ahead window, so this is the one - /// that can actually fill the device queue. - WillNeed, -} - -/// Issue one `posix_fadvise` call, or do nothing where it does not exist. -/// -/// Errors are deliberately ignored: `fadvise` is a *hint*, and a kernel that -/// declines it (memory pressure, an unsupported filesystem) must not fail the -/// read. A hint that was dropped shows up as a throughput number that did not -/// move, which is exactly the result the experiment is testing for. -#[cfg(target_os = "linux")] -fn advise(file: &File, offset: u64, len: u64, advice: nix::fcntl::PosixFadviseAdvice) { - let offset = i64::try_from(offset).unwrap_or(i64::MAX); - let len = i64::try_from(len).unwrap_or(i64::MAX); - let _ = nix::fcntl::posix_fadvise(file, offset, len, advice); -} - -#[cfg(target_os = "linux")] -fn advise_sequential(file: &File, offset: u64, len: u64) { - advise(file, offset, len, nix::fcntl::PosixFadviseAdvice::POSIX_FADV_SEQUENTIAL); -} - -#[cfg(target_os = "linux")] -fn advise_willneed(file: &File, offset: u64, len: u64) { - advise(file, offset, len, nix::fcntl::PosixFadviseAdvice::POSIX_FADV_WILLNEED); -} - -#[cfg(not(target_os = "linux"))] -fn advise_sequential(_file: &File, _offset: u64, _len: u64) {} - -#[cfg(not(target_os = "linux"))] -fn advise_willneed(_file: &File, _offset: u64, _len: u64) {} - -struct Args { - rung: String, - path: PathBuf, - buf_bytes: usize, - streams: usize, - limit_bytes: u64, - pattern: Pattern, - chunk_bytes: u64, - writer_mbps: u64, - writer_dir: Option, - advice: Advice, - /// Bytes ahead of the cursor to keep advised, for `Advice::WillNeed`. - advise_ahead: u64, -} - -fn parse_args() -> Result { - let mut it = std::env::args().skip(1); - let rung = it - .next() - .context("usage: read_ladder [--buf N] [--streams N] [--limit-gb G]")?; - let path = PathBuf::from(it.next().context("missing ")?); - let mut buf_bytes = DEFAULT_BUF_BYTES; - let mut streams = 1usize; - let mut limit_bytes = u64::MAX; - let mut pattern = Pattern::Contiguous; - let mut chunk_bytes = 4 * 1024 * 1024u64; - let mut writer_mbps = 0u64; - let mut writer_dir = None; - let mut advice = Advice::None; - // 8 MiB: four 2 MiB reads of lead, comfortably more than the ~400 KB the - // device's bandwidth-delay product asks for, so a null result cannot be - // blamed on advising too little. - let mut advise_ahead = 8 * 1024 * 1024u64; - while let Some(flag) = it.next() { - let value = it.next().with_context(|| format!("{flag} needs a value"))?; - match flag.as_str() { - "--buf" => buf_bytes = value.parse().context("--buf")?, - "--streams" => streams = value.parse().context("--streams")?, - "--limit-gb" => { - let gb: f64 = value.parse().context("--limit-gb")?; - #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] - { - limit_bytes = (gb * 1e9) as u64; - } - } - "--pattern" => { - pattern = match value.as_str() { - "contiguous" => Pattern::Contiguous, - "interleaved" => Pattern::Interleaved, - other => bail!("--pattern must be contiguous|interleaved, got {other}"), - }; - } - "--chunk" => chunk_bytes = value.parse().context("--chunk")?, - "--writer-mbps" => writer_mbps = value.parse().context("--writer-mbps")?, - "--fadvise" => { - advice = match value.as_str() { - "none" => Advice::None, - "sequential" => Advice::Sequential, - "willneed" => Advice::WillNeed, - other => bail!("--fadvise must be none|sequential|willneed, got {other}"), - }; - } - "--advise-ahead" => advise_ahead = value.parse().context("--advise-ahead")?, - "--writer-dir" => writer_dir = Some(PathBuf::from(value)), - other => bail!("unknown flag {other}"), - } - } - if streams == 0 { - bail!("--streams must be at least 1"); - } - if chunk_bytes == 0 { - bail!("--chunk must be positive"); - } - if buf_bytes == 0 { - bail!("--buf must be positive"); - } - if writer_mbps > 0 && writer_dir.is_none() { - bail!("--writer-mbps needs --writer-dir"); - } - // Only `raw` spawns the concurrent Writer; the pipeline rungs share a single - // reader loop with no writer, so accepting these flags there would silently - // ignore them (and, unthrottled, could write a multi-GB temp file for - // nothing). Reject them instead of pretending to honor them. - if (writer_mbps > 0 || writer_dir.is_some()) && rung != "raw" { - bail!("--writer-mbps/--writer-dir only apply to the `raw` rung"); - } - Ok(Args { - rung, - path, - buf_bytes, - streams, - limit_bytes, - pattern, - chunk_bytes, - writer_mbps, - writer_dir, - advice, - advise_ahead, - }) -} - -#[derive(Default)] -struct Tally { - bytes: u64, - blocks: u64, - records: u64, - wrote_bytes: u64, -} - -fn main() -> Result<()> { - let args = parse_args()?; - let started = Instant::now(); - let tally = match args.rung.as_str() { - "raw" => rung_raw(&args)?, - "blocks" | "decompress" | "records" | "key" | "push" => rung_pipeline(&args)?, - other => bail!("unknown rung {other}"), - }; - report(&args, &tally, started.elapsed().as_secs_f64()); - Ok(()) -} - -#[allow(clippy::cast_precision_loss, reason = "byte and record counts stay far below 2^52")] -fn report(args: &Args, tally: &Tally, secs: f64) { - let mb_per_sec = if secs > 0.0 { tally.bytes as f64 / secs / 1e6 } else { 0.0 }; - let per = |n: u64| if n == 0 { 0.0 } else { secs * 1e9 / n as f64 }; - let write_mb_per_sec = if secs > 0.0 { tally.wrote_bytes as f64 / secs / 1e6 } else { 0.0 }; - let pattern = match args.pattern { - Pattern::Contiguous => "contiguous", - Pattern::Interleaved => "interleaved", - }; - println!( - "rung={rung} streams={streams} pattern={pattern} chunk={chunk} buf={buf} \ - secs={secs:.2} MBps={mb:.0} write_MBps={wmb:.0} \ - bytes={bytes} blocks={blocks} ns_per_block={npb:.0} records={records} \ - ns_per_record={npr:.1} resource=\"{res}\"", - rung = args.rung, - streams = args.streams, - chunk = args.chunk_bytes, - buf = args.buf_bytes, - wmb = write_mb_per_sec, - mb = mb_per_sec, - bytes = tally.bytes, - blocks = tally.blocks, - npb = per(tally.blocks), - records = tally.records, - npr = per(tally.records), - res = resource_for(&args.rung), - ); -} - -/// `raw`: how fast the file's bytes can be pulled into this process at all. -/// -/// `streams > 1` splits the file into contiguous ranges and reads each on its own -/// thread with `read_at`, which is what a parallel reader would do — the volume -/// measured 605 MB/s at one stream and 1110 MB/s at four, and nothing in the -/// current reader can exploit that. -fn rung_raw(args: &Args) -> Result { - let file = - Arc::new(File::open(&args.path).with_context(|| format!("open {}", args.path.display()))?); - let total = file.metadata()?.len().min(args.limit_bytes); - let streams = u64::try_from(args.streams).expect("stream count fits in u64"); - let span = total.div_ceil(streams); - - // Production reads and writes the same volume at once: Phase 1 spills 52.6 GB - // while it reads 43.1 GB. A read rate measured on an idle device is therefore - // an upper bound that production can never see, and the gap is the whole - // question -- the volume is provisioned near 1000 MB/s, so a 1048 MB/s read - // plus a ~363 MB/s spill stream cannot both be served. - let writer = args.writer_dir.as_ref().map(|dir| Writer::spawn(dir, args.writer_mbps)); - - let mut handles = Vec::with_capacity(args.streams); - for stream in 0..streams { - let file = Arc::clone(&file); - let buf_bytes = args.buf_bytes; - let pattern = args.pattern; - let chunk = args.chunk_bytes; - let start = stream * span; - let end = ((stream + 1) * span).min(total); - let advice = args.advice; - let advise_ahead = args.advise_ahead; - handles.push(std::thread::spawn(move || -> Result { - let mut buf = vec![0u8; buf_bytes]; - let mut read_total = 0u64; - if advice == Advice::Sequential { - advise_sequential(&file, start, end.saturating_sub(start)); - } - match pattern { - Pattern::Contiguous => { - let mut offset = start; - // Prime the queue before the first read, so the very first - // call is not the one uncovered case. - if advice == Advice::WillNeed { - advise_willneed(&file, offset, advise_ahead.min(end - offset)); - } - while offset < end { - let want = usize::try_from((end - offset).min(buf_bytes as u64)) - .expect("clamped to buffer size"); - // Advise the window *beyond* what this read will consume, - // so the device is working on the next span while this - // one is being copied out. - if advice == Advice::WillNeed { - let ahead_start = offset + want as u64; - if ahead_start < end { - advise_willneed( - &file, - ahead_start, - advise_ahead.min(end - ahead_start), - ); - } - } - let n = file.read_at(&mut buf[..want], offset)?; - if n == 0 { - break; - } - offset += n as u64; - read_total += n as u64; - } - } - Pattern::Interleaved => { - // Stream k takes chunk k, k + N, k + 2N ... so the file is - // covered in order across the pool rather than in N disjoint - // spans. Each stream strides by `streams * chunk`, which is - // what makes this a different ask of kernel read-ahead. - let stride = chunk * streams; - let mut chunk_start = stream * chunk; - while chunk_start < total { - let chunk_end = (chunk_start + chunk).min(total); - let mut offset = chunk_start; - while offset < chunk_end { - let want = usize::try_from((chunk_end - offset).min(buf_bytes as u64)) - .expect("clamped to buffer size"); - let n = file.read_at(&mut buf[..want], offset)?; - if n == 0 { - break; - } - offset += n as u64; - read_total += n as u64; - } - chunk_start += stride; - } - } - } - Ok(read_total) - })); - } - - // Join every reader before propagating a failure, so the writer is always - // stopped and its temp file removed by `finish` -- returning through `??` - // here would leave the writer thread running and `read_ladder_writer.tmp` - // (written unthrottled when --writer-mbps is 0) on disk. - let mut tally = Tally::default(); - let mut read_result = Ok(()); - for handle in handles { - match handle.join().map_err(|_| anyhow::anyhow!("reader thread panicked"))? { - Ok(n) => tally.bytes += n, - Err(e) => read_result = Err(e), - } - } - if let Some(writer) = writer { - tally.wrote_bytes = writer.finish()?; - } - read_result?; - Ok(tally) -} - -/// A rate-limited background writer, so the read rungs can be measured against -/// the spill load production actually runs alongside them. -struct Writer { - stop: Arc, - handle: std::thread::JoinHandle>, - path: PathBuf, -} - -impl Writer { - fn spawn(dir: &Path, target_mbps: u64) -> Self { - use std::io::Write; - let path = dir.join("read_ladder_writer.tmp"); - let stop = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let thread_stop = Arc::clone(&stop); - let thread_path = path.clone(); - let handle = std::thread::spawn(move || -> Result { - let mut file = File::create(&thread_path) - .with_context(|| format!("create {}", thread_path.display()))?; - let block = vec![0x5Au8; 4 * 1024 * 1024]; - let started = Instant::now(); - let mut written = 0u64; - while !thread_stop.load(std::sync::atomic::Ordering::Relaxed) { - file.write_all(&block)?; - written += block.len() as u64; - if target_mbps > 0 { - // Token bucket: hold the writer at the target rate rather than - // letting it race the reader for the device, which would - // measure contention at some arbitrary ratio instead of the - // one production generates. - #[allow(clippy::cast_precision_loss, reason = "byte totals stay below 2^52")] - let due = written as f64 / (target_mbps as f64 * 1e6); - let elapsed = started.elapsed().as_secs_f64(); - if due > elapsed { - std::thread::sleep(std::time::Duration::from_secs_f64(due - elapsed)); - } - } - } - file.sync_all()?; - Ok(written) - }); - Self { stop, handle, path } - } - - fn finish(self) -> Result { - self.stop.store(true, std::sync::atomic::Ordering::Relaxed); - let written = self.handle.join().map_err(|_| anyhow::anyhow!("writer thread panicked"))?; - let _ = std::fs::remove_file(&self.path); - written - } -} - -/// Everything from `blocks` up, sharing one loop so the rungs differ by exactly the -/// stage each adds and nothing else. -fn rung_pipeline(args: &Args) -> Result { - let rung = args.rung.as_str(); - let want_decompress = matches!(rung, "decompress" | "records" | "key" | "push"); - let want_records = matches!(rung, "records" | "key" | "push"); - let want_key = matches!(rung, "key" | "push"); - let want_push = rung == "push"; - - // The header is read through the crate's own opener, so the key rung sees the - // same library ordinals production does rather than an invented mapping. - let lookup = if want_key { - let (_reader, header) = fgumi_sort::open_raw_bam_record_reader_with_header(&args.path)?; - Some(LibraryLookup::from_header(&header)) - } else { - None - }; - let hasher = cb_hasher(); - - let file = File::open(&args.path).with_context(|| format!("open {}", args.path.display()))?; - let mut reader = BufReader::with_capacity(args.buf_bytes, file); - - let mut tally = Tally::default(); - let mut decompressed = Vec::with_capacity(64 * 1024); - // Reused across every block, as the worker pool does -- constructing one per - // block would put allocator noise in the decompress rung's delta. - let mut decompressor = Decompressor::new(); - let mut framer = RecordFramer::default(); - let mut arena: Vec = Vec::new(); - let mut refs: Vec<(u64, usize)> = Vec::new(); - - loop { - let batch = read_raw_blocks(&mut reader, BLOCKS_PER_BATCH)?; - if batch.is_empty() { - break; - } - for block in &batch { - tally.bytes += block.len() as u64; - tally.blocks += 1; - if !want_decompress { - continue; - } - // `decompress_block_into` *appends*. Without this the buffer accumulates - // the whole file, which costs the decompress rung a realloc storm and - // makes the records rung quadratic -- it re-frames every prior block, - // which is ~2 TB of copying on a 43 GB input and eats RAM until the - // OOM killer intervenes. This cost a workstation once; leave it here. - decompressed.clear(); - decompress_block_into(block, &mut decompressor, &mut decompressed)?; - if !want_records { - continue; - } - framer.push(&decompressed); - while let Some(record) = framer.next_record() { - tally.records += 1; - if want_key { - let key = extract_template_key_inline( - record, - lookup.as_ref().expect("key rung builds a lookup"), - None, - &hasher, - ); - if want_push { - refs.push((key.primary, arena.len())); - arena.extend_from_slice(record); - // The real ingest spills at a memory bound; here the arena - // is reset instead, so the rung measures the copy rather - // than this process's willingness to hold 40 GB. - if arena.len() >= 512 * 1024 * 1024 { - arena.clear(); - refs.clear(); - } - } else { - std::hint::black_box(&key); - } - } - } - } - if tally.bytes >= args.limit_bytes { - break; - } - } - Ok(tally) -} - -/// Frames BAM records out of the decompressed byte stream. -/// -/// Records straddle BGZF block boundaries, so this carries the tail of one block -/// into the next — the same job `PooledInputStream` does in production, minus the -/// reorder buffer, which only exists because production decompresses out of order. -#[derive(Default)] -struct RecordFramer { - carry: Vec, - start: usize, - header_skipped: bool, -} - -/// The carry may hold one partial record plus one decompressed block. A BGZF -/// block is at most 64 KiB and a BAM record in any real file is far below this, -/// so exceeding the bound does not mean "unusual input" -- it means the consumer -/// has stopped draining and the buffer is growing without limit. Asserting it is -/// cheaper than discovering it through the OOM killer. -const CARRY_LIMIT: usize = 16 * 1024 * 1024; - -impl RecordFramer { - fn push(&mut self, bytes: &[u8]) { - // Reclaim consumed bytes before growing, so the carry stays near one block - // rather than the whole file. - if self.start > 0 { - self.carry.drain(..self.start); - self.start = 0; - } - self.carry.extend_from_slice(bytes); - assert!( - self.carry.len() <= CARRY_LIMIT, - "record framer carry reached {} bytes, past the {CARRY_LIMIT}-byte bound: \ - the consumer is not draining and this is about to exhaust memory", - self.carry.len(), - ); - } - - fn next_record(&mut self) -> Option<&[u8]> { - if !self.header_skipped && !self.try_skip_header() { - return None; - } - let available = &self.carry[self.start..]; - if available.len() < 4 { - return None; - } - let len = u32::from_le_bytes(available[..4].try_into().ok()?) as usize; - // A real BAM record is far below CARRY_LIMIT (a BGZF block is <= 64 KiB), so - // an oversized length is a corrupt or misframed record, not one still - // waiting for more data. Without this guard `available.len() < 4 + len` - // never clears, `next_record` returns None on every call, and `push` grows - // the carry until it asserts "the consumer is not draining" -- blaming the - // wrong cause. Surface the real one here instead. - assert!( - len <= CARRY_LIMIT, - "record framer read a record length of {len} bytes, past the {CARRY_LIMIT}-byte \ - bound: the input is corrupt or misframed", - ); - if available.len() < 4 + len { - return None; - } - let from = self.start + 4; - self.start += 4 + len; - Some(&self.carry[from..from + len]) - } - - /// Skip `magic | l_text | text | n_ref | (l_name, name, l_ref) * n_ref`. - /// - /// Returns false while the header is still incomplete, so the caller simply - /// feeds another block. - fn try_skip_header(&mut self) -> bool { - let buf = &self.carry[self.start..]; - if buf.len() < 12 { - return false; - } - assert_eq!(&buf[..4], BAM_MAGIC, "input is not a BAM stream"); - let l_text = u32::from_le_bytes(buf[4..8].try_into().expect("4 bytes")) as usize; - let Some(after_text) = 8usize.checked_add(l_text) else { return false }; - if buf.len() < after_text + 4 { - return false; - } - let n_ref = u32::from_le_bytes(buf[after_text..after_text + 4].try_into().expect("4 bytes")) - as usize; - let mut at = after_text + 4; - for _ in 0..n_ref { - if buf.len() < at + 4 { - return false; - } - let l_name = u32::from_le_bytes(buf[at..at + 4].try_into().expect("4 bytes")) as usize; - // Parity with `next_record`: an oversized name length is corruption, not - // a header still arriving, so reject it with the real cause. - assert!( - l_name <= CARRY_LIMIT, - "BAM header declared a reference-name length of {l_name} bytes, past the \ - {CARRY_LIMIT}-byte bound: the input is corrupt or misframed", - ); - // Checked like `l_text` above: a length that overflows the cursor is - // treated as an incomplete header (feed another block) rather than - // wrapping past the buffer. - let Some(next_at) = at - .checked_add(4) - .and_then(|x| x.checked_add(l_name)) - .and_then(|x| x.checked_add(4)) - else { - return false; - }; - at = next_at; - if buf.len() < at { - return false; - } - } - self.start += at; - self.header_skipped = true; - true - } -} diff --git a/crates/fgumi-sort/src/arena_pool.rs b/crates/fgumi-sort/src/arena_pool.rs new file mode 100644 index 000000000..63ce3f822 --- /dev/null +++ b/crates/fgumi-sort/src/arena_pool.rs @@ -0,0 +1,354 @@ +#![deny(unsafe_code)] +//! Bounded, reusable pool of [`SegmentedBuf`] sort arenas — an RSS fix. +//! +//! The Phase-1 spill path fills a multi-GB `SegmentedBuf`, hands it downstream +//! as an `Arc`, and (pre-pool) minted a *fresh* buffer for the next fill via +//! `mem::take`. Combined with the block-parallel spill running async, two full +//! arenas were live at the Phase-1→Phase-2 boundary (the in-flight chunk — +//! pinned by slow compression through `SpillGather`'s bounded `pending` — plus +//! the next fill), inflating peak RSS to ~2× base. +//! +//! This pool **bounds** the live arenas to `capacity` (default 1, matching +//! legacy's one-arena-at-a-time `drain_pending_spill` model) and **reuses** +//! their storage: [`try_acquire`](ArenaPool::try_acquire) hands out a +//! reset-for-reuse buffer or `None` when all `capacity` arenas are in flight +//! (the caller backpressures with `NoProgress` until an in-flight chunk's `Arc` +//! drops and returns its arena via [`PooledSegmentedBuf`]'s `Drop`). Capping at +//! 1 means the next fill cannot start until the prior chunk is spilled and its +//! arena freed — exactly legacy's behaviour, and the RSS-gate-safe footprint. + +use std::ops::{Deref, DerefMut}; +use std::sync::{Arc, Mutex}; + +use crate::segmented_buf::SegmentedBuf; + +/// A bounded free-list of reusable [`SegmentedBuf`] arenas. At most `capacity` +/// buffers ever exist; `try_acquire` returns `None` when all are in flight. +pub struct ArenaPool { + inner: Mutex, + /// Maximum number of live arenas (`N_arena`); clamped to ≥ 1. + capacity: usize, + /// Segment size for freshly-allocated arenas. + segment_size: usize, +} + +struct PoolInner { + /// Returned, reset-for-reuse buffers ready to hand out. + free: Vec, + /// Total arenas ever allocated (never decremented; bounds at `capacity`). + made: usize, +} + +impl ArenaPool { + /// Create a pool bounded to `capacity` (≥1) arenas, each `segment_size`. + #[must_use] + pub fn new(capacity: usize, segment_size: usize) -> Arc { + Arc::new(Self { + inner: Mutex::new(PoolInner { free: Vec::new(), made: 0 }), + capacity: capacity.max(1), + // Clamped for the same reason `capacity` is, and to match + // `SegmentedBuf::with_capacity`, which clamps it again downstream: a + // 0 here would otherwise be silently accepted at this boundary and + // corrected somewhere else. + segment_size: segment_size.max(1), + }) + } + + /// Acquire a reset-for-reuse arena, or `None` if all `capacity` arenas are + /// in flight (the caller must backpressure and retry once one returns). + /// + /// Returns the RAII wrapper, never a bare [`SegmentedBuf`], and that is + /// load-bearing rather than stylistic. `made` is only ever incremented; the + /// sole path back to the free-list is [`PooledSegmentedBuf::drop`]. If a + /// caller could hold the arena unwrapped, then *any* drop on that path — a + /// `?` on a malformed record, an early return, an unwind — would retire the + /// slot permanently. At the default `capacity == 1` that leaves the pool + /// empty forever, and because the documented response to `None` is to + /// backpressure until an in-flight arena returns, the caller would spin + /// forever instead of failing: a hang, not an error. Handing back the + /// wrapper makes the return unconditional and unforgeable. + /// + /// # Panics + /// + /// Panics if the internal mutex is poisoned. That is a deliberate asymmetry + /// with `release`, which recovers the same guard: the + /// protected state is never torn, so recovery is *sound* anywhere — but + /// `release` runs from a `Drop` and must not panic during unwinding, while + /// this path has no such constraint and a poisoned pool is worth surfacing + /// loudly rather than papering over. + #[must_use] + pub fn try_acquire(self: &Arc) -> Option { + let mut g = self.inner.lock().expect("arena pool mutex poisoned"); + let buf = if let Some(buf) = g.free.pop() { + buf + } else if g.made < self.capacity { + g.made += 1; + SegmentedBuf::with_capacity(0, self.segment_size) + } else { + return None; + }; + drop(g); + Some(PooledSegmentedBuf::pooled(buf, Arc::clone(self))) + } + + /// Return an arena to the free-list, reset for reuse (capacity retained). + /// + /// Reached only from [`PooledSegmentedBuf::drop`], so it must not panic: if + /// the arena is dropped while a panic is already unwinding, a second panic + /// here aborts the process outright — no catch, no message. No path poisons + /// this mutex *today* — `try_acquire` is the only one that runs non-trivial + /// code under it, and it allocates with `capacity = 0`, so neither the + /// first-segment nor the segment-vector allocation scales with + /// `segment_size` and neither can overflow. That is a property of the + /// current body, not of the lock: any future panic under this guard would + /// poison it, and this path would be the one that turned that into an abort. + /// + /// Recovering the guard is sound here because the state it protects cannot + /// be torn: `free` holds reset-for-reuse buffers and `made` is a counter, so + /// a thread that died mid-critical-section leaves the free-list structurally + /// valid (at worst `made` counts an arena that was never handed out, which + /// costs a slot but breaks nothing). `try_acquire` deliberately keeps its + /// `expect` — it is not on a drop path, and a poisoned pool should surface + /// there rather than be papered over. + fn release(&self, mut buf: SegmentedBuf) { + // Catch a buffer that never came from this pool's `try_acquire` (wrong + // `segment_size`): reusing it would silently defeat the `capacity`-based + // RSS bound this pool exists to enforce. Debug-only — the check is a + // developer guard, not a production cost. + debug_assert_eq!( + buf.segment_size(), + self.segment_size, + "released arena's segment_size doesn't match this pool's configured size" + ); + buf.reset_for_reuse(); + self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner).free.push(buf); + } + + /// Number of arenas currently available in the free-list (test/diagnostic). + /// + /// Does not panic on a poisoned mutex — see the comment in the body. + #[cfg(test)] + #[must_use] + pub fn free_len(&self) -> usize { + // Recovers like `release` rather than panicking: this exists to observe + // pool state, including after a poisoning, which is exactly when a test + // most needs to read it. + self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner).free.len() + } +} + +/// A [`SegmentedBuf`] that returns to its [`ArenaPool`] on drop (when `pool` +/// is `Some`). Used as the `Arc`-shared backing store of an in-memory sort +/// chunk: the arena is reclaimed once the last `Arc` clone drops (after the +/// gather has framed a spilled chunk). Buffers that are not pool-managed — the +/// residual chunk, and any buffer built outside the pool — use +/// [`unpooled`](Self::unpooled), so they drop normally and never reach the +/// free-list. +pub struct PooledSegmentedBuf { + /// `Some` until `Drop`. `Option` only so `Drop` can move the buffer out. + buf: Option, + /// `Some` → return to the pool on drop; `None` → drop normally. + pool: Option>, +} + +impl PooledSegmentedBuf { + /// Wrap a buffer that should return to `pool` on drop. + /// + /// Deliberately private: [`ArenaPool::try_acquire`] is the only way to get a + /// pooled wrapper, so every buffer on the free-list came from this pool. + /// Were this public, a caller could wrap a foreign buffer and `release` + /// would push it onto the free-list — `try_acquire` pops from `free` before + /// consulting `made`, so that directly raises the number of simultaneously + /// live arenas above `capacity`, defeating the RSS bound the pool exists to + /// enforce. The `debug_assert_eq!` on `segment_size` in `release` catches + /// only the mismatched-size case, and only in debug. + #[must_use] + fn pooled(buf: SegmentedBuf, pool: Arc) -> Self { + Self { buf: Some(buf), pool: Some(pool) } + } + + /// Wrap a buffer that is NOT pool-managed (drops normally). + #[must_use] + pub fn unpooled(buf: SegmentedBuf) -> Self { + Self { buf: Some(buf), pool: None } + } +} + +impl Default for PooledSegmentedBuf { + /// An empty, NON-pooled buffer. + /// + /// This exists so `std::mem::take(&mut some_pooled_buf)` moves the WRAPPER + /// out. Without it, `take` deref-coerces through `DerefMut` and steals the + /// inner `SegmentedBuf`, leaving the wrapper holding a default-sized buffer + /// that it still believes belongs to the pool: the real arena is orphaned + /// (never returned) and a wrong-`segment_size` buffer is later pushed onto + /// the free list. Pooled-ness must travel with the arena, and it only does + /// if the whole wrapper moves. + fn default() -> Self { + Self::unpooled(SegmentedBuf::default()) + } +} + +impl DerefMut for PooledSegmentedBuf { + /// Fill happens *through* the wrapper. Without this the caller would have to + /// unwrap to write, which is exactly the window in which a dropped arena + /// silently retires its pool slot — see [`ArenaPool::try_acquire`]. + fn deref_mut(&mut self) -> &mut SegmentedBuf { + self.buf.as_mut().expect("PooledSegmentedBuf used after drop") + } +} + +impl Deref for PooledSegmentedBuf { + type Target = SegmentedBuf; + fn deref(&self) -> &SegmentedBuf { + self.buf.as_ref().expect("PooledSegmentedBuf used after drop") + } +} + +impl Drop for PooledSegmentedBuf { + fn drop(&mut self) { + if let Some(buf) = self.buf.take() + && let Some(pool) = &self.pool + { + pool.release(buf); + } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[test] + fn bounds_to_capacity_and_reuses() { + let pool = ArenaPool::new(1, 1024); + let a = pool.try_acquire().expect("first acquire"); + assert!(pool.try_acquire().is_none(), "capacity 1 → second acquire blocked"); + drop(a); // the wrapper's Drop is the only return path + assert_eq!(pool.free_len(), 1, "released arena back in the free-list"); + let _b = pool.try_acquire().expect("acquire after release reuses the arena"); + assert!(pool.try_acquire().is_none(), "still capacity-bound"); + } + + /// `capacity` must actually bound the pool. Every other test here uses + /// `capacity == 1`, where "hand out one arena" and "respect the capacity" + /// are indistinguishable — an implementation ignoring the field entirely + /// passes them all. This one separates the two. + #[rstest] + #[case::two(2)] + #[case::five(5)] + fn capacity_bounds_the_number_of_live_arenas(#[case] capacity: usize) { + let pool = ArenaPool::new(capacity, 1024); + let live: Vec<_> = (0..capacity) + .map(|i| pool.try_acquire().unwrap_or_else(|| panic!("acquire {i} within capacity"))) + .collect(); + assert!(pool.try_acquire().is_none(), "the {capacity}+1'th acquire must be refused"); + drop(live); + assert_eq!(pool.free_len(), capacity, "every arena returns on drop"); + } + + /// The leak this API shape exists to prevent: an arena dropped on an error + /// path must still return. `made` is never decremented, so a slot lost here + /// is lost forever — at the default capacity 1 that empties the pool, and + /// since the documented response to `None` is to backpressure until an arena + /// returns, the caller would spin forever rather than fail. + #[test] + fn arena_dropped_on_an_error_path_still_returns_to_the_pool() { + // A fill that bails part-way: the arena is live and partly written when + // the `?` fires, and nothing wraps it on the way out. + fn fill_then_fail(pool: &Arc) -> Result<(), &'static str> { + let mut arena = pool.try_acquire().ok_or("pool empty")?; + arena.extend_from_slice(&[0xAB; 32]); + Err("malformed record")?; + unreachable!() + } + + let pool = ArenaPool::new(1, 1024); + // Discriminating: `fill_then_fail` can also fail with "pool empty", which + // would mean the acquire never happened and the rest of the test proves + // nothing about the drop path. + assert_eq!( + fill_then_fail(&pool), + Err("malformed record"), + "must reach the mid-fill bail, not fail to acquire", + ); + assert_eq!(pool.free_len(), 1, "the arena came back despite the early return"); + let reused = pool.try_acquire().expect("pool is still usable after the error path"); + assert!(reused.is_empty(), "and the returned arena was reset"); + } + + /// `Deref`/`DerefMut` are the only way a consumer reaches the arena — the + /// wrapper exposes no accessor — so a broken impl would make every pooled + /// read or write unreachable. Nothing else here exercises them directly. + #[test] + fn deref_reaches_the_wrapped_arena_for_both_read_and_write() { + let pool = ArenaPool::new(1, 1024); + let mut arena = pool.try_acquire().expect("acquire"); + + // DerefMut: the fill path. + let off = arena.extend_from_slice(b"pooled bytes"); + // Deref: the read path the consumer uses. + assert_eq!(arena.slice(off, 12), b"pooled bytes"); + assert_eq!(arena.len(), 12, "Deref reaches the wrapped buffer's own state"); + + drop(arena); + let reused = pool.try_acquire().expect("reacquire"); + assert!(reused.is_empty(), "and the arena was reset on the way back"); + } + + /// A poisoned pool mutex must not make `PooledSegmentedBuf::drop` panic. + /// + /// This one cannot be written as a regression test after the fact: if + /// `release` panics while an arena drops during unwinding, the process + /// aborts, so there is no failing assertion to observe — the test binary + /// just dies. Pinning the recovery here is the only way to keep it. + #[test] + fn release_recovers_a_poisoned_lock_rather_than_panicking_in_drop() { + let pool = ArenaPool::new(1, 1024); + let arena = pool.try_acquire().expect("acquire"); + + // Poison the mutex the only way it can be poisoned: a thread dies while + // holding it. No current path does that (see `release`'s doc), so the + // panic here is deliberately synthetic — the point is to pin `release`'s + // recovery, which exists so a *future* panic under this guard cannot + // turn an arena drop during unwinding into a process abort. + let poisoner = Arc::clone(&pool); + let previous_hook = std::panic::take_hook(); + std::panic::set_hook(Box::new(|_| {})); // keep the deliberate panic out of test output + let died = std::thread::spawn(move || { + let _guard = poisoner.inner.lock().expect("lock before poisoning"); + panic!("poison the pool mutex"); + }) + .join(); + std::panic::set_hook(previous_hook); + assert!(died.is_err(), "the poisoning thread must actually have panicked"); + + drop(arena); // must not panic + assert_eq!(pool.free_len(), 1, "the arena still returned to the poisoned pool"); + } + + #[test] + fn unpooled_does_not_return() { + let pool = ArenaPool::new(1, 1024); + drop(PooledSegmentedBuf::unpooled(SegmentedBuf::with_capacity(0, 1024))); + assert_eq!(pool.free_len(), 0, "unpooled buffer never reaches the pool"); + } + + #[test] + fn reused_arena_retains_capacity() { + let pool = ArenaPool::new(1, 16); + let mut buf = pool.try_acquire().unwrap(); + for i in 0..4u8 { + buf.extend_from_slice(&[i; 10]); + } + let (segs, cap) = (buf.num_segments(), buf.allocated_capacity()); + assert!(segs >= 2, "grew multiple segments"); + drop(buf); + let reused = pool.try_acquire().unwrap(); + assert!(reused.is_empty(), "reset for reuse"); + assert_eq!(reused.num_segments(), segs, "retained segments"); + assert_eq!(reused.allocated_capacity(), cap, "no realloc"); + } +} diff --git a/crates/fgumi-sort/src/bgzf_io.rs b/crates/fgumi-sort/src/bgzf_io.rs index 3c95dd657..b2be41a43 100644 --- a/crates/fgumi-sort/src/bgzf_io.rs +++ b/crates/fgumi-sort/src/bgzf_io.rs @@ -15,38 +15,9 @@ use std::collections::BTreeMap; use std::io::{BufWriter, Write}; use std::sync::Arc; -/// Padding beyond the frame size for the staging buffer capacity. +/// Padding beyond `BGZF_MAX_BLOCK_SIZE` for the staging buffer capacity. const STAGING_PADDING: usize = 4096; -/// Uncompressed bytes per zstd **spill** frame. -/// -/// BGZF mandates blocks of at most 64 KiB; zstd has no such limit, and spill -/// files are a private temporary format that nothing outside this crate reads. -/// Pinning zstd frames to the BGZF ceiling was therefore vestigial, and it set -/// the merge's block count -- 5,368,249 blocks on a measured 89-way merge, of -/// which the consumer takes roughly one per round trip. Fewer, larger frames cut -/// round trips proportionally. -/// -/// Raising this costs memory in two places, both worth watching: the per-worker -/// zstd decompress buffer ([`crate::worker_pool::zstd_decomp_cap`]) and the -/// per-file reorder buffer, which holds up to `PHASE2_DECOMP_CAP` *uncompressed* -/// frames for each of K files. Read-ahead itself does not scale with it, because -/// the refill allowance is sized in bytes rather than blocks. -pub(crate) const SPILL_FRAME_BYTES: usize = BGZF_MAX_BLOCK_SIZE; - -/// Bytes a staging buffer accumulates before it submits a compress job. -/// -/// Keyed on the codec so BAM output stays inside the BGZF format: output always -/// uses [`SpillCodec::Bgzf`], so it cannot pick up a spill-only frame size by -/// accident. -#[must_use] -pub(crate) const fn spill_frame_bytes(codec: SpillCodec) -> usize { - match codec { - SpillCodec::Bgzf => BGZF_MAX_BLOCK_SIZE, - SpillCodec::Zstd => SPILL_FRAME_BYTES, - } -} - /// Per-block position notification emitted by the I/O writer loop when index /// generation is enabled. /// @@ -77,17 +48,6 @@ pub(crate) struct StagingBuffer { } impl StagingBuffer { - /// The permit pool that carries this writer's histograms. - /// - /// Exposed so a caller can retain the [`Arc`] and harvest - /// [`PermitPool::writer_stats`] *after* the output drain. The writer's own - /// `finish` consumes the staging to run that drain, so the pool is the only - /// handle that outlives it. See that method for why the drain must be - /// included. - pub(crate) fn permit_pool(&self) -> &Arc { - &self.permit_pool - } - /// Seconds this buffer's producer spent blocked waiting for an output /// permit, and the number of waits. See [`PermitPool::blocked`]. pub(crate) fn write_backpressure(&self) -> (f64, u64) { @@ -105,7 +65,7 @@ impl StagingBuffer { ) -> Self { Self { pool, - buf: Vec::with_capacity(spill_frame_bytes(codec) + STAGING_PADDING), + buf: Vec::with_capacity(BGZF_MAX_BLOCK_SIZE + STAGING_PADDING), next_serial: 0, result_tx, permit_pool, @@ -125,7 +85,7 @@ impl StagingBuffer { /// Returns true if the staging buffer has reached the BGZF block size threshold. #[inline] pub(crate) fn is_full(&self) -> bool { - self.buf.len() >= spill_frame_bytes(self.codec) + self.buf.len() >= BGZF_MAX_BLOCK_SIZE } /// Current uncompressed length of the pending (not-yet-flushed) block. @@ -165,9 +125,8 @@ impl StagingBuffer { self.permit_pool.acquire()?; let data = std::mem::replace(&mut self.buf, self.pool.buffer_pool.checkout()); - let want = spill_frame_bytes(self.codec) + STAGING_PADDING; - if self.buf.capacity() < want { - self.buf.reserve(want - self.buf.capacity()); + if self.buf.capacity() < BGZF_MAX_BLOCK_SIZE + STAGING_PADDING { + self.buf.reserve(BGZF_MAX_BLOCK_SIZE + STAGING_PADDING - self.buf.capacity()); } let serial = self.next_serial; @@ -204,7 +163,7 @@ impl StagingBuffer { pub(crate) fn write_chunked(&mut self, data: &[u8]) -> anyhow::Result<()> { let mut remaining = data; while !remaining.is_empty() { - let space = spill_frame_bytes(self.codec).saturating_sub(self.buf.len()); + let space = BGZF_MAX_BLOCK_SIZE.saturating_sub(self.buf.len()); let n = remaining.len().min(space); self.buf.extend_from_slice(&remaining[..n]); remaining = &remaining[n..]; @@ -233,7 +192,7 @@ fn write_block_in_order( // Best-effort: the indexing consumer keeps the receiver alive until finish. let _ = tx.send(BlockOffset { serial, compressed_start: *compressed_offset }); } - permit_pool.write_dur.time(|| writer.write_all(data))?; + writer.write_all(data)?; *compressed_offset += data.len() as u64; permit_pool.release(); Ok(()) @@ -293,10 +252,7 @@ fn io_writer_loop_inner( block_offset_tx: Option<&Sender>, ) -> Result<()> { let mut next_expected: u64 = 0; - // Value carries the arrival instant so the wait for an earlier serial is - // measurable; a block written straight through never enters the map and so - // records no wait, which is the correct reading. - let mut reorder_buf: BTreeMap, std::time::Instant)> = BTreeMap::new(); + let mut reorder_buf: BTreeMap> = BTreeMap::new(); let mut compressed_offset: u64 = 0; let tx = block_offset_tx; @@ -314,8 +270,7 @@ fn io_writer_loop_inner( )?; next_expected += 1; - while let Some((data, arrived)) = reorder_buf.remove(&next_expected) { - permit_pool.write_reorder_wait.record(crate::merge_trace::elapsed_nanos(arrived)); + while let Some(data) = reorder_buf.remove(&next_expected) { write_block_in_order( writer, next_expected, @@ -327,22 +282,15 @@ fn io_writer_loop_inner( next_expected += 1; } } else { - reorder_buf.insert(result.serial, (result.compressed, std::time::Instant::now())); + reorder_buf.insert(result.serial, result.compressed); // Permit held: released when this block is written in the cascade above. } - // Sampled on every arrival, in-order or not, so the depth reflects the - // queue the writer is actually carrying rather than only its bad moments. - // A block count rides a duration histogram, so `record_count` scales it - // into the microsecond lane; recording the raw count would bucket every - // depth below `BLOCKS_TO_NANOS` to zero and the report would read zero. - permit_pool.write_reorder_depth.record_count(reorder_buf.len() as u64); } // Drain remaining buffered blocks — any gap means a worker dropped a result. while let Some((&serial, _)) = reorder_buf.first_key_value() { if serial == next_expected { - let (data, arrived) = reorder_buf.remove(&serial).expect("key just checked"); - permit_pool.write_reorder_wait.record(crate::merge_trace::elapsed_nanos(arrived)); + let data = reorder_buf.remove(&serial).expect("key just checked"); write_block_in_order( writer, next_expected, @@ -387,7 +335,7 @@ mod tests { /// without the `ZSPILL_MAGIC` prefix (which a real chunk writer would write /// before invoking the loop). fn roundtrip_data(data: &[u8], codec: SpillCodec) -> Vec { - let pool = Arc::new(SortWorkerPool::new(2, 1, 6, codec, false)); + let pool = Arc::new(SortWorkerPool::new(2, 1, 6, codec)); let (result_tx, result_rx) = pool.compress_result_channel(); let buffer_pool = pool.buffer_pool.clone(); let permit_pool = make_permit_pool(&pool); @@ -425,7 +373,7 @@ mod tests { #[case(SpillCodec::Bgzf)] #[case(SpillCodec::Zstd)] fn test_staging_buffer_flush_empty_is_noop(#[case] codec: SpillCodec) { - let pool = Arc::new(SortWorkerPool::new(1, 1, 6, codec, false)); + let pool = Arc::new(SortWorkerPool::new(1, 1, 6, codec)); let (result_tx, _result_rx) = pool.compress_result_channel(); let permit_pool = make_permit_pool(&pool); @@ -453,7 +401,7 @@ mod tests { #[case(SpillCodec::Bgzf)] #[case(SpillCodec::Zstd)] fn test_staging_buffer_is_full(#[case] codec: SpillCodec) { - let pool = Arc::new(SortWorkerPool::new(1, 1, 6, codec, false)); + let pool = Arc::new(SortWorkerPool::new(1, 1, 6, codec)); let (result_tx, _result_rx) = pool.compress_result_channel(); let permit_pool = make_permit_pool(&pool); let mut staging = StagingBuffer::new( @@ -465,127 +413,21 @@ mod tests { ); assert!(!staging.is_full(), "empty buffer should not be full"); - staging.buf().extend(vec![0u8; spill_frame_bytes(codec)]); - assert!(staging.is_full(), "buffer at the codec's frame size should be full"); - - if let Ok(p) = Arc::try_unwrap(pool) { - p.shutdown(); - } - } - - /// BGZF's 64 KiB block ceiling is a *format* requirement; zstd has none. - /// - /// Spill files are a private temporary format, so pinning their zstd frames - /// to the BGZF limit is vestigial — and it sets the block count, which is - /// what the merge pays per round trip (measured: 5,368,249 blocks on an - /// 89-way merge, ~1 consumed per consumer round trip). BAM output must stay - /// at the BGZF size regardless, and it always uses the BGZF codec, so keying - /// the threshold on the codec keeps output correct by construction. - #[test] - fn test_only_zstd_spill_frames_escape_the_bgzf_block_ceiling() { - assert_eq!( - spill_frame_bytes(SpillCodec::Bgzf), - BGZF_MAX_BLOCK_SIZE, - "BGZF blocks are capped by the format and must not grow" - ); - assert!( - spill_frame_bytes(SpillCodec::Zstd) >= BGZF_MAX_BLOCK_SIZE, - "zstd frames must never be smaller than the BGZF block they replace" - ); - assert!( - crate::worker_pool::zstd_decomp_cap() >= spill_frame_bytes(SpillCodec::Zstd), - "the decompress buffer must hold the largest frame the writer can emit, or every \ - frame at the new size fails to decompress" - ); - assert!( - crate::worker_pool::MAX_ZSTD_FRAME_BYTES >= spill_frame_bytes(SpillCodec::Zstd), - "the read-side length guard must admit a frame the writer can emit" - ); - } + staging.buf().extend(vec![0u8; BGZF_MAX_BLOCK_SIZE]); + assert!(staging.is_full(), "buffer at BGZF_MAX_BLOCK_SIZE should be full"); - /// Both zstd frame caps must admit the largest frame the writer can emit. - /// - /// There are two, on two different read paths: the pool's, used by the merge, - /// and `zspill_stream`'s, used by consolidation. They were previously held - /// equal by a comment saying "kept in sync", which is not enforcement -- and - /// the consolidation path is exercised only by the spill-heavy - /// configurations, so a mismatch would pass the standard matrix and fail - /// exactly where spill volume is largest. - #[test] - fn test_frame_caps_admit_the_largest_frame_the_writer_emits() { - let frame = spill_frame_bytes(SpillCodec::Zstd); - assert!( - crate::worker_pool::zstd_decomp_cap() >= frame, - "the merge path's decompress cap is below the writer's frame size" - ); - assert!( - crate::zspill_stream::frame_decomp_cap() >= frame, - "the consolidation path's decompress cap is below the writer's frame size" - ); - } - - /// The spill writer's pre-flush budget must *be* the staging buffer's frame - /// size, not a copy of it. - /// - /// `PooledChunkWriter` pre-flushes so a record never straddles a frame - /// boundary -- which is what lets the merge borrow most records in place. That - /// budget was `BGZF_MAX_BLOCK_SIZE` outright, so it flushed at 64 KiB no - /// matter what the staging buffer was configured for, and raising the frame - /// size changed the block count *not at all*: 5,368,249 blocks measured at - /// both 64 KiB and 256 KiB. The sweep looked like a 4% regression rather than - /// a no-op, so nothing about the wall time revealed that the knob was inert. - /// - /// Third duplicated threshold found this way, after the two zstd decompress - /// caps. Derive, then pin. - #[rstest] - #[case(SpillCodec::Bgzf)] - #[case(SpillCodec::Zstd)] - fn test_spill_writer_pre_flushes_at_the_staging_frame_size(#[case] codec: SpillCodec) { - let pool = Arc::new(SortWorkerPool::new(1, 1, 6, codec, false)); - let dir = tempfile::tempdir().expect("tempdir"); - let writer = - crate::pooled_chunk_writer::PooledChunkWriter::::new( - Arc::clone(&pool), - &dir.path().join("c.keyed"), - codec, - ) - .expect("writer"); - assert_eq!( - writer.frame_bytes(), - spill_frame_bytes(codec), - "the writer's pre-flush budget must equal the frame size the staging buffer flushes at, or the frame size has no effect on the block count" - ); - drop(writer); if let Ok(p) = Arc::try_unwrap(pool) { p.shutdown(); } } - /// At the default frame size the derived cap must equal the 256 KiB constant - /// it replaced. - /// - /// This buffer is per-worker scratch touched once per decompressed frame, so - /// its size is a cache parameter. Deriving it with a *fixed* 4 MiB of slack - /// instead of a proportional 4x inflated it 16x at the default frame size and - /// cost 25% of merge wall (249.5s against a 199.2s baseline) with peak RSS - /// essentially unchanged -- a regression invisible to any memory check. - #[test] - fn test_default_frame_size_preserves_the_original_decompress_cap() { - assert_eq!(SPILL_FRAME_BYTES, BGZF_MAX_BLOCK_SIZE, "guard: default frame size"); - assert_eq!( - crate::worker_pool::zstd_decomp_cap(), - 256 * 1024, - "the derived cap must reproduce the tuned constant at the default frame size" - ); - } - #[rstest] #[case(SpillCodec::Bgzf)] #[case(SpillCodec::Zstd)] fn test_staging_buffer_write_chunked_large_data(#[case] codec: SpillCodec) { // Data larger than BGZF_MAX_BLOCK_SIZE must be split into multiple compress jobs. let large = vec![b'A'; BGZF_MAX_BLOCK_SIZE * 2 + 1000]; - let pool = Arc::new(SortWorkerPool::new(2, 1, 6, codec, false)); + let pool = Arc::new(SortWorkerPool::new(2, 1, 6, codec)); let (result_tx, result_rx) = pool.compress_result_channel(); let buffer_pool = pool.buffer_pool.clone(); let permit_pool = make_permit_pool(&pool); @@ -631,7 +473,7 @@ mod tests { let data1 = b"first block data".to_vec(); let data2 = b"second block data".to_vec(); - let pool = Arc::new(SortWorkerPool::new(2, 1, 6, codec, false)); + let pool = Arc::new(SortWorkerPool::new(2, 1, 6, codec)); let (result_tx, result_rx) = pool.compress_result_channel(); let buffer_pool = pool.buffer_pool.clone(); let permit_pool = Arc::new(PermitPool::new(4)); diff --git a/crates/fgumi-sort/src/chunk_sorter.rs b/crates/fgumi-sort/src/chunk_sorter.rs new file mode 100644 index 000000000..784d23aa7 --- /dev/null +++ b/crates/fgumi-sort/src/chunk_sorter.rs @@ -0,0 +1,1189 @@ +//! Lean in-memory buffering sorter for the P6 `SortBuffer` step. +//! +//! [`CoordinateChunkSorter`] is the buffer-management slice of the former +//! `CoordinateSortStream` (the streaming sort engine retired in P7) with **all +//! the disk and pool machinery removed**: it ingests records into a +//! [`RecordBuffer`], and on +//! demand par-sorts the buffer and drains it into an arena-backed sorted chunk +//! (`InMemoryChunk`, zero body copies — the arena is moved, +//! not copied). It never writes to disk, owns no +//! `SortWorkerPool`, no temp dirs, and no spill files — compressing and writing +//! a chunk is the `CompressSpill` step's job (P6 Phase-1 split). +//! +//! The driver (`SortBuffer`) calls [`push`](CoordinateChunkSorter::push) per +//! record; when it returns `true` (buffer at the memory limit), the driver takes +//! a sorted chunk via [`take_sorted_chunk`](CoordinateChunkSorter::take_sorted_chunk) +//! and emits it as a spill chunk. At end of input the driver takes one final +//! residual chunk the same way. +//! +//! # Parity +//! +//! `take_sorted_chunk` uses the global stable [`RecordBuffer::par_sort`] plus a +//! zero-copy arena drain — exactly the no-spill residual path in +//! `CoordinateSortStream::into_slot_setup`. The legacy with-spill multi-thread +//! path used `par_sort_into_chunks` to emit *k* sub-runs; emitting one +//! globally-sorted chunk instead is record-order identical (the sub-runs are +//! merged back by key with a lower-source-index tie-break, which for a stable +//! global sort is the same original-insertion order), and it keeps the residual +//! a single merge source so a `Parallel` `CompressSpill` reordering chunks can +//! never perturb the tie-break. + +use std::sync::Arc; + +use anyhow::Result; +#[cfg(test)] +use fgumi_raw_bam::RawRecord; +use fgumi_raw_bam::SamTag; +use rayon::ThreadPool; + +use crate::arena_pool::ArenaPool; +use crate::external::{ + KeyTypesSpec, LibraryLookup, TemplateKeyVariant, dropped_lane_error, + extract_template_key_inline, select_template_variant, verify_dropped_lanes, +}; +use crate::inline::{ + CbKey32, InMemoryChunk, RecordBuffer, TemplateKey, TemplateKey24, TemplateKey40, + TemplateLaneKey, TemplateRecordBuffer, TertKey32, +}; +use crate::keys::RawCoordinateKey; +use crate::memory_probe::force_mi_collect; + +/// In-memory coordinate buffering sorter that emits owned sorted chunks. +pub struct CoordinateChunkSorter { + buffer: RecordBuffer, + /// Private rayon pool sized to `--threads`, used for `par_sort` and the + /// parallel chunk materialization. + rayon_pool: ThreadPool, + /// Memory-usage threshold (bytes) at which `push` signals a chunk is due. + memory_limit: usize, + total_records: u64, + /// Bounded pool of reusable sort arenas (lever-1 RSS fix). The buffer's + /// backing `SegmentedBuf` is acquired here and returns (reset-for-reuse) + /// when the emitted chunk's `Arc` drops, so at most `N_arena` arenas are + /// ever live. Capacity 1 (default) matches legacy's one-arena-at-a-time + /// `drain_pending_spill` model. + arena_pool: Arc, + /// `true` once an arena is installed in `buffer` and not yet drained. + has_arena: bool, +} + +impl CoordinateChunkSorter { + /// Construct from a pre-sized buffer + rayon pool + arena pool. Built via + /// [`RawExternalSorter::into_coordinate_chunk_sorter`](crate::RawExternalSorter::into_coordinate_chunk_sorter). + /// `buffer` starts drained; the first [`ensure_arena`](Self::ensure_arena) + /// installs one from `arena_pool`. + pub(crate) fn new( + buffer: RecordBuffer, + rayon_pool: ThreadPool, + memory_limit: usize, + arena_pool: Arc, + ) -> Self { + Self { buffer, rayon_pool, memory_limit, total_records: 0, arena_pool, has_arena: false } + } + + /// Ensure the buffer has a backing arena before a fill. Idempotent while + /// one is installed. + /// + /// `false` has **two** meanings and they call for opposite responses, so a + /// caller that treats every `false` as backpressure will spin: + /// + /// - **Pool exhausted** — all `N_arena` arenas are in flight. This is the + /// backpressure signal: the caller (`SortBuffer`) reports `NoProgress` and + /// retries once an in-flight chunk is spilled and its arena returns. + /// Retrying is what makes progress. + /// - **Buffer non-empty** — a fill is already underway in the buffer's + /// current storage. Retrying cannot change this until the buffer is + /// drained, so the caller must keep filling and take a chunk, NOT + /// backpressure. See the comment on the check below for why an arena is + /// never swapped in mid-fill. + /// + /// Call this only at the start of a fill, where the second case cannot + /// arise, and the return value is unambiguous. + #[must_use] + pub fn ensure_arena(&mut self) -> bool { + if self.has_arena { + return true; + } + // Only acquire at the START of a fill (drained buffer). Never swap a + // partially-filled (default) buffer for a pool arena mid-fill — the + // arena install replaces the backing store, which would orphan the + // records already buffered. A caller that began filling the default + // (because the bounded pool was momentarily empty) finishes into it and + // drains it unpooled. + if !self.buffer.is_empty() { + return false; + } + if let Some(arena) = self.arena_pool.try_acquire() { + self.buffer.install_arena(arena); + self.has_arena = true; + true + } else { + false + } + } + + /// Push one raw BAM record into the buffer. + /// + /// Returns `true` when the buffer's memory usage has reached the configured + /// limit and the caller should take a sorted (spill) chunk via + /// [`Self::take_sorted_chunk`] before pushing more. The caller must have + /// called [`ensure_arena`](Self::ensure_arena) (returning `true`) first. + /// + /// # Errors + /// + /// Returns an error if the record cannot be parsed for its coordinate key + /// (e.g. truncated record). + pub fn push(&mut self, bam_bytes: &[u8]) -> Result { + // Best-effort: acquire a pool arena if one is free. A bounded, RSS- + // critical driver would call `ensure_arena` explicitly first and + // backpressure on `false`, so on that path the arena is always installed + // here. Callers (the `SortBuffer` RecordBatch path, unit tests) that + // don't pre-check simply fill the buffer's current (possibly pool, + // possibly default) storage — `take_sorted_chunk` returns it to the pool + // only when it is a pool arena, which the buffer's own wrapper records + // (`has_arena` gates acquisition here, not the return). + let _ = self.ensure_arena(); + // Count only after the record is successfully buffered, so a truncated + // record (push_coordinate error) never bumps `total_records` for a + // record that was never stored. + self.buffer.push_coordinate(bam_bytes)?; + self.total_records += 1; + Ok(self.buffer.memory_usage() >= self.memory_limit) + } + + /// Par-sort the current buffer, drain it into an arena-backed sorted chunk, + /// and leave the buffer drained for reuse. + /// + /// Returns an empty chunk if the buffer is empty (so the caller can skip + /// emitting an empty chunk). + /// + /// **Zero-copy materialisation.** The sorted buffer is drained into an + /// [`InMemoryChunk`] that *moves* the arena (`Arc`) and keeps + /// per-record `(key, offset, len)` references into it — no per-record byte + /// copy or allocation. (The pre-fix path materialised an owned + /// `Vec<(K, RawRecord)>`, copying + `malloc`-ing every record and doubling + /// peak RSS; the A/B smoke campaign flagged that regression vs the legacy + /// path, which already used this primitive.) + #[must_use] + pub fn take_sorted_chunk(&mut self) -> InMemoryChunk { + if self.buffer.is_empty() { + return InMemoryChunk::empty(); + } + { + let buffer = &mut self.buffer; + let rayon_pool = &self.rayon_pool; + rayon_pool.install(|| buffer.par_sort()); + } + // One drain for both cases: the buffer holds a `PooledSegmentedBuf`, + // which already records whether it came from the pool, and the drain + // moves that wrapper whole. A pool arena therefore returns + // (reset-for-reuse) when the chunk's last `Arc` drops — the reuse and + // cap-1 bound on the standalone coordinate path — while a buffer that + // bypassed the pool simply drops. Branching on a parallel `has_arena` + // bool would be a second source of truth for something the data already + // knows. `has_arena` is still cleared so the next fill re-acquires. + let chunk = self.buffer.drain_into_single_chunk(); + self.has_arena = false; + chunk + } + + /// The arena pool this sorter draws from (test/diagnostic). + #[cfg(test)] + #[must_use] + pub(crate) fn arena_pool(&self) -> &Arc { + &self.arena_pool + } + + /// Total records pushed so far. + #[must_use] + pub fn total_records(&self) -> u64 { + self.total_records + } + + /// `true` iff the buffer currently holds no records. + #[must_use] + pub fn buffer_is_empty(&self) -> bool { + self.buffer.is_empty() + } + + /// Construct a minimal `CoordinateChunkSorter` for unit tests. + /// + /// Mirrors the construction in + /// [`RawExternalSorter::into_coordinate_chunk_sorter`](crate::RawExternalSorter::into_coordinate_chunk_sorter): + /// a zero-capacity `RecordBuffer` (the pool hands out the first arena on the + /// initial `ensure_arena`/`push`), a single-thread rayon pool, and a + /// capacity-1 `ArenaPool` with the standard segment size. + /// + /// # Panics + /// + /// Panics if rayon cannot build a single-thread pool (which never occurs in + /// practice but is the contract of `ThreadPoolBuilder::build`). + #[cfg(any(test, feature = "test-utils"))] + #[must_use] + pub fn for_test(memory_limit: usize, n_ref: u32) -> Self { + let buffer = crate::inline::RecordBuffer::with_capacity(0, 0, n_ref); + let rayon_pool = + rayon::ThreadPoolBuilder::new().num_threads(1).build().expect("rayon pool for test"); + let arena_pool = crate::arena_pool::ArenaPool::new(1, crate::inline::SORT_SEGMENT_SIZE); + Self::new(buffer, rayon_pool, memory_limit, arena_pool) + } +} + +// ============================================================================ +// TemplateChunkSorter — template-coordinate (stable radix, like coordinate) +// ============================================================================ + +/// The narrowed lane-key buffer the [`TemplateChunkSorter`] sorts on. The +/// variant is chosen on the first `push` from `--key-types` (via +/// [`select_template_variant`]) and matches the batch `RawExternalSorter::sort` +/// dispatch one-for-one: `(cb, tertiary)` → `(false,false)`=`TemplateKey24` +/// (3 lanes), `(true,false)`=`CbKey32`, `(false,true)`=`TertKey32` (4 lanes), +/// `(true,true)`=`TemplateKey40` (5 lanes, the full key). The radix sort runs +/// on the narrower key for speed; the emitted chunk re-widens to the full +/// [`TemplateKey`] (see [`TemplateChunkSorter::take_sorted_chunk`]). +enum TemplateBuffer { + K24(TemplateRecordBuffer), + Cb32(TemplateRecordBuffer), + Tert32(TemplateRecordBuffer), + K40(TemplateRecordBuffer), +} + +/// Run `$body` against the inner `TemplateRecordBuffer` of whichever variant +/// `$self` holds, with the buffer bound to `$buf`. `$body` is monomorphized per +/// arm, so type inference resolves `K` (e.g. `TemplateLaneKey::from_full`, +/// `b.push`) from the matched buffer type. +macro_rules! with_template_buffer { + ($self:expr, $buf:ident => $body:expr) => { + match $self { + TemplateBuffer::K24($buf) => $body, + TemplateBuffer::Cb32($buf) => $body, + TemplateBuffer::Tert32($buf) => $body, + TemplateBuffer::K40($buf) => $body, + } + }; +} + +impl TemplateBuffer { + /// Build the buffer for the chosen `variant`, sized from the shared + /// (key-width-agnostic) capacity hints. The slight per-`K` ref-size + /// difference is immaterial for a capacity hint — the buffer grows as + /// needed — so a single estimate is used across variants. + fn from_variant( + variant: TemplateKeyVariant, + estimated_records: usize, + estimated_data_bytes: usize, + ) -> Self { + match (variant.cb, variant.tertiary) { + (false, false) => Self::K24(TemplateRecordBuffer::with_capacity( + estimated_records, + estimated_data_bytes, + )), + (true, false) => Self::Cb32(TemplateRecordBuffer::with_capacity( + estimated_records, + estimated_data_bytes, + )), + (false, true) => Self::Tert32(TemplateRecordBuffer::with_capacity( + estimated_records, + estimated_data_bytes, + )), + (true, true) => Self::K40(TemplateRecordBuffer::with_capacity( + estimated_records, + estimated_data_bytes, + )), + } + } + + /// Push a record under the narrowed key derived from its full key. + fn push_full(&mut self, bam_bytes: &[u8], full: &TemplateKey) -> Result<()> { + with_template_buffer!(self, b => b.push(bam_bytes, TemplateLaneKey::from_full(full))) + } + + fn memory_usage(&self) -> usize { + with_template_buffer!(self, b => b.memory_usage()) + } + + fn is_empty(&self) -> bool { + with_template_buffer!(self, b => b.is_empty()) + } + + fn par_sort(&mut self) { + with_template_buffer!(self, b => b.par_sort()); + } + + // Test-only: only the owned oracle path (`take_sorted_chunk_owned`) clears + // the buffer explicitly; the production arena drain empties it in place. + #[cfg(test)] + fn clear(&mut self) { + with_template_buffer!(self, b => b.clear()); + } + + /// Materialize the (already narrow-key-sorted) buffer into an owned chunk + /// keyed by the **full** [`TemplateKey`], re-extracted per record. Narrow-K + /// order equals full-key order because every dropped lane is verified + /// constant, so the chunk is correctly ordered for the downstream merge. + /// Must be invoked inside the sorter's `rayon_pool.install`. + /// + /// Test-only: the owned materialisation is retained purely as the byte-parity + /// oracle for [`drain_into_full_key_chunk`](Self::drain_into_full_key_chunk), + /// which is the production path. + #[cfg(test)] + fn materialize_full( + &self, + lib_lookup: &LibraryLookup, + cell_tag: Option, + cb_hasher: &ahash::RandomState, + ) -> Vec<(TemplateKey, RawRecord)> { + use rayon::prelude::*; + with_template_buffer!(self, b => b + .refs() + .par_iter() + .map(|r| { + let bam_bytes = b.get_record(r); + let full = extract_template_key_inline(bam_bytes, lib_lookup, cell_tag, cb_hasher); + (full, RawRecord::from(bam_bytes.to_vec())) + }) + .collect::>()) + } + + /// Zero-copy analogue of `materialize_full`: drain + /// the (already narrow-key-sorted) buffer into an arena-backed + /// `InMemoryChunk` keyed by the re-extracted full key, WITHOUT + /// copying record bodies (the arena is moved into the chunk). Produces the + /// same sorted order, keys, and body bytes as `materialize_full` — the only + /// difference is the records reference the moved arena instead of owned + /// `RawRecord`s. Must be invoked inside the sorter's `rayon_pool.install`. + fn drain_into_full_key_chunk( + &mut self, + lib_lookup: &LibraryLookup, + cell_tag: Option, + cb_hasher: &ahash::RandomState, + ) -> InMemoryChunk { + with_template_buffer!(self, b => b.drain_into_full_key_chunk(|body| { + extract_template_key_inline(body, lib_lookup, cell_tag, cb_hasher) + })) + } +} + +/// State that exists only once the first record has been seen: the chosen +/// narrowed-key variant, the full key of the first record (the dropped-lane +/// verify baseline), and the matching buffer. Bundling the three behind one +/// `Option` captures their shared "set together on the first push" invariant in +/// the type, so the hot path never unwraps. +struct Provisioned { + /// Full key of the first record; later records' dropped lanes are verified + /// against it. + first_key: TemplateKey, + /// The chosen narrowed-key variant (`--key-types` + first record). + variant: TemplateKeyVariant, + /// The narrowed-key buffer matching `variant`. + buffer: TemplateBuffer, +} + +/// In-memory template-coordinate buffering sorter that emits owned sorted +/// chunks. The template analogue of [`CoordinateChunkSorter`]: a +/// `TemplateRecordBuffer` + private rayon pool, with the library / cell-barcode +/// / MI key-extraction state the template key needs. Uses the same global stable +/// radix `TemplateRecordBuffer::par_sort`, so the single-residual-chunk parity +/// argument is identical to coordinate's. +/// +/// **Narrowed radix key (`--key-types`).** Like the batch `RawExternalSorter::sort` +/// path, the in-memory radix sort runs on the `KeyTypesSpec`-narrowed key +/// (`TemplateBuffer` — `TemplateKey24`/`CbKey32`/`TertKey32`/`TemplateKey40`), +/// chosen lazily on the first `push` once the first record is available to seed +/// `select_template_variant`. The narrower key is a **speed** optimization +/// (fewer radix passes). The emitted chunk re-widens to the full [`TemplateKey`] +/// in [`take_sorted_chunk`](Self::take_sorted_chunk) so the buffer-chain protocol +/// (`CompressSpill`/`SortMerge`) stays on the full key. +/// +/// Narrowing is order-preserving and validated: `--key-types` may only drop a +/// lane the dropped-lane validation (`verify_dropped_lanes`) proves constant, so +/// the narrow-key sort yields the identical order to a full-key sort. This path +/// therefore produces byte-for-byte the same output as the batch path AND rejects +/// the same dropped-lane violations. (The pre-P6 streaming path skipped this +/// validation entirely; restoring it makes `--key-types` actually function in +/// production.) +pub struct TemplateChunkSorter { + /// Variant + first-key + buffer, provisioned on the first `push` (`None` + /// until then; an empty stream leaves it `None`). + state: Option, + rayon_pool: ThreadPool, + memory_limit: usize, + /// Capacity hints for the lazily-built `TemplateBuffer` (key-width-agnostic). + estimated_records: usize, + estimated_data_bytes: usize, + lib_lookup: LibraryLookup, + cb_hasher: ahash::RandomState, + cell_tag: Option, + /// `--key-types` spec; selects which optional lanes may be dropped. + key_types: KeyTypesSpec, + /// Whether the header realizes >1 library ordinal (informs `Auto` selection). + header_library_varies: bool, + total_records: u64, +} + +impl TemplateChunkSorter { + /// Built via + /// [`RawExternalSorter::into_template_chunk_sorter`](crate::RawExternalSorter::into_template_chunk_sorter). + /// The narrowed-key buffer is deferred to the first `push` (the variant is + /// chosen from the first record), so this takes capacity hints rather than a + /// pre-built buffer. + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( + rayon_pool: ThreadPool, + memory_limit: usize, + estimated_records: usize, + estimated_data_bytes: usize, + lib_lookup: LibraryLookup, + cb_hasher: ahash::RandomState, + cell_tag: Option, + key_types: KeyTypesSpec, + header_library_varies: bool, + ) -> Self { + Self { + state: None, + rayon_pool, + memory_limit, + estimated_records, + estimated_data_bytes, + lib_lookup, + cb_hasher, + cell_tag, + key_types, + header_library_varies, + total_records: 0, + } + } + + /// Push one record (extracting the template key); returns `true` when the + /// buffer has reached the memory limit. + /// + /// On the first record the dropped-lane variant is provisioned from + /// `--key-types`; every record then has its dropped lanes verified against + /// the first record's, matching the batch `RawExternalSorter::sort` path. + /// + /// # Errors + /// + /// Returns an error if the record cannot be buffered (e.g. truncated), or if + /// a record carries a dropped-lane value (CB / MI / library) absent from the + /// input's first record. + pub fn push(&mut self, bam_bytes: &[u8]) -> Result { + let key = extract_template_key_inline( + bam_bytes, + &self.lib_lookup, + self.cell_tag, + &self.cb_hasher, + ); + + let memory_usage = if let Some(state) = self.state.as_mut() { + // Subsequent records: verify the lanes the chosen variant drops are + // constant relative to the first record, then buffer. + if let Some(violation) = verify_dropped_lanes(&state.first_key, &key, state.variant) { + let name = fgumi_raw_bam::RawRecordView::new(bam_bytes).read_name(); + return Err(dropped_lane_error(&String::from_utf8_lossy(name), violation)); + } + state.buffer.push_full(bam_bytes, &key)?; + state.buffer.memory_usage() + } else { + // First record: provision the variant from --key-types and build the + // matching narrowed-key buffer. Buffer the record BEFORE committing + // `self.state`, so a `push_full` error leaves the sorter unprovisioned + // rather than holding a first_key/variant from a never-buffered record. + let variant = + select_template_variant(Some(&key), self.key_types, self.header_library_varies); + let mut buffer = TemplateBuffer::from_variant( + variant, + self.estimated_records, + self.estimated_data_bytes, + ); + buffer.push_full(bam_bytes, &key)?; + let memory_usage = buffer.memory_usage(); + self.state = Some(Provisioned { first_key: key, variant, buffer }); + memory_usage + }; + + // Count only after the record is successfully buffered (past the + // dropped-lane check and push_full), so an error never bumps + // `total_records` for a record that was never stored. + self.total_records += 1; + Ok(memory_usage >= self.memory_limit) + } + + /// Par-sort (on the narrowed key) + drain the buffer into one **arena-backed** + /// `InMemoryChunk` keyed by the full [`TemplateKey`] — zero body + /// copies (the sort arena is moved into the chunk, the full key re-extracted + /// per record). Symmetric with [`CoordinateChunkSorter::take_sorted_chunk`]. + /// Both `data` and `refs` are cleared by the drain. + #[must_use] + pub fn take_sorted_chunk(&mut self) -> InMemoryChunk { + let Some(state) = self.state.as_mut() else { + return InMemoryChunk::default(); + }; + if state.buffer.is_empty() { + return InMemoryChunk::default(); + } + self.rayon_pool.install(|| state.buffer.par_sort()); + let lib_lookup = &self.lib_lookup; + let cell_tag = self.cell_tag; + let cb_hasher = &self.cb_hasher; + let chunk = self + .rayon_pool + .install(|| state.buffer.drain_into_full_key_chunk(lib_lookup, cell_tag, cb_hasher)); + force_mi_collect(); + chunk + } + + /// Test-only owned materialisation — the byte-parity oracle reference for + /// [`take_sorted_chunk`](Self::take_sorted_chunk). Par-sorts then copies each + /// record into an owned `Vec<(TemplateKey, RawRecord)>` via `materialize_full` + /// (the pre-arena behaviour), so a test can assert the arena chunk is + /// byte-identical to what the owned path produced. + #[cfg(test)] + #[must_use] + pub(crate) fn take_sorted_chunk_owned(&mut self) -> Vec<(TemplateKey, RawRecord)> { + let Some(state) = self.state.as_mut() else { + return Vec::new(); + }; + if state.buffer.is_empty() { + return Vec::new(); + } + self.rayon_pool.install(|| state.buffer.par_sort()); + let buffer_ref = &state.buffer; + let lib_lookup = &self.lib_lookup; + let cell_tag = self.cell_tag; + let cb_hasher = &self.cb_hasher; + let chunk = self + .rayon_pool + .install(|| buffer_ref.materialize_full(lib_lookup, cell_tag, cb_hasher)); + state.buffer.clear(); + force_mi_collect(); + chunk + } + + /// Total records pushed so far. + #[must_use] + pub fn total_records(&self) -> u64 { + self.total_records + } + + /// The narrowed-key variant chosen on the first push (`None` before any + /// push). Exposed for tests asserting `--key-types` narrowing selects the + /// expected lane width. + #[cfg(test)] + pub(crate) fn chosen_variant(&self) -> Option { + self.state.as_ref().map(|s| s.variant) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use crate::RawExternalSorter; + use crate::inline::InMemoryChunk; + use crate::keys::{RawCoordinateKey, SortOrder}; + use fgumi_raw_bam::testutil::make_bam_bytes; + use noodles::sam::Header; + + /// Build `n` coordinate records at descending positions on tid 0, so a + /// correct sort must reorder them to ascending position. + fn descending_pos_records(n: usize) -> Vec> { + (0..n) + .map(|i| { + let pos = i32::try_from(n - 1 - i).expect("pos fits i32"); + let name = format!("r{i:05}"); + make_bam_bytes(0, pos, 0, name.as_bytes(), &[], 80, -1, -1, &[]) + }) + .collect() + } + + /// Assert a chunk's keys are non-decreasing (each chunk is a sorted run). + fn assert_chunk_sorted(chunk: &InMemoryChunk) { + for i in 1..chunk.len() { + assert!( + chunk.key_at(i - 1).cmp(chunk.key_at(i)).is_le(), + "chunk not sorted by coordinate key" + ); + } + } + + /// A generous memory limit keeps everything in one residual chunk; `push` + /// never signals full, and the single chunk is fully sorted. + #[test] + fn no_spill_single_sorted_residual_chunk() { + let records = descending_pos_records(500); + let mut sorter = RawExternalSorter::new(SortOrder::Coordinate) + .memory_limit(256 * 1024 * 1024) + .threads(2) + .into_coordinate_chunk_sorter(&Header::default()) + .expect("build chunk sorter"); + + for r in &records { + assert!(!sorter.push(r).expect("push"), "no spill expected under a large memory limit"); + } + let residual = sorter.take_sorted_chunk(); + assert_eq!(residual.len(), records.len(), "residual must hold every record"); + assert_chunk_sorted(&residual); + assert_eq!(sorter.total_records(), records.len() as u64); + assert!(sorter.buffer_is_empty(), "buffer cleared after take"); + } + + /// The pooled arena must complete a full cycle: acquire, fill, drain into a + /// chunk, and return to the free list once that chunk is dropped — then be + /// handed out again. + /// + /// The pool's own tests cover it in isolation, and the spill test above + /// holds every chunk alive, so nothing observed the arena actually coming + /// back. That gap hid a real defect: the drain used to re-wrap the buffer + /// *inside* the pooled wrapper, which deref-coerced past it, orphaned the + /// arena, and left a wrong-sized default behind for the pool to reclaim. At + /// the default `capacity == 1` that retires the only slot permanently. + #[test] + fn pooled_arena_returns_to_the_pool_when_its_chunk_drops() { + let records = descending_pos_records(64); + let mut sorter = RawExternalSorter::new(SortOrder::Coordinate) + .memory_limit(8 * 1024 * 1024) + .threads(1) + .into_coordinate_chunk_sorter(&Header::default()) + .expect("build chunk sorter"); + let pool = Arc::clone(sorter.arena_pool()); + + assert!(sorter.ensure_arena(), "the pool has a free arena to install"); + assert_eq!(pool.free_len(), 0, "it is in flight, not on the free list"); + for r in &records { + sorter.push(r).expect("push"); + } + + let chunk = sorter.take_sorted_chunk(); + assert_eq!(chunk.len(), records.len()); + assert_chunk_sorted(&chunk); + assert_eq!(pool.free_len(), 0, "still held by the live chunk"); + + drop(chunk); + assert_eq!(pool.free_len(), 1, "the arena returned once its chunk dropped"); + assert!(sorter.ensure_arena(), "and the pool can hand it out again"); + } + + /// The pool-exhausted `false`: at `capacity == 1`, a live chunk holds the + /// only arena, so the next `ensure_arena` must refuse. This is the cap-1 + /// backpressure signal the RSS bound rests on — if it returned `true` here + /// a second arena would be minted and the bound would be gone. Only the + /// success path was covered. + #[test] + fn ensure_arena_is_false_while_the_only_arena_is_in_flight() { + let records = descending_pos_records(64); + let mut sorter = RawExternalSorter::new(SortOrder::Coordinate) + .memory_limit(8 * 1024 * 1024) + .threads(1) + .into_coordinate_chunk_sorter(&Header::default()) + .expect("build chunk sorter"); + + assert!(sorter.ensure_arena(), "the fresh pool hands out its one arena"); + for r in &records { + sorter.push(r).expect("push"); + } + let chunk = sorter.take_sorted_chunk(); + + // The chunk is alive, so the arena has not returned to the pool. + assert!(!sorter.ensure_arena(), "capacity 1: no arena while the chunk holds it"); + drop(chunk); + assert!(sorter.ensure_arena(), "and it is available again once the chunk drops"); + } + + /// The buffer-non-empty `false`, which is a *different* condition with the + /// opposite correct response: retrying cannot help until the buffer drains. + /// + /// It guards a correctness property, not just an optimisation — installing + /// an arena mid-fill replaces the backing store out from under records + /// already buffered, orphaning them. `push` calls `ensure_arena` + /// best-effort, so after the first record into a pool-less buffer the + /// buffer is non-empty and the refusal must hold. + #[test] + fn ensure_arena_is_false_mid_fill_so_the_backing_store_is_never_swapped() { + let records = descending_pos_records(4); + let mut sorter = RawExternalSorter::new(SortOrder::Coordinate) + .memory_limit(8 * 1024 * 1024) + .threads(1) + .into_coordinate_chunk_sorter(&Header::default()) + .expect("build chunk sorter"); + let pool = Arc::clone(sorter.arena_pool()); + + // Drain the pool first so `push`'s best-effort acquire finds nothing and + // the fill lands in the buffer's default storage. + let held = pool.try_acquire().expect("take the only arena"); + sorter.push(&records[0]).expect("push into the default storage"); + assert!(!sorter.ensure_arena(), "pool empty AND buffer non-empty"); + + // Returning the arena is not enough: the buffer is now non-empty, so the + // refusal must persist for the mid-fill reason alone. + drop(held); + assert_eq!(pool.free_len(), 1, "the arena is back on the free list"); + assert!(!sorter.ensure_arena(), "an available arena must still not be installed mid-fill",); + + // Draining the buffer clears the condition. + drop(sorter.take_sorted_chunk()); + assert!(sorter.ensure_arena(), "and once drained the arena installs"); + } + + /// A tiny memory limit forces mid-stream spill chunks; every chunk is a + /// sorted run and the union covers all records exactly once. + #[test] + fn tiny_limit_emits_multiple_sorted_chunks_covering_all_records() { + let records = descending_pos_records(2_000); + let mut sorter = RawExternalSorter::new(SortOrder::Coordinate) + .memory_limit(16 * 1024) // tiny → frequent spills + .threads(2) + .into_coordinate_chunk_sorter(&Header::default()) + .expect("build chunk sorter"); + + let mut chunks = Vec::new(); + for r in &records { + if sorter.push(r).expect("push") { + let chunk = sorter.take_sorted_chunk(); + assert!(!chunk.is_empty()); + assert_chunk_sorted(&chunk); + chunks.push(chunk); + } + } + let residual = sorter.take_sorted_chunk(); + if !residual.is_empty() { + assert_chunk_sorted(&residual); + chunks.push(residual); + } + + assert!(chunks.len() >= 2, "tiny limit should force at least one spill plus residual"); + // Assert identity, not just count: a spill path that drops record A and + // duplicates B keeps the total unchanged, so compare the multiset of + // record bytes across all chunks against the input. + let mut seen: Vec> = chunks + .iter() + .flat_map(|c| (0..c.len()).map(move |i| c.record_bytes(i).to_vec())) + .collect(); + seen.sort(); + let mut expected: Vec> = records.clone(); + expected.sort(); + assert_eq!(seen, expected, "chunk union must equal the input set (no drops/dups)"); + } + + /// Equal-key stability: when every record shares one coordinate, the global + /// stable `par_sort` must emit them in input order. This is the + /// output-identity-critical tie-break the single-residual design relies on + /// (vs samtools sort), so pin it directly — not just sortedness. + #[test] + fn equal_keys_preserve_input_order_within_chunk() { + // All at tid 0, pos 0 → identical coordinate key; distinct names make + // each record's bytes unique so input order is observable. + let records: Vec> = (0..400) + .map(|i| { + let name = format!("r{i:05}"); + make_bam_bytes(0, 0, 0, name.as_bytes(), &[], 60, -1, -1, &[]) + }) + .collect(); + let mut sorter = RawExternalSorter::new(SortOrder::Coordinate) + .memory_limit(256 * 1024 * 1024) + .threads(2) + .into_coordinate_chunk_sorter(&Header::default()) + .expect("build chunk sorter"); + for r in &records { + assert!(!sorter.push(r).expect("push")); + } + let chunk = sorter.take_sorted_chunk(); + assert_eq!(chunk.len(), records.len()); + for (i, rec) in records.iter().enumerate() { + assert_eq!( + chunk.record_bytes(i), + rec.as_slice(), + "equal-key record {i} emerged out of input order" + ); + } + } + + /// `--key-types` is honored in the buffer path: dropping the tertiary (MI) + /// lane and then feeding a record whose MI differs from the first record's is + /// rejected with the same actionable error the batch `sort()` path produces. + /// (Pre-P6 the streaming production path skipped this validation entirely.) + #[test] + fn template_chunk_sorter_rejects_dropped_lane_violation() { + use crate::external::KeyTypesSpec; + + // Raw BAM aux for `MI:i:` (tag bytes + 'i' type + i32 LE value). + fn mi_aux(v: i32) -> Vec { + let mut a = vec![b'M', b'I', b'i']; + a.extend_from_slice(&v.to_le_bytes()); + a + } + + let mut sorter = RawExternalSorter::new(SortOrder::TemplateCoordinate) + .memory_limit(256 * 1024 * 1024) + .threads(1) + .key_types(KeyTypesSpec::None) // drop all optional lanes incl. tertiary/MI + .into_template_chunk_sorter(&Header::default()) + .expect("build template chunk sorter"); + + let r1 = make_bam_bytes(0, 10, 0, b"r1", &[], 40, -1, -1, &mi_aux(1)); + let r2 = make_bam_bytes(0, 10, 0, b"r2", &[], 40, -1, -1, &mi_aux(2)); + assert!(!sorter.push(&r1).expect("first push provisions the variant")); + let err = + sorter.push(&r2).expect_err("a differing MI under --key-types none must be rejected"); + let msg = err.to_string(); + assert!( + msg.contains("MI") && msg.contains("--key-types mi"), + "unexpected dropped-lane error: {msg}" + ); + } + + /// With a consistent input (the dropped lanes are constant), `--key-types` + /// validation passes and the sort proceeds — the common case. + #[test] + fn template_chunk_sorter_accepts_constant_dropped_lanes() { + use crate::external::KeyTypesSpec; + let mut sorter = RawExternalSorter::new(SortOrder::TemplateCoordinate) + .memory_limit(256 * 1024 * 1024) + .threads(1) + .key_types(KeyTypesSpec::None) + .into_template_chunk_sorter(&Header::default()) + .expect("build"); + // No optional lanes present at all → nothing to violate. + for i in 0..200i32 { + let pos = 200 - i; + let rec = make_bam_bytes(0, pos, 0, format!("r{i}").as_bytes(), &[], 40, -1, -1, &[]); + sorter.push(&rec).expect("constant (absent) dropped lanes accepted"); + } + assert_eq!(sorter.take_sorted_chunk().len(), 200); + } + + /// Plain template records (no cell barcode, no MI, single library) under + /// `Auto` select the narrowest 3-lane [`TemplateKey24`] variant — confirming + /// the radix sort runs on the narrowed key, not the full key — while still + /// producing a correctly ordered chunk keyed by the full [`TemplateKey`]. + #[test] + fn template_chunk_sorter_selects_narrow_variant_without_cb_or_mi() { + use crate::external::KeyTypesSpec; + let mut sorter = RawExternalSorter::new(SortOrder::TemplateCoordinate) + .memory_limit(256 * 1024 * 1024) + .threads(2) + .key_types(KeyTypesSpec::Auto) + .into_template_chunk_sorter(&Header::default()) + .expect("build template chunk sorter"); + assert!(sorter.chosen_variant().is_none(), "no variant before the first push"); + + // Descending positions on tid 0 → a correct sort reorders to ascending. + for i in 0..300i32 { + let pos = 300 - i; + let rec = + make_bam_bytes(0, pos, 0, format!("r{i:05}").as_bytes(), &[], 40, -1, -1, &[]); + sorter.push(&rec).expect("plain record accepted"); + } + + let variant = sorter.chosen_variant().expect("variant provisioned on first push"); + assert_eq!( + variant.lanes(), + 3, + "no cb / no MI / single library must narrow to the 3-lane TemplateKey24" + ); + + let chunk = sorter.take_sorted_chunk_owned(); + assert_eq!(chunk.len(), 300, "every record retained"); + for w in chunk.windows(2) { + assert!( + w[0].0.cmp(&w[1].0).is_le(), + "narrow-key sort must yield a full-key-ordered chunk" + ); + } + } + + /// Output-identity: the narrowed-key (3-lane `TemplateKey24`) sort must emit + /// byte-identical record order to the full-key (5-lane `TemplateKey40`) sort, + /// **including equal-key ties** — template-coordinate order is + /// output-identity-critical, so pin identity, not just sortedness. The input + /// deliberately contains records sharing a full template key (same tid / pos / + /// name) but differing in a non-key field (mapq), so a stable sort must keep + /// each tie group in input order under both key widths. + #[test] + fn template_chunk_sorter_narrow_matches_full_key_order() { + use crate::external::KeyTypesSpec; + + let mut records: Vec> = Vec::new(); + for i in 0..120i32 { + let pos = (i % 8) + 1; // 8 distinct positions → coordinate-level ties + let name = format!("t{:03}", i % 30); // 30 names → full-key ties within a pos + for seq_len in [20usize, 40usize] { + // Same (tid, pos, name) but distinct seq length → identical + // template key, distinct bytes: an observable equal-key tie. + records.push(make_bam_bytes(0, pos, 0, name.as_bytes(), &[], seq_len, -1, -1, &[])); + } + } + + let sort_with = |key_types: KeyTypesSpec| -> (usize, Vec>) { + let mut sorter = RawExternalSorter::new(SortOrder::TemplateCoordinate) + .memory_limit(256 * 1024 * 1024) + .threads(2) + .key_types(key_types) + .into_template_chunk_sorter(&Header::default()) + .expect("build template chunk sorter"); + for r in &records { + sorter.push(r).expect("push"); + } + let lanes = sorter.chosen_variant().expect("variant provisioned").lanes(); + let out = sorter + .take_sorted_chunk_owned() + .into_iter() + .map(|(_, rec)| rec.as_ref().to_vec()) + .collect(); + (lanes, out) + }; + + let (narrow_lanes, narrow) = sort_with(KeyTypesSpec::Auto); + let (full_lanes, full) = sort_with(KeyTypesSpec::Full); + assert_eq!(narrow_lanes, 3, "Auto over cb/MI-free data must narrow to TemplateKey24"); + assert_eq!(full_lanes, 5, "Full must force the 5-lane TemplateKey40"); + assert_eq!(narrow.len(), records.len(), "every record retained"); + assert_eq!( + narrow, full, + "narrowed-key sort must emit byte-identical record order (incl. equal-key ties) \ + to the full-key sort" + ); + } + + /// Byte-parity oracle: the arena-backed production drain + /// (`take_sorted_chunk`, zero body copies) must produce byte-for-byte the + /// same sorted (full-key, body) sequence as the owned `materialize_full` + /// path (`take_sorted_chunk_owned`). Same records + the same deterministic stable + /// radix on two identical sorters ⇒ identical output, so any divergence in + /// the arena offset math, full-key re-extraction, or ordering fails here. + /// Exercises equal-key ties (same tid/pos/name, distinct seq length) under + /// both the narrowed (`Auto` → `TemplateKey24`) and full (`Full` → + /// `TemplateKey40`) key widths. + #[test] + fn template_arena_drain_matches_owned_materialize() { + use crate::external::KeyTypesSpec; + + let mut records: Vec> = Vec::new(); + for i in 0..120i32 { + let pos = (i % 8) + 1; + let name = format!("t{:03}", i % 30); + for seq_len in [20usize, 40usize] { + records.push(make_bam_bytes(0, pos, 0, name.as_bytes(), &[], seq_len, -1, -1, &[])); + } + } + + let build = |key_types: KeyTypesSpec| { + RawExternalSorter::new(SortOrder::TemplateCoordinate) + .memory_limit(256 * 1024 * 1024) + .threads(2) + .key_types(key_types) + .into_template_chunk_sorter(&Header::default()) + .expect("build template chunk sorter") + }; + + for key_types in [KeyTypesSpec::Auto, KeyTypesSpec::Full] { + let mut owned_sorter = build(key_types); + let mut arena_sorter = build(key_types); + for r in &records { + owned_sorter.push(r).expect("push owned"); + arena_sorter.push(r).expect("push arena"); + } + let owned = owned_sorter.take_sorted_chunk_owned(); + let arena = arena_sorter.take_sorted_chunk(); + + assert_eq!(arena.len(), owned.len(), "record count must match ({key_types:?})"); + assert_eq!(owned.len(), records.len(), "every record retained ({key_types:?})"); + for (i, (key, rec)) in owned.iter().enumerate() { + assert_eq!( + arena.key_at(i), + key, + "full key at {i} must match owned path ({key_types:?})" + ); + assert_eq!( + arena.record_bytes(i), + rec.as_ref(), + "record bytes at {i} must be byte-identical to owned path ({key_types:?})" + ); + } + } + } + + /// A worker copy must inherit the parent's provisioned lane variant. + /// + /// `fresh()` builds the per-worker copy of the arena front. Provisioning + /// (variant + dropped-lane baseline) is meant to happen once per sort, but + /// `fresh()` used to reset it to `None`, which has three consequences once + /// there is more than one worker: each `Auto` worker re-selects a lane from + /// its own first record, so two workers in one sort can emit different key + /// widths; each validates dropped lanes against its own baseline, so a lane + /// that varies only *across* workers goes unnoticed; and a worker that + /// receives no records seals to `K40` regardless of what the others chose. + /// + /// The empty worker is the sharpest of the three to assert, because it + /// isolates the inheritance: it pushes nothing, so any variant it reports + /// can only have come from the parent. + #[test] + #[allow(unsafe_code)] + fn fresh_worker_copy_inherits_the_provisioned_variant() { + use crate::arena_pool::PooledSegmentedBuf; + use crate::external::KeyTypesSpec; + use crate::segmented_buf::SegmentedBuf; + use crate::template_arena::{TemplateArenaAccumulator, TemplateMemChunk}; + use std::sync::Arc; + + let rec = make_bam_bytes(0, 1, 0, b"t000", &[], 20, -1, -1, &[]); + let mut parent = + TemplateArenaAccumulator::from_header(&Header::default(), None, KeyTypesSpec::Auto); + + let mut arena = SegmentedBuf::with_capacity(0, 1 << 20); + arena.reserve_full_capacity(); + // SAFETY: the slot is fully written (copy_from_slice) before any read. + let off = unsafe { arena.grow_uninit(rec.len()) }; + unsafe { arena.slice_mut(off, rec.len()) }.copy_from_slice(&rec); + parent.push(&rec, off as u64, u32::try_from(rec.len()).unwrap()).expect("parent push"); + + // Clone the worker AFTER the parent is provisioned — the order the + // "provision once per sort" contract requires. + let mut worker = parent.fresh(); + + let parent_chunk = parent.seal(Arc::new(PooledSegmentedBuf::unpooled(arena)), 1); + assert!( + matches!(parent_chunk, TemplateMemChunk::K24(_)), + "precondition: this record must make `Auto` select the narrow K24 lane", + ); + + // The worker received no records. Without inheritance it falls back to + // K40 and the sort emits two different key widths. + let worker_chunk = worker.seal( + Arc::new(PooledSegmentedBuf::unpooled(SegmentedBuf::with_capacity(0, 1 << 20))), + 1, + ); + assert!( + matches!(worker_chunk, TemplateMemChunk::K24(_)), + "an empty worker must seal to the parent's lane, not fall back to K40", + ); + assert_eq!(worker_chunk.len(), 0, "and it is still empty"); + } + + /// Worker copies must share ONE bounded sort pool, not build one each. + /// + /// The pool exists to keep the radix + gather on exactly `sort_threads` + /// threads instead of the global rayon pool, so that a spill does not + /// oversubscribe the pipeline's own workers. `sort_threads` is only known at + /// the first `seal`, which is *after* `fresh()` has made the copies — so if + /// each copy held its own holder, N workers would build N pools and put + /// `N × sort_threads` threads on the box, reintroducing the very + /// oversubscription the pool prevents, one level up. + /// + /// Asserted by pointer identity: both copies must hand back the same pool. + #[test] + fn worker_copies_share_one_bounded_sort_pool() { + use crate::external::KeyTypesSpec; + use crate::template_arena::TemplateArenaAccumulator; + + let parent = + TemplateArenaAccumulator::from_header(&Header::default(), None, KeyTypesSpec::Auto); + let mut a = parent.fresh(); + let mut b = parent.fresh(); + + // Sealing an unprovisioned accumulator returns early without touching the + // pool, so drive the pool getter directly. + let pa = std::ptr::from_ref(a.sort_pool_for_test(2)); + let pb = std::ptr::from_ref(b.sort_pool_for_test(2)); + assert!( + std::ptr::eq(pa, pb), + "worker copies must share one pool; separate pools mean N x sort_threads threads", + ); + } + + /// Byte-parity: the arena-front `TemplateArenaAccumulator` must produce the + /// same sorted (full-key, body) sequence as the owned `TemplateChunkSorter` + /// — same library/CB provisioning, same `--key-types` narrowed-lane radix, + /// same full-key re-extraction — for a mix of tied and distinct keys, under + /// both the narrowed (`Auto` → `TemplateKey24`) and full (`Full` → + /// `TemplateKey40`) widths. This pins the arena front (Inc 3) to the legacy + /// path before it is wired into the pipeline. + #[test] + #[allow(unsafe_code)] + fn template_arena_accumulator_matches_owned_sorter() { + use crate::arena_pool::PooledSegmentedBuf; + use crate::external::KeyTypesSpec; + use crate::segmented_buf::SegmentedBuf; + use crate::template_arena::TemplateArenaAccumulator; + use std::sync::Arc; + + let mut records: Vec> = Vec::new(); + for i in 0..120i32 { + let pos = (i % 8) + 1; + let name = format!("t{:03}", i % 30); + for seq_len in [20usize, 40usize] { + records.push(make_bam_bytes(0, pos, 0, name.as_bytes(), &[], seq_len, -1, -1, &[])); + } + } + + for key_types in [KeyTypesSpec::Auto, KeyTypesSpec::Full] { + // ---- Owned oracle ---- + let mut owned = RawExternalSorter::new(SortOrder::TemplateCoordinate) + .memory_limit(256 * 1024 * 1024) + .threads(2) + .key_types(key_types) + .into_template_chunk_sorter(&Header::default()) + .expect("build owned template sorter"); + for r in &records { + owned.push(r).expect("owned push"); + } + let owned_chunk = owned.take_sorted_chunk_owned(); + + // ---- Arena accumulator (bodies resident in a shared arena) ---- + let mut arena = SegmentedBuf::with_capacity(0, 1 << 20); + arena.reserve_full_capacity(); + let mut acc = + TemplateArenaAccumulator::from_header(&Header::default(), None, key_types); + for r in &records { + // SAFETY: each slot is fully written (copy_from_slice) before any read. + let off = unsafe { arena.grow_uninit(r.len()) }; + unsafe { arena.slice_mut(off, r.len()) }.copy_from_slice(r); + acc.push(r, off as u64, u32::try_from(r.len()).unwrap()).expect("arena push"); + } + let arena_arc = Arc::new(PooledSegmentedBuf::unpooled(arena)); + let arena_chunk = acc.seal(arena_arc, 2); + + assert_eq!(arena_chunk.len(), owned_chunk.len(), "record count ({key_types:?})"); + assert_eq!(owned_chunk.len(), records.len(), "every record retained ({key_types:?})"); + // Byte-parity is the invariant: the arena chunk's records, in sorted + // order, must be byte-identical to the owned oracle's. The arena chunk + // now carries the chosen narrow lane key (variant-agnostic here), and + // narrow-lane order equals full-key order, so record order matches the + // owned full-key sort exactly. + for (i, (_key, rec)) in owned_chunk.iter().enumerate() { + assert_eq!( + arena_chunk.record_bytes(i), + rec.as_ref(), + "record body at {i} must be byte-identical to owned path ({key_types:?})" + ); + } + } + } + + /// An empty template stream never builds a buffer and yields an empty + /// residual (the lazy `Option` stays `None`). + #[test] + fn template_chunk_sorter_empty_stream_yields_empty_residual() { + use crate::external::KeyTypesSpec; + let mut sorter = RawExternalSorter::new(SortOrder::TemplateCoordinate) + .memory_limit(256 * 1024 * 1024) + .threads(1) + .key_types(KeyTypesSpec::Auto) + .into_template_chunk_sorter(&Header::default()) + .expect("build template chunk sorter"); + assert!(sorter.take_sorted_chunk().is_empty()); + assert_eq!(sorter.total_records(), 0); + assert!(sorter.chosen_variant().is_none(), "no push → no variant"); + } + + /// An empty stream yields an empty residual (the caller emits no chunk). + #[test] + fn empty_stream_yields_empty_residual() { + let mut sorter = RawExternalSorter::new(SortOrder::Coordinate) + .memory_limit(256 * 1024 * 1024) + .threads(1) + .into_coordinate_chunk_sorter(&Header::default()) + .expect("build chunk sorter"); + assert!(sorter.take_sorted_chunk().is_empty()); + assert_eq!(sorter.total_records(), 0); + } +} diff --git a/crates/fgumi-sort/src/external.rs b/crates/fgumi-sort/src/external.rs index 2950a298d..6107b8418 100644 --- a/crates/fgumi-sort/src/external.rs +++ b/crates/fgumi-sort/src/external.rs @@ -32,6 +32,7 @@ use crate::keys::{QuerynameComparator, RawSortKey, SortOrder}; use crate::memory_probe::{ BufferProbeStats, ConsumerProbeStats, MergeProbe, SpillProbe, force_mi_collect, log_snapshot, }; +use crate::merge_slots::SortMergeSlot; use crate::pooled_chunk_writer::PooledChunkWriter; use crate::read_ahead::{PooledInputStream, RawReadAheadReader, RecordSource}; use crate::tmp_dir_alloc::TmpDirAllocator; @@ -61,7 +62,7 @@ use tempfile::TempDir; /// Maximum number of records held in one in-memory chunk. /// /// Each key carries its ingest position within the chunk (see -/// [`RawSortKey::set_position`](crate::keys::RawSortKey::set_position)) so that +/// [`RawSortKey::set_position`]) so that /// the key is a *total* order and the chunk sort can be unstable. That position /// is a `u32`, so a chunk that grew past `u32::MAX` records would have to stamp /// two records with the same position — exact name/flag ties would stop being @@ -116,7 +117,7 @@ pub fn format_thread_counts(sort_threads: usize, merge_threads: usize) -> String /// Used to identify bottlenecks and validate thread architecture changes. /// All times are cumulative (multiple spill cycles accumulate). #[derive(Debug, Default)] -struct SortPhaseTimer { +pub(crate) struct SortPhaseTimer { /// Time reading records from input BAM (includes BGZF decompression). read_secs: f64, /// Time sorting in-memory buffers (rayon parallel sort or single-threaded). @@ -139,19 +140,13 @@ struct SortPhaseTimer { overall_start: Option, /// Tracks the start of the current read span (between spills). read_span_start: Option, - /// Whether the `--sort-stats` diagnostics are enabled for this run (run-scoped copy of - /// [`RawExternalSorter::sort_stats`]); consulted by this timer's `stat!` emitters. - sort_stats: bool, } impl SortPhaseTimer { - /// `sort_stats` is whether `--sort-stats` diagnostics are on for this run (the trailing - /// `bool` at call sites; tests pass `false`). - fn new(sort_stats: bool) -> Self { + fn new() -> Self { Self { overall_start: Some(Instant::now()), read_span_start: Some(Instant::now()), - sort_stats, ..Default::default() } } @@ -251,14 +246,7 @@ impl SortPhaseTimer { /// `max_temp_files` is reported so a run that consolidated says which limit /// it consolidated against. #[allow(clippy::cast_precision_loss)] - fn log_summary( - &self, - sort_threads: usize, - merge_threads: usize, - max_temp_files: usize, - phase1: &Phase1FloorInputs, - ) { - let sort_stats = self.sort_stats; + fn log_summary(&self, sort_threads: usize, merge_threads: usize, max_temp_files: usize) { let overall = self.overall_start.map_or(0.0, |s| s.elapsed().as_secs_f64()); // Guard against division by zero when sort completes in negligible time. let overall_nonzero = if overall > 0.0 { overall } else { f64::EPSILON }; @@ -270,15 +258,11 @@ impl SortPhaseTimer { let sort_secs = self.sort_secs; let spill_secs = self.spill_write_secs; - stat!(sort_stats, "=== Sort Phase Timing ==="); - stat!(sort_stats, " Read + decompress: {read_secs:.1}s ({read_pct:.0}%)"); - stat!( - sort_stats, - " In-memory sort: {sort_secs:.1}s ({sort_pct:.0}%) [{spill_count} spills]" - ); + info!("=== Sort Phase Timing ==="); + info!(" Read + decompress: {read_secs:.1}s ({read_pct:.0}%)"); + info!(" In-memory sort: {sort_secs:.1}s ({sort_pct:.0}%) [{spill_count} spills]"); let spill_mb = self.total_spill_bytes as f64 / (1024.0 * 1024.0); - stat!( - sort_stats, + info!( " Spill write: {spill_secs:.1}s ({spill_pct:.0}%) [{spill_count} writes, {spill_mb:.1} MB total]" ); if self.consolidate_count > 0 { @@ -290,266 +274,23 @@ impl SortPhaseTimer { // what it was up against. Reporting the number and not what to do // about it keeps the recommendation with the caller that chose the // limit -- the engine cannot know whether it was requested. - stat!( - sort_stats, + info!( " Consolidation: {cons_secs:.1}s ({cons_pct:.0}%) [{cons_count} merges, limit {max_temp_files}]" ); } if self.merge_secs > 0.0 { let merge_secs = self.merge_secs; let merge_pct = 100.0 * merge_secs / overall_nonzero; - stat!(sort_stats, " K-way merge: {merge_secs:.1}s ({merge_pct:.0}%)"); + info!(" K-way merge: {merge_secs:.1}s ({merge_pct:.0}%)"); } if self.write_output_secs > 0.0 { let write_secs = self.write_output_secs; let write_pct = 100.0 * write_secs / overall_nonzero; - stat!(sort_stats, " Write output: {write_secs:.1}s ({write_pct:.0}%)"); - } - stat!(sort_stats, " Total wall clock: {overall:.1}s"); - stat!(sort_stats, " Threads: {}", format_thread_counts(sort_threads, merge_threads)); - self.log_phase1_floor(phase1); - stat!(sort_stats, "========================="); - } - - /// Which of three limits Phase 1's *ingest* is against, and what is - /// recoverable without doing less work. - /// - /// Scoped to the read span rather than the whole phase, deliberately. The - /// in-memory sort is parallel (rayon) and the spill write-out overlaps the - /// next read, so folding them in would put parallel work on the same side of - /// the comparison as one thread's serial CPU and report the difference as - /// "coordination" -- naming a limit that is not there. The read span is the - /// part where one thread reads every record while the pool feeds it, which is - /// exactly the shape [`crate::merge_headroom`] models. - /// - /// Externally sampled, this phase is 60% of a whole-genome sort's wall clock - /// with its main thread 91% busy while 16 cores average 5.3. If that holds - /// in-process, the binding limit is the ingest thread and no amount of - /// additional worker capacity moves it. - fn log_phase1_floor(&self, phase1: &Phase1FloorInputs) { - let sort_stats = self.sort_stats; - if self.read_secs <= 0.0 { - return; - } - let ingest = phase1.ingest; - let floors = crate::merge_headroom::MergeFloors { - loop_secs: self.read_secs, - consumer_secs: (self.read_secs - ingest.park_secs).max(0.0), - worker_busy_secs: phase1.input_busy_secs, - threads: phase1.threads.max(1), - }; - stat!(sort_stats, " Phase 1 ingest floor: {} is the limit", floors.binding().label()); - stat!( - sort_stats, - " ingest serial {:.1}s | worker capacity {:.1}s ({} threads) | read span {:.1}s", - floors.consumer_secs, - floors.worker_floor_secs(), - floors.threads, - self.read_secs - ); - stat!( - sort_stats, - " recoverable without doing less work: {:.1}s ({:.0}% of the read span)", - floors.recoverable_secs(), - 100.0 * floors.recoverable_share() - ); - if ingest.parks > 0 { - // Mean park separates a supply problem from a handoff problem: many - // short parks and few long ones need opposite fixes, and the totals - // alone cannot tell them apart. - stat!( - sort_stats, - " ingest parked {:.1}s over {} parks ({:.0} us each): {} starved, {} head-of-line", - ingest.park_secs, - ingest.parks, - ingest.mean_park_micros().unwrap_or(0.0), - ingest.parks_starved, - ingest.parks_head_of_line - ); - } else { - stat!(sort_stats, " ingest never parked: the pool always had the next block ready"); - } - if ingest.spill_waits > 0 { - stat!( - sort_stats, - " waited {:.1}s over {} spill handoffs (outside every phase bucket)", - ingest.spill_wait_secs, - ingest.spill_waits - ); - } - Self::log_reader_partition(&phase1.reader, sort_stats); - self.log_ingest_partition(phase1); - } - - /// What the input reader's exclusively-owned thread spent its time on. - /// - /// The reader is Phase 1's *second* serial resource and the one the floor - /// line cannot see: `phase1_input_busy_secs` folds it in with decompression - /// and divides by the thread count, which is right for a step any worker may - /// run and wrong for one only worker 0 ever runs. So a reader at its limit - /// disappears into a worker-capacity figure that looks comfortable. - /// - /// The split that matters is refill against framing, because they have - /// nothing in common: refill time is the disk (compare it to the volume's - /// measured ceiling) and framing time is header parse, per-block allocation - /// and the body copy (compare it to the `raw_block_read` bench). Reporting - /// only the step total gives a number -- 24.3 us/block on the production - /// cell -- that is consistent with either being the whole cost. - #[allow(clippy::cast_precision_loss, reason = "call and byte counts stay below 2^52")] - fn log_reader_partition(reader: &crate::phase1_stats::ReaderReport, sort_stats: bool) { - if reader.batches == 0 { - return; - } - stat!( - sort_stats, - " Phase 1 input reader: {:.1}s over {} batches ({:.1} us/block, {} blocks)", - reader.step_secs, - reader.batches, - reader.per_block_micros(reader.step_secs), - reader.blocks - ); - if reader.refills > 0 { - stat!( - sort_stats, - " refill reads: {:.1}s over {} calls ({:.2} ms each, {:.0} MB/s, {:.1} GB)", - reader.refill_secs, - reader.refills, - reader.refill_secs * 1000.0 / reader.refills as f64, - reader.refill_mb_per_sec(), - reader.refill_bytes as f64 / 1e9 - ); - } else { - stat!( - sort_stats, - " refill reads: none on this thread (async reader, or input already buffered)" - ); - } - stat!( - sort_stats, - " framing: {:.1}s ({:.1} us/block)", - reader.framing_secs(), - reader.per_block_micros(reader.framing_secs()) - ); - stat!( - sort_stats, - " dispatch: {:.1}s ({:.1} us/block)", - reader.dispatch_secs, - reader.per_block_micros(reader.dispatch_secs) - ); - stat!(sort_stats, " refill latency: {}", reader.refill_latency.summary()); - // The discriminator. A `read()` that got slower because the thread kept - // losing the CPU shows runqueue wait and extra timeslices; one that got - // slower because the device or the memory system delivered fewer bytes - // per second shows neither, because the thread was blocked on I/O - // throughout. Wall-clock timing cannot separate those, and the whole - // question of what a busier worker pool does to the reader turns on it. - if reader.refill_sched_samples > 0 { - stat!( - sort_stats, - " refill scheduling: runqueue wait {:.2}s ({:.1}% of refill), on-CPU {:.2}s, \ - {} timeslices over {} sampled calls", - reader.refill_runqueue_secs, - 100.0 * reader.refill_runqueue_share(), - reader.refill_oncpu_secs, - reader.refill_timeslices, - reader.refill_sched_samples, - ); - } else { - stat!( - sort_stats, - " refill scheduling: unavailable (needs /proc/thread-self/schedstat)" - ); - } - stat!( - sort_stats, - " unattributed: {:.1}s ({:.0}% of the step)", - reader.residual_secs(), - 100.0 * reader.residual_share() - ); - } - - /// What the ingest thread's serial CPU is made of, per segment. - /// - /// Correction happens at the sampled scale and scaling second: subtracting - /// one clock pair from the *scaled* total instead would understate the - /// correction by the scale factor, which on a 1-in-1021 sample is three - /// orders of magnitude. - #[allow(clippy::cast_precision_loss, reason = "record counts stay below 2^52")] - fn log_ingest_partition(&self, phase1: &Phase1FloorInputs) { - let sort_stats = self.sort_stats; - if phase1.samples == 0 || phase1.records == 0 { - return; - } - let scale = phase1.records as f64 / phase1.samples as f64; - let segments = - phase1.sample.corrected(phase1.samples, phase1.clock_overhead_nanos).scaled(scale); - let partition = crate::phase1_stats::IngestPartition { - segments, - read_secs: self.read_secs, - park_secs: phase1.ingest.park_secs, - }; - let per_record = |secs: f64| 1e9 * secs / phase1.records as f64; - stat!( - sort_stats, - " Ingest segments ({} samples of {} records, scaled {scale:.0}x, clock {}ns/pair)", - phase1.samples, - phase1.records, - phase1.clock_overhead_nanos - ); - stat!( - sort_stats, - " fetch next record: {:.1}s ({:.0} ns/rec) [includes {:.1}s parked]", - segments.fetch, - per_record(segments.fetch), - partition.park_secs - ); - stat!( - sort_stats, - " push to arena: {:.1}s ({:.0} ns/rec)", - segments.push, - per_record(segments.push) - ); - stat!( - sort_stats, - " progress tick: {:.1}s ({:.0} ns/rec)", - segments.tick, - per_record(segments.tick) - ); - stat!( - sort_stats, - " probe + mem check: {:.1}s ({:.0} ns/rec)", - segments.probe, - per_record(segments.probe) - ); - // Exact, not sampled, and reported outside the partition above: both - // costs are per-batch or per-chunk, so the per-record sampler cannot see - // them honestly. `dispatch` is what deferral costs the serial thread; - // `barrier` is the extraction the pool failed to hide. - let census = phase1.key_overlap; - if let Some(pct) = census.overlap_percent() { - stat!( - sort_stats, - " Deferred key extraction: {pct:.1}% overlapped ({} of {} records; \ - {} keyed at a barrier)", - census.overlapped_records, - census.total_records(), - census.barrier_records, - ); - stat!( - sort_stats, - " dispatch (exact): {:.1}s barrier wait (exact): {:.1}s", - census.dispatch_secs, - census.barrier_secs, - ); + info!(" Write output: {write_secs:.1}s ({write_pct:.0}%)"); } - // Signed: a negative residual means the segments over-attribute, which is - // the failure mode this partition exists to make visible. - stat!( - sort_stats, - " unattributed: {:+.1}s ({:+.0}% of the read span)", - partition.residual_secs(), - 100.0 * partition.residual_share() - ); + info!(" Total wall clock: {overall:.1}s"); + info!(" Threads: {}", format_thread_counts(sort_threads, merge_threads)); + info!("========================="); } } @@ -557,35 +298,6 @@ impl SortPhaseTimer { // Library Lookup for Template-Coordinate Sort // ============================================================================ -/// What the Phase 1 floor line needs that [`SortPhaseTimer`] cannot see: the -/// pool's worker time and the ingest thread's waits. -#[derive(Debug, Clone, Copy, Default)] -pub(crate) struct Phase1FloorInputs { - /// Worker seconds reading and decompressing input blocks. - pub(crate) input_busy_secs: f64, - /// Active Phase 1 worker threads the busy total is spread over. - pub(crate) threads: usize, - /// What the ingest thread waited for. - pub(crate) ingest: crate::phase1_stats::Phase1IngestReport, - /// Raw (unscaled, uncorrected) sub-phase sample from the ingest loop, when - /// that loop is instrumented. Left at zero by the orders that are not. - pub(crate) sample: crate::phase1_stats::IngestSample, - /// Records timed for `sample`. - pub(crate) samples: u64, - /// Records the loop processed, the numerator of the sampling scale. - pub(crate) records: u64, - /// Measured cost of one `Instant::now()`/`elapsed()` pair, subtracted once - /// per segment per sample. - pub(crate) clock_overhead_nanos: u64, - /// How the exclusively-owned input reader spent its time. Unlike `sample`, - /// this is order-independent -- every sort order reads the same way -- so it - /// is filled in on all four paths. - pub(crate) reader: crate::phase1_stats::ReaderReport, - /// What deferred key extraction achieved, on the orders that defer it - /// (template-coordinate). Default (all zero) elsewhere. - pub(crate) key_overlap: crate::phase1_keys::KeyOverlapCensus, -} - /// Deterministic hasher for cell barcode hashing in template-coordinate sort. /// /// Uses arbitrary fixed seeds so that hash values are reproducible across runs. @@ -793,41 +505,14 @@ fn process_umask() -> u32 { u32::from(previous) } -/// RG-id to library-ordinal map, hashed with `ahash` rather than std's `SipHash`. -/// -/// `ordinal_from_rg` runs once per record on the sort's serial Phase 1 thread -- -/// the thread that sets 60% of a spill-heavy sort's wall clock -- and hashing a -/// ~40-byte RG id with `SipHash` measured **6.1% of that thread** on -/// `1kg-wgs-HG00096` (`core::hash::sip::Hasher::write`, third-largest entry in -/// its profile). The keys come from the header of a BAM the caller chose to -/// sort, so there is no adversarial-input exposure that would argue for -/// `SipHash`'s collision resistance here. -/// -/// Unlike [`cb_hasher`] and [`LibraryLookup::hasher`], this one is left at -/// `ahash`'s default randomly-seeded state rather than `with_seeds`. Those two -/// feed hash values *into the sort key*, so they must be identical across -/// processes or the same input would sort differently run to run. This map's -/// hasher is never observed: it is probed only by `get`, and -/// `distinct_header_ordinals` collects its values into a set, so neither the -/// hash values nor the iteration order reaches an output. -type RgOrdinalMap = HashMap, u32, ahash::RandomState>; - /// Maps read group ID -> library ordinal for O(1) comparison. /// /// Pre-computes ordinals by sorting library names alphabetically. /// Empty/unknown library sorts first (ordinal 0). -/// -/// `Clone` is derived so the deferred key-extraction context can own a copy -/// rather than borrow one, and it is **seed-preserving on purpose**: cloning -/// `ahash::RandomState` copies its keys rather than drawing new ones, so a clone -/// hashes read names identically to its original. Anything that reseeded here -/// would change the template-coordinate sort key and break byte-identity -/// silently — the output would still be correctly sorted, just not the same -/// order twice. #[derive(Clone)] pub struct LibraryLookup { /// RG ID -> library ordinal - rg_to_ordinal: RgOrdinalMap, + rg_to_ordinal: HashMap, u32>, /// Deterministic hasher for read name hashing, constructed once for reuse. hasher: ahash::RandomState, } @@ -859,7 +544,7 @@ impl LibraryLookup { } // Build RG ID -> ordinal mapping - let rg_to_ordinal: RgOrdinalMap = header + let rg_to_ordinal: HashMap, u32> = header .read_groups() .iter() .map(|(id, rg)| { @@ -1385,8 +1070,11 @@ impl GenericKeyedChunkReader { /// individual `RawRecord`s and produces `Owned` chunks — the /// per-record allocation is sunk cost, so we preserve the original /// zero-copy `mem::swap` merge bridge. -pub(crate) enum MemorySources { +pub enum MemorySources { + /// Inline-buffer (coordinate / template) residual chunks sharing an + /// `Arc` — zero per-record allocation. Shared(Vec>), + /// Queryname-style residual chunks (each record an owned `RawRecord`). Owned(Vec>), } @@ -1559,42 +1247,10 @@ struct MergeConsumerDiag { reassembled: u64, } -/// Per-merge tallies of the record-fetch fast/slow split. -/// -/// [`RECORD_BORROWED`] and [`RECORD_REASSEMBLED`] used to be incremented per -/// record, directly from the merge consumer. That is the one serial thread that -/// touches every record of the merge, and it is the merge's binding floor on the -/// measured cell (`consumer serial 112.9s` against `worker capacity 89.4s` and a -/// 157.1s loop). A relaxed `fetch_add` there is not free: on aarch64 it is an -/// outline-atomics call into `__aarch64_ldadd8_relax`, the same helper that -/// measured 28% of this thread's cycles while the progress counter used it -/// per-record (see [`crate::progress_batch`]). -/// -/// Counting into locals and publishing once per merge keeps every number the -/// reports read -- they consume the statics as a before/after delta around the -/// loop -- and removes two atomics per record from the critical path. Publish -/// *before* that delta is read, or the merge reports zero. -#[derive(Default)] -struct RecordFetchCounts { - /// Records handed over borrowed from the current decompressed block. - borrowed: u64, - /// Records reassembled into scratch because they straddled a block boundary. - reassembled: u64, -} - -impl RecordFetchCounts { - /// Fold this merge's tallies into the process-wide totals. - fn publish(&self) { - RECORD_BORROWED.fetch_add(self.borrowed, std::sync::atomic::Ordering::Relaxed); - RECORD_REASSEMBLED.fetch_add(self.reassembled, std::sync::atomic::Ordering::Relaxed); - } -} - #[inline] fn winner_record_bytes<'a, K: RawSortKey + Default + 'static>( source: &'a ChunkSource, consumer: Option<&'a MainThreadChunkConsumer>, - counts: &mut RecordFetchCounts, ) -> Result<&'a [u8]> { match source { ChunkSource::PoolDisk { source_id, scratch } => { @@ -1611,14 +1267,12 @@ fn winner_record_bytes<'a, K: RawSortKey + Default + 'static>( // that touches every record, so a clock read here would cost more // than the step. What was unknown is how often the copy path fires // at all -- a frequency answers that, and the per-record cost is a - // memcpy of a ~100-byte record either way. The tally is a local - // (see [`RecordFetchCounts`]); an atomic per record on this thread - // is itself measurable. + // memcpy of a ~100-byte record either way. if let Some(borrowed) = consumer.current_record_bytes(*source_id) { - counts.borrowed += 1; + RECORD_BORROWED.fetch_add(1, std::sync::atomic::Ordering::Relaxed); Ok(borrowed) } else { - counts.reassembled += 1; + RECORD_REASSEMBLED.fetch_add(1, std::sync::atomic::Ordering::Relaxed); Ok(scratch.as_slice()) } } @@ -1708,9 +1362,6 @@ pub(crate) struct MainThreadChunkConsumer { /// Shared pool state, for the epoch clock and the trace counters that /// record both halves of a producer/consumer handoff. shared: Arc, - /// Decompression state, so this thread can serve itself a block that has been - /// read but not yet claimed instead of parking for a worker to do it. - decomp: crate::worker_pool::DecompressorSet, /// Source the previous block came from, and how many consecutive blocks /// have now come from it. Says whether the merge dwells on one run at a /// time -- in which case lookahead on that run would pay -- or hops. @@ -1849,11 +1500,9 @@ impl MainThreadChunkConsumer { ) -> Self { let parser_state = (0..files.len()).map(|_| SourceParserState::new()).collect(); let stalls = crate::merge_stalls::ConsumerStallTracker::new(files.len()); - let codec = files.first().map_or(crate::codec::SpillCodec::Bgzf, |f| f.codec); Self { files, parser_state, - decomp: crate::worker_pool::DecompressorSet::for_codec(codec), decompression_error, chunk_read_error, worker_panicked, @@ -2130,125 +1779,19 @@ impl MainThreadChunkConsumer { self.shared .phase2_awaited_source .store(source_id, std::sync::atomic::Ordering::Relaxed); - // Arm the park handshake. Workers stamp when the awaited block is - // first claimed and when the consumer is woken; the difference from - // `park_start` splits this park into additive stages. Cleared here - // rather than after so a stamp can only ever be attributed to a park - // that had already begun. - let ord = std::sync::atomic::Ordering::Relaxed; - self.shared.awaited_claim_nanos.store(0, ord); - self.shared.awaited_publish_nanos.store(0, ord); - let park_started_at = self.shared.now_nanos(); - // Sampled BEFORE the park, not after: the question is why nobody had - // already started on this block, and the pool's state on *resume* is - // the answer to a different question -- by then somebody has. - // Before sleeping: if the block this thread is waiting for has already - // been read and nobody has claimed it, decompress it here rather than - // parking. That trades ~53us of work for a ~190us wait on a worker - // that has to be woken first, and it removes the producer-to-consumer - // handoff instead of trying to schedule around it -- which is what - // eight scheduling interventions failed to do. - // - // Loops while it keeps succeeding: the same argument applies to the - // blocks after the one immediately needed, and serving them here is - // depth on exactly the file being drained, which is where >=94% of - // consumed blocks come from. Bounded by the FIFO running dry or the - // reorder cap, both of which `try_pop_raw_for_decompress` enforces, so - // this cannot spin. - // - // Only on the way into a park, never in place of merging: this thread - // is the serial bottleneck at 76s of a 166s merge, and taking work it - // could have delegated while it still has records to emit would be - // strictly worse. - if self.serve_self_instead_of_parking(source_id) { - // Self-service is not a park: leave `parked_yet` untouched so the - // first real park of this pull still records a `stalled_pull`. - // Serving before the first park would otherwise suppress that - // increment and understate the stall rate. - continue; - } - - let supply = self.shared.park_supply_now(); - // Published for `get_sort_priorities`: while this is set, a worker - // serves the awaited file before it drains output compression. - // Ordering is Release/Acquire-free on purpose -- a worker reading it - // one iteration late costs one deferred block, not correctness. - self.shared.consumer_parked.store(true, ord); - let park_start = Instant::now(); std::thread::park(); - self.shared.consumer_parked.store(false, ord); let parked_ns = crate::merge_trace::elapsed_nanos(park_start); - self.shared.park_supply.record(supply, parked_ns); self.stalls.record_park(source_id, parked_ns, !parked_yet); self.shared.consumer_trace.record_park(state, parked_ns, in_flight); - - // Split the park. `now_nanos` is the same clock the workers stamp - // with, so the segments are comparable; `parked_ns` comes from - // `Instant` and is only used for the existing counters. - let resumed_at = self.shared.now_nanos(); - let claim = self.shared.awaited_claim_nanos.load(ord); - let publish = self.shared.awaited_publish_nanos.load(ord); - let segments = crate::merge_stalls::split_park( - park_started_at, - (claim != 0).then_some(claim), - (publish != 0).then_some(publish), - resumed_at, - ); - // Depth on the critical path: how many blocks the awaited file had - // ready the moment the consumer could run again. A mean near 1 means - // every block is fetched on demand, which is a different problem - // from a slow fetch. - let (_, ready, _) = self.files[source_id].depths(); - self.shared.park_attribution.record(segments, claim != 0, ready as u64); parked_yet = true; } } - /// Decompress this source's already-read blocks on this thread, returning - /// whether anything was published. - /// - /// Called only on the way into a park. If the block the merge is waiting for - /// has been read and nobody has claimed it -- 12% of parks -- decompressing it - /// here trades ~53us of work for a ~190us wait on a worker that has to be - /// woken first. It removes the producer-to-consumer handoff instead of trying - /// to schedule around it, which is what eight scheduling interventions failed - /// to do. The decisive evidence: raising `PHASE2_DECOMP_CAP` 8 -> 128 cut parks - /// from 978,325 to 400,487 and left total park time at 89.6s against 89.7s, so - /// the cost is per block, not per park, and no amount of buffering reaches it. - /// - /// Loops while it keeps succeeding, because the same argument applies to the - /// blocks after the one immediately needed, and serving them here is depth on - /// exactly the file being drained -- where >=94% of consumed blocks come from - /// (median run is 1 block, but the mean is 31.7 and p99 is 512). It cannot - /// spin: each iteration either pops a block or stops, bounded by the FIFO - /// emptying or the reorder cap, both enforced inside - /// `try_pop_raw_for_decompress`. - /// - /// **Never in place of merging.** This thread is the serial bottleneck, 76s of - /// a 166s merge, so taking work it could have delegated while it still has - /// records to emit would be strictly worse. The only safe moment is the one - /// where it was about to sleep anyway. - fn serve_self_instead_of_parking(&mut self, source_id: usize) -> bool { - let mut served = false; - while crate::worker_pool::SortWorkerPool::consumer_decompress_one( - &self.shared, - &mut self.decomp, - &self.files[source_id], - source_id, - ) { - served = true; - } - if served { - self.shared.consumer_self_served.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - } - served - } - /// Parse the next record from a source's byte stream. /// - /// Handles the format: for `EMBEDDED_IN_RECORD` keys, reads [len(4)][record(len)]. - /// For keyed format, reads [key][len(4)][record(len)]. + /// Handles the format: for `EMBEDDED_IN_RECORD` keys, reads `[len(4)][record(len)]`. + /// For keyed format, reads `[key][len(4)][record(len)]`. fn parse_next_record(&mut self, source_id: usize, buf: &mut Vec) -> Result> { let mut len_buf = [0u8; 4]; @@ -2529,9 +2072,8 @@ struct RunFormer { /// the surviving file count can be far lower than the run count for reasons that /// have nothing to do with how ordered the input was -- which is the one thing /// this line exists to report. -fn log_run_formation(chunks_spilled: usize, runs: usize, sort_stats: bool) { - stat!( - sort_stats, +fn log_run_formation(chunks_spilled: usize, runs: usize) { + info!( "Spill runs: {runs} from {chunks_spilled} chunks ({} extended an existing run)", chunks_spilled.saturating_sub(runs) ); @@ -2606,6 +2148,45 @@ fn probe_stats(buf: &impl ProbeableBuffer) -> BufferProbeStats { } } +/// Read-stream policy for the Phase-1 input reader. +/// +/// R10-sync compatibility shim: `main`'s `--read-streams` CLI (#838/#846) is carried on the +/// arena engine so `commands/sort.rs` compiles and the flag keeps parsing; the arena reader +/// does not yet honor it. Forward-porting the measured multi-stream reader onto the arena +/// engine is tracked as R7b. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum ReadStreams { + /// Measure and grow. The default. + #[default] + Auto, + /// Exactly this many; `1` is the plain sequential reader. + Fixed(usize), +} + +impl std::fmt::Display for ReadStreams { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Auto => f.write_str("auto"), + Self::Fixed(n) => write!(f, "{n}"), + } + } +} + +impl std::str::FromStr for ReadStreams { + type Err = String; + + fn from_str(text: &str) -> Result { + if text.eq_ignore_ascii_case("auto") { + return Ok(Self::Auto); + } + match text.parse::() { + Ok(0) => Err("--read-streams must be `auto` or at least 1".to_string()), + Ok(n) => Ok(Self::Fixed(n)), + Err(_) => Err(format!("expected `auto` or a positive number, got `{text}`")), + } + } +} + /// Raw-bytes external sorter for BAM files. /// /// This sorter uses lazy record parsing to minimize memory usage and avoid @@ -2653,20 +2234,13 @@ pub struct RawExternalSorter { /// a modest allocation and let `Vec` grow on demand, while explicit limits /// pre-allocate the full budget upfront (preserving prior behavior). initial_capacity: Option, - /// How the input and the merge's spill files are read. See - /// `spill_reader` for why more than one stream is needed. - read_streams: ReadStreams, + /// When true, wrap input in a `PrefetchReader` for async I/O. + async_reader: bool, /// Which optional template-key lanes to retain (template-coordinate only). /// /// Defaults to [`KeyTypesSpec::Auto`], which provisions the narrowest key /// that fits the first record's optional lanes. key_types: KeyTypesSpec, - /// Whether to emit the `--sort-stats` performance diagnostics for this run. - /// - /// Run-scoped (a field on the sorter) rather than a process-global static, so two sorts - /// in one process (fgumi-sort is a library) keep independent settings. Copied into the - /// [`SortPhaseTimer`] and [`SortWorkerPool`] this sorter builds, which own the emitters. - sort_stats: bool, } /// RAII guard that ensures Phase 2 teardown runs on every exit path between @@ -2751,6 +2325,141 @@ impl Drop for Phase2Guard<'_, K> { } impl RawExternalSorter { + /// The configured sort order. Read by the typed-step `SortAndSpill` adapter + /// to pick the matching `*SortStream` variant. + #[must_use] + pub fn sort_order(&self) -> SortOrder { + self.sort_order + } + + /// The configured cell-barcode tag (template-coordinate CB hashing), if any. + /// Read when building a record-input arena strategy so it matches the + /// block-input template strategy exactly. + #[must_use] + pub fn cell_tag_value(&self) -> Option { + self.cell_tag + } + + /// The configured `--key-types` narrowing spec (template-coordinate only). + #[must_use] + pub fn key_types_spec(&self) -> KeyTypesSpec { + self.key_types + } + + /// The configured in-memory sort budget in bytes (per-run seal threshold). + #[must_use] + pub fn memory_limit_bytes(&self) -> usize { + self.memory_limit + } + + /// The configured base thread count (`--threads`). + /// + /// NOT the Phase-1 sort width — that is `phase1_threads()`, which applies + /// the `--sort-threads` override. Reading this as the sort width is the + /// D1.1 regression, pinned by `test_raw_sorter_phase_threads_builders`. + #[must_use] + pub fn num_threads(&self) -> usize { + self.threads + } + + /// Convert this configured sorter into a lean coordinate buffering sorter + /// for the P6 `SortBuffer` step. + /// + /// Same buffer sizing / rayon pool as the retired `into_coordinate_stream`, + /// but with **no** `SortWorkerPool`, temp dirs, or spill files — the + /// `SortBuffer` step only ingests + par-sorts + materializes chunks; the + /// `CompressSpill` step owns compression and disk I/O. + /// + /// # Errors + /// + /// Returns an error if `self.sort_order` is not `Coordinate`, the reference + /// count exceeds `u32::MAX`, or the rayon sort pool cannot be built. + pub fn into_coordinate_chunk_sorter( + self, + header: &Header, + ) -> Result { + anyhow::ensure!( + matches!(self.sort_order, SortOrder::Coordinate), + "into_coordinate_chunk_sorter requires SortOrder::Coordinate (got {:?})", + self.sort_order, + ); + let rayon_pool = self.build_sort_rayon_pool()?; + let nref = u32::try_from(header.reference_sequences().len()) + .context("reference sequence count exceeds u32::MAX")?; + let init_cap = self.effective_initial_capacity(); + // Per-record footprint estimate matches `into_coordinate_stream`. + let estimated_records = init_cap / 240; + // The record *bodies* arena comes from the `ArenaPool` on the first + // `ensure_arena`, so no data-byte capacity is requested here. That is not + // the same as allocating nothing: `with_capacity` still sizes the segment + // for the inline headers (`estimated_records * HEADER_SIZE`), and + // `install_arena` replaces and drops that buffer when the pooled arena + // lands. Transient header storage, in other words — not zero. + let buffer = RecordBuffer::with_capacity(estimated_records, 0, nref); + // Bounded reusable-arena pool (lever-1 RSS fix). Capacity 1 = legacy's + // one-arena-at-a-time footprint: the next chunk cannot fill until the + // prior chunk is spilled and its arena returns, bounding peak RSS to + // ~one arena ≈ base. (Tunable via `--sort-arenas` in a later increment.) + let arena_pool = crate::arena_pool::ArenaPool::new(1, crate::inline::SORT_SEGMENT_SIZE); + Ok(crate::chunk_sorter::CoordinateChunkSorter::new( + buffer, + rayon_pool, + self.memory_limit, + arena_pool, + )) + } + + /// Convert this configured sorter into a lean template-coordinate buffering + /// sorter for the P6 `SortBuffer` step (the template analogue of + /// [`Self::into_coordinate_chunk_sorter`] — same sizing as the retired + /// `into_template_coordinate_stream`, minus the pool/temp-dirs). + /// + /// # Errors + /// + /// Returns an error if `self.sort_order` is not `TemplateCoordinate` or the + /// rayon sort pool cannot be built. + pub fn into_template_chunk_sorter( + self, + header: &Header, + ) -> Result { + anyhow::ensure!( + matches!(self.sort_order, SortOrder::TemplateCoordinate), + "into_template_chunk_sorter requires SortOrder::TemplateCoordinate (got {:?})", + self.sort_order, + ); + let rayon_pool = self.build_sort_rayon_pool()?; + let lib_lookup = LibraryLookup::from_header(header); + let header_library_varies = lib_lookup.distinct_header_ordinals() > 1; + let cb_hasher = cb_hasher(); + // Capacity hints for the lazily-built narrowed-key buffer, derived the + // SAME way as `sort_template_coordinate_impl` so the two move together + // when `EST_BAM_BYTES_PER_TEMPLATE_RECORD` changes. That sibling is + // pinned by `estimator_bytes_per_record_matches_ref_size`; a hardcoded + // number here would drift silently past it. + // + // The variant — and so the exact key width — is not known until the first + // record, so this uses the WIDEST ref (`TemplateKey40`, 56 B). Over- + // estimating bytes-per-record under-estimates the record count, which + // only under-sizes a capacity hint: the safe direction. + let ref_bytes = std::mem::size_of::>(); + let data_bytes_per_record = TEMPLATE_HEADER_SIZE + EST_BAM_BYTES_PER_TEMPLATE_RECORD; + let bytes_per_record = data_bytes_per_record + ref_bytes; + let init_cap = self.effective_initial_capacity(); + let estimated_records = (init_cap / bytes_per_record).max(1); + let estimated_data_bytes = init_cap * data_bytes_per_record / bytes_per_record; + Ok(crate::chunk_sorter::TemplateChunkSorter::new( + rayon_pool, + self.memory_limit, + estimated_records, + estimated_data_bytes, + lib_lookup, + cb_hasher, + self.cell_tag, + self.key_types, + header_library_varies, + )) + } + /// Create a new raw external sorter with the given sort order. #[must_use] pub fn new(sort_order: SortOrder) -> Self { @@ -2769,19 +2478,11 @@ impl RawExternalSorter { max_temp_files: crate::fd_limit::FALLBACK_MAX_TEMP_FILES, cell_tag: None, initial_capacity: None, - read_streams: ReadStreams::default(), + async_reader: false, key_types: KeyTypesSpec::default(), - sort_stats: false, } } - /// Enable or disable the `--sort-stats` performance diagnostics for this run. - #[must_use] - pub fn sort_stats(mut self, enabled: bool) -> Self { - self.sort_stats = enabled; - self - } - /// Set the memory limit for in-memory sorting. #[must_use] pub fn memory_limit(mut self, limit: usize) -> Self { @@ -2841,6 +2542,10 @@ impl RawExternalSorter { } /// Effective Phase-1 worker count: `sort_threads` if set, else `threads`. + /// + /// Public so the streaming arena front (`SortBuffer::from_sorter`) sizes its + /// per-chunk sort with the resolved `--sort-threads` override rather than the + /// raw base `--threads` (`num_threads`), which silently ignores the override. #[must_use] pub fn phase1_threads(&self) -> usize { self.sort_threads.unwrap_or(self.threads).max(1) @@ -2934,6 +2639,26 @@ impl RawExternalSorter { self } + /// Set the read-stream policy (R10-sync compatibility shim — see [`ReadStreams`]). + /// + /// Accepted and discarded: the arena reader does not yet honor `--read-streams`. + /// Forward-porting `main`'s measured multi-stream reader (#838/#846) onto the arena + /// engine is tracked as R7b. + #[must_use] + pub fn read_streams(self, _streams: ReadStreams) -> Self { + self + } + + /// Enable `--sort-stats` phase-timing diagnostics (R10-sync compatibility shim). + /// + /// Accepted and discarded: `main`'s #855 rewired `commands/sort.rs` to a `.sort_stats(bool)` + /// builder (replacing the removed `set_sort_stats` global). The arena engine emits its own + /// diagnostics; wiring this flag to gate them is tracked as R7b. + #[must_use] + pub fn sort_stats(self, _enabled: bool) -> Self { + self + } + /// Set maximum temp files before consolidation. /// /// When the number of temp files exceeds this limit, the oldest files @@ -2981,10 +2706,13 @@ impl RawExternalSorter { self } - /// How to read the input and the merge's spill files. + /// Enable/disable the async prefetch reader on input. + /// + /// When enabled, the input BAM is wrapped in a `PrefetchReader` before the + /// BGZF layer, which overlaps block I/O with decompression. #[must_use] - pub fn read_streams(mut self, streams: ReadStreams) -> Self { - self.read_streams = streams; + pub fn async_reader(mut self, enabled: bool) -> Self { + self.async_reader = enabled; self } @@ -3071,16 +2799,7 @@ impl RawExternalSorter { pool: &std::sync::Arc, ) -> Result<()> { if let Some(prev) = pending.take() { - // Timed because it lands in no phase bucket: this sits between - // `end_read_span` and `time_sort`, so without its own counter the - // ingest thread's wall clock and its own CPU cannot be reconciled and - // the difference shows up only as an unexplained residual against - // total wall clock. - let waited_at = Instant::now(); prev.handle.wait()?; - pool.phase1_ingest_stats().record_spill_wait( - u64::try_from(waited_at.elapsed().as_nanos()).unwrap_or(u64::MAX), - ); timer.record_spill_growth(&prev.chunk_path, prev.size_before); // A chunk that extended an existing run is already represented in // `chunk_files`; only a chunk that started a run adds a merge source. @@ -3214,7 +2933,7 @@ impl RawExternalSorter { ); let (record_source, header) = { let (reader, header) = - create_raw_bam_reader_pool_integrated(input, &pool, self.read_streams)?; + create_raw_bam_reader_pool_integrated(input, &pool, self.async_reader)?; (RecordSource::direct(reader), header) }; @@ -3292,14 +3011,12 @@ impl RawExternalSorter { // One worker pool spans both phases, so size it to the wider of the two // and cap the active count per phase via `set_active_workers`. - let mut pool = SortWorkerPool::new( + let pool = SortWorkerPool::new( self.phase1_threads().max(self.phase2_threads()), self.temp_compression, self.output_compression, self.spill_codec, - self.sort_stats, ); - pool.read_streams = self.read_streams; // Phase 1 runs first; the merge raises this to `phase2_threads()`. pool.set_active_workers(self.phase1_threads()); Ok(Arc::new(pool)) @@ -3547,7 +3264,6 @@ impl RawExternalSorter { )?; let mut records_merged = 0u64; - let mut merge_progress_batch = crate::progress_batch::BatchedProgress::new(); let merge_progress = ProgressTracker::new("Merged records").with_interval(1_000_000); // Set when an input is found not to be monotonic in the merge order: // (source index into `inputs`, 1-based record number within that input). @@ -3562,7 +3278,7 @@ impl RawExternalSorter { writer.write_raw_record(&records[winner])?; records_merged += 1; - merge_progress_batch.tick(&merge_progress); + merge_progress.log_if_needed(1); let reader_idx = source_map[winner]; if let Some(raw_record) = readers[reader_idx].next() { @@ -3608,9 +3324,6 @@ impl RawExternalSorter { writer.finish()?; out_target.persist()?; - // Must precede log_final: the tracker has not seen the last partial batch, - // and without this the reported total comes up short by up to one batch. - merge_progress_batch.flush(&merge_progress); merge_progress.log_final(); Ok(records_merged) @@ -3648,7 +3361,7 @@ impl RawExternalSorter { use crate::keys::RawCoordinateKey; let mut stats = RawSortStats::default(); - let mut timer = SortPhaseTimer::new(self.sort_stats); + let mut timer = SortPhaseTimer::new(); // Get number of references (unmapped reads map to nref) let nref = header.reference_sequences().len() as u32; @@ -3671,13 +3384,6 @@ impl RawExternalSorter { let rayon_pool = self.build_sort_rayon_pool()?; let progress = ProgressTracker::new("Read records").with_interval(1_000_000); - // Count records locally and forward in batches: `log_if_needed` does a - // relaxed `fetch_add` per call, which on aarch64 is an outline-atomics - // call into `__aarch64_ldadd8_relax`. On the ingest thread -- the serial - // thread that sets 60% of a spill-heavy sort's wall clock -- that helper - // measured 3.4% of the profile. The merge loops already batch for the - // same reason. - let mut progress_batch = crate::progress_batch::BatchedProgress::new(); debug!("Phase 1: Reading and sorting chunks (inline buffer, keyed output)..."); let mut probe = SpillProbe::new("phase1"); @@ -3687,7 +3393,7 @@ impl RawExternalSorter { // call, which is fine — `push_coordinate` copies the bytes into the buffer). while let Some(record) = record_source.next_record_borrowed()? { stats.total_records += 1; - progress_batch.tick(&progress); + progress.log_if_needed(1); // Push directly to buffer - key extracted inline from raw bytes buffer.push_coordinate(record)?; @@ -3740,9 +3446,7 @@ impl RawExternalSorter { }, )?); - // Keep the arena's pages mapped for the next chunk; `clear` would - // hand them back and cost a minor fault per page on the refill. - buffer.reset_for_reuse(); + buffer.clear(); force_mi_collect(); probe.post_spill(Some(pool.phase1_queue_depths())); timer.begin_read_span(); @@ -3750,7 +3454,6 @@ impl RawExternalSorter { } timer.end_read_span(); - progress_batch.flush(&progress); progress.log_final(); if let Some(err) = record_source.take_error() { return Err(anyhow::Error::from(err)); @@ -3768,16 +3471,6 @@ impl RawExternalSorter { probe.phase1_end(buffer.memory_usage() as u64); // Ingest is done: everything from here is Phase 2 (merge/write). - // Snapshot Phase 1's floor inputs before the pool is handed to the merge: - // the counters describe the phase that has just ended, and reading them - // later would both borrow a moved value and risk folding in Phase 2 work. - let phase1_floor = Phase1FloorInputs { - input_busy_secs: pool.phase1_input_busy_secs(), - threads: self.phase1_threads(), - ingest: pool.phase1_ingest_stats().snapshot(), - reader: pool.phase1_reader_report(), - ..Phase1FloorInputs::default() - }; self.enter_output_phase(&pool); if chunk_files.is_empty() { @@ -3820,7 +3513,7 @@ impl RawExternalSorter { let memory_chunks = MemorySources::Shared(memory_chunks); let n_memory = memory_chunks.num_non_empty(); - log_run_formation(timer.spill_count(), stats.runs_written, self.sort_stats); + log_run_formation(timer.spill_count(), stats.runs_written); debug!( "Phase 2: Merging {} chunks (keyed O(1) comparisons)...", chunk_files.len() + n_memory @@ -3842,12 +3535,7 @@ impl RawExternalSorter { if let Ok(pool) = Arc::try_unwrap(pool) { pool.shutdown(); } - timer.log_summary( - self.phase1_threads(), - self.phase2_threads(), - self.max_temp_files, - &phase1_floor, - ); + timer.log_summary(self.phase1_threads(), self.phase2_threads(), self.max_temp_files); debug!("Sort complete: {} records processed", stats.total_records); Ok(stats) @@ -3875,7 +3563,7 @@ impl RawExternalSorter { debug!("Indexing enabled: will write BAM index alongside output"); let mut stats = RawSortStats::default(); - let mut timer = SortPhaseTimer::new(self.sort_stats); + let mut timer = SortPhaseTimer::new(); let nref = header.reference_sequences().len() as u32; let init_cap = self.effective_initial_capacity(); @@ -3895,17 +3583,9 @@ impl RawExternalSorter { debug!("Phase 1: Reading and sorting chunks (inline buffer, keyed output)..."); let mut probe = SpillProbe::new("phase1"); - // Borrow each record's bytes straight out of the decompressed block, as the - // non-indexed sibling does. `RecordSource`'s `Iterator` impl yields an owned - // `RawRecord`, which is a heap allocation plus a full-record memcpy per - // record -- freed again as soon as `push_coordinate` has copied the bytes - // into the arena -- on the serial thread that sets 60% of a spill-heavy - // sort's wall clock. Nothing here needs the record to outlive the push: the - // BAI is built by the writer from BGZF offsets during the merge, not from - // this loop. - while let Some(record) = record_source.next_record_borrowed()? { + for record in record_source.by_ref() { stats.total_records += 1; - buffer.push_coordinate(record)?; + buffer.push_coordinate(record.as_ref())?; if probe.should_sample_read(stats.total_records) { probe.log_mid_read(probe_stats(&buffer), Some(pool.phase1_queue_depths())); @@ -3953,9 +3633,7 @@ impl RawExternalSorter { }, )?); - // Keep the arena's pages mapped for the next chunk; `clear` would - // hand them back and cost a minor fault per page on the refill. - buffer.reset_for_reuse(); + buffer.clear(); force_mi_collect(); probe.post_spill(Some(pool.phase1_queue_depths())); timer.begin_read_span(); @@ -3982,16 +3660,6 @@ impl RawExternalSorter { let output_header = self.create_output_header(header); // Ingest is done: everything from here is Phase 2 (merge/write). - // Snapshot Phase 1's floor inputs before the pool is handed to the merge: - // the counters describe the phase that has just ended, and reading them - // later would both borrow a moved value and risk folding in Phase 2 work. - let phase1_floor = Phase1FloorInputs { - input_busy_secs: pool.phase1_input_busy_secs(), - threads: self.phase1_threads(), - ingest: pool.phase1_ingest_stats().snapshot(), - reader: pool.phase1_reader_report(), - ..Phase1FloorInputs::default() - }; self.enter_output_phase(&pool); if chunk_files.is_empty() { @@ -4040,7 +3708,7 @@ impl RawExternalSorter { let memory_chunks = MemorySources::Shared(memory_chunks); let n_memory = memory_chunks.num_non_empty(); - log_run_formation(timer.spill_count(), stats.runs_written, self.sort_stats); + log_run_formation(timer.spill_count(), stats.runs_written); debug!( "Phase 2: Merging {} chunks with index generation...", chunk_files.len() + n_memory @@ -4066,12 +3734,7 @@ impl RawExternalSorter { if let Ok(pool) = Arc::try_unwrap(pool) { pool.shutdown(); } - timer.log_summary( - self.phase1_threads(), - self.phase2_threads(), - self.max_temp_files, - &phase1_floor, - ); + timer.log_summary(self.phase1_threads(), self.phase2_threads(), self.max_temp_files); debug!("Sort complete: {} records processed", stats.total_records); Ok(stats) @@ -4124,7 +3787,7 @@ impl RawExternalSorter { use crate::keys::SortContext; let mut stats = RawSortStats::default(); - let mut timer = SortPhaseTimer::new(self.sort_stats); + let mut timer = SortPhaseTimer::new(); let ctx = SortContext::from_header(header); @@ -4143,19 +3806,12 @@ impl RawExternalSorter { let rayon_pool = self.build_sort_rayon_pool()?; let progress = ProgressTracker::new("Read records").with_interval(1_000_000); - // Count records locally and forward in batches: `log_if_needed` does a - // relaxed `fetch_add` per call, which on aarch64 is an outline-atomics - // call into `__aarch64_ldadd8_relax`. On the ingest thread -- the serial - // thread that sets 60% of a spill-heavy sort's wall clock -- that helper - // measured 3.4% of the profile. The merge loops already batch for the - // same reason. - let mut progress_batch = crate::progress_batch::BatchedProgress::new(); debug!("Phase 1: Reading and sorting chunks (keyed output)..."); let mut probe = SpillProbe::new("phase1"); for record in record_source.by_ref() { stats.total_records += 1; - progress_batch.tick(&progress); + progress.log_if_needed(1); // Extract key from raw bytes. Stamp the ingest position within this // chunk so the key is totally ordered: read name + flags alone is not @@ -4232,7 +3888,6 @@ impl RawExternalSorter { } timer.end_read_span(); - progress_batch.flush(&progress); progress.log_final(); if let Some(err) = record_source.take_error() { return Err(anyhow::Error::from(err)); @@ -4250,16 +3905,6 @@ impl RawExternalSorter { probe.phase1_end(memory_used as u64); // Ingest is done: everything from here is Phase 2 (merge/write). - // Snapshot Phase 1's floor inputs before the pool is handed to the merge: - // the counters describe the phase that has just ended, and reading them - // later would both borrow a moved value and risk folding in Phase 2 work. - let phase1_floor = Phase1FloorInputs { - input_busy_secs: pool.phase1_input_busy_secs(), - threads: self.phase1_threads(), - ingest: pool.phase1_ingest_stats().snapshot(), - reader: pool.phase1_reader_report(), - ..Phase1FloorInputs::default() - }; self.enter_output_phase(&pool); if chunk_files.is_empty() { @@ -4342,7 +3987,7 @@ impl RawExternalSorter { let memory_chunks = MemorySources::Owned(keyed_chunks); let n_memory = memory_chunks.num_non_empty(); - log_run_formation(timer.spill_count(), stats.runs_written, self.sort_stats); + log_run_formation(timer.spill_count(), stats.runs_written); debug!( "Phase 2: Merging {} chunks (keyed comparisons)...", chunk_files.len() + n_memory @@ -4364,12 +4009,7 @@ impl RawExternalSorter { if let Ok(pool) = Arc::try_unwrap(pool) { pool.shutdown(); } - timer.log_summary( - self.phase1_threads(), - self.phase2_threads(), - self.max_temp_files, - &phase1_floor, - ); + timer.log_summary(self.phase1_threads(), self.phase2_threads(), self.max_temp_files); debug!("Sort complete: {} records processed", stats.total_records); Ok(stats) @@ -4494,7 +4134,7 @@ impl RawExternalSorter { cb_hasher: &ahash::RandomState, ) -> Result { let mut stats = RawSortStats::default(); - let mut timer = SortPhaseTimer::new(self.sort_stats); + let mut timer = SortPhaseTimer::new(); // The full key the decode-time verify compares dropped lanes against. // On empty input, `first_record` is `None` so the pre-loop push is @@ -4528,120 +4168,55 @@ impl RawExternalSorter { let rayon_pool = self.build_sort_rayon_pool()?; let progress = ProgressTracker::new("Read records").with_interval(1_000_000); - // Count records locally and forward in batches: `log_if_needed` does a - // relaxed `fetch_add` per call, which on aarch64 is an outline-atomics - // call into `__aarch64_ldadd8_relax`. On the ingest thread -- the serial - // thread that sets 60% of a spill-heavy sort's wall clock -- that helper - // measured 3.4% of the profile. The merge loops already batch for the - // same reason. - let mut progress_batch = crate::progress_batch::BatchedProgress::new(); debug!("Phase 1: Reading and sorting chunks (inline buffer)..."); let mut probe = SpillProbe::new("phase1"); - // Key extraction runs on the pool, not here: at 120 ns/record it was the - // single largest cost on this thread (93.6s of a 137.5s ingest against a - // 145.4s read span), and it has no ordering requirement. The ingest - // thread now pushes bytes and hands out batches; `deferred` cuts them, - // the pool runs them, and the keys are spliced back before any sort. - let mut deferred = crate::phase1_keys::DeferredKeys::::new( - Arc::new(crate::phase1_keys::KeyContext { - lib_lookup: lib_lookup.clone(), - cell_tag: self.cell_tag, - cb_hasher: cb_hasher.clone(), - first_key: first, - variant, - }), - self.phase1_threads(), - ); - - // Process the captured first record before draining the rest. Its key is - // deferred like every other record's, including the dropped-lane verify - // (trivially satisfied for the first record, since it *is* the baseline). - // A single record cannot exceed the memory limit, so no spill check is - // needed here. + // Process the captured first record before draining the rest: extract its + // full key, verify the dropped lanes match `first` (trivially true for the + // first record itself), and push the narrowed key. A single record cannot + // exceed the memory limit, so no spill check is needed here. if let Some(record) = first_record { stats.total_records += 1; - progress_batch.tick(&progress); - deferred.push(&mut buffer, &pool, record.as_ref())?; - } - - // Sub-phase timing for the ingest thread's serial CPU, which the floor - // line identifies as this phase's binding limit. Sampled 1-in-N and - // clock-corrected, per `crate::phase1_stats::IngestSample`; the segments - // are checked against the measured read span with a signed residual, so a - // partition that over-attributes says so instead of looking tidy. - let ingest_sample_interval = crate::phase1_stats::INGEST_SAMPLE_INTERVAL; - let clock_overhead_nanos = crate::merge_headroom::measure_clock_overhead_nanos(); - let mut ingest_raw = crate::phase1_stats::IngestSample::default(); - let mut ingest_samples: u64 = 0; - let mut sample_countdown: u64 = 0; + progress.log_if_needed(1); + + let bam_bytes = record.as_ref(); + let full = extract_template_key_inline(bam_bytes, lib_lookup, self.cell_tag, cb_hasher); + if let Some(violation) = verify_dropped_lanes(&first, &full, variant) { + let name = String::from_utf8_lossy( + fgumi_raw_bam::RawRecordView::new(bam_bytes).read_name(), + ) + .into_owned(); + return Err(dropped_lane_error(&name, violation)); + } + buffer.push(bam_bytes, K::from_full(&full))?; + } // Borrow each record's bytes in place (see the coordinate ingest loop); // the key is extracted and the bytes copied into the buffer before the // borrow ends, so no owned `RawRecord` is needed here. - loop { - // Decide sampling before the fetch, so every segment below is timed - // on the same records or on none. Timing a subset would bias the - // partition toward whichever step happened to be measured, and the - // partition's whole value is that its segments sum to the span. - let sample_this = sample_countdown == 0; - if sample_this { - sample_countdown = ingest_sample_interval - 1; - ingest_samples += 1; - } else { - sample_countdown -= 1; - } - - let t = sample_this.then(Instant::now); - let Some(bam_bytes) = record_source.next_record_borrowed()? else { break }; - if let Some(t0) = t { - ingest_raw.fetch += t0.elapsed().as_secs_f64(); - } - + while let Some(bam_bytes) = record_source.next_record_borrowed()? { stats.total_records += 1; - let t = sample_this.then(Instant::now); - progress_batch.tick(&progress); - if let Some(t0) = t { - ingest_raw.tick += t0.elapsed().as_secs_f64(); - } - - // Push the bytes and get a key onto them. With workers to spare the - // key is extracted on the pool and spliced in before anything sorts - // this buffer; with a single worker it is extracted right here, as - // it always was. `push` therefore covers extraction on the - // single-worker path and not on the deferred one -- read it against - // the deferred-extraction census below, not on its own. - let t = sample_this.then(Instant::now); - deferred.push(&mut buffer, &pool, bam_bytes)?; - if let Some(t0) = t { - ingest_raw.push += t0.elapsed().as_secs_f64(); + progress.log_if_needed(1); + + // Extract the full template key, verify the lanes the chosen variant + // dropped are constant relative to the first record, then push the + // narrowed key. + let full = extract_template_key_inline(bam_bytes, lib_lookup, self.cell_tag, cb_hasher); + if let Some(violation) = verify_dropped_lanes(&first, &full, variant) { + let name = String::from_utf8_lossy( + fgumi_raw_bam::RawRecordView::new(bam_bytes).read_name(), + ) + .into_owned(); + return Err(dropped_lane_error(&name, violation)); } + buffer.push(bam_bytes, K::from_full(&full))?; - let t = sample_this.then(Instant::now); - let should_probe = probe.should_sample_read(stats.total_records); - // Records whose keys are still in flight are charged here, so a - // lagging pool cannot let the buffer overrun the memory limit. - let over_limit = - buffer.memory_usage() + deferred.in_flight_bytes() >= self.memory_limit; - if let Some(t0) = t { - ingest_raw.probe += t0.elapsed().as_secs_f64(); - } - if should_probe { + if probe.should_sample_read(stats.total_records) { probe.log_mid_read(probe_stats(&buffer), Some(pool.phase1_queue_depths())); } // Check memory usage - if over_limit { - // Finish the chunk's keys *inside* the read span. This is ingest - // work for records already read, and accounting it to the spill - // region instead would let a pool that never keeps up still - // report a shorter read span -- the one way this change could - // look successful while doing nothing. - let violation = deferred.finish(&mut buffer, &pool)?; - if let Some(v) = violation { - return Err(dropped_lane_error(&v.name, v.violation)); - } - + if buffer.memory_usage() >= self.memory_limit { timer.end_read_span(); let bstats = probe_stats(&buffer); let depths = Some(pool.phase1_queue_depths()); @@ -4679,39 +4254,14 @@ impl RawExternalSorter { }, )?); - // Keep the arena's pages mapped for the next chunk; `clear` would - // hand them back and cost a minor fault per page on the refill. - buffer.reset_for_reuse(); - deferred.reset(); + buffer.clear(); force_mi_collect(); probe.post_spill(Some(pool.phase1_queue_depths())); timer.begin_read_span(); } } - // Drain the tail before closing the read span: the records still held in - // the live arena segment have no keys yet, and everything below this - // point either sorts or spills the buffer. - // - // Inside the span on purpose. This is ingest work for records already - // read, and the read span is the number this whole change is trying to - // move -- excluding the drain would let a pool that never kept up still - // report a shorter span, which is the one way this optimization could - // look successful while doing nothing. - let violation = deferred.finish(&mut buffer, &pool)?; - - // Ingest is over, so the arena segments held for the next chunk have no - // next chunk. Holding them through the merge costs peak RSS and buys - // nothing. - buffer.release_retained(); - timer.end_read_span(); - - if let Some(v) = violation { - return Err(dropped_lane_error(&v.name, v.violation)); - } - - progress_batch.flush(&progress); progress.log_final(); if let Some(err) = record_source.take_error() { return Err(anyhow::Error::from(err)); @@ -4729,20 +4279,6 @@ impl RawExternalSorter { probe.phase1_end(buffer.memory_usage() as u64); // Ingest is done: everything from here is Phase 2 (merge/write). - // Snapshot Phase 1's floor inputs before the pool is handed to the merge: - // the counters describe the phase that has just ended, and reading them - // later would both borrow a moved value and risk folding in Phase 2 work. - let phase1_floor = Phase1FloorInputs { - input_busy_secs: pool.phase1_input_busy_secs(), - threads: self.phase1_threads(), - ingest: pool.phase1_ingest_stats().snapshot(), - reader: pool.phase1_reader_report(), - key_overlap: deferred.overlap_census(), - sample: ingest_raw, - samples: ingest_samples, - records: stats.total_records, - clock_overhead_nanos, - }; self.enter_output_phase(&pool); if chunk_files.is_empty() { @@ -4784,7 +4320,7 @@ impl RawExternalSorter { let memory_chunks = MemorySources::Shared(memory_chunks); let n_memory = memory_chunks.num_non_empty(); - log_run_formation(timer.spill_count(), stats.runs_written, self.sort_stats); + log_run_formation(timer.spill_count(), stats.runs_written); debug!("Phase 2: Merging {} chunks...", chunk_files.len() + n_memory); // Merge using O(1) key comparisons @@ -4803,12 +4339,7 @@ impl RawExternalSorter { if let Ok(pool) = Arc::try_unwrap(pool) { pool.shutdown(); } - timer.log_summary( - self.phase1_threads(), - self.phase2_threads(), - self.max_temp_files, - &phase1_floor, - ); + timer.log_summary(self.phase1_threads(), self.phase2_threads(), self.max_temp_files); debug!("Sort complete: {} records processed", stats.total_records); Ok(stats) @@ -4819,7 +4350,7 @@ impl RawExternalSorter { /// Disk chunks become `PoolDisk` sources: the shared worker pool reads and /// decompresses them in the background while the main thread parses records, /// so no per-source threads are spawned. Both the plain - /// ([`merge_chunks_generic`]) and indexed ([`merge_chunks_with_index`]) + /// ([`Self::merge_chunks_generic`]) and indexed ([`Self::merge_chunks_with_index`]) /// merges go through this pool-integrated path. fn build_chunk_sources( chunk_files: &[PathBuf], @@ -4848,8 +4379,8 @@ impl RawExternalSorter { /// Build the pooled merge sources and activate Phase 2. /// - /// Shared by both the plain ([`merge_chunks_generic`]) and indexed - /// ([`merge_chunks_with_index`]) merges so the Phase-2 lifecycle is + /// Shared by both the plain ([`Self::merge_chunks_generic`]) and indexed + /// ([`Self::merge_chunks_with_index`]) merges so the Phase-2 lifecycle is /// single-sourced and the two paths cannot drift. Returns the sources plus /// an RAII [`Phase2Guard`] (borrowing `pool`); callers finish through /// [`Phase2Guard::finish_output`], and `Drop` resets Phase 2 on any path @@ -4918,20 +4449,18 @@ impl RawExternalSorter { /// diagnosis -- someone raising `--threads` on an I/O-bound sort -- is worse /// than the cost of a vague right one. The numbers above the verdict are the /// evidence; this line only points at them. - fn log_merge_verdict(utilization: Option, fetch_fraction: f64, sort_stats: bool) { + fn log_merge_verdict(utilization: Option, fetch_fraction: f64) { use crate::merge_phases::{MergeVerdict, classify_merge}; let Some(utilization) = utilization else { return }; let (util_pct, fetch_pct) = (100.0 * utilization, 100.0 * fetch_fraction); match classify_merge(utilization, fetch_fraction) { - MergeVerdict::CpuBound => stat!( - sort_stats, + MergeVerdict::CpuBound => info!( " Verdict: worker pool saturated ({util_pct:.0}%); the merge is CPU-bound, so \ more threads may help" ), - MergeVerdict::IoBound => stat!( - sort_stats, + MergeVerdict::IoBound => info!( " Verdict: workers idle ({util_pct:.0}% of capacity) while the consumer spent \ {fetch_pct:.0}% of its loop waiting for data -- neither side is the constraint. \ More threads are unlikely to help. Storage is one candidate; the pipeline's own \ @@ -4939,8 +4468,7 @@ impl RawExternalSorter { workers notice new work -- is another, and the Merge Stalls block below is what \ separates them." ), - MergeVerdict::Mixed => stat!( - sort_stats, + MergeVerdict::Mixed => info!( " Verdict: worker pool {util_pct:.0}% utilized, consumer {fetch_pct:.0}% \ waiting on data -- neither clearly saturated" ), @@ -4976,19 +4504,11 @@ impl RawExternalSorter { /// exactly. The breakdown then partitions the loop by construction rather /// than overshooting it by 30-50%. fn log_consumer_cpu( - loop_total: f64, - scaled: crate::merge_headroom::ConsumerSample, + sampled: (f64, f64, f64, f64), records: (u64, f64), stalls: Option<&crate::merge_stalls::ConsumerStallReport>, - sort_stats: bool, ) { - let crate::merge_headroom::ConsumerSample { - publish: est_publish, - present: est_present, - write: est_write, - advance: est_read, - tree: est_tree, - } = scaled; + let (loop_total, est_read, est_tree, est_write) = sampled; let (records_merged, blocked_secs) = records; if records_merged == 0 { return; @@ -5003,9 +4523,8 @@ impl RawExternalSorter { // waits would be counted on both sides and inflate their shares. let fetch_cpu = (est_read - park_secs).max(0.0); let write_cpu = (est_write - blocked_secs).max(0.0); - let ratio_total = est_publish + est_present + fetch_cpu + est_tree + write_cpu; - stat!( - sort_stats, + let ratio_total = fetch_cpu + est_tree + write_cpu; + info!( " Consumer CPU: {consumer_cpu:.1}s exact (loop {loop_total:.1}s - park \ {park_secs:.1}s - backpressure {blocked_secs:.1}s) = {:.0} ns/record", consumer_cpu * 1e9 / records @@ -5015,55 +4534,26 @@ impl RawExternalSorter { } let share = |part: f64| consumer_cpu * part / ratio_total; let ns = |part: f64| share(part) * 1e9 / records; - stat!( - sort_stats, - " fetch + decompress: {:.1}s ({:.0} ns/rec)", - share(fetch_cpu), - ns(fetch_cpu) - ); - stat!( - sort_stats, - " present record: {:.1}s ({:.0} ns/rec)", - share(est_present), - ns(est_present) - ); - stat!( - sort_stats, - " loser tree: {:.1}s ({:.0} ns/rec)", - share(est_tree), - ns(est_tree) - ); - stat!( - sort_stats, - " enqueue write: {:.1}s ({:.0} ns/rec)", - share(write_cpu), - ns(write_cpu) - ); - stat!( - sort_stats, - " next-source predict: {:.1}s ({:.0} ns/rec)", - share(est_publish), - ns(est_publish) - ); - stat!( - sort_stats, + info!(" parse + key extract: {:.1}s ({:.0} ns/rec)", share(fetch_cpu), ns(fetch_cpu)); + info!(" loser tree: {:.1}s ({:.0} ns/rec)", share(est_tree), ns(est_tree)); + info!(" enqueue write: {:.1}s ({:.0} ns/rec)", share(write_cpu), ns(write_cpu)); + info!( " (totals exact; the split between them is sampled, so read the shares as \ proportions rather than to the tenth of a second)" ); } - #[allow(clippy::too_many_lines, reason = "a single cohesive merge-phase diagnostic block")] fn log_merge_sub_phases( walls: (f64, f64), - consumer: crate::merge_headroom::ConsumerSample, + consumer: (f64, f64, f64), sampling: (u64, u64), active_workers: usize, pool: &Arc, stalls: Option, consumer_diag: MergeConsumerDiag, ) { - let sort_stats = pool.sort_stats(); let (loop_total, merge_total) = walls; + let (write_secs, read_secs, tree_secs) = consumer; let (samples_taken, records_merged) = sampling; if records_merged == 0 { return; @@ -5071,27 +4561,44 @@ impl RawExternalSorter { #[allow(clippy::cast_precision_loss, reason = "sample counts stay far below 2^52")] let scale = if samples_taken > 0 { records_merged as f64 / samples_taken as f64 } else { 1.0 }; - let scaled = consumer.scaled(scale); - let est_read = scaled.advance; - let park_secs = stalls.as_ref().map_or(0.0, |x| x.park_secs); - - stat!(sort_stats, "=== Merge Sub-Phase Timing ==="); - let MergeConsumerDiag { backpressure_secs: blocked_secs, .. } = consumer_diag; - Self::log_consumer_rows( - scaled, - (samples_taken, records_merged), - scale, - loop_total, - park_secs, - consumer_diag, + let (est_write, est_read, est_tree) = + (write_secs * scale, read_secs * scale, tree_secs * scale); + + info!("=== Merge Sub-Phase Timing ==="); + info!( + " Consumer (main thread; {samples_taken} samples of {records_merged} records, scaled {scale:.0}x)" ); + info!(" Fetch next record: {est_read:.1}s (includes waiting on decompressed blocks)"); + let MergeConsumerDiag { + backpressure_secs: blocked_secs, + backpressure_waits: blocked_waits, + borrowed, + reassembled, + } = consumer_diag; + if blocked_waits > 0 { + #[allow(clippy::cast_precision_loss, reason = "wait counts stay far below 2^52")] + let mean_us = blocked_secs * 1e6 / blocked_waits as f64; + info!( + " Output backpressure: {blocked_secs:.1}s over {blocked_waits} waits \ + (mean {mean_us:.0} us, exact) -- the consumer blocked for an output permit" + ); + } else { + info!(" Output backpressure: none -- the compressors always had a permit ready"); + } + if borrowed + reassembled > 0 { + #[allow(clippy::cast_precision_loss, reason = "record counts stay far below 2^52")] + let pct = 100.0 * reassembled as f64 / (borrowed + reassembled) as f64; + info!( + " Record presentation: {borrowed} borrowed zero-copy, {reassembled} \ + reassembled across a block boundary ({pct:.2}%, exact)" + ); + } + info!(" Loser tree: {est_tree:.1}s"); Self::log_consumer_cpu( - loop_total, - scaled, + (loop_total, est_read, est_tree, est_write), (records_merged, blocked_secs), stalls.as_ref(), - sort_stats, ); // Parking is a subset of "fetch next record", so exact park time cannot @@ -5103,8 +4610,7 @@ impl RawExternalSorter { if let Some(s) = stalls && s.park_secs > est_read * 1.05 { - stat!( - sort_stats, + info!( " NOTE: exact park time is {:.1}s, above the sampled fetch estimate of \ {est_read:.1}s -- the sample under-caught stalls; prefer the Merge Stalls block", s.park_secs @@ -5120,15 +4626,10 @@ impl RawExternalSorter { let (comp_s, comp_n) = workers.output_compress; let (spill_s, spill_n) = workers.spill_compress; let workers_n = active_workers; - stat!( - sort_stats, + info!( " Workers ({workers_n} active threads; busy time, overlaps the above and itself)" ); - stat!( - sort_stats, - " Spill disk read: {read_s:.1}s ({:.0}%) [{read_n} batches]", - pct(read_s) - ); + info!(" Spill disk read: {read_s:.1}s ({:.0}%) [{read_n} batches]", pct(read_s)); // Blocks-per-batch, split by allowance. A deep share near zero means // the deep-read gate (drain frontier or awaited source) is not // firing, which reads identically to "the deeper read-ahead did not @@ -5136,33 +4637,20 @@ impl RawExternalSorter { let (deep_b, deep_blk, shal_b, shal_blk) = pool.read_batch_split(); #[allow(clippy::cast_precision_loss, reason = "block counts stay far below 2^52")] let per = |blk: u64, b: u64| if b == 0 { 0.0 } else { blk as f64 / b as f64 }; - stat!( - sort_stats, + info!( " deep {deep_b} batches / {deep_blk} blocks ({:.1} per batch), \ other {shal_b} batches / {shal_blk} blocks ({:.1} per batch)", per(deep_blk, deep_b), per(shal_blk, shal_b) ); - stat!( - sort_stats, - " Spill decompress: {dec_s:.1}s ({:.0}%) [{dec_n} blocks]", - pct(dec_s) - ); - stat!( - sort_stats, - " Output compress: {comp_s:.1}s ({:.0}%) [{comp_n} blocks]", - pct(comp_s) - ); - stat!( - sort_stats, - " Total worker busy: {busy:.1}s (NOT comparable to loop wall clock)" - ); + info!(" Spill decompress: {dec_s:.1}s ({:.0}%) [{dec_n} blocks]", pct(dec_s)); + info!(" Output compress: {comp_s:.1}s ({:.0}%) [{comp_n} blocks]", pct(comp_s)); + info!(" Total worker busy: {busy:.1}s (NOT comparable to loop wall clock)"); // Utilization is the thread-efficiency question: well below 100% // means the pool idled, the merge was bound by something other than // worker CPU, and adding compression threads cannot help. if let Some(util) = workers.worker_utilization(merge_total, workers_n) { - stat!( - sort_stats, + info!( " Worker utilization: {:.0}% of {workers_n} active threads x \ {merge_total:.1}s", 100.0 * util @@ -5170,8 +4658,8 @@ impl RawExternalSorter { } // Phase 1 spill compression rides the same worker step, so it is // reported for context but excluded from the merge totals above. - stat!(sort_stats, " Phase 1 (not part of the merge)"); - stat!(sort_stats, " Spill compress: {spill_s:.1}s [{spill_n} blocks]"); + info!(" Phase 1 (not part of the merge)"); + info!(" Spill compress: {spill_s:.1}s [{spill_n} blocks]"); // Utilization over the full merge window; the fetch fraction stays // relative to the consumer loop, which is the only thing it is a @@ -5179,126 +4667,10 @@ impl RawExternalSorter { Self::log_merge_verdict( workers.worker_utilization(merge_total, workers_n), if loop_total > 0.0 { est_read / loop_total } else { 0.0 }, - sort_stats, - ); - Self::log_merge_headroom( - loop_total, - busy, - workers_n, - park_secs, - blocked_secs, - sort_stats, ); } Self::log_merge_stalls(loop_total, merge_total, active_workers, stalls, pool); - stat!(sort_stats, "=============================="); - } - - /// The consumer's sampled sub-phase rows, and their reconciliation against the - /// loop they partition. - /// - /// Extracted from [`Self::log_merge_sub_phases`] to keep that function inside - /// the line limit; it is one cohesive block of reporting rather than a - /// mechanical split. - fn log_consumer_rows( - scaled: crate::merge_headroom::ConsumerSample, - sampling: (u64, u64), - scale: f64, - loop_total: f64, - park_secs: f64, - consumer_diag: MergeConsumerDiag, - ) { - let (samples_taken, records_merged) = sampling; - let est_tree = scaled.tree; - let est_read = scaled.advance; - debug!( - " Consumer (main thread; {samples_taken} samples of {records_merged} records, scaled {scale:.0}x)" - ); - debug!(" Fetch next record: {est_read:.1}s (includes waiting on decompressed blocks)"); - let MergeConsumerDiag { - backpressure_secs: blocked_secs, - backpressure_waits: blocked_waits, - borrowed, - reassembled, - } = consumer_diag; - if blocked_waits > 0 { - #[allow(clippy::cast_precision_loss, reason = "wait counts stay far below 2^52")] - let mean_us = blocked_secs * 1e6 / blocked_waits as f64; - debug!( - " Output backpressure: {blocked_secs:.1}s over {blocked_waits} waits \ - (mean {mean_us:.0} us, exact) -- the consumer blocked for an output permit" - ); - } else { - debug!(" Output backpressure: none -- the compressors always had a permit ready"); - } - if borrowed + reassembled > 0 { - #[allow(clippy::cast_precision_loss, reason = "record counts stay far below 2^52")] - let pct = 100.0 * reassembled as f64 / (borrowed + reassembled) as f64; - debug!( - " Record presentation: {borrowed} borrowed zero-copy, {reassembled} \ - reassembled across a block boundary ({pct:.2}%, exact)" - ); - } - debug!(" Loser tree: {est_tree:.1}s"); - // Reconcile the sampled segments against the loop they are meant to - // partition. Three of these rows were reported for a long time without - // ever being summed against the loop, so there was no way to tell whether - // they accounted for most of it or a third. A signed residual is the check: - // positive means real time is unaccounted for, negative means the sampled - // regions over-attribute (their own clock overhead, or a sample biased - // toward expensive records). - let partition = crate::merge_headroom::LoopPartition { - segments: scaled, - loop_secs: loop_total, - park_secs, - }; - debug!( - " Sampled segments sum to {:.1}s of a {loop_total:.1}s loop \ - ({:+.0}% unattributed); {:.1}s of the fetch bucket was the consumer \ - working rather than waiting", - scaled.total(), - 100.0 * partition.unattributed_share(), - partition.advance_work_secs() - ); - } - - /// Name the wall this merge is against, and what is recoverable without - /// doing less work. - /// - /// Three limits, three different fixes, and they are routinely confused: the - /// serial consumer's own CPU, worker capacity, and coordination. Measured on - /// one cell, the same build at 8 and 16 threads gave opposite advice -- 1.8% - /// recoverable against 28% -- and nothing in the wall clock distinguished them. - /// Printing the floor turns "is my sort slow?" into "which limit am I on?", - /// which is the only version of the question that has an action attached. - fn log_merge_headroom( - loop_total: f64, - worker_busy_secs: f64, - threads: usize, - park_secs: f64, - blocked_secs: f64, - sort_stats: bool, - ) { - let floors = crate::merge_headroom::MergeFloors { - loop_secs: loop_total, - consumer_secs: (loop_total - park_secs - blocked_secs).max(0.0), - worker_busy_secs, - threads, - }; - stat!(sort_stats, " Merge floor: {} is the limit", floors.binding().label()); - stat!( - sort_stats, - " consumer serial {:.1}s | worker capacity {:.1}s ({threads} threads) | \ - loop {loop_total:.1}s", - floors.consumer_secs, - floors.worker_floor_secs() - ); - stat!( - sort_stats, - " recoverable without doing less work: {:.1}s ({:.0}% of the merge)", - floors.recoverable_secs(), - 100.0 * floors.recoverable_share() - ); + info!("=============================="); } /// Log why the merge stalled, as opposed to where its time went. @@ -5324,7 +4696,6 @@ impl RawExternalSorter { pool: &Arc, ) { use crate::merge_stalls::{Phase2Skip, ScanVerdict}; - let sort_stats = pool.sort_stats(); // Utilization gates the stall shape: "no worker on it" means a // scheduling defect on an idle pool and plain saturation on a busy one. @@ -5340,10 +4711,10 @@ impl RawExternalSorter { let wake = pool.wake_latency_report(); let stalls = stalls.filter(|s| !s.is_empty()); if !merge_stalls_are_silent(stalls.as_ref(), &scans, &wake) { - debug!("=== Merge Stalls ==="); + info!("=== Merge Stalls ==="); if let Some(s) = stalls { - Self::log_consumer_stalls(loop_total, utilization, s, sort_stats); + Self::log_consumer_stalls(loop_total, utilization, s); } if !scans.is_empty() { @@ -5352,8 +4723,8 @@ impl RawExternalSorter { .map(|&r| format!("{}={}", r.label(), scans.skips[r as usize])) .collect::>() .join(" "); - debug!(" Worker scans finding no work: {} ({reasons})", scans.scans); - debug!( + info!(" Worker scans finding no work: {} ({reasons})", scans.scans); + info!( " Of those: {:.0}% backpressured, {:.0}% waiting on a peer's read, {:.0}% \ contended", 100.0 * scans.verdict_share(ScanVerdict::Backpressured), @@ -5363,7 +4734,7 @@ impl RawExternalSorter { } if !wake.is_empty() { - debug!( + info!( " Worker discovery lag: ~{:.1}s of {:.1}s deep-sleep worker-seconds; {} \ sleeps ended in a find, {:.0}% of them after >=320us", wake.estimated_discovery_lag_secs(), @@ -5371,70 +4742,17 @@ impl RawExternalSorter { wake.productive_sleep_count(), 100.0 * wake.deep_sleep_wake_share() ); - debug!( + info!( " (the consumer unparks one worker when a reorder buffer drains, so this \ bounds how late work arriving any other way is noticed; it delays the merge \ only when every worker is asleep at once)" ); - // Why the discovery lag above is as large as it is. A wake aimed at a - // worker that is already running does nothing, and the consumer then - // waits for some other worker's backoff to expire -- bounded by - // MAX_BACKOFF_US, not by unpark cost. `recoverable` is the subset where a - // parked worker existed and the rotating target passed it over; the rest - // is a genuinely saturated pool and is nobody's fault. - let wakes = pool.wake_accounting(); - if wakes.issued > 0 { - debug!( - " Wake targeting: {} of {} wakes hit an already-running worker ({:.0}%), {} of \ - those had a parked worker available ({:.0}% recoverable)", - wakes.on_running, - wakes.issued, - Self::percent(wakes.on_running, wakes.issued), - wakes.recoverable, - Self::percent(wakes.recoverable, wakes.issued) - ); - } - debug!( - " Wakes issued: {} (backoff 10us doubling to {}us ceiling; one worker woken \ - per wake)", - pool.wakes_issued(), - crate::worker_pool::MAX_BACKOFF_US - ); } - // What the pool looked like when the consumer parked, which is the - // question "why did nobody start on this block yet" reduced to three - // mutually exclusive answers with three different fixes. Splitting park - // *time* as well as park counts matters: a rare-but-long class is - // invisible in counts alone. Outside the wake gate above so it stands - // on its own park count: a heavy-parking consumer whose wake report is - // empty would otherwise collect this census and never print it. - let supply = pool.park_supply_report(); - if supply.total_parks() > 0 { - const LABELS: [&str; crate::merge_stalls::ParkSupply::COUNT] = - ["a worker was asleep", "all busy, compress queued", "all busy merging"]; - stat!( - sort_stats, - " Why nobody had started on the awaited block, at the moment of the park:" - ); - for (i, label) in LABELS.iter().enumerate() { - stat!( - sort_stats, - " {label:<26} {:>10} parks ({:>4.1}%) {:>7.1}s ({:>4.1}% of park time)", - supply.counts[i], - Self::percent(supply.counts[i], supply.total_parks()), - Self::secs(supply.nanos[i]), - Self::percent(supply.nanos[i], supply.total_nanos()) - ); - } - } - - Self::log_park_attribution(pool, loop_total); - // Close this block before delegating: `log_block_lifecycle` opens and // closes its own, so without a terminator here the lifecycle block // reads as nested inside the stall block rather than following it. - debug!("===================="); + info!("===================="); } // Outside the gate above, not inside it. The lifecycle reports record the @@ -5442,253 +4760,6 @@ impl RawExternalSorter { // never stalled is precisely the run with a complete lifecycle trace and // nothing to say above. `log_block_lifecycle` carries its own gate. Self::log_block_lifecycle(pool); - - // Also outside the gate: these two counters say whether the targeted-depth - // paths ran at all, which matters most on the runs that did not stall. - Self::log_merge_validity_gates(pool); - } - - /// Nanoseconds as seconds. - #[expect(clippy::cast_precision_loss, reason = "nanosecond totals stay below 2^52")] - fn secs(nanos: u64) -> f64 { - nanos as f64 / 1e9 - } - - /// `part` as a percentage of `whole`, or 0.0 when `whole` is zero. - #[expect(clippy::cast_precision_loss, reason = "counts stay far below 2^52")] - fn percent(part: u64, whole: u64) -> f64 { - if whole == 0 { - return 0.0; - } - 100.0 * part as f64 / whole as f64 - } - - /// Validity gates for the two targeted-depth paths in the merge. - /// - /// Both changes are inert in a way that a wall-clock number cannot reveal: the - /// consumer may never find an unclaimed block to decompress, and the loser tree - /// may never name a next source. Either would land as "no measurable change" - /// rather than as "the mechanism did not run", so each gets a counter and each - /// counter is printed. - fn log_merge_validity_gates(pool: &SortWorkerPool) { - let sort_stats = pool.sort_stats(); - let self_served = pool.consumer_self_served(); - if self_served > 0 { - stat!( - sort_stats, - " Consumer served itself: {self_served} parks avoided by decompressing an \ - already-read block inline" - ); - } - let predictions = pool.phase2_predictions(); - if predictions > 0 { - stat!(sort_stats, " Next-source predictions published: {predictions}"); - } - } - - /// `total` divided over `claims`, or 0.0 when there are none. - #[expect(clippy::cast_precision_loss, reason = "counts stay far below 2^52")] - fn per_claim(total: u64, claims: u64) -> f64 { - if claims == 0 { - return 0.0; - } - total as f64 / claims as f64 - } - - /// Log a latency distribution for every stage a block passes through. - /// - /// [`crate::merge_phases::MergePhaseBreakdown`] already gives each stage a - /// busy total and a block count, which yields a mean. A mean cannot settle - /// an argument: output compression averages ~187 us over five million blocks - /// and a uniform 187 us behaves nothing like a bimodal mix with a long tail, - /// yet only the tail explains a worker being unavailable when the consumer - /// needs one. So report count, total, mean and three percentiles for each, - /// including the writer -- previously visible only as the consumer's - /// backpressure wait, which jumped 0.0s to 29.9s across one sweep with no - /// way to attribute it. - /// - /// `wasted visits` is the other half: files a worker passed over before one - /// gave it work. The scan tally is published only when a scan finds nothing, - /// so the walk to a *successful* claim was never counted -- and on an 89-way - /// merge that walk is most of what an "idle" worker is doing. - /// - /// The distributions log at **debug** and the wasted-visit line at **info**, - /// the same split [`Self::log_block_lifecycle`] uses: a percentile table is - /// for an investigation, not for someone who just sorted a BAM. Collection - /// is unconditional either way. - fn log_stage_latency( - pool: &SortWorkerPool, - writer: &( - crate::merge_trace::HistogramReport, - crate::merge_trace::HistogramReport, - crate::merge_trace::HistogramReport, - ), - ) { - let sort_stats = pool.sort_stats(); - let stage = pool.stage_latency(); - let &(write_dur, reorder_wait, reorder_depth) = writer; - let rows: [(&str, crate::merge_trace::HistogramReport); 6] = [ - ("read (batched)", stage.read.snapshot()), - ("decompress spill", stage.decompress.snapshot()), - ("compress output", stage.output_compress.snapshot()), - ("compress spill (ph1)", stage.spill_compress.snapshot()), - ("write block", write_dur), - ("write reorder wait", reorder_wait), - ]; - if rows.iter().all(|(_, r)| r.is_empty()) { - return; - } - // Distributions go to debug and decision-changing numbers stay at info, - // matching `log_block_lifecycle` above: a six-row percentile table is - // what an investigation needs and far more than someone who just sorted - // a BAM should have to scroll past. Collection stays unconditional -- - // gating collection is what made these questions unanswerable from logs - // already in hand. - debug!("=== Stage Latency ==="); - debug!( - " {:<22} {:>10} {:>9} {:>9} {:>9} {:>9} {:>9}", - "stage", "count", "total", "mean", "p50", "p90", "p99" - ); - for (label, r) in rows { - if r.is_empty() { - continue; - } - debug!( - " {label:<22} {:>10} {:>8.1}s {:>8.0}us {:>8}us {:>8}us {:>8}us", - r.count, - r.total_secs(), - r.mean_micros(), - r.percentile_micros(0.50), - r.percentile_micros(0.90), - r.percentile_micros(0.99) - ); - } - if !reorder_depth.is_empty() { - debug!( - " Writer reorder depth: mean {:.1} blocks, p90 {}, p99 {} (blocks held waiting \ - for an earlier serial)", - reorder_depth.mean_micros(), - reorder_depth.percentile_micros(0.90), - reorder_depth.percentile_micros(0.99) - ); - } - // Closes the debug block above, so the separator does not outlive its - // header when only info is enabled. - debug!("====================="); - // Stays at info: this one is a decision-changer, not a distribution. It - // is how much of an "idle" worker is scanning rather than waiting, and - // reading it wrong sends an investigation at the scan loop -- filling - // the walk with more work per claim was measured at +36% (see - // `PHASE2_DECOMP_CAP`). - let claims = stage.useful_claims.load(std::sync::atomic::Ordering::Relaxed); - if claims > 0 { - stat!( - sort_stats, - " Wasted file visits: {:.1} per claim over {} claims ({} visits produced \ - nothing) -- what an idle worker is actually doing", - stage.wasted_visits_per_claim(), - claims, - stage.wasted_visits.load(std::sync::atomic::Ordering::Relaxed) - ); - } - } - - /// Log where consumer park time went, and what each worker did. - /// - /// The park table is the one figure in this report that partitions an exact - /// total by construction, which is why it is the one to read first. Earlier - /// attempts to explain park were shares of park *events* (blind to a - /// rare-but-long cause) or sums over workers (which overcount, because the - /// consumer waits for whichever worker arrives first). `unattributed` is the - /// honesty term: a large value means the model is incomplete, not that the - /// merge is idle for no reason. - #[expect( - clippy::cast_precision_loss, - reason = "nanosecond and count totals are within f64's exact-integer range here" - )] - fn log_park_attribution(pool: &SortWorkerPool, loop_total: f64) { - let sort_stats = pool.sort_stats(); - let park = pool.park_attribution_report(); - if park.is_empty() { - return; - } - let total = park.total_nanos(); - if total == 0 { - return; - } - let secs = |ns: u64| ns as f64 / 1e9; - let share = |ns: u64| 100.0 * ns as f64 / total as f64; - let per_park = |ns: u64| ns as f64 / park.parks as f64 / 1e3; - - stat!(sort_stats, " Consumer park, by stage (exact, partitions the park):"); - stat!(sort_stats, " {:<28} {:>8} {:>7} {:>10}", "stage", "time", "share", "per park"); - for (label, ns) in [ - ("waiting for a worker", park.to_claim_nanos), - ("read + decompress work", park.work_nanos), - ("waiting for its own wake", park.to_resume_nanos), - ("unattributed", park.unattributed_nanos), - ] { - stat!( - sort_stats, - " {label:<28} {:>7.1}s {:>6.0}% {:>9.0}us", - secs(ns), - share(ns), - per_park(ns) - ); - } - stat!( - sort_stats, - " {:<28} {:>7.1}s {:>6.0}% {:>9.0}us over {} parks ({:.0}% of loop wall)", - "TOTAL", - secs(total), - 100.0, - per_park(total), - park.parks, - if loop_total > 0.0 { 100.0 * secs(total) / loop_total } else { 0.0 } - ); - stat!( - sort_stats, - " Blocks ready on the awaited file at resume: mean {:.2} (1.0 = every block \ - fetched on demand, one round trip per block)", - park.mean_ready_on_resume() - ); - stat!( - sort_stats, - " Parks with no claim during them: {} of {} ({:.0}%) -- the block was already \ - in flight or already done", - park.unclaimed_parks, - park.parks, - 100.0 * park.unclaimed_parks as f64 / park.parks as f64 - ); - - let threads = pool.per_thread_report(); - let claims_total: u64 = threads.iter().map(|&(_, _, c)| c).sum(); - if claims_total == 0 { - return; - } - stat!(sort_stats, " Per worker (merge + phase 1 combined):"); - stat!( - sort_stats, - " {:>3} {:>9} {:>9} {:>7} {:>10} {:>5}", - "wid", - "busy", - "idle", - "busy%", - "claims", - "share" - ); - for (w, &(busy, idle, claims)) in threads.iter().enumerate() { - let denom = busy + idle; - stat!( - sort_stats, - " {w:>3} {:>8.1}s {:>8.1}s {:>6.0}% {:>10} {:>4.0}%", - secs(busy), - secs(idle), - if denom > 0 { 100.0 * busy as f64 / denom as f64 } else { 0.0 }, - claims, - 100.0 * claims as f64 / claims_total as f64 - ); - } } /// Log every stage of a spill block's journey, and the refill cycle. @@ -5713,7 +4784,6 @@ impl RawExternalSorter { ) { use crate::merge_stalls::AwaitedState; use crate::merge_trace::MAX_TRACKED_IN_FLIGHT; - let sort_stats = pool.sort_stats(); // A merge that never stalled still has source runs to report, so this // block stands on its own park count rather than on the report being @@ -5730,31 +4800,6 @@ impl RawExternalSorter { debug!(" {:<14} {}", state.label(), hist.summary()); } } - // The same five states weighted by TIME, not by park count, and at info - // because the two disagree and only one of them is a cost. - // - // The park-supply census made this concrete: by count its largest class - // was "all busy, compress queued" at 41%, by time that class was 13% - // while "a worker was asleep" was 81%. Counts named the wrong fix, - // confidently. These histograms have always carried the nanoseconds -- - // they were simply reported one level down, so every arm collected the - // answer and none printed it. - let state_secs: [f64; AwaitedState::COUNT] = - std::array::from_fn(|i| consumer.park_by_state[i].total_secs()); - let state_total: f64 = state_secs.iter().sum(); - if state_total > 0.0 { - let share = |i: usize| 100.0 * state_secs[i] / state_total; - stat!( - sort_stats, - " The awaited file by park TIME ({state_total:.1}s): {:.0}% gap-filling, \ - {:.0}% gap-stalled, {:.0}% decompressing, {:.0}% raw-queued, {:.0}% starved", - share(AwaitedState::ReorderGapFilling as usize), - share(AwaitedState::ReorderGapStalled as usize), - share(AwaitedState::Decompressing as usize), - share(AwaitedState::RawQueued as usize), - share(AwaitedState::Starved as usize) - ); - } let parks_pct = |count: u64| { #[allow(clippy::cast_precision_loss, reason = "park counts stay far below 2^52")] @@ -5762,8 +4807,7 @@ impl RawExternalSorter { pct }; let mean_depth = consumer.mean_in_flight(); - stat!( - sort_stats, + info!( " Workers on the awaited file at a park: none {:.0}%, exactly one {:.0}%, \ two or more {:.0}% (mean {mean_depth:.1}, tracked to {MAX_TRACKED_IN_FLIGHT})", parks_pct(consumer.idle_file_parks()), @@ -5785,8 +4829,7 @@ impl RawExternalSorter { // batch produced it. let (mean_block_bytes, derived_cap, derived_batch) = pool.awaited_sizing(); if mean_block_bytes > 0 { - stat!( - sort_stats, + info!( " Hot-file refill sized from measured blocks: {mean_block_bytes} B/block \ -> batch {derived_batch}, cap {derived_cap}" ); @@ -5796,8 +4839,7 @@ impl RawExternalSorter { if skip_total > 0 { #[allow(clippy::cast_precision_loss, reason = "skip counts stay far below 2^52")] let pct = |n: u64| 100.0 * n as f64 / skip_total as f64; - stat!( - sort_stats, + info!( " Why the pool passed over the awaited file: raw-lock {:.0}%, raw-empty \ {:.0}%, decomp-lock {:.0}%, decomp-capped {:.0}% (of {skip_total})", pct(skips[0]), @@ -5807,15 +4849,13 @@ impl RawExternalSorter { ); } if mean_depth >= cap / 2.0 { - stat!( - sort_stats, + info!( " -> the pool is running near its tracked depth on the file the merge is \ blocked on, so these parks are the head block's decompress latency; more \ concurrency on that file cannot shorten them" ); } else if consumer.multi_worker_parks() > 0 { - stat!( - sort_stats, + info!( " -> the pool is on the awaited file but only {mean_depth:.1} deep of \ {MAX_TRACKED_IN_FLIGHT} tracked, so capacity is going unused on the file the \ merge is blocked on -- supply to that file, not decompress latency, is the \ @@ -5826,7 +4866,6 @@ impl RawExternalSorter { fn log_block_lifecycle(pool: &Arc) { use crate::merge_trace::EmptyCause; - let sort_stats = pool.sort_stats(); let life = pool.block_lifecycle_report(); let refill = pool.refill_report(); @@ -5842,7 +4881,7 @@ impl RawExternalSorter { // change a decision stay at info. Collection is unconditional either // way -- gating collection is what made the original question // unanswerable from logs we had already collected. - stat!(sort_stats, "=== Merge Block Lifecycle ==="); + info!("=== Merge Block Lifecycle ==="); debug!(" disk read -> {}", life.read_batch.summary()); debug!(" raw dwell -> {} (queued, waiting for a worker)", life.raw_dwell.summary()); debug!(" decompress -> {}", life.decompress.summary()); @@ -5850,8 +4889,7 @@ impl RawExternalSorter { " reorder -> {} (decompressed, waiting for the consumer)", life.reorder_dwell.summary() ); - stat!( - sort_stats, + info!( " Per block: {:.0}us in the raw FIFO unclaimed, {:.0}us decompressing, {:.0}us \ buffered before use (p50)", life.raw_dwell.percentile_micros(0.50), @@ -5859,8 +4897,7 @@ impl RawExternalSorter { life.reorder_dwell.percentile_micros(0.50) ); if life.reorder_is_pass_through() { - stat!( - sort_stats, + info!( " NOTE: blocks are consumed almost as fast as they are inserted, so the \ reorder buffer is a pass-through and PHASE2_DECOMP_CAP is not the binding \ constraint -- however full the other files look" @@ -5868,13 +4905,8 @@ impl RawExternalSorter { } if !refill.is_empty() { - stat!( - sort_stats, - " Refill cycle ({} times a file's buffer ran dry)", - refill.empties() - ); - stat!( - sort_stats, + info!(" Refill cycle ({} times a file's buffer ran dry)", refill.empties()); + info!( " At the moment it emptied: {:.0}% had raw blocks unclaimed, {:.0}% already \ decompressing, {:.0}% nothing at all", 100.0 * refill.cause_share(EmptyCause::RawReady), @@ -5886,8 +4918,7 @@ impl RawExternalSorter { if !refill.read_lag.is_empty() { debug!(" empty -> read {}", refill.read_lag.summary()); } - stat!( - sort_stats, + info!( " -> {:.0}% of refill latency is spent waiting for a worker to START, {:.0}% \ doing the work", 100.0 * refill.claim_share(), @@ -5902,14 +4933,12 @@ impl RawExternalSorter { // it and divides by zero to fill the line below. Self::log_awaited_file_depth(&consumer, pool); if !consumer.source_run_length.is_empty() { - stat!( - sort_stats, + info!( " Consecutive blocks per source: {}", consumer.source_run_length.summary_blocks() ); if consumer.source_run_length.percentile_micros(0.90) <= 1 { - stat!( - sort_stats, + info!( " -> the merge switches source almost every block, so there is no hot \ file to prioritise; demand is spread across all runs at once" ); @@ -5920,7 +4949,7 @@ impl RawExternalSorter { if !scans.is_empty() { debug!(" Fruitless worker scan cost: {}", scans.summary()); } - stat!(sort_stats, "============================="); + info!("============================="); } /// Log where the merge loop blocked and what the other files were doing. @@ -5933,20 +4962,17 @@ impl RawExternalSorter { loop_total: f64, utilization: f64, s: crate::merge_stalls::ConsumerStallReport, - sort_stats: bool, ) { use crate::merge_stalls::{StallShape, classify_stall}; let park_fraction = if loop_total > 0.0 { s.park_secs / loop_total } else { 0.0 }; - stat!( - sort_stats, + info!( " Consumer parked: {:.1}s ({:.0}% of loop wall, exact) over {} parks", s.park_secs, 100.0 * park_fraction, s.parks ); - stat!( - sort_stats, + info!( " Block pulls that had to wait: {}/{} ({:.0}%), {:.1} parks each (1.0 = no wasted \ wake-ups)", s.stalled_pulls, @@ -5954,8 +4980,7 @@ impl RawExternalSorter { 100.0 * s.stall_rate(), s.parks_per_stall() ); - stat!( - sort_stats, + info!( " Worst source: #{} at {:.1}s ({:.0}% of park time; {} sources parked on)", s.top_source, s.top_source_park_secs, @@ -5963,8 +4988,7 @@ impl RawExternalSorter { s.sources_parked_on ); if s.censuses > 0 { - stat!( - sort_stats, + info!( " Other files at a park ({} parks sampled): {:.0}% at cap, {:.0}% starved, \ {:.0}% unreadable", s.censuses, @@ -5972,8 +4996,7 @@ impl RawExternalSorter { 100.0 * s.starved_share, 100.0 * s.contended_share ); - stat!( - sort_stats, + info!( " The awaited file: {:.0}% gap-filling, {:.0}% gap-stalled, {:.0}% \ decompressing, {:.0}% raw-queued, {:.0}% starved", 100.0 * s.awaited.reorder_gap_filling, @@ -5982,8 +5005,7 @@ impl RawExternalSorter { 100.0 * s.awaited.raw_queued, 100.0 * s.awaited.starved ); - stat!( - sort_stats, + info!( " -> block not read yet {:.0}%, exists but unclaimed {:.0}%, being produced \ {:.0}%", 100.0 * s.awaited.starved, @@ -5992,40 +5014,33 @@ impl RawExternalSorter { ); } match classify_stall(park_fraction, utilization, s.contended_share, s.awaited) { - StallShape::NotStalled => { - stat!(sort_stats, " Shape: the consumer is not waiting on blocks"); - } - StallShape::PoolSaturated => stat!( - sort_stats, + StallShape::NotStalled => info!(" Shape: the consumer is not waiting on blocks"), + StallShape::PoolSaturated => info!( " Shape: pool saturated -- the consumer waits because every worker is busy, \ which is what a healthy CPU-bound merge looks like. Fewer bytes to compress or \ more threads would help; nothing here is misscheduled" ), - StallShape::HeadOfLine => stat!( - sort_stats, + StallShape::HeadOfLine => info!( " Shape: head-of-line -- the awaited file has nothing anywhere in its \ pipeline, so the block has not been read from disk yet. The constraint is \ upstream of the pool: storage, or read concurrency" ), - StallShape::WorkUnclaimed => stat!( - sort_stats, + StallShape::WorkUnclaimed => info!( " Shape: work unclaimed -- the block the consumer needs already exists and no \ worker is on it. Capacity is not the problem; scheduling and wake latency are. \ Compare the discovery-lag line below" ), - StallShape::DecompressLatency => stat!( - sort_stats, + StallShape::DecompressLatency => info!( " Shape: decompression latency -- a worker is already producing the needed \ block, so the consumer is paying the per-block cost serially. Check the reorder \ dwell below before reaching for a deeper cap: if blocks are consumed as fast as \ they are inserted, the buffer is not what the pipeline is running into" ), - StallShape::Contended => stat!( - sort_stats, + StallShape::Contended => info!( " Shape: lock contention -- a large share of file state could not be read \ without blocking, so the shares above understate what was available" ), - StallShape::Mixed => stat!(sort_stats, " Shape: no single candidate dominates"), + StallShape::Mixed => info!(" Shape: no single candidate dominates"), } } @@ -6046,7 +5061,6 @@ impl RawExternalSorter { ) -> Result { use crate::loser_tree::LoserTree; use crate::pooled_bam_writer::PooledBamWriter; - let sort_stats = self.sort_stats; let (mut sources, mut guard) = Self::setup_phase2_merge::(chunk_files, memory_chunks, pool)?; @@ -6080,7 +5094,6 @@ impl RawExternalSorter { let mut writer = PooledBamWriter::new(Arc::clone(pool), output, &output_header)?; let mut records_merged = 0u64; - let mut merge_progress_batch = crate::progress_batch::BatchedProgress::new(); let merge_progress = ProgressTracker::new("Merged records") .with_interval(1_000_000) .with_total(total_records); @@ -6101,13 +5114,6 @@ impl RawExternalSorter { // than the loop wall clock they are supposed to partition. A prime // interval decorrelates the sample from any block-size-derived period. let merge_sample_interval: u64 = 1021; - // Calibrated once per merge, not assumed: the clock pair that brackets each - // sampled segment costs 15-35ns depending on host, which is the same order - // as the segments themselves. Without subtracting it the rows are more clock - // than work -- measured at -70% residual before this existed. - let clock_overhead_nanos = crate::merge_headroom::measure_clock_overhead_nanos(); - let mut merge_publish_secs = 0.0f64; - let mut merge_present_secs = 0.0f64; let mut merge_write_secs = 0.0f64; let mut merge_read_secs = 0.0f64; let mut merge_tree_secs = 0.0f64; @@ -6131,59 +5137,31 @@ impl RawExternalSorter { // Snapshot the process-wide presentation counters so the sub-phase log // reports only this merge's delta, not totals accumulated by any prior // (sequential or concurrent) sort sharing the process. - let mut fetch_counts = RecordFetchCounts::default(); let borrowed_before = RECORD_BORROWED.load(std::sync::atomic::Ordering::Relaxed); let reassembled_before = RECORD_REASSEMBLED.load(std::sync::atomic::Ordering::Relaxed); let loop_start = Instant::now(); - let mut published_src: Option = None; while tree.winner_is_active() { - // Decide sampling first, so every segment below is timed on the same - // records or on none. Timing a subset would bias the partition toward - // whichever step happened to be measured, and the partition's whole - // value is that its segments sum to the loop. + let winner = tree.winner(); + let src_idx = source_map[winner]; let sample_this = sample_countdown == 0; if sample_this { sample_countdown = merge_sample_interval - 1; - samples_taken += 1; } else { sample_countdown -= 1; } - let t = sample_this.then(Instant::now); - let winner = tree.winner(); - let src_idx = source_map[winner]; - // Publish the next likely source on run change. That is not once per - // record but it is far from rare: 25,003,410 publications against - // 779,820,469 records on the measured cell, one per ~31 records, each - // calling `runner_up()` at ~11.6ns. The pool gives the named file the - // deep read allowance, so its first read starts while the consumer is - // still draining the current file instead of after it has stalled. A - // wrong prediction costs one deep read on a file the merge reaches - // eventually; there is no correctness component. - if published_src != Some(src_idx) { - published_src = Some(src_idx); - pool.set_phase2_next_source(tree.runner_up().map(|w| source_map[w])); - } - if let Some(t0) = t { - merge_publish_secs += t0.elapsed().as_secs_f64(); - } - - let t = sample_this.then(Instant::now); - let record_bytes = - winner_record_bytes(&sources[src_idx], guard.consumer_ref(), &mut fetch_counts)?; - if let Some(t0) = t { - merge_present_secs += t0.elapsed().as_secs_f64(); - } - - let t = sample_this.then(Instant::now); - writer.write_raw_record(record_bytes)?; - if let Some(t0) = t { + let record_bytes = winner_record_bytes(&sources[src_idx], guard.consumer_ref())?; + if sample_this { + let t0 = Instant::now(); + writer.write_raw_record(record_bytes)?; merge_write_secs += t0.elapsed().as_secs_f64(); + } else { + writer.write_raw_record(record_bytes)?; } records_merged += 1; - merge_progress_batch.tick(&merge_progress); + merge_progress.log_if_needed(1); if merge_probe.should_sample(records_merged) { let depths = pool.phase1_queue_depths(); @@ -6191,26 +5169,29 @@ impl RawExternalSorter { merge_probe.log_mid_with_depths(depths, consumer_stats); } - let t = sample_this.then(Instant::now); - let next = sources[src_idx].advance(guard.consumer_mut())?; - if let Some(t0) = t { + if sample_this { + let t0 = Instant::now(); + let next = sources[src_idx].advance(guard.consumer_mut())?; merge_read_secs += t0.elapsed().as_secs_f64(); - } - let t = sample_this.then(Instant::now); - if let Some(key) = next { - tree.replace_winner(key); - } else { - tree.remove_winner(); - } - if let Some(t0) = t { + let t0 = Instant::now(); + if let Some(key) = next { + tree.replace_winner(key); + } else { + tree.remove_winner(); + } merge_tree_secs += t0.elapsed().as_secs_f64(); + samples_taken += 1; + } else { + let next = sources[src_idx].advance(guard.consumer_mut())?; + if let Some(key) = next { + tree.replace_winner(key); + } else { + tree.remove_winner(); + } } } - // Before the delta below is read, or this merge reports zero. - fetch_counts.publish(); - let loop_total = loop_start.elapsed().as_secs_f64(); let borrowed_this_merge = RECORD_BORROWED.load(std::sync::atomic::Ordering::Relaxed) - borrowed_before; @@ -6233,13 +5214,6 @@ impl RawExternalSorter { // compressors are behind, which the sampled breakdown charged to "enqueue // write" -- a bucket documented as a handoff that excludes compression. let (write_blocked_secs, write_blocked_waits) = writer.write_backpressure(); - // Retain the permit pool so the writer histograms can be harvested *after* - // `finish` drains the output queue: the pool outlives the writer, and the - // block writes and reorder waits incurred during that drain belong in the - // "write block" and "write reorder wait" rows. A snapshot taken here, - // before the drain, would omit the tail -- the same reason the stall - // report below is harvested pre-finalize but the histograms are not. - let writer_permit_pool = writer.permit_pool(); // Harvest the consumer's stall report before finalizing: `finish_output` // releases the merge sources and with them the consumer, and the report @@ -6261,32 +5235,9 @@ impl RawExternalSorter { // before the drain. guard.finish_output(|| writer.finish())?; - // Harvest now, after the drain, from the retained pool. Empty only if the - // writer was already finalized when the pool was retained, which cannot - // happen on this path. - let writer_stats = writer_permit_pool - .as_deref() - .map_or_else(Default::default, crate::worker_pool::PermitPool::writer_stats); - - let raw_sample = crate::merge_headroom::ConsumerSample { - publish: merge_publish_secs, - present: merge_present_secs, - write: merge_write_secs, - advance: merge_read_secs, - tree: merge_tree_secs, - }; - let corrected_sample = raw_sample.corrected(samples_taken, clock_overhead_nanos); - stat!( - sort_stats, - " Consumer sampling: {samples_taken} samples, clock {clock_overhead_nanos}ns per \ - segment; sampled total {:.1}ms -> {:.1}ms after removing measurement overhead \ - (sampled basis, before scaling to the full merge)", - raw_sample.total() * 1e3, - corrected_sample.total() * 1e3 - ); Self::log_merge_sub_phases( (loop_total, loop_start.elapsed().as_secs_f64()), - corrected_sample, + (merge_write_secs, merge_read_secs, merge_tree_secs), (samples_taken, records_merged), active_workers, pool, @@ -6298,11 +5249,7 @@ impl RawExternalSorter { reassembled: reassembled_this_merge, }, ); - Self::log_stage_latency(pool, &writer_stats); - // Must precede log_final: the tracker has not seen the last partial batch, - // and without this the reported total comes up short by up to one batch. - merge_progress_batch.flush(&merge_progress); merge_progress.log_final(); log_snapshot("phase2.end", 0); @@ -6314,7 +5261,7 @@ impl RawExternalSorter { /// Identical input/output pipeline to `merge_chunks_generic` — the shared /// worker pool decompresses the input runs (`PoolDisk` sources) and /// compresses the output blocks — but the output goes through - /// [`PooledBamWriter::new_indexing`], which tracks each record's virtual + /// [`crate::pooled_bam_writer::PooledBamWriter::new_indexing`], which tracks each record's virtual /// file offset and returns the generated BAI index. A single pool of /// workers therefore serves both decompression and compression; there is no /// separate writer thread pool and no serialized single-reader input. @@ -6358,7 +5305,6 @@ impl RawExternalSorter { let mut tree = LoserTree::new(initial_keys); let mut writer = PooledBamWriter::new_indexing(Arc::clone(pool), output, &output_header)?; - let mut merge_progress_batch = crate::progress_batch::BatchedProgress::new(); let merge_progress = ProgressTracker::new("Merged records") .with_interval(1_000_000) .with_total(total_records); @@ -6374,28 +5320,13 @@ impl RawExternalSorter { } let loop_start = Instant::now(); let mut records_merged: u64 = 0; - let mut published_src: Option = None; - let mut fetch_counts = RecordFetchCounts::default(); while tree.winner_is_active() { let winner = tree.winner(); let src_idx = source_map[winner]; - // Publish the next likely source, but only when the run changes. The - // pool gives it the deep read allowance, so its first read starts - // while the consumer is still draining the current file instead of - // after it has already stalled. A wrong prediction costs one deep - // read on a file the merge will reach anyway; there is no correctness - // component. Mirrors `merge_chunks_generic`; without it the indexed - // path -- the default coordinate sort -- gets no predictive - // read-ahead. - if published_src != Some(src_idx) { - published_src = Some(src_idx); - pool.set_phase2_next_source(tree.runner_up().map(|w| source_map[w])); - } - let record_bytes = - winner_record_bytes(&sources[src_idx], guard.consumer_ref(), &mut fetch_counts)?; + let record_bytes = winner_record_bytes(&sources[src_idx], guard.consumer_ref())?; writer.write_raw_record(record_bytes)?; records_merged += 1; - merge_progress_batch.tick(&merge_progress); + merge_progress.log_if_needed(1); if let Some(key) = sources[src_idx].advance(guard.consumer_mut())? { tree.replace_winner(key); @@ -6414,7 +5345,6 @@ impl RawExternalSorter { // cannot change it, and the consumer's report is harvested before // `finish_output` releases the merge sources and with them the // consumer. Both describe the loop that has just ended. - fetch_counts.publish(); let loop_total = loop_start.elapsed().as_secs_f64(); let active_workers = pool.active_workers(); let stalls = { @@ -6426,14 +5356,6 @@ impl RawExternalSorter { guard.consumer_ref().map(MainThreadChunkConsumer::stall_report) }; - // Retain the permit pool so the writer histograms can be harvested *after* - // `finish_index` drains the output queue, exactly as the generic path - // does: the pool outlives the writer, and the block writes and reorder - // waits incurred during that drain belong in the "write block" and - // "write reorder wait" rows. A snapshot taken here, before the drain, - // would omit the tail. - let writer_permit_pool = writer.permit_pool(); - // Finalize before logging, for the reason the generic path does: `finish` // drains the output queue, and every block still in it is compressed by // the same workers `log_merge_stalls` divides into `merge_total`. Logging @@ -6444,14 +5366,6 @@ impl RawExternalSorter { // number. `loop_total` stays the consumer's park-fraction denominator, // which is a fraction of the merge loop alone. let index = guard.finish_output(|| writer.finish_index())?; - - // Harvest now, after the drain, from the retained pool. Empty only if the - // writer was already finalized when the pool was retained, which cannot - // happen on this path. - let writer_stats = writer_permit_pool - .as_deref() - .map_or_else(Default::default, crate::worker_pool::PermitPool::writer_stats); - Self::log_merge_stalls( loop_total, loop_start.elapsed().as_secs_f64(), @@ -6459,15 +5373,7 @@ impl RawExternalSorter { stalls, pool, ); - // `--write-index` is the default for a coordinate sort, so wiring the - // stage-latency table here -- matching `merge_chunks_generic` -- is what - // gives most production merges any stage-latency and writer-histogram - // rows at all. - Self::log_stage_latency(pool, &writer_stats); - - // Must precede log_final: the tracker has not seen the last partial batch, - // and without this the reported total comes up short by up to one batch. - merge_progress_batch.flush(&merge_progress); + merge_progress.log_final(); Ok((index, records_merged)) } @@ -6477,6 +5383,22 @@ impl RawExternalSorter { super::create_output_header(self.sort_order, header) } + /// Create the spill temp dirs + allocator for the P6 `CompressSpill` step. + /// + /// Public wrapper over `Self::create_temp_dirs`: honors `--temp-dir` / + /// `FGUMI_TMP_DIRS` (the configured `temp_dirs`) exactly as the legacy + /// streaming path does. The caller hands the [`Vec`] RAII handles to + /// `CompressSpill` (held for the step's lifetime) and the [`TmpDirAllocator`] + /// behind a shared mutex for per-file base-dir allocation. + /// + /// # Errors + /// + /// Returns an error if a temp directory cannot be created or has insufficient + /// free space. + pub fn create_spill_dirs(&self) -> Result<(Vec, TmpDirAllocator)> { + self.create_temp_dirs() + } + /// Create per-base temp directories and an allocator over their subdirs. /// /// For each user-supplied base directory, a fresh sort-run subdirectory is @@ -6685,104 +5607,16 @@ pub(crate) use crate::SortStats as RawSortStats; // rather than in `fgumi-bam-io`. // ============================================================================ -/// How many concurrent streams to read a file with. -/// -/// `Auto` starts at one and grows only when fills are demonstrably occupying -/// the span (see `spill_reader`); a fixed value pins it. One measured -/// device wants four and another wants one, and the difference is 28% on the -/// first and -1.8% on the second, which is why the default measures rather -/// than guesses. -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] -pub enum ReadStreams { - /// Measure and grow. The default. - #[default] - Auto, - /// Exactly this many; `1` is the plain sequential reader. - Fixed(usize), -} - -impl std::fmt::Display for ReadStreams { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Auto => f.write_str("auto"), - Self::Fixed(n) => write!(f, "{n}"), - } - } -} - -impl ReadStreams { - /// Whether this asks for the plain sequential reader, which is what every - /// phase did before scattered reads existed. - #[must_use] - pub fn is_sequential(self) -> bool { - matches!(self, Self::Fixed(1)) - } - - /// Streams to start with. `Auto` starts at one and grows from what it - /// measures, so it costs nothing until it has evidence. - #[must_use] - pub fn initial(self) -> usize { - match self { - Self::Auto => 1, - Self::Fixed(n) => n.max(1), - } - } - - /// Whether the reader should tune itself as it goes. - #[must_use] - pub fn is_auto(self) -> bool { - matches!(self, Self::Auto) - } -} - -impl std::str::FromStr for ReadStreams { - type Err = String; - - fn from_str(text: &str) -> Result { - if text.eq_ignore_ascii_case("auto") { - return Ok(Self::Auto); - } - match text.parse::() { - Ok(0) => Err("--read-streams must be `auto` or at least 1".to_string()), - Ok(n) => Ok(Self::Fixed(n)), - Err(_) => Err(format!("expected `auto` or a positive number, got `{text}`")), - } - } -} - -/// Fills the phase-1 input reader keeps in flight ahead of the framer. -/// -/// One, and deeper does not help. Swept in a single boot on -/// `1kg-wgs-HG00096` (t16, 4 streams), read span against total wall: -/// -/// | depth | read span | wall | -/// | --- | --- | --- | -/// | 0 | 69.4s | 271.8s | -/// | 1 | 66.9s | 267.7s | -/// | 2 | 66.7s | 269.4s | -/// | 4 | 66.4s | 268.5s | -/// -/// Depths 2 and 4 move the span by less than a second and the wall clock not -/// at all -- all three sit inside the +/-0.7% in-boot noise floor -- and depth -/// 2 cost 277 MB of peak RSS for it. One is enough because the span is not -/// fetch-bound: the fetch runs at 1327 MB/s and takes 32.5s of a 66.4s span -/// against an ingest-serial floor of 53.0s, so what is left is coordination -/// between the reader, the decompress workers and the ingest thread, which no -/// amount of read-ahead touches. -/// -/// Costs `(n + 1) * FILL_BYTES` of buffers on one reader, which is also why -/// the merge's spill readers -- there are K of them, and the merge is already -/// at its consumer-serial floor -- get none. -const PHASE1_LOOKAHEAD_FILLS: usize = 1; - /// Create a raw BAM reader using the pool's Phase 1 integrated reading. /// /// Workers in the pool do `ReadInputBlocks` + `DecompressInput`. The main /// thread consumes decompressed bytes via `PooledInputStream`. /// -/// No extra threads are spawned either way: with one stream the pool's block -/// reader reads directly from the input file, and with several it offers byte -/// slices to the pool (see `spill_reader`). +/// When `async_reader` is false, no extra threads are spawned: the pool's +/// block reader reads directly from the input file. When `async_reader` is +/// true, the input file is wrapped in a `PrefetchReader`, which spawns one +/// dedicated OS thread (`fgumi-prefetch`) that reads raw bytes ahead into a +/// bounded queue so the pool's block reader never blocks on disk I/O. /// /// # Flow /// @@ -6800,84 +5634,45 @@ const PHASE1_LOOKAHEAD_FILLS: usize = 1; fn create_raw_bam_reader_pool_integrated>( path: P, pool: &Arc, - read_streams: ReadStreams, + async_reader: bool, ) -> Result<(fgumi_raw_bam::RawBamReader, Header)> { use crate::worker_pool::phase; use std::io; let path_ref = path.as_ref(); - // Times the reads that refill the buffer below, so that `framing_secs` -- - // time inside `read_raw_blocks` that was *not* spent fetching -- means what - // it says. Both arms report into it: the sequential one through - // `TimedReader`, the scattered one by timing its own fills. Leaving the - // scattered arm out booked its fill waits as framing and showed 10.9 - // us/block against a true 0.9. - let reader_stats = pool.reader_stats(); - let opened: Box = if is_stdin_path(path_ref) { - // `io::Stdin` re-acquires a mutex and reads through an 8 KiB buffer on - // every call; the pool's block reader wants far bigger gulps than that, - // so give the stdin path the same 2 MiB buffer the file path gets. - // Scattered reads are not available here -- a pipe has no offsets. - Box::new(io::BufReader::with_capacity( - SORT_INPUT_BUFFER_SIZE, - crate::phase1_stats::TimedReader::new(io::stdin(), reader_stats), - )) + if async_reader { + // `--async-reader` is about decoupling the read from the block + // reader, which stdin needs at least as much as a file does: the + // prefetch thread also subsumes the buffering below, reading ahead + // in chunks into a bounded queue. + log::debug!("async sort reader enabled: spawning fgumi-prefetch thread for stdin"); + Box::new(fgumi_bam_io::prefetch_reader::PrefetchReader::new(io::stdin())) + } else { + // `io::Stdin` re-acquires a mutex and reads through an 8 KiB buffer + // on every call; the pool's block reader wants far bigger gulps than + // that, so give the stdin path the same 2 MiB buffer the file path + // gets. + Box::new(io::BufReader::with_capacity(SORT_INPUT_BUFFER_SIZE, io::stdin())) + } } else { let file = std::fs::File::open(path_ref) .with_context(|| format!("Failed to open input BAM: {}", path_ref.display()))?; - // Grow the per-fd readahead window. Measured worth almost nothing on - // its own (359 MB/s against 360 with no hint), but it is free and - // applies however the bytes are subsequently read. + // Grow the per-fd readahead window. This is the plain sequential hint + // and applies however the bytes are subsequently read; the WILLNEED + // hints that `PrefetchReader` issues are a separate, async-only extra. fgumi_bam_io::os_hints::advise_sequential(&file); - if read_streams.is_sequential() { - // `--read-streams 1`: the plain buffered reader every phase used - // before scattered reads existed. - Box::new(io::BufReader::with_capacity( - SORT_INPUT_BUFFER_SIZE, - crate::phase1_stats::TimedReader::new(file, reader_stats), - )) + if async_reader { + log::debug!( + "async sort reader enabled: spawning fgumi-prefetch thread for {}", + path_ref.display() + ); + Box::new(fgumi_bam_io::prefetch_reader::PrefetchReader::from_file(file)) } else { - // Scattered positional reads, the same mechanism the merge uses on - // spill files: slices are offered to the pool so the device sees - // real queue depth, which one blocking `read()` cannot produce - // however large its buffer. The framer downstream is unchanged -- - // this still presents a sequential `Read`. - // - // Only reachable on a real file: stdin and other non-seekable inputs - // take the branch above, where positional reads do not exist. - match crate::spill_reader::ScatterReader::for_streams( - file, - 0, - read_streams, - Some(pool.fetch_queue()), - ) - .map(|reader| { - reader.timed(Arc::clone(&reader_stats)).looking_ahead(PHASE1_LOOKAHEAD_FILLS) - }) { - Ok(reader) => { - log::debug!( - "scattered sort reader: {read_streams} over {}", - path_ref.display() - ); - Box::new(reader) - } - Err(e) => { - // A file whose length we cannot read is one we cannot slice. - // Fall back rather than fail a sort over a prefetch tweak. - log::warn!("scattered sort reader unavailable ({e}); using one stream"); - let file = std::fs::File::open(path_ref).with_context(|| { - format!("Failed to reopen input BAM: {}", path_ref.display()) - })?; - Box::new(io::BufReader::with_capacity( - SORT_INPUT_BUFFER_SIZE, - crate::phase1_stats::TimedReader::new(file, reader_stats), - )) - } - } + Box::new(io::BufReader::with_capacity(SORT_INPUT_BUFFER_SIZE, file)) } }; @@ -6898,7 +5693,6 @@ fn create_raw_bam_reader_pool_integrated>( pool.decompressed_input_done_flag(), pool.input_read_error_flag(), pool.decompress_error_flag(), - pool.phase1_ingest_stats(), ); // Deliberately not phrased as a header failure. The header was parsed @@ -7108,6 +5902,668 @@ pub(crate) fn dropped_lane_error(name: &str, v: DroppedLaneViolation) -> anyhow: ) } +/// Open `path` as a single [`SortMergeSlot`] with the given `file_id`. +/// +/// The spill codec is auto-detected from the file magic: for zstd (`ZSP1`) the +/// 4-byte magic is consumed here, leaving the reader at the first `[len][frame]` +/// record; for BGZF the `1f 8b` is part of the first block, so the reader is +/// rewound to byte 0. `SortSpillDecompress` reads `slot.codec` to decompress +/// accordingly. +/// +/// `file_id` is the slot's stable merge-order identifier — `SortMerge` orders +/// sources by it so the `LoserTree` tie-break for equal sort keys matches the +/// legacy chunk-files order. The P6 `CompressSpill` step passes a chunk's +/// **logical spill index** here so the tie-break is independent of which worker +/// happened to write the file. +/// +/// # Errors +/// +/// Returns an error if the file cannot be opened, its magic cannot be read, or +/// (for BGZF) the rewind fails. +pub fn open_spill_slot(path: &std::path::Path, file_id: u32) -> Result> { + let mut file = std::fs::File::open(path) + .with_context(|| format!("failed to open spill chunk {}", path.display()))?; + let mut magic = [0u8; 4]; + let filled = read_exact_or_eof(&mut file, &mut magic) + .with_context(|| format!("failed to read spill chunk magic {}", path.display()))?; + let codec = if filled { + crate::codec::SpillCodec::from_magic(&magic).unwrap_or(crate::codec::SpillCodec::Bgzf) + } else { + crate::codec::SpillCodec::Bgzf + }; + if matches!(codec, crate::codec::SpillCodec::Bgzf) { + use std::io::Seek; + file.seek(std::io::SeekFrom::Start(0)) + .with_context(|| format!("failed to rewind spill chunk {}", path.display()))?; + } + let reader = std::io::BufReader::new(file); + Ok(Arc::new(SortMergeSlot::new(file_id, reader, codec))) +} + +// ============================================================================ +// CoordinateSortStream — streaming coordinate-sort handle +// ============================================================================ + +// ============================================================================ +// QuerynameSortStream — streaming queryname-sort handle (lex or natural) +// ============================================================================ + +// ============================================================================ +// TemplateCoordinateSortStream — streaming template-coordinate-sort handle +// ============================================================================ + +// ============================================================================ +// Streaming-sort merge path (unified pipeline `SortMerge` step) +// ============================================================================ +// +// Ported from the issue-#330 branch onto main's sort engine. This is the +// slot-backed, *non-blocking* merge driver used exclusively by the typed-step +// pipeline (`MergeDriver::from_slots`). It reads records out of +// `SortMergeSlot` queues, which the `SortSpillDecompress` step has *already* +// filled with decompressed `Vec` blocks — so this path is **codec-agnostic** +// (it never touches disk or the spill codec) and carries no chunk-corruption +// risk. The eager pool-integrated `MergeDriver::new` path from #330 is +// intentionally NOT ported: standalone `fgumi sort` uses main's own +// `merge_chunks_*` engine, and the pipeline only ever calls `from_slots`. + +/// Result of a non-blocking source read. `WouldBlock` is only ever produced by +/// `Slot` sources whose decompressed queue is momentarily empty and not yet at +/// EOF; the in-memory source always returns `Ready`. +pub(crate) enum TryRead { + /// A record was read (`Some`) or the source is at clean EOF (`None`). + Ready(Option), + /// The source has no bytes available right now but is not at EOF; the + /// caller should yield and retry on a later dispatch. + WouldBlock, +} + +/// Per-slot byte-stream parser state for the non-blocking merge reader. +/// +/// As decompressed blocks are popped from a `SortMergeSlot`, their bytes are +/// stashed in `current_buf` and consumed left-to-right by the record parser. +/// Named distinctly from main's own `SourceParserState` (the pool-merge path) +/// because this one additionally carries resumable cross-block framing state. +struct SlotParserState { + /// Current decompressed block being consumed. + current_buf: Vec, + /// Read position within `current_buf`. + current_pos: usize, + /// Resumable framing state for [`slot_try_next_record`]. `Some` while a + /// record spans a block boundary and the next block was not yet available, + /// so the read returned `WouldBlock` mid-record; the next call resumes from + /// here. + pending: Option, +} + +/// Partially-read record, retained across a `WouldBlock` so the non-blocking +/// slot reader can resume. The wire layout is `[key:ksize][len:4][body:len]` +/// for non-embedded keys (`ksize = K::SERIALIZED_SIZE`) and `[len:4][body:len]` +/// for embedded keys (`ksize = 0`, key extracted from the body). We collect the +/// key prefix (if any), then the length prefix, then the body. +#[derive(Default)] +struct PendingRecord { + /// Expected serialized key size (`0` for embedded keys). + key_size: usize, + /// Key-prefix bytes collected so far (length grows to `key_size`). + key_buf: Vec, + /// Length-prefix bytes collected so far (`0..=4`). + len_buf: [u8; 4], + /// Count of `len_buf` bytes filled. + len_got: usize, + /// Total body length, known once `len_got == 4`. The body accumulates + /// directly into the caller's record buffer; progress is `out.len()`. + body_len: Option, +} + +impl SlotParserState { + fn new() -> Self { + Self { current_buf: Vec::new(), current_pos: 0, pending: None } + } + + fn remaining(&self) -> usize { + self.current_buf.len() - self.current_pos + } +} + +/// Outcome of a non-blocking attempt to pull the next decompressed block from a +/// slot into `parser.current_buf`. +enum BlockLoad { + /// A block was moved into `parser.current_buf`. + Loaded, + /// The slot is fully delivered (queue empty AND `queue_eof`). + Eof, + /// The queue is empty but the producer is still feeding this slot. + WouldBlock, +} + +/// Pop one decompressed block (FIFO) into `parser`, surfacing a decompression +/// error in preference to EOF, reporting clean EOF, or `WouldBlock` while the +/// producer is still active. Never parks. +fn slot_try_load_block( + slot: &Arc, + parser: &mut SlotParserState, +) -> Result { + use std::sync::atomic::Ordering; + + let mut guard = slot.decompressed.lock().expect("slot.decompressed mutex poisoned"); + if let Some(data) = guard.pop_front() { + drop(guard); + parser.current_buf = data; + parser.current_pos = 0; + return Ok(BlockLoad::Loaded); + } + // Queue empty. Both atomics are visible under the lock release-acquire + // chain (producer stores them while holding this same mutex). Surface + // error in preference to EOF. + if slot.decomp_error.load(Ordering::Acquire) { + drop(guard); + anyhow::bail!("spill decompression error on slot {}", slot.file_id); + } + if slot.queue_eof.load(Ordering::Acquire) { + return Ok(BlockLoad::Eof); + } + Ok(BlockLoad::WouldBlock) +} + +/// Serialized key size on the wire for a slot source: `0` for embedded keys +/// (key extracted from the body), else `K::SERIALIZED_SIZE`. Variable-length +/// non-embedded keys are unsupported on the slot path (none exist in +/// production — every slot-path key is embedded). +fn slot_key_size() -> Result { + if K::EMBEDDED_IN_RECORD { + Ok(0) + } else { + K::SERIALIZED_SIZE.ok_or_else(|| { + anyhow::anyhow!("non-embedded slot key must have a fixed serialized size") + }) + } +} + +/// Parse a sort key from `bytes`: extracted from the body for embedded keys, or +/// deserialized from the `key_size`-byte prefix for non-embedded keys. +fn slot_parse_key(key_bytes: &[u8], body: &[u8]) -> Result { + if K::EMBEDDED_IN_RECORD { + // An embedded key is read out of the body, so an empty body cannot + // carry one. Every extractor indexes straight into the slice + // (`RawCoordinateKey` reads `ref_id` off the front), so without this an + // empty body panics mid-merge instead of surfacing a corrupt-slot + // error — and it is reachable: a corrupt length prefix of 0 makes the + // body loop break immediately with `out` still empty. + // + // This is a floor, not a full bounds check. The minimum body an + // extractor needs is key-specific (8 bytes here, 16 for the coordinate + // key, more for queryname), and `RawSortKey` does not expose it, so a + // corrupt *small but non-zero* length can still panic inside an + // extractor. Closing that needs a per-key minimum on the trait. + anyhow::ensure!( + !body.is_empty(), + "record body is empty and cannot carry an embedded sort key (corrupt slot)", + ); + Ok(K::extract_from_record(body)) + } else { + let mut r = key_bytes; + K::read_from(&mut r).map_err(|e| anyhow::anyhow!("malformed slot key: {e}")) + } +} + +/// Non-blocking, resumable read of the next record from a slot. +/// +/// Returns `Ready(Some(key))` with the record body in `out`, `Ready(None)` at +/// clean source EOF, or `WouldBlock` when a needed block is not yet available +/// (progress is saved in `parser.pending` so the next call resumes). +/// +/// # Errors +/// +/// Truncation (EOF mid-record), a malformed key, or a producer-side +/// decompression error. +fn slot_try_next_record( + slot: &Arc, + parser: &mut SlotParserState, + out: &mut Vec, +) -> Result> { + let key_size = slot_key_size::()?; + + if parser.pending.is_none() { + // Ensure at least one byte is available, or detect clean EOF at a + // record boundary. + if parser.remaining() == 0 { + match slot_try_load_block(slot, parser)? { + BlockLoad::Loaded => {} + BlockLoad::Eof => return Ok(TryRead::Ready(None)), + BlockLoad::WouldBlock => return Ok(TryRead::WouldBlock), + } + } + // Fast path: the entire record lives in the current block → parse in + // place with a single copy of the body into `out`. + let header = key_size + 4; + if parser.remaining() >= header { + let p = parser.current_pos; + let len = u32::from_le_bytes( + parser.current_buf[p + key_size..p + header].try_into().expect("4-byte slice"), + ) as usize; + if parser.remaining() >= header + len { + out.clear(); + out.extend_from_slice(&parser.current_buf[p + header..p + header + len]); + let key = slot_parse_key::(&parser.current_buf[p..p + key_size], out)?; + parser.current_pos += header + len; + return Ok(TryRead::Ready(Some(key))); + } + } + // Slow path: record spans a block boundary (or fewer than `header` + // bytes remain). Begin a resumable pending record. + parser.pending = Some(PendingRecord { key_size, ..PendingRecord::default() }); + } + + slot_collect_pending::(slot, parser, out) +} + +/// Resumable continuation of [`slot_try_next_record`]'s slow path. Collects the +/// key prefix (if any), the 4-byte length prefix, then the body across as many +/// blocks as needed, returning `WouldBlock` (state retained in +/// `parser.pending`) whenever the next block isn't ready yet. +fn slot_collect_pending( + slot: &Arc, + parser: &mut SlotParserState, + out: &mut Vec, +) -> Result> { + // Stage 0: key prefix (non-embedded keys only; `key_size == 0` is a no-op). + loop { + let (key_size, key_have) = { + let pending = parser.pending.as_ref().expect("pending set"); + (pending.key_size, pending.key_buf.len()) + }; + let need = key_size - key_have; + if need == 0 { + break; + } + if parser.current_pos == parser.current_buf.len() { + match slot_try_load_block(slot, parser)? { + BlockLoad::Loaded => {} + BlockLoad::Eof => anyhow::bail!( + "truncated record key in slot {} ({key_have} of {key_size} key bytes at EOF)", + slot.file_id, + ), + BlockLoad::WouldBlock => return Ok(TryRead::WouldBlock), + } + } + let avail = parser.current_buf.len() - parser.current_pos; + let take = need.min(avail); + let src = &parser.current_buf[parser.current_pos..parser.current_pos + take]; + parser.pending.as_mut().expect("pending set").key_buf.extend_from_slice(src); + parser.current_pos += take; + } + + // Stage 1: length prefix. + loop { + let len_got = parser.pending.as_ref().expect("pending set").len_got; + if len_got == 4 { + break; + } + if parser.current_pos == parser.current_buf.len() { + match slot_try_load_block(slot, parser)? { + BlockLoad::Loaded => {} + BlockLoad::Eof => anyhow::bail!( + "truncated record length in slot {} ({len_got} of 4 length bytes at EOF)", + slot.file_id, + ), + BlockLoad::WouldBlock => return Ok(TryRead::WouldBlock), + } + } + let avail = parser.current_buf.len() - parser.current_pos; + let take = (4 - len_got).min(avail); + let src = &parser.current_buf[parser.current_pos..parser.current_pos + take]; + let pending = parser.pending.as_mut().expect("pending set"); + pending.len_buf[len_got..len_got + take].copy_from_slice(src); + pending.len_got += take; + parser.current_pos += take; + } + + // Compute body length once and reset `out` for accumulation. + { + let pending = parser.pending.as_mut().expect("pending set"); + if pending.body_len.is_none() { + let len = u32::from_le_bytes(pending.len_buf) as usize; + // The prefix is untrusted — it comes straight off a decompressed + // spill block — and the next statement hands it to `reserve`, which + // answers an allocation failure by aborting the process rather than + // returning an error. Bound it first, so a corrupt slot surfaces + // `Err` like the `decomp_error` and truncation arms do. The ceiling + // is the one a single record already cannot exceed: a record must + // fit within one sort segment (see `SORT_SEGMENT_SIZE`). + anyhow::ensure!( + len <= crate::inline::SORT_SEGMENT_SIZE, + "implausible record length {len} in slot {} (exceeds the \ + {}-byte segment a single record must fit in)", + slot.file_id, + crate::inline::SORT_SEGMENT_SIZE, + ); + pending.body_len = Some(len); + out.clear(); + out.reserve(len); + } + } + + // Stage 2: body. + loop { + let body_len = parser.pending.as_ref().expect("pending set").body_len.expect("len known"); + if out.len() >= body_len { + break; + } + if parser.current_pos == parser.current_buf.len() { + match slot_try_load_block(slot, parser)? { + BlockLoad::Loaded => {} + BlockLoad::Eof => anyhow::bail!( + "truncated record body in slot {} ({} of {body_len} bytes at EOF)", + slot.file_id, + out.len(), + ), + BlockLoad::WouldBlock => return Ok(TryRead::WouldBlock), + } + } + let avail = parser.current_buf.len() - parser.current_pos; + let take = (body_len - out.len()).min(avail); + out.extend_from_slice(&parser.current_buf[parser.current_pos..parser.current_pos + take]); + parser.current_pos += take; + } + + let key = { + let pending = parser.pending.as_ref().expect("pending set"); + slot_parse_key::(&pending.key_buf, out)? + }; + parser.pending = None; + Ok(TryRead::Ready(Some(key))) +} + +/// Merge source for the slot-backed `MergeDriver`: either a decompressing slot +/// (read non-blocking) or an already-sorted in-memory residual chunk. +enum SlotMergeSource { + /// Pipeline-integrated source — `SortSpillDecompress` fills + /// `slot.decompressed` with decompressed blocks; this driver reads them + /// in-order via the per-source [`SlotParserState`]. + Slot { slot: Arc, parser: SlotParserState }, + /// In-memory sorted residual records from a queryname-style sort (each + /// record an owned `RawRecord`). Read via a zero-copy `mem::swap` bridge. + Memory { records: Vec<(K, fgumi_raw_bam::RawRecord)>, idx: usize }, + /// In-memory sorted residual from an inline-buffer (coordinate / template) + /// sort, sharing an `Arc` — zero per-record allocation. Read + /// copies each record's bytes into the caller's reused `buf`. + MemoryShared { chunk: InMemoryChunk, idx: usize }, +} + +impl SlotMergeSource { + /// Non-blocking read of the next record into `buf`, returning the sort key. + /// `Slot` sources may return `WouldBlock`; `Memory` always returns `Ready`. + fn try_next_record(&mut self, buf: &mut Vec) -> Result> { + match self { + SlotMergeSource::Slot { slot, parser } => slot_try_next_record::(slot, parser, buf), + SlotMergeSource::Memory { records, idx } => { + if *idx < records.len() { + let (ref mut key, ref mut data) = records[*idx]; + // Bridge: RawRecord wraps Vec; swap via the inner vec to + // avoid re-allocating. The caller's buf is a plain Vec. + std::mem::swap(buf, data.as_mut_vec()); + let key = std::mem::take(key); + *idx += 1; + Ok(TryRead::Ready(Some(key))) + } else { + Ok(TryRead::Ready(None)) + } + } + SlotMergeSource::MemoryShared { chunk, idx } => { + if *idx < chunk.len() { + // Bytes are borrowed from the shared `Arc`, so + // copy them into the caller's reused `buf` (one memcpy, no + // per-record allocation — the materialise-time copy is gone). + buf.clear(); + buf.extend_from_slice(chunk.record_bytes(*idx)); + let key = chunk.take_key(*idx); + *idx += 1; + Ok(TryRead::Ready(Some(key))) + } else { + Ok(TryRead::Ready(None)) + } + } + } + } +} + +/// One step of a non-blocking k-way merge. +pub enum MergeStep<'a> { + /// The next merged record. The bytes borrow the driver's internal winner + /// buffer and are valid only until the next `try_step` call — the caller + /// must copy them out before stepping again. The borrow checker enforces + /// this, which is what makes the deferred-refill protocol sound. + Produced(&'a [u8]), + /// A source's decompressed queue is momentarily empty (and not at EOF), or + /// a source hasn't been primed yet. No record is available right now; the + /// caller should yield and retry on a later dispatch, by which point the + /// producer will have refilled the slot. + Stalled, + /// The merge is exhausted. + Done, +} + +impl std::fmt::Debug for MergeStep<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Produced(bytes) => write!(f, "Produced({} bytes)", bytes.len()), + Self::Stalled => f.write_str("Stalled"), + Self::Done => f.write_str("Done"), + } + } +} + +/// Object-safe view over `MergeDriver` so the pipeline `SortMerge` step can +/// hold a single concrete type regardless of sort order. +pub trait MergeDriverDyn: Send { + /// Produce the next merged record **without blocking**. Any step that needs + /// a not-yet-ready slot block returns [`MergeStep::Stalled`]. + /// + /// # Errors + /// + /// Propagates spill-file corruption or decompression errors from the + /// underlying slot sources. + fn try_step(&mut self) -> Result>; + /// Total records emitted since the driver was constructed. + fn records_merged(&self) -> u64; +} + +/// Internal phase of the [`MergeDriver`] non-blocking state machine. +enum MergePhase { + /// Lazily priming the loser tree: each source's first record is pulled one + /// at a time (non-blocking). `next` is the index of the next source to + /// prime; `keys` accumulates the primed initial keys; `rec` is the partial + /// record buffer for the source currently being primed (retained across a + /// `Stalled` so a spanning first record resumes correctly). + Priming { next: usize, keys: Vec, rec: Vec }, + /// Active k-way merge. `pending_refill` is set after a record is emitted + /// (`Produced`) and means the winner's source must be refilled before the + /// next winner is produced. + Merging { tree: crate::loser_tree::LoserTree, pending_refill: bool }, + /// Exhausted. + Done, +} + +/// Pausable state machine for a non-blocking k-way merge over slot + in-memory +/// sources. Constructed via [`MergeDriver::from_slots`]; driven by the pipeline +/// `SortMerge` step via [`MergeDriverDyn::try_step`]. +pub struct MergeDriver { + sources: Vec>, + /// Per-active-source current record buffer, indexed by loser-tree leaf + /// (parallel to `source_map`). Grows during `Priming`. + records: Vec>, + /// Loser-tree leaf → `sources` index. Built during `Priming`. + source_map: Vec, + /// Non-blocking merge state machine. + phase: MergePhase, + records_merged: u64, + progress: ProgressTracker, +} + +impl MergeDriver { + /// Construct a driver from `SortMergeSlot`s and already-sorted in-memory + /// chunks. Used by the cooperative pipeline `SortMerge` step. + /// + /// Each slot becomes a slot source read **non-blocking** by + /// [`MergeDriverDyn::try_step`]. Priming is **lazy**: the loser tree is + /// built on the first `try_step` call(s) as each source yields its first + /// record, so construction never blocks on a not-yet-decompressed slot. + /// Construction is infallible (read errors surface from `try_step`), and + /// the empty-input case is reported by the first `try_step` returning + /// [`MergeStep::Done`]. + #[must_use] + pub fn from_slots( + slots: Vec>, + memory: MemorySources, + total_records: u64, + ) -> Self { + let num_slots = slots.len(); + let num_memory = memory.num_non_empty(); + let num_sources = num_slots + num_memory; + + if num_slots > 0 { + info!( + "Pipeline-integrated merge: {num_slots} slot sources + {num_memory} memory sources" + ); + } + + let mut sources: Vec> = Vec::with_capacity(num_sources); + // Order by `file_id`, do not trust the caller's `Vec` order. `file_id` + // is documented on `SortMergeSlot` as the stable merge-order identity, + // and the `LoserTree` breaks ties between equal sort keys by leaf index + // — so if the leaf order came from push order instead, two records with + // the same key could be emitted in a different order than the legacy + // chunk-files path produces. That is an output-identity divergence in a + // tool whose sort output is checked byte-for-byte, and it would be + // invisible except on inputs that happen to contain ties. Sorting a + // handful of slots once per merge costs nothing next to that. + let mut slots = slots; + slots.sort_by_key(|s| s.file_id); + for slot in slots { + sources.push(SlotMergeSource::Slot { slot, parser: SlotParserState::new() }); + } + match memory { + MemorySources::Owned(chunks) => { + for chunk in chunks { + if !chunk.is_empty() { + sources.push(SlotMergeSource::Memory { records: chunk, idx: 0 }); + } + } + } + MemorySources::Shared(chunks) => { + for chunk in chunks { + if !chunk.is_empty() { + sources.push(SlotMergeSource::MemoryShared { chunk, idx: 0 }); + } + } + } + } + + let progress = ProgressTracker::new("Merged records") + .with_interval(1_000_000) + .with_total(total_records); + + Self { + sources, + records: Vec::with_capacity(num_sources), + source_map: Vec::with_capacity(num_sources), + // Lazy priming: tree built on first try_step as sources yield. + phase: MergePhase::Priming { + next: 0, + keys: Vec::with_capacity(num_sources), + rec: Vec::new(), + }, + records_merged: 0, + progress, + } + } +} + +impl MergeDriverDyn for MergeDriver { + fn try_step(&mut self) -> Result> { + // Move the phase out so the body works with owned `tree`/`keys`/`rec` + // and can borrow `self.{sources,records,source_map}` freely without + // aliasing the phase; the phase is written back before every return. + // `Done` is the placeholder; any early `?` leaves the driver `Done`, + // which is safe because the whole merge aborts on error. + loop { + match std::mem::replace(&mut self.phase, MergePhase::Done) { + MergePhase::Priming { mut next, mut keys, mut rec } => { + let mut stalled = false; + while next < self.sources.len() { + match self.sources[next].try_next_record(&mut rec)? { + TryRead::WouldBlock => { + stalled = true; + break; + } + TryRead::Ready(Some(key)) => { + keys.push(key); + self.records.push(std::mem::take(&mut rec)); + self.source_map.push(next); + next += 1; + } + TryRead::Ready(None) => { + rec.clear(); + next += 1; + } + } + } + if stalled { + self.phase = MergePhase::Priming { next, keys, rec }; + return Ok(MergeStep::Stalled); + } + if keys.is_empty() { + return Ok(MergeStep::Done); + } + info!("Merging from {} sources...", keys.len()); + let tree = crate::loser_tree::LoserTree::new(keys); + self.phase = MergePhase::Merging { tree, pending_refill: false }; + // Loop to emit the first winner. + } + MergePhase::Merging { mut tree, pending_refill } => { + // Resolve the refill deferred from the previous `Produced` + // (the read that would overwrite the just-emitted winner's + // buffer). On `WouldBlock`, yield with the refill still + // pending. + if pending_refill && tree.winner_is_active() { + let winner = tree.winner(); + let src = self.source_map[winner]; + match self.sources[src].try_next_record(&mut self.records[winner])? { + TryRead::WouldBlock => { + self.phase = MergePhase::Merging { tree, pending_refill: true }; + return Ok(MergeStep::Stalled); + } + TryRead::Ready(Some(key)) => tree.replace_winner(key), + TryRead::Ready(None) => tree.remove_winner(), + } + } + if !tree.winner_is_active() { + return Ok(MergeStep::Done); + } + let winner = tree.winner(); + self.records_merged += 1; + self.progress.log_if_needed(1); + // Defer this winner's refill to the next call so the + // returned borrow stays valid until the caller copies it. + self.phase = MergePhase::Merging { tree, pending_refill: true }; + return Ok(MergeStep::Produced(&self.records[winner])); + } + MergePhase::Done => return Ok(MergeStep::Done), + } + } + } + + fn records_merged(&self) -> u64 { + self.records_merged + } +} + +impl Drop for MergeDriver { + fn drop(&mut self) { + self.progress.log_final(); + } +} + #[cfg(test)] mod tests { use super::*; @@ -7741,177 +7197,172 @@ mod tests { } #[test] - fn test_raw_sorter_memory_limit() { - let sorter = RawExternalSorter::new(SortOrder::Coordinate).memory_limit(256 * 1024 * 1024); - assert_eq!(sorter.memory_limit, 256 * 1024 * 1024); - } + fn test_raw_sorter_phase_threads_builders() { + let new = || RawExternalSorter::new(SortOrder::Coordinate).threads(8); + // Default: both phases fall back to threads. + let base = new(); + assert_eq!((base.sort_threads, base.merge_threads), (None, None)); + assert_eq!((base.phase1_threads(), base.phase2_threads()), (8, 8)); + // sort_threads overrides only Phase 1. + let sort_only = new().sort_threads(2); + assert_eq!((sort_only.phase1_threads(), sort_only.phase2_threads()), (2, 8)); + // The public `num_threads()` accessor returns the BASE count, unaffected + // by the Phase-1 override — so the streaming front must read + // `phase1_threads()`, not `num_threads()`, to honor `--sort-threads` + // (the D1.1 regression was reading the base here). Pin the distinction. + assert_eq!(sort_only.num_threads(), 8); + assert_eq!(sort_only.phase1_threads(), 2); + // merge_threads overrides only Phase 2. + let merge_only = new().merge_threads(3); + assert_eq!((merge_only.phase1_threads(), merge_only.phase2_threads()), (8, 3)); + // Both override independently. + let both = new().sort_threads(2).merge_threads(16); + assert_eq!((both.phase1_threads(), both.phase2_threads()), (2, 16)); + // Clamp to >= 1. + let zero = new().sort_threads(0).merge_threads(0); + assert_eq!((zero.phase1_threads(), zero.phase2_threads()), (1, 1)); + } + + /// Read every record's raw bytes from a BAM, in file order. + fn records_bytes(path: &Path) -> Vec> { + let (mut reader, _h) = create_raw_bam_reader(path, 1).expect("open bam"); + let mut out = Vec::new(); + let mut rec = fgumi_raw_bam::RawRecord::new(); + while reader.read_record(&mut rec).expect("read record") != 0 { + out.push(rec.view().as_bytes().to_vec()); + } + out + } + + /// Splitting Phase-1 (`sort_threads`) and Phase-2 (`merge_threads`) from the + /// base `threads` is a pure scheduling knob: output must be byte-identical to + /// the unsplit sort, with spills forced. Matrixed over every order (`#[values]`) + /// and every thread-split variant (`#[case]`) so a failure names the exact order + /// and split. The cases mirror the prior inline matrix: `sort_low` (Phase-1 on + /// 1), `merge_low` (Phase-2 on 1), `both` (Phase-1 on 1 / Phase-2 on 2), and + /// `sort_high` (Phase-1 above base). + #[rstest::rstest] + #[case::sort_low(Some(1), None)] + #[case::merge_low(None, Some(1))] + #[case::both(Some(1), Some(2))] + #[case::sort_high(Some(8), Some(2))] + fn phase_thread_split_is_byte_identical( + #[values( + SortOrder::Coordinate, + SortOrder::Queryname(QuerynameComparator::Natural), + SortOrder::Queryname(QuerynameComparator::Lexicographic), + SortOrder::TemplateCoordinate + )] + order: SortOrder, + #[case] sort_t: Option, + #[case] merge_t: Option, + ) { + use fgumi_sam::SamBuilder; + let mut builder = SamBuilder::new(); + for i in 0..2000 { + let _ = builder + .add_pair() + .name(&format!("read{i}")) + .start1(i * 50 + 1) + .start2(i * 50 + 101) + .build(); + } + let dir = tempfile::tempdir().expect("temp dir"); + let input = dir.path().join("in.bam"); + builder.write_bam(&input).expect("write bam"); - #[test] - fn test_raw_sorter_temp_compression() { - let sorter = RawExternalSorter::new(SortOrder::Coordinate).temp_compression(0); - assert_eq!(sorter.temp_compression, 0); + let mk = |out: &Path, sort_t: Option, merge_t: Option| { + let mut s = RawExternalSorter::new(order) + .threads(4) + .memory_limit(32 * 1024) // force spills + .spill_codec(crate::codec::SpillCodec::Bgzf) + .temp_compression(0) + .output_compression(0); + if let Some(j) = sort_t { + s = s.sort_threads(j); + } + if let Some(k) = merge_t { + s = s.merge_threads(k); + } + s.sort(&input, out).expect("sort"); + }; + // Base: both phases on the unsplit `threads(4)`. + let base = dir.path().join("base.bam"); + mk(&base, None, None); + let golden = records_bytes(&base); + // The split variant must reproduce the base output exactly. + let out = dir.path().join("variant.bam"); + mk(&out, sort_t, merge_t); + assert_eq!( + records_bytes(&out), + golden, + "order {order:?} split sort_t={sort_t:?} merge_t={merge_t:?}: \ + thread split must not change output" + ); } - /// `--sort-stats` is run-scoped: off by default on a fresh sorter (a plain `fgumi sort` - /// prints none of the ~99 diagnostic lines) and set per-sorter through the builder, so two - /// sorts in one process keep independent settings without any process-global flag. - #[test] - fn test_raw_sorter_sort_stats_default_off_and_builder() { - let sorter = RawExternalSorter::new(SortOrder::Coordinate); - assert!(!sorter.sort_stats, "sort stats must default to off"); - assert!(sorter.sort_stats(true).sort_stats); - assert!(!RawExternalSorter::new(SortOrder::Coordinate).sort_stats(false).sort_stats); - } - - /// End-to-end run-scoped `--sort-stats`: a spilling, consolidating, - /// k-way-merged sort with `sort_stats(true)` drives the diagnostic emitters - /// (phase timing, the Phase 1 ingest floor with its reader/ingest partitions, - /// and the merge sub-phase / consumer-CPU reporting) through the real path - /// rather than a hand-built struct. - /// - /// Two details make the diagnostic expressions actually execute: the record - /// count clears the 1-in-1021 ingest sampler so the ingest partition has - /// samples, and the max log level is raised to `Info` so `stat!`'s inner - /// `log::info!` evaluates and formats its arguments (with the default no-op - /// logger the formatted record is discarded, but every diagnostic expression - /// still runs -- which is exactly the run-scoped path under test). A sort with - /// the flag off leaves those expressions unevaluated. - /// - /// `--sort-stats` is diagnostics-only, so both settings must sort the same - /// records to the same output -- asserted here, not merely that the emitters - /// did not panic. - #[test] - fn test_sort_stats_true_drives_diagnostics_through_a_spilling_merge() { + /// Regression guard for the residual Phase-1 sort honoring `sort_threads` + /// rather than the base `threads`. With `threads(1).sort_threads(4)` the + /// residual sort/chunk decision must take the *parallel* `phase1_threads()` + /// branch (previously it gated on `threads`, so it ran serially and ignored + /// `sort_threads`). Spills are forced so the keyed-chunk residual path + /// (`chunk_files` non-empty) is exercised; output must be byte-identical to + /// the serial `threads(1)` sort for every order. + #[rstest::rstest] + #[case::coordinate(SortOrder::Coordinate)] + #[case::queryname_natural(SortOrder::Queryname(QuerynameComparator::Natural))] + #[case::queryname_lex(SortOrder::Queryname(QuerynameComparator::Lexicographic))] + #[case::template_coordinate(SortOrder::TemplateCoordinate)] + fn residual_phase1_honors_sort_threads_with_single_base_thread(#[case] order: SortOrder) { use fgumi_sam::SamBuilder; - - // Above INGEST_SAMPLE_INTERVAL (1021 records) so the ingest sampler fires; - // descending coordinates yield many chunks under the tiny memory limit. - let num_pairs = 2000; let mut builder = SamBuilder::new(); - for i in 0..num_pairs { - let descending = num_pairs - 1 - i; + for i in 0..2000 { let _ = builder .add_pair() - .name(&format!("read{descending:06}")) - .start1(descending * 200 + 1) - .start2(descending * 200 + 101) + .name(&format!("read{i}")) + .start1(i * 50 + 1) + .start2(i * 50 + 101) .build(); } - let dir = tempfile::tempdir().expect("failed to create temp directory"); - let input = dir.path().join("input.bam"); - builder.write_bam(&input).expect("failed to write BAM"); - - // Info so the `stat!` emitters evaluate (and format) their arguments. - log::set_max_level(log::LevelFilter::Info); + let dir = tempfile::tempdir().expect("temp dir"); + let input = dir.path().join("in.bam"); + builder.write_bam(&input).expect("write bam"); - let run = |out: &Path, stats_on: bool| { - RawExternalSorter::new(SortOrder::Coordinate) - .memory_limit(1024) // 1 KB — forces many chunks - .max_temp_files(4) // forces consolidation + a k-way merge - .output_compression(0) - .threads(4) // parallel Phase 1 exercises deferred key extraction - .sort_stats(stats_on) - .sort(input.as_path(), out) - .expect("sort should succeed") + let mk = |out: &Path, sort_t: Option| { + let mut s = RawExternalSorter::new(order) + .threads(1) + .memory_limit(32 * 1024) // force spills -> exercise residual chunk path + .spill_codec(crate::codec::SpillCodec::Bgzf) + .temp_compression(0) + .output_compression(0); + if let Some(j) = sort_t { + s = s.sort_threads(j); + } + s.sort(&input, out).expect("sort"); }; + let serial = dir.path().join("serial.bam"); + mk(&serial, None); // threads(1): residual sort runs serially + let golden = records_bytes(&serial); - let out_on = dir.path().join("stats_on.bam"); - let out_off = dir.path().join("stats_off.bam"); - let stats_on = run(&out_on, true); - run(&out_off, false); - - // A real spilling merge actually ran, so the merge diagnostics were reached. - assert!( - stats_on.runs_written >= 5, - "expected a spilling merge to exercise the diagnostics, got {} runs", - stats_on.runs_written - ); - - // Diagnostics-only: the flag must not change what the sort produces. - let expected = (num_pairs * 2) as u64; - assert_eq!(count_bam_records(&out_on), expected, "records lost with --sort-stats on"); + let parallel = dir.path().join("parallel.bam"); + mk(¶llel, Some(4)); // sort_threads(4): residual sort now parallel assert_eq!( - count_bam_records(&out_off), - expected, - "sort_stats must not change the sorted output" + records_bytes(¶llel), + golden, + "order {order:?}: threads(1).sort_threads(4) must match the serial sort output" ); } - /// `log_consumer_stalls` renders a distinct `Shape:` line for each - /// [`crate::merge_stalls::StallShape`] the classifier can return. A real merge - /// lands on exactly one shape, and which one is timing-dependent, so driving - /// the emitter through every arm here both pins the `classify_stall` -> shape - /// mapping (the real behavior asserted) and exercises each `stat!(sort_stats, - /// ...)` arm run-scoped. Info level is raised so the emitter formats its - /// arguments (the no-op logger discards the record). - #[rstest::rstest] - #[case::not_stalled(0.1, 0.0, 0.5, 0.0, 0.0, 0.0, crate::merge_stalls::StallShape::NotStalled)] - #[case::contended(5.0, 0.4, 0.5, 0.0, 0.0, 0.0, crate::merge_stalls::StallShape::Contended)] - #[case::pool_saturated( - 5.0, - 0.0, - 0.9, - 0.0, - 0.0, - 0.0, - crate::merge_stalls::StallShape::PoolSaturated - )] - #[case::head_of_line(5.0, 0.0, 0.5, 0.6, 0.0, 0.0, crate::merge_stalls::StallShape::HeadOfLine)] - #[case::work_unclaimed( - 5.0, - 0.0, - 0.5, - 0.0, - 0.6, - 0.0, - crate::merge_stalls::StallShape::WorkUnclaimed - )] - #[case::decompress_latency( - 5.0, - 0.0, - 0.5, - 0.0, - 0.0, - 0.6, - crate::merge_stalls::StallShape::DecompressLatency - )] - #[case::mixed(5.0, 0.0, 0.5, 0.0, 0.0, 0.0, crate::merge_stalls::StallShape::Mixed)] - fn test_log_consumer_stalls_renders_every_shape( - #[case] park_secs: f64, - #[case] contended_share: f64, - #[case] utilization: f64, - #[case] starved: f64, - #[case] raw_queued: f64, - #[case] decompressing: f64, - #[case] expected: crate::merge_stalls::StallShape, - ) { - use crate::merge_stalls::{AwaitedShares, ConsumerStallReport, classify_stall}; - - log::set_max_level(log::LevelFilter::Info); - let loop_total = 10.0; - let awaited = AwaitedShares { decompressing, raw_queued, starved, ..Default::default() }; - let report = ConsumerStallReport { - block_pulls: 10, - stalled_pulls: 5, - parks: 8, - park_secs, - capped_share: 0.1, - starved_share: 0.1, - contended_share, - awaited, - censuses: 4, // > 0 so the per-file census rows render too - top_source: 1, - top_source_park_secs: 2.0, - sources_parked_on: 3, - }; + #[test] + fn test_raw_sorter_memory_limit() { + let sorter = RawExternalSorter::new(SortOrder::Coordinate).memory_limit(256 * 1024 * 1024); + assert_eq!(sorter.memory_limit, 256 * 1024 * 1024); + } - // The classifier maps these inputs to the expected shape ... - assert_eq!( - classify_stall(park_secs / loop_total, utilization, contended_share, awaited), - expected - ); - // ... and the emitter renders that arm run-scoped, without panicking. - RawExternalSorter::log_consumer_stalls(loop_total, utilization, report, true); + #[test] + fn test_raw_sorter_temp_compression() { + let sorter = RawExternalSorter::new(SortOrder::Coordinate).temp_compression(0); + assert_eq!(sorter.temp_compression, 0); } /// `RawExternalSorter::sort` rejects `temp_compression=0` + `SpillCodec::Zstd` @@ -8418,76 +7869,6 @@ mod tests { aux } - /// Build `MC:Z:` aux tag bytes. - fn mc_aux(value: &[u8]) -> Vec { - let mut aux = Vec::new(); - aux.extend_from_slice(b"MCZ"); - aux.extend_from_slice(value); - aux.push(0); // null terminator - aux - } - - /// Overwrite the mate position of a record built by `build_mapped_bam`, - /// which otherwise sets it equal to the record's own position. - fn with_mate_pos(mut bam: Vec, mate_pos: i32) -> Vec { - bam[24..28].copy_from_slice(&mate_pos.to_le_bytes()); - bam - } - - /// The mate lane resolves through `MC`, and it lands where a record whose - /// mate is already at the unclipped position lands. - /// - /// The template-coordinate key is output-identity-critical against - /// `samtools sort`, and the mate lane is the one part of it that comes from - /// parsing a tag rather than from a fixed field offset. Asserting the two - /// keys are equal pins the whole lane -- packing included -- rather than - /// just the parser, which `cigar.rs` already covers. - #[test] - fn test_extract_template_key_mate_lane_comes_from_mc() { - let header = Header::builder().build(); - let lib_lookup = LibraryLookup::from_header(&header); - - // Mate at 100 with 10 leading soft clips: unclipped 5' is 90. - let with_mc = with_mate_pos(build_mapped_bam(0, 50, b"read1", &mc_aux(b"10S40M")), 100); - // The same record whose mate is already reported at 90, and no MC to parse. - let without_mc = with_mate_pos(build_mapped_bam(0, 50, b"read1", &[]), 90); - - let keyed = extract_template_key_inline(&with_mc, &lib_lookup, None, &test_cb_hasher()); - let expected = - extract_template_key_inline(&without_mc, &lib_lookup, None, &test_cb_hasher()); - assert_eq!(keyed, expected, "MC-derived mate lane must equal the unclipped position"); - } - - /// A non-UTF-8 `MC` reaches the parser and its valid prefix still sets the - /// mate lane. - /// - /// Extraction hands `MC` over as raw bytes rather than validating it as - /// UTF-8 first, so this value is parsed where it was previously discarded - /// (leaving the mate lane at the raw mate position). This pins the change at - /// the key level, which is the level the sort order is defined at. - #[test] - fn test_extract_template_key_mate_lane_parses_a_non_utf8_mc_prefix() { - let header = Header::builder().build(); - let lib_lookup = LibraryLookup::from_header(&header); - - let with_bad_mc = - with_mate_pos(build_mapped_bam(0, 50, b"read1", &mc_aux(b"10S40M\xff")), 100); - let unclipped = with_mate_pos(build_mapped_bam(0, 50, b"read1", &[]), 90); - let raw_mate = with_mate_pos(build_mapped_bam(0, 50, b"read1", &[]), 100); - - let keyed = extract_template_key_inline(&with_bad_mc, &lib_lookup, None, &test_cb_hasher()); - assert_eq!( - keyed, - extract_template_key_inline(&unclipped, &lib_lookup, None, &test_cb_hasher()), - "the valid CIGAR prefix must still be applied" - ); - assert_ne!( - keyed, - extract_template_key_inline(&raw_mate, &lib_lookup, None, &test_cb_hasher()), - "discarding the tag would leave the mate lane at the raw mate position" - ); - } - #[test] fn test_extract_template_key_cb_present_has_nonzero_hash() { let header = Header::builder().build(); @@ -10006,6 +9387,31 @@ mod tests { assert_eq!(count, 40); assert_eq!(count_bam_records(&merged), 40); + + // Verify template-coordinate order/identity via an independent oracle + // (mirrors `narrow_sort_output_passes_full_width_verify`): re-extract the + // template key from every merged record and assert non-decreasing core + // order. A count-only check would pass even if the merge reordered or + // duplicated records. + let (_, hdr) = create_raw_bam_reader(&merged, 1).expect("header"); + let lib = LibraryLookup::from_header(&hdr); + let hasher = cb_hasher(); + let file = std::fs::File::open(&merged).expect("open merged"); + let mut raw_reader = crate::reader::RawBamRecordReader::new(file).expect("reader"); + raw_reader.skip_header().expect("skip header"); + let (total, violations, first) = crate::verify::verify_sort_order( + raw_reader, + |bam| extract_template_key_inline(bam, &lib, None, &hasher), + |cur: &TemplateKey40, prev: &TemplateKey40| { + cur.core_cmp(prev) == std::cmp::Ordering::Less + }, + ) + .expect("verify runs"); + assert_eq!(total, 40, "verify must see every merged record"); + assert_eq!( + violations, 0, + "template-coordinate order violated after merge (first={first:?}, total={total})" + ); } #[test] @@ -10160,7 +9566,7 @@ mod tests { #[test] fn test_sort_phase_timer_all_methods() { - let mut timer = SortPhaseTimer::new(false); + let mut timer = SortPhaseTimer::new(); assert!(timer.overall_start.is_some()); assert!(timer.read_span_start.is_some()); @@ -10224,7 +9630,7 @@ mod tests { // log_summary must not panic (output goes to log sink). `consolidate_count` // is 1 here, so the consolidation branch is exercised too. - timer.log_summary(4, 4, 64, &Phase1FloorInputs::default()); + timer.log_summary(4, 4, 64); } // ======================================================================== @@ -11647,3 +11053,760 @@ mod tests { ); } } + +#[cfg(test)] +mod from_slots_merge_tests { + //! Direct unit tests for the slot-backed `MergeDriver::from_slots` merge + //! driver: cross-block record parsing (incl. records spanning many blocks), + //! non-blocking stall/resume on empty-non-EOF slots, the embedded and + //! non-embedded sort-key arms, and memory-chunk mixing. Ported verbatim + //! from the issue-#330 source branch (the slimmed `from_slots` API is + //! identical at this layer). These exercise the parser/state-machine paths + //! that the end-to-end `three_step_chain_*` tests do not isolate. + use super::*; + use std::io::{BufReader, Read, Write}; + use std::sync::Arc as StdArc; + + #[derive(Default, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Debug)] + struct TestKey(u64); + + impl RawSortKey for TestKey { + const SERIALIZED_SIZE: Option = Some(8); + const EMBEDDED_IN_RECORD: bool = false; + fn extract(_bam: &[u8], _ctx: &crate::keys::SortContext) -> Self { + unimplemented!( + "TestKey is only used in MergeDriver::from_slots tests, not for ingestion" + ) + } + fn write_to(&self, w: &mut W) -> std::io::Result<()> { + w.write_all(&self.0.to_le_bytes()) + } + fn read_from(r: &mut R) -> std::io::Result { + let mut buf = [0u8; 8]; + r.read_exact(&mut buf)?; + Ok(TestKey(u64::from_le_bytes(buf))) + } + } + + /// Serialize a sequence of `(TestKey, record_bytes)` pairs in the spill + /// file's wire format: `[key(8)][len(4)][record(len)]` per entry. + fn serialize_records(records: &[(TestKey, Vec)]) -> Vec { + let mut out = Vec::new(); + for (key, rec) in records { + key.write_to(&mut out).unwrap(); + #[allow(clippy::cast_possible_truncation)] + let len = rec.len() as u32; + out.write_all(&len.to_le_bytes()).unwrap(); + out.write_all(rec).unwrap(); + } + out + } + + /// Build a populated `SortMergeSlot` with `file_id` whose decompressed + /// queue contains `records` partitioned across `n_blocks` "blocks" + /// (so the slot parser exercises cross-block reads). Sets `queue_eof` + /// so the slot reports drained once consumer empties it. + fn populated_slot( + file_id: u32, + records: &[(TestKey, Vec)], + n_blocks: usize, + ) -> StdArc { + assert!(n_blocks >= 1, "n_blocks must be >= 1"); + let bytes = serialize_records(records); + let block_size = bytes.len().div_ceil(n_blocks); + let slot = StdArc::new(SortMergeSlot::new( + file_id, + BufReader::new(tempfile::tempfile().expect("tempfile")), + crate::codec::SpillCodec::Bgzf, + )); + { + let mut dec = slot.decompressed.lock().unwrap(); + for chunk in bytes.chunks(block_size) { + dec.push_back(chunk.to_vec()); + } + slot.queue_eof.store(true, std::sync::atomic::Ordering::Release); + } + slot + } + + /// Drive a `MergeDriver` to exhaustion via the non-blocking `try_step`, + /// returning the emitted record bytes in emission order. + /// + /// Tests validate ordering by inspecting the human-readable record bytes + /// (e.g. `b"AAA-1"`, `b"BBB-2"`) — the keys themselves are consumed by + /// the parser before the record is exposed, so we have no direct access to + /// them at this layer. These slots are pre-populated and EOF-marked before + /// the merge, so `try_step` never returns `Stalled` here. + fn drain_merge_driver_bytes( + mut driver: MergeDriver, + ) -> Vec> { + let mut out = Vec::new(); + loop { + match driver.try_step().expect("try_step") { + MergeStep::Produced(bytes) => out.push(bytes.to_vec()), + MergeStep::Done => break, + MergeStep::Stalled => panic!("unexpected Stalled on pre-populated EOF slots"), + } + } + out + } + + #[test] + fn from_slots_merges_three_pre_populated_slots() { + // Three slots, each with sorted records by TestKey. The merged output + // should be globally sorted. + let file0: Vec<(TestKey, Vec)> = vec![ + (TestKey(1), b"AAA-1".to_vec()), + (TestKey(4), b"AAA-4".to_vec()), + (TestKey(7), b"AAA-7".to_vec()), + ]; + let file1: Vec<(TestKey, Vec)> = vec![ + (TestKey(2), b"BBB-2".to_vec()), + (TestKey(5), b"BBB-5".to_vec()), + (TestKey(8), b"BBB-8".to_vec()), + ]; + let file2: Vec<(TestKey, Vec)> = vec![ + (TestKey(3), b"CCC-3".to_vec()), + (TestKey(6), b"CCC-6".to_vec()), + (TestKey(9), b"CCC-9".to_vec()), + ]; + + let slots = vec![ + populated_slot(0, &file0, 1), + populated_slot(1, &file1, 1), + populated_slot(2, &file2, 1), + ]; + + let driver = MergeDriver::::from_slots(slots, MemorySources::Owned(Vec::new()), 9); + let emitted_bytes = drain_merge_driver_bytes(driver); + let expected: Vec> = vec![ + b"AAA-1".to_vec(), + b"BBB-2".to_vec(), + b"CCC-3".to_vec(), + b"AAA-4".to_vec(), + b"BBB-5".to_vec(), + b"CCC-6".to_vec(), + b"AAA-7".to_vec(), + b"BBB-8".to_vec(), + b"CCC-9".to_vec(), + ]; + assert_eq!(emitted_bytes, expected, "merged sequence not globally sorted by key"); + } + + #[test] + fn from_slots_handles_records_spanning_block_boundaries() { + // Same data, but each slot's bytes are split across 4 "decompressed + // blocks" so the slot parser's read_exact + advance_to_next_block + // path is exercised. The key+len header itself may straddle a block + // boundary depending on byte counts. + let records: Vec<(TestKey, Vec)> = (0u64..20) + .map(|i| (TestKey(i), format!("record-{i:02}-with-some-padding").into_bytes())) + .collect(); + + let slots = vec![populated_slot(0, &records, 4)]; + let driver = + MergeDriver::::from_slots(slots, MemorySources::Owned(Vec::new()), 20); + let emitted_bytes = drain_merge_driver_bytes(driver); + let expected: Vec> = records.into_iter().map(|(_, b)| b).collect(); + assert_eq!(emitted_bytes, expected, "cross-block parse corrupted record sequence"); + } + + #[test] + fn from_slots_empty_sources_drain_to_zero_records() { + let empty_slot = StdArc::new(SortMergeSlot::new( + 0, + BufReader::new(tempfile::tempfile().expect("tempfile")), + crate::codec::SpillCodec::Bgzf, + )); + // Mark drained without inserting any blocks. + empty_slot.queue_eof.store(true, std::sync::atomic::Ordering::Release); + // Lazy priming: `from_slots` always yields a driver; the empty + // determination happens on the first `try_step`, which returns `Done`. + let driver = MergeDriver::::from_slots( + vec![empty_slot], + MemorySources::Owned(Vec::new()), + 0, + ); + let emitted = drain_merge_driver_bytes(driver); + assert!(emitted.is_empty(), "all-empty merge must emit zero records"); + } + + /// The length prefix is untrusted: it comes straight off a decompressed + /// spill block. The slow path used to hand it to `out.reserve(len)` with no + /// check, so a corrupt slot could request up to 4 GiB — and `reserve` + /// answers an allocation failure by aborting the process, not by returning + /// the truncation error the next loop iteration would have produced. + /// + /// The length here is over the bound but still allocatable, so this test + /// discriminates on the *error message* rather than on surviving: before + /// the bound existed it reserved happily and then failed as a truncated + /// body, which is the wrong diagnosis for a corrupt length. + #[test] + fn implausible_record_length_is_rejected_before_reserving() { + let bogus_len = u32::try_from(crate::inline::SORT_SEGMENT_SIZE).unwrap() + 1; + let mut bytes = Vec::new(); + TestKey(1).write_to(&mut bytes).unwrap(); + bytes.write_all(&bogus_len.to_le_bytes()).unwrap(); + bytes.write_all(b"only-a-few-body-bytes").unwrap(); + + let slot = StdArc::new(SortMergeSlot::new( + 5, + BufReader::new(tempfile::tempfile().expect("tempfile")), + crate::codec::SpillCodec::Bgzf, + )); + slot.decompressed.lock().unwrap().push_back(bytes); + slot.queue_eof.store(true, std::sync::atomic::Ordering::Release); + + let mut driver = + MergeDriver::::from_slots(vec![slot], MemorySources::Owned(Vec::new()), 1); + let err = driver.try_step().expect_err("an implausible length must surface as Err"); + assert!( + err.to_string().contains("implausible record length") && err.to_string().contains('5'), + "the error must name the bad length and the slot; got: {err}", + ); + } + + /// An embedded key is read out of the record body itself, so an empty body + /// cannot carry one. `RawCoordinateKey::extract_from_record` reads `ref_id` + /// off the front, which on an empty body indexes out of bounds — a panic in + /// the middle of a merge rather than a corrupt-slot error. A zero length is + /// reachable from a corrupt length prefix: the body loop sees + /// `out.len() >= 0` and breaks immediately with `out` still empty. + #[test] + fn embedded_key_on_an_empty_body_is_an_error_not_a_panic() { + let err = slot_parse_key::(&[], &[]) + .expect_err("an empty body cannot carry an embedded key"); + assert!( + err.to_string().contains("empty and cannot carry an embedded sort key"), + "the error must name the empty body; got: {err}", + ); + } + + /// Truncation at EOF must be an error, not a stall — at **every** stage of + /// the framer, not just the last one. + /// + /// Every other test either completes its record or stalls on a non-EOF slot, + /// so the three `BlockLoad::Eof` arms in `slot_collect_pending` were + /// unexercised. They are siblings: swapping any one for `WouldBlock` hangs + /// the pipeline forever on a truncated spill while the suite stays green, so + /// covering only the body arm left two of the three live. + /// + /// A serialized record is `[key(8)][len(4)][body(len)]`, so keeping 4 and 10 + /// bytes lands mid-key and mid-length; `None` means "all but the last 5", + /// which lands mid-body. + #[rstest::rstest] + #[case::mid_key(Some(4), "truncated record key in slot 3")] + #[case::mid_length(Some(10), "truncated record length in slot 3")] + #[case::mid_body(None, "truncated record body in slot 3")] + fn truncation_at_eof_is_an_error_not_a_stall_at_every_stage( + #[case] keep: Option, + #[case] expected: &str, + ) { + let serialized = serialize_records(&[(TestKey(1), b"a-body-that-is-cut-short".to_vec())]); + let keep = keep.unwrap_or(serialized.len() - 5); + assert!(keep < serialized.len(), "the case must actually truncate"); + + let slot = StdArc::new(SortMergeSlot::new( + 3, + BufReader::new(tempfile::tempfile().expect("tempfile")), + crate::codec::SpillCodec::Bgzf, + )); + slot.decompressed.lock().unwrap().push_back(serialized[..keep].to_vec()); + slot.queue_eof.store(true, std::sync::atomic::Ordering::Release); + + let mut driver = + MergeDriver::::from_slots(vec![slot], MemorySources::Owned(Vec::new()), 1); + let err = driver.try_step().expect_err("truncation must surface as Err"); + assert!(err.to_string().contains(expected), "expected `{expected}`, got: {err}"); + } + + /// `from_slots` sorts slots by `file_id` so equal-key ties resolve the same + /// way the legacy chunk-files path resolves them — the `LoserTree` breaks + /// ties by leaf index, so leaf order must follow `file_id` and not the + /// caller's `Vec` order. Nothing pinned that, so the sort could vanish in a + /// later cleanup and the divergence would be invisible except on inputs + /// that contain ties. Pass the slots reversed and require `file_id` to win. + #[test] + fn key_ties_break_by_file_id_not_by_caller_push_order() { + let file0 = populated_slot(0, &[(TestKey(1), b"from-file-0".to_vec())], 1); + let file1 = populated_slot(1, &[(TestKey(1), b"from-file-1".to_vec())], 1); + let driver = MergeDriver::::from_slots( + vec![file1, file0], // reversed on purpose + MemorySources::Owned(Vec::new()), + 2, + ); + assert_eq!( + drain_merge_driver_bytes(driver), + vec![b"from-file-0".to_vec(), b"from-file-1".to_vec()], + "equal-key ties must resolve by file_id, not by caller push order", + ); + } + + /// The core BUG #3 regression: a slot whose decompressed queue is empty + /// and `!queue_eof` must make `try_step` return `Stalled` (yield), NOT + /// block the caller. The old blocking consumer (`block_ready.wait()`) + /// would park here forever at a single worker. After the producer fills + /// the slot and marks EOF, the merge resumes and drains to completion. + #[test] + fn from_slots_try_step_stalls_on_empty_non_eof_slot_then_resumes() { + let slot = StdArc::new(SortMergeSlot::new( + 0, + BufReader::new(tempfile::tempfile().expect("tempfile")), + crate::codec::SpillCodec::Bgzf, + )); + // Empty queue, NOT eof — the producer is "still feeding". + let mut driver = MergeDriver::::from_slots( + vec![StdArc::clone(&slot)], + MemorySources::Owned(Vec::new()), + 0, + ); + + // Priming cannot read the slot's first record → Stalled, not a hang. + assert!( + matches!(driver.try_step().expect("try_step"), MergeStep::Stalled), + "empty non-eof slot must Stall, not block" + ); + // Repeated calls keep stalling (idempotent yield). + assert!(matches!(driver.try_step().expect("try_step"), MergeStep::Stalled)); + + // Producer fills the slot and marks EOF. + { + let bytes = serialize_records(&[(TestKey(1), b"only-1".to_vec())]); + let mut dec = slot.decompressed.lock().unwrap(); + dec.push_back(bytes); + slot.queue_eof.store(true, std::sync::atomic::Ordering::Release); + } + + // The merge now resumes: emits the record, then is Done. + match driver.try_step().expect("try_step") { + MergeStep::Produced(bytes) => assert_eq!(bytes, b"only-1"), + MergeStep::Stalled => panic!("should have resumed, not Stalled"), + MergeStep::Done => panic!("should have produced the record, not Done"), + } + assert!(matches!(driver.try_step().expect("try_step"), MergeStep::Done)); + } + + /// `try_step` must also stall (not block) when a slot drains to empty + /// **mid-merge** (after priming) while still `!queue_eof`, then resume + /// once the producer pushes the next block and marks EOF. + #[test] + fn from_slots_try_step_stalls_mid_merge_then_resumes() { + let slot = StdArc::new(SortMergeSlot::new( + 0, + BufReader::new(tempfile::tempfile().expect("tempfile")), + crate::codec::SpillCodec::Bgzf, + )); + // One record present, NOT eof (more "coming"). + { + let bytes = serialize_records(&[(TestKey(1), b"rec-A".to_vec())]); + slot.decompressed.lock().unwrap().push_back(bytes); + } + let mut driver = MergeDriver::::from_slots( + vec![StdArc::clone(&slot)], + MemorySources::Owned(Vec::new()), + 0, + ); + + // Prime + emit the first record. + match driver.try_step().expect("try_step") { + MergeStep::Produced(bytes) => assert_eq!(bytes, b"rec-A"), + other => panic!("expected Produced(rec-A), got {other:?}"), + } + // Deferred refill finds the slot drained + not eof → Stalled. + assert!( + matches!(driver.try_step().expect("try_step"), MergeStep::Stalled), + "drained mid-merge non-eof slot must Stall" + ); + + // Producer pushes the next record and marks EOF. + { + let bytes = serialize_records(&[(TestKey(2), b"rec-B".to_vec())]); + let mut dec = slot.decompressed.lock().unwrap(); + dec.push_back(bytes); + slot.queue_eof.store(true, std::sync::atomic::Ordering::Release); + } + + match driver.try_step().expect("try_step") { + MergeStep::Produced(bytes) => assert_eq!(bytes, b"rec-B"), + other => panic!("expected Produced(rec-B), got {other:?}"), + } + assert!(matches!(driver.try_step().expect("try_step"), MergeStep::Done)); + } + + /// Resumable-framer regression: feed a single **non-embedded** record one + /// byte per block, asserting `try_step` `Stalled`s (retaining + /// `parser.pending`) until the whole record is available, then reassembles + /// the exact bytes. With 1-byte blocks the framer takes the slow path and + /// resumes a `WouldBlock` across the key prefix, the length prefix, AND the + /// body — the partial-record path no other test covers byte-for-byte. + #[test] + fn from_slots_try_step_resumes_record_fed_one_byte_at_a_time() { + let rec_body = b"multi-stage-resumable-record-body".to_vec(); + let serialized = serialize_records(&[(TestKey(7), rec_body.clone())]); + assert!(serialized.len() > 12, "must span the 8-byte key + 4-byte len header"); + + let slot = StdArc::new(SortMergeSlot::new( + 0, + BufReader::new(tempfile::tempfile().expect("tempfile")), + crate::codec::SpillCodec::Bgzf, + )); + slot.decompressed.lock().unwrap().push_back(vec![serialized[0]]); + + let mut driver = MergeDriver::::from_slots( + vec![StdArc::clone(&slot)], + MemorySources::Owned(Vec::new()), + 1, + ); + + // Each step before the last byte consumes one byte into the pending + // record and stalls (queue empty, not EOF). + for (i, &b) in serialized.iter().enumerate().skip(1) { + assert!( + matches!(driver.try_step().expect("try_step"), MergeStep::Stalled), + "expected Stalled with only {i} of {} bytes available", + serialized.len(), + ); + slot.decompressed.lock().unwrap().push_back(vec![b]); + } + + // The step that consumes the final byte completes the record. + match driver.try_step().expect("try_step") { + MergeStep::Produced(bytes) => assert_eq!(bytes, rec_body.as_slice()), + other => panic!("expected Produced(reassembled), got {other:?}"), + } + // No more bytes and not yet EOF → Stalled; EOF → Done. + assert!(matches!(driver.try_step().expect("try_step"), MergeStep::Stalled)); + slot.queue_eof.store(true, std::sync::atomic::Ordering::Release); + assert!(matches!(driver.try_step().expect("try_step"), MergeStep::Done)); + } + + /// Same resumable-framer property for an **embedded**-key record split mid + /// body across two blocks (the key lives in the body, so this exercises the + /// embedded `extract_from_record` path resuming after a `WouldBlock`). + #[test] + fn from_slots_try_step_resumes_embedded_record_split_mid_body() { + // Embedded record: [len:4][body], key = first 8 body bytes. + let record = embedded_record(123, b"-embedded-tail-bytes"); + let serialized = serialize_embedded(std::slice::from_ref(&record)); + // Cut mid-body: 4-byte len + 10 body bytes in the first block. + let split = 4 + 10; + assert!(split < serialized.len() && split > 4 + 8, "split mid-body, past the key"); + let (head, tail) = serialized.split_at(split); + + let slot = StdArc::new(SortMergeSlot::new( + 0, + BufReader::new(tempfile::tempfile().expect("tempfile")), + crate::codec::SpillCodec::Bgzf, + )); + slot.decompressed.lock().unwrap().push_back(head.to_vec()); + + let mut driver = MergeDriver::::from_slots( + vec![StdArc::clone(&slot)], + MemorySources::Owned(Vec::new()), + 1, + ); + + // Body incomplete in the first block → Stalled mid-record. + assert!(matches!(driver.try_step().expect("try_step"), MergeStep::Stalled)); + + // Deliver the rest + EOF; the record reassembles exactly. + { + let mut dec = slot.decompressed.lock().unwrap(); + dec.push_back(tail.to_vec()); + slot.queue_eof.store(true, std::sync::atomic::Ordering::Release); + } + match driver.try_step().expect("try_step") { + MergeStep::Produced(bytes) => assert_eq!(bytes, record.as_slice()), + other => panic!("expected Produced(reassembled embedded), got {other:?}"), + } + assert!(matches!(driver.try_step().expect("try_step"), MergeStep::Done)); + } + + #[test] + fn from_slots_merges_slots_with_memory_chunks() { + let slot_records: Vec<(TestKey, Vec)> = vec![ + (TestKey(1), b"slot-1".to_vec()), + (TestKey(3), b"slot-3".to_vec()), + (TestKey(5), b"slot-5".to_vec()), + ]; + let memory_records: Vec<(TestKey, fgumi_raw_bam::RawRecord)> = vec![ + (TestKey(2), fgumi_raw_bam::RawRecord::from(b"mem-2".to_vec())), + (TestKey(4), fgumi_raw_bam::RawRecord::from(b"mem-4".to_vec())), + (TestKey(6), fgumi_raw_bam::RawRecord::from(b"mem-6".to_vec())), + ]; + + let slots = vec![populated_slot(0, &slot_records, 2)]; + let driver = MergeDriver::::from_slots( + slots, + MemorySources::Owned(vec![memory_records]), + 6, + ); + let emitted_bytes = drain_merge_driver_bytes(driver); + let expected: Vec> = vec![ + b"slot-1".to_vec(), + b"mem-2".to_vec(), + b"slot-3".to_vec(), + b"mem-4".to_vec(), + b"slot-5".to_vec(), + b"mem-6".to_vec(), + ]; + assert_eq!(emitted_bytes, expected, "slot+memory merge not globally sorted"); + } + + /// The same merge through `MemorySources::Shared`, which is the arm + /// production actually uses: `sort_coordinate_optimized` and + /// `sort_template_coordinate_impl` both build `Shared`, and `Owned` is the + /// queryname arm only. Every other `from_slots` test drove `Owned`, so the + /// arm carrying two of the three sort orders had no coverage here at all. + /// + /// The two arms differ in mechanism, not just in type: `Owned` swaps the + /// record vec out, while `Shared` copies `chunk.record_bytes(idx)` into the + /// caller's buffer and takes the key by index. A defect in that copy, or in + /// the `idx`/`len` pairing, is invisible to the `Owned` tests. + #[test] + fn from_slots_merges_slots_with_shared_memory_chunks() { + let slot_records: Vec<(TestKey, Vec)> = vec![ + (TestKey(1), b"slot-1".to_vec()), + (TestKey(3), b"slot-3".to_vec()), + (TestKey(5), b"slot-5".to_vec()), + ]; + let memory_chunk = crate::inline::InMemoryChunk::from_owned_records(vec![ + (TestKey(2), b"mem-2".to_vec()), + (TestKey(4), b"mem-4".to_vec()), + (TestKey(6), b"mem-6".to_vec()), + ]); + + let slots = vec![populated_slot(0, &slot_records, 2)]; + let driver = + MergeDriver::::from_slots(slots, MemorySources::Shared(vec![memory_chunk]), 6); + let expected: Vec> = vec![ + b"slot-1".to_vec(), + b"mem-2".to_vec(), + b"slot-3".to_vec(), + b"mem-4".to_vec(), + b"slot-5".to_vec(), + b"mem-6".to_vec(), + ]; + assert_eq!( + drain_merge_driver_bytes(driver), + expected, + "slot + shared-memory merge not globally sorted", + ); + } + + // (v4: `from_slots_bails_on_unfilled_gap_in_decompressed` removed. + // v3.1's `decompressed: Mutex>>` is now a + // `Mutex>>` (FIFO). There is no concept of + // "ordinals" or "gaps" — the producer pushes in order under the + // reader lock, so a missing-block-in-the-middle state is no + // longer representable.) + + // ------------------------------------------------------------------------ + // Coverage for the EMBEDDED_IN_RECORD = true parser arm. + // + // All three production sort keys (RawCoordinateKey, RawQuerynameKey, + // TemplateKey) set EMBEDDED_IN_RECORD = true, but TestKey above uses + // the non-embedded format. This test pins the embedded-format slot path + // with a synthetic embedded key (TestEmbeddedKey) whose value lives at + // bytes 0..8 of the record itself. + + /// Embedded-key test type: serialized format is `[len(4)][record(len)]` + /// (no separate key prefix). The key is the first 8 LE bytes of the record. + #[derive(Default, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Debug)] + struct TestEmbeddedKey(u64); + + impl RawSortKey for TestEmbeddedKey { + const SERIALIZED_SIZE: Option = Some(0); + const EMBEDDED_IN_RECORD: bool = true; + fn extract(_bam: &[u8], _ctx: &crate::keys::SortContext) -> Self { + unimplemented!("TestEmbeddedKey is only used in MergeDriver::from_slots tests") + } + fn extract_from_record(bam: &[u8]) -> Self { + assert!(bam.len() >= 8, "embedded-key test record must have ≥8 bytes"); + let mut buf = [0u8; 8]; + buf.copy_from_slice(&bam[..8]); + TestEmbeddedKey(u64::from_le_bytes(buf)) + } + fn write_to(&self, _w: &mut W) -> std::io::Result<()> { + // No-op: embedded keys are not written separately; they live + // inside the record bytes. + Ok(()) + } + fn read_from(_r: &mut R) -> std::io::Result { + unreachable!( + "EMBEDDED_IN_RECORD = true keys take the slot_try_next_record \ + embedded arm, which never calls read_from" + ) + } + } + + /// Build a record whose first 8 bytes encode the embedded sort key. + fn embedded_record(key: u64, tail: &[u8]) -> Vec { + let mut v = Vec::with_capacity(8 + tail.len()); + v.extend_from_slice(&key.to_le_bytes()); + v.extend_from_slice(tail); + v + } + + /// Serialize records in the embedded-key spill format: `[len(4)][record(len)]`. + fn serialize_embedded(records: &[Vec]) -> Vec { + let mut out = Vec::new(); + for rec in records { + #[allow(clippy::cast_possible_truncation)] + let len = rec.len() as u32; + out.write_all(&len.to_le_bytes()).unwrap(); + out.write_all(rec).unwrap(); + } + out + } + + fn embedded_populated_slot( + file_id: u32, + records: &[Vec], + n_blocks: usize, + ) -> StdArc { + let bytes = serialize_embedded(records); + let block_size = bytes.len().div_ceil(n_blocks.max(1)); + let slot = StdArc::new(SortMergeSlot::new( + file_id, + BufReader::new(tempfile::tempfile().expect("tempfile")), + crate::codec::SpillCodec::Bgzf, + )); + { + let mut dec = slot.decompressed.lock().unwrap(); + for chunk in bytes.chunks(block_size) { + dec.push_back(chunk.to_vec()); + } + slot.queue_eof.store(true, std::sync::atomic::Ordering::Release); + } + slot + } + + #[test] + fn from_slots_embedded_key_path_merges_correctly() { + // Three slots, each with sorted records by embedded key. The merged + // output must be globally sorted. Exercises `slot_try_next_record`'s + // EMBEDDED_IN_RECORD arm + the `extract_from_record` key-from-bytes + // path. + let file0 = + vec![embedded_record(1, b"-A"), embedded_record(4, b"-A"), embedded_record(7, b"-A")]; + let file1 = + vec![embedded_record(2, b"-B"), embedded_record(5, b"-B"), embedded_record(8, b"-B")]; + let file2 = + vec![embedded_record(3, b"-C"), embedded_record(6, b"-C"), embedded_record(9, b"-C")]; + let slots = vec![ + embedded_populated_slot(0, &file0, 1), + embedded_populated_slot(1, &file1, 1), + embedded_populated_slot(2, &file2, 1), + ]; + + let driver = + MergeDriver::::from_slots(slots, MemorySources::Owned(Vec::new()), 9); + let emitted = drain_merge_driver_bytes(driver); + + // Verify keys are emitted in sorted order. Tail bytes prove the + // identity of each source. + let expected = vec![ + embedded_record(1, b"-A"), + embedded_record(2, b"-B"), + embedded_record(3, b"-C"), + embedded_record(4, b"-A"), + embedded_record(5, b"-B"), + embedded_record(6, b"-C"), + embedded_record(7, b"-A"), + embedded_record(8, b"-B"), + embedded_record(9, b"-C"), + ]; + assert_eq!(emitted, expected); + } + + #[test] + fn from_slots_skips_empty_slots_mixed_with_populated() { + // Verify a slot whose decompressed queue is empty (and reader marked EOF) + // is correctly elided from the LoserTree. Without correct handling + // the priming step would either panic or produce a phantom source. + let populated = vec![ + (TestKey(1), b"one".to_vec()), + (TestKey(2), b"two".to_vec()), + (TestKey(3), b"three".to_vec()), + ]; + + let empty_slot = StdArc::new(SortMergeSlot::new( + 99, + BufReader::new(tempfile::tempfile().expect("tempfile")), + crate::codec::SpillCodec::Bgzf, + )); + empty_slot.queue_eof.store(true, std::sync::atomic::Ordering::Release); + // Insert empty-slot first, then a populated slot. The driver should + // skip the empty one during priming and only emit records from the + // populated slot. + let slots = vec![empty_slot, populated_slot(0, &populated, 1)]; + + let driver = MergeDriver::::from_slots(slots, MemorySources::Owned(Vec::new()), 3); + let emitted = drain_merge_driver_bytes(driver); + assert_eq!(emitted, vec![b"one".to_vec(), b"two".to_vec(), b"three".to_vec()]); + } + + /// `slot_try_load_block` checks `decomp_error` BEFORE `queue_eof`, so a slot + /// whose decompression failed surfaces `Err` rather than the clean EOF that + /// an empty queue plus `queue_eof` looks like. Nothing exercised the error + /// arm: every existing test drives `queue_eof`, so swapping the two checks + /// would silently convert a failed spill into a short read — records + /// dropped, exit status zero — with the whole suite green. + #[test] + fn decomp_error_surfaces_as_an_error_not_a_clean_eof() { + let slot = StdArc::new(SortMergeSlot::new( + 7, + BufReader::new(tempfile::tempfile().expect("tempfile")), + crate::codec::SpillCodec::Bgzf, + )); + // The shape a failed decompression leaves behind: queue drained, EOF + // set, and the error flag raised. Ordering matters — both are set, so + // only the check order decides the outcome. + slot.decomp_error.store(true, std::sync::atomic::Ordering::Release); + slot.queue_eof.store(true, std::sync::atomic::Ordering::Release); + + let mut parser = SlotParserState::new(); + let msg = match slot_try_load_block(&slot, &mut parser) { + Err(e) => e.to_string(), + Ok(_) => panic!("a slot with decomp_error must not report a clean EOF"), + }; + // The whole phrase, not a bare `7`: a lone digit matches any number that + // happens to appear in the message, which is a weaker claim than "this + // error named the slot that failed". + assert!( + msg.contains("spill decompression error on slot 7"), + "the error must name the failure and the slot; got: {msg}", + ); + } + + #[test] + fn from_slots_records_merged_equals_emitted_total() { + // `records_merged()` (surfaced by `SortMerge` at completion) must equal + // the number of records actually emitted across all sources. + let f0: Vec<(TestKey, Vec)> = + (0u64..15).map(|i| (TestKey(i * 2), format!("a-{i:02}").into_bytes())).collect(); + let f1: Vec<(TestKey, Vec)> = + (0u64..15).map(|i| (TestKey(i * 2 + 1), format!("b-{i:02}").into_bytes())).collect(); + let total = (f0.len() + f1.len()) as u64; + + let slots = vec![populated_slot(0, &f0, 3), populated_slot(1, &f1, 2)]; + let mut driver = + MergeDriver::::from_slots(slots, MemorySources::Owned(Vec::new()), total); + + let mut emitted = 0u64; + loop { + match driver.try_step().expect("try_step") { + MergeStep::Produced(_) => emitted += 1, + MergeStep::Done => break, + MergeStep::Stalled => panic!("unexpected Stalled on pre-populated EOF slots"), + } + } + assert_eq!(emitted, total, "should emit every input record"); + assert_eq!( + driver.records_merged(), + total, + "records_merged must equal the emitted record count" + ); + } +} diff --git a/crates/fgumi-sort/src/inline.rs b/crates/fgumi-sort/src/inline.rs index 03bd74d5a..7126807cb 100644 --- a/crates/fgumi-sort/src/inline.rs +++ b/crates/fgumi-sort/src/inline.rs @@ -14,6 +14,7 @@ #![allow(dead_code)] +use crate::arena_pool::PooledSegmentedBuf; use crate::keys::{RawCoordinateKey, RawSortKey, SortContext}; use crate::radix::bytes_needed_u64; use crate::segmented_buf::SegmentedBuf; @@ -44,7 +45,7 @@ use std::sync::Arc; /// In exchange, materialization is allocation-free and the peak /// memory of the sort drops by ~1× (the materialization no longer /// transiently doubles the buffer). -pub(crate) struct InMemoryChunk { +pub struct InMemoryChunk { /// Shared backing store for record bytes (the original sort /// buffer's `SegmentedBuf`). All sibling chunks from one /// `par_sort_into_chunks` call share this Arc, so the segments @@ -55,7 +56,11 @@ pub(crate) struct InMemoryChunk { /// merge is slightly higher, but the materialization-peak /// memory drops by ~1× (pre-PR transiently held both the buffer /// and a full copy in per-record `Vec`s). - data: Arc, + /// Wrapped in [`PooledSegmentedBuf`] so a chunk drained from a pooled + /// coordinate arena returns its storage to the + /// [`ArenaPool`](crate::arena_pool::ArenaPool) when the last `Arc` clone + /// drops. Non-pooled constructors use `PooledSegmentedBuf::unpooled`. + data: Arc, /// `(sort_key, byte_offset_in_data, len)` per record, in sorted /// order. `offset` is `u64` because a single `SegmentedBuf` can /// exceed 4 GiB at high `--max-memory` × `--threads`; `len` is @@ -67,39 +72,79 @@ impl InMemoryChunk { /// Construct an empty chunk backed by an empty shared buffer. #[must_use] pub(crate) fn empty() -> Self { - Self { data: Arc::new(SegmentedBuf::default()), records: Vec::new() } + Self { + data: Arc::new(PooledSegmentedBuf::unpooled(SegmentedBuf::default())), + records: Vec::new(), + } } /// Construct a chunk holding the given records, all referencing - /// the same shared data buffer. + /// the same shared (possibly pooled) data buffer. #[must_use] - pub(crate) fn from_parts(data: Arc, records: Vec<(K, u64, u32)>) -> Self { + pub(crate) fn from_parts(data: Arc, records: Vec<(K, u64, u32)>) -> Self { Self { data, records } } + /// Build a chunk from owned `(key, bytes)` records by packing the bytes into + /// a fresh `SegmentedBuf`. This **copies** every record, so it is for callers + /// that already hold owned records (tests, or future owned→shared bridging) — + /// the production sort path uses `RecordBuffer::drain_into_single_chunk`, + /// which moves the arena without copying. + /// + /// # Panics + /// + /// Panics if any record's length exceeds `u32::MAX`. BAM records are always + /// < 4 GiB, so this cannot happen for real records. + #[must_use] + pub fn from_owned_records(records: Vec<(K, Vec)>) -> Self { + let mut data = SegmentedBuf::new(); + let records = records + .into_iter() + .map(|(key, bytes)| { + let offset = data.extend_from_slice(&bytes) as u64; + let len = u32::try_from(bytes.len()) + .expect("InMemoryChunk record length exceeds u32 (BAM records are < 4 GiB)"); + (key, offset, len) + }) + .collect(); + Self { data: Arc::new(PooledSegmentedBuf::unpooled(data)), records } + } + /// Number of records in the chunk. #[must_use] - pub(crate) fn len(&self) -> usize { + pub fn len(&self) -> usize { self.records.len() } /// Whether the chunk is empty. #[must_use] - pub(crate) fn is_empty(&self) -> bool { + pub fn is_empty(&self) -> bool { self.records.is_empty() } + /// Total record-payload bytes (the sum of record lengths; excludes keys and + /// index overhead). Used for byte-budget accounting at the chunk boundary. + #[must_use] + pub fn payload_bytes(&self) -> usize { + self.records.iter().map(|(_, _, len)| *len as usize).sum() + } + /// Borrow the `i`th record's bytes from the shared data buffer. + /// + /// `pub` so out-of-crate consumers (e.g. the block-parallel spill serializer + /// in `fgumi-pipeline-io`) can iterate a chunk's records zero-copy. #[must_use] #[allow(clippy::cast_possible_truncation)] // offset/len fit in usize on all supported targets - pub(crate) fn record_bytes(&self, i: usize) -> &[u8] { + pub fn record_bytes(&self, i: usize) -> &[u8] { let (_, offset, len) = &self.records[i]; self.data.slice(*offset as usize, *len as usize) } /// Borrow the `i`th record's sort key. + /// + /// `pub` for the same reason as [`record_bytes`](Self::record_bytes). #[must_use] - pub(crate) fn key_at(&self, i: usize) -> &K { + pub fn key_at(&self, i: usize) -> &K { &self.records[i].0 } @@ -283,7 +328,14 @@ const _: () = assert!( /// Sorting only reorders the index; records stay in place. pub struct RecordBuffer { /// Segmented byte storage for all records (headers + BAM data). - data: SegmentedBuf, + /// + /// Held as the pool wrapper, not a bare `SegmentedBuf`: an arena acquired + /// from the [`ArenaPool`](crate::arena_pool::ArenaPool) lives here for its + /// whole fill, so if this were + /// bare, dropping the buffer on any error path would retire that pool slot + /// permanently. Non-pooled buffers hold an `unpooled` wrapper and drop + /// normally. + data: PooledSegmentedBuf, /// Index of record references for sorting. refs: Vec, /// Number of reference sequences (for unmapped handling). @@ -294,16 +346,22 @@ pub struct RecordBuffer { /// /// Both `RecordBuffer` and `TemplateRecordBuffer` use this segment size so that /// a single BAM record (≤128 MiB in practice) always fits within one segment. -const SORT_SEGMENT_SIZE: usize = 256 * 1024 * 1024; +pub const SORT_SEGMENT_SIZE: usize = 256 * 1024 * 1024; /// Shared implementation for `par_sort_into_chunks` on both buffer types. /// /// Sorts refs in place (parallel radix sort, partitioned by `chunk_size`), -/// then drains `self.data` into a single `Arc` shared +/// then drains `self.data` into a single `Arc` shared /// across all produced chunks. Each chunk's `records` Vec is built from /// its refs' `(key, offset + header_size, len)` triples — no per-record /// allocation or memcpy. /// +/// The drain goes through each buffer's own `take_data_arc`, never through +/// `mem::take` here: the two buffer types store `data` at different types +/// (`PooledSegmentedBuf` vs bare `SegmentedBuf`), so a `take` written in the +/// macro would deref-coerce past `RecordBuffer`'s pool wrapper and orphan a +/// pooled arena. See `RecordBuffer::take_data_arc`. +/// /// `$header_size` is the byte distance from each ref's `offset` field /// to the start of its actual record bytes inside `data` /// (`HEADER_SIZE` for `RecordBuffer`, `TEMPLATE_HEADER_SIZE` for @@ -318,7 +376,7 @@ macro_rules! par_sort_into_chunks_impl { if $threads <= 1 || n < RADIX_THRESHOLD * 2 || n <= 10_000 { $sort_fn(&mut $self.refs); - let data = Arc::new(std::mem::take(&mut $self.data)); + let data = $self.take_data_arc(); let records: Vec<_> = $self.refs.iter().map(|r| ($key_fn(r), r.offset + header_size, r.len)).collect(); // Clear refs to leave the buffer in a consistent drained @@ -336,7 +394,7 @@ macro_rules! par_sort_into_chunks_impl { $sort_fn(chunk); }); - let data = Arc::new(std::mem::take(&mut $self.data)); + let data = $self.take_data_arc(); let chunks: Vec<_> = $self .refs @@ -364,10 +422,10 @@ impl RecordBuffer { #[must_use] pub fn with_capacity(estimated_records: usize, estimated_bytes: usize, nref: u32) -> Self { Self { - data: SegmentedBuf::with_capacity( + data: PooledSegmentedBuf::unpooled(SegmentedBuf::with_capacity( estimated_bytes + estimated_records * HEADER_SIZE, SORT_SEGMENT_SIZE, - ), + )), refs: Vec::with_capacity(estimated_records), nref, } @@ -484,7 +542,7 @@ impl RecordBuffer { /// index-panic against the empty `SegmentedBuf`. #[must_use] pub(crate) fn drain_into_single_chunk(&mut self) -> InMemoryChunk { - let data = Arc::new(std::mem::take(&mut self.data)); + let data = self.take_data_arc(); let records: Vec<_> = self .refs .iter() @@ -497,6 +555,36 @@ impl RecordBuffer { InMemoryChunk::from_parts(data, records) } + /// Take the backing arena out of the buffer as the shared store for the + /// chunks drained from it, leaving an empty placeholder behind. + /// + /// Moves the **wrapper**, not the inner `SegmentedBuf`, so whether this + /// arena is pool-managed travels with it: a pooled arena returns to its + /// [`ArenaPool`](crate::arena_pool::ArenaPool) when the chunk's last `Arc` + /// drops, an unpooled one just drops. Re-wrapping the inner buffer instead + /// would deref-coerce past the wrapper, orphan the arena, and leave a + /// wrong-sized default behind for the pool to reclaim — at the default + /// `capacity == 1` that retires the only slot permanently. + /// + /// Every drain on this type goes through here so that reasoning lives in + /// one place rather than being re-derived at each call site. + fn take_data_arc(&mut self) -> Arc { + Arc::new(std::mem::take(&mut self.data)) + } + + /// Install a (pool-acquired) arena as this buffer's backing store, replacing + /// the empty placeholder left by `drain_into_single_chunk`. Requires the + /// buffer to be drained (no refs). + pub(crate) fn install_arena(&mut self, arena: PooledSegmentedBuf) { + // Unconditional (not `debug_assert!`): replacing `self.data` while stale + // `refs` still index the old arena is a silent data-corruption path for a + // bit-identical sort engine. This is a state-transition guard, not a + // hot-path bounds check — the `is_empty()` cost is negligible next to + // record processing. + assert!(self.refs.is_empty(), "install_arena over a non-drained buffer"); + self.data = arena; + } + /// Get record bytes by reference. #[inline] #[must_use] @@ -547,25 +635,6 @@ impl RecordBuffer { } /// Clear the buffer for reuse. - /// Give back the arena segments held for reuse; see - /// [`SegmentedBuf::release_retained`](crate::segmented_buf::SegmentedBuf::release_retained). - /// Call once ingest is finished — holding them through the merge costs peak - /// RSS and buys nothing, because nothing more is pushed. - pub fn release_retained(&mut self) { - self.data.release_retained(); - } - - /// Empty the buffer between chunks, **keeping the arena's segment - /// allocations** so the next chunk does not fault every page back in. - /// - /// See [`SegmentedBuf::reset_for_reuse`](crate::segmented_buf::SegmentedBuf::reset_for_reuse) - /// for the measurement. Use this rather than [`clear`](Self::clear) - /// anywhere the buffer is about to be refilled; `clear` is for teardown. - pub fn reset_for_reuse(&mut self) { - self.data.reset_for_reuse(); - self.refs.clear(); - } - pub fn clear(&mut self) { self.data.clear(); self.refs.clear(); @@ -840,7 +909,7 @@ impl RawSortKey for TemplateKey { /// # Panics /// - /// Always panics. `TemplateKey` extraction requires a [`LibraryLookup`](crate::external::LibraryLookup) + /// Always panics. `TemplateKey` extraction requires a [`LibraryLookup`](crate::LibraryLookup) /// context not available through the `RawSortKey` trait interface. All /// callers must use `extract_template_key_inline()` instead. fn extract(_bam: &[u8], _ctx: &SortContext) -> Self { @@ -1431,25 +1500,9 @@ impl TemplateRecordBuffer { /// Create a new buffer with estimated capacity. #[must_use] pub fn with_capacity(estimated_records: usize, estimated_bytes: usize) -> Self { - Self::with_segment_size(estimated_records, estimated_bytes, SORT_SEGMENT_SIZE) - } - - /// Create a new buffer with an explicit arena segment size. - /// - /// Production always uses [`SORT_SEGMENT_SIZE`]. This exists so tests can - /// drive the behaviour that only appears once the arena has *sealed* a - /// segment — deferred key batches are cut per sealed segment, and at 256 MiB - /// a segment no realistic test input ever fills, leaving that path - /// unexercised by everything short of a quarter-gigabyte fixture. - #[must_use] - pub fn with_segment_size( - estimated_records: usize, - estimated_bytes: usize, - segment_size: usize, - ) -> Self { let header_bytes = estimated_records * TEMPLATE_HEADER_SIZE; Self { - data: SegmentedBuf::with_capacity(estimated_bytes + header_bytes, segment_size), + data: SegmentedBuf::with_capacity(estimated_bytes + header_bytes, SORT_SEGMENT_SIZE), refs: Vec::with_capacity(estimated_records), } } @@ -1491,72 +1544,6 @@ impl TemplateRecordBuffer { Ok(()) } - /// Push a record whose sort key will be supplied later by - /// [`fill_keys`](Self::fill_keys). - /// - /// The ref lands with `K::default()` in its key field, which is **not** a - /// valid sort position — sorting before every deferred key has been filled - /// would silently collate records under a constant key rather than fail. - /// The ingest thread's barrier is what guarantees that cannot happen; it - /// counts filled keys against `refs.len()` before it sorts. - /// - /// # Errors - /// - /// Returns an error if the record (plus header) exceeds the segment size - /// (256 MiB) or if the record length exceeds `u32::MAX`. - #[inline] - pub fn push_deferred(&mut self, record: &[u8]) -> anyhow::Result<()> { - self.push(record, K::default()) - } - - /// Write keys for the refs at `first_ref .. first_ref + keys.len()`. - /// - /// # Panics - /// - /// Panics if the range runs past the end of the ref array — that means a - /// batch outlived the chunk it belonged to, which would otherwise corrupt a - /// later chunk's keys silently. - pub fn fill_keys(&mut self, first_ref: usize, keys: &[K]) { - let end = first_ref + keys.len(); - assert!( - end <= self.refs.len(), - "key batch covers refs {first_ref}..{end} but the buffer holds only {}", - self.refs.len(), - ); - for (r, k) in self.refs[first_ref..end].iter_mut().zip(keys) { - r.key = *k; - } - } - - /// Number of arena segments that are sealed — finished and safe to share - /// with a worker. See [`SegmentedBuf::sealed_len`](crate::segmented_buf::SegmentedBuf::sealed_len). - #[must_use] - pub fn sealed_segments(&self) -> usize { - self.data.sealed_len() - } - - /// A shared handle to sealed arena segment `idx`, or `None` if it is the - /// live segment. - #[must_use] - pub fn sealed_segment(&self, idx: usize) -> Option>> { - self.data.sealed_segment(idx) - } - - /// Seal the live arena segment, so every record pushed so far lives in a - /// segment that can be shared with a worker. - /// - /// Called at the chunk barrier, where the buffer is about to be sorted and - /// cleared, so the tail of the sealed segment is never wasted in practice. - pub fn seal_arena_segment(&mut self) { - self.data.seal_current(); - } - - /// The arena's segment size, for mapping a global offset to its segment. - #[must_use] - pub fn segment_size(&self) -> usize { - self.data.segment_size() - } - /// Sort the index by cached key using stable LSD radix sort. /// /// Uses multi-field radix sort which is stable (preserves relative order @@ -1639,25 +1626,6 @@ impl TemplateRecordBuffer { } /// Clear the buffer for reuse. - /// Give back the arena segments held for reuse; see - /// [`SegmentedBuf::release_retained`](crate::segmented_buf::SegmentedBuf::release_retained). - /// Call once ingest is finished — holding them through the merge costs peak - /// RSS and buys nothing, because nothing more is pushed. - pub fn release_retained(&mut self) { - self.data.release_retained(); - } - - /// Empty the buffer between chunks, **keeping the arena's segment - /// allocations** so the next chunk does not fault every page back in. - /// - /// See [`SegmentedBuf::reset_for_reuse`](crate::segmented_buf::SegmentedBuf::reset_for_reuse) - /// for the measurement. Use this rather than [`clear`](Self::clear) - /// anywhere the buffer is about to be refilled; `clear` is for teardown. - pub fn reset_for_reuse(&mut self) { - self.data.reset_for_reuse(); - self.refs.clear(); - } - pub fn clear(&mut self) { self.data.clear(); self.refs.clear(); @@ -1688,17 +1656,28 @@ impl TemplateRecordBuffer { ) } + /// Take the backing buffer out as the shared store for the chunks drained + /// from it, leaving an empty placeholder behind. + /// + /// Mirrors [`RecordBuffer::take_data_arc`] so both buffer types drain the + /// same way through `par_sort_into_chunks_impl!`. This buffer is never + /// pool-managed — `data` is a bare `SegmentedBuf`, not a wrapper — so the + /// shared handle is `unpooled` and simply drops with its last `Arc`. + fn take_data_arc(&mut self) -> Arc { + Arc::new(PooledSegmentedBuf::unpooled(std::mem::take(&mut self.data))) + } + /// Drain the (already-sorted) buffer into a single in-memory chunk /// whose records share an `Arc` backing store. /// - /// See [`RecordBuffer::drain_into_single_chunk`] for usage and + /// See `RecordBuffer::drain_into_single_chunk` for usage and /// lifetime notes — both `data` and `refs` are cleared after this /// call. `iter_sorted_keyed()` would yield zero records and /// `get_record()` called with a stale `TemplateRecordRef` would /// index-panic against the empty `SegmentedBuf`. #[must_use] pub(crate) fn drain_into_single_chunk(&mut self) -> InMemoryChunk { - let data = Arc::new(std::mem::take(&mut self.data)); + let data = self.take_data_arc(); let records: Vec<_> = self .refs .iter() @@ -1707,6 +1686,48 @@ impl TemplateRecordBuffer { self.refs.clear(); InMemoryChunk::from_parts(data, records) } + + /// Drain the (already-sorted) buffer into a single in-memory chunk keyed by + /// the **full** [`TemplateKey`] (re-extracted per record via `extract`), + /// whose records share an `Arc` backing store — NO record + /// bytes are copied (the arena is moved, not the bodies). + /// + /// This is the zero-copy analogue of the owned materialisation the legacy + /// `TemplateChunkSorter` performs: the narrow lane key `K` drove the sort, + /// and the full key is re-widened here for the downstream merge (whose + /// `MergeDriver` orders on the full key), but the record bodies + /// stay resident in the moved arena instead of being copied into owned + /// `RawRecord`s. `extract` receives each record's body bytes (the inline + /// header already skipped, exactly what [`get_record`](Self::get_record) + /// returns) and returns its full `TemplateKey`; it MUST be the same + /// extraction the owned path uses, so the produced chunk is byte-for-byte + /// identical to the owned `Vec<(TemplateKey, RawRecord)>` (same sorted order, + /// same keys, same bodies). + /// + /// Runs the per-record extraction in parallel (like the owned path's + /// `par_iter`), so call it inside the sorter's `rayon_pool.install`. Both + /// `data` and `refs` are cleared after this call (the arena is moved into the + /// returned chunk); see [`drain_into_single_chunk`](Self::drain_into_single_chunk). + #[must_use] + #[allow(clippy::cast_possible_truncation)] // body_off/len fit usize on all supported targets (BAM < 4 GiB, arena < address space) + pub(crate) fn drain_into_full_key_chunk( + &mut self, + extract: impl Fn(&[u8]) -> TemplateKey + Sync, + ) -> InMemoryChunk { + use rayon::prelude::*; + let data = self.take_data_arc(); + let records: Vec<(TemplateKey, u64, u32)> = self + .refs + .par_iter() + .map(|r| { + let body_off = r.offset + TEMPLATE_HEADER_SIZE as u64; + let body = data.slice(body_off as usize, r.len as usize); + (extract(body), body_off, r.len) + }) + .collect(); + self.refs.clear(); + InMemoryChunk::from_parts(data, records) + } } impl ProbeableBuffer for TemplateRecordBuffer { @@ -2730,53 +2751,46 @@ mod tests { } } + /// Parallel/serial parity at the arena dispatch cutoff. The arena front + /// (`template_chunk_from_arena_refs`) routes to `parallel_radix_sort_template_refs` + /// only when `refs.len() >= PARALLEL_SORT_THRESHOLD`; the arena-vs-owned parity + /// test runs far below that, so the parallel branch is never compared against a + /// reference there. Mirror `parallel_coordinate_sort_matches_serial_radix_at_threshold`: + /// build a threshold-sized input with deliberate key ties and pin the parallel + /// radix's ordering against the single-threaded stable radix. Since the + /// arena-vs-owned test already proves `radix_sort_template_refs` matches the + /// owned sorter, parallel-matches-serial here pins parallel-matches-owned at + /// threshold by transitivity. #[test] - fn test_deferring_a_key_and_filling_it_matches_pushing_it_outright() { - // Deferred extraction is only safe if the buffer it produces is - // indistinguishable from the one the serial push produced — same bytes - // at the same offsets, same keys on the same refs. - let keys: Vec = (0..16i32) + fn parallel_template_radix_matches_serial_at_threshold() { + let n = crate::ref_sort::PARALLEL_SORT_THRESHOLD; // crosses the arena cutoff + // Vary the library-name hash (a full-key radix lane) over a small modulus + // → ~260 ties per value across the input, so a stability or lane-ordering + // bug in the parallel path would reorder equal-key records relative to the + // serial radix. `offset` ascends with input order to witness stability. + let refs: Vec> = (0..n) .map(|i| { - let name_hash = u64::from(i.unsigned_abs()); - TemplateKey::new( - 0, - 100 + i, - false, - 0, - 200, - false, - 0, - 0, - (1, true), - name_hash, - false, - ) + let hash = (i as u64).wrapping_mul(2_654_435_761) % 1009; + let key = + TemplateKey::new(0, 100, false, 0, 200, false, 0, 0, (1, true), hash, false); + TemplateRecordRef { key, offset: i as u64, len: 10, padding: 0 } }) .collect(); - let records: Vec> = (0..16u8).map(|i| vec![i; 40 + usize::from(i)]).collect(); - - let mut direct = TemplateRecordBuffer::::with_capacity(16, 4096); - for (record, key) in records.iter().zip(&keys) { - direct.push(record, TemplateLaneKey::from_full(key)).expect("push"); - } - let mut deferred = TemplateRecordBuffer::::with_capacity(16, 4096); - for record in &records { - deferred.push_deferred(record).expect("push_deferred"); - } - // Fill in two uneven batches, as the pool would. - let narrowed: Vec = keys.iter().map(TemplateLaneKey::from_full).collect(); - deferred.fill_keys(0, &narrowed[..5]); - deferred.fill_keys(5, &narrowed[5..]); + let mut parallel = refs.clone(); + let mut serial = refs; + parallel_radix_sort_template_refs(&mut parallel); // n > 10_000 → real parallel path + radix_sort_template_refs(&mut serial); // single-threaded stable reference - assert_eq!(direct.refs().len(), deferred.refs().len()); - for (a, b) in direct.refs().iter().zip(deferred.refs()) { - assert_eq!(a.key, b.key, "keys must match"); - assert_eq!(a.offset, b.offset, "record bytes must land at the same offset"); - assert_eq!(a.len, b.len); - } - for (i, r) in deferred.refs().iter().enumerate() { - assert_eq!(deferred.get_record(r), &records[i][..], "record bytes must survive"); + assert_eq!(parallel.len(), serial.len()); + for (i, (p, s)) in parallel.iter().zip(&serial).enumerate() { + assert!( + p.key == s.key && p.offset == s.offset, + "parallel template radix diverged from serial radix at index {i} \ + (parallel offset {}, serial offset {})", + p.offset, + s.offset, + ); } } @@ -3477,7 +3491,7 @@ mod tests { fn test_in_memory_chunk_from_parts_reads_records() { let (buf, offsets) = build_segmented_buf_with_records(&[b"first", b"second-record", b"third"]); - let data = Arc::new(buf); + let data = Arc::new(PooledSegmentedBuf::unpooled(buf)); let records = vec![ (10u32, offsets[0].0, offsets[0].1), (20u32, offsets[1].0, offsets[1].1), @@ -3499,7 +3513,7 @@ mod tests { fn test_in_memory_chunk_take_key_replaces_with_default() { let (buf, offsets) = build_segmented_buf_with_records(&[b"payload"]); let chunk = InMemoryChunk::::from_parts( - Arc::new(buf), + Arc::new(PooledSegmentedBuf::unpooled(buf)), vec![(42u32, offsets[0].0, offsets[0].1)], ); let mut chunk = chunk; @@ -3517,7 +3531,7 @@ mod tests { // Arc, so the data isn't duplicated and is freed // only when the last chunk drops. let (buf, offsets) = build_segmented_buf_with_records(&[b"alpha", b"beta", b"gamma"]); - let data = Arc::new(buf); + let data = Arc::new(PooledSegmentedBuf::unpooled(buf)); let chunk_a = InMemoryChunk::::from_parts( Arc::clone(&data), vec![(1, offsets[0].0, offsets[0].1)], @@ -3620,6 +3634,44 @@ mod tests { assert_eq!(Arc::strong_count(&chunks[0].data), chunks.len()); } + /// `par_sort_into_chunks` must move the pool WRAPPER out of the buffer, the + /// same way `drain_into_single_chunk` does. Re-wrapping the *inner* + /// `SegmentedBuf` deref-coerces past the wrapper: the real arena is orphaned + /// (the chunks drop it instead of returning it) while the buffer keeps a + /// wrong-sized default that still believes it is pooled. At the default + /// `capacity == 1` that retires the only slot, and since the documented + /// response to `try_acquire() == None` is to backpressure until an arena + /// returns, the caller hangs rather than errors. + /// + /// `buffer` is deliberately still alive at the assert: the arena must come + /// back from the dropped chunks, not from the drained buffer's leftover + /// wrapper. Both exits of `par_sort_into_chunks_impl!` are covered — the + /// single-chunk early return and the multi-chunk path. + #[rstest::rstest] + #[case::single_chunk_exit(1_000)] + #[case::multi_chunk_exit(10_500)] + fn par_sort_into_chunks_returns_the_pooled_arena_when_its_chunks_drop(#[case] n: usize) { + let arena_pool = crate::arena_pool::ArenaPool::new(1, SORT_SEGMENT_SIZE); + let mut buffer = RecordBuffer::with_capacity(n, n * 64, 4); + buffer.install_arena(arena_pool.try_acquire().expect("the fresh pool hands out its arena")); + + for i in 0..n { + // Descending position so the sort is non-trivial. `#[allow]` on the + // test fn would not reach the per-case fns `rstest` generates, so + // the cast is done with `try_from` rather than suppressed. + let pos = i32::try_from(n - i).expect("test record count fits i32"); + buffer.push_coordinate(&make_coordinate_bam_record(0, pos)).expect("push_coordinate"); + } + + let rayon_pool = + rayon::ThreadPoolBuilder::new().num_threads(4).build().expect("rayon pool build"); + let chunks = rayon_pool.install(|| buffer.par_sort_into_chunks(4)); + assert_eq!(arena_pool.free_len(), 0, "the arena is in flight while the chunks hold it"); + + drop(chunks); + assert_eq!(arena_pool.free_len(), 1, "dropping the chunks returns the arena to the pool"); + } + mod template_key32 { use super::*; diff --git a/crates/fgumi-sort/src/keys.rs b/crates/fgumi-sort/src/keys.rs index 7ad97c899..0ef41ee3c 100644 --- a/crates/fgumi-sort/src/keys.rs +++ b/crates/fgumi-sort/src/keys.rs @@ -9,7 +9,7 @@ //! - [`RawCoordinateKey`]: Fixed-size genomic coordinate key (tid, pos, strand) //! - [`RawQuerynameKey`]: Read name with natural numeric ordering //! - [`RawQuerynameLexKey`]: Read name with lexicographic ordering -//! - [`TemplateKey`](crate::inline::TemplateKey): Template-level position for UMI grouping +//! - [`TemplateKey`](crate::TemplateKey): Template-level position for UMI grouping //! //! # Generic Sorting Abstraction //! @@ -23,11 +23,58 @@ //! - samtools' `bam1_tag` union (C) use noodles::sam::Header; +use smallvec::SmallVec; use std::cmp::Ordering; use fgumi_raw_bam::RawRecordView; use std::io::{Read, Write}; +/// Inline-capacity buffer for queryname key name bytes. An inline `SmallVec` +/// keeps the name **contiguous inside the sort ref array** instead of a +/// per-record heap allocation. During the comparison sort this eliminates the +/// pointer-chase to scattered heap that dominates on memory-latency-bound cores +/// (e.g. Graviton), where the queryname sort regressed vs the owned-`Vec` key. +/// Names longer than the inline capacity spill to the heap transparently. +/// +/// # Why 44 +/// +/// 44 is the smallest capacity that keeps **every** read name we have measured +/// inline. It is *not* free relative to 40: on 64-bit targets `SmallVec<[u8; 40]>` +/// is a 48-byte `NameBuf` (a 56-byte `RawQuerynameKey`), whereas `SmallVec<[u8; 44]>` +/// is a 56-byte `NameBuf` (a 64-byte key) — so 44 costs 8 bytes per key over 40. +/// That is a deliberate trade to keep inline the ~16% of native Illumina names +/// that 40 would spill to the heap (the pointer-chase this inline buffer exists +/// to avoid); the `position_fits_in_existing_key_padding` assertion pins the +/// resulting layout. +/// +/// Read names in practice fall into two regimes, and the choice was measured on +/// both rather than on one: +/// +/// - **Native Illumina** (`A00132:53:HFHJKDSXX:1:1646:26467:33332`) — 37-41 +/// bytes with the NUL. This is raw instrument output. At capacity 24 *every* +/// name spills. +/// - **SRA-normalized / synthetic** (`SRR36097899.1`, `mol000000_read0001`) — +/// 13-22 bytes. These fit at 24 already, so a larger inline buffer is pure +/// cost for them. +/// +/// End-to-end `fgumi sort --order queryname` (15-16M records, `-m 512m -t 8`, +/// 4-6 reps, CV ≈ 3%), capacity 44 against capacity 24: +/// +/// | input | CPU | peak RSS | +/// | --- | --- | --- | +/// | native Illumina | −9.5% | −4.4% | +/// | SRA-normalized | −0.6% (neutral) | +3.9% | +/// +/// The long-name win is not a memory-for-speed trade, which is the intuition +/// this measurement overturned: a million individually heap-allocated names cost +/// *more* resident memory than a million 16-byte-larger inline keys, so 44 is +/// both faster and smaller there. Isolated key-level cost on native names is +/// 2.4x faster to extract and 1.95x faster to sort. The price is ~4% peak RSS on +/// short-name inputs, where CPU is unchanged. +pub(crate) const NAME_INLINE_CAP: usize = 44; + +type NameBuf = SmallVec<[u8; NAME_INLINE_CAP]>; + // ============================================================================ // Generic Sorting Abstraction (Trait-based, inspired by fgbio/samtools) // ============================================================================ @@ -519,13 +566,18 @@ fn write_queryname_key(name: &[u8], flags: u16, writer: &mut W) -> std } /// Deserialize a queryname key from `[name_len: u16][name: bytes][flags: u16]`. +/// +/// Reads the name straight into an inline [`NameBuf`] so the spill-merge read +/// path allocates no per-record heap for short names — the merge-side analogue +/// of the SSO key change on the ingest path. #[inline] -fn read_queryname_key(reader: &mut R) -> std::io::Result<(Vec, u16)> { +fn read_queryname_key(reader: &mut R) -> std::io::Result<(NameBuf, u16)> { let mut len_buf = [0u8; 2]; reader.read_exact(&mut len_buf)?; let name_len = u16::from_le_bytes(len_buf) as usize; - let mut name = vec![0u8; name_len]; + let mut name: NameBuf = SmallVec::new(); + name.resize(name_len, 0); reader.read_exact(&mut name)?; let mut flags_buf = [0u8; 2]; @@ -543,10 +595,11 @@ fn read_queryname_key(reader: &mut R) -> std::io::Result<(Vec, u16) /// Unlike fixed-size keys, serialization size depends on name length. /// /// Serialization format: `[name_len: u16][name: bytes][flags: u16]` -#[derive(Clone, Debug, Default)] +#[derive(Clone, Debug)] pub struct RawQuerynameKey { - /// Read name bytes, null-terminated for `natural_compare_nul`. - name: Vec, + /// Read name bytes, null-terminated for `natural_compare_nul`. Inline for + /// short names (see [`NameBuf`]) so the comparison sort stays cache-local. + name: NameBuf, /// Flags for segment ordering (R1 before R2). flags: u16, /// Ingest position within the current in-memory chunk; see @@ -555,6 +608,20 @@ pub struct RawQuerynameKey { pos: u32, } +impl Default for RawQuerynameKey { + /// A default key must still satisfy the null-terminated invariant that + /// `cmp`'s `natural_compare_nul` relies on — a derived `Default` would leave + /// `name` empty (no terminator), so a default-constructed key would feed a + /// non-terminated pointer into the comparator (out-of-bounds read). This is + /// required because generic merge code bounds `K: RawSortKey + Default` + /// (e.g. `MergeDriver`, `MemorySources`). The default is the empty + /// name `"\0"` — it orders before any non-empty name and terminates + /// immediately. + fn default() -> Self { + Self { name: SmallVec::from_slice(&[0]), flags: 0, pos: 0 } + } +} + impl PartialEq for RawQuerynameKey { fn eq(&self, other: &Self) -> bool { self.cmp(other) == Ordering::Equal @@ -572,7 +639,7 @@ impl RawQuerynameKey { if name.last() != Some(&0) { name.push(0); } - Self { name, flags, pos: 0 } + Self { name: SmallVec::from_vec(name), flags, pos: 0 } } /// Returns the read name bytes (including the null terminator). @@ -589,12 +656,13 @@ impl RawQuerynameKey { let flags = queryname_flag_order(u16::from_le_bytes([bam[14], bam[15]])); match raw_name_with_nul(bam) { // BAM stores the name NUL-terminated, so copy those bytes directly - // instead of stripping the NUL and re-appending it. `pos` is filled in - // later by `set_position`; a merge-time rebuild leaves it 0. - Some(name) => Self { name: name.to_vec(), flags, pos: 0 }, + // (inline when short) instead of stripping the NUL and re-appending + // it. `pos` is filled in later by `set_position`; a merge-time + // rebuild leaves it 0. + Some(name) => Self { name: SmallVec::from_slice(name), flags, pos: 0 }, // Truncated or non-terminated record: fall back to the stripped-name // path (`new` re-adds the terminator and sets `pos: 0`), preserving - // the prior bytes. + // the prior bytes exactly. None => Self::new(extract_raw_name_and_flags(bam).0.to_vec(), flags), } } @@ -604,7 +672,11 @@ impl Ord for RawQuerynameKey { #[inline] #[allow(unsafe_code)] fn cmp(&self, other: &Self) -> Ordering { - // SAFETY: `name` is always null-terminated (see `extract_queryname_key` and `new`). + // SAFETY: `name` is always null-terminated by every constructor + // (`extract_queryname_key`, `new`, `read_from`, and the hand-written + // `Default`, which exists precisely to uphold this invariant -- a + // derived one would leave `name` empty and this walk would run off the + // end). unsafe { natural_compare_nul(self.name.as_ptr(), other.name.as_ptr()) } .then_with(|| self.flags.cmp(&other.flags)) .then_with(|| self.pos.cmp(&other.pos)) @@ -644,8 +716,18 @@ impl RawSortKey for RawQuerynameKey { #[inline] fn read_from(reader: &mut R) -> std::io::Result { - let (name, flags) = read_queryname_key(reader)?; - Ok(Self::new(name, flags)) + // A well-formed frame (written by `write_to`) already carries the + // trailing NUL that `cmp`'s `natural_compare_nul` unsafe relies on, so + // the common path constructs directly rather than via `new`. But + // `read_from` is a public entry point that could be handed a truncated + // or malformed frame; re-append the terminator if it is missing so the + // NUL-termination invariant holds unconditionally and the comparator can + // never walk off the end. + let (mut name, flags) = read_queryname_key(reader)?; + if name.last() != Some(&0) { + name.push(0); + } + Ok(Self { name, flags, pos: 0 }) } } @@ -661,8 +743,9 @@ impl RawSortKey for RawQuerynameKey { /// Serialization format: `[name_len: u16][name: bytes][flags: u16]` #[derive(Clone, Eq, PartialEq, Debug, Default)] pub struct RawQuerynameLexKey { - /// Read name bytes. - name: Vec, + /// Read name bytes. Inline for short names (see [`NameBuf`]) so the + /// comparison sort stays cache-local. + name: NameBuf, /// Flags for segment ordering (R1 before R2). flags: u16, /// Ingest position within the current in-memory chunk; see @@ -675,7 +758,7 @@ impl RawQuerynameLexKey { /// Create a new lexicographic queryname key. #[must_use] pub fn new(name: Vec, flags: u16) -> Self { - Self { name, flags, pos: 0 } + Self { name: SmallVec::from_vec(name), flags, pos: 0 } } /// Returns the read name bytes. @@ -689,7 +772,7 @@ impl RawQuerynameLexKey { #[must_use] fn extract_queryname_key(bam: &[u8]) -> Self { let (raw_name, flags) = extract_raw_name_and_flags(bam); - Self { name: raw_name.to_vec(), flags, pos: 0 } + Self { name: SmallVec::from_slice(raw_name), flags, pos: 0 } } } @@ -751,22 +834,31 @@ mod tests { /// The ingest position must land in the key's existing tail padding. /// - /// On a 64-bit target `Vec` (24 bytes, align 8) plus a `u16` leaves 6 bytes of - /// padding, so the `u32` is free. If this ever fails the tiebreak has started - /// costing memory per record, and the trade against a stable sort needs + /// On a 64-bit target [`NameBuf`] measures 56 bytes (align 8); adding the `u16` + /// flags reaches 58, which rounds up to 64 and leaves 6 bytes of padding. The + /// `u32` position therefore costs nothing. If this ever fails the tiebreak has + /// started costing memory per record, and the trade against a stable sort needs /// re-evaluating: the whole point is that a total order is cheaper than /// stable-sorting. /// - /// The assertion is scoped to 64-bit because the padding is. A 32-bit `Vec` is - /// 12 bytes, which leaves only 2 bytes after the `u16`, so the key grows from 16 to - /// 20 bytes there. That is a cost question, not a correctness one — the total order - /// the unstable sort depends on holds on every target — and fgumi's sort is a - /// 64-bit workload, so the zero-growth claim is made for 64-bit only. + /// The size moved from 32 to 48 when the name became an inline `SmallVec` rather + /// than a `Vec`, and from 48 to 64 when that buffer grew from 24 to 44 bytes + /// so native Illumina names stop spilling (see [`NameBuf`]). Both steps are the + /// SSO buffer being paid for, not the position, which was free before either + /// change and is free after both. + /// + /// The assertion is scoped to 64-bit because the padding is. That is a cost + /// question, not a correctness one — the total order the unstable sort depends on + /// holds on every target — and fgumi's sort is a 64-bit workload, so the + /// zero-growth claim is made for 64-bit only. #[test] #[cfg(target_pointer_width = "64")] fn position_fits_in_existing_key_padding() { - assert_eq!(std::mem::size_of::(), 32); - assert_eq!(std::mem::size_of::(), 32); + assert_eq!(std::mem::size_of::(), 64); + assert_eq!(std::mem::size_of::(), 64); + // The claim is that `pos` is free, so pin the padding it lives in rather + // than only the total: name + flags alone already round up to 64. + assert_eq!(std::mem::size_of::() + std::mem::size_of::(), 58); } /// Name and flags alone are not a total order; position makes them one. @@ -816,6 +908,42 @@ mod tests { assert_eq!(a, b); } + /// `RawQuerynameKey::default()` must produce a null-terminated `name` so the + /// `unsafe` `natural_compare_nul` in `cmp` never walks off the end of an + /// unterminated buffer (a derived `Default` would leave `name` empty). The + /// default orders before any non-empty name. + #[test] + fn default_queryname_key_is_null_terminated() { + let def = RawQuerynameKey::default(); + assert_eq!(def.name(), &[0u8], "default name must be a single NUL terminator"); + // cmp against a real key must not read out of bounds and must order the + // empty default first. + let real = RawQuerynameKey::new(b"read1".to_vec(), 0); + assert_eq!(def.cmp(&real), Ordering::Less); + assert_eq!(def.cmp(&RawQuerynameKey::default()), Ordering::Equal); + } + + /// `read_from` must restore the NUL terminator that `cmp`'s `unsafe` + /// `natural_compare_nul` relies on, even for a truncated or hand-crafted + /// frame whose serialized name lacks it — the branch a `write_to`-produced + /// frame never exercises (its name already carries the terminator). + #[test] + fn read_from_terminates_a_frame_whose_name_lacks_the_nul() { + // Hand-built frame: [u16 len][name bytes, no NUL][u16 flags]. + let name = b"readX"; + let mut buf = Vec::new(); + buf.extend_from_slice(&u16::try_from(name.len()).unwrap().to_le_bytes()); + buf.extend_from_slice(name); + buf.extend_from_slice(&7u16.to_le_bytes()); + + let mut cursor = std::io::Cursor::new(&buf); + let key = RawQuerynameKey::read_from(&mut cursor).expect("read_from should succeed"); + assert_eq!(key.name(), b"readX\0", "read_from must restore the terminator"); + // cmp must not read past the buffer and must compare equal to the key + // the terminated constructor produces from the same bytes. + assert_eq!(key.cmp(&RawQuerynameKey::new(b"readX".to_vec(), 7)), Ordering::Equal); + } + // ======================================================================== // queryname_flag_order tests // ======================================================================== @@ -1837,4 +1965,113 @@ mod tests { "extracted name must be NUL-terminated for natural_compare_nul" ); } + + // NameBuf heap-fallback (SSO spill) tests + // + // `NameBuf` inlines names up to `NAME_INLINE_CAP` bytes and spills longer ones to the + // heap. Read names in real data are almost always short, so the common + // unit tests above only exercise the inline path. These tests explicitly + // drive names past the inline capacity for BOTH key kinds (lexicographic + // and natural) across every surface that touches a `NameBuf`: `new`, + // `extract_from_record`, the `write_to`/`read_from` serialization + // roundtrip, and ordering — so the heap-backed path can never silently + // regress. + // ======================================================================== + + /// A read name comfortably longer than `NameBuf`'s inline capacity, so the + /// backing `SmallVec` is forced onto the heap. The trailing `9` / `10` lets a + /// paired name below expose the lex-vs-natural ordering split. + /// + /// The prefix alone must exceed `NAME_INLINE_CAP`, so every name derived from + /// it spills regardless of suffix — `heap_name_prefix_actually_spills` pins + /// that. Raising the capacity without lengthening this would silently turn + /// the whole heap-fallback block below into a second set of inline tests. + const HEAP_NAME_PREFIX: &[u8] = b"A00123:45:HGVWXDSXY:1:1101:12345:67890:24680:13579:"; + + /// The heap-fallback block below is only meaningful while its names spill. + #[test] + fn heap_name_prefix_actually_spills() { + assert!( + HEAP_NAME_PREFIX.len() > NAME_INLINE_CAP, + "HEAP_NAME_PREFIX is {} bytes but NameBuf inlines {NAME_INLINE_CAP}; the \ + heap-fallback tests would run on the inline path", + HEAP_NAME_PREFIX.len(), + ); + } + + #[test] + fn test_natural_key_heap_name_serialization_roundtrip() { + let name = [HEAP_NAME_PREFIX, b"67890"].concat(); + assert!(name.len() > NAME_INLINE_CAP, "test name must exceed NameBuf inline capacity"); + + let key = RawQuerynameKey::new(name.clone(), 42); + // `new` appends the NUL terminator that `natural_compare_nul` relies on. + assert_eq!(key.name(), [name.as_slice(), b"\0"].concat().as_slice()); + + let mut buf = Vec::new(); + key.write_to(&mut buf).expect("write_to should succeed"); + let mut cursor = std::io::Cursor::new(&buf); + let restored = RawQuerynameKey::read_from(&mut cursor).expect("read_from should succeed"); + + assert_eq!(key, restored); + assert_eq!(restored.name(), key.name(), "heap name (incl. NUL) must survive roundtrip"); + } + + #[test] + fn test_lex_key_heap_name_serialization_roundtrip() { + let name = [HEAP_NAME_PREFIX, b"67890"].concat(); + assert!(name.len() > NAME_INLINE_CAP, "test name must exceed NameBuf inline capacity"); + + let key = RawQuerynameLexKey::new(name.clone(), 42); + assert_eq!(key.name(), name.as_slice()); + + let mut buf = Vec::new(); + key.write_to(&mut buf).expect("write_to should succeed"); + let mut cursor = std::io::Cursor::new(&buf); + let restored = + RawQuerynameLexKey::read_from(&mut cursor).expect("read_from should succeed"); + + assert_eq!(key, restored); + assert_eq!(restored.name(), key.name(), "heap name must survive roundtrip"); + } + + #[test] + fn test_queryname_keys_extract_from_record_heap_name() { + use fgumi_raw_bam::testutil::make_bam_bytes; + + let name = [HEAP_NAME_PREFIX, b"67890"].concat(); + assert!(name.len() > NAME_INLINE_CAP, "test name must exceed NameBuf inline capacity"); + let bam = make_bam_bytes(0, 0, 0, &name, &[], 100, -1, -1, &[]); + + // Lexicographic: the extracted name is the raw name, no NUL. + let lex = RawQuerynameLexKey::extract_from_record(&bam); + assert_eq!(lex.name(), name.as_slice()); + + // Natural: the extracted name is NUL-terminated for `natural_compare_nul`. + let nat = RawQuerynameKey::extract_from_record(&bam); + assert_eq!(nat.name(), [name.as_slice(), b"\0"].concat().as_slice()); + } + + #[test] + fn test_heap_name_lex_vs_natural_ordering_difference() { + // Both names spill to the heap; they differ only in the numeric suffix + // `9` vs `10`, so the two comparators disagree exactly as they do for + // short names — proving the split is preserved on the heap-backed path. + let name_9 = [HEAP_NAME_PREFIX, b"9"].concat(); + let name_10 = [HEAP_NAME_PREFIX, b"10"].concat(); + assert!( + name_9.len() > NAME_INLINE_CAP && name_10.len() > NAME_INLINE_CAP, + "names must exceed inline capacity", + ); + + // Lexicographic: '1' < '9', so `...:10` sorts before `...:9`. + let lex_9 = RawQuerynameLexKey::new(name_9.clone(), 0); + let lex_10 = RawQuerynameLexKey::new(name_10.clone(), 0); + assert!(lex_10 < lex_9, "lexicographic: '...:10' should be < '...:9'"); + + // Natural: 9 < 10 numerically, so `...:9` sorts before `...:10`. + let nat_9 = RawQuerynameKey::new(name_9, 0); + let nat_10 = RawQuerynameKey::new(name_10, 0); + assert!(nat_9 < nat_10, "natural: '...:9' should be < '...:10'"); + } } diff --git a/crates/fgumi-sort/src/lib.rs b/crates/fgumi-sort/src/lib.rs index 91075ac16..5c3167024 100644 --- a/crates/fgumi-sort/src/lib.rs +++ b/crates/fgumi-sort/src/lib.rs @@ -39,63 +39,41 @@ use noodles::sam::header::record::value::Map; use noodles::sam::header::record::value::map::header::tag as header_tag; use tempfile::TempDir; -// All sub-modules are crate-private. Items intended for external consumers are -// re-exported at the crate root below. +// Sub-modules are crate-private except where a consumer needs to name the type +// itself rather than just receive it (`arena_pool`, `segmented_buf`, `codec`, +// `ref_sort`, `template_arena`). +// Everything else is re-exported at the crate root below. +pub mod arena_pool; pub(crate) mod bgzf_io; +pub(crate) mod chunk_sorter; pub mod codec; -/// Whether to emit the sort's performance diagnostics. -/// -/// A plain `fgumi sort` used to print ~99 diagnostic lines at INFO -- spill -/// geometry, per-phase breakdowns, park census tables, stall histograms -- and -/// none of it was opt-in. Those lines exist to answer "which of three limits is -/// this merge on, and where did the time go" during a performance investigation. -/// They are read from a log file with a grep, never from a terminal, and they -/// buried the handful of lines an end user can act on. -/// -/// The flag is run-scoped, not a process-global: it lives on [`RawExternalSorter`] (the sort's -/// options struct) and is copied into the `SortPhaseTimer` and `SortWorkerPool` that sorter -/// builds, so two sorts in one process (fgumi-sort is a library) keep independent settings. The -/// emitters therefore read it from the `self`/`pool` they already hold, or take it as a -/// `sort_stats: bool` parameter where they are free functions. -/// -/// Emit a performance diagnostic, at INFO, only when `--sort-stats` is on for the run. -/// -/// The first argument is the run-scoped `sort_stats` flag (a `bool`); the rest is a `log::info!` -/// format string and arguments. -/// -/// Deliberately not `debug!`: these lines are the harness's data, and a benchmark that has to -/// raise the global log level to collect them also collects every other crate's debug output, -/// which is how a 1,700-line log became a 60,000-line one. `--sort-stats` selects *these* lines -/// and nothing else. -macro_rules! stat { - ($enabled:expr, $($arg:tt)*) => { - if $enabled { - log::info!($($arg)*); - } - }; -} - pub(crate) mod external; pub(crate) mod fd_limit; pub(crate) mod inline; pub(crate) mod keys; pub(crate) mod loser_tree; pub(crate) mod memory_probe; -pub(crate) mod merge_headroom; pub(crate) mod merge_phases; +pub(crate) mod merge_slots; pub(crate) mod merge_stalls; pub(crate) mod merge_trace; -pub(crate) mod phase1_keys; -pub(crate) mod phase1_stats; pub(crate) mod pipeline; pub(crate) mod pooled_bam_writer; pub(crate) mod pooled_chunk_writer; -pub(crate) mod progress_batch; pub(crate) mod radix; pub(crate) mod read_ahead; pub(crate) mod reader; -pub(crate) mod segmented_buf; -pub(crate) mod spill_reader; +pub mod ref_sort; +pub mod segmented_buf; +pub mod template_arena; +// Block-granular spill compression kernel for the block-parallel spill steps +// (`SpillGather`/`SpillCompress`/`SpillWrite`). Surfaced via the re-exports +// below. +pub(crate) mod spill_block; +pub(crate) mod spill_block_reader; +// Synchronous inline spill compress for the P6 `CompressSpill` step (retires the +// `SortWorkerPool` compress path). Surfaced via `write_sorted_chunk` below. +pub(crate) mod sync_spill_writer; pub(crate) mod tmp_dir_alloc; pub(crate) mod verify; pub(crate) mod worker_pool; @@ -234,27 +212,43 @@ fn create_temp_dir(base: Option<&Path>) -> Result { } } +pub use arena_pool::{ArenaPool, PooledSegmentedBuf}; +pub use chunk_sorter::{CoordinateChunkSorter, TemplateChunkSorter}; pub use codec::SpillCodec; +pub use external::MemorySources; pub use external::{ - KeyTypesSpec, LibraryLookup, RawExternalSorter, ReadStreams, cb_hasher, - extract_template_key_inline, format_thread_counts, + KeyTypesSpec, LibraryLookup, MergeDriver, MergeDriverDyn, MergeStep, RawExternalSorter, + ReadStreams, cb_hasher, extract_template_key_inline, format_thread_counts, open_spill_slot, }; pub use fd_limit::{ FALLBACK_MAX_TEMP_FILES, fits_nofile_budget, resolve_temp_file_limit, soft_nofile, temp_file_limit_from_nofile, }; pub use inline::{ - PackedCoordinateKey, RecordRef, TemplateKey, extract_coordinate_key_inline, - radix_sort_record_refs, radix_sort_record_refs_with_max, + CbKey32, InMemoryChunk, PackedCoordinateKey, RecordBuffer, RecordRef, SORT_SEGMENT_SIZE, + TemplateKey, TemplateKey24, TertKey32, extract_coordinate_key_inline, radix_sort_record_refs, + radix_sort_record_refs_with_max, }; pub use keys::{ QuerynameComparator, RawCoordinateKey, RawQuerynameKey, RawQuerynameLexKey, RawSortKey, SortContext, SortOrder, natural_compare, natural_compare_nul, normalize_natural_key, }; +pub use merge_slots::{PHASE2_DECOMP_CAP, SortMergeReader, SortMergeSlot}; pub use reader::{ OwnedRawBamRecordReader, RawBamRecordReader, open_raw_bam_record_reader, open_raw_bam_record_reader_with_header, }; +pub use ref_sort::{ + coordinate_chunk_from_arena_refs, coordinate_chunk_from_refs, queryname_chunk_from_arena_refs, +}; +pub use segmented_buf::SegmentedBuf; +pub use spill_block::{SpillBlockCompressor, frame_keyed_record_into, spill_magic, spill_trailer}; +pub use spill_block_reader::SpillBlockDecompressor; +pub use sync_spill_writer::{write_sorted_chunk, write_sorted_chunk_inmem}; +pub use template_arena::{ + TemplateArenaAccumulator, TemplateMemChunk, template_chunk_from_arena_refs, +}; +pub use tmp_dir_alloc::TmpDirAllocator; pub use verify::{VerifySummary, verify_sort_order}; #[cfg(test)] diff --git a/crates/fgumi-sort/src/loser_tree.rs b/crates/fgumi-sort/src/loser_tree.rs index 26d1c7eaa..a463d2c1d 100644 --- a/crates/fgumi-sort/src/loser_tree.rs +++ b/crates/fgumi-sort/src/loser_tree.rs @@ -134,47 +134,6 @@ impl LoserTree { self.losers[0] } - /// The source most likely to be consumed after the winner, or `None` when - /// fewer than two sources remain. - /// - /// The runner-up is always among the losers stored along the winner's - /// root-to-leaf path -- anything eliminated elsewhere in the tree lost to a - /// source that itself lost to the winner. So this is O(log k) comparisons over - /// existing state, ~6 for a 44-way merge, and needs no new bookkeeping. - /// - /// `losers[1]` alone is NOT the answer, which is the tempting O(1) shortcut: - /// after a replay that node holds whoever lost the final comparison of *that* - /// replay, not the global second-smallest. With keys `[50, 100, 30, 20]` the - /// winner is source 3 and `losers[1]` is source 0 (key 50), while the true - /// runner-up is source 2 (key 30). - /// - /// It exists to be read *predictively*. The merge starves at run transitions: - /// only 2,475 of 167,624 source switches starve, but each costs ~20ms against - /// an 8.4ms read latency -- the read had not been started when the consumer - /// arrived. Depth cannot fix that (doubling the read-ahead cap moved the - /// starved share of park time not at all, 99% to 99%), but knowing which file - /// is next, ~300 blocks before it is needed, can. - #[must_use] - pub fn runner_up(&self) -> Option { - if self.num_active < 2 { - return None; - } - let winner = self.losers[0]; - let mut node = self.leaf_to_node(winner); - let mut best: Option = None; - while node > 0 { - let candidate = self.losers[node]; - if candidate != EMPTY && self.active[candidate] { - best = Some(match best { - Some(b) if self.is_greater(candidate, b) => b, - _ => candidate, - }); - } - node >>= 1; - } - best - } - /// Check if the winner is still an active source. #[inline] #[must_use] @@ -426,41 +385,4 @@ mod tests { assert_eq!(result, (1..=10).collect::>()); } - /// The runner-up predicts the source the merge will consume after the current - /// winner run. It starts that source's disk read ~300 blocks early. Getting it - /// wrong prefetches the wrong file, which is how grab-N cost 36%. - #[test] - fn test_runner_up_is_the_second_smallest_key() { - let tree = LoserTree::new(vec![50u32, 10, 30, 20]); - assert_eq!(tree.winner(), 1, "guard: 10 is the winner"); - assert_eq!(tree.runner_up(), Some(3), "20 is next, at source 3"); - } - - /// Must track the tree, not a stale snapshot: after the winner advances the - /// prediction has to move with it. - #[test] - fn test_runner_up_follows_the_tree_as_the_winner_advances() { - let mut tree = LoserTree::new(vec![50u32, 10, 30, 20]); - tree.replace_winner(100); - assert_eq!(tree.winner(), 3, "20 now wins"); - assert_eq!(tree.runner_up(), Some(2), "30 is next, at source 2"); - } - - /// One source has no next source; predicting one would send a read at the file - /// already being drained, which already has the deep allowance. - #[test] - fn test_runner_up_is_none_with_a_single_source() { - assert_eq!(LoserTree::new(vec![7u32]).runner_up(), None); - } - - /// Exhausted sources must never be predicted: a read issued for a drained file - /// is pure waste, and `active` is what distinguishes them. - #[test] - fn test_runner_up_skips_an_exhausted_source() { - let mut tree = LoserTree::new(vec![10u32, 20]); - assert_eq!(tree.winner(), 0); - tree.remove_winner(); - assert_eq!(tree.winner(), 1, "guard: source 0 is gone"); - assert_eq!(tree.runner_up(), None, "no live source remains behind the winner"); - } } diff --git a/crates/fgumi-sort/src/memory_probe.rs b/crates/fgumi-sort/src/memory_probe.rs index e70213aea..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 { @@ -488,7 +487,7 @@ impl MergeProbe { /// threshold. /// /// `pool_depths` is the `(raw_input, decompressed_input, buffer_pool)` - /// triple from [`SortWorkerPool::phase1_queue_depths`]. + /// triple from [`crate::worker_pool::SortWorkerPool::phase1_queue_depths`]. pub fn log_mid_with_depths( &mut self, pool_depths: (usize, usize, usize), diff --git a/crates/fgumi-sort/src/merge_headroom.rs b/crates/fgumi-sort/src/merge_headroom.rs deleted file mode 100644 index 06cf664e9..000000000 --- a/crates/fgumi-sort/src/merge_headroom.rs +++ /dev/null @@ -1,509 +0,0 @@ -//! Which wall the merge is against, and how much of its wall clock is -//! recoverable without doing less work. -//! -//! A spill-heavy sort's merge has exactly three limits, and every optimisation -//! either moves one or moves nothing: -//! -//! 1. **The serial consumer.** One thread owns the loser tree and emits records -//! in a strict global order, so its own CPU time is a hard floor. Nothing in -//! the pool can shorten it. -//! 2. **Worker capacity.** Total worker busy divided by the active thread count. -//! Read, decompress and output-compress are all parallel, so this is the floor -//! they collectively impose. -//! 3. **Neither.** If the loop is longer than both floors the difference is -//! coordination -- the consumer waiting for work that exists, or for work that -//! has not been started. -//! -//! Reporting this matters because the three imply completely different fixes and -//! are routinely confused. Measured on one cell, 8 threads against 16: at 8 the -//! consumer was 98% of the loop and the recoverable share was 1.8%, so a -//! scheduling change could not have paid however well designed; at 16 the same -//! build left 28% of the loop recoverable. Same engine, same data, opposite -//! advice -- and no way to tell from wall clock alone. - -/// Which limit the merge is actually against. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum Binding { - /// The serial consumer's own CPU is the largest floor. - Consumer, - /// Worker capacity is the largest floor. - Workers, - /// Both floors are well below the loop: the gap is coordination. - Coordination, -} - -impl Binding { - pub(crate) const fn label(self) -> &'static str { - match self { - Self::Consumer => "consumer serial CPU", - Self::Workers => "worker capacity", - Self::Coordination => "coordination (neither floor is close)", - } - } -} - -/// The measured merge loop against the two floors it cannot go below. -#[derive(Debug, Clone, Copy)] -pub(crate) struct MergeFloors { - /// Measured merge loop wall clock. - pub(crate) loop_secs: f64, - /// The consumer thread's own CPU: `loop - park - backpressure`, exact. - pub(crate) consumer_secs: f64, - /// Total worker busy seconds across every phase-2 worker. - pub(crate) worker_busy_secs: f64, - /// Active worker threads the busy total is spread over. - pub(crate) threads: usize, -} - -impl MergeFloors { - /// Worker capacity floor: busy time spread perfectly over the active threads. - /// - /// Zero threads is reported as no floor rather than a division by zero: a - /// merge with no workers is bounded by the consumer alone. - #[must_use] - pub(crate) fn worker_floor_secs(&self) -> f64 { - if self.threads == 0 { - return 0.0; - } - #[expect(clippy::cast_precision_loss, reason = "thread counts are tiny")] - let threads = self.threads as f64; - self.worker_busy_secs / threads - } - - /// The larger of the two floors -- the shortest this merge could run without - /// doing less work. - #[must_use] - pub(crate) fn floor_secs(&self) -> f64 { - self.consumer_secs.max(self.worker_floor_secs()) - } - - /// Which floor is closest to the measured loop. - /// - /// `Coordination` when neither floor is within 10% of the loop: at that point - /// the gap is larger than either limit and naming a "binding" floor would be - /// misleading. - #[must_use] - pub(crate) fn binding(&self) -> Binding { - let floor = self.floor_secs(); - if self.loop_secs > 0.0 && floor < 0.9 * self.loop_secs { - return Binding::Coordination; - } - if self.consumer_secs >= self.worker_floor_secs() { - Binding::Consumer - } else { - Binding::Workers - } - } - - /// Wall clock that could be recovered without reducing total work. - /// - /// Never negative: a floor above the measured loop means the estimate is off, - /// not that time is owed. - #[must_use] - pub(crate) fn recoverable_secs(&self) -> f64 { - (self.loop_secs - self.floor_secs()).max(0.0) - } - - /// [`Self::recoverable_secs`] as a share of the loop, 0.0 when the loop is empty. - #[must_use] - pub(crate) fn recoverable_share(&self) -> f64 { - if self.loop_secs <= 0.0 { - return 0.0; - } - self.recoverable_secs() / self.loop_secs - } -} - -/// Nanoseconds one `Instant::now()` / `elapsed()` pair costs on this machine. -/// -/// Measured rather than assumed because it varies by platform and by clock source -/// -- 15-35ns across the hosts this engine has been profiled on -- and because it -/// is the same order as the segments it is used to correct. A hard-coded constant -/// would silently stop being right. -/// -/// Timed with an empty body, so the result is exactly what a segment's interval -/// picks up on top of the work it is measuring. -#[must_use] -pub(crate) fn measure_clock_overhead_nanos() -> u64 { - const ITERATIONS: u64 = 4096; - let mut total: u64 = 0; - for _ in 0..ITERATIONS { - let started = std::time::Instant::now(); - total += u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX); - } - total / ITERATIONS -} - -/// Sampled seconds for each step the merge consumer takes per record, in loop -/// order. -/// -/// Sampled rather than timed on every record because the consumer runs at -/// 144-239 ns/record and one `Instant::now()` costs ~20-25 ns on aarch64: timing -/// five steps on every record would add more overhead than the steps it measures. -/// One in `MERGE_SAMPLE_INTERVAL` records is timed and the result scaled, so the -/// partition holds over the sample and the scale factor is reported alongside it. -#[derive(Debug, Clone, Copy, Default)] -pub(crate) struct ConsumerSample { - /// Reading the winner and publishing the predicted next source. - pub(crate) publish: f64, - /// Presenting the winning record: borrowed in place, or reassembled across a - /// block boundary. - pub(crate) present: f64, - /// Handing the record to the output writer. - pub(crate) write: f64, - /// Advancing that source, which is where the consumer parks or decompresses a - /// block itself. **Includes wait time**, so it is not pure CPU. - pub(crate) advance: f64, - /// Replaying the loser tree, or removing an exhausted source. - pub(crate) tree: f64, -} - -impl ConsumerSample { - /// Every segment multiplied by the sampling scale. - #[must_use] - pub(crate) fn scaled(self, scale: f64) -> Self { - Self { - publish: self.publish * scale, - present: self.present * scale, - write: self.write * scale, - advance: self.advance * scale, - tree: self.tree * scale, - } - } - - /// Every segment with its own measurement overhead removed. - /// - /// Each sampled segment is bracketed by an `Instant::now()` / `elapsed()` pair, - /// and that pair's cost lands *inside* the interval it is timing. On aarch64 it - /// runs 15-35ns against segments of 2-100ns, so the raw numbers can be more - /// clock than work: the first run of this instrument reported five segments - /// summing to 321.5s of a 189.3s loop, and a `next-source predict` row of - /// 21 ns/record whose true cost is 0.37 ns/record. - /// - /// The correction is one pair per segment per sample, because every segment is - /// timed on every sampled record. Clamped at zero: a segment cheaper than the - /// clock that measures it cannot be resolved by this method, and zero says that - /// where a negative would just look like a bug. - #[must_use] - pub(crate) fn corrected(self, samples: u64, overhead_nanos: u64) -> Self { - if samples == 0 || overhead_nanos == 0 { - return self; - } - #[expect( - clippy::cast_precision_loss, - reason = "sample and nanosecond counts stay below 2^52" - )] - let per_segment = (samples * overhead_nanos) as f64 / 1e9; - let fix = |v: f64| (v - per_segment).max(0.0); - Self { - publish: fix(self.publish), - present: fix(self.present), - write: fix(self.write), - advance: fix(self.advance), - tree: fix(self.tree), - } - } - - /// The five segments summed. - #[must_use] - pub(crate) fn total(self) -> f64 { - self.publish + self.present + self.write + self.advance + self.tree - } -} - -/// A scaled consumer sample checked against the loop it is supposed to partition. -/// -/// The check is the point. Three of these segments were already collected and -/// scaled before this type existed, and were reported without ever being summed -/// against the measured loop -- so there was no way to know whether they -/// accounted for most of it or a third of it. A partition whose residual is not -/// reported is an estimate wearing a budget's clothes. -#[derive(Debug, Clone, Copy)] -pub(crate) struct LoopPartition { - /// Scaled per-segment seconds. - pub(crate) segments: ConsumerSample, - /// Measured loop wall clock, exact. - pub(crate) loop_secs: f64, - /// Consumer park within `segments.advance`, measured exactly and separately. - pub(crate) park_secs: f64, -} - -impl LoopPartition { - /// Loop time the segments do not account for. - /// - /// Signed on purpose: a negative residual means the sample over-attributes - /// (clock overhead inside the timed regions, or a sampling bias), and hiding - /// that behind a clamp would hide the one number that says the partition is - /// unsound. - #[must_use] - pub(crate) fn unattributed_secs(self) -> f64 { - self.loop_secs - self.segments.total() - } - - /// `unattributed` as a share of the loop, for judging whether the partition is - /// trustworthy at a glance. - #[must_use] - pub(crate) fn unattributed_share(self) -> f64 { - if self.loop_secs <= 0.0 { - return 0.0; - } - self.unattributed_secs() / self.loop_secs - } - - /// The part of `advance` that was the consumer working rather than waiting. - /// - /// `advance` covers both parking on a block and decompressing one inline, and - /// only the second is CPU the consumer could shed. Park is measured exactly, so - /// subtracting it separates the two. Clamped at zero: park is exact while - /// `advance` is sampled, so at small sample counts the estimate can land below - /// it without meaning anything. - #[must_use] - pub(crate) fn advance_work_secs(self) -> f64 { - (self.segments.advance - self.park_secs).max(0.0) - } -} - -#[cfg(test)] -mod tests { - use super::{Binding, ConsumerSample, LoopPartition, MergeFloors}; - use rstest::rstest; - - /// Assertions use an epsilon rather than exact float equality: every field - /// here is a measured duration, so the arithmetic is inherently approximate. - const EPS: f64 = 1e-9; - - /// The measured t8 cell after the targeted-depth changes: consumer 186.5s of a - /// 190.0s loop, workers 1412.3s over 8 threads. The consumer binds and almost - /// nothing is recoverable -- which is what says a scheduling fix cannot pay - /// here however well designed. - #[test] - fn test_a_saturated_consumer_binds_and_leaves_almost_nothing() { - let f = MergeFloors { - loop_secs: 190.0, - consumer_secs: 186.5, - worker_busy_secs: 1412.3, - threads: 8, - }; - assert!((f.worker_floor_secs() - 176.537_5).abs() < 1e-4); - assert_eq!(f.binding(), Binding::Consumer); - assert!((f.floor_secs() - 186.5).abs() < EPS, "the larger floor is the consumer's"); - assert!((f.recoverable_secs() - 3.5).abs() < 1e-9); - assert!(f.recoverable_share() < 0.02, "under 2%: got {}", f.recoverable_share()); - } - - /// The same build at 16 threads: consumer 112.1s of a 156.0s loop, workers - /// 1420.2s over 16. Neither floor is near the loop, so the honest verdict is - /// coordination and 28% is recoverable. - #[test] - fn test_a_loop_far_above_both_floors_reads_as_coordination() { - let f = MergeFloors { - loop_secs: 156.0, - consumer_secs: 112.1, - worker_busy_secs: 1420.2, - threads: 16, - }; - assert!((f.worker_floor_secs() - 88.762_5).abs() < 1e-4); - assert_eq!(f.binding(), Binding::Coordination); - assert!((f.recoverable_secs() - 43.9).abs() < 1e-9); - assert!((f.recoverable_share() - 0.281_41).abs() < 1e-4, "got {}", f.recoverable_share()); - } - - /// Worker capacity can be the larger floor, and must be named when it is: - /// the fix for it (do less work per record) is nothing like the fix for a - /// slow consumer. - #[test] - fn test_worker_capacity_binds_when_it_is_the_larger_floor() { - let f = MergeFloors { - loop_secs: 200.0, - consumer_secs: 60.0, - worker_busy_secs: 1560.0, // 195s over 8 - threads: 8, - }; - assert!((f.worker_floor_secs() - 195.0).abs() < EPS); - assert_eq!(f.binding(), Binding::Workers); - assert!((f.recoverable_secs() - 5.0).abs() < EPS); - } - - /// A floor above the measured loop means the inputs disagree, which must clamp - /// to zero rather than report negative recoverable time. - #[test] - fn test_a_floor_above_the_loop_never_reports_negative_headroom() { - let f = MergeFloors { - loop_secs: 100.0, - consumer_secs: 130.0, - worker_busy_secs: 0.0, - threads: 8, - }; - assert!((f.recoverable_secs() - 0.0).abs() < EPS); - assert!((f.recoverable_share() - 0.0).abs() < EPS); - } - - /// Degenerate inputs must not divide by zero or panic: a merge with no workers - /// is bounded by its consumer, and an empty loop has no share to report. - #[rstest] - #[case::no_workers(100.0, 90.0, 500.0, 0)] - #[case::empty_loop(0.0, 0.0, 0.0, 8)] - fn test_degenerate_inputs_are_finite( - #[case] loop_secs: f64, - #[case] consumer_secs: f64, - #[case] worker_busy_secs: f64, - #[case] threads: usize, - ) { - let f = MergeFloors { loop_secs, consumer_secs, worker_busy_secs, threads }; - assert!(f.worker_floor_secs().is_finite()); - assert!(f.floor_secs().is_finite()); - assert!(f.recoverable_secs().is_finite()); - assert!((0.0..=1.0).contains(&f.recoverable_share())); - } - - /// The calibration must produce a plausible figure on whatever host the tests - /// run on. Not asserting a specific value -- that is hardware -- and zero is - /// valid: a coarse-resolution clock can round an empty body's interval to 0, - /// which `ConsumerSample::corrected` already handles as a no-op rather than a - /// bug. Only the upper bound is a real signal: a huge number would mean the loop - /// is measuring something other than the clock. - #[test] - fn test_clock_calibration_is_plausible() { - let ns = super::measure_clock_overhead_nanos(); - assert!(ns < 10_000, "{ns}ns per clock pair is not a clock read"); - } - - /// Every timed segment carries one clock-read pair inside its interval, so the - /// correction is per segment per sample -- not per sample. - /// - /// This is the difference between a usable partition and an unusable one. The - /// first run of this instrument reported segments summing to 321.5s of a 189.3s - /// loop (-70%), with `next-source predict` at 21 ns/record against a true cost - /// the loser-tree benchmark puts at 0.37 ns/record. Almost the entire row was - /// the clock. - #[test] - fn test_correction_removes_one_clock_pair_per_segment_per_sample() { - // 1000 samples, 20ns of clock per pair: each segment loses 20us. - let raw = ConsumerSample { - publish: 0.000_030, - present: 0.000_040, - write: 0.000_050, - advance: 0.000_060, - tree: 0.000_070, - }; - let c = raw.corrected(1000, 20); - assert!((c.publish - 0.000_010).abs() < 1e-12, "got {}", c.publish); - assert!((c.present - 0.000_020).abs() < 1e-12); - assert!((c.write - 0.000_030).abs() < 1e-12); - assert!((c.advance - 0.000_040).abs() < 1e-12); - assert!((c.tree - 0.000_050).abs() < 1e-12); - } - - /// A segment smaller than its own measurement overhead must clamp to zero - /// rather than go negative: that is the signature of a row that is entirely - /// clock, and reporting it as negative time would be worse than reporting none. - #[test] - fn test_a_segment_smaller_than_its_overhead_clamps_to_zero() { - let raw = ConsumerSample { publish: 0.000_005, ..ConsumerSample::default() }; - let c = raw.corrected(1000, 20); - assert!((c.publish - 0.0).abs() < 1e-12, "got {}", c.publish); - } - - /// No samples, or an unmeasurable clock, must leave the sample untouched rather - /// than divide by zero or subtract a guess. - #[rstest] - #[case::no_samples(0, 20)] - #[case::no_measurable_overhead(1000, 0)] - fn test_correction_is_a_no_op_without_a_calibration( - #[case] samples: u64, - #[case] overhead_nanos: u64, - ) { - let raw = ConsumerSample { advance: 1.5, ..ConsumerSample::default() }; - let c = raw.corrected(samples, overhead_nanos); - assert!((c.advance - 1.5).abs() < 1e-12); - } - - /// The residual is the whole reason this type exists, so it must be reported - /// and not clamped. Segments summing to less than the loop means real time is - /// unaccounted for. - #[test] - fn test_an_incomplete_partition_reports_the_time_it_cannot_account_for() { - let p = LoopPartition { - segments: ConsumerSample { - publish: 5.0, - present: 10.0, - write: 20.0, - advance: 50.0, - tree: 15.0, - }, - loop_secs: 156.0, - park_secs: 43.3, - }; - assert!((p.segments.total() - 100.0).abs() < EPS); - assert!((p.unattributed_secs() - 56.0).abs() < EPS, "156 - 100"); - assert!((p.unattributed_share() - 0.358_97).abs() < 1e-4); - } - - /// Over-attribution must surface as a negative residual rather than be hidden: - /// it means the sampled timings include their own clock overhead, or the sample - /// is biased toward expensive records. Either way the partition is unsound and - /// the reader has to be told. - #[test] - fn test_over_attribution_stays_negative_instead_of_clamping() { - let p = LoopPartition { - segments: ConsumerSample { - publish: 0.0, - present: 0.0, - write: 0.0, - advance: 120.0, - tree: 0.0, - }, - loop_secs: 100.0, - park_secs: 0.0, - }; - assert!(p.unattributed_secs() < 0.0, "got {}", p.unattributed_secs()); - assert!((p.unattributed_secs() + 20.0).abs() < EPS); - } - - /// `advance` mixes waiting with working and only the working half is sheddable, - /// so park -- which is exact -- is subtracted out. - #[test] - fn test_advance_separates_the_consumer_working_from_the_consumer_waiting() { - let p = LoopPartition { - segments: ConsumerSample { - publish: 0.0, - present: 0.0, - write: 0.0, - advance: 60.0, - tree: 0.0, - }, - loop_secs: 156.0, - park_secs: 43.3, - }; - assert!((p.advance_work_secs() - 16.7).abs() < 1e-9); - } - - /// Park is exact and `advance` is sampled, so the estimate can fall below it - /// without that meaning negative work. - #[test] - fn test_advance_work_clamps_when_the_sample_lands_below_exact_park() { - let p = LoopPartition { - segments: ConsumerSample { - publish: 0.0, - present: 0.0, - write: 0.0, - advance: 30.0, - tree: 0.0, - }, - loop_secs: 156.0, - park_secs: 43.3, - }; - assert!((p.advance_work_secs() - 0.0).abs() < EPS); - } - - /// Scaling multiplies every segment by the same factor and leaves the ratios - /// between them alone. - #[test] - fn test_scaling_preserves_the_shape_of_the_sample() { - let raw = - ConsumerSample { publish: 1.0, present: 2.0, write: 3.0, advance: 4.0, tree: 5.0 }; - let scaled = raw.scaled(1021.0); - assert!((scaled.total() - 15.0 * 1021.0).abs() < 1e-6); - assert!((scaled.publish / scaled.tree - raw.publish / raw.tree).abs() < 1e-9); - } -} diff --git a/crates/fgumi-sort/src/merge_phases.rs b/crates/fgumi-sort/src/merge_phases.rs index 760a01e41..eb0523907 100644 --- a/crates/fgumi-sort/src/merge_phases.rs +++ b/crates/fgumi-sort/src/merge_phases.rs @@ -85,61 +85,6 @@ pub(crate) struct MergePhaseCounters { pub(crate) spill_compress: ComponentCounter, } -/// Latency distributions for the four stages a block passes through, plus the -/// writer, plus how much scanning was wasted getting there. -/// -/// [`MergePhaseCounters`] already gives each stage a busy total and a block -/// count, which yields a *mean*. That is not enough to settle an argument: a -/// stage averaging 187 us could be uniform or bimodal with a long tail, and only -/// the tail explains why a worker is unavailable when the consumer needs one. -/// Output compression is 69% of all worker busy on the measured cell and had no -/// distribution at all; the writer had none either, so a jump in consumer -/// backpressure from 0.0s to 29.9s could not be attributed. -#[derive(Debug, Default)] -pub(crate) struct StageLatency { - /// One batched read of compressed blocks from a spill file. - pub(crate) read: crate::merge_trace::DurationHistogram, - /// Decompressing one spill block. - pub(crate) decompress: crate::merge_trace::DurationHistogram, - /// Compressing one output block. - pub(crate) output_compress: crate::merge_trace::DurationHistogram, - /// Compressing one Phase 1 spill block. - pub(crate) spill_compress: crate::merge_trace::DurationHistogram, - // The writer's own histograms live on `PermitPool`: the I/O writer thread - // already holds that `Arc`, and one pool per writer keeps the output - // writer's stats separate from a spill writer's. - /// Files a worker passed over before one gave it work, summed over all - /// productive scans. - /// - /// The scan tally is published only when a scan finds *nothing*, so the - /// files skipped on the way to a successful claim were invisible. On an - /// 89-way merge a worker can walk most of the file set before finding work, - /// and that walk is the cost this counts. - pub(crate) wasted_visits: std::sync::atomic::AtomicU64, - /// Scans that ended in a claim. The denominator for `wasted_visits`. - pub(crate) useful_claims: std::sync::atomic::AtomicU64, -} - -impl StageLatency { - /// Record a productive scan that skipped `skipped` files before claiming. - pub(crate) fn record_claim(&self, skipped: u64) { - use std::sync::atomic::Ordering; - self.wasted_visits.fetch_add(skipped, Ordering::Relaxed); - self.useful_claims.fetch_add(1, Ordering::Relaxed); - } - - /// Mean files visited fruitlessly per unit of work claimed. - #[expect(clippy::cast_precision_loss, reason = "counts stay far below 2^52")] - pub(crate) fn wasted_visits_per_claim(&self) -> f64 { - use std::sync::atomic::Ordering; - let claims = self.useful_claims.load(Ordering::Relaxed); - if claims == 0 { - return 0.0; - } - self.wasted_visits.load(Ordering::Relaxed) as f64 / claims as f64 - } -} - /// A read-only view of the counters, for logging. #[derive(Debug, Clone, Copy)] pub struct MergePhaseBreakdown { diff --git a/crates/fgumi-sort/src/merge_slots.rs b/crates/fgumi-sort/src/merge_slots.rs new file mode 100644 index 000000000..f9da8d125 --- /dev/null +++ b/crates/fgumi-sort/src/merge_slots.rs @@ -0,0 +1,937 @@ +//! Per-spill-file shared state for the unified-pipeline merge phase. +//! +//! `SortMergeSlot` is the per-spill-file shared state shuttled between +//! `SortSpillDecompress` (producer side: reads raw BGZF blocks from +//! disk, decompresses inline, pushes to the slot's bounded queue) and +//! `SortMerge` (consumer side: pops decompressed blocks from the +//! queue, parses records, drives the k-way merge) via +//! `Arc` clones. +//! +//! # Per-slot bounded queue design (v4 — see commit `9c39dea` / PR #389 and +//! `docs/design/sort-phase2-unification-deferral.md`) +//! +//! Each slot carries a bounded queue of decompressed BGZF blocks +//! (`PHASE2_DECOMP_CAP` entries). Backpressure lives here — the +//! producer is **non-blocking**: pushes only when the queue has +//! space; otherwise skips this slot and tries the next. The consumer +//! is also **non-blocking**: when the queue is empty but the slot is +//! not yet `queue_eof` it reports `WouldBlock` (see +//! `external.rs::slot_try_load_block`) and the cooperative `SortMerge` +//! step yields, so the framework re-dispatches it after the producer +//! has refilled the slot. No condvar, no parked thread. +//! +//! ## Why "queue has space OR consumer has a block OR slot EOF" is +//! the full state space +//! +//! For any slot at any wall-clock time, one of the following is true: +//! +//! 1. `decompressed.len() < PHASE2_DECOMP_CAP` — producer can push. +//! 2. `decompressed.len() > 0` — consumer can pop. +//! 3. `queue_eof == true` — slot is done; consumer returns EOF. +//! +//! The "consumer would block" path is reachable only when +//! `decompressed.len() == 0 && !queue_eof`, in which case the +//! producer will eventually flip the state to either (1) (push more +//! blocks) or (3) (set `queue_eof` on the read returning fewer +//! bytes than asked), and the next consumer dispatch observes it. No +//! "transient cap with all-workers-Skip" window. +//! +//! ## Atomic ordering (`decomp_error` / `queue_eof`) +//! +//! Producers' BOTH success and error paths MUST hold the +//! `decompressed` mutex while storing `queue_eof` (and the error +//! path additionally stores `decomp_error`). The consumer always +//! acquires the same mutex at the top of its poll loop. The +//! mutex's release-acquire chain establishes happens-before for +//! BOTH atomics simultaneously — regardless of which the consumer +//! loads first. Without this discipline, a stale `decomp_error` +//! load can race a fresh `queue_eof` load and produce silent +//! truncation. +//! +//! ## What used to live here, and why it's gone (v4 vs v3.1) +//! +//! Pre-v4 the slot also carried a `raw_blocks: Mutex` +//! queue and a `decomp_in_flight: AtomicUsize` counter. Pre-v4 +//! workers split work into a separate "read raw" step and a "claim +//! and decompress" step, with cap+gap-filler admission. That +//! design deadlocked at production scale because the framework's +//! drain protocol could Skip workers during a transient "all slots +//! at cap simultaneously" window. v4 collapses read-and-decompress +//! into one inline operation per worker per slot per `try_run`, +//! which for that **inline path** eliminates the `raw_blocks` queue, +//! the in-flight counter, the reorder buffer (a plain FIFO suffices, +//! because one worker reads per slot at a time via the reader lock, +//! so blocks decompress strictly in read order), and the gap-filler +//! escape. +//! +//! The reorder buffer and in-flight counter are **back** on this +//! struct — [`SortMergeSlot::reorder`] and [`SortMergeSlot::in_flight`] — +//! for the later **block-parallel** path, where several workers +//! decompress one file's blocks concurrently and results complete out +//! of order. The `bp_*` methods drive that path. Read the paragraph +//! above as "the inline path does not need them", not as "this module +//! does not have them": the state space argument above likewise covers +//! the inline path, while the block-parallel path additionally requires +//! `in_flight == 0` and a drained `reorder` before `queue_eof` finalizes. +//! +//! ## Status in this tree, and the OTHER Phase-2 implementation +//! +//! **`SortMergeSlot` has in-crate callers but no production caller.** The +//! arena engine gave it two: `external.rs::open_spill_slot`, which opens a +//! spill file as a slot, and `MergeDriver::from_slots`, which merges a set +//! of them. Both are reachable only from tests in this tree — their +//! production consumer is the typed-step `SortSpillDecompress` / +//! `SortMerge` pair, which arrives with `fgumi-pipeline-io` in a later +//! phase. Until then the crate's live Phase-2 is +//! `worker_pool::Phase2FileState`, driven through `RawExternalSorter::sort` +//! — which is what `fgumi sort` and `fgumi merge` actually call. +//! +//! That "no production driver yet" fact is now stated in four places — twice +//! here, once on `worker_pool::Phase2FileState`, and once in +//! `docs/design/sort-phase2-unification-deferral.md`, which is **canonical**. +//! All four have to change together when the pipeline steps land, so start +//! from the design doc. (This PR already had to correct that doc for a stale +//! method name, which is the failure mode.) +//! +//! `worker_pool::Phase2FileState` keeps its own reorder buffer and +//! in-flight counter because its single-reader/**multi-decompressor** +//! topology needs them, and it retains the gap-filler this module +//! dropped. Once the rewrite lands, this module becomes the production +//! Phase-2 for both standalone `fgumi sort` and the fused `runall` +//! sort, and `Phase2FileState` is retained as the +//! `RawExternalSorter::sort` library path and the `#[cfg(test)]` +//! parity oracle. (History: commit `9d6d7e9` / PR #395 and +//! `docs/design/sort-phase2-unification-deferral.md`.) + +use std::collections::VecDeque; +use std::fs::File; +use std::io::BufReader; + +// Concurrency primitives are sourced from `loom` under `--cfg loom` so the +// model-checking test (`tests/loom_merge_slots.rs`) exercises the REAL +// `SortMergeSlot` atomics/mutexes — every interleaving and memory reordering of +// the block-parallel EOF/in-flight/finalize protocol — instead of a hand-copied +// re-implementation. Under a normal build these are the `std` types verbatim. +#[cfg(loom)] +use loom::sync::Mutex; +#[cfg(loom)] +use loom::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +#[cfg(not(loom))] +use std::sync::Mutex; +#[cfg(not(loom))] +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + +use fgumi_bam_io::reorder::ReorderBuffer; + +use crate::codec::SpillCodec; + +/// Per-slot decompressed-block queue cap. Bounds in-flight +/// decompressed memory: `num_slots × PHASE2_DECOMP_CAP ×` per-entry size, +/// where each entry is a decompressed BGZF block or zstd frame ranging from +/// ~64 KB (BGZF) up to 256 KB (zstd worst case). +/// +/// Raised from 8 to 32 (increment 1a) to give the work-stealing decompressor +/// more runway ahead of the `Detached` merge, which was input-starved +/// (`MergeDiag stalls`) at the spill-heavy operating point. The `worker_pool.rs` +/// copy (the `cfg(test)`/library oracle path) is intentionally left at 8 — these +/// two are no longer equal. +/// +/// Hard cap — there is no admission escape. Producers skip a slot +/// when its queue is at this cap, returning to it on a later +/// `try_run` after the consumer has drained. It stays a per-slot, independent +/// bound (deadlock-safety is unchanged — see the module header). +/// +/// # Not the bound in force today +/// +/// `worker_pool.rs` declares a *different* constant of the same name +/// (`pub(crate) const PHASE2_DECOMP_CAP: usize = 8`) for its own Phase-2, and +/// that is the one governing this crate's live merge path: `fgumi sort` and +/// `fgumi merge` run through `RawExternalSorter`, i.e. the pool. This module's +/// callers (`open_spill_slot`, `MergeDriver::from_slots`) are reachable only +/// from tests until the typed-step `SortMerge` consumer lands. So +/// `fgumi_sort::PHASE2_DECOMP_CAP` resolves to 32 while the merge that actually +/// runs is bounded at 8. The two are deliberately unequal (see this constant's +/// history above); the confusable part is only which is in force, and that +/// flips when the pipeline steps land — see +/// `docs/design/sort-phase2-unification-deferral.md`, the canonical statement +/// of when that happens. +pub const PHASE2_DECOMP_CAP: usize = 32; + +/// Disk reader state for a single spill file. Mutex'd separately +/// from `decompressed` so the producer can hold the reader during +/// disk I/O without blocking the consumer's pop. +pub struct SortMergeReader { + /// Buffered file handle. + pub inner: BufReader, + /// Next per-slot sequence number to assign to a raw block read from this + /// file. Used ONLY by the block-parallel `SortSpillDecompress` path: each + /// raw block read under the reader lock is stamped with a monotonically + /// increasing sequence number so the out-of-order decompression results can + /// be reassembled in read order by the slot's [`SortMergeSlot::reorder`] + /// buffer. Because every read is serialized under the reader lock, the + /// sequence numbers are dense and assigned in strict read order. Unused by + /// the inline (file-granularity) path, which never reorders. + pub next_seq: u64, +} + +/// Per-spill-file shared state. +/// +/// Constructed by `SortAndSpill` (one per closed spill file via +/// `slots_for_chunk_files`), embedded in +/// `SortPhase1Event::SpillReady` as `Arc`, forwarded +/// verbatim by `SortSpillDecompress`, and finally installed in +/// `SortMerge`'s slot table. Drops when the last `Arc` is released +/// — typically after `SortMerge`'s merge driver is exhausted. +pub struct SortMergeSlot { + /// Stable identifier — the index this slot occupies in the + /// merge driver's source list. `SortMerge` orders sources by + /// `file_id` so the `LoserTree` tie-break for equal sort keys + /// is deterministic and matches the legacy chunk-files order. + pub file_id: u32, + /// Spill codec of this chunk's file, detected from the file magic when the + /// slot is opened (`slots_for_chunk_files`). The `SortSpillDecompress` step + /// reads it to decompress BGZF blocks or zstd frames; `reader` is already + /// positioned past any codec file-magic. + pub codec: SpillCodec, + /// Disk reader. Held only while reading raw bytes from disk. + /// One worker at a time per slot via `try_lock`. + pub reader: Mutex, + /// Bounded queue of decompressed blocks (each a BGZF block or a zstd + /// frame, per this slot's `codec`), FIFO. Pushed by the producer + /// (`SortSpillDecompress`) after read + inline decompress; popped by the + /// consumer (`SortMerge` via `slot_try_load_block`). Bounded at + /// `PHASE2_DECOMP_CAP`; producer skips when at cap. + pub decompressed: Mutex>>, + /// Set true once the producer detects EOF on the disk reader + /// AND has pushed any final batch of decompressed blocks to + /// `decompressed`. After this transition, the slot will never + /// receive another push. Consumer surfaces this as a clean EOF + /// when its poll finds `decompressed.is_empty() && queue_eof`. + /// + /// **Atomic ordering:** producer must hold the `decompressed` + /// mutex while storing this. Consumer reads it under the same + /// mutex. Mutex release-acquire creates happens-before. + pub queue_eof: AtomicBool, + /// Set true if BGZF decompression of a raw block fails. Consumer + /// surfaces this as `Err` rather than the silent `Ok(false)` that + /// an empty queue + `queue_eof` would look like. + /// + /// **Atomic ordering:** identical to `queue_eof`. Producer + /// stores while holding `decompressed`; consumer reads under + /// the same lock. + pub decomp_error: AtomicBool, + + // ── Block-parallel decompression state (file_granularity == false) ─────── + // + // These three fields are inert in the inline (file-granularity) path. In + // the block-parallel path multiple workers decompress one file's blocks + // concurrently: the READ is serialized under `reader` (sequence-tagged via + // `SortMergeReader::next_seq`), but the decompression happens outside the + // lock, so results complete out of order and are reassembled here. + /// Number of raw blocks that have been READ (under the reader lock) but not + /// yet inserted into `reorder`. Incremented under the reader lock when a + /// batch is read; decremented after the worker inserts that batch into + /// `reorder`. The slot may declare EOF only once this reaches zero — a + /// worker that observes reader-EOF must not truncate the merge while another + /// worker still holds an in-flight (read-but-undelivered) block. + pub in_flight: AtomicUsize, + /// Set true once a read returns fewer raw blocks than requested, i.e. the + /// disk reader reached a clean EOF. Distinct from `queue_eof`: `reader_eof` + /// means "no more blocks will be read", whereas `queue_eof` means "every + /// block has been read, decompressed, reassembled, and delivered to the + /// FIFO". `queue_eof` is set only when `reader_eof && in_flight == 0 && + /// reorder.is_empty()`. Stored under the reader lock (Release), read in the + /// finalize path (Acquire); being an atomic it does not participate in lock + /// ordering. + pub reader_eof: AtomicBool, + /// Per-slot reorder buffer that reassembles out-of-order decompression + /// results back into read (sequence) order before they are drained into the + /// FIFO. Lock order: acquire `reorder` BEFORE `decompressed` (never the + /// reverse); the reader lock, when held, is outermost. The consumer never + /// touches this — it only pops the in-order FIFO. + pub reorder: Mutex>>, +} + +impl SortMergeSlot { + /// Construct an empty slot for `file_id` backed by `reader` (positioned + /// past any codec file-magic) with the detected `codec`. + #[must_use] + pub fn new(file_id: u32, reader: BufReader, codec: SpillCodec) -> Self { + Self { + file_id, + codec, + reader: Mutex::new(SortMergeReader { inner: reader, next_seq: 0 }), + decompressed: Mutex::new(VecDeque::with_capacity(PHASE2_DECOMP_CAP)), + queue_eof: AtomicBool::new(false), + decomp_error: AtomicBool::new(false), + in_flight: AtomicUsize::new(0), + reader_eof: AtomicBool::new(false), + reorder: Mutex::new(ReorderBuffer::new()), + } + } + + /// Returns `true` when this slot has cleanly delivered all of its + /// output: `queue_eof` is set, the decompressed queue is empty, and + /// no decompression error was recorded. + /// + /// A slot whose `decomp_error` flag is set is **never** reported as + /// drained — an errored slot must surface as an error to its + /// consumer, not be mistaken for clean EOF. Callers using + /// `is_drained()` as a completion predicate must check + /// [`Self::has_error`] (or the source-level error flags) to + /// distinguish "still producing" from "failed". + /// + /// `decomp_error` is read while holding the `decompressed` mutex, honoring + /// the file's read-under-lock discipline. `queue_eof` is read *before* the + /// lock, on the fast path, and that is safe for the reverse reason: a + /// producer only stores `queue_eof` while holding this mutex, so observing + /// it `true` means the storing critical section has already released, and + /// the lock taken immediately below then establishes happens-before for + /// `decomp_error`. Observing it `false` unlocked can only be stale in the + /// direction that returns `false` early, which is the answer a + /// not-yet-finalized slot warrants anyway. + /// + /// # Panics + /// + /// Panics if `decompressed` mutex is poisoned. + #[must_use] + pub fn is_drained(&self) -> bool { + // Fast path: not yet at EOF — no need to take the lock. + if !self.queue_eof.load(Ordering::Acquire) { + return false; + } + let guard = self.decompressed.lock().expect("SortMergeSlot decompressed mutex poisoned"); + // Read `decomp_error` under the same lock the producer held when + // storing it. An errored slot is not a clean drain. + if self.decomp_error.load(Ordering::Acquire) { + return false; + } + guard.is_empty() + } + + /// Returns `true` if this slot recorded a decompression error. + /// + /// Read under the `decompressed` mutex to honor the file's + /// read-under-lock discipline. Use alongside [`Self::is_drained`] to + /// distinguish a clean EOF (`is_drained() == true`) from a failed + /// slot (`has_error() == true`), since `is_drained()` returns + /// `false` in both the "still producing" and "errored" cases. + /// + /// # Panics + /// + /// Panics if `decompressed` mutex is poisoned. + #[must_use] + pub fn has_error(&self) -> bool { + let _guard = self.decompressed.lock().expect("SortMergeSlot decompressed mutex poisoned"); + self.decomp_error.load(Ordering::Acquire) + } + + /// Gather probe statistics for this slot: `(pending_blocks, + /// pending_bytes, active)`. + /// + /// `pending_blocks` is the count of decompressed blocks waiting + /// for the consumer. `pending_bytes` is the sum of their byte + /// lengths. `active` is `!queue_eof` (the slot is still being + /// fed by the producer). + /// + /// # Panics + /// + /// Panics if `decompressed` mutex is poisoned. + #[must_use] + pub fn probe_stats(&self) -> (u64, u64, bool) { + let dec = self.decompressed.lock().expect("SortMergeSlot decompressed mutex poisoned"); + #[allow(clippy::cast_possible_truncation)] + let pending_blocks = dec.len() as u64; + let pending_bytes: u64 = dec.iter().map(|buf| buf.len() as u64).sum(); + drop(dec); + // `Acquire` for uniformity with every other `queue_eof` reader, even + // though this is a best-effort diagnostics probe. + let active = !self.queue_eof.load(Ordering::Acquire); + (pending_blocks, pending_bytes, active) + } + + /// Current number of decompressed blocks resident in the FIFO. Locks + /// `decompressed` for an O(1) `len()` read — used by the decompressor's + /// emptiest-first refill order (most-starved slot first). Cheaper than + /// [`Self::probe_stats`], which also sums per-block byte lengths. + /// + /// # Panics + /// + /// Panics if the `decompressed` mutex is poisoned. + #[must_use] + pub fn fifo_len(&self) -> usize { + self.decompressed.lock().expect("SortMergeSlot decompressed mutex poisoned").len() + } + + // ── Block-parallel decompression helpers (file_granularity == false) ───── + + /// Block-parallel admission control: may a worker read another batch of raw + /// blocks (whose first block would be tagged `next_seq`) into this slot's + /// reorder window? + /// + /// Combines [`ReorderBuffer::would_accept`] (the deadlock-free predicate the + /// pipeline uses elsewhere) with a hard `heap_bytes < window_budget` + /// backstop. The backstop is what actually bounds memory: `would_accept` + /// alone returns `true` (accept-all) while the buffer is *stuck* (the front + /// sequence not yet decompressed), which would let a slow straggler balloon + /// the window. The backstop is deadlock-safe **in this topology** because + /// reads are serialized and densely sequenced, so the front gap is ALWAYS an + /// already-read in-flight block (some worker is decompressing it) — never an + /// unread block that a new read would be required to fetch. Refusing new + /// reads therefore cannot wedge progress. + /// + /// `window_budget == 0` means unlimited. + /// + /// # Panics + /// + /// Panics if the `reorder` mutex is poisoned. + #[must_use] + pub fn bp_reorder_admits(&self, next_seq: u64, window_budget: u64) -> bool { + let rb = self.reorder.lock().expect("reorder mutex poisoned"); + if !rb.would_accept(next_seq, window_budget) { + return false; + } + window_budget == 0 || rb.heap_bytes() < window_budget + } + + /// Reserve `count` in-flight blocks just read under the reader lock. + /// + /// **Ordering requirement:** the caller MUST call this BEFORE + /// [`Self::bp_set_reader_eof`] for the EOF-carrying batch. Both stores are + /// `Release`; the lock-free finalizer in `drain_locked_and_finalize` + /// reads `reader_eof` (Acquire) then `in_flight` (Acquire) WITHOUT the + /// reader lock, so only this publish order guarantees that observing + /// `reader_eof == true` also makes this batch's `in_flight` increment + /// visible — otherwise the finalizer can declare a clean EOF that truncates + /// the EOF read's own still-in-flight block (loom-verified; see + /// `tests/loom_merge_slots.rs`). + pub(crate) fn bp_add_in_flight(&self, count: usize) { + self.in_flight.fetch_add(count, Ordering::Release); + } + + /// Mark the disk reader as having reached a clean EOF (called under the + /// reader lock when a read returns fewer blocks than requested). + /// + /// **Ordering requirement:** call this AFTER [`Self::bp_add_in_flight`] has + /// reserved the current batch — see that method's note. Publishing + /// `reader_eof` before the in-flight reservation reopens a truncation race. + pub(crate) fn bp_set_reader_eof(&self) { + self.reader_eof.store(true, Ordering::Release); + } + + /// Publish the accounting for a batch of `count` raw blocks just read under + /// the reader lock, in the one order that is correct: reserve the in-flight + /// blocks FIRST, then (on a short read) set `reader_eof`. + /// + /// This is the single source of truth for the publish order — both the + /// production worker (`SortSpillDecompress::try_fill_block_parallel_slot`) + /// and the loom model (`tests/loom_merge_slots.rs`) call it, so the ordering + /// is model-checked against the real code and the two cannot drift. + /// + /// **Why this order (loom-verified).** The lock-free finalizer in + /// `drain_locked_and_finalize` reads `reader_eof` (Acquire) then + /// `in_flight` (Acquire) WITHOUT the reader lock, so the reader lock does not + /// order this batch's accounting against it. Both stores are `Release`; a + /// finalizer that observes `reader_eof == true` therefore also observes this + /// (possibly EOF-carrying) batch's `in_flight` increment, and so cannot + /// finalize `queue_eof` while the batch is still in flight. Setting + /// `reader_eof` first would let a finalizer see `reader_eof == true` with + /// `in_flight == 0` after earlier blocks drained, finalizing a clean EOF that + /// silently truncates this batch's own block. Swapping the two lines makes + /// `tests/loom_merge_slots.rs` fail (as it did before the fix in `ac0a2ad`). + pub fn bp_commit_read(&self, count: usize, hit_eof: bool) { + self.bp_add_in_flight(count); + if hit_eof { + self.bp_set_reader_eof(); + } + } + + /// Number of additional decompressed blocks the FIFO can accept before it + /// hits [`PHASE2_DECOMP_CAP`]. + /// + /// # Panics + /// + /// Panics if the `decompressed` mutex is poisoned. + #[must_use] + pub fn bp_fifo_room(&self) -> usize { + let dec = self.decompressed.lock().expect("decompressed mutex poisoned"); + PHASE2_DECOMP_CAP.saturating_sub(dec.len()) + } + + /// Tracked reorder-window heap bytes (for tests / diagnostics). + /// + /// # Panics + /// + /// Panics if the `reorder` mutex is poisoned. + #[must_use] + pub fn bp_reorder_heap_bytes(&self) -> u64 { + self.reorder.lock().expect("reorder mutex poisoned").heap_bytes() + } + + /// Insert a freshly-decompressed batch `[start_seq, start_seq + count)` into + /// the reorder buffer, release the in-flight reservation, drain any now-ready + /// (in-order) blocks into the FIFO (bounded by [`PHASE2_DECOMP_CAP`]), and + /// finalize `queue_eof` if the slot is fully delivered. + /// + /// `count` is the reservation being released and **must** equal + /// `blocks.len()`. The one exception is the empty final read + /// (`count == 0`, `blocks` empty), which still calls through here so the + /// EOF can be finalized once the last in-flight block drains. Returns + /// `true` if it made progress (drained at least one block or finalized + /// EOF). + /// + /// Delivering fewer blocks than reserved is **silent truncation**, not a + /// tolerated shape: the short batch releases the reservation for sequence + /// numbers that were never inserted, so `in_flight` reaches zero and + /// `reorder` reports empty with those sequences simply absent. The slot + /// then finalizes a *clean* `queue_eof` — `is_drained()` true, + /// `has_error()` false — and the consumer treats a truncated spill file as + /// a complete one. Releasing *more* than was reserved is worse still: the + /// `fetch_sub` below wraps `in_flight` to `usize::MAX`, after which the + /// finalize predicate's `in_flight == 0` can never hold and the slot wedges + /// — the merge consumer polls `WouldBlock` forever rather than failing. + /// + /// Both are caller contract violations with no in-band signal, so they are + /// asserted here. A caller that decompresses fewer blocks than it reserved + /// must store `true` into [`SortMergeSlot::decomp_error`] while holding the + /// `decompressed` mutex — see the module header's ordering rules — rather + /// than short-batching this call. + /// + /// # Panics + /// + /// Panics if `blocks.len() != count`, or if `count` exceeds the outstanding + /// in-flight reservation. Both checks are **always on**, release included: + /// the failure they catch is silent record loss or a wedged merge, neither + /// of which has an in-band signal, and release is the configuration that + /// ships. The cost is two compares per *batch* — not per record — on a path + /// that decompresses up to [`PHASE2_DECOMP_CAP`] blocks per call. Both run + /// before any lock is taken, so a violation cannot poison the `reorder` or + /// `decompressed` mutex. + /// + /// Also panics if the `reorder` or `decompressed` mutex is already poisoned. + pub fn bp_insert_drain_finalize( + &self, + start_seq: u64, + blocks: Vec>, + count: usize, + ) -> bool { + assert_eq!( + blocks.len(), + count, + "bp_insert_drain_finalize would release a reservation of {count} while delivering \ + {} blocks; a short batch finalizes a clean EOF over the missing sequences", + blocks.len() + ); + // Load once: the message must report the value that actually failed the + // check, not a re-read that another worker may have moved in between. + let in_flight = self.in_flight.load(Ordering::Acquire); + assert!( + in_flight >= count, + "bp_insert_drain_finalize would release {count} in-flight blocks with only \ + {in_flight} reserved; the fetch_sub below would wrap and wedge the slot", + ); + let mut rb = self.reorder.lock().expect("reorder mutex poisoned"); + for (i, block) in blocks.into_iter().enumerate() { + let size = block.len(); + // `start_seq + i` cannot exceed the number of blocks read from one + // spill file, which is far below u64::MAX. + rb.insert_with_size(start_seq + i as u64, block, size); + } + // The reservation now lives in `reorder`, so release it. AcqRel so the + // finalize load below sees a consistent count. + self.in_flight.fetch_sub(count, Ordering::AcqRel); + self.drain_locked_and_finalize(&mut rb) + } + + /// Drain-only block-parallel pass: move any in-order ready blocks from the + /// reorder buffer into the FIFO (used when the FIFO had no room earlier, or + /// post-reader-EOF to flush stragglers other workers inserted) and finalize + /// `queue_eof`. Returns `true` if it made progress. + /// + /// # Panics + /// + /// Panics if the `reorder` or `decompressed` mutex is poisoned. + pub fn bp_drain_and_finalize(&self) -> bool { + let mut rb = self.reorder.lock().expect("reorder mutex poisoned"); + self.drain_locked_and_finalize(&mut rb) + } + + /// Shared drain + finalize, called with the `reorder` lock held. Acquires + /// `decompressed` (lock order: reorder → decompressed) so `queue_eof` is + /// stored under the same mutex the consumer reads it under. + fn drain_locked_and_finalize(&self, rb: &mut ReorderBuffer>) -> bool { + let mut dec = self.decompressed.lock().expect("decompressed mutex poisoned"); + let mut room = PHASE2_DECOMP_CAP.saturating_sub(dec.len()); + let mut drained = 0usize; + while room > 0 { + let Some(block) = rb.try_pop_next() else { break }; + dec.push_back(block); + room -= 1; + drained += 1; + } + // Finalize EOF: reader is exhausted, every read block has been inserted + // (in_flight == 0), and the reorder buffer is fully drained. Stored + // under `decompressed` so the consumer's release-acquire chain sees it. + let mut finalized = false; + if !self.queue_eof.load(Ordering::Acquire) + && self.reader_eof.load(Ordering::Acquire) + && self.in_flight.load(Ordering::Acquire) == 0 + && rb.is_empty() + { + self.queue_eof.store(true, Ordering::Release); + finalized = true; + } + drained > 0 || finalized + } +} + +// These exercise `SortMergeSlot` outside `loom::model`, which is illegal once +// the primitives are loom's, so they compile only in a normal (non-loom) build. +// The `--cfg loom` invocation runs `tests/loom_merge_slots.rs` exclusively. +#[cfg(all(test, not(loom)))] +mod tests { + use rstest::rstest; + + use super::*; + + fn empty_reader() -> BufReader { + BufReader::new(tempfile::tempfile().expect("create tempfile")) + } + + #[test] + fn new_slot_starts_empty_and_not_drained() { + let slot = SortMergeSlot::new(0, empty_reader(), SpillCodec::Bgzf); + assert_eq!(slot.file_id, 0); + assert!(!slot.is_drained(), "fresh slot not drained until queue_eof"); + let (pending_blocks, pending_bytes, active) = slot.probe_stats(); + assert_eq!(pending_blocks, 0); + assert_eq!(pending_bytes, 0); + assert!(active, "active until queue_eof flips"); + } + + #[test] + fn drained_when_queue_eof_and_empty() { + let slot = SortMergeSlot::new(0, empty_reader(), SpillCodec::Bgzf); + + // Push a decompressed block; not drained even if queue_eof. + slot.decompressed.lock().unwrap().push_back(vec![0xAB, 0xCD]); + slot.queue_eof.store(true, Ordering::Release); + assert!(!slot.is_drained(), "not drained while decompressed non-empty"); + + // Consume the block. + let popped = slot.decompressed.lock().unwrap().pop_front(); + assert_eq!(popped, Some(vec![0xAB, 0xCD])); + assert!(slot.is_drained(), "drained once queue empty AND queue_eof"); + } + + #[test] + fn fifo_len_reports_block_count() { + let slot = SortMergeSlot::new(0, empty_reader(), SpillCodec::Bgzf); + assert_eq!(slot.fifo_len(), 0); + slot.decompressed.lock().unwrap().push_back(vec![0u8; 10]); + slot.decompressed.lock().unwrap().push_back(vec![0u8; 20]); + assert_eq!(slot.fifo_len(), 2, "counts blocks, not bytes"); + } + + #[test] + fn probe_stats_counts_decompressed_only() { + let slot = SortMergeSlot::new(0, empty_reader(), SpillCodec::Bgzf); + slot.decompressed.lock().unwrap().push_back(vec![0u8; 1024]); + slot.decompressed.lock().unwrap().push_back(vec![0u8; 512]); + + let (pending_blocks, pending_bytes, active) = slot.probe_stats(); + assert_eq!(pending_blocks, 2); + assert_eq!(pending_bytes, 1024 + 512); + assert!(active, "not yet queue_eof"); + } + + /// Blocks reach the FIFO in READ order even when decompression completes + /// out of order — the reorder buffer's entire purpose. Hand-pushing into + /// `decompressed` and popping it back would assert that `VecDeque` is a + /// queue, which is true of `VecDeque` and says nothing about the slot. + #[test] + fn fifo_order() { + let slot = SortMergeSlot::new(7, empty_reader(), SpillCodec::Bgzf); + slot.bp_commit_read(4, true); + + // The later half of the file finishes decompressing first. Nothing may + // drain past the gap at seq 0, however ready seqs 2 and 3 are. + slot.bp_insert_drain_finalize(2, vec![vec![2], vec![3]], 2); + assert_eq!(slot.fifo_len(), 0, "nothing drains while seq 0 is missing"); + + // The front of the file lands; now the whole run drains, in read order. + slot.bp_insert_drain_finalize(0, vec![vec![0], vec![1]], 2); + let popped: Vec = slot.decompressed.lock().unwrap().drain(..).map(|b| b[0]).collect(); + assert_eq!(popped, vec![0, 1, 2, 3], "read order, not completion order"); + assert!(slot.queue_eof.load(Ordering::Acquire), "all 4 delivered ⇒ EOF finalizes"); + } + + /// `has_error` must report the flag through its own lock-and-load, not + /// merely echo the atomic. Asserting on `decomp_error` directly would test + /// `AtomicBool` rather than anything belonging to `SortMergeSlot`. + #[test] + fn has_error_reports_the_decomp_error_flag() { + let slot = SortMergeSlot::new(0, empty_reader(), SpillCodec::Bgzf); + assert!(!slot.has_error(), "a fresh slot has no error"); + slot.decomp_error.store(true, Ordering::Release); + assert!(slot.has_error(), "has_error must observe a stored decomp_error"); + } + + #[test] + fn errored_slot_is_not_drained() { + let slot = SortMergeSlot::new(0, empty_reader(), SpillCodec::Bgzf); + + // Mark EOF with an empty queue but a decompression error set: this + // must NOT be reported as a clean drain, otherwise a caller using + // is_drained() as a completion predicate would treat the failure as + // successful EOF and silently truncate the merge. + slot.queue_eof.store(true, Ordering::Release); + slot.decomp_error.store(true, Ordering::Release); + assert!(!slot.is_drained(), "errored slot must not be reported as drained"); + assert!(slot.has_error(), "has_error must surface the recorded decomp error"); + } + + #[test] + fn clean_eof_reports_no_error() { + let slot = SortMergeSlot::new(0, empty_reader(), SpillCodec::Bgzf); + slot.queue_eof.store(true, Ordering::Release); + assert!(slot.is_drained(), "clean empty + queue_eof is drained"); + assert!(!slot.has_error(), "clean drain has no error"); + } + + // ── Block-parallel helper tests ───────────────────────────────────────── + + /// A single in-order batch drains straight into the FIFO and (because the + /// reader is at EOF with no other in-flight blocks) finalizes `queue_eof`. + #[test] + fn bp_single_batch_drains_and_finalizes() { + let slot = SortMergeSlot::new(0, empty_reader(), SpillCodec::Bgzf); + slot.bp_add_in_flight(2); + slot.bp_set_reader_eof(); + let progressed = slot.bp_insert_drain_finalize(0, vec![vec![1u8], vec![2u8]], 2); + assert!(progressed); + assert!(slot.queue_eof.load(Ordering::Acquire), "reader_eof + drained ⇒ queue_eof"); + let mut dec = slot.decompressed.lock().unwrap(); + assert_eq!(dec.pop_front(), Some(vec![1u8])); + assert_eq!(dec.pop_front(), Some(vec![2u8])); + } + + /// EOF-with-stragglers: a worker reads the FINAL (short) batch and hits + /// reader-EOF while another worker still holds an earlier in-flight batch. + /// The EOF-observing worker must NOT finalize `queue_eof` (which would + /// truncate the merge); only after the straggler is delivered, in order, + /// does the slot finalize. No block is dropped or reordered. + #[test] + fn bp_eof_with_straggler_does_not_truncate() { + let slot = SortMergeSlot::new(0, empty_reader(), SpillCodec::Bgzf); + + // Worker A reserved seqs 0,1; worker B reserved seqs 2,3 on the final + // (short) read and set reader_eof. + slot.bp_add_in_flight(2); // A: seq 0,1 + slot.bp_add_in_flight(2); // B: seq 2,3 (final batch) + slot.bp_set_reader_eof(); + + // Worker B finishes decompressing FIRST and delivers seqs 2,3. Gap at + // 0,1 ⇒ nothing drains, and in_flight (A's 2) is non-zero ⇒ no EOF. + let b_progress = slot.bp_insert_drain_finalize(2, vec![vec![2u8], vec![3u8]], 2); + assert!(!b_progress, "straggler ahead-of-gap insert drains nothing and can't finalize"); + assert!(!slot.queue_eof.load(Ordering::Acquire), "must NOT declare EOF with A in flight"); + assert!(slot.decompressed.lock().unwrap().is_empty(), "nothing delivered yet"); + + // Worker A finishes and delivers seqs 0,1 ⇒ all four drain in order and + // EOF finalizes. + let a_progress = slot.bp_insert_drain_finalize(0, vec![vec![0u8], vec![1u8]], 2); + assert!(a_progress); + assert!(slot.queue_eof.load(Ordering::Acquire), "EOF finalizes once straggler delivered"); + + let mut dec = slot.decompressed.lock().unwrap(); + let order: Vec = std::iter::from_fn(|| dec.pop_front()).map(|b| b[0]).collect(); + assert_eq!(order, vec![0, 1, 2, 3], "blocks delivered in read order, none truncated"); + } + + /// The reorder window stays bounded when the front sequence straggles: a + /// worker keeps decompressing ahead-of-gap blocks, but `bp_reorder_admits` + /// refuses new reads once `heap_bytes` reaches the window budget, so the + /// buffer never balloons past `budget + one batch`. + #[test] + fn bp_reorder_window_is_bounded_under_straggler() { + const BLOCK: usize = 1024; + const BUDGET: u64 = 4 * BLOCK as u64; // 4 blocks + let slot = SortMergeSlot::new(0, empty_reader(), SpillCodec::Bgzf); + + // Stamp each block with its sequence number. Asserting only a count at + // the end would pass for a drain that dropped one block and duplicated + // another, or that delivered the window out of order. + let block_for = |seq: u64| { + let mut block = vec![0u8; BLOCK]; + block[0] = u8::try_from(seq).expect("the window budget keeps seqs far below 256"); + block + }; + + // Seq 0 is the straggler — it is reserved but never decompressed/inserted, + // so the buffer can never pop and keeps a permanent front gap. + slot.bp_add_in_flight(1); // seq 0 in flight forever (the straggler) + + // Workers race ahead delivering seqs 1,2,3,… as long as admission allows. + let mut next = 1u64; + let mut admitted = 0; + while slot.bp_reorder_admits(next, BUDGET) { + slot.bp_add_in_flight(1); + slot.bp_insert_drain_finalize(next, vec![block_for(next)], 1); + // Front gap at seq 0 ⇒ nothing drains. + assert!( + slot.decompressed.lock().unwrap().is_empty(), + "no block may reach the FIFO while the seq-0 gap is open", + ); + assert!( + slot.bp_reorder_heap_bytes() <= BUDGET, + "reorder window must stay within budget, got {} > {BUDGET}", + slot.bp_reorder_heap_bytes(), + ); + next += 1; + admitted += 1; + assert!(admitted < 1000, "admission must eventually backpressure, not loop forever"); + } + assert!(admitted > 0, "should admit at least some ahead-of-gap blocks"); + assert!( + slot.bp_reorder_heap_bytes() >= BUDGET.saturating_sub(BLOCK as u64), + "should have filled the window before backpressuring", + ); + + // Now the part that makes the EOF check mean something. Until here + // `reader_eof` was never set, so asserting `!queue_eof` inside the loop + // was vacuous — the finalize predicate could not fire for want of + // `reader_eof`, whatever the straggler did. Set it, and the predicate + // becomes falsifiable: the only thing still holding EOF back is seq 0, + // which is in flight and absent from `reorder`. + slot.bp_set_reader_eof(); + assert!(!slot.bp_drain_and_finalize(), "nothing can drain past the seq-0 gap"); + assert!( + !slot.queue_eof.load(Ordering::Acquire), + "must not finalize EOF while the straggler at seq 0 is still in flight", + ); + + // Deliver the straggler; now the whole window drains in order and EOF + // finalizes, confirming the blocks held behind the gap were retained + // rather than dropped. + slot.bp_insert_drain_finalize(0, vec![block_for(0)], 1); + assert!( + slot.queue_eof.load(Ordering::Acquire), + "EOF must finalize once the straggler lands and the window drains", + ); + let delivered: Vec = + slot.decompressed.lock().unwrap().drain(..).map(|block| block[0]).collect(); + let expected: Vec = (0..=u8::try_from(admitted) + .expect("admitted is bounded by the window budget")) + .collect(); + assert_eq!( + delivered, expected, + "the straggler and every admitted block must reach the FIFO exactly once, in read \ + order — not merely in the right quantity", + ); + } + + /// The drain stops at [`PHASE2_DECOMP_CAP`], so a slot can sit at + /// reader-EOF with `in_flight == 0` and *still* owe blocks that did not fit + /// in the FIFO. Finalizing `queue_eof` there would strand them: the + /// consumer stops at `is_drained()`, and the blocks left in `reorder` are + /// never popped. Completion instead depends on the driver calling + /// [`SortMergeSlot::bp_drain_and_finalize`] again once the consumer has made + /// room — the one drain path whose correctness lives outside this module. + /// + /// Two boundaries: a FIFO exactly at cap (`room == 0`, the loop never runs) + /// and one block short of it (`room` hits zero *inside* the loop, after a + /// partial delivery). Both must defer the finalize. + #[rstest] + #[case::fifo_exactly_at_cap(PHASE2_DECOMP_CAP, 0, false)] + #[case::fifo_one_short_of_cap(PHASE2_DECOMP_CAP - 1, 1, true)] + fn bp_full_fifo_defers_finalize_until_the_consumer_makes_room( + #[case] prefill: usize, + #[case] drained_now: usize, + #[case] progressed_now: bool, + ) { + const DEFERRED: [u8; 2] = [0xAA, 0xBB]; + let slot = SortMergeSlot::new(0, empty_reader(), SpillCodec::Bgzf); + + // Occupy the FIFO so the incoming batch cannot fully land. The prefill + // bytes stay below 0xAA, so they never alias the delivered blocks. + // A constant byte, distinct from DEFERRED: the prefill is only there to + // occupy room, and keying it to `i` would make the test panic the day + // PHASE2_DECOMP_CAP is tuned above 256. + { + let mut dec = slot.decompressed.lock().unwrap(); + for _ in 0..prefill { + dec.push_back(vec![0u8]); + } + } + + // The final read: two blocks, reader at EOF, nothing else in flight. + // Every finalize precondition except `reorder.is_empty()` now holds. + slot.bp_commit_read(2, true); + let progressed = + slot.bp_insert_drain_finalize(0, DEFERRED.iter().map(|&b| vec![b]).collect(), 2); + + assert_eq!(progressed, progressed_now, "progress is reported iff a block actually drained"); + assert_eq!(slot.fifo_len(), prefill + drained_now, "the FIFO fills only to the cap"); + assert!( + !slot.queue_eof.load(Ordering::Acquire), + "{} block(s) are still owed from `reorder`; finalizing EOF here strands them", + DEFERRED.len() - drained_now, + ); + assert!(!slot.is_drained(), "a slot that still owes blocks is not drained"); + assert_eq!(slot.bp_fifo_room(), 0, "the drain ran until the FIFO had no room left"); + + // The consumer pops; the driver's follow-up drain must deliver the + // remainder in read order and only then finalize. + { + let mut dec = slot.decompressed.lock().unwrap(); + for _ in 0..DEFERRED.len() { + dec.pop_front(); + } + } + assert!(slot.bp_drain_and_finalize(), "room freed ⇒ the deferred blocks drain"); + assert!( + slot.queue_eof.load(Ordering::Acquire), + "EOF finalizes once `reorder` empties and nothing is in flight", + ); + + let delivered: Vec> = slot.decompressed.lock().unwrap().drain(..).collect(); + assert_eq!( + delivered.len(), + prefill, + "prefill + {} delivered, less the {} the consumer popped: a dropped or duplicated \ + block moves this count", + DEFERRED.len(), + DEFERRED.len(), + ); + assert_eq!( + delivered[delivered.len() - DEFERRED.len()..], + DEFERRED.map(|b| vec![b]), + "the deferred blocks arrive behind the prefill, in read order", + ); + } + + /// The two caller-contract violations the assertions exist to catch: a + /// short batch (silent truncation) and an over-release (wraps `in_flight` + /// and wedges the slot). The assertions are always on, so these tests run + /// in release too — which is the configuration whose behaviour they pin. + #[test] + #[should_panic(expected = "while delivering")] + fn bp_short_batch_is_rejected() { + let slot = SortMergeSlot::new(0, empty_reader(), SpillCodec::Bgzf); + slot.bp_commit_read(2, true); + // Deliver 1 of the 2 reserved blocks: without the assertion this + // finalizes a clean queue_eof over the missing seq 1. + slot.bp_insert_drain_finalize(0, vec![vec![0xAAu8; 8]], 2); + } + + #[test] + #[should_panic(expected = "would wrap")] + fn bp_over_release_is_rejected() { + let slot = SortMergeSlot::new(0, empty_reader(), SpillCodec::Bgzf); + slot.bp_commit_read(1, true); + // Release 2 with only 1 reserved: the fetch_sub would wrap in_flight + // to usize::MAX and the slot could never finalize. + slot.bp_insert_drain_finalize(0, vec![vec![1u8; 8], vec![2u8; 8]], 2); + } +} diff --git a/crates/fgumi-sort/src/merge_stalls.rs b/crates/fgumi-sort/src/merge_stalls.rs index 8583c03c5..619b4449c 100644 --- a/crates/fgumi-sort/src/merge_stalls.rs +++ b/crates/fgumi-sort/src/merge_stalls.rs @@ -659,172 +659,6 @@ impl WakeLatencyReport { /// First bucket counted as a "deep" sleep (>= 320 µs). const DEEP_SLEEP_FIRST_BUCKET: usize = 5; -// ============================================================================ -// Consumer park decomposition -// ============================================================================ - -/// One consumer park, split into the stages of fetching the block it wanted. -/// -/// The four fields partition the park by construction — see [`split_park`]. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub(crate) struct ParkSegments { - /// Park spent before any worker claimed the needed block: the cost of - /// getting *a* worker onto it, whichever one arrives first. - pub to_claim: u64, - /// Park spent between the claim and the block being published: the read and - /// decompress work itself. - pub work: u64, - /// Park spent after the block was published, before the consumer resumed: - /// the consumer's own wake latency. - pub to_resume: u64, - /// Park not attributable to fetching this block — the honesty term. - pub unattributed: u64, -} - -/// Split a park into [`ParkSegments`]. -/// -/// `claim` and `publish` are `None` when that stamp did not land inside this -/// park: the block may have been claimed before the consumer parked, or not yet -/// claimed when it resumed. Both are ordinary, and neither may be allowed to -/// inflate a segment — an unstamped stage contributes to `unattributed` -/// instead, so the reader can see how much of park the model fails to explain. -/// -/// Stamps come from other threads and are only `Relaxed`, so they can be -/// observed out of order relative to `park_start`. Every subtraction therefore -/// saturates and every segment is clamped to what remains, which is what keeps -/// the sum exactly equal to `resume - park_start` rather than merely close. -pub(crate) fn split_park( - park_start: u64, - claim: Option, - publish: Option, - resume: u64, -) -> ParkSegments { - let total = resume.saturating_sub(park_start); - let mut seg = ParkSegments::default(); - let mut spent = 0u64; - - // Each stage is measured from the later of its own start and the previous - // stage's end, so a stamp that predates the park contributes zero rather - // than borrowing time from a neighbour. - let mut cursor = park_start; - if let Some(claim) = claim { - seg.to_claim = claim.saturating_sub(cursor).min(total - spent); - spent += seg.to_claim; - cursor = cursor.max(claim); - } - if let Some(publish) = publish { - seg.work = publish.saturating_sub(cursor).min(total - spent); - spent += seg.work; - cursor = cursor.max(publish); - seg.to_resume = resume.saturating_sub(cursor).min(total - spent); - spent += seg.to_resume; - } - seg.unattributed = total - spent; - seg -} - -/// Where the merge consumer's park time actually goes. -/// -/// Every earlier attempt to explain park was either a share of park *events* — -/// which hides a rare-but-long cause — or a sum over workers, which overcounts -/// because the consumer waits for the *first* worker to deliver while the sum -/// counts all of them. Measured on one 16-thread merge, the worker-side sum read -/// 96.7s of critical-path lag against 97.0s of park, then a change that removed -/// 62.9s of that lag moved park by 7.8s. These counters are consumer-side and -/// additive instead, so the segments cannot exceed the park they came from. -#[derive(Debug, Default)] -pub(crate) struct ParkAttribution { - to_claim_nanos: AtomicU64, - work_nanos: AtomicU64, - to_resume_nanos: AtomicU64, - unattributed_nanos: AtomicU64, - parks: AtomicU64, - /// Parks where no worker claimed the block during the park, so the whole - /// wait was for a stage that had already started or already finished. - unclaimed_parks: AtomicU64, - /// Blocks the awaited file had ready when the consumer resumed, summed. - /// - /// Divided by [`Self::parks`] this is pipeline depth on the critical path. A - /// mean near 1 means every block the consumer needs is fetched on demand and - /// costs a full round trip, which is a different problem from a slow round - /// trip and has different fixes. - ready_on_resume: AtomicU64, -} - -impl ParkAttribution { - /// Record one park, split by [`split_park`]. - pub(crate) fn record(&self, seg: ParkSegments, claimed: bool, ready_on_resume: u64) { - self.to_claim_nanos.fetch_add(seg.to_claim, Ordering::Relaxed); - self.work_nanos.fetch_add(seg.work, Ordering::Relaxed); - self.to_resume_nanos.fetch_add(seg.to_resume, Ordering::Relaxed); - self.unattributed_nanos.fetch_add(seg.unattributed, Ordering::Relaxed); - self.parks.fetch_add(1, Ordering::Relaxed); - if !claimed { - self.unclaimed_parks.fetch_add(1, Ordering::Relaxed); - } - self.ready_on_resume.fetch_add(ready_on_resume, Ordering::Relaxed); - } - - /// Snapshot for logging. - pub(crate) fn snapshot(&self) -> ParkAttributionReport { - ParkAttributionReport { - to_claim_nanos: self.to_claim_nanos.load(Ordering::Relaxed), - work_nanos: self.work_nanos.load(Ordering::Relaxed), - to_resume_nanos: self.to_resume_nanos.load(Ordering::Relaxed), - unattributed_nanos: self.unattributed_nanos.load(Ordering::Relaxed), - parks: self.parks.load(Ordering::Relaxed), - unclaimed_parks: self.unclaimed_parks.load(Ordering::Relaxed), - ready_on_resume: self.ready_on_resume.load(Ordering::Relaxed), - } - } -} - -/// Read-only view of [`ParkAttribution`]. -#[derive(Debug, Clone, Copy)] -pub struct ParkAttributionReport { - /// Park spent waiting for any worker to claim the needed block. - pub to_claim_nanos: u64, - /// Park spent on the read and decompress themselves. - pub work_nanos: u64, - /// Park spent after publication, waiting for the consumer's own wake. - pub to_resume_nanos: u64, - /// Park the model does not explain. - pub unattributed_nanos: u64, - /// Parks recorded. - pub parks: u64, - /// Parks in which no claim landed. - pub unclaimed_parks: u64, - /// Summed blocks ready on the awaited file at resume. - pub ready_on_resume: u64, -} - -impl ParkAttributionReport { - /// Whether anything was recorded. - #[must_use] - pub fn is_empty(self) -> bool { - self.parks == 0 - } - - /// Total park accounted for, which must match the exact park clock. - #[must_use] - pub fn total_nanos(self) -> u64 { - self.to_claim_nanos + self.work_nanos + self.to_resume_nanos + self.unattributed_nanos - } - - /// Mean blocks ready on the awaited file when the consumer resumed. - /// - /// Near 1.0 means the critical path has no pipeline depth: each block is - /// fetched on demand, so the merge pays one full round trip per block. - #[must_use] - #[expect(clippy::cast_precision_loss, reason = "counts are well within f64's exact range")] - pub fn mean_ready_on_resume(self) -> f64 { - if self.parks == 0 { - return 0.0; - } - self.ready_on_resume as f64 / self.parks as f64 - } -} - // ============================================================================ // Consumer side: where the merge loop blocks // ============================================================================ @@ -1377,141 +1211,6 @@ pub fn classify_stall( } } -// ============================================================================ -// Worker recruitment: why the consumer waits for a worker it already woke -// ============================================================================ - -/// Why nobody had already started on the block the consumer is about to wait for. -/// -/// The three cases are not degrees of one problem -- they imply different fixes, -/// which is why they are counted apart: -/// -/// - `SleeperAvailable`: capacity was idle and unused. Pure coordination loss; -/// the pool should have been on it. -/// - `AllBusyCompressing`: every worker was busy and output compression was -/// queued. Priority inversion -- `get_sort_priorities` puts `Compress` ahead -/// of `Phase2FileWork` whenever the compress queue is non-empty, so the block -/// the consumer is blocked on waits behind output work. -/// - `AllBusyMerging`: every worker was busy on merge work. Genuine capacity; -/// nothing to schedule better. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum ParkSupply { - /// At least one worker was parked. - SleeperAvailable, - /// Nobody parked, and output compression was queued. - AllBusyCompressing, - /// Nobody parked, and the compress queue was empty. - AllBusyMerging, -} - -impl ParkSupply { - /// Number of variants, so the census arrays indexed by [`Self::index`] - /// cannot fall behind the enum. Sizing them with a literal instead lets a - /// new variant compile and then index out of bounds inside - /// [`ParkSupplyCensus::record`] -- on a worker-adjacent path, where it - /// surfaces only as a panicked sort worker. - pub(crate) const COUNT: usize = 3; - - /// Stable index for the census arrays. - fn index(self) -> usize { - match self { - Self::SleeperAvailable => 0, - Self::AllBusyCompressing => 1, - Self::AllBusyMerging => Self::COUNT - 1, - } - } -} - -// Keeps `COUNT` honest: the match is exhaustive, so a fourth variant fails to -// compile here instead of indexing past the end of the census arrays at run -// time. -const _: () = { - const fn assert_count(supply: ParkSupply) -> usize { - match supply { - ParkSupply::SleeperAvailable => 0, - ParkSupply::AllBusyCompressing => 1, - ParkSupply::AllBusyMerging => ParkSupply::COUNT - 1, - } - } - assert!(assert_count(ParkSupply::AllBusyMerging) == 2); -}; - -/// Classify the pool's state at the instant the consumer parks. -pub(crate) fn classify_park_supply(parked_workers: usize, compress_depth: usize) -> ParkSupply { - // A sleeper outranks a queued compress deliberately: the two fixes are not - // interchangeable. Waking an idle worker costs nothing, while reordering - // priorities trades output throughput for merge latency, so a park with both - // conditions true belongs in the cheaper bucket. - if parked_workers > 0 { - ParkSupply::SleeperAvailable - } else if compress_depth > 0 { - ParkSupply::AllBusyCompressing - } else { - ParkSupply::AllBusyMerging - } -} - -/// Park time and park counts split by [`ParkSupply`]. -#[derive(Debug, Default)] -pub(crate) struct ParkSupplyCensus { - counts: [AtomicU64; ParkSupply::COUNT], - nanos: [AtomicU64; ParkSupply::COUNT], -} - -impl ParkSupplyCensus { - /// Attribute one park of `nanos` to `supply`. - pub(crate) fn record(&self, supply: ParkSupply, nanos: u64) { - let i = supply.index(); - self.counts[i].fetch_add(1, Ordering::Relaxed); - self.nanos[i].fetch_add(nanos, Ordering::Relaxed); - } - - /// Counts and nanoseconds per class, indexed as [`ParkSupply::index`]. - pub(crate) fn snapshot(&self) -> ParkSupplyReport { - ParkSupplyReport { - counts: std::array::from_fn(|i| self.counts[i].load(Ordering::Relaxed)), - nanos: std::array::from_fn(|i| self.nanos[i].load(Ordering::Relaxed)), - } - } -} - -/// Read-only view of [`ParkSupplyCensus`]. -#[derive(Debug, Clone, Copy, Default)] -pub struct ParkSupplyReport { - /// Parks per class: sleeper-available, all-busy-compressing, all-busy-merging. - pub counts: [u64; ParkSupply::COUNT], - /// Park nanoseconds per class, same order. - pub nanos: [u64; ParkSupply::COUNT], -} - -impl ParkSupplyReport { - /// Total parks censused. - #[must_use] - pub fn total_parks(self) -> u64 { - self.counts.iter().sum() - } - - /// Total park nanoseconds censused. - #[must_use] - pub fn total_nanos(self) -> u64 { - self.nanos.iter().sum() - } -} - -/// Takes a predicate rather than a slice so the wake path can read the atomics -/// directly: this runs on every wake -- millions of them -- and materializing a -/// `Vec` there would allocate on the hottest coordination path in the -/// merge. Production and tests call the same function. -pub(crate) fn first_parked_from(cursor: usize, width: usize, is_parked: F) -> Option -where - F: Fn(usize) -> bool, -{ - if width == 0 { - return None; - } - (0..width).map(|offset| (cursor + offset) % width).find(|&i| is_parked(i)) -} - #[cfg(test)] mod tests { use super::*; @@ -1699,78 +1398,6 @@ mod tests { assert_eq!(report.sleeps[7], 2000); } - /// The four segments must partition the park exactly, whatever order the - /// stamps arrive in. - /// - /// This is the property that makes the decomposition trustworthy where the - /// worker-side lag sum was not: a segment cannot be inflated without another - /// shrinking, so no single number can be read as larger than the park it - /// came from. - #[rstest::rstest] - // Ordinary case: claimed then published inside the park. - #[case(1_000, Some(1_100), Some(1_150), 1_200, (100, 50, 50, 0))] - // Claimed but not yet published when the consumer resumed: no work segment, - // and the remainder is unattributed rather than silently folded into work. - #[case(1_000, Some(1_100), None, 1_200, (100, 0, 0, 100))] - // Already claimed before the park: nothing to wait for a worker on. - #[case(1_000, None, Some(1_050), 1_200, (0, 50, 150, 0))] - // Neither stamp: the block was already available, so the park is not - // attributable to fetching it at all. - #[case(1_000, None, None, 1_200, (0, 0, 0, 200))] - // Both stamps predate the park: the block was already published when the - // consumer parked, so the whole wait is the consumer failing to notice -- - // charged to `to_resume`, not to a fetch that had already finished. Stamps - // arrive from other threads under `Relaxed`, so this ordering is reachable - // and must not yield a negative segment or a sum over the park. - #[case(1_000, Some(900), Some(950), 1_200, (0, 0, 200, 0))] - fn test_park_segments_always_partition_the_park( - #[case] park_start: u64, - #[case] claim: Option, - #[case] publish: Option, - #[case] resume: u64, - #[case] want: (u64, u64, u64, u64), - ) { - let (want_to_claim, want_work, want_to_resume, want_unattributed) = want; - let seg = split_park(park_start, claim, publish, resume); - assert_eq!(seg.to_claim, want_to_claim, "to_claim"); - assert_eq!(seg.work, want_work, "work"); - assert_eq!(seg.to_resume, want_to_resume, "to_resume"); - assert_eq!(seg.unattributed, want_unattributed, "unattributed"); - assert_eq!( - seg.to_claim + seg.work + seg.to_resume + seg.unattributed, - resume - park_start, - "segments must sum to the measured park" - ); - } - - proptest::proptest! { - /// Property: the segments partition the park for *any* arrangement of - /// stamps, including ones that predate the park or arrive reversed. - /// - /// This is the whole basis for trusting the decomposition over the - /// worker-side lag sum it replaces, and the case list cannot cover the - /// orderings that `Relaxed` loads across threads make reachable. - #[test] - fn prop_park_segments_partition_the_park( - park_start in 0u64..1_000_000, - span in 1u64..1_000_000, - claim_off in proptest::option::of(-500_000i64..1_500_000), - publish_off in proptest::option::of(-500_000i64..1_500_000), - ) { - let resume = park_start + span; - let stamp = |off: Option| { - off.map(|o| u64::try_from(i64::try_from(park_start).unwrap_or(i64::MAX) + o) - .unwrap_or(0)) - }; - let seg = split_park(park_start, stamp(claim_off), stamp(publish_off), resume); - proptest::prop_assert_eq!( - seg.to_claim + seg.work + seg.to_resume + seg.unattributed, - span, - "segments must sum to the park exactly" - ); - } - } - #[test] fn test_deep_sleep_wake_share() { let stats = WakeLatencyStats::default(); @@ -2136,91 +1763,4 @@ mod tests { let total = shares.starved + shares.unclaimed() + shares.in_progress(); assert!((total - 1.0).abs() < 1e-9); } - // ======================================================================== - // Worker recruitment - // ======================================================================== - - /// The three classes imply three different fixes, so the boundaries matter - /// more than the counts: a parked worker means coordination lost the block, - /// a queued compress means priority did, and neither means capacity did. - #[rstest] - #[case::one_sleeper_is_coordination_loss(1, 0, ParkSupply::SleeperAvailable)] - #[case::a_sleeper_outranks_a_full_compress_queue(3, 9, ParkSupply::SleeperAvailable)] - #[case::nobody_free_with_compress_queued_is_priority(0, 5, ParkSupply::AllBusyCompressing)] - #[case::nobody_free_and_nothing_queued_is_capacity(0, 0, ParkSupply::AllBusyMerging)] - fn test_park_supply_separates_coordination_from_priority_from_capacity( - #[case] parked_workers: usize, - #[case] compress_depth: usize, - #[case] expected: ParkSupply, - ) { - assert_eq!(classify_park_supply(parked_workers, compress_depth), expected); - } - - /// A sleeper is reported even when compression is also backed up, because the - /// fixes are not interchangeable: waking the sleeper costs nothing, whereas - /// reordering priorities trades output throughput for merge latency. - #[test] - fn test_a_sleeper_is_not_masked_by_a_busy_compress_queue() { - assert_eq!(classify_park_supply(1, 100), ParkSupply::SleeperAvailable); - } - - #[test] - fn test_census_attributes_parks_and_time_by_class() { - let census = ParkSupplyCensus::default(); - census.record(ParkSupply::SleeperAvailable, 100); - census.record(ParkSupply::SleeperAvailable, 50); - census.record(ParkSupply::AllBusyCompressing, 700); - census.record(ParkSupply::AllBusyMerging, 5); - - let report = census.snapshot(); - assert_eq!(report.counts, [2, 1, 1]); - assert_eq!(report.nanos, [150, 700, 5]); - assert_eq!(report.total_parks(), 4); - assert_eq!(report.total_nanos(), 855); - } - - #[test] - fn test_census_starts_empty() { - let report = ParkSupplyCensus::default().snapshot(); - assert_eq!(report.total_parks(), 0); - assert_eq!(report.total_nanos(), 0); - } - - #[test] - fn test_first_parked_starts_at_the_cursor_and_wraps() { - let parked = [false, false, true, false]; - let at = |i: usize| parked[i]; - assert_eq!(first_parked_from(0, 4, at), Some(2), "scans forward from the cursor"); - assert_eq!(first_parked_from(3, 4, at), Some(2), "and wraps around to find it"); - assert_eq!(first_parked_from(2, 4, at), Some(2), "the cursor itself counts"); - } - - /// The point of the counter this feeds: when nobody is parked a wake has - /// nowhere useful to go, and must be recorded as wasted rather than silently - /// spent on a running worker. - #[test] - fn test_first_parked_reports_none_when_every_worker_is_running() { - assert_eq!(first_parked_from(0, 3, |_| false), None); - } - - /// Wakes stay inside the active window, exactly as `wake_target` does -- a - /// capped worker will not take Phase 2 work, so waking it is the same lost - /// wake by another route. - #[test] - fn test_first_parked_ignores_workers_outside_the_active_limit() { - let parked = [false, false, true, true]; - let at = |i: usize| parked[i]; - assert_eq!( - first_parked_from(0, 2, at), - None, - "workers 2 and 3 are parked but capped out of Phase 2" - ); - assert_eq!(first_parked_from(0, 3, at), Some(2)); - } - - #[test] - fn test_first_parked_tolerates_an_empty_pool() { - assert_eq!(first_parked_from(0, 0, |_| true), None); - assert_eq!(first_parked_from(5, 0, |_| true), None, "a zero limit admits nobody"); - } } diff --git a/crates/fgumi-sort/src/merge_trace.rs b/crates/fgumi-sort/src/merge_trace.rs index c0ec2c8f8..6bb71924f 100644 --- a/crates/fgumi-sort/src/merge_trace.rs +++ b/crates/fgumi-sort/src/merge_trace.rs @@ -105,18 +105,6 @@ impl DurationHistogram { self.total_nanos.fetch_add(nanos, Ordering::Relaxed); } - /// Record one *count* observation (a queue depth, a run length) into a - /// histogram built for durations, scaling one unit to [`BLOCKS_TO_NANOS`] - /// so counts land in the microsecond lane the log2 bucketing is built - /// around. Without the scale every count below `BLOCKS_TO_NANOS` buckets to - /// zero and both the mean and the tail read as zero. - /// - /// Read the result back with [`HistogramReport::summary_blocks`], never - /// [`HistogramReport::summary`] -- the latter labels these as durations. - pub(crate) fn record_count(&self, count: u64) { - self.record(count.saturating_mul(BLOCKS_TO_NANOS)); - } - /// Time `f`, recording how long it took, and return its value. pub(crate) fn time(&self, f: impl FnOnce() -> R) -> R { let start = Instant::now(); @@ -546,28 +534,12 @@ impl ConsumerTraceStats { /// Record that the merge drew `blocks` consecutive blocks from one source. /// /// This is the test for whether the merge has a *hot file* worth steering - /// workers toward, and it must be read weighted by **blocks, not by runs** -- - /// the two give opposite answers. - /// - /// By run count the merge looks perfectly interleaved: the measured median run - /// is 1 block, so most of the time the loser tree does pull the next record - /// from some other source. An earlier version of this comment stopped there and - /// concluded that any fix framed as "prefetch the file the consumer needs next" - /// was answering a question the workload does not pose. - /// - /// Weighted by blocks the same histogram says the opposite. On a measured - /// 44-way merge (`n=167624 total=5306698 mean=31.7 p50=1 p90=2 p99=512`), if - /// the bottom 90% of runs average two blocks or fewer they account for at most - /// 302k of 5.31M blocks -- so **at least 94% of all blocks come from the top - /// decile of runs**, which are hundreds of consecutive blocks from one source. - /// Steering read-ahead at the source the merge is draining, and at the one it - /// will drain next, is worth 2.8-5.2% and 2.7% of merge wall respectively. - /// - /// Both readings are true; only the block-weighted one predicted a result. It - /// is the same count-versus-time trap that made the park census point at a - /// scheduling change that then measured 0.06%. A short floor for the frequent - /// one-block runs still has to exist -- what does not work is spreading depth - /// uniformly, which has been measured negative four separate times. + /// workers toward. A run length of 1 means it does not: the loser tree + /// interleaves records from every run at once, so by the time one source's + /// block is exhausted the merge has already pulled from many others. + /// Demand is spread evenly and continuously across all K files, and any fix + /// framed as "prefetch the file the consumer needs next" is answering a + /// question the workload does not pose. /// /// Reuses the duration histogram's bucketing for a count: run lengths are as /// heavy-tailed as the timings and want the same log2 treatment. One block @@ -577,7 +549,7 @@ impl ConsumerTraceStats { /// Read the result back with [`HistogramReport::summary_blocks`], never /// [`HistogramReport::summary`] -- the latter labels these as durations. pub(crate) fn record_source_run(&self, blocks: u64) { - self.source_run_length.record_count(blocks); + self.source_run_length.record(blocks.saturating_mul(BLOCKS_TO_NANOS)); } pub(crate) fn snapshot(&self) -> ConsumerTraceReport { @@ -732,44 +704,6 @@ mod tests { assert_eq!(report.percentile_micros(1.00), 8192); } - /// A queue depth is a single-digit count. Recording it raw drops every - /// observation into bucket 0 and the report reads zero -- the merge-trace - /// bug this guards. [`DurationHistogram::record_count`] scales one unit to - /// [`BLOCKS_TO_NANOS`] first, so the mean and the tail read as the depths - /// themselves. - #[test] - fn test_record_count_scales_small_counts_out_of_the_zero_bucket() { - let depths = [2_u64, 3, 4, 9]; - - // Raw counts collapse: every depth below `BLOCKS_TO_NANOS` buckets to - // zero, so the mean rounds to 0.0 and the tail is a flat zero. - let raw = DurationHistogram::default(); - for &depth in &depths { - raw.record(depth); - } - let raw = raw.snapshot(); - assert_eq!(raw.count, 4); - assert_eq!(format!("{:.1}", raw.mean_micros()), "0.0", "raw counts read as zero"); - assert_eq!(raw.percentile_micros(1.0), 0, "raw counts collapse to a zero tail"); - - // Scaled through `record_count`, the same depths read back as blocks. - let scaled = DurationHistogram::default(); - for &depth in &depths { - scaled.record_count(depth); - } - let scaled = scaled.snapshot(); - assert_eq!(scaled.count, 4); - // (2 + 3 + 4 + 9) / 4 = 4.5 blocks, reported exactly. - assert!((scaled.mean_micros() - 4.5).abs() < 0.01, "{}", scaled.mean_micros()); - // Bucketed as if microseconds, so the tail reads as blocks (log2 floor). - assert_eq!(scaled.percentile_micros(1.0), 8); - assert_eq!( - scaled.summary_blocks(), - "n=4 total=18 blocks mean=4.5 p50=2 p90=8 p99=8", - "the depth line must read as blocks" - ); - } - /// A mean alone cannot tell these two apart, and they have different /// causes: one is uniformly slow, the other is fast with a bad tail. #[test] @@ -913,46 +847,6 @@ mod tests { assert_eq!(report.park_by_state[AwaitedState::Decompressing as usize].count, 3); } - /// Park time per state, not just park count. - /// - /// The two disagree and only one is a cost: the park-supply census had a class - /// that was 41% of parks and 13% of park time, and reading the counts picked a - /// fix that then measured at 0.06%. These nanoseconds were collected from the - /// start and reported one log level down, so every experiment in that - /// investigation had the answer and none printed it. - #[test] - fn test_consumer_trace_carries_park_time_per_state_not_only_counts() { - let stats = ConsumerTraceStats::default(); - // Many cheap parks in one state, one expensive park in another: the count - // majority and the time majority are deliberately opposite. - for _ in 0..10 { - stats.record_park(AwaitedState::Decompressing, 1_000, 1); - } - stats.record_park(AwaitedState::RawQueued, 500_000, 0); - - let report = stats.snapshot(); - let decomp = report.park_by_state[AwaitedState::Decompressing as usize]; - let raw = report.park_by_state[AwaitedState::RawQueued as usize]; - - assert_eq!(decomp.count, 10, "counts say Decompressing dominates 10:1"); - assert_eq!(raw.count, 1); - assert!( - (decomp.total_secs() - 10e-6).abs() < 1e-12, - "10 parks of 1us is 10us, got {}", - decomp.total_secs() - ); - assert!( - (raw.total_secs() - 500e-6).abs() < 1e-12, - "one park of 500us is 500us, got {}", - raw.total_secs() - ); - assert!( - raw.total_secs() > decomp.total_secs() * 40.0, - "time says RawQueued dominates 50:1 -- the opposite of the counts, which is \ - precisely the reading that must be available at info level" - ); - } - #[test] fn test_consumer_trace_reports_depth_beyond_one_worker() { // `none` and `exactly one` cannot distinguish a pool running two deep diff --git a/crates/fgumi-sort/src/phase1_keys.rs b/crates/fgumi-sort/src/phase1_keys.rs deleted file mode 100644 index d07b29f18..000000000 --- a/crates/fgumi-sort/src/phase1_keys.rs +++ /dev/null @@ -1,1021 +0,0 @@ -//! Deferred sort-key extraction — the ingest thread's largest serial cost, -//! moved onto the worker pool. -//! -//! Phase 1's ingest thread is a serial consumer, and the floor line names it as -//! the phase's binding limit: on a 16-thread whole-genome sort it is 137.5s of a -//! 145.4s read span, against a worker-capacity floor of 22.2s. Of that 137.5s, -//! **93.6s is `extract_template_key_inline`** — one aux-tag scan, one name hash -//! and one unclipped-5' computation per record, at 120 ns each over 780M -//! records. It is pure per-record CPU with no ordering requirement, so the only -//! reason it sat on the serial thread was that it ran where the record bytes -//! happened to be in hand. -//! -//! This module moves it. The ingest thread pushes a record's bytes into the -//! arena and records its extent, leaving the ref's key unset; batches of extents -//! are handed to the pool as [`KeyExtractionJob`]s, and the filled keys come -//! back over a channel to be spliced into the ref array. The pool has the -//! capacity to absorb it — 22.2s of worker busy across a 145s span — so the work -//! disappears rather than relocating. -//! -//! # Why batches are cut at arena segment boundaries -//! -//! A worker reads record bytes straight out of the arena while the ingest thread -//! is still appending to it. That is only sound against a segment the writer has -//! finished with, which is why [`SegmentedBuf`](crate::segmented_buf::SegmentedBuf) -//! splits its sealed segments from its live one: a batch names exactly one -//! sealed segment, holds it alive by `Arc`, and the writer never touches it -//! again. Records still in the live segment keep their keys unextracted until it -//! seals, or until the chunk barrier, whichever comes first. - -use std::sync::Arc; -use std::sync::mpsc::{Receiver, Sender}; - -use fgumi_raw_bam::{RawRecordView, SamTag}; - -use crate::external::{ - DroppedLaneViolation, LibraryLookup, TemplateKeyVariant, extract_template_key_inline, - verify_dropped_lanes, -}; -use crate::inline::{TEMPLATE_HEADER_SIZE, TemplateKey, TemplateLaneKey, TemplateRecordBuffer}; -use crate::worker_pool::SortWorkerPool; - -/// Everything a deferred batch needs to reproduce the serial extraction exactly. -/// -/// Built once per sort and shared by `Arc`, so a batch carries one pointer -/// rather than a copy of the library table and hasher. -pub(crate) struct KeyContext { - /// Read-group to library-ordinal table, from the BAM header. - pub(crate) lib_lookup: LibraryLookup, - /// The cell-barcode tag to fold into the key, when the sort was asked for one. - pub(crate) cell_tag: Option, - /// Fixed-seed hasher for cell-barcode values. **Never reseed this** — it - /// feeds the sort key, so a random seed would break byte-identity. - pub(crate) cb_hasher: ahash::RandomState, - /// The first record's full key: the baseline every later record's dropped - /// lanes are verified against. - pub(crate) first_key: TemplateKey, - /// Which lanes the chosen narrowed key retains. - pub(crate) variant: TemplateKeyVariant, -} - -/// A dropped-lane violation, tagged with the record that carried it. -/// -/// Carries the record's index within the whole chunk (not within the batch) so -/// the ingest thread can report the *first* offending record deterministically, -/// however the batches happened to be scheduled. -pub(crate) struct KeyViolation { - /// Index of the offending record in the chunk's ref array. - pub(crate) ref_index: usize, - /// Which lane disagreed with the first record. - pub(crate) violation: DroppedLaneViolation, - /// The offending record's read name, for the error message. - pub(crate) name: String, -} - -/// One completed batch of keys, addressed by where they belong. -pub(crate) struct KeyBatchResult { - /// Index of this batch's first record in the chunk's ref array. - pub(crate) first_ref: usize, - /// Keys for `first_ref .. first_ref + keys.len()`, in ref order. - pub(crate) keys: Box<[K]>, - /// The lowest-indexed violation this batch found, if any. - pub(crate) violation: Option, -} - -/// A unit of deferred key extraction that any pool worker can run. -/// -/// Type-erased on purpose: the worker pool is not generic over the sort key, and -/// making it so would thread `K` through every step, queue and stats array for -/// the benefit of one step. The batch carries its own sender instead, so the -/// pool only ever sees "run this". -pub(crate) trait KeyExtractionJob: Send { - /// Extract every key in the batch and publish the result. - /// - /// Consumes the job: a batch runs exactly once. - fn run(self: Box); -} - -/// Bytes ahead of the current record to software-prefetch while extracting keys. -/// -/// A batch reads record bodies out of a **sealed** arena segment — bytes the -/// ingest thread wrote up to a segment ago and which are long gone from cache by -/// the time a worker gets to them. The scan is therefore latency-bound on cache -/// misses rather than compute-bound, and prefetching ahead hides them. -/// -/// 2 KiB is `main-runall`'s measured value for the same shape of scan (a cold -/// ~2.6 GiB arena at ~220 B/record): 2 KiB gave the best speedup (~15%), ≤1 KiB -/// was negligible (too little lead time) and 4 KiB matched 2 KiB. At our ~250 -/// B/record that is ~8 records of lead. -const KEY_PREFETCH_DISTANCE: usize = 2048; - -/// Software-prefetch (read, into L1, temporal) the cache line containing `byte`. -/// -/// SAFETY: both `prfm pldl1keep` (`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`, so the pointer is -/// valid to name. Cribbed from `main-runall`'s `prefetch_read_l1`; a no-op on -/// other architectures. -#[inline] -fn prefetch_read_l1(byte: &u8) { - let ptr: *const u8 = byte; - #[cfg(target_arch = "aarch64")] - #[allow(unsafe_code)] - // SAFETY: 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: 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; the hint is a no-op elsewhere - } -} - -/// Extract one record's full key and check the lanes the chosen variant drops. -/// -/// The single extraction implementation, shared by the batched path and the -/// immediate one, so the two cannot drift into producing different keys. -fn extract_one(ctx: &KeyContext, bam: &[u8]) -> (TemplateKey, Option) { - let full = extract_template_key_inline(bam, &ctx.lib_lookup, ctx.cell_tag, &ctx.cb_hasher); - let violation = verify_dropped_lanes(&ctx.first_key, &full, ctx.variant); - (full, violation) -} - -/// A batch of template-coordinate records awaiting key extraction. -pub(crate) struct TemplateKeyBatch { - /// Shared extraction context. - pub(crate) ctx: Arc, - /// The sealed arena segment holding every record body in this batch. - pub(crate) segment: Arc>, - /// Byte offset of `segment`'s first byte in the arena's global address space. - pub(crate) segment_base: u64, - /// Index of this batch's first record in the chunk's ref array. - pub(crate) first_ref: usize, - /// Per record, in ref order: the global offset of its BAM bytes (the inline - /// header already skipped) and their length. - pub(crate) extents: Box<[(u64, u32)]>, - /// Where the filled keys are published. - pub(crate) results: Sender>, -} - -impl KeyExtractionJob for TemplateKeyBatch -where - K: Send, -{ - fn run(self: Box) { - let me = *self; - let ctx = &me.ctx; - let mut keys: Vec = Vec::with_capacity(me.extents.len()); - let mut violation: Option = None; - - for (i, &(offset, len)) in me.extents.iter().enumerate() { - let local = usize::try_from(offset - me.segment_base) - .expect("record offset precedes its own segment"); - // Pull the bytes this scan will reach shortly into L1 while the - // current record is being parsed. - if let Some(ahead) = me.segment.get(local + KEY_PREFETCH_DISTANCE) { - prefetch_read_l1(ahead); - } - let bam = &me.segment[local..local + len as usize]; - - let (full, lane) = extract_one(ctx, bam); - - // First violation wins: batches are scheduled in whatever order the - // pool picks them up, so only the lowest index is reproducible. - if violation.is_none() - && let Some(v) = lane - { - let name = - String::from_utf8_lossy(RawRecordView::new(bam).read_name()).into_owned(); - violation = Some(KeyViolation { ref_index: me.first_ref + i, violation: v, name }); - } - - keys.push(K::from_full(&full)); - } - - // A closed channel means the ingest thread has already failed and gone - // away; there is nobody left to report to and nothing to clean up. - drop(me.results.send(KeyBatchResult { - first_ref: me.first_ref, - keys: keys.into_boxed_slice(), - violation, - })); - } -} - -// ============================================================================ -// Ingest-side driver -// ============================================================================ - -/// What deferral actually achieved over a whole sort. -/// -/// Reported because a deferral that silently stops engaging is invisible any -/// other way: the keys are correct whether a worker extracted them or the -/// barrier did, so no output check, equivalence test or record count can tell -/// the difference. Only this ratio can. -#[derive(Debug, Clone, Copy, Default)] -pub(crate) struct KeyOverlapCensus { - /// Records whose keys were extracted while the ingest thread kept reading. - pub(crate) overlapped_records: u64, - /// Records whose keys the ingest thread had to wait at a barrier for. - pub(crate) barrier_records: u64, - /// Exact seconds spent cutting and handing out batches. - pub(crate) dispatch_secs: f64, - /// Exact seconds spent waiting at barriers for keys. - pub(crate) barrier_secs: f64, -} - -impl KeyOverlapCensus { - /// Share of records keyed off the serial thread's critical path, as a - /// percentage. `None` when nothing was deferred at all. - pub(crate) fn overlap_percent(self) -> Option { - let total = self.total_records(); - if total == 0 { - return None; - } - #[allow(clippy::cast_precision_loss)] - Some(100.0 * self.overlapped_records as f64 / total as f64) - } - - /// Total records accounted for. - pub(crate) fn total_records(self) -> u64 { - self.overlapped_records + self.barrier_records - } -} - -/// Records per deferred batch. -/// -/// Big enough that the per-batch cost (one `Box`, one queue push, one channel -/// send, 48 KiB of extents) is noise against ~4k extractions — each batch is -/// still ~0.5 ms of work — and small enough that a batch's key array (4k x up to -/// 56 B = 224 KiB) stays inside L2, so splicing it back into the ref array hits -/// warm cache rather than a second trip to memory. Batches are also cut at arena -/// segment boundaries, so this is an upper bound rather than an exact size. -const KEY_BATCH_RECORDS: usize = 4 * 1024; - -/// Drives deferred key extraction for one sort: cuts batches, hands them to the -/// pool, and splices the results back into the buffer's ref array. -/// -/// One instance spans the whole sort and is [`reset`](Self::reset) per chunk, so -/// the channel and its context are allocated once rather than per spill. -pub(crate) struct DeferredKeys { - /// Shared extraction context, cloned into every batch by `Arc`. - ctx: Arc, - /// Kept alive for the whole sort so the receiver never disconnects while - /// batches are outstanding. - tx: Sender>, - /// Completed batches come back here. - rx: Receiver>, - /// Refs `0..dispatched` have been handed to some batch. - dispatched: usize, - /// Batches handed out whose result has not been absorbed yet. - outstanding: usize, - /// Refs whose key has actually been written back. - filled: usize, - /// Sealed arena segment count as of the last dispatch check. - sealed_seen: usize, - /// Lowest-indexed dropped-lane violation seen across all batches. - violation: Option, - /// Records dispatched *before* their chunk's barrier, summed over the whole - /// sort. Not cleared by [`reset`](Self::reset). - overlapped_records: u64, - /// Records still unbatched when their chunk's barrier ran, so their keys - /// were extracted with the ingest thread waiting. Not cleared by `reset`. - barrier_records: u64, - /// Seconds the ingest thread spent cutting and handing out batches. - /// - /// Timed exactly rather than sampled: dispatch fires once per sealed arena - /// segment — a few thousand times across a whole sort, against hundreds of - /// millions of records — so a 1-in-N per-record sample would be estimating a - /// rare event with a method built for a uniform one. A few thousand clock - /// pairs is unmeasurable overhead. - dispatch_secs: f64, - /// Seconds the ingest thread spent at a barrier waiting for keys. - /// - /// The number to watch: it is the part of extraction the pool did not manage - /// to hide, and so the part that is still serial. Exact, per chunk. - barrier_secs: f64, - /// Whether to defer at all. - /// - /// False when the pool has fewer than two workers, because then there is no - /// second thread for a batch to run on — deferral could only ever move the - /// same extraction onto the same thread, later, having paid for an extent - /// list, a key array and a splice on the way. That is not a judgement call - /// about which is faster; with one worker the batched path is strictly more - /// work. The single-threaded fast path (`--threads` absent) is exactly this - /// case, so it keeps its original inline extraction. - defer: bool, -} - -impl DeferredKeys { - /// Build a driver over `ctx`, deferring when Phase 1 has the workers to make - /// deferral worth its bookkeeping. - /// - /// `phase1_threads` is the count of workers *active during Phase 1*, not the - /// pool's size. The pool is sized to the wider of the two phases and capped - /// per phase, so `--threads 1 --merge-threads 16` builds 16 workers of which - /// exactly one is awake while ingest runs — gating on the pool size there - /// would defer batches to a pool that is not allowed to run them. - pub(crate) fn new(ctx: Arc, phase1_threads: usize) -> Self { - let defer = phase1_threads >= 2; - let (tx, rx) = std::sync::mpsc::channel(); - Self { - ctx, - tx, - rx, - dispatched: 0, - outstanding: 0, - filled: 0, - sealed_seen: 0, - violation: None, - overlapped_records: 0, - barrier_records: 0, - dispatch_secs: 0.0, - barrier_secs: 0.0, - defer, - } - } - - /// Get one record into the buffer with a sort key, now or later. - /// - /// The ingest loop's single entry point, so the deferred and immediate modes - /// cannot diverge at the call site. Immediate mode extracts and pushes in one - /// step exactly as the pre-deferral loop did; deferred mode pushes a - /// placeholder and lets the pool fill it in. - /// - /// # Errors - /// - /// Returns an error if the record cannot be pushed into the arena, or (in - /// immediate mode) if it carries a value in a lane the chosen key drops. - pub(crate) fn push( - &mut self, - buffer: &mut TemplateRecordBuffer, - pool: &SortWorkerPool, - bam: &[u8], - ) -> anyhow::Result<()> { - if !self.defer { - let (full, lane) = extract_one(&self.ctx, bam); - if let Some(v) = lane { - let name = - String::from_utf8_lossy(RawRecordView::new(bam).read_name()).into_owned(); - return Err(crate::external::dropped_lane_error(&name, v)); - } - buffer.push(bam, K::from_full(&full))?; - // Counted as overlapped-with-nothing so the census still sums to the - // record count and a single-worker run reads as 0% overlapped rather - // than as no deferral information at all. - self.barrier_records += 1; - self.filled += 1; - self.dispatched += 1; - return Ok(()); - } - buffer.push_deferred(bam)?; - self.after_push(buffer, pool); - Ok(()) - } - - /// Records whose keys were extracted while the ingest thread kept reading, - /// and records whose keys it had to wait at the barrier for. - /// - /// This is how you tell the difference between deferral *working* and - /// deferral silently not engaging: the second is indistinguishable from the - /// first by output, by test, and by every equivalence check — the keys are - /// right either way, they were just paid for serially. A run where - /// `overlapped` is near zero has the cost of the old code plus the - /// bookkeeping of the new. - pub(crate) fn overlap_census(&self) -> KeyOverlapCensus { - KeyOverlapCensus { - overlapped_records: self.overlapped_records, - barrier_records: self.barrier_records, - dispatch_secs: self.dispatch_secs, - barrier_secs: self.barrier_secs, - } - } - - /// Bytes currently held by in-flight batches, to be counted against the - /// sort's memory limit. - /// - /// A record whose key is still outstanding is charged twice — once for the - /// placeholder key already in its ref, once for the copy the batch is - /// filling — plus its extent. Leaving this out would let the buffer overrun - /// the limit by however much extraction happens to be lagging. - pub(crate) fn in_flight_bytes(&self) -> usize { - let per_record = std::mem::size_of::() + std::mem::size_of::<(u64, u32)>(); - self.dispatched.saturating_sub(self.filled) * per_record - } - - /// Hand out batches for every record that has become shareable, if the - /// arena has sealed a segment since the last call. - /// - /// Cheap enough to call after every push: the common case is one integer - /// comparison. - pub(crate) fn after_push( - &mut self, - buffer: &mut TemplateRecordBuffer, - pool: &SortWorkerPool, - ) { - let sealed = buffer.sealed_segments(); - if !self.defer || sealed == self.sealed_seen { - return; - } - self.sealed_seen = sealed; - let started = std::time::Instant::now(); - // The record that triggered the seal landed in the new live segment, so - // it is not shareable yet; everything before it is. - let shareable = buffer.refs().len().saturating_sub(1); - self.dispatch_upto(buffer, pool, shareable); - self.absorb_ready(buffer); - self.dispatch_secs += started.elapsed().as_secs_f64(); - } - - /// Cut and dispatch batches covering refs `self.dispatched..end`. - /// - /// Batches never span an arena segment: a batch names exactly one sealed - /// segment and holds it by `Arc`, which is what makes reading it concurrent - /// with the ingest thread's appends sound. - fn dispatch_upto( - &mut self, - buffer: &TemplateRecordBuffer, - pool: &SortWorkerPool, - end: usize, - ) { - let segment_size = buffer.segment_size() as u64; - while self.dispatched < end { - let refs = buffer.refs(); - let start = self.dispatched; - let segment_index = refs[start].offset / segment_size; - let limit = end.min(start + KEY_BATCH_RECORDS); - let mut stop = start; - while stop < limit && refs[stop].offset / segment_size == segment_index { - stop += 1; - } - - let segment_index = usize::try_from(segment_index).expect("segment index fits usize"); - let Some(segment) = buffer.sealed_segment(segment_index) else { - // The segment is still live, so nothing from here on is - // shareable yet. Leave it for the next seal or the barrier. - return; - }; - - let extents: Box<[(u64, u32)]> = refs[start..stop] - .iter() - .map(|r| (r.offset + TEMPLATE_HEADER_SIZE as u64, r.len)) - .collect(); - let batch = TemplateKeyBatch:: { - ctx: Arc::clone(&self.ctx), - segment, - segment_base: segment_index as u64 * segment_size, - first_ref: start, - extents, - results: self.tx.clone(), - }; - - self.dispatched = stop; - self.outstanding += 1; - // A full queue means the pool is saturated, not that the batch can - // be dropped: run it here rather than lose its keys. - if let Err(batch) = pool.submit_key_job(Box::new(batch)) { - batch.run(); - } - } - } - - /// Absorb every batch that has already finished, without waiting. - fn absorb_ready(&mut self, buffer: &mut TemplateRecordBuffer) { - while let Ok(result) = self.rx.try_recv() { - self.absorb(buffer, result); - } - } - - /// Write one batch's keys into the buffer and fold in its violation. - fn absorb(&mut self, buffer: &mut TemplateRecordBuffer, result: KeyBatchResult) { - buffer.fill_keys(result.first_ref, &result.keys); - self.filled += result.keys.len(); - self.outstanding -= 1; - if let Some(v) = result.violation - && self.violation.as_ref().is_none_or(|seen| v.ref_index < seen.ref_index) - { - self.violation = Some(v); - } - } - - /// Wait until every record in the buffer has its real key, then report the - /// first dropped-lane violation if there was one. - /// - /// Must be called before the buffer is sorted or spilled. Runs queued - /// batches on the calling thread while it waits, so it cannot deadlock - /// against a pool whose workers have already parked for the phase. - /// - /// # Errors - /// - /// Returns an error if a worker panicked while holding a batch. That batch's - /// sender is dropped by the unwind without ever publishing, so its keys can - /// never arrive — and because this driver holds its own sender, the channel - /// never disconnects and a plain blocking wait would hang forever. The wait - /// is therefore bounded and re-checks the pool's panic flag. - /// - /// # Panics - /// - /// Panics if the buffer still holds a record whose key was never filled — - /// that would otherwise sort those records under a constant placeholder key - /// and produce a wrong order that still looks like a valid BAM. - pub(crate) fn finish( - &mut self, - buffer: &mut TemplateRecordBuffer, - pool: &SortWorkerPool, - ) -> anyhow::Result> { - let started = std::time::Instant::now(); - let all = buffer.refs().len(); - if self.defer { - // Everything left is in the live segment; seal it so it can be - // shared. Skipped in immediate mode, where it would only churn a - // fresh full-size segment allocation per chunk for no reader. - buffer.seal_arena_segment(); - self.sealed_seen = buffer.sealed_segments(); - let overlapped = self.dispatched.min(all); - self.overlapped_records += overlapped as u64; - self.barrier_records += (all - overlapped) as u64; - self.dispatch_upto(buffer, pool, all); - } - - while self.outstanding > 0 { - match self.rx.try_recv() { - Ok(result) => { - self.absorb(buffer, result); - continue; - } - Err(std::sync::mpsc::TryRecvError::Empty) => {} - // Unreachable: `self.tx` is alive for as long as `self` is. - Err(std::sync::mpsc::TryRecvError::Disconnected) => break, - } - if pool.run_one_key_job() { - continue; - } - // Nothing ready and nothing queued: every outstanding batch is - // running on a worker right now, so waiting is the cheap move. The - // timeout exists only so a panicked worker surfaces as an error - // instead of an indefinite hang. - match self.rx.recv_timeout(std::time::Duration::from_millis(50)) { - Ok(result) => self.absorb(buffer, result), - Err(std::sync::mpsc::RecvTimeoutError::Timeout) => { - anyhow::ensure!( - !pool.worker_panicked(), - "a sort worker panicked while extracting sort keys; {} key \ - batches can never complete", - self.outstanding, - ); - } - Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => break, - } - } - - self.barrier_secs += started.elapsed().as_secs_f64(); - assert_eq!( - self.filled, - all, - "{} of {all} records were never given a sort key; sorting now would collate them \ - under the placeholder", - all - self.filled, - ); - Ok(self.violation.take()) - } - - /// Forget the previous chunk's accounting, after the buffer was cleared. - pub(crate) fn reset(&mut self) { - self.dispatched = 0; - self.outstanding = 0; - self.filled = 0; - self.sealed_seen = 0; - self.violation = None; - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::external::cb_hasher; - use crate::inline::TemplateKey40; - use noodles::sam::Header; - - /// Minimal mapped BAM record: paired, forward, mate at the same locus, with - /// `aux` appended verbatim. - #[allow(clippy::cast_possible_truncation)] - fn mapped_bam(tid: i32, pos: i32, name: &[u8], aux: &[u8]) -> Vec { - let mut bam = vec![0u8; 32]; - bam[0..4].copy_from_slice(&tid.to_le_bytes()); - bam[4..8].copy_from_slice(&pos.to_le_bytes()); - bam[8] = (name.len() + 1) as u8; - bam[14..16].copy_from_slice(&3u16.to_le_bytes()); - bam[20..24].copy_from_slice(&tid.to_le_bytes()); - bam[24..28].copy_from_slice(&pos.to_le_bytes()); - bam.extend_from_slice(name); - bam.push(0); - bam.extend_from_slice(aux); - bam - } - - /// `CB:Z:` aux bytes. - fn cb_aux(value: &[u8]) -> Vec { - let mut aux = b"CBZ".to_vec(); - aux.extend_from_slice(value); - aux.push(0); - aux - } - - /// Pack `records` into one segment and describe them as a batch would see - /// them, returning `(segment, extents)`. - /// A packed fixture segment and the extents describing the records in it. - type PackedSegment = (Arc>, Box<[(u64, u32)]>); - - fn pack(records: &[Vec]) -> PackedSegment { - let mut segment = Vec::new(); - let mut extents = Vec::new(); - for r in records { - let len = u32::try_from(r.len()).expect("fixture records are small"); - extents.push((segment.len() as u64, len)); - segment.extend_from_slice(r); - } - (Arc::new(segment), extents.into_boxed_slice()) - } - - /// A context whose baseline is `records[0]` and which retains every lane, so - /// no record can trip the dropped-lane check. - fn context(records: &[Vec], cell_tag: Option) -> Arc { - let lib_lookup = LibraryLookup::from_header(&Header::default()); - let hasher = cb_hasher(); - let first_key = extract_template_key_inline(&records[0], &lib_lookup, cell_tag, &hasher); - Arc::new(KeyContext { - lib_lookup, - cell_tag, - cb_hasher: hasher, - first_key, - variant: TemplateKeyVariant { cb: true, tertiary: true }, - }) - } - - fn run_batch( - ctx: Arc, - segment: Arc>, - extents: Box<[(u64, u32)]>, - first_ref: usize, - ) -> Option> { - let (tx, rx) = std::sync::mpsc::channel(); - let job: Box = Box::new(TemplateKeyBatch:: { - ctx, - segment, - segment_base: 0, - first_ref, - extents, - results: tx, - }); - job.run(); - rx.try_recv().ok() - } - - #[test] - fn test_a_batch_extracts_the_keys_serial_extraction_would_have() { - // The whole change is only safe if deferring extraction cannot alter the - // key, so this compares against the serial call the ingest loop used to - // make, record for record. - let records: Vec> = - (0..8).map(|i| mapped_bam(1, 100 + i, format!("read{i}").as_bytes(), &[])).collect(); - let ctx = context(&records, None); - let (segment, extents) = pack(&records); - - let result: KeyBatchResult = - run_batch(Arc::clone(&ctx), segment, extents, 0).expect("batch publishes a result"); - - let expected: Vec = records - .iter() - .map(|r| { - let full = - extract_template_key_inline(r, &ctx.lib_lookup, ctx.cell_tag, &ctx.cb_hasher); - TemplateLaneKey::from_full(&full) - }) - .collect(); - - assert_eq!(result.first_ref, 0); - assert_eq!(result.keys.len(), 8); - assert_eq!(&result.keys[..], &expected[..]); - assert!(result.violation.is_none()); - } - - #[test] - fn test_a_batch_reports_the_lowest_indexed_violation_not_the_last() { - // Batches run in whatever order the pool picks them up, so the reported - // record has to be the first offending one or the error message is not - // reproducible across runs. Two records violate here; the earlier must win. - let mut records = vec![mapped_bam(1, 100, b"r0", &cb_aux(b"AAAA"))]; - records.push(mapped_bam(1, 101, b"r1", &cb_aux(b"AAAA"))); - records.push(mapped_bam(1, 102, b"r2", &cb_aux(b"CCCC"))); - records.push(mapped_bam(1, 103, b"r3", &cb_aux(b"AAAA"))); - records.push(mapped_bam(1, 104, b"r4", &cb_aux(b"GGGG"))); - - // Force the CB lane to be dropped so the differing barcodes violate. - let lib_lookup = LibraryLookup::from_header(&Header::default()); - let hasher = cb_hasher(); - let first_key = - extract_template_key_inline(&records[0], &lib_lookup, Some(SamTag::CB), &hasher); - let ctx = Arc::new(KeyContext { - lib_lookup, - cell_tag: Some(SamTag::CB), - cb_hasher: hasher, - first_key, - variant: TemplateKeyVariant { cb: false, tertiary: false }, - }); - - let (segment, extents) = pack(&records); - // `first_ref` is deliberately non-zero: the reported index must be the - // record's position in the chunk, not in the batch. - let result: KeyBatchResult = - run_batch(ctx, segment, extents, 1000).expect("batch publishes a result"); - - let v = result.violation.expect("a dropped CB lane must be reported"); - assert_eq!(v.ref_index, 1002, "the first violating record is r2, at chunk index 1002"); - assert_eq!(v.name, "r2"); - } - - #[test] - fn test_a_multi_segment_ingest_keys_every_record_exactly_once() { - // The end-to-end contract, and the only test that reaches the - // dispatch-on-seal path: drive a real buffer and a real pool across - // several arena segments, then check every key against what the serial - // extraction this replaced would have produced. A batch that read the - // wrong segment, double-counted a range, or silently skipped the live - // segment's tail all show up here as a wrong or placeholder key. - use crate::codec::SpillCodec; - use crate::inline::TemplateRecordBuffer; - use crate::worker_pool::SortWorkerPool; - - // Small enough that a few hundred records fill several segments. - const SEGMENT: usize = 4096; - - let records: Vec> = (0..600i32) - .map(|i| { - let name = format!("read{i:05}"); - // Vary length so records straddle segment boundaries unevenly. - let filler = vec![b'N'; usize::try_from(i % 37).expect("non-negative")]; - let mut aux = b"MCZ".to_vec(); - aux.extend_from_slice(&filler); - aux.push(0); - mapped_bam(1, 1000 + (i % 97), name.as_bytes(), &aux) - }) - .collect(); - let ctx = context(&records, None); - - let pool = SortWorkerPool::new(4, 1, 6, SpillCodec::Bgzf, false); - pool.set_phase(crate::worker_pool::phase::PHASE1); - - let mut buffer = - TemplateRecordBuffer::::with_segment_size(600, SEGMENT, SEGMENT); - let mut deferred = DeferredKeys::::new(Arc::clone(&ctx), 4); - for record in &records { - buffer.push_deferred(record).expect("push"); - deferred.after_push(&mut buffer, &pool); - } - assert!( - buffer.sealed_segments() > 2, - "the fixture must span several segments to exercise seal-driven dispatch, got {}", - buffer.sealed_segments(), - ); - - // Most records must have been dispatched *before* the barrier. Without - // this the test passes with mid-stream dispatch disabled entirely — the - // barrier would extract everything and the keys would still be right, - // which is exactly how this optimization could silently do nothing. - assert_eq!( - deferred.overlap_census().total_records(), - 0, - "the census is only taken at the barrier" - ); - - let violation = deferred.finish(&mut buffer, &pool).expect("no worker panics"); - assert!(violation.is_none(), "no record drops a retained lane here"); - - let census = deferred.overlap_census(); - let overlapped = census.overlapped_records; - assert_eq!(census.total_records(), records.len() as u64); - assert!( - overlapped > (records.len() as u64) / 2, - "most keys must be extracted while ingest is still reading; \ - only {overlapped} of {} were", - records.len(), - ); - - let expected: Vec = records - .iter() - .map(|r| { - let full = - extract_template_key_inline(r, &ctx.lib_lookup, ctx.cell_tag, &ctx.cb_hasher); - TemplateLaneKey::from_full(&full) - }) - .collect(); - let got: Vec = buffer.refs().iter().map(|r| r.key).collect(); - assert_eq!(got, expected, "every record must carry the key serial extraction gives it"); - - // And the bytes must still be retrievable at the offsets the refs claim. - for (i, r) in buffer.refs().iter().enumerate() { - assert_eq!(buffer.get_record(r), &records[i][..], "record {i} bytes"); - } - } - - #[test] - fn test_a_single_worker_pool_keys_inline_and_reports_no_overlap() { - // With one worker there is no second thread for a batch to run on, so - // deferral is skipped and extraction happens in the push, as it did - // before this change. The keys must be identical either way, and the - // census must say plainly that nothing overlapped rather than going - // silent — a run reporting no deferral information at all would be - // indistinguishable from deferral that broke. - use crate::codec::SpillCodec; - use crate::inline::TemplateRecordBuffer; - use crate::worker_pool::SortWorkerPool; - - const SEGMENT: usize = 4096; - - let records: Vec> = (0..200i32) - .map(|i| mapped_bam(1, 500 + i, format!("q{i:04}").as_bytes(), &[])) - .collect(); - let ctx = context(&records, None); - - let pool = SortWorkerPool::new(1, 1, 6, SpillCodec::Bgzf, false); - pool.set_phase(crate::worker_pool::phase::PHASE1); - - let mut buffer = - TemplateRecordBuffer::::with_segment_size(200, SEGMENT, SEGMENT); - let mut deferred = DeferredKeys::::new(Arc::clone(&ctx), 1); - for record in &records { - deferred.push(&mut buffer, &pool, record).expect("push"); - } - let violation = deferred.finish(&mut buffer, &pool).expect("no worker panics"); - assert!(violation.is_none()); - - let expected: Vec = records - .iter() - .map(|r| { - let full = - extract_template_key_inline(r, &ctx.lib_lookup, ctx.cell_tag, &ctx.cb_hasher); - TemplateLaneKey::from_full(&full) - }) - .collect(); - let got: Vec = buffer.refs().iter().map(|r| r.key).collect(); - assert_eq!(got, expected, "inline extraction must give the same keys"); - - let census = deferred.overlap_census(); - assert_eq!(census.total_records(), records.len() as u64, "the census still sums"); - assert_eq!(census.overlap_percent(), Some(0.0), "nothing can overlap with one worker"); - } - - #[test] - fn test_a_single_worker_pool_still_rejects_a_dropped_lane() { - // The immediate path has its own error return, so it needs its own - // coverage: a violation must surface as an error from the push rather - // than silently keying the record under a lane the sort dropped. - use crate::codec::SpillCodec; - use crate::inline::{TemplateKey24, TemplateRecordBuffer}; - use crate::worker_pool::SortWorkerPool; - - let records = [ - mapped_bam(1, 100, b"okay", &cb_aux(b"AAAA")), - mapped_bam(1, 101, b"bad", &cb_aux(b"TTTT")), - ]; - let lib_lookup = LibraryLookup::from_header(&Header::default()); - let hasher = cb_hasher(); - let first_key = - extract_template_key_inline(&records[0], &lib_lookup, Some(SamTag::CB), &hasher); - let ctx = Arc::new(KeyContext { - lib_lookup, - cell_tag: Some(SamTag::CB), - cb_hasher: hasher, - first_key, - variant: TemplateKeyVariant { cb: false, tertiary: false }, - }); - - let pool = SortWorkerPool::new(1, 1, 6, SpillCodec::Bgzf, false); - pool.set_phase(crate::worker_pool::phase::PHASE1); - let mut buffer = TemplateRecordBuffer::::with_capacity(2, 4096); - let mut deferred = DeferredKeys::::new(ctx, 1); - - deferred.push(&mut buffer, &pool, &records[0]).expect("the baseline record is fine"); - let err = deferred - .push(&mut buffer, &pool, &records[1]) - .expect_err("the offending record must be rejected"); - assert!(err.to_string().contains("bad"), "the error names the record: {err}"); - } - - #[test] - fn test_a_multi_segment_ingest_reports_a_violation_in_a_sealed_segment() { - // A dropped-lane violation found by a *worker*, mid-stream, has to reach - // the ingest thread — the serial path could return the error inline, and - // this one cannot. - use crate::codec::SpillCodec; - use crate::inline::{TemplateKey24, TemplateRecordBuffer}; - use crate::worker_pool::SortWorkerPool; - - const SEGMENT: usize = 4096; - - let mut records: Vec> = (0..400) - .map(|i| mapped_bam(1, 1000 + i, format!("r{i:04}").as_bytes(), &cb_aux(b"AAAA"))) - .collect(); - // One offender, early enough to land in a sealed segment. - records[7] = mapped_bam(1, 1007, b"offender", &cb_aux(b"TTTT")); - - let lib_lookup = LibraryLookup::from_header(&Header::default()); - let hasher = cb_hasher(); - let first_key = - extract_template_key_inline(&records[0], &lib_lookup, Some(SamTag::CB), &hasher); - let ctx = Arc::new(KeyContext { - lib_lookup, - cell_tag: Some(SamTag::CB), - cb_hasher: hasher, - first_key, - variant: TemplateKeyVariant { cb: false, tertiary: false }, - }); - - let pool = SortWorkerPool::new(4, 1, 6, SpillCodec::Bgzf, false); - pool.set_phase(crate::worker_pool::phase::PHASE1); - - let mut buffer = - TemplateRecordBuffer::::with_segment_size(400, SEGMENT, SEGMENT); - let mut deferred = DeferredKeys::::new(ctx, 4); - for record in &records { - buffer.push_deferred(record).expect("push"); - deferred.after_push(&mut buffer, &pool); - } - let violation = deferred - .finish(&mut buffer, &pool) - .expect("no worker panics") - .expect("the offender must be reported"); - assert_eq!(violation.ref_index, 7); - assert_eq!(violation.name, "offender"); - } - - #[test] - fn test_a_batch_whose_receiver_is_gone_finishes_quietly() { - // The ingest thread drops the receiver when it fails, and outstanding - // batches must not panic on the way out or a real error is replaced by a - // worker panic. - let records = vec![mapped_bam(1, 100, b"r0", &[])]; - let ctx = context(&records, None); - let (segment, extents) = pack(&records); - - let (tx, rx) = std::sync::mpsc::channel::>(); - drop(rx); - let job: Box = Box::new(TemplateKeyBatch:: { - ctx, - segment, - segment_base: 0, - first_ref: 0, - extents, - results: tx, - }); - job.run(); - } - - #[test] - fn test_a_batch_reads_at_its_segments_base_not_the_arenas() { - // Extents are global arena offsets while the segment handle is local, so - // the batch has to subtract its base. Getting this wrong reads the wrong - // record and silently produces a plausible-but-wrong key rather than - // failing, which is why it is pinned separately. - const BASE: u64 = 4096; - - let records: Vec> = - (0..4).map(|i| mapped_bam(1, 200 + i, format!("s{i}").as_bytes(), &[])).collect(); - let ctx = context(&records, None); - let (segment, local_extents) = pack(&records); - let shifted: Box<[(u64, u32)]> = - local_extents.iter().map(|&(o, l)| (o + BASE, l)).collect(); - - let (tx, rx) = std::sync::mpsc::channel(); - let job: Box = Box::new(TemplateKeyBatch:: { - ctx: Arc::clone(&ctx), - segment, - segment_base: BASE, - first_ref: 0, - extents: shifted, - results: tx, - }); - job.run(); - let result = rx.try_recv().expect("batch publishes a result"); - - let expected: Vec = records - .iter() - .map(|r| { - let full = - extract_template_key_inline(r, &ctx.lib_lookup, ctx.cell_tag, &ctx.cb_hasher); - TemplateLaneKey::from_full(&full) - }) - .collect(); - assert_eq!(&result.keys[..], &expected[..]); - } -} diff --git a/crates/fgumi-sort/src/phase1_stats.rs b/crates/fgumi-sort/src/phase1_stats.rs deleted file mode 100644 index 1b0193d0a..000000000 --- a/crates/fgumi-sort/src/phase1_stats.rs +++ /dev/null @@ -1,721 +0,0 @@ -//! What Phase 1's serial ingest thread waits for. -//! -//! Phase 2 has a floor line ([`crate::merge_headroom`]) because its three -//! limits -- serial consumer, worker capacity, coordination -- imply unrelated -//! fixes and are routinely confused. Phase 1 has the same three limits and had -//! none of the accounting: its report was four wall-clock spans (read, sort, -//! spill write, merge) with no way to tell a thread that is busy from one that -//! is waiting. -//! -//! That gap matters more than it did for the merge. External sampling of a -//! 16-thread whole-genome sort put Phase 1 at **60% of total wall clock with its -//! main thread 91% busy** while all 16 cores averaged 5.3 -- so the phase is -//! bound by one thread's serial CPU, and every worker-side change is pushing on -//! a wall that is not there. In-process numbers should say that without needing -//! a `/proc` sampler attached from outside. -//! -//! # The two waits -//! -//! The ingest thread blocks in exactly two places, and neither was measured: -//! -//! 1. **Waiting for a decompressed block.** [`crate::read_ahead::PooledInputStream`] -//! parks when the next serial it needs has not arrived. Because blocks are -//! consumed in serial order through a reorder buffer, this can fire while -//! other blocks are ready -- head-of-line blocking -- which is a different -//! problem from an empty queue and is counted separately here. -//! 2. **Waiting for the previous spill to finish.** `drain_pending_spill` waits -//! on the prior chunk's write handle between the read span ending and the -//! in-memory sort starting, so that time lands in *no* phase bucket at all -//! and shows up only as an unexplained residual against total wall clock. -//! -//! Timing is exact rather than sampled: both waits are milliseconds against a -//! ~30 ns clock read, so the clock is 5 orders of magnitude below the quantity -//! and costs nothing to read (the same argument [`crate::merge_trace`] makes for -//! the merge's block pull, where an exact timer cost 0.15%). - -use std::sync::atomic::{AtomicU64, Ordering}; - -/// Why the ingest thread parked waiting for a block. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum ParkCause { - /// Nothing decompressed was available at all: the workers are behind. - Starved, - /// Blocks were buffered, but not the serial the consumer needs next. The - /// pipeline has work in flight and the consumer still cannot use it. - HeadOfLine, -} - -/// Counters for the ingest thread's waits, shared with the pool's input stream. -#[derive(Debug, Default)] -pub(crate) struct Phase1IngestStats { - /// Times the ingest thread parked waiting for its next block. - parks: AtomicU64, - /// Total nanoseconds parked waiting for a block. - park_nanos: AtomicU64, - /// Parks where the reorder buffer held nothing at all. - parks_starved: AtomicU64, - /// Parks where the reorder buffer held blocks, but not the next serial. - parks_head_of_line: AtomicU64, - /// Times the ingest thread waited on the previous spill to complete. - spill_waits: AtomicU64, - /// Total nanoseconds waiting on a previous spill. - spill_wait_nanos: AtomicU64, -} - -impl Phase1IngestStats { - /// Record one park and what the reorder buffer looked like when it happened. - pub(crate) fn record_park(&self, elapsed_nanos: u64, cause: ParkCause) { - self.parks.fetch_add(1, Ordering::Relaxed); - self.park_nanos.fetch_add(elapsed_nanos, Ordering::Relaxed); - match cause { - ParkCause::Starved => self.parks_starved.fetch_add(1, Ordering::Relaxed), - ParkCause::HeadOfLine => self.parks_head_of_line.fetch_add(1, Ordering::Relaxed), - }; - } - - /// Record one wait on the previous chunk's spill write. - pub(crate) fn record_spill_wait(&self, elapsed_nanos: u64) { - self.spill_waits.fetch_add(1, Ordering::Relaxed); - self.spill_wait_nanos.fetch_add(elapsed_nanos, Ordering::Relaxed); - } - - /// A consistent view of the counters, for reporting. - pub(crate) fn snapshot(&self) -> Phase1IngestReport { - Phase1IngestReport { - parks: self.parks.load(Ordering::Relaxed), - park_secs: secs(self.park_nanos.load(Ordering::Relaxed)), - parks_starved: self.parks_starved.load(Ordering::Relaxed), - parks_head_of_line: self.parks_head_of_line.load(Ordering::Relaxed), - spill_waits: self.spill_waits.load(Ordering::Relaxed), - spill_wait_secs: secs(self.spill_wait_nanos.load(Ordering::Relaxed)), - } - } -} - -#[allow(clippy::cast_precision_loss, reason = "nanosecond totals stay far below 2^52")] -fn secs(nanos: u64) -> f64 { - nanos as f64 / 1_000_000_000.0 -} - -/// What the ingest thread waited for, as reported. -#[derive(Debug, Clone, Copy, Default)] -pub(crate) struct Phase1IngestReport { - pub(crate) parks: u64, - pub(crate) park_secs: f64, - pub(crate) parks_starved: u64, - pub(crate) parks_head_of_line: u64, - pub(crate) spill_waits: u64, - pub(crate) spill_wait_secs: f64, -} - -impl Phase1IngestReport { - /// Mean park in microseconds, or `None` when it never parked. - /// - /// The mean is the discriminant between "parked rarely and long" (a supply - /// problem) and "parked constantly and briefly" (a handoff problem), which - /// the merge campaign found to be the difference between a fixable stall and - /// an unfixable one. - pub(crate) fn mean_park_micros(&self) -> Option { - #[allow(clippy::cast_precision_loss, reason = "park counts stay far below 2^52")] - (self.parks > 0).then(|| self.park_secs * 1_000_000.0 / self.parks as f64) - } -} - -/// One record in this many is timed for the sub-phase partition. -/// -/// Prime, so the sampled set cannot align with any periodic structure in the -/// input (read groups, tile boundaries, alternating mate records) and bias the -/// partition toward whichever records happen to be cheap. -pub(crate) const INGEST_SAMPLE_INTERVAL: u64 = 1021; - -/// Where the ingest thread's serial CPU goes, per record. -/// -/// **Per-record segments only.** Deferred key extraction is deliberately absent: -/// its costs are per-*batch* and per-*chunk*, not per-record, so sampling one -/// record in [`INGEST_SAMPLE_INTERVAL`] and scaling would be measuring a rare -/// event with a method built for a uniform one. Those two are timed exactly -/// instead and reported beside this partition, the same way `park_secs` is. -/// -/// The floor line says this thread *is* the limit -- on a 16-thread whole-genome -/// sort it is 137.2s of a 145.7s read span, against a worker-capacity floor of -/// 22.4s -- so the only question left is what the 137.2s is made of. Nothing else -/// in the phase can answer it: worker counters describe the pool, and wall-clock -/// spans describe the phase, and neither looks inside the loop. -/// -/// Sampled rather than timed on every record, for the same reason -/// [`crate::merge_headroom::ConsumerSample`] is: the loop runs at ~175 ns/record -/// and an `Instant::now()` pair costs 15-35 ns on aarch64, so timing five -/// segments on every record would cost more than several of the segments it -/// measures. One record in [`INGEST_SAMPLE_INTERVAL`] is timed and scaled, and -/// the scale is reported next to the result. -/// -/// Each field is **exactly one** timed region in the loop, which is what makes -/// [`Self::corrected`] valid: it subtracts one clock pair per field per sample, -/// so a field spanning two bracketed regions would be under-corrected by a -/// whole pair. That is why progress counting and spill probing are separate -/// fields rather than one "bookkeeping" bucket -- they sit at opposite ends of -/// the loop body and cannot share a bracket. -#[derive(Debug, Clone, Copy, Default)] -pub(crate) struct IngestSample { - /// Pulling the next record's bytes from the pool's decompressed stream. - /// **Includes park time**, so it is not pure CPU -- `park_secs` measures that - /// part exactly and separately. - pub(crate) fetch: f64, - /// Copying the record into the arena and appending its ref. - pub(crate) push: f64, - /// Counting the record toward the progress log. - pub(crate) tick: f64, - /// The spill probe's sample check and the memory-limit test that follows it. - pub(crate) probe: f64, -} - -impl IngestSample { - /// Every segment multiplied by the sampling scale. - #[must_use] - pub(crate) fn scaled(self, scale: f64) -> Self { - Self { - fetch: self.fetch * scale, - push: self.push * scale, - tick: self.tick * scale, - probe: self.probe * scale, - } - } - - /// Every segment with its own measurement overhead removed. - /// - /// One `Instant::now()`/`elapsed()` pair per segment per sampled record, and - /// that pair's cost lands inside the interval it times. Clamped at zero: a - /// segment cheaper than the clock measuring it cannot be resolved this way, - /// and zero says so where a negative would read as a bug. - #[must_use] - pub(crate) fn corrected(self, samples: u64, overhead_nanos: u64) -> Self { - if samples == 0 || overhead_nanos == 0 { - return self; - } - #[expect(clippy::cast_precision_loss, reason = "sample counts stay below 2^52")] - let per_segment = (samples * overhead_nanos) as f64 / 1e9; - let fix = |v: f64| (v - per_segment).max(0.0); - Self { - fetch: fix(self.fetch), - push: fix(self.push), - tick: fix(self.tick), - probe: fix(self.probe), - } - } - - /// The six segments summed. - #[must_use] - pub(crate) fn total(self) -> f64 { - self.fetch + self.push + self.tick + self.probe - } -} - -/// A scaled ingest sample checked against the read span it should partition. -#[derive(Debug, Clone, Copy)] -pub(crate) struct IngestPartition { - /// Scaled, corrected per-segment seconds. - pub(crate) segments: IngestSample, - /// Measured read span, exact. - pub(crate) read_secs: f64, - /// Park time inside `segments.fetch`, measured exactly and separately. - pub(crate) park_secs: f64, -} - -impl IngestPartition { - /// Read-span time the segments do not account for. - /// - /// **Signed on purpose.** A negative residual means the sample - /// over-attributes -- clock overhead left inside the timed regions, or a - /// sampling bias -- and the merge's first partition did exactly that, - /// summing to 321.5s of a 189.3s loop. Only the sign made it visible; a - /// clamped residual would have reported a tidy zero and the partition would - /// have been believed. - #[must_use] - pub(crate) fn residual_secs(self) -> f64 { - self.read_secs - self.segments.total() - } - - /// Residual as a share of the read span, for judging whether the partition - /// is trustworthy at all. - #[must_use] - pub(crate) fn residual_share(self) -> f64 { - if self.read_secs > 0.0 { self.residual_secs() / self.read_secs } else { 0.0 } - } -} - -// ============================================================================ -// The input reader -// ============================================================================ - -/// What the input reader's exclusive thread spends its time on. -/// -/// `ReadInputBlocks` is a *second* serial resource in Phase 1, and its cost is -/// not explained by the disk. Worker 0 owns the step exclusively -/// (`exclusive_step_for` returns `None` for every other worker), so one thread -/// frames the whole file: on the production cell that is 124.5s across 5,114,408 -/// blocks -- **24.3 us per 8.4 KB block** -- while the same 43.1 GB at that -/// volume's measured single-stream ceiling (605 MB/s, direct I/O, caches -/// dropped) is only ~71s. The step is not disk-bound, and nothing accounts for -/// the remaining ~55s. -/// -/// It cannot be split from outside the reader, because the disk wait is *inside* -/// it: `read_raw_blocks` frames blocks out of a 2 MiB `BufReader`, so whichever -/// block happens to exhaust the buffer pays for the refill and every other block -/// pays nothing. Timing the step measures framing and I/O fused together, and no -/// per-block average can separate them. [`TimedReader`] sits one layer *below* -/// the buffer and times the underlying `read()` itself, which is what makes -/// [`ReaderReport::framing_secs`] -- time inside `read_raw_blocks` that was not -/// spent in a `read()` -- mean anything. -/// -/// Timed exactly rather than sampled: ~320k batches and ~21k refills against a -/// ~30 ns clock is well under 30 ms of measurement on a 124.5s step, five orders -/// of magnitude below the quantity. [`IngestSample`] has to sample only because -/// its loop runs 780 million times; this one does not. -/// One thread's scheduler accounting: `(on_cpu_nanos, runqueue_wait_nanos, -/// timeslices)`, read from `/proc/thread-self/schedstat`. -/// -/// This is the measurement that separates the two stories a slower `read()` -/// could be telling. If the call got slower because the thread lost the CPU — -/// to the fifteen workers now extracting keys, or to anything else — the -/// runqueue-wait delta across the call grows and the timeslice count rises. If -/// it got slower because the *device* or the memory system delivered bytes more -/// slowly, the thread was blocked on I/O the whole time and both stay flat. -/// Wall-clock timing alone cannot tell those apart. -/// -/// `/proc/thread-self` is Linux-only and needs no thread-id lookup; elsewhere -/// this returns `None` and the report simply omits the section, which is why -/// nothing downstream treats a missing sample as a zero. -#[cfg(target_os = "linux")] -pub(crate) fn thread_schedstat() -> Option<(u64, u64, u64)> { - let text = std::fs::read_to_string("/proc/thread-self/schedstat").ok()?; - parse_schedstat(&text) -} - -/// Non-Linux stub: no scheduler accounting is available. -#[cfg(not(target_os = "linux"))] -pub(crate) fn thread_schedstat() -> Option<(u64, u64, u64)> { - None -} - -/// Parse a `schedstat` line: on-CPU nanos, runqueue-wait nanos, timeslices. -/// -/// Split out from the read so it can be tested off Linux, where the proc file -/// does not exist and the whole path would otherwise be unexercised — which is -/// also why it is `dead_code`-exempt there: only the tests call it. -#[cfg_attr(not(target_os = "linux"), allow(dead_code))] -pub(crate) fn parse_schedstat(text: &str) -> Option<(u64, u64, u64)> { - let mut fields = text.split_whitespace(); - let on_cpu = fields.next()?.parse().ok()?; - let runqueue = fields.next()?.parse().ok()?; - let timeslices = fields.next()?.parse().ok()?; - Some((on_cpu, runqueue, timeslices)) -} - -#[derive(Debug, Default)] -pub(crate) struct ReaderStats { - /// `read()` calls made against the file, below the 2 MiB buffer. - refills: AtomicU64, - /// Nanoseconds inside those calls. - refill_nanos: AtomicU64, - /// Bytes they returned. - refill_bytes: AtomicU64, - /// Successful `ReadInputBlocks` batches. - batches: AtomicU64, - /// Blocks those batches framed. - blocks: AtomicU64, - /// Nanoseconds inside `read_raw_blocks`, refill time included. - read_raw_nanos: AtomicU64, - /// Nanoseconds pushing framed blocks onto the decompress queue. - dispatch_nanos: AtomicU64, - /// Per-`read()` latency distribution. - /// - /// A uniform shift and a heavy tail are different diagnoses: the first says - /// every call got slower (the device or the memory system is delivering less - /// bandwidth), the second says most calls were fine and some stalled - /// (reclaim, contention spikes). The mean alone cannot tell them apart. - refill_latency: crate::merge_trace::DurationHistogram, - /// Runqueue wait accumulated by the reading thread *inside* `read()` calls. - refill_runqueue_nanos: AtomicU64, - /// On-CPU time accumulated by the reading thread inside `read()` calls. - refill_oncpu_nanos: AtomicU64, - /// Timeslices the reading thread was granted inside `read()` calls. - refill_timeslices: AtomicU64, - /// Refills for which a scheduler sample was actually available. - /// - /// Separate from `refills` on purpose: on a platform without - /// `/proc/thread-self/schedstat` this stays zero and the report omits the - /// section, rather than printing a zero runqueue wait that would read as - /// evidence the reader never lost the CPU. - refill_sched_samples: AtomicU64, -} - -impl ReaderStats { - /// Record one `read()` against the underlying file. - pub(crate) fn record_refill(&self, elapsed_nanos: u64, bytes: usize) { - self.refills.fetch_add(1, Ordering::Relaxed); - self.refill_nanos.fetch_add(elapsed_nanos, Ordering::Relaxed); - self.refill_bytes.fetch_add(bytes as u64, Ordering::Relaxed); - self.refill_latency.record(elapsed_nanos); - } - - /// Record what the scheduler did to the reading thread during one `read()`. - /// - /// `on_cpu` and `runqueue` are deltas in nanoseconds and `timeslices` a - /// delta in count, taken across the call. - pub(crate) fn record_refill_sched(&self, on_cpu: u64, runqueue: u64, timeslices: u64) { - self.refill_oncpu_nanos.fetch_add(on_cpu, Ordering::Relaxed); - self.refill_runqueue_nanos.fetch_add(runqueue, Ordering::Relaxed); - self.refill_timeslices.fetch_add(timeslices, Ordering::Relaxed); - self.refill_sched_samples.fetch_add(1, Ordering::Relaxed); - } - - /// Record one framed batch: how long `read_raw_blocks` took and what it returned. - pub(crate) fn record_batch(&self, elapsed_nanos: u64, blocks: usize) { - self.batches.fetch_add(1, Ordering::Relaxed); - self.blocks.fetch_add(blocks as u64, Ordering::Relaxed); - self.read_raw_nanos.fetch_add(elapsed_nanos, Ordering::Relaxed); - } - - /// Record one dispatch of a framed batch onto the queue. - pub(crate) fn record_dispatch(&self, elapsed_nanos: u64) { - self.dispatch_nanos.fetch_add(elapsed_nanos, Ordering::Relaxed); - } - - /// A consistent view of the counters, against the step total that should - /// contain them. - pub(crate) fn snapshot(&self, step_secs: f64) -> ReaderReport { - ReaderReport { - refills: self.refills.load(Ordering::Relaxed), - refill_secs: secs(self.refill_nanos.load(Ordering::Relaxed)), - refill_bytes: self.refill_bytes.load(Ordering::Relaxed), - batches: self.batches.load(Ordering::Relaxed), - blocks: self.blocks.load(Ordering::Relaxed), - read_raw_secs: secs(self.read_raw_nanos.load(Ordering::Relaxed)), - dispatch_secs: secs(self.dispatch_nanos.load(Ordering::Relaxed)), - step_secs, - refill_latency: self.refill_latency.snapshot(), - refill_runqueue_secs: secs(self.refill_runqueue_nanos.load(Ordering::Relaxed)), - refill_oncpu_secs: secs(self.refill_oncpu_nanos.load(Ordering::Relaxed)), - refill_timeslices: self.refill_timeslices.load(Ordering::Relaxed), - refill_sched_samples: self.refill_sched_samples.load(Ordering::Relaxed), - } - } -} - -/// Where the input reader's serial time went, as reported. -#[derive(Debug, Clone, Copy, Default)] -pub(crate) struct ReaderReport { - pub(crate) refills: u64, - pub(crate) refill_secs: f64, - pub(crate) refill_bytes: u64, - pub(crate) batches: u64, - pub(crate) blocks: u64, - /// Time inside `read_raw_blocks`, **including** the refills it triggered. - pub(crate) read_raw_secs: f64, - pub(crate) dispatch_secs: f64, - /// Exact `ReadInputBlocks` busy total the parts should add up to. - pub(crate) step_secs: f64, - /// Per-`read()` latency distribution. - pub(crate) refill_latency: crate::merge_trace::HistogramReport, - /// Seconds the reading thread spent waiting for a CPU inside `read()`. - pub(crate) refill_runqueue_secs: f64, - /// Seconds it spent on-CPU inside `read()`. - pub(crate) refill_oncpu_secs: f64, - /// Timeslices granted to it inside `read()`. - pub(crate) refill_timeslices: u64, - /// Refills that produced a scheduler sample; zero means unavailable here. - pub(crate) refill_sched_samples: u64, -} - -impl ReaderReport { - /// Share of refill wall time the reading thread spent waiting for a CPU. - /// - /// The number that decides whether a slower `read()` is a scheduling - /// problem. Near zero means the thread was blocked on I/O the whole time and - /// the bytes simply arrived more slowly; a material share means it was - /// runnable and waiting, and the fix is scheduling rather than bandwidth. - pub(crate) fn refill_runqueue_share(&self) -> f64 { - if self.refill_secs <= 0.0 { - return 0.0; - } - self.refill_runqueue_secs / self.refill_secs - } - - /// Time inside `read_raw_blocks` that was **not** spent in a `read()`: - /// header parse, validation, per-block allocation, and the body copy. - /// - /// This is the number the whole struct exists to produce. It is signed, and - /// a small negative is a known, benign case rather than a bug: the header - /// parse refills the buffer once before Phase 1 starts, so a few - /// milliseconds of refill can sit outside every timed batch. A *large* - /// negative means the reader was rebuilt or the counters were shared across - /// runs, and should not be explained away. - pub(crate) fn framing_secs(&self) -> f64 { - self.read_raw_secs - self.refill_secs - } - - /// Step time the parts do not account for -- the `try_lock`, the serial - /// reservation, and dropping the guard. - /// - /// **Signed on purpose**, for the reason [`IngestPartition::residual_secs`] - /// gives: the merge's first partition summed to 321.5s of a 189.3s loop, and - /// only the sign made that visible. - pub(crate) fn residual_secs(&self) -> f64 { - self.step_secs - self.read_raw_secs - self.dispatch_secs - } - - /// Residual as a share of the step, for judging the partition at a glance. - pub(crate) fn residual_share(&self) -> f64 { - if self.step_secs > 0.0 { self.residual_secs() / self.step_secs } else { 0.0 } - } - - /// Spread `secs` over the blocks framed, in microseconds -- the unit the - /// 24.3 us/block question is asked in. - #[allow(clippy::cast_precision_loss, reason = "block counts stay far below 2^52")] - pub(crate) fn per_block_micros(&self, secs: f64) -> f64 { - if self.blocks == 0 { 0.0 } else { secs * 1_000_000.0 / self.blocks as f64 } - } - - /// Throughput the refills actually achieved, in MB/s. - /// - /// The discriminant against the volume's measured ceiling: at the ceiling the - /// refills are disk-bound and the remaining time is framing; well under it, - /// the reader is leaving bandwidth on the floor and the fix is upstream of - /// framing entirely. - #[allow(clippy::cast_precision_loss, reason = "byte totals stay far below 2^52")] - pub(crate) fn refill_mb_per_sec(&self) -> f64 { - if self.refill_secs > 0.0 { self.refill_bytes as f64 / self.refill_secs / 1e6 } else { 0.0 } - } -} - -/// Times every `read()` made against whatever it wraps. -/// -/// Belongs *below* the input `BufReader`, so it sees buffer refills rather than -/// per-block reads -- see [`ReaderStats`] for why that placement is the whole -/// point. -pub(crate) struct TimedReader { - inner: R, - stats: std::sync::Arc, -} - -impl TimedReader { - pub(crate) fn new(inner: R, stats: std::sync::Arc) -> Self { - Self { inner, stats } - } -} - -impl std::io::Read for TimedReader { - fn read(&mut self, buf: &mut [u8]) -> std::io::Result { - // Two ~2 us proc reads per refill, against ~21k refills on a ~125s step: - // under 0.1% of the quantity measured, and the only way to tell a slower - // device from a thread that keeps losing the CPU. - let sched_before = thread_schedstat(); - let started = std::time::Instant::now(); - let result = self.inner.read(buf); - let elapsed = u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX); - if let (Some(before), Some(after)) = (sched_before, thread_schedstat()) { - self.stats.record_refill_sched( - after.0.saturating_sub(before.0), - after.1.saturating_sub(before.1), - after.2.saturating_sub(before.2), - ); - } - // A failed read still cost time; counting zero bytes for it keeps the - // throughput figure honest rather than crediting the failure with bytes. - self.stats.record_refill(elapsed, result.as_ref().copied().unwrap_or(0)); - result - } -} - -#[cfg(test)] -mod tests { - #[test] - fn test_schedstat_parses_the_kernel_format() { - // Format is " "; the kernel - // appends a trailing newline. - assert_eq!(parse_schedstat("12345 678 9\n"), Some((12345, 678, 9))); - assert_eq!(parse_schedstat("0 0 0"), Some((0, 0, 0))); - } - - #[test] - fn test_schedstat_refuses_a_line_it_does_not_understand() { - // A silently-wrong parse would report zero runqueue wait, which reads as - // "the reader never lost the CPU" -- the exact conclusion this - // measurement exists to test. Better to have no sample than a false one. - assert_eq!(parse_schedstat(""), None, "empty"); - assert_eq!(parse_schedstat("12345 678"), None, "truncated: only two fields"); - assert_eq!(parse_schedstat("version 15"), None, "not numeric"); - } - - use super::*; - - /// A reader report built from the production cell's shape: 124.5s of step - /// time over 5,114,408 blocks, of which ~71s is disk. - fn production_shape() -> ReaderReport { - ReaderReport { - refills: 20_560, - refill_secs: 71.0, - refill_bytes: 43_131_188_552, - batches: 319_651, - blocks: 5_114_408, - read_raw_secs: 121.4, - dispatch_secs: 3.1, - step_secs: 124.5, - ..ReaderReport::default() - } - } - - #[test] - fn test_framing_is_read_raw_time_with_the_refills_taken_out() { - // The whole point of putting a timer below the BufReader: the refill - // happens inside `read_raw_blocks`, so the raw span is framing and I/O - // fused. Reporting the raw span as framing would credit userspace with - // every second the disk spent. - let report = production_shape(); - assert!((report.framing_secs() - 50.4).abs() < 1e-9, "got {}", report.framing_secs()); - } - - #[test] - fn test_a_reader_report_partitions_the_step_with_a_signed_residual() { - let report = production_shape(); - assert!((report.residual_secs() - 0.0).abs() < 1e-9, "got {}", report.residual_secs()); - - // Over-attribution must be visible. If the parts claim more than the - // step, a clamped residual reports a tidy zero and the partition gets - // believed -- exactly how the merge's first partition summed to 321.5s - // of a 189.3s loop and survived a reading. - let over = ReaderReport { read_raw_secs: 200.0, ..production_shape() }; - assert!(over.residual_secs() < 0.0, "got {}", over.residual_secs()); - assert!(over.residual_share() < 0.0, "got {}", over.residual_share()); - } - - #[test] - fn test_per_block_micros_reproduces_the_number_under_investigation() { - // 124.5s over 5,114,408 blocks is the 24.3 us/block the campaign is - // trying to explain; the report has to say so in that unit. - let report = production_shape(); - let per_block = report.per_block_micros(report.step_secs); - assert!((per_block - 24.34).abs() < 0.01, "got {per_block}"); - assert!((report.per_block_micros(0.0) - 0.0).abs() < 1e-9); - } - - #[test] - fn test_refill_throughput_is_measured_against_the_bytes_actually_returned() { - // 43.1 GB in 71.0s is ~607 MB/s, which is what makes "the disk is at its - // ceiling" a claim rather than an assumption. - let report = production_shape(); - assert!( - (report.refill_mb_per_sec() - 607.5).abs() < 1.0, - "got {}", - report.refill_mb_per_sec() - ); - assert!((ReaderReport::default().refill_mb_per_sec() - 0.0).abs() < 1e-9); - } - - #[test] - fn test_a_failed_read_is_timed_but_credited_no_bytes() { - use std::io::Read; - struct Failing; - impl Read for Failing { - fn read(&mut self, _buf: &mut [u8]) -> std::io::Result { - Err(std::io::Error::other("device fell off")) - } - } - let stats = std::sync::Arc::new(ReaderStats::default()); - let mut reader = TimedReader::new(Failing, std::sync::Arc::clone(&stats)); - let mut buf = [0u8; 8]; - assert!(reader.read(&mut buf).is_err()); - - let report = stats.snapshot(0.0); - assert_eq!(report.refills, 1, "the attempt still cost time and must be counted"); - assert_eq!(report.refill_bytes, 0, "a failed read must not inflate throughput"); - } - - #[test] - fn test_the_residual_is_signed_so_over_attribution_is_visible() { - // The merge's first partition summed to 321.5s of a 189.3s loop. A - // clamped residual would have shown 0.0 and the numbers would have been - // believed; the sign is what exposed the clock overhead inside them. - let over = IngestPartition { - segments: IngestSample { fetch: 200.0, ..IngestSample::default() }, - read_secs: 100.0, - park_secs: 0.0, - }; - assert!(over.residual_secs() < 0.0, "got {}", over.residual_secs()); - assert!((over.residual_share() + 1.0).abs() < 1e-9, "got {}", over.residual_share()); - } - - #[test] - fn test_clock_correction_subtracts_one_pair_per_segment_per_sample() { - // Ten samples, 20 ns per pair, four segments: each segment carries - // 10 x 20 ns = 200 ns of clock, and each is corrected independently. - // `tick` is deliberately *below* that, to exercise the clamp. - let raw = IngestSample { fetch: 1e-6, push: 1e-6, tick: 1e-7, probe: 1e-6 }; - let fixed = raw.corrected(10, 20); - assert!((fixed.fetch - 0.8e-6).abs() < 1e-12, "got {}", fixed.fetch); - // A segment cheaper than the clock that measured it clamps to zero rather - // than going negative, which would read as a bug rather than as - // "unresolvable by this method". - assert!((fixed.tick - 0.0).abs() < 1e-12, "got {}", fixed.tick); - } - - #[test] - fn test_scaling_happens_before_correction_is_meaningful() { - // Scale multiplies the sampled segments up to the whole loop; correction - // works on the sampled scale. Applying them in the wrong order would - // subtract one pair's cost from the *scaled* total rather than from each - // sample, understating the correction by the scale factor. - let raw = IngestSample { push: 2e-6, ..IngestSample::default() }; - let corrected_then_scaled = raw.corrected(10, 20).scaled(1000.0); - let scaled_then_corrected = raw.scaled(1000.0).corrected(10, 20); - assert!(corrected_then_scaled.push < scaled_then_corrected.push); - } - - #[test] - fn test_park_causes_are_counted_separately() { - let stats = Phase1IngestStats::default(); - stats.record_park(1_000, ParkCause::Starved); - stats.record_park(3_000, ParkCause::HeadOfLine); - stats.record_park(6_000, ParkCause::HeadOfLine); - - let report = stats.snapshot(); - assert_eq!(report.parks, 3); - assert_eq!(report.parks_starved, 1); - assert_eq!(report.parks_head_of_line, 2); - // Starvation and head-of-line blocking have different fixes, so a report - // that only totalled them would not distinguish the two. - assert_eq!(report.parks_starved + report.parks_head_of_line, report.parks); - assert!((report.park_secs - 10e-6).abs() < 1e-12, "got {}", report.park_secs); - } - - #[test] - fn test_the_spill_handoff_wait_is_counted_on_its_own() { - let stats = Phase1IngestStats::default(); - stats.record_park(2_000_000_000, ParkCause::Starved); - stats.record_spill_wait(3_000_000_000); - - let report = stats.snapshot(); - // The spill wait sits between the read span ending and the sort starting, - // so it is in no phase bucket. It is reported separately from the park - // rather than summed into it: one is the pool failing to keep the ingest - // thread fed, the other is the previous chunk's write not being done, and - // a single "blocked" total would hide which. - assert!((report.spill_wait_secs - 3.0).abs() < 1e-9, "got {}", report.spill_wait_secs); - assert!((report.park_secs - 2.0).abs() < 1e-9, "got {}", report.park_secs); - assert_eq!(report.spill_waits, 1); - } - - #[test] - fn test_mean_park_is_absent_rather_than_zero_when_it_never_parked() { - let idle = Phase1IngestStats::default().snapshot(); - // A reported 0 us mean would read as "parked, instantly", which is the - // opposite of what never parking means. - assert_eq!(idle.mean_park_micros(), None); - - let stats = Phase1IngestStats::default(); - stats.record_park(4_000_000, ParkCause::Starved); - stats.record_park(6_000_000, ParkCause::Starved); - let mean = stats.snapshot().mean_park_micros().expect("parked twice"); - assert!((mean - 5_000.0).abs() < 1e-6, "got {mean}"); - } -} diff --git a/crates/fgumi-sort/src/pooled_bam_writer.rs b/crates/fgumi-sort/src/pooled_bam_writer.rs index 13f8de13f..b7880b2b9 100644 --- a/crates/fgumi-sort/src/pooled_bam_writer.rs +++ b/crates/fgumi-sort/src/pooled_bam_writer.rs @@ -63,18 +63,6 @@ struct IndexState { } impl PooledBamWriter { - /// The permit pool carrying this writer's histograms, if the writer has not - /// yet been finalized. - /// - /// Retain this [`Arc`] before [`finish`](Self::finish) and read - /// [`PermitPool::writer_stats`] afterwards: the pool outlives the writer, so - /// the snapshot then includes the block writes and reorder waits performed - /// during the output drain that `finish` runs — a snapshot taken through the - /// live writer would omit that tail. - pub(crate) fn permit_pool(&self) -> Option> { - self.staging.as_ref().map(|staging| Arc::clone(staging.permit_pool())) - } - /// Seconds the producer spent blocked waiting for an output permit, and the /// number of waits. /// @@ -442,7 +430,7 @@ mod tests { // Oversized (unmapped) record that spans multiple BGZF blocks. records.push(make_test_record(b"oversized", BGZF_MAX_BLOCK_SIZE + 500)); - let pool = Arc::new(SortWorkerPool::new(4, 1, 6, crate::codec::SpillCodec::Bgzf, false)); + let pool = Arc::new(SortWorkerPool::new(4, 1, 6, crate::codec::SpillCodec::Bgzf)); let plain_path = dir.path().join("plain.bam"); { @@ -493,7 +481,7 @@ mod tests { let dir = tempfile::TempDir::new().expect("tempdir"); let bam_path = dir.path().join("test.bam"); let header = test_header(); - let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf, false)); + let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf)); let num_records = 200; let records: Vec> = (0..num_records) @@ -540,7 +528,7 @@ mod tests { #[case::dev_stdout("/dev/stdout")] fn test_pooled_bam_writer_indexing_rejects_stdout(#[case] spelling: &str) { let header = test_header(); - let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf, false)); + let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf)); // Matched rather than `expect_err`: the writer is deliberately not // `Debug`, so unwrapping the error out of the `Result` does not compile. @@ -569,7 +557,7 @@ mod tests { let dir = tempfile::TempDir::new().expect("tempdir"); let bam_path = dir.path().join("empty.bam"); let header = test_header(); - let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf, false)); + let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf)); { let writer = @@ -599,7 +587,7 @@ mod tests { let dir = tempfile::TempDir::new().expect("tempdir"); let bam_path = dir.path().join("many.bam"); let header = test_header(); - let pool = Arc::new(SortWorkerPool::new(4, 1, 6, crate::codec::SpillCodec::Bgzf, false)); + let pool = Arc::new(SortWorkerPool::new(4, 1, 6, crate::codec::SpillCodec::Bgzf)); let num_records = 5000; { @@ -628,57 +616,12 @@ mod tests { } } - /// The writer histograms must be harvested *after* `finish` drains the output - /// queue, not before: the drain flushes the final partial block (and any - /// still-queued blocks) through the I/O thread, so a snapshot taken while the - /// writer is still alive omits those drain-time writes. Retaining the permit - /// pool and reading [`PermitPool::writer_stats`] after `finish` captures the - /// full write count. - #[test] - fn test_writer_stats_include_finish_drain() { - let dir = tempfile::TempDir::new().expect("tempdir"); - let bam_path = dir.path().join("drain.bam"); - let header = test_header(); - let pool = Arc::new(SortWorkerPool::new(4, 1, 6, crate::codec::SpillCodec::Bgzf, false)); - - // Retain the permit pool before finalizing, exactly as the merge path does, - // and snapshot the block-write count both before and after the drain. - let (before, after) = { - let mut writer = - PooledBamWriter::new(Arc::clone(&pool), &bam_path, &header).expect("create writer"); - for i in 0..5000 { - let rec = make_test_record(format!("read_{i:06}").as_bytes(), 100); - writer.write_raw_record(&rec).expect("write record"); - } - let permit_pool = writer.permit_pool().expect("permit pool present before finish"); - - let before = permit_pool.writer_stats().0.count; - writer.finish().expect("finish writer"); - let after = permit_pool.writer_stats().0.count; - (before, after) - }; - - // The final flush and drain happen inside `finish`, so the post-drain - // count must exceed the pre-drain count -- the exact regression the - // pre-finalize snapshot silently dropped. - assert!(after > 0, "drain-inclusive snapshot must record block writes, got {after}"); - assert!( - after > before, - "finish drain must add block writes the pre-finish snapshot missed: \ - before={before}, after={after}" - ); - - if let Ok(pool) = Arc::try_unwrap(pool) { - pool.shutdown(); - } - } - #[test] fn test_pooled_bam_writer_raw_bytes_match() { let dir = tempfile::TempDir::new().expect("tempdir"); let bam_path = dir.path().join("raw_match.bam"); let header = test_header(); - let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf, false)); + let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf)); let records: Vec> = (0..50).map(|i| make_test_record(format!("r{i}").as_bytes(), 30)).collect(); @@ -715,7 +658,7 @@ mod tests { let dir = tempfile::TempDir::new().expect("tempdir"); let bam_path = dir.path().join("oversized.bam"); let header = test_header(); - let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf, false)); + let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf)); // A sequence of BGZF_MAX_BLOCK_SIZE bytes exceeds the threshold. let oversized_rec = make_test_record(b"oversized_read", BGZF_MAX_BLOCK_SIZE); @@ -750,7 +693,7 @@ mod tests { let dir = tempfile::TempDir::new().expect("tempdir"); let bam_path = dir.path().join("dropped_writer.bam"); let header = test_header(); - let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf, false)); + let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf)); { let mut writer = diff --git a/crates/fgumi-sort/src/pooled_chunk_writer.rs b/crates/fgumi-sort/src/pooled_chunk_writer.rs index 46706e520..fdd43221e 100644 --- a/crates/fgumi-sort/src/pooled_chunk_writer.rs +++ b/crates/fgumi-sort/src/pooled_chunk_writer.rs @@ -28,7 +28,7 @@ use crate::keys::RawSortKey; use crate::worker_pool::{CompressResult, CompressTarget, PermitPool, SortWorkerPool}; use anyhow::Result; use crossbeam_channel::bounded; -use fgumi_bgzf::BGZF_EOF; +use fgumi_bgzf::{BGZF_EOF, BGZF_MAX_BLOCK_SIZE}; use std::io::BufWriter; use std::marker::PhantomData; use std::path::Path; @@ -45,16 +45,6 @@ pub struct PooledChunkWriter { staging: Option, /// Reusable scratch buffer for key serialization (non-embedded keys only). key_buf: Vec, - /// Bytes a frame may hold, resolved from the codec once. - /// - /// This writer pre-flushes so a record never straddles a frame boundary -- - /// which is what lets the merge borrow most records in place instead of - /// reassembling them. That budget must be the size the staging buffer will - /// actually flush at: when this was `BGZF_MAX_BLOCK_SIZE` outright, it - /// pre-flushed at 64 KiB no matter what the staging buffer was configured - /// for, so raising the frame size changed the block count not at all - /// (measured: 5,368,249 blocks at both 64 KiB and 256 KiB). - frame_bytes: usize, io_handle: Option>>, _phantom: PhantomData, } @@ -171,7 +161,6 @@ impl PooledChunkWriter { CompressTarget::Spill, )), key_buf: Vec::new(), - frame_bytes: crate::bgzf_io::spill_frame_bytes(codec), io_handle: Some(io_handle), _phantom: PhantomData, }) @@ -196,12 +185,11 @@ impl PooledChunkWriter { // Fast path: key is part of the record bytes, no extra serialization. // Budget: 4-byte length prefix + record bytes. let needed = 4 + record.len(); - let frame = self.frame_bytes; - if staging.buf().len() + needed > frame { + if staging.buf().len() + needed > BGZF_MAX_BLOCK_SIZE { staging.flush()?; } staging.buf().extend_from_slice(&(record.len() as u32).to_le_bytes()); - if record.len() > frame.saturating_sub(4) { + if record.len() > BGZF_MAX_BLOCK_SIZE.saturating_sub(4) { staging.write_chunked(record)?; } else { staging.buf().extend_from_slice(record); @@ -213,10 +201,10 @@ impl PooledChunkWriter { self.key_buf.clear(); key.write_to(&mut self.key_buf)?; let needed = self.key_buf.len() + 4 + record.len(); - // No size limit check: records larger than one frame are handled by - // write_chunked(), which splits them across multiple blocks. The + // No size limit check: records larger than one BGZF block are handled + // by write_chunked(), which splits them across multiple blocks. The // reader uses streaming read_exact() that transparently spans blocks. - if staging.buf().len() + needed > self.frame_bytes { + if staging.buf().len() + needed > BGZF_MAX_BLOCK_SIZE { staging.flush()?; } staging.buf().extend_from_slice(&self.key_buf); @@ -226,16 +214,6 @@ impl PooledChunkWriter { Ok(()) } - /// The frame budget this writer pre-flushes against. - /// - /// Exposed so it can be pinned to [`crate::bgzf_io::spill_frame_bytes`] - /// rather than trusted to match it. - #[cfg(test)] - #[must_use] - pub(crate) fn frame_bytes(&self) -> usize { - self.frame_bytes - } - /// Finish writing: flush remaining data, wait for I/O thread. /// /// # Errors @@ -347,7 +325,7 @@ mod tests { fn test_appending_leaves_one_bgzf_terminator() { let dir = TempDir::new().expect("tempdir"); let path = dir.path().join("run.keyed"); - let pool = Arc::new(SortWorkerPool::new(2, 1, 6, SpillCodec::Bgzf, false)); + let pool = Arc::new(SortWorkerPool::new(2, 1, 6, SpillCodec::Bgzf)); // Two chunks, the second appended to the first. for (chunk, appending) in [(0u64, false), (1u64, true)] { @@ -402,7 +380,7 @@ mod tests { // auto-detects the magic and routes to `ZspillStreamReader`. let dir = TempDir::new().unwrap(); let chunk_path = dir.path().join("test_chunk_zstd.keyed"); - let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Zstd, false)); + let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Zstd)); let records: Vec<(TemplateKey, Vec)> = (0..100) .map(|i| { @@ -463,7 +441,7 @@ mod tests { fn test_pooled_writer_roundtrip() { let dir = TempDir::new().unwrap(); let chunk_path = dir.path().join("test_chunk.keyed"); - let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf, false)); + let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf)); let records: Vec<(TemplateKey, Vec)> = (0..100) .map(|i| { @@ -513,7 +491,7 @@ mod tests { fn test_pooled_writer_empty() { let dir = TempDir::new().unwrap(); let chunk_path = dir.path().join("empty_chunk.keyed"); - let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf, false)); + let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf)); { let writer = PooledChunkWriter::::new( @@ -539,7 +517,7 @@ mod tests { fn test_pooled_writer_large_records() { let dir = TempDir::new().unwrap(); let chunk_path = dir.path().join("large_chunk.keyed"); - let pool = Arc::new(SortWorkerPool::new(4, 1, 6, crate::codec::SpillCodec::Bgzf, false)); + let pool = Arc::new(SortWorkerPool::new(4, 1, 6, crate::codec::SpillCodec::Bgzf)); let records: Vec<(TemplateKey, Vec)> = (0..500) .map(|i| { @@ -587,7 +565,7 @@ mod tests { // `handle.wait()` must join it and surface any errors. let dir = TempDir::new().unwrap(); let chunk_path = dir.path().join("pipelined_chunk.keyed"); - let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf, false)); + let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf)); let records: Vec<(TemplateKey, Vec)> = (0..50).map(|i| (make_key(i), vec![(i % 256) as u8; 100])).collect(); @@ -630,7 +608,7 @@ mod tests { // the `Drop` impl joins the thread and logs any error. let dir = TempDir::new().unwrap(); let chunk_path = dir.path().join("dropped_chunk.keyed"); - let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf, false)); + let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf)); let handle = { let mut writer = PooledChunkWriter::::new( @@ -661,7 +639,7 @@ mod tests { // deadlock — the Drop impl signals the I/O thread and joins it. let dir = TempDir::new().unwrap(); let chunk_path = dir.path().join("dropped_writer.keyed"); - let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf, false)); + let pool = Arc::new(SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf)); { let mut writer = PooledChunkWriter::::new( diff --git a/crates/fgumi-sort/src/progress_batch.rs b/crates/fgumi-sort/src/progress_batch.rs deleted file mode 100644 index 67831a5c1..000000000 --- a/crates/fgumi-sort/src/progress_batch.rs +++ /dev/null @@ -1,106 +0,0 @@ -//! Batched progress ticks for per-record loops. - -use fgumi_bam_io::progress::ProgressTracker; - -/// Records counted locally before touching the tracker's atomic. -/// -/// A merge emits hundreds of millions of records on ONE serial thread, and -/// `ProgressTracker::log_if_needed` does a relaxed `fetch_add` per call. On -/// aarch64 Rust defaults to outline-atomics, so that is not an inline `LDADD` -/// but a call into `__aarch64_ldadd8_relax`, which runtime-checks for LSE and -/// can fall back to an `ldxr`/`stxr` loop. Profiling the merge consumer on -/// `c7g.4xlarge` put that helper at **28% of the thread's cycles** -- the single -/// largest entry, ahead of `LoserTree::replay` at 11% and the record memcpy at -/// 8% -- on a thread whose 145 ns/record now sets the merge's wall clock. -/// -/// 4096 against a 1,000,000-record logging interval means a milestone is noticed -/// at most 0.4% late, which no progress line cares about, and removes 99.98% of -/// the atomic traffic. -const TICK_BATCH: u64 = 4096; - -/// Accumulates record counts and forwards them to a [`ProgressTracker`] in -/// batches. -/// -/// [`Self::flush`] must be called before the tracker's final log, or the last -/// partial batch is never counted and the totals come up short. -pub(crate) struct BatchedProgress { - pending: u64, -} - -impl BatchedProgress { - pub(crate) const fn new() -> Self { - Self { pending: 0 } - } - - /// Count one record, forwarding to `tracker` once a batch has accumulated. - #[inline] - pub(crate) fn tick(&mut self, tracker: &ProgressTracker) { - self.pending += 1; - if self.pending >= TICK_BATCH { - tracker.log_if_needed(self.pending); - self.pending = 0; - } - } - - /// Forward whatever is left. Idempotent. - pub(crate) fn flush(&mut self, tracker: &ProgressTracker) { - if self.pending > 0 { - tracker.log_if_needed(self.pending); - self.pending = 0; - } - } - - /// Records not yet forwarded. For tests. - #[cfg(test)] - pub(crate) const fn pending(&self) -> u64 { - self.pending - } -} - -#[cfg(test)] -mod tests { - use super::*; - - /// The whole point is that the tracker's total is unchanged: batching may - /// notice a milestone late, but it must never lose a record. - #[test] - fn test_batched_ticks_preserve_the_total() { - let tracker = ProgressTracker::new("test").with_interval(1_000_000); - let mut batch = BatchedProgress::new(); - for _ in 0..10_000 { - batch.tick(&tracker); - } - batch.flush(&tracker); - assert_eq!(tracker.count(), 10_000, "every record must reach the tracker"); - } - - /// Without the flush the last partial batch is silently dropped, which is the - /// one way this change can corrupt a total. - #[test] - fn test_unflushed_remainder_is_the_only_loss_and_flush_recovers_it() { - let tracker = ProgressTracker::new("test").with_interval(1_000_000); - let mut batch = BatchedProgress::new(); - for _ in 0..(TICK_BATCH + 5) { - batch.tick(&tracker); - } - assert_eq!(tracker.count(), TICK_BATCH, "one full batch forwarded"); - assert_eq!(batch.pending(), 5, "remainder still held locally"); - - batch.flush(&tracker); - assert_eq!(tracker.count(), TICK_BATCH + 5); - batch.flush(&tracker); - assert_eq!(tracker.count(), TICK_BATCH + 5, "flush is idempotent"); - } - - #[test] - fn test_atomic_is_touched_once_per_batch_not_once_per_record() { - let tracker = ProgressTracker::new("test").with_interval(1_000_000); - let mut batch = BatchedProgress::new(); - for _ in 0..(TICK_BATCH - 1) { - batch.tick(&tracker); - } - assert_eq!(tracker.count(), 0, "nothing forwarded until a batch fills"); - batch.tick(&tracker); - assert_eq!(tracker.count(), TICK_BATCH); - } -} diff --git a/crates/fgumi-sort/src/read_ahead.rs b/crates/fgumi-sort/src/read_ahead.rs index f3cfa69b3..ab5e5118a 100644 --- a/crates/fgumi-sort/src/read_ahead.rs +++ b/crates/fgumi-sort/src/read_ahead.rs @@ -291,8 +291,6 @@ pub struct PooledInputStream { current_buf: Vec, /// Read position within `current_buf`. current_pos: usize, - /// Where this stream records the ingest thread's waits. - stats: std::sync::Arc, /// Reusable scratch buffer for records (or their length prefixes) that /// straddle a decompressed-block boundary and therefore cannot be borrowed /// directly out of `current_buf`. See [`PooledInputStream::next_record_borrowed`]. @@ -307,14 +305,12 @@ impl PooledInputStream { decompressed_input_done: std::sync::Arc, input_read_error: std::sync::Arc, decompression_error: std::sync::Arc, - stats: std::sync::Arc, ) -> Self { Self { decompressed_input, decompressed_input_done, input_read_error, decompression_error, - stats, reorder: fgumi_bam_io::ReorderBuffer::new(), current_buf: Vec::new(), current_pos: 0, @@ -387,27 +383,8 @@ impl PooledInputStream { return None; } - // Park until a worker pushes a block and calls unpark(). - // - // Timed exactly rather than sampled: a park is microseconds to - // milliseconds against a ~30 ns clock read, so the clock is orders of - // magnitude below the quantity it measures. The cause is captured - // before parking because it is not recoverable afterwards -- and the - // two causes have different fixes. An empty reorder buffer means the - // workers are behind; a non-empty one means blocks are ready and the - // serial order will not let the consumer have them, which no amount - // of extra decompression capacity would help. - let cause = if self.reorder.buffer_len() == 0 { - crate::phase1_stats::ParkCause::Starved - } else { - crate::phase1_stats::ParkCause::HeadOfLine - }; - let parked_at = std::time::Instant::now(); + // Park until a worker pushes a block and calls unpark() std::thread::park(); - self.stats.record_park( - u64::try_from(parked_at.elapsed().as_nanos()).unwrap_or(u64::MAX), - cause, - ); // After waking, check for errors before looping back to drain. // A worker may have set an error flag instead of pushing a block. @@ -422,7 +399,7 @@ impl PooledInputStream { /// Read the next raw BAM record, borrowing its bytes from the current /// decompressed block when possible. /// - /// This is the borrow-in-place counterpart to [`read_raw_record`]: it removes + /// This is the borrow-in-place counterpart to [`fgumi_raw_bam::read_raw_record`]: it removes /// the per-record `read_exact` copy into a `RawRecord` on the common path /// where the record body lies wholly within the current decompressed block. /// @@ -437,7 +414,7 @@ impl PooledInputStream { /// The returned slice borrows `self`; it is invalidated by the next call to /// any method on this stream. Returns `Ok(None)` at clean EOF. /// - /// A `block_size` of 0 is treated as EOF, mirroring [`read_raw_record`]. + /// A `block_size` of 0 is treated as EOF, mirroring [`fgumi_raw_bam::read_raw_record`]. /// /// # Errors /// @@ -829,7 +806,6 @@ mod tests { Arc::new(AtomicBool::new(true)), // decompressed_input_done Arc::new(AtomicBool::new(false)), // input_read_error Arc::new(AtomicBool::new(false)), // decompression_error - Arc::default(), // ingest wait counters ) } @@ -937,7 +913,6 @@ mod tests { Arc::new(AtomicBool::new(true)), Arc::new(AtomicBool::new(false)), Arc::new(AtomicBool::new(false)), - Arc::default(), ); let err = pooled.next_record_borrowed().expect_err("truncated body should error"); assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof); diff --git a/crates/fgumi-sort/src/ref_sort.rs b/crates/fgumi-sort/src/ref_sort.rs new file mode 100644 index 000000000..a2f154ced --- /dev/null +++ b/crates/fgumi-sort/src/ref_sort.rs @@ -0,0 +1,415 @@ +//! Build a sorted coordinate [`InMemoryChunk`] from record bodies already resident +//! in a shared arena, WITHOUT copying record bytes — the core of the parallel-inflate +//! ingest. The chunk's `(offset, len)` point at record bodies in the supplied arena; +//! the existing Phase-2 merge consumes it unchanged. + +use std::cmp::Ordering; +use std::sync::Arc; + +use voracious_radix_sort::{RadixSort, Radixable}; + +use crate::arena_pool::PooledSegmentedBuf; +use crate::inline::{ + InMemoryChunk, RecordRef, extract_coordinate_key_inline, radix_sort_record_refs, +}; +use crate::keys::{RawCoordinateKey, RawSortKey}; + +/// Arrays smaller than this sort faster single-threaded (the parallel radix's +/// partition + thread-coordination overhead exceeds its benefit on small inputs; +/// microbench-tuned). Shared with the template arena builder +/// (`template_arena::template_chunk_from_arena_refs`) so both order-specific +/// radix fronts fall back to the serial path at the same input size. +pub(crate) const PARALLEL_SORT_THRESHOLD: usize = 256 * 1024; + +/// `repr(transparent)` view of a [`RecordRef`] that orders by the +/// `(sort_key, offset)` **composite** rather than by `sort_key` alone. This lets +/// `voracious`'s (unstable) parallel radix produce a STABLE-equivalent coordinate +/// order: within one chunk every record has a distinct `offset` that increases in +/// input order, so `(sort_key, offset)` is a total order whose ties resolve to +/// input order — byte-identical to the stable single-threaded radix. Wrapping +/// (instead of changing `RecordRef`'s own `PartialEq`) keeps `RecordRef`'s +/// existing key-only equality untouched for the rest of the engine. +#[repr(transparent)] +#[derive(Copy, Clone)] +struct CoordSortRef(RecordRef); + +/// `sort_coordinate_refs` reinterprets `*mut [RecordRef]` as `*mut [CoordSortRef]`. +/// A fat pointer carries an element count, not a byte length, so the cast is +/// sound only while the two types share size and alignment — a stride mismatch +/// would walk the slice off its end. `repr(transparent)` guarantees that today; +/// this pins it, so adding a field or dropping the attribute is a compile error +/// rather than silent unsoundness. Worth the two lines because this crate is not +/// covered by the Miri job (`miri.yml` runs only `-p fgumi-raw-bam sort`). +const _: () = { + assert!(size_of::() == size_of::()); + assert!(align_of::() == align_of::()); +}; + +impl CoordSortRef { + #[inline] + fn composite(self) -> u128 { + (u128::from(self.0.sort_key) << 64) | u128::from(self.0.offset) + } +} + +impl PartialOrd for CoordSortRef { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.composite().cmp(&other.composite())) + } +} + +impl PartialEq for CoordSortRef { + fn eq(&self, other: &Self) -> bool { + self.composite() == other.composite() + } +} + +impl Radixable for CoordSortRef { + type Key = u128; + #[inline] + fn key(&self) -> u128 { + self.composite() + } +} + +/// Sort coordinate `RecordRef`s in place, stably by `sort_key`. +/// +/// For large inputs with `sort_threads > 1`, uses `voracious`'s parallel radix +/// over the `(sort_key, offset)` composite (see [`CoordSortRef`]) — measured ~4× +/// faster than the single-threaded radix at realistic coordinate-key widths. For +/// small inputs or a single thread it falls back to the single-threaded stable LSD +/// radix (the parallel path's overhead does not pay off there). Both produce +/// byte-identical output (verified by the ref-sort oracle tests) — but only +/// under the precondition below. +/// +/// # Precondition: `offset` ascends with input order +/// +/// The two branches break equal-key ties differently. The serial radix is +/// *stable*, so ties keep input order. The parallel radix sorts the +/// `(sort_key, offset)` composite, so ties order by `offset`. They agree only +/// while `offset` ascends with input order — which holds for every caller today, +/// because refs are built by appending each record to an arena in input order. +/// +/// A caller that broke it (refs assembled from a reordered source, or from two +/// arenas) would get a different record order depending on `sort_threads` and +/// input size, and only above `PARALLEL_SORT_THRESHOLD` — an output-identity +/// divergence invisible to any test whose refs happen to be built in order, +/// which includes `parallel_coordinate_sort_matches_serial_radix_at_threshold` +/// (it constructs `offset: i`). Hence the debug assertion rather than a comment. +fn sort_coordinate_refs(refs: &mut [RecordRef], sort_threads: usize) { + if sort_threads > 1 && refs.len() >= PARALLEL_SORT_THRESHOLD { + debug_assert!( + refs.windows(2).all(|w| w[0].offset <= w[1].offset), + "parallel coordinate sort requires `offset` to ascend with input order: it breaks \ + equal-key ties by `offset`, while the serial radix breaks them by input order, so \ + the two diverge on ties when this does not hold", + ); + // SAFETY: `CoordSortRef` is `#[repr(transparent)]` over `RecordRef`, so the + // two have identical size, alignment, and layout; reinterpreting the slice + // (via a pointer cast — clippy rejects a ref-to-ref `transmute`) is sound. + // This is the only `unsafe` in this module's production path. + #[allow(unsafe_code)] + let view: &mut [CoordSortRef] = + unsafe { &mut *(std::ptr::from_mut::<[RecordRef]>(refs) as *mut [CoordSortRef]) }; + view.voracious_mt_sort(sort_threads); + } else { + radix_sort_record_refs(refs); + } +} + +/// Build a sorted coordinate chunk from record bodies already resident in `arena`. +/// +/// `bodies[i] = (body_offset, body_len)` locate each record's BAM body (the 4-byte +/// `block_size` prefix already excluded) within `arena`. The coordinate key is +/// extracted from each body exactly as `RecordBuffer::push_coordinate` does, the +/// refs are radix-sorted, and the same `arena` is wrapped into the returned chunk +/// (no record bytes are copied). The chunk's `(offset, len)` therefore point at the +/// bodies in `arena`, byte-identical to what the copy-based sorter materializes. +#[must_use] +#[allow(clippy::cast_possible_truncation)] // offset/len fit in usize on all supported targets (BAM is < 4 GiB; arenas < address space) +pub fn coordinate_chunk_from_arena_refs( + arena: Arc, + bodies: &[(u64, u32)], + n_ref: u32, + sort_threads: usize, +) -> InMemoryChunk { + let mut refs: Vec = Vec::with_capacity(bodies.len()); + for &(offset, len) in bodies { + let body = arena.slice(offset as usize, len as usize); + let sort_key = extract_coordinate_key_inline(body, n_ref); + refs.push(RecordRef::new(sort_key, offset, len)); + } + coordinate_chunk_from_refs(arena, refs, sort_threads) +} + +/// Build a sorted coordinate chunk from already-extracted `RecordRef`s (each +/// carries its coordinate `sort_key` and the `(offset, len)` of its body in +/// `arena`). Radix-sorts the refs and wraps `arena` into the returned chunk — +/// NO walk of the arena and NO record-byte copies. +/// +/// This is the fused entry point: the caller (`FindBoundariesAndSort`) extracts +/// the key during its single boundary scan and hands the refs here, so the arena +/// is walked exactly once for the whole run (the previous two-step +/// [`coordinate_chunk_from_arena_refs`] re-walked it to re-derive each key). +#[must_use] +pub fn coordinate_chunk_from_refs( + arena: Arc, + mut refs: Vec, + sort_threads: usize, +) -> InMemoryChunk { + sort_coordinate_refs(&mut refs, sort_threads); + let records: Vec<(RawCoordinateKey, u64, u32)> = + refs.iter().map(|r| (RawCoordinateKey { sort_key: r.sort_key }, r.offset, r.len)).collect(); + InMemoryChunk::from_parts(arena, records) +} + +/// Build a sorted queryname chunk from already-extracted `(key, offset, len)` +/// refs pointing into `arena`. Unlike the coordinate/template builders, the +/// queryname key is a variable-length read name (not radix-able), so this sorts +/// by the key comparator with `par_sort_unstable_by` — matching the legacy +/// queryname path's unstable comparator sort. No record bytes are copied: the +/// returned chunk references its bodies in `arena`. +/// +/// **Threading contract (differs from the coordinate/template radix builders on +/// purpose).** Those take an explicit `sort_threads` and gate the parallel radix +/// on both `sort_threads > 1` and `PARALLEL_SORT_THRESHOLD`. This builder takes +/// no `sort_threads`: the caller (`QuerynameStrategy::seal`, in `fgumi-pipeline-io`) +/// runs it inside a rayon pool already sized to `sort_threads`, so parallelism is +/// bounded there, +/// and `rayon::par_sort_unstable_by` has its own internal sequential cutoff for +/// small slices — so a tiny run does not pay a partition/coordination cost even +/// without an explicit size threshold here. +#[must_use] +pub fn queryname_chunk_from_arena_refs( + arena: Arc, + mut records: Vec<(K, u64, u32)>, +) -> InMemoryChunk { + use rayon::prelude::*; + // `par_sort_unstable_by` gives no order guarantee for records whose keys + // compare equal (exact name+flag-class ties). samtools resolves those + // remaining ties by input order; the arena assigns offsets in ingest order, + // so break exact ties by ascending offset to preserve that (and keep output + // byte-identical run to run). + records.par_sort_unstable_by(|a, b| a.0.cmp(&b.0).then_with(|| a.1.cmp(&b.1))); + InMemoryChunk::from_parts(arena, records) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The parallel `voracious_mt_sort` coordinate path only fires at + /// `sort_threads > 1 && refs.len() >= PARALLEL_SORT_THRESHOLD`; the oracle + /// tests use small, single-thread inputs and never reach it. Build a + /// threshold-sized input with deliberate key ties and pin the parallel path's + /// ordering against the single-threaded stable radix (`radix_sort_record_refs`, + /// which the oracle already proves matches the copy sorter). Both order by + /// `sort_key`, breaking ties by input order (== ascending `offset` here), so + /// they must agree on `(sort_key, offset)` for every record — a divergence + /// would mean the parallel composite mis-sorts or loses stability. + #[test] + fn parallel_coordinate_sort_matches_serial_radix_at_threshold() { + let n = PARALLEL_SORT_THRESHOLD; // exactly at the parallel cutoff + // Keys in a small range → ~260 ties per key across the input, so a + // stability bug in the parallel path would reorder equal-key records. + let refs: Vec = (0..n) + .map(|i| { + let i = i as u64; + RecordRef::new(i.wrapping_mul(2_654_435_761) % 997, i, 1) + }) + .collect(); + let mut parallel = refs.clone(); + let mut serial = refs; + sort_coordinate_refs(&mut parallel, 4); // threads > 1 + len >= threshold → parallel path + radix_sort_record_refs(&mut serial); // single-threaded stable reference + assert_eq!(parallel.len(), serial.len()); + for (p, s) in parallel.iter().zip(&serial) { + assert_eq!( + (p.sort_key, p.offset), + (s.sort_key, s.offset), + "parallel voracious path diverged from the serial radix" + ); + } + } + use crate::arena_pool::PooledSegmentedBuf; + use crate::chunk_sorter::CoordinateChunkSorter; + use crate::segmented_buf::SegmentedBuf; + use std::sync::Arc; + + // Build a minimal valid BAM record body (no block_size prefix) with the given + // refId/pos and a one-character read name `name`. The read name is NOT part of + // the coordinate sort key (which reads only refId/pos/flags at bytes 0-15), so + // two records with the same (refId,pos) but different `name` are a coordinate + // TIE whose relative order is observable in the output bytes — that is what lets + // these tests actually witness a stable-vs-unstable sort difference. Returns the + // body bytes. + fn coord_body(ref_id: i32, pos: i32, name: u8) -> Vec { + // Minimal BAM record body: refID(4) pos(4) l_read_name(1) mapq(1) bin(2) + // n_cigar_op(2) flag(2) l_seq(4) next_refID(4) next_pos(4) tlen(4) + read_name. + 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 (distinguishes tied records) + b.push(0); // read_name NUL terminator + b + } + + #[allow(unsafe_code)] + #[test] + fn ref_sort_matches_copy_based_sorter_byte_for_byte() { + // Records deliberately out of coordinate order. + let n_ref = 4u32; + // Distinct read-name bytes make every record byte-unique, so a stability + // divergence (reordering the (0,50) tie at indices 1 and 3) would change the + // output byte sequence and fail the assertion below. + 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 index 1 → stable sort must keep b before d + coord_body(1, 999, b'e'), + ]; + + // ---- Path A: the existing copy-based sorter (the ORACLE) ---- + let mut oracle = CoordinateChunkSorter::for_test(usize::MAX, n_ref); // see note below + for r in &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(); + + // ---- Path B: the ref-sort over an arena holding verbatim [block_size][body] ---- + let mut arena = SegmentedBuf::with_capacity(0, 1 << 20); + arena.reserve_full_capacity(); + let mut bodies = Vec::new(); + for r in &recs { + // store verbatim record: 4-byte block_size (= body len) then body. + let block_size = u32::try_from(r.len()).unwrap(); + // SAFETY: each slot fully written before any read. + let prefix_off = unsafe { arena.grow_uninit(4) }; + unsafe { arena.slice_mut(prefix_off, 4) }.copy_from_slice(&block_size.to_le_bytes()); + let body_off = unsafe { arena.grow_uninit(r.len()) }; + unsafe { arena.slice_mut(body_off, r.len()) }.copy_from_slice(r); + bodies.push((body_off as u64, block_size)); // ref points at the BODY, len = block_size + } + let arena_arc = Arc::new(PooledSegmentedBuf::unpooled(arena)); + let ref_chunk = coordinate_chunk_from_arena_refs(arena_arc, &bodies, n_ref, 1); + let ref_bytes: Vec> = + (0..ref_chunk.len()).map(|i| ref_chunk.record_bytes(i).to_vec()).collect(); + + assert_eq!( + ref_bytes, oracle_bytes, + "ref-sort output must be byte-identical to the copy-based sorter" + ); + } + + use proptest::prelude::*; + + proptest! { + // For any multiset of (refId, pos) records — including ties and unmapped — + // the ref-sort produces the same sorted body sequence as the copy sorter. + #[test] + #[allow(unsafe_code)] + fn prop_ref_sort_matches_copy_sorter( + // `-1` tids exercise the unmapped → `nref` sentinel (`u64::MAX`) + // ordering path, which `0..5` alone never reaches despite the header. + coords in prop::collection::vec((-1i32..5, 0i32..2000), 0..60), + ) { + let n_ref = 5u32; + // A unique read-name per record (by input index) makes every record + // byte-distinct, so any reordering of equal-key (tied) records would + // change the output and fail the assertion — witnessing stability. + let recs: Vec> = coords + .iter() + .enumerate() + .map(|(i, &(t, p))| coord_body(t, p, u8::try_from(i).unwrap())) + .collect(); + + let mut oracle = CoordinateChunkSorter::for_test(usize::MAX, n_ref); + for r in &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(); + + let mut arena = SegmentedBuf::with_capacity(0, 1 << 20); + arena.reserve_full_capacity(); + let mut bodies = Vec::new(); + for r in &recs { + let bs = u32::try_from(r.len()).unwrap(); + // SAFETY: each slot is fully written (via `slice_mut`) before any + // read — the 4-byte prefix and the body are copied in immediately + // after each `grow_uninit`, satisfying the write-before-read contract. + let po = unsafe { arena.grow_uninit(4) }; + unsafe { arena.slice_mut(po, 4) }.copy_from_slice(&bs.to_le_bytes()); + // SAFETY: as above — the body slot is fully written before any read. + let bo = unsafe { arena.grow_uninit(r.len()) }; + unsafe { arena.slice_mut(bo, r.len()) }.copy_from_slice(r); + bodies.push((bo as u64, bs)); + } + let arena_arc = Arc::new(PooledSegmentedBuf::unpooled(arena)); + let ref_chunk = coordinate_chunk_from_arena_refs(arena_arc, &bodies, n_ref, 1); + let ref_bytes: Vec> = + (0..ref_chunk.len()).map(|i| ref_chunk.record_bytes(i).to_vec()).collect(); + + prop_assert_eq!(ref_bytes, oracle_bytes); + } + } + + /// Exact queryname ties (identical name + flag class) must emerge in ingest + /// order — samtools' final tie-break — which the arena encodes as ascending + /// offset. `par_sort_unstable_by` on the key alone gives no such guarantee, + /// so feed equal-key refs in scrambled order and assert the sort restores + /// ascending-offset (== ingest) order. + #[test] + #[allow(unsafe_code)] + fn queryname_arena_sort_breaks_exact_ties_by_offset() { + use crate::keys::RawQuerynameKey; + + let n = 2_000usize; + let key = RawQuerynameKey::new(b"dup/name".to_vec(), 0); // one shared key → all ties + + // Each body encodes its ingest index; offsets ascend with the index, so a + // correct sort must decode 0,1,2,… in order. + let mut arena = SegmentedBuf::with_capacity(0, 1 << 20); + arena.reserve_full_capacity(); + let mut by_index: Vec<(u64, u32)> = Vec::with_capacity(n); + for i in 0..n { + let body = (i as u64).to_le_bytes(); + // SAFETY: each slot is fully written (copy_from_slice) before any read. + let off = unsafe { arena.grow_uninit(body.len()) }; + unsafe { arena.slice_mut(off, body.len()) }.copy_from_slice(&body); + by_index.push((off as u64, u32::try_from(body.len()).unwrap())); + } + + // Feed refs in a scrambled order (odds then evens) so an unstable sort with + // no offset tie-break would not reliably yield ascending offsets. + let mut refs: Vec<(RawQuerynameKey, u64, u32)> = Vec::with_capacity(n); + for i in (0..n).filter(|i| i % 2 == 1).chain((0..n).filter(|i| i % 2 == 0)) { + let (off, len) = by_index[i]; + refs.push((key.clone(), off, len)); + } + + let arena_arc = Arc::new(PooledSegmentedBuf::unpooled(arena)); + let chunk = queryname_chunk_from_arena_refs(arena_arc, refs); + assert_eq!(chunk.len(), n); + + let mut prev: Option = None; + for i in 0..chunk.len() { + let v = u64::from_le_bytes(chunk.record_bytes(i).try_into().unwrap()); + if let Some(p) = prev { + assert!(v > p, "exact-tie records not in ascending ingest order: {p} then {v}"); + } + prev = Some(v); + } + } +} diff --git a/crates/fgumi-sort/src/segmented_buf.rs b/crates/fgumi-sort/src/segmented_buf.rs index 5b82babdd..1f03aee67 100644 --- a/crates/fgumi-sort/src/segmented_buf.rs +++ b/crates/fgumi-sort/src/segmented_buf.rs @@ -9,15 +9,21 @@ //! Designed as a drop-in replacement for `Vec` in the sort buffer, where //! records are appended sequentially and later accessed by `(offset, len)`. -use std::sync::Arc; - /// Default segment size: 256 MiB. const DEFAULT_SEGMENT_SIZE: usize = 256 * 1024 * 1024; /// A growable byte buffer backed by fixed-size segments. /// -/// Each segment is an independent heap allocation. Appending data never moves -/// existing bytes — when the current segment is full, a new one is allocated. +/// Each segment is an independent heap allocation, so growing the buffer never +/// copies data already written to *earlier* segments — when the current segment +/// is full, a new one is allocated rather than the whole buffer being doubled. +/// Within the current segment the backing `Vec` may still reallocate and move +/// its own bytes unless it was pre-sized (see +/// [`reserve_full_capacity`](Self::reserve_full_capacity)). Pre-sizing fixes the +/// *inner* buffer only; it does NOT make it sound to run a growth while a +/// `slice_mut` borrow is live, because the growth may also push to the OUTER +/// segment vector. See the `# Safety` note on +/// [`grow_uninit`](Self::grow_uninit) for the supported shape. /// /// Records written via [`extend_from_slice`](Self::extend_from_slice) are /// guaranteed not to span segment boundaries: if the current segment lacks @@ -25,31 +31,23 @@ const DEFAULT_SEGMENT_SIZE: usize = 256 * 1024 * 1024; pub struct SegmentedBuf { /// The fixed capacity of each segment. segment_size: usize, - /// Segments that are complete and will never be appended to again. - /// - /// Held behind `Arc` so a reader can be handed one segment while the writer - /// keeps appending to [`current`](Self::current). Splitting the finished - /// segments from the live one is what makes that safe *by construction* - /// rather than by asserting that a `Vec>` will not reallocate: the - /// writer never names a sealed segment again, so no `&mut` can alias the - /// `&` a reader holds. - sealed: Vec>>, - /// The segment currently being appended to, with capacity `segment_size` - /// (the first may be smaller — see [`with_capacity`](Self::with_capacity)). - current: Vec, - /// Segment allocations reclaimed by [`reset_for_reuse`](Self::reset_for_reuse), - /// waiting to be used again. - /// - /// This is the whole point of `reset_for_reuse`: keeping the pages mapped - /// across a fill→spill→reset cycle instead of returning them to the - /// allocator, which returns them to the OS, which makes the next chunk fault - /// every page back in one at a time. - free: Vec>, + /// Backing storage: one `Vec` per segment, each with capacity `segment_size`. + segments: Vec>, /// Total bytes stored across all segments. total_len: usize, + /// Index of the segment currently being filled (the write cursor). Writes + /// land in `segments[cur]`; on overflow the cursor advances and **reuses** + /// `segments[cur+1]` if it already exists (retained by `reset_for_reuse`), + /// else a new segment is pushed. The overflow path pads `total_len` to the + /// segment boundary, so `total_len / segment_size == cur` holds at every + /// segment *transition* — which is what keeps `locate()` exact. It is not a + /// standing invariant: mid-segment `total_len` is `cur * segment_size + + /// (bytes written into segment cur)`, so the quotient equals `cur` only + /// because that remainder is always < `segment_size`. For an actively-growing buffer `cur` is the last + /// index, so behaviour is identical to the pre-cursor `segments.last()` path. + 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. /// @@ -63,13 +61,9 @@ impl SegmentedBuf { // re-allocations as the buffer grows. let first_cap = capacity.min(segment_size); let estimated_segments = (capacity / segment_size).max(1); - Self { - segment_size, - sealed: Vec::with_capacity(estimated_segments), - current: Vec::with_capacity(first_cap), - free: Vec::new(), - total_len: 0, - } + let mut segments = Vec::with_capacity(estimated_segments); + segments.push(Vec::with_capacity(first_cap)); + Self { segment_size, segments, total_len: 0, cur: 0 } } /// Create a new buffer with a default segment size of 256 MiB. @@ -78,6 +72,33 @@ impl SegmentedBuf { Self::with_capacity(0, DEFAULT_SEGMENT_SIZE) } + /// Ensure `needed` bytes fit contiguously in the current segment, starting a + /// fresh (gap-padded) segment if they don't. + /// + /// This is the single source of truth for the "pad `total_len` to the next + /// segment boundary, then [`advance_segment`](Self::advance_segment)" offset + /// arithmetic shared by [`extend_from_slice`](Self::extend_from_slice), + /// [`reserve_contiguous`](Self::reserve_contiguous), and + /// [`grow_uninit`](Self::grow_uninit). Keeping it in one place is a + /// correctness requirement, not just DRY: the three call sites must agree + /// byte-for-byte on where a record lands, and a divergence between them + /// would silently corrupt sort output identity. Anything that needs to + /// predict these offsets ahead of the write must drive this buffer rather + /// than re-derive the arithmetic. Callers read `self.total_len` as the write + /// offset *after* this returns. + #[inline] + fn make_room(&mut self, needed: usize) { + let seg = &self.segments[self.cur]; + if seg.len() + needed > self.segment_size { + // Pad total_len to the next segment boundary so locate() works. + let remainder = self.total_len % self.segment_size; + if remainder > 0 { + self.total_len += self.segment_size - remainder; + } + self.advance_segment(); + } + } + /// Append bytes to the buffer, returning the global offset of the write. /// /// If the current segment does not have enough remaining capacity for the @@ -100,16 +121,65 @@ impl SegmentedBuf { self.segment_size, ); - if self.current.len() + data.len() > self.segment_size { - self.advance_segment(); - } + self.make_room(data.len()); let offset = self.total_len; - self.current.extend_from_slice(data); + self.segments[self.cur].extend_from_slice(data); self.total_len += data.len(); offset } + /// Advance the write cursor to the next segment, **reusing** a retained + /// (already-cleared) segment if one exists, else allocating a fresh one. + /// Reuse is what makes `reset_for_reuse` + refill allocation-free. + #[inline] + fn advance_segment(&mut self) { + self.cur += 1; + if self.cur == self.segments.len() { + self.segments.push(Vec::with_capacity(self.segment_size)); + } + // Always on, release included: `cur` and `total_len` are kept in step by + // the gap padding in `make_room`, and a non-empty segment here breaks + // that silently — `locate` then resolves every later offset into the + // wrong segment and the sort emits wrong bytes with no error. One + // `is_empty()` per segment transition (one per `segment_size` bytes, so + // per 256 MiB on the sort path) is not a measurable cost next to that. + assert!( + self.segments[self.cur].is_empty(), + "advance_segment landed on a non-empty segment (reset_for_reuse must clear all)" + ); + // Uphold the pointer-stability invariant `grow_uninit`'s SAFETY note + // relies on: the segment the cursor lands in always holds `segment_size` + // capacity, so a later `grow_uninit` inside it cannot reallocate and move + // a slot `slice_mut` already handed out. `reserve_full_capacity` primes + // only `segments[cur]`, and this transition happens inside `make_room` + // where no caller can observe it and re-prime. A no-op on every path + // today — segments are pushed `with_capacity(segment_size)` and retained + // ones were built the same way — but nothing else enforces it. + let segment_size = self.segment_size; + let seg = &mut self.segments[self.cur]; + if seg.capacity() < segment_size { + seg.reserve_exact(segment_size); + } + } + + /// Ensure the current segment's backing `Vec` has capacity == `segment_size`. + /// + /// After this, any sequence of [`grow_uninit`](Self::grow_uninit) / + /// [`extend_from_slice`](Self::extend_from_slice) calls that stay within this + /// segment never reallocate it — so a `&mut [u8]` previously handed out by + /// [`slice_mut`](Self::slice_mut) for an earlier slot cannot be invalidated by + /// a later grow. This is what makes the parallel-inflate arena sound: the + /// serial admit path calls this once per fresh segment before handing any slot + /// to an inflate worker. No-op if the segment already holds `segment_size`. + pub fn reserve_full_capacity(&mut self) { + let seg = &mut self.segments[self.cur]; + let have = seg.capacity(); + if have < self.segment_size { + seg.reserve_exact(self.segment_size - seg.len()); + } + } + /// Ensure at least `additional` bytes fit in the current segment. /// /// If the remaining capacity in the current segment is less than @@ -130,9 +200,7 @@ impl SegmentedBuf { self.segment_size, ); - if self.current.len() + additional > self.segment_size { - self.advance_segment(); - } + self.make_room(additional); self.total_len } @@ -143,17 +211,189 @@ impl SegmentedBuf { /// /// # Panics /// - /// Panics (in debug builds) if the segments vec is empty. + /// Panics (in debug builds) if the write would overflow the current + /// segment — call [`reserve_contiguous`](Self::reserve_contiguous) first to + /// start a fresh one. In release the `debug_assert` is compiled out and the + /// segment simply grows past `segment_size`, which breaks the invariant that + /// a record never spans a segment boundary. #[inline] pub fn extend_in_place(&mut self, data: &[u8]) { debug_assert!( - self.current.len() + data.len() <= self.segment_size, + self.segments[self.cur].len() + data.len() <= self.segment_size, "extend_in_place exceeds segment capacity; use reserve_contiguous first" ); - self.current.extend_from_slice(data); + self.segments[self.cur].extend_from_slice(data); self.total_len += data.len(); } + /// Reserve a contiguous slot of `additional` **uninitialized-but-live** bytes, + /// returning the global offset of its first byte. The slot is guaranteed to lie + /// within a single segment (gap-padding is applied exactly as + /// [`reserve_contiguous`](Self::reserve_contiguous)), so it can later be written + /// through [`slice_mut`](Self::slice_mut) and read through [`slice`](Self::slice). + /// + /// Unlike `reserve_contiguous` + `extend_in_place`, this grows the segment's + /// length WITHOUT copying or zero-filling — it is the admit-path primitive for + /// the parallel-inflate arena, where an inflate worker writes every byte of the + /// slot exactly once. The returned offsets reproduce `reserve_contiguous`'s + /// accounting byte-for-byte. + /// + /// # Safety + /// + /// On return the `additional` bytes at the returned offset are LIVE but + /// UNINITIALIZED. The caller MUST fully initialize the entire slot (e.g. via + /// [`slice_mut`](Self::slice_mut)) before reading any of those bytes through + /// [`slice`](Self::slice) or [`slice_mut`](Self::slice_mut). Reading the slot before it is + /// fully written is undefined behavior. + /// + /// **Pointer stability, and why it is not sufficient for concurrency:** + /// `grow_uninit` calls `Vec::reserve`, which *can* reallocate and move a + /// segment's storage — invalidating any `&mut [u8]` previously handed out by + /// [`slice_mut`](Self::slice_mut) for an earlier slot in the same segment. + /// [`reserve_full_capacity`](Self::reserve_full_capacity), called once per + /// fresh segment, fixes that segment at `segment_size` so no subsequent + /// `grow_uninit` within it reallocates. + /// + /// That addresses the *inner* buffer only, and it does NOT make it sound to + /// run `grow_uninit` while a `slice_mut` borrow is live on another thread. + /// Two further hazards remain, and neither is fixed by pre-sizing: + /// + /// - `grow_uninit` may reach `advance_segment`, which pushes to + /// `self.segments`. That push reallocates and frees the OUTER + /// `Vec>` whenever it is at capacity, leaving a concurrent + /// `slice_mut` indexing freed memory — a use-after-free, not merely a + /// stale read. It is capacity-dependent, so it fires intermittently. + /// - `grow_uninit`'s `set_len` writes the inner `Vec`'s length field while a + /// concurrent `slice_mut` reads it for its bounds assert — a data race, + /// on every call rather than only at a capacity boundary. + /// + /// Neither is closed by pre-sizing, and neither can be closed at all through + /// this API: `grow_uninit` takes `&mut self` and `slice_mut` takes `&self`, + /// so overlapping them is an aliasing violation whatever the caller does. + /// That is why the overlap is forbidden outright by `slice_mut`'s third + /// precondition rather than made safe — Miri reports UB on exactly this + /// pattern. The supported shape is therefore to reserve ALL slots + /// for a segment first, and only then hand them to workers — see + /// `slice_mut_concurrent_disjoint_writes_are_sound`, which does that. In the + /// single-threaded case the borrow checker already forbids the overlap. + /// + /// # Panics + /// + /// Panics if `additional > segment_size`. + #[allow(unsafe_code)] + #[allow(clippy::uninit_vec)] // intentional: slot is live-but-uninitialized; contract requires caller to fully write it + pub unsafe fn grow_uninit(&mut self, additional: usize) -> usize { + assert!( + additional <= self.segment_size, + "grow of {} bytes exceeds segment size {}", + additional, + self.segment_size, + ); + + self.make_room(additional); + + let offset = self.total_len; + let seg = &mut self.segments[self.cur]; + let new_len = seg.len() + additional; + seg.reserve(additional); + // SAFETY: after `reserve(additional)`, `capacity() >= new_len`, so + // `set_len(new_len)` only extends `len` over already-allocated bytes. The + // grown region `old_len..new_len` is intentionally left uninitialized; the + // documented `# Safety` contract requires the caller to fully write it via + // `slice_mut` before any read. `u8` has no drop glue and no validity + // invariant, so growing `len` over uninitialized bytes is not itself UB — + // only a read-before-write would be, which the contract forbids. + // + // NB: this `reserve` MAY reallocate and move the segment's storage; that + // is safe here only because no `slice_mut` borrow is live across this call + // — either the borrow checker forbids it (`&mut self`), or the caller + // pre-sized the segment via `reserve_full_capacity` so this `reserve` is a + // no-op (see the `# Safety` "Pointer stability" note). + unsafe { + seg.set_len(new_len); + } + self.total_len += additional; + offset + } + + /// Obtain a mutable view of a previously-reserved slot by global `(offset, len)`. + /// + /// Takes `&self` (not `&mut self`) so that multiple **disjoint** slots can be + /// written concurrently from different threads, each holding its own + /// `&mut [u8]` into the shared buffer — the parallel-inflate use case. + /// + /// # Safety + /// + /// The caller MUST guarantee that: + /// 1. `(offset, len)` was produced by [`grow_uninit`](Self::grow_uninit) (or + /// `reserve_contiguous` + a matching in-place fill), so it lies within a + /// single segment's live region and is in bounds; and + /// 2. no other `slice`/`slice_mut` borrow of ANY overlapping range is live for + /// the duration of the returned reference — callers partition the buffer + /// into non-overlapping slots and write each exactly once; and + /// 3. no `&mut self` method runs for the duration of the returned reference. + /// That is every mutating method on this type — + /// [`grow_uninit`](Self::grow_uninit), + /// [`extend_from_slice`](Self::extend_from_slice), + /// [`reserve_contiguous`](Self::reserve_contiguous), + /// [`extend_in_place`](Self::extend_in_place), + /// [`reserve_full_capacity`](Self::reserve_full_capacity), + /// [`reset_for_reuse`](Self::reset_for_reuse), and [`clear`](Self::clear). + /// + /// Violating (1) or (2) is undefined behavior (an aliasing `&mut`, or an + /// out-of-bounds reference). Violating (3) is a data race and can be a + /// use-after-free: this method reads `self.segments[..]` and the segment's + /// `len`, both of which a concurrent `grow_uninit` may reallocate or write. + /// Pre-sizing the segment does not rescue it — see `grow_uninit`'s pointer- + /// stability note. The supported concurrent shape is to reserve every slot + /// for a segment *before* handing any of them out, which is what the + /// parallel-inflate path does. + /// + /// # Panics + /// + /// Panics if the range spans a segment boundary or exceeds the segment's + /// live length. + #[must_use] + #[allow(unsafe_code)] + #[allow(clippy::mut_from_ref)] // intentional: &self lets disjoint slots be written concurrently; caller's disjointness contract is the safety invariant + pub unsafe fn slice_mut(&self, offset: usize, len: usize) -> &mut [u8] { + // A zero-length slot never needs to be located: `offset` may legitimately + // sit on a segment boundary (e.g. after `grow_uninit(0)` when the current + // segment is exactly full), where `locate` would index past the last + // segment and panic. Return an empty slice instead. + if len == 0 { + // Narrow to the case this exists for: `grow_uninit(0)` at an exactly + // full segment, where `offset` legitimately sits one past the last + // byte and `locate` would index past the last segment. A zero-length + // read at a genuinely bogus offset must still fail loudly rather than + // silently return empty. + assert!( + offset <= self.total_len, + "slice_mut offset {offset} is past the end of the buffer ({})", + self.total_len, + ); + return &mut []; + } + let (seg_idx, seg_offset) = self.locate(offset); + let seg = &self.segments[seg_idx]; + assert!( + seg_offset + len <= seg.len(), + "slice_mut ({offset}..{}) spans segment boundary (seg {seg_idx}, seg_offset {seg_offset}, seg_len {})", + offset + len, + seg.len(), + ); + // SAFETY: `seg_offset + len <= seg.len()` (asserted) keeps the range inside + // the segment's live region, so the pointer + length are in bounds. The + // `&mut [u8]` synthesized from a shared `&self` is sound only under the + // caller's contract (point 2) that this range is disjoint from every other + // concurrently-borrowed range, so the produced `&mut` aliases no other + // `&`/`&mut`. The bytes are live (slot reserved via grow_uninit). + unsafe { + let base = seg.as_ptr().add(seg_offset).cast_mut(); + std::slice::from_raw_parts_mut(base, len) + } + } + /// Total bytes stored (including gap padding at segment boundaries). #[inline] #[must_use] @@ -180,9 +420,22 @@ impl SegmentedBuf { #[inline] #[must_use] pub fn slice(&self, offset: usize, len: usize) -> &[u8] { + // A zero-length range never needs to be located: `offset` may legitimately + // sit on a segment boundary (e.g. after a zero-length reservation when the + // current segment is exactly full), where `locate` would index past the + // last segment and panic. Return an empty slice instead. + if len == 0 { + // Same narrowing as `slice_mut`: permit the one-past-the-end offset a + // zero-length reservation produces, reject anything beyond it. + assert!( + offset <= self.total_len, + "slice offset {offset} is past the end of the buffer ({})", + self.total_len, + ); + return &[]; + } let (seg_idx, seg_offset) = self.locate(offset); - let seg: &[u8] = - if seg_idx == self.sealed.len() { &self.current } else { &self.sealed[seg_idx] }; + let seg = &self.segments[seg_idx]; assert!( seg_offset + len <= seg.len(), "slice ({offset}..{}) spans segment boundary (seg {seg_idx}, seg_offset {seg_offset}, seg_len {})", @@ -195,120 +448,38 @@ impl SegmentedBuf { /// Total allocated capacity in bytes across all segments. #[must_use] pub fn allocated_capacity(&self) -> usize { - self.sealed.iter().map(|s| s.capacity()).sum::() - + self.free.iter().map(Vec::capacity).sum::() - + self.current.capacity() - } - - /// Give back every retained segment allocation. - /// - /// Call once ingest is done. [`reset_for_reuse`](Self::reset_for_reuse) - /// keeps segments mapped so the *next* chunk does not re-fault them, but - /// after the last chunk there is no next fill and holding them just carries - /// the arena's peak footprint through the merge — measured at +1.68 GB of - /// peak RSS on the production cell, for no benefit. - pub fn release_retained(&mut self) { - self.free = Vec::new(); + self.segments.iter().map(Vec::capacity).sum() } - /// Segment allocations held for reuse but not currently part of the buffer. - #[must_use] - pub fn retained_segments(&self) -> usize { - self.free.len() - } - - /// Empty the buffer **keeping every segment allocation** for the next fill. - /// - /// The counterpart to [`clear`](Self::clear), and the one a sort should use - /// between chunks. `clear` drops the sealed segments, which hands their - /// pages to the allocator and ultimately back to the OS; the next chunk then - /// takes a minor fault on every page it touches again. Measured on the - /// production cell (780M records, 44 chunks, ~8 GB of arena per chunk): - /// **25.9M minor faults** and ~25s of wall clock with `clear`, against - /// **18k faults** when the pages are simply kept. - /// - /// A segment a reader still holds cannot be reclaimed — its `Arc` is - /// dropped instead and the next fill allocates a replacement. That is a - /// correctness requirement, not an optimization: reusing storage a - /// key-extraction batch was still reading would overwrite its records - /// mid-flight. - pub fn reset_for_reuse(&mut self) { - for segment in std::mem::take(&mut self.sealed) { - if let Ok(mut owned) = Arc::try_unwrap(segment) { - owned.clear(); - self.free.push(owned); - } - } - self.current.clear(); - self.total_len = 0; - } - - /// Number of allocated segments, sealed plus the one being appended to. + /// Number of allocated segments. #[must_use] pub fn num_segments(&self) -> usize { - self.sealed.len() + 1 + self.segments.len() } - /// Number of sealed segments — those that will never be appended to again. - /// - /// Segment indices `0..sealed_len()` are stable and shareable; index - /// `sealed_len()` is the live segment and is not. - #[must_use] - pub fn sealed_len(&self) -> usize { - self.sealed.len() - } - - /// A shared handle to sealed segment `idx`, or `None` if `idx` is the live - /// segment (or past the end). - /// - /// The returned handle stays valid and byte-stable for as long as it is - /// held, including across [`clear`](Self::clear) — which is what lets a - /// worker read record bytes out of it while the writer fills later segments. - #[must_use] - pub fn sealed_segment(&self, idx: usize) -> Option>> { - self.sealed.get(idx).map(Arc::clone) - } - - /// Clear all data, retaining the live segment's allocation. - /// - /// Sealed segments are released here. A handle handed out by - /// [`sealed_segment`](Self::sealed_segment) keeps its own segment alive - /// past this call, so the storage is freed when the last reader drops it - /// rather than while it is still being read. + /// Clear all data, retaining only the first segment's allocation (drops + /// segments `1..`). Used by callers that want memory released between + /// unrelated sorts (legacy `.sort()`, the streaming path). For + /// allocation-free arena reuse across spill chunks use + /// [`reset_for_reuse`](Self::reset_for_reuse) instead. pub fn clear(&mut self) { - self.sealed.clear(); - self.free.clear(); - self.current.clear(); + self.segments.truncate(1); + self.segments[0].clear(); self.total_len = 0; + self.cur = 0; } - /// Seal the live segment so its bytes can be shared, if it holds anything. - /// - /// A no-op on an empty live segment, so calling it twice does not push an - /// empty segment or open a spurious gap. The writer's next append starts a - /// fresh segment, which wastes the tail of the sealed one — acceptable only - /// because the caller (the ingest thread's chunk barrier) is about to - /// [`clear`](Self::clear) the whole buffer anyway. - pub fn seal_current(&mut self) { - if !self.current.is_empty() { - self.advance_segment(); - } - } - - /// Seal the live segment and start a fresh one. - /// - /// Pads `total_len` to the next segment boundary first, so - /// [`locate`](Self::locate)'s division arithmetic stays exact. - fn advance_segment(&mut self) { - let remainder = self.total_len % self.segment_size; - if remainder > 0 { - self.total_len += self.segment_size - remainder; + /// Reset for reuse **retaining every segment allocation** (capacity kept). + /// The next fill reuses the retained segments via the write cursor, so a + /// pooled arena cycles through fill→spill→reset without reallocating its + /// segments. Unlike [`clear`](Self::clear), `num_segments()` and + /// `allocated_capacity()` are unchanged (they hold the peak across reuses). + pub fn reset_for_reuse(&mut self) { + for seg in &mut self.segments { + seg.clear(); } - // Reuse a retained allocation when there is one; only ask the OS for - // fresh pages when there is not. - let next = self.free.pop().unwrap_or_else(|| Vec::with_capacity(self.segment_size)); - let filled = std::mem::replace(&mut self.current, next); - self.sealed.push(Arc::new(filled)); + self.total_len = 0; + self.cur = 0; } /// Translate a global byte offset to `(segment_index, offset_within_segment)`. @@ -324,9 +495,9 @@ impl SegmentedBuf { let seg_idx = offset / self.segment_size; let seg_offset = offset % self.segment_size; debug_assert!( - seg_idx <= self.sealed.len(), + seg_idx < self.segments.len(), "locate({offset}): seg_idx {seg_idx} out of bounds (len {})", - self.sealed.len() + 1 + self.segments.len() ); (seg_idx, seg_offset) } @@ -346,6 +517,8 @@ impl Default for SegmentedBuf { #[cfg(test)] mod tests { + use rstest::rstest; + use super::*; #[test] @@ -356,6 +529,41 @@ mod tests { assert_eq!(buf.num_segments(), 1); // one pre-allocated segment } + #[test] + fn reset_for_reuse_retains_segments_and_reuses_storage() { + // segment_size 16; ten-byte writes each start a fresh segment (10+10 > 16). + let mut buf = SegmentedBuf::with_capacity(0, 16); + let mut offs = Vec::new(); + for i in 0..4u8 { + offs.push(buf.extend_from_slice(&[i; 10])); + } + assert_eq!(buf.num_segments(), 4, "four 10B writes into 16B segments → 4 segments"); + for (i, &o) in offs.iter().enumerate() { + assert_eq!(buf.slice(o, 10), &[u8::try_from(i).unwrap(); 10][..]); + } + let cap_before = buf.allocated_capacity(); + + // Reset for reuse: data cleared, but the 4 segment allocations are RETAINED. + buf.reset_for_reuse(); + assert!(buf.is_empty()); + assert_eq!(buf.len(), 0); + assert_eq!(buf.num_segments(), 4, "reset retains the segment allocations"); + assert_eq!(buf.allocated_capacity(), cap_before, "no realloc on reset"); + + // Refill: must REUSE the retained segments (count stays 4, capacity unchanged), + // and locate()/slice() must remain correct over the reused segments. + let mut offs2 = Vec::new(); + for i in 0..4u8 { + offs2.push(buf.extend_from_slice(&[100 + i; 10])); + } + assert_eq!(buf.num_segments(), 4, "refill reuses retained segments (no growth)"); + assert_eq!(buf.allocated_capacity(), cap_before, "refill reused storage (no new alloc)"); + assert_eq!(offs2, offs, "offsets reproduce identically after reuse"); + for (i, &o) in offs2.iter().enumerate() { + assert_eq!(buf.slice(o, 10), &[100 + u8::try_from(i).unwrap(); 10][..]); + } + } + #[test] fn test_extend_and_len() { let mut buf = SegmentedBuf::with_capacity(0, 1024); @@ -420,98 +628,6 @@ mod tests { assert_eq!(buf.slice(o2, 5), b"abcde"); } - #[test] - fn test_reset_for_reuse_keeps_the_segment_allocations() { - // The whole point: a sort cycles fill -> spill -> reset once per chunk, - // and dropping the segments each time hands their pages back to the - // allocator, which returns them to the OS. The next chunk then faults - // every page back in. Measured on the production cell: 25.9M minor - // faults and ~25s of wall clock, against 18k faults when the pages are - // simply kept. - let mut buf = SegmentedBuf::with_capacity(0, 10); - for _ in 0..4 { - buf.extend_from_slice(b"0123456789"); - } - let segments_before = buf.num_segments(); - let capacity_before = buf.allocated_capacity(); - assert!(segments_before >= 4, "fixture must span several segments"); - - buf.reset_for_reuse(); - - assert_eq!(buf.len(), 0, "the buffer is logically empty"); - assert_eq!( - buf.allocated_capacity(), - capacity_before, - "reset must not give the allocations back", - ); - assert_eq!( - buf.retained_segments(), - segments_before - 1, - "every sealed segment is retained for the next fill", - ); - - // Refilling consumes the retained segments rather than allocating. - for _ in 0..4 { - buf.extend_from_slice(b"0123456789"); - } - assert_eq!(buf.retained_segments(), 0, "the refill reused them"); - assert_eq!(buf.allocated_capacity(), capacity_before, "and allocated nothing new"); - assert_eq!(buf.slice(0, 10), b"0123456789", "and the data is correct"); - } - - #[test] - fn test_reset_for_reuse_lets_go_of_a_segment_a_reader_still_holds() { - // A reclaimed segment must be exclusively owned. A key-extraction batch - // that outlived its chunk would otherwise have its bytes overwritten by - // the next chunk's records while it read them. - let mut buf = SegmentedBuf::with_capacity(0, 10); - buf.extend_from_slice(b"0123456789"); - buf.extend_from_slice(b"abcde"); - let held = buf.sealed_segment(0).expect("segment 0 is sealed"); - - buf.reset_for_reuse(); - - assert_eq!(buf.retained_segments(), 0, "the held segment cannot be reused"); - assert_eq!(&held[..], b"0123456789", "and the reader still sees its bytes"); - } - - #[test] - fn test_a_segment_seals_only_once_a_later_one_starts() { - let mut buf = SegmentedBuf::with_capacity(0, 10); - buf.extend_from_slice(b"0123456789"); - - // The segment being appended to is never sealed: a worker handed it - // could observe a partial write, and the ingest thread still owns it. - assert_eq!(buf.sealed_len(), 0); - assert!(buf.sealed_segment(0).is_none()); - - // Starting segment 1 seals segment 0. - buf.extend_from_slice(b"abcde"); - assert_eq!(buf.sealed_len(), 1); - assert_eq!(&buf.sealed_segment(0).expect("segment 0 is sealed")[..], b"0123456789"); - assert!(buf.sealed_segment(1).is_none(), "the live segment is not sealed"); - } - - #[test] - fn test_a_sealed_segment_is_readable_while_later_segments_are_written() { - // The invariant the parallel key-extraction path rests on: a handle to a - // sealed segment stays valid and byte-stable no matter how much is - // appended afterwards. Without it, a worker reading record bytes races - // the ingest thread's growth. - let mut buf = SegmentedBuf::with_capacity(0, 10); - buf.extend_from_slice(b"0123456789"); - buf.extend_from_slice(b"abcde"); - - let held = buf.sealed_segment(0).expect("segment 0 is sealed"); - - // Fill several more segments; segment 0's bytes must not move or change. - for _ in 0..4 { - buf.extend_from_slice(b"0123456789"); - } - assert_eq!(&held[..], b"0123456789"); - assert!(buf.num_segments() >= 5); - } - #[test] fn test_offset_accounting_with_gaps() { // When a write spills to a new segment, total_len must include the @@ -554,6 +670,109 @@ mod tests { assert_eq!(buf.num_segments(), 1); } + /// `clear` truncates to one segment, so it MUST also rewind the write cursor + /// — otherwise `cur` still points at a segment that no longer exists and the + /// next write indexes out of bounds. A single-segment buffer cannot catch + /// this: `cur` is already 0 there, so dropping the reset passes. + #[test] + fn clear_rewinds_the_write_cursor_after_multiple_segments() { + let mut buf = SegmentedBuf::with_capacity(0, 16); + for _ in 0..4 { + buf.extend_from_slice(&[7u8; 10]); // 10 > 16/2, so each forces a new segment + } + assert!(buf.num_segments() > 1, "precondition: the buffer spans several segments"); + + buf.clear(); + assert_eq!(buf.num_segments(), 1, "clear drops the extra segments"); + + // The write that would panic if `cur` still pointed past segment 0. + let off = buf.extend_from_slice(b"after clear"); + assert_eq!(off, 0, "the first post-clear write lands at offset 0"); + assert_eq!(buf.slice(off, 11), b"after clear"); + } + + #[allow(unsafe_code)] + #[test] + fn grow_uninit_then_slice_mut_round_trips() { + // segment_size 100; first slot fits, second forces a gap to a new segment. + let mut buf = SegmentedBuf::with_capacity(0, 100); + + // SAFETY: each slot is fully written via slice_mut before any slice() read. + let o0 = unsafe { buf.grow_uninit(60) }; + assert_eq!(o0, 0); + assert_eq!(buf.len(), 60); + assert_eq!(buf.num_segments(), 1); + unsafe { buf.slice_mut(o0, 60) }.fill(0xAB); + + // 60 used, a 60-byte slot won't fit in the remaining 40 → gap to seg 1. + let o1 = unsafe { buf.grow_uninit(60) }; + assert_eq!(o1, 100, "gap-padded to the segment boundary (matches reserve_contiguous)"); + assert_eq!(buf.num_segments(), 2); + assert_eq!(buf.len(), 160); + unsafe { buf.slice_mut(o1, 60) }.fill(0xCD); + + // Both slots read back exactly through the safe slice() path. + assert_eq!(buf.slice(o0, 60), &[0xAB; 60][..]); + assert_eq!(buf.slice(o1, 60), &[0xCD; 60][..]); + } + + #[allow(unsafe_code)] + #[test] + fn grow_uninit_offsets_match_reserve_contiguous() { + // grow_uninit must reproduce reserve_contiguous's offset/gap accounting so + // the ISIZE prefix-sum planner and the existing reserve path agree. + let sizes = [3usize, 3, 5, 90, 7]; + let mut a = SegmentedBuf::with_capacity(0, 100); + let mut b = SegmentedBuf::with_capacity(0, 100); + for &n in &sizes { + let ra = a.reserve_contiguous(n); + a.extend_in_place(&vec![0u8; n]); + // SAFETY: slot fully written immediately below. + let rb = unsafe { b.grow_uninit(n) }; + unsafe { b.slice_mut(rb, n) }.fill(0); + assert_eq!(ra, rb, "grow_uninit offset diverged from reserve_contiguous for n={n}"); + } + assert_eq!(a.len(), b.len()); + assert_eq!(a.num_segments(), b.num_segments()); + } + + #[allow(unsafe_code)] + #[test] + #[should_panic(expected = "exceeds segment size")] + fn grow_uninit_oversize_panics() { + let mut buf = SegmentedBuf::with_capacity(0, 10); + // SAFETY: call panics before any slot is reserved. + let _ = unsafe { buf.grow_uninit(11) }; + } + + /// Regression: a zero-length reservation whose offset lands exactly on a + /// segment boundary (the current segment is full) must not make `slice` / + /// `slice_mut` index past the last segment and panic. Both accessors return + /// an empty slice for a zero-length range. + #[allow(unsafe_code)] + #[test] + fn zero_length_slice_at_full_segment_boundary_does_not_panic() { + let mut buf = SegmentedBuf::with_capacity(0, 10); + + // Fill the (only) segment exactly full. + buf.extend_from_slice(b"0123456789"); + assert_eq!(buf.num_segments(), 1); + assert_eq!(buf.len(), 10); + + // A zero-length grow returns an offset on the boundary (== total_len), + // which `locate` would map to a not-yet-allocated segment. + // SAFETY: the returned slot has length 0, so no byte is ever read or + // written through it. + let off = unsafe { buf.grow_uninit(0) }; + assert_eq!(off, 10, "zero-length grow lands on the segment boundary"); + assert_eq!(buf.num_segments(), 1, "zero-length grow allocates nothing"); + + // Neither accessor may panic; both yield an empty slice. + assert!(buf.slice(off, 0).is_empty()); + // SAFETY: zero-length slot — no bytes are aliased, read, or written. + assert!(unsafe { buf.slice_mut(off, 0) }.is_empty()); + } + #[test] fn test_many_segments() { let mut buf = SegmentedBuf::with_capacity(0, 100); @@ -650,4 +869,206 @@ mod tests { // Can read the entire contiguous range assert_eq!(buf.slice(0, 12), b"aaaabbbbcccc"); } + + #[allow(unsafe_code)] + #[test] + fn slice_mut_concurrent_disjoint_writes_are_sound() { + use std::thread; + + // Reserve N disjoint slots on the serial admit path, then write them all + // concurrently — each thread owns exactly one slot. This is the + // parallel-inflate access pattern. + // + // What this does NOT prove: `miri.yml` runs only `-p fgumi-raw-bam sort` + // and `-p fgumi-pipeline-core erased`, so nothing checks this crate under + // miri. The test exercises the pattern and would catch a slot-arithmetic + // bug via the content assertions below, but it is not evidence of + // race-freedom — a benign-looking interleaving that violates `slice_mut`'s + // disjointness contract would still pass here. + let n_slots = 8usize; + let slot_len = 40usize; // 40-byte slots in a 100-byte segment → 2 slots/segment + let mut buf = SegmentedBuf::with_capacity(0, 100); + let mut slots = Vec::with_capacity(n_slots); + for _ in 0..n_slots { + // SAFETY: every reserved slot is fully written exactly once below before + // any read; offsets are distinct so the ranges are non-overlapping. + slots.push(unsafe { buf.grow_uninit(slot_len) }); + } + + let buf_ref = &buf; + thread::scope(|scope| { + for (i, &offset) in slots.iter().enumerate() { + let byte = u8::try_from(i).unwrap(); + scope.spawn(move || { + // SAFETY: each thread writes a distinct slot (distinct offset, + // fixed len); the ISIZE-prefix-sum analogue here is the distinct + // grow_uninit offsets, so the ranges are pairwise disjoint and no + // two &mut alias. + let dst = unsafe { buf_ref.slice_mut(offset, slot_len) }; + dst.fill(byte); + }); + } + }); + + // After the scope joins (write→read happens-before), every slot reads back + // its writer's byte through the safe slice() path. + for (i, &offset) in slots.iter().enumerate() { + let byte = u8::try_from(i).unwrap(); + assert_eq!(buf.slice(offset, slot_len), &vec![byte; slot_len][..]); + } + } + + /// `is_err()` alone is satisfied by ANY panic, including one from unrelated + /// bounds arithmetic, so it cannot tell you the guard under test fired. Both + /// cases therefore assert on the message. The second case is the one the name + /// promises: a range that starts in one segment and runs into the next — the + /// original test only over-read within a single segment, so the true + /// boundary-spanning path was never exercised. + #[rstest] + // One 40-byte slot in a 100-byte segment; asking for 80 runs past the live + // region of the only segment. + #[case::past_the_live_region(&[40], 80)] + // 60 + 60 > 100, so the second slot lands in segment 1 and segment 0 stays + // 60 bytes live. Asking for 61 from offset 0 crosses out of segment 0 — + // which is what the name promises and what the old single-slot test could + // not reach. + #[case::across_a_segment_boundary(&[60, 60], 61)] + fn slice_mut_outside_one_segments_live_region_panics( + #[case] slots: &[usize], + #[case] ask: usize, + ) { + let msg = slice_mut_panic_message(slots, ask); + assert!( + msg.contains("spans segment boundary"), + "must fail slice_mut's own bounds guard, not something incidental; got: {msg}", + ); + } + + /// Reserve two slots, then ask for `ask` bytes from the first and return the + /// panic message. Split out because `#[allow(unsafe_code)]` does not reach + /// the per-case functions `#[rstest]` generates. + #[allow(unsafe_code)] + fn slice_mut_panic_message(slots: &[usize], ask: usize) -> String { + let mut buf = SegmentedBuf::with_capacity(0, 100); + let mut first = 0; + for (i, &len) in slots.iter().enumerate() { + // SAFETY: every reserved slot is fully written immediately below, + // before anything reads it. + let off = unsafe { buf.grow_uninit(len) }; + unsafe { buf.slice_mut(off, len) }.fill(u8::try_from(i).expect("few slots")); + if i == 0 { + first = off; + } + } + let a = first; + + let err = std::panic::catch_unwind(|| { + // SAFETY: the call asserts and panics before producing a reference. + let _ = unsafe { buf.slice_mut(a, ask) }; + }) + .expect_err("slice_mut outside the segment's live region must panic"); + err.downcast_ref::() + .map_or_else(|| (*err.downcast_ref::<&str>().unwrap_or(&"")).to_string(), Clone::clone) + } + + #[test] + fn reserve_full_capacity_makes_segment_zero_full() { + // A pool-fresh arena has a capacity-0 segment 0; reserve_full_capacity must + // bring it to segment_size so the first grow_uninit cannot reallocate it. + let seg = 4096usize; + let mut buf = SegmentedBuf::with_capacity(0, seg); + buf.reserve_full_capacity(); + assert!(buf.num_segments() == 1); + assert!(buf.allocated_capacity() >= seg, "segment 0 must hold >= segment_size capacity"); + } + + /// `reserve_full_capacity` reserves `segment_size - seg.len()`, and the + /// subtraction matters: `Vec::reserve_exact` takes an amount ADDITIONAL to + /// the current length, so reserving the full `segment_size` on a partly + /// filled segment over-allocates by whatever is already written — up to 2x + /// on a nearly-full segment. Testing only the empty case cannot see that. + #[test] + fn reserve_full_capacity_does_not_over_allocate_a_partly_filled_segment() { + let seg = 4096usize; + let mut buf = SegmentedBuf::with_capacity(0, seg); + buf.extend_from_slice(&[0u8; 3000]); + buf.reserve_full_capacity(); + assert!( + buf.allocated_capacity() >= seg, + "must still reach segment_size so later grows cannot realloc", + ); + assert!( + buf.allocated_capacity() < seg + 3000, + "must reserve segment_size - len, not segment_size on top of len; got {}", + buf.allocated_capacity(), + ); + } + + #[allow(unsafe_code)] + #[test] + fn grow_uninit_is_realloc_free_after_reserve_full_capacity() { + // After reserve_full_capacity, growing many blocks that sum to <= segment_size + // must (a) stay in ONE segment (no gap/advance) and (b) never move the segment's + // backing buffer — so a slice_mut handed out for an early block stays valid. + let seg = 4096usize; + let mut buf = SegmentedBuf::with_capacity(0, seg); + buf.reserve_full_capacity(); + + // SAFETY: each slot is fully written before any read; offsets are distinct. + let off0 = unsafe { buf.grow_uninit(100) }; + let ptr0 = unsafe { buf.slice_mut(off0, 100) }.as_ptr() as usize; + unsafe { buf.slice_mut(off0, 100) }.fill(0xA1); + + // Grow several more blocks; total stays under segment_size → one segment. + let mut prev_end = off0 + 100; + for i in 0..20u8 { + let n = 100usize; + let off = unsafe { buf.grow_uninit(n) }; + assert_eq!(off, prev_end, "no gap: offsets are contiguous within one segment"); + unsafe { buf.slice_mut(off, n) }.fill(0xB0 + i); + prev_end = off + n; + } + assert_eq!(buf.num_segments(), 1, "all growth stayed in the pre-sized segment 0"); + + // The early block's backing pointer is unchanged → no realloc moved it. + let ptr0_after = buf.slice(off0, 100).as_ptr() as usize; + assert_eq!(ptr0, ptr0_after, "segment 0 buffer must not have been reallocated"); + assert_eq!(buf.slice(off0, 100), &[0xA1; 100][..]); + } + + #[allow(unsafe_code)] + #[test] + fn grow_uninit_reuses_arena_after_reset_without_stale_reads() { + // Increment-5 pattern: a pooled arena is filled, reset_for_reuse'd (len → 0, + // segment capacity RETAINED), then refilled. The post-reset grow_uninit takes + // the seg.reserve path on an already-grown segment; confirm offsets reproduce + // and reads return the fresh bytes, never the retained bytes from the first fill. + let mut buf = SegmentedBuf::with_capacity(0, 100); + + // First fill: two 40-byte slots (both in segment 0; 40 + 40 <= 100). + // SAFETY: each slot is fully written via slice_mut before any read. + let a0 = unsafe { buf.grow_uninit(40) }; + unsafe { buf.slice_mut(a0, 40) }.fill(0x11); + let a1 = unsafe { buf.grow_uninit(40) }; + unsafe { buf.slice_mut(a1, 40) }.fill(0x22); + assert_eq!(buf.slice(a0, 40), &[0x11; 40][..]); + assert_eq!(buf.slice(a1, 40), &[0x22; 40][..]); + + // Reset for reuse: data cleared, segment allocations retained. + buf.reset_for_reuse(); + assert!(buf.is_empty()); + + // Refill: offsets must reproduce the first fill, and reads must return the + // freshly-written bytes — never the retained 0x11/0x22 still physically present + // in the reused segment's capacity. + // SAFETY: each slot is fully written before any read. + let b0 = unsafe { buf.grow_uninit(40) }; + unsafe { buf.slice_mut(b0, 40) }.fill(0x33); + let b1 = unsafe { buf.grow_uninit(40) }; + unsafe { buf.slice_mut(b1, 40) }.fill(0x44); + assert_eq!(b0, a0, "post-reset grow_uninit reproduces the first-fill offset"); + assert_eq!(b1, a1, "post-reset grow_uninit reproduces the first-fill offset"); + assert_eq!(buf.slice(b0, 40), &[0x33; 40][..]); + assert_eq!(buf.slice(b1, 40), &[0x44; 40][..]); + } } diff --git a/crates/fgumi-sort/src/spill_block.rs b/crates/fgumi-sort/src/spill_block.rs new file mode 100644 index 000000000..8bcba442c --- /dev/null +++ b/crates/fgumi-sort/src/spill_block.rs @@ -0,0 +1,436 @@ +//! Block-granular spill compression kernel shared by the Phase-1 spill steps +//! (`SpillGather` → `SpillCompress` → `SpillWrite`). +//! +//! The legacy [`SyncSpillWriter`](crate::sync_spill_writer) streams a whole +//! sorted chunk through one stateful compressor on a single worker. To compress +//! a spill across the framework pool instead, the chunk is cut into raw blocks by +//! the (Serial) serialize step, each block is compressed **independently** by a +//! (Parallel) compress worker, and the writer concatenates the results per file. +//! This module holds the three pieces that keep that byte-compatible with the +//! existing spill readers: +//! +//! - [`frame_keyed_record_into`] — append one `[key?][u32 LE len][record]` record +//! to a raw block buffer (the identical on-disk record layout +//! `SyncSpillWriter::write_record` produces). +//! - [`SpillBlockCompressor`] — a per-worker, codec-fixed compressor whose +//! [`compress_block`](SpillBlockCompressor::compress_block) turns one raw block +//! into a self-contained, independently-decodable unit: framed BGZF block(s) +//! for bgzf, or a `[u32 LE len][zstd frame]` for zstd. +//! - [`spill_magic`] / [`spill_trailer`] — the per-codec file prologue/epilogue +//! the writer brackets the compressed blocks with. +//! +//! Because a BGZF/zstd spill stream is just independent blocks concatenated and +//! the readers stream across boundaries, **any** blocking of the same raw +//! `[key?][len][record]…` byte stream reads back to the identical records (see +//! the round-trip + re-blocking-independence tests below and the format notes in +//! [`sync_spill_writer`](crate::sync_spill_writer)). Block boundaries are +//! therefore the serialize step's free choice, and compression is a pure function +//! of each block — so the result is byte-for-byte independent of how many workers +//! compress concurrently. The one bound on that freedom is +//! [`MAX_SPILL_BLOCK_LEN`]. + +use std::io; + +use fgumi_bgzf::BGZF_EOF; +use fgumi_bgzf::writer::InlineBgzfCompressor; +use zstd::bulk::Compressor as ZstdCompressor; + +/// Largest raw block [`SpillBlockCompressor::compress_block`] will accept, and +/// the ceiling the reader sizes its decompression scratch to. +/// +/// This is one contract with two enforcement points, which is why it lives here +/// (the writer side) and is imported by `spill_block_reader`, rather than being +/// declared next to the buffer it bounds. The reader has to cap *something*: the +/// decompressed size it reads out of a zstd frame header is attacker- (or +/// corruption-) controlled, so an uncapped allocation would take the process +/// out. Once the reader caps, the writer must reject the same size, or a block +/// above the cap writes successfully and is unreadable — a spill file this +/// process produced and cannot consume. +/// +/// 64 MiB is far above any single BAM record in practice (long reads are ~1 MiB +/// at the extreme), so it bounds corruption without bounding real data. +pub(crate) const MAX_SPILL_BLOCK_LEN: usize = 64 * 1024 * 1024; + +use crate::codec::{SpillCodec, ZSPILL_MAGIC}; +use crate::keys::RawSortKey; + +/// Append one keyed record to `block` in the spill record framing: +/// `[key bytes if !K::EMBEDDED_IN_RECORD][u32 LE record-len][record body]`. +/// +/// This is the exact layout `SyncSpillWriter::write_record` writes; the +/// serialize step accumulates records into a raw block with this helper, and the +/// merge / Phase-2 readers parse it back unchanged. +/// +/// # Errors +/// +/// Returns an error if the record is longer than `u32::MAX` (cannot fit the +/// length prefix) or key serialization fails. +pub fn frame_keyed_record_into( + block: &mut Vec, + key: &K, + record: &[u8], +) -> io::Result<()> { + // Validate the length BEFORE writing the key, so an oversized record fails + // loud without leaving partial (key) bytes in the block. + let record_len = u32::try_from(record.len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("BAM record too large ({} bytes) for a u32 length prefix", record.len()), + ) + })?; + if !K::EMBEDDED_IN_RECORD { + // `Vec: Write`, so the key serializes straight into the block buffer. + key.write_to(block)?; + } + block.extend_from_slice(&record_len.to_le_bytes()); + block.extend_from_slice(record); + Ok(()) +} + +/// Per-worker, codec-fixed block compressor. Constructed once per compress +/// worker (like `BgzfCompress`'s `InlineBgzfCompressor`) and reused across the +/// blocks that worker handles. +pub enum SpillBlockCompressor { + /// bgzf: each block becomes one or more framed BGZF blocks (header + deflate + /// + footer). No EOF marker — that is the file trailer (see [`spill_trailer`]). + Bgzf(InlineBgzfCompressor), + /// zstd: each block becomes one `[u32 LE frame-len][zstd frame]` unit. + Zstd(ZstdCompressor<'static>), +} + +impl SpillBlockCompressor { + /// Build a compressor for `codec` at `compression` (bgzf level — `0` writes + /// framed *stored* blocks; or zstd level, which must be ≥ 1). + /// + /// # Errors + /// + /// Returns an error if the zstd compressor cannot be initialized. + pub fn new(codec: SpillCodec, compression: u32) -> io::Result { + match codec { + SpillCodec::Bgzf => Ok(Self::Bgzf(InlineBgzfCompressor::new(compression))), + SpillCodec::Zstd => { + #[allow(clippy::cast_possible_wrap)] + let compressor = ZstdCompressor::new(compression as i32).map_err(|e| { + io::Error::other(format!("zstd compressor init (level {compression}): {e}")) + })?; + Ok(Self::Zstd(compressor)) + } + } + } + + /// Compress one raw block into a self-contained, independently-decodable + /// unit. Concatenating these units (bracketed by [`spill_magic`] / + /// [`spill_trailer`]) reproduces the same decompressed stream as the + /// streaming `SyncSpillWriter`, regardless of where the block boundaries fall. + /// + /// An empty block yields an empty `Vec` (no empty BGZF block / zstd frame is + /// emitted), matching the streaming writer's "no empty frames" invariant. + /// + /// # Errors + /// + /// Returns an error if compression fails, if a zstd frame exceeds `u32::MAX`, + /// or if `raw` exceeds `MAX_SPILL_BLOCK_LEN` — see that constant for why + /// the write side is where that has to fail. + pub fn compress_block(&mut self, raw: &[u8]) -> io::Result> { + if raw.is_empty() { + return Ok(Vec::new()); + } + // Fail here rather than let the reader fail later. The block boundary is + // the serialize step's free choice (see the module doc), so nothing in + // the type system stops it picking one the reader cannot decompress — + // and a spill file that writes successfully and cannot be read back is + // the worst shape this failure can take. Same reason the length check in + // `frame_keyed_record_into` runs before any bytes are written. + if raw.len() > MAX_SPILL_BLOCK_LEN { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "spill block of {} bytes exceeds the {MAX_SPILL_BLOCK_LEN}-byte limit the \ + reader can decompress", + raw.len(), + ), + )); + } + match self { + Self::Bgzf(compressor) => { + compressor.write_all(raw)?; + // Flush forces a block boundary at the end of this raw block, so + // the unit is independently decodable. + compressor.flush()?; + let mut out = Vec::new(); + for block in compressor.take_blocks() { + out.extend_from_slice(&block.data); + // The bytes are copied into `out`, so the block buffer is + // free; hand it back so the next block reuses the allocation + // instead of allocating a fresh output `Vec` per 64 KiB block. + compressor.recycle_buffer(block.data); + } + Ok(out) + } + Self::Zstd(compressor) => { + let frame = compressor + .compress(raw) + .map_err(|e| io::Error::other(format!("zstd compress: {e}")))?; + let frame_len = u32::try_from(frame.len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidData, + "zstd frame larger than 4 GiB cannot fit a u32 length prefix", + ) + })?; + let mut out = Vec::with_capacity(4 + frame.len()); + out.extend_from_slice(&frame_len.to_le_bytes()); + out.extend_from_slice(&frame); + Ok(out) + } + } + } +} + +/// File prologue for `codec`: `ZSPILL_MAGIC` for zstd, empty for bgzf (a BGZF +/// stream needs no prologue — each block self-identifies via the gzip magic). +#[must_use] +pub fn spill_magic(codec: SpillCodec) -> &'static [u8] { + match codec { + SpillCodec::Bgzf => &[], + SpillCodec::Zstd => &ZSPILL_MAGIC, + } +} + +/// File epilogue for `codec`: the `BGZF_EOF` empty-block terminator for bgzf, +/// empty for zstd (the zstd spill format has no trailing marker). +#[must_use] +pub fn spill_trailer(codec: SpillCodec) -> &'static [u8] { + match codec { + SpillCodec::Bgzf => &BGZF_EOF, + SpillCodec::Zstd => &[], + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::external::GenericKeyedChunkReader; + use crate::inline::TemplateKey; + use std::fs::File; + use std::io::{BufWriter, Write as _}; + use std::path::Path; + use tempfile::TempDir; + + #[allow(clippy::cast_possible_truncation)] + fn make_key(i: u64) -> TemplateKey { + TemplateKey::new( + i as i32, + i as i32, + false, + i32::MAX, + i32::MAX, + false, + 0, + 0, + (0, false), + i, + false, + ) + } + + #[allow(clippy::cast_possible_truncation)] + fn sample_records(n: u64) -> Vec<(TemplateKey, Vec)> { + (0..n).map(|i| (make_key(i), vec![(i % 256) as u8; 200 + (i as usize % 50)])).collect() + } + + /// The writer must reject exactly what the reader cannot decompress. + /// + /// These two limits are one contract with two enforcement points, and the + /// failure mode when they disagree is asymmetric: a writer that accepts more + /// than the reader accepts produces a spill file this process wrote and + /// cannot read back. Rejecting at write time turns that into a loud error + /// on the block that caused it. + /// + /// Both cases allocate a buffer of the size under test, so the at-limit case + /// really does build and compress 64 MiB. That is deliberate — it is the + /// only way to show the limit is inclusive rather than off by one — and it + /// is affordable because the buffer is zeroed, which zstd collapses almost + /// instantly. The over-limit case is the cheap one: `compress_block` checks + /// the length before touching the compressor, so it short-circuits. + #[rstest::rstest] + #[case::at_the_limit(MAX_SPILL_BLOCK_LEN, true)] + #[case::one_over_the_limit(MAX_SPILL_BLOCK_LEN + 1, false)] + fn compress_block_accepts_exactly_what_the_reader_can_decompress( + #[case] len: usize, + #[case] accepted: bool, + ) { + // `compress_block` checks the length before touching the compressor, so + // a zeroed buffer of the boundary size exercises the branch without + // needing meaningful content. + let raw = vec![0u8; len]; + let mut compressor = SpillBlockCompressor::new(SpillCodec::Zstd, 1).expect("compressor"); + let result = compressor.compress_block(&raw); + assert_eq!( + result.is_ok(), + accepted, + "a {len}-byte block should {} at the {MAX_SPILL_BLOCK_LEN}-byte limit", + if accepted { "be accepted" } else { "be rejected" }, + ); + if !accepted { + let err = result.unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::InvalidInput); + assert!( + err.to_string().contains("exceeds the"), + "the error must name the limit; got: {err}", + ); + } + } + + /// A zero-length body must frame as `[key][0u32]` and read back as an empty + /// record, on BOTH key branches. + /// + /// The failure this guards is not an error but a silent desync: the length + /// prefix is what the reader advances by, so a framer or reader that + /// mishandles a zero length leaves the stream misaligned and every + /// *subsequent* record decodes as garbage. The neighbouring record proves + /// alignment survived — asserting only that the empty record round-trips + /// would miss exactly that. + #[rstest::rstest] + #[case::non_embedded_key(false)] + #[case::embedded_key(true)] + fn zero_length_record_body_frames_and_keeps_the_stream_aligned(#[case] embedded: bool) { + let mut block = Vec::new(); + if embedded { + let key = crate::keys::RawCoordinateKey { sort_key: 7 }; + frame_keyed_record_into(&mut block, &key, &[]).unwrap(); + // `[u32 LE 0]` and nothing else — no key prefix, no body. + assert_eq!(block, 0u32.to_le_bytes().to_vec(), "embedded zero-length framing"); + frame_keyed_record_into(&mut block, &key, &[0xAB, 0xCD]).unwrap(); + let mut expected = 0u32.to_le_bytes().to_vec(); + expected.extend_from_slice(&2u32.to_le_bytes()); + expected.extend_from_slice(&[0xAB, 0xCD]); + assert_eq!(block, expected, "the next record must start right after the empty one"); + } else { + let key = make_key(3); + frame_keyed_record_into(&mut block, &key, &[]).unwrap(); + let mut key_bytes = Vec::new(); + key.write_to(&mut key_bytes).unwrap(); + let mut expected = key_bytes.clone(); + expected.extend_from_slice(&0u32.to_le_bytes()); + assert_eq!(block, expected, "non-embedded zero-length framing"); + frame_keyed_record_into(&mut block, &key, &[0xAB, 0xCD]).unwrap(); + expected.extend_from_slice(&key_bytes); + expected.extend_from_slice(&2u32.to_le_bytes()); + expected.extend_from_slice(&[0xAB, 0xCD]); + assert_eq!(block, expected, "the next record must start right after the empty one"); + } + } + + /// An embedded key (`RawCoordinateKey`, `EMBEDDED_IN_RECORD == true`) writes + /// no key prefix — the framed block is exactly `[u32 LE len][record]`. Locks + /// the embedded layout at the kernel boundary (the streaming-writer tests use + /// the non-embedded `TemplateKey`, which exercises the key-prefix branch). + #[test] + fn frame_omits_key_prefix_for_embedded_keys() { + let mut block = Vec::new(); + let key = crate::keys::RawCoordinateKey { sort_key: 0x0102_0304_0506_0708 }; + frame_keyed_record_into(&mut block, &key, &[0xAB, 0xCD]).unwrap(); + // Just the u32 LE length (2) + the 2 record bytes — no key prefix. + assert_eq!(block, [2, 0, 0, 0, 0xAB, 0xCD]); + } + + /// Write a spill file via the block kernel: frame records into raw blocks + /// (cutting a new block every `records_per_block` records, to force many + /// independent compressed units), compress each block, and bracket with the + /// codec magic/trailer. + fn write_via_kernel( + path: &Path, + codec: SpillCodec, + compression: u32, + records: &[(TemplateKey, Vec)], + records_per_block: usize, + ) { + let mut comp = SpillBlockCompressor::new(codec, compression).unwrap(); + let mut out = BufWriter::with_capacity(256 * 1024, File::create(path).unwrap()); + out.write_all(spill_magic(codec)).unwrap(); + let mut block = Vec::new(); + let mut in_block = 0usize; + let flush = + |block: &mut Vec, comp: &mut SpillBlockCompressor, out: &mut BufWriter| { + out.write_all(&comp.compress_block(block).unwrap()).unwrap(); + block.clear(); + }; + for (k, r) in records { + frame_keyed_record_into(&mut block, k, r).unwrap(); + in_block += 1; + if in_block >= records_per_block { + flush(&mut block, &mut comp, &mut out); + in_block = 0; + } + } + flush(&mut block, &mut comp, &mut out); + out.write_all(spill_trailer(codec)).unwrap(); + out.flush().unwrap(); + } + + fn read_back(path: &Path) -> Vec<(TemplateKey, Vec)> { + let mut reader = + GenericKeyedChunkReader::::open(path, None).expect("open reader"); + let mut buf = Vec::new(); + let mut out = Vec::new(); + while let Some(key) = reader.next_record(&mut buf).expect("read record") { + out.push((key, buf.clone())); + } + out + } + + /// Kernel-built files read back to the original records for every codec/level + /// — including a small block size that forces many independent compressed + /// units, proving block-independent compression is reader-compatible. + #[test] + fn kernel_blocks_read_back_records() { + let dir = TempDir::new().unwrap(); + let records = sample_records(400); + for (codec, level, name) in + [(SpillCodec::Zstd, 3, "z3"), (SpillCodec::Bgzf, 1, "b1"), (SpillCodec::Bgzf, 0, "b0")] + { + let path = dir.path().join(format!("kernel_{name}.keyed")); + write_via_kernel(&path, codec, level, &records, 7); + assert_eq!(read_back(&path), records, "kernel round-trip mismatch for {name}"); + } + } + + /// Kernel output matches the streaming `write_sorted_chunk` oracle, read back + /// through the production reader — locks block-granular ≡ streaming. + #[test] + fn kernel_matches_streaming_oracle_read_back() { + let dir = TempDir::new().unwrap(); + let records = sample_records(500); + let keyed: Vec<(TemplateKey, fgumi_raw_bam::RawRecord)> = + records.iter().map(|(k, r)| (*k, fgumi_raw_bam::RawRecord::from(r.clone()))).collect(); + for (codec, level) in [(SpillCodec::Zstd, 3), (SpillCodec::Bgzf, 1)] { + // Unique paths per codec: the spill writers now create files + // exclusively (`create_new`), so a path may not be reused across + // iterations. + let kernel_path = dir.path().join(format!("kernel_{codec:?}.keyed")); + let oracle_path = dir.path().join(format!("oracle_{codec:?}.keyed")); + write_via_kernel(&kernel_path, codec, level, &records, 11); + crate::write_sorted_chunk(&oracle_path, codec, level, &keyed).unwrap(); + assert_eq!( + read_back(&kernel_path), + read_back(&oracle_path), + "kernel vs streaming-oracle read-back differ ({codec:?})" + ); + } + } + + /// Different block boundaries produce files that read back identically — the + /// invariant that lets the (Serial) serialize step choose boundaries freely. + #[test] + fn reblocking_is_boundary_independent() { + let dir = TempDir::new().unwrap(); + let records = sample_records(300); + for codec in [SpillCodec::Zstd, SpillCodec::Bgzf] { + let a = dir.path().join("a.keyed"); + let b = dir.path().join("b.keyed"); + write_via_kernel(&a, codec, 1, &records, 3); + write_via_kernel(&b, codec, 1, &records, 64); + assert_eq!(read_back(&a), read_back(&b), "re-blocking changed readback ({codec:?})"); + } + } +} diff --git a/crates/fgumi-sort/src/spill_block_reader.rs b/crates/fgumi-sort/src/spill_block_reader.rs new file mode 100644 index 000000000..6a4846c26 --- /dev/null +++ b/crates/fgumi-sort/src/spill_block_reader.rs @@ -0,0 +1,605 @@ +//! Codec-aware streaming block decompressor for the typed-step +//! `SortSpillDecompress` pipeline step. +//! +//! A spill chunk is either a BGZF block-stream (self-framed `1f 8b` blocks) or +//! a zstd "ZSP1" stream (`[u32 LE frame-len][zstd frame]` records after the +//! 4-byte file magic). Both standalone `fgumi sort` (via the worker pool) and +//! the fused streaming sort write spill chunks with whichever +//! [`SpillCodec`] the sorter is configured for; this +//! decompressor lets the streaming `SortSpillDecompress` step read either. +//! +//! Each decompressed BGZF block (or zstd frame) is returned as one `Vec`. +//! The block/frame boundaries are immaterial: the downstream `MergeDriver` +//! parses `[len][record]` records out of a slot's decompressed-block queue and +//! reassembles records that span block boundaries, so any chunking of the +//! decompressed byte stream is correct. + +use std::io::{self, Read}; + +use libdeflater::Decompressor as BgzfDecompressor; +use zstd::bulk::Decompressor as ZstdDecompressor; + +use crate::codec::SpillCodec; +use crate::worker_pool::read_length_prefix; + +/// Output staging-buffer capacity for a single decompressed block/frame. Mirrors +/// the worker pool's `ZSTD_FRAME_DECOMP_CAP` / `BgzfDecompress` scratch sizing so +/// the mimalloc size-class reuse pattern matches. +const SCRATCH_CAP: usize = 256 * 1024; + +/// Ceiling for sizing the scratch buffer up from [`SCRATCH_CAP`]. +/// +/// A single BAM record can legitimately exceed `SCRATCH_CAP` (long reads), so +/// the reader sizes up to fit rather than refusing — but not without limit: the +/// size comes from the frame header, which a corrupt frame can claim to be +/// anything, so an unbounded allocation on garbage would take the process out. +/// +/// This is deliberately the *same* constant the writer enforces, imported rather +/// than redeclared: if the reader capped lower than the writer accepted, a block +/// in between would write successfully and be unreadable. +use crate::spill_block::MAX_SPILL_BLOCK_LEN as MAX_SCRATCH_CAP; + +/// Per-worker codec-aware decompressor. Holds a libdeflate decompressor (BGZF) +/// and a zstd decompressor plus reusable scratch buffers; one instance per +/// `SortSpillDecompress` worker copy. +pub struct SpillBlockDecompressor { + bgzf: BgzfDecompressor, + zstd: ZstdDecompressor<'static>, + /// Reused output buffer for the BGZF path (`mem::replace`d into the result). + bgzf_scratch: Vec, + /// Reused output buffer for the zstd path (decompressed-into, then copied). + zstd_buf: Vec, + /// Reused input buffer for the zstd path (one compressed frame at a time). + zstd_frame: Vec, +} + +impl SpillBlockDecompressor { + /// Construct a fresh decompressor. + /// + /// # Panics + /// + /// Panics if the zstd decompressor cannot be created (allocation failure). + #[must_use] + pub fn new() -> Self { + Self { + bgzf: BgzfDecompressor::new(), + zstd: ZstdDecompressor::new().expect("zstd decompressor init"), + bgzf_scratch: Vec::with_capacity(SCRATCH_CAP), + zstd_buf: Vec::new(), + zstd_frame: Vec::new(), + } + } + + /// Read and decompress up to `max` blocks from `reader` using `codec`, + /// returning the decompressed block bytes. A result with fewer than `max` + /// entries (including an empty `Vec`) signals that the reader reached a + /// clean EOF. + /// + /// The reader must be positioned at a block boundary — for zstd, at the + /// start of a `[len][frame]` record (i.e. the `ZSP1` file magic already + /// consumed); for BGZF, at the start of a block. `slots_for_chunk_files` + /// detects the codec from the file magic and positions the reader + /// accordingly when opening each slot. + /// + /// # Errors + /// + /// Propagates I/O errors, BGZF/zstd decompression failures, and truncation + /// (a partial length prefix or frame body at EOF). + pub fn read_blocks( + &mut self, + reader: &mut R, + codec: SpillCodec, + max: usize, + ) -> io::Result>> { + match codec { + SpillCodec::Bgzf => self.read_bgzf_blocks(reader, max), + SpillCodec::Zstd => self.read_zstd_frames(reader, max), + } + } + + fn read_bgzf_blocks( + &mut self, + reader: &mut R, + max: usize, + ) -> io::Result>> { + let raw_blocks = fgumi_bgzf::reader::read_raw_blocks(reader, max)?; + let mut out = Vec::with_capacity(raw_blocks.len()); + for raw in raw_blocks { + fgumi_bgzf::reader::decompress_block_slice_into( + &raw.data, + &mut self.bgzf, + &mut self.bgzf_scratch, + )?; + // mem::replace the filled scratch out and re-allocate a fresh one, + // so the consumer owns the bytes and the next decompress reuses a + // same-size-class allocation (matches `BgzfDecompress`). + out.push(std::mem::replace(&mut self.bgzf_scratch, Vec::with_capacity(SCRATCH_CAP))); + } + Ok(out) + } + + /// Decompress one zstd frame into the reusable scratch buffer, sizing it up + /// first if the output would not fit, and return the decompressed length. + /// + /// Both read paths funnel through here so the sizing policy cannot drift + /// between them — they previously each held their own copy of the + /// resize-then-decompress sequence, and only one of the two would have been + /// fixed by a change to either. + /// + /// The size comes from the frame header rather than from retrying on a + /// too-small-destination error. `decompress_to_buffer` surfaces zstd's + /// `get_error_name` text through `io::Error`, and that wording is not an + /// API contract — matching on it means a zstd bump could silently turn + /// "size up and succeed" into a hard failure on a spill file this very + /// process wrote. Reading the header also allocates exactly once instead of + /// doubling into place. + fn decompress_zstd_frame(&mut self, frame: &[u8]) -> io::Result { + // Both spill writers compress each block in a single + // `Compressor::compress` call, which records the content size in the + // frame header, so a missing size means the frame did not come from us. + let content_size = zstd::zstd_safe::get_frame_content_size(frame) + .map_err(|e| io::Error::other(format!("zstd spill frame header: {e}")))? + .ok_or_else(|| { + io::Error::other("zstd spill frame carries no content size (not written by fgumi)") + })?; + let needed = usize::try_from(content_size).map_err(|_| { + io::Error::other(format!("zstd spill frame content size {content_size} exceeds usize")) + })?; + if needed > MAX_SCRATCH_CAP { + return Err(io::Error::other(format!( + "zstd spill frame decompresses to {needed} bytes, over the \ + {MAX_SCRATCH_CAP}-byte scratch ceiling (or the frame is corrupt)" + ))); + } + // Reuse by CAPACITY, not by length. `decompress_to_buffer` sets the + // vector's length to the decompressed byte count, so after any frame + // `len` is the size of that frame's output, not the buffer size — a + // `resize(want, 0)` guarded on `len` would therefore zero-fill the + // difference on essentially every frame, immediately before zstd + // overwrites it. `clear` + `reserve` keeps the allocation and zeroes + // nothing; `reserve` is a no-op once capacity is high enough, so the + // buffer still never shrinks. + let want = needed.max(SCRATCH_CAP); + self.zstd_buf.clear(); + self.zstd_buf.reserve(want); + self.zstd + .decompress_to_buffer(frame, &mut self.zstd_buf) + .map_err(|e| io::Error::other(format!("zstd spill frame decompress: {e}"))) + } + + fn read_zstd_frames( + &mut self, + reader: &mut R, + max: usize, + ) -> io::Result>> { + let mut out = Vec::with_capacity(max); + for _ in 0..max { + // `read_length_prefix` returns `Ok(None)` only at a clean frame + // boundary EOF; a 1–3 byte partial prefix surfaces as an error. + let Some(frame_len) = read_length_prefix(reader)? else { + break; + }; + // Read into spare capacity rather than `resize(frame_len, 0)` + + // `read_exact`: the resize zero-fills the whole frame on every + // iteration only to overwrite it immediately. `read_to_end` appends, + // so the reuse is by capacity and nothing is zeroed. It also stops + // short instead of erroring, so the short read is checked here to + // keep `read_exact`'s `UnexpectedEof` semantics. + self.zstd_frame.clear(); + self.zstd_frame.reserve(frame_len); + Read::take(&mut *reader, frame_len as u64).read_to_end(&mut self.zstd_frame)?; + if self.zstd_frame.len() != frame_len { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!( + "truncated zstd spill frame: expected {frame_len} bytes, got {}", + self.zstd_frame.len(), + ), + )); + } + // Move the frame out so `decompress_zstd_frame` can take `&mut self`, + // then put it back to keep its allocation across frames. An error + // here abandons the whole read, so not restoring it on that path + // costs nothing. + let frame = std::mem::take(&mut self.zstd_frame); + let n = self.decompress_zstd_frame(&frame)?; + self.zstd_frame = frame; + out.push(self.zstd_buf[..n].to_vec()); + } + Ok(out) + } + + /// Read up to `max` *raw* (still-compressed) blocks from `reader` using + /// `codec`, returning the compressed payloads **without** decompressing them. + /// + /// This is the read half of the block-parallel `SortSpillDecompress` path: + /// the caller acquires the per-slot reader lock, calls `read_raw` to pull a + /// batch of compressed blocks (which it sequence-tags), releases the lock, + /// and then decompresses each block via [`Self::decompress_one`] *outside* + /// the lock so multiple workers decompress one file's blocks concurrently. + /// The read and the decompression of a given block still happen within a + /// single `try_run` of a single worker — only the lock is released between + /// them — which preserves the read-and-decompress-together invariant that + /// the FIFO inline path also upholds. + /// + /// For BGZF each returned `Vec` is one complete raw block (header + + /// compressed data + footer), exactly what [`Self::decompress_one`] expects; + /// EOF-marker blocks are skipped. For zstd each is one raw frame body (the + /// `[u32 LE len]` prefix is consumed here). A result shorter than `max` + /// (including empty) signals a clean EOF, matching [`Self::read_blocks`]. + /// + /// # Errors + /// + /// Propagates I/O errors and truncation (a partial BGZF block, or a partial + /// zstd length prefix / frame body at EOF). + pub fn read_raw( + &mut self, + reader: &mut R, + codec: SpillCodec, + max: usize, + ) -> io::Result>> { + match codec { + SpillCodec::Bgzf => { + let raw_blocks = fgumi_bgzf::reader::read_raw_blocks(reader, max)?; + Ok(raw_blocks.into_iter().map(|b| b.data).collect()) + } + SpillCodec::Zstd => { + let mut out = Vec::with_capacity(max); + for _ in 0..max { + let Some(frame_len) = read_length_prefix(reader)? else { + break; + }; + // Same zero-fill removal as `read_zstd_frames`: `vec![0u8; n]` + // zeroes the frame immediately before the read overwrites it. + // The *reuse* half of that fix does not transfer — `read_raw` + // hands each frame out as an owned `Vec` — but the zero-fill + // half does. `read_to_end` stops short instead of erroring, so + // the short read is checked to keep `read_exact`'s + // `UnexpectedEof` semantics. + let mut frame = Vec::with_capacity(frame_len); + Read::take(&mut *reader, frame_len as u64).read_to_end(&mut frame)?; + if frame.len() != frame_len { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!( + "truncated zstd spill frame: expected {frame_len} bytes, got {}", + frame.len(), + ), + )); + } + out.push(frame); + } + Ok(out) + } + } + } + + /// Decompress a single raw block/frame previously read by [`Self::read_raw`]. + /// + /// `raw` is one BGZF raw block (header + compressed data + footer) or one + /// zstd frame body, per `codec`. Returns the decompressed bytes. Uses the + /// worker's reusable scratch buffers, so this is cheap to call in a loop over + /// a freshly-read batch. + /// + /// # Errors + /// + /// Propagates BGZF/zstd decompression failures (bad CRC, size mismatch, or a + /// malformed frame). + pub fn decompress_one(&mut self, codec: SpillCodec, raw: &[u8]) -> io::Result> { + match codec { + SpillCodec::Bgzf => { + fgumi_bgzf::reader::decompress_block_slice_into( + raw, + &mut self.bgzf, + &mut self.bgzf_scratch, + )?; + Ok(std::mem::replace(&mut self.bgzf_scratch, Vec::with_capacity(SCRATCH_CAP))) + } + SpillCodec::Zstd => { + let n = self.decompress_zstd_frame(raw)?; + Ok(self.zstd_buf[..n].to_vec()) + } + } + } +} + +impl Default for SpillBlockDecompressor { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::pooled_chunk_writer::PooledChunkWriter; + use crate::worker_pool::SortWorkerPool; + use std::io::Cursor; + use std::sync::Arc; + + /// A length prefix that promises more bytes than follow must surface + /// `UnexpectedEof` from BOTH zstd read paths. + /// + /// `read_zstd_frames` and `read_raw` each read the frame body via + /// `Read::take(..).read_to_end(..)`, which — unlike the `read_exact` it + /// replaced — *stops short* instead of erroring. The explicit length check is + /// what restores the loud failure, and it exists in two places, so this is a + /// guard-set-parity case: drop the check in either path and a truncated spill + /// gets silently decompressed as a short frame with every test still green. + #[rstest::rstest] + #[case::read_zstd_frames(false)] + #[case::read_raw(true)] + fn a_truncated_zstd_frame_is_unexpected_eof(#[case] via_read_raw: bool) { + // `[u32 LE 64]` promising 64 body bytes, followed by only 8. + let mut bytes = 64u32.to_le_bytes().to_vec(); + bytes.extend_from_slice(&[0xEE; 8]); + + let mut d = SpillBlockDecompressor::new(); + let mut cursor = Cursor::new(bytes); + let err = if via_read_raw { + d.read_raw(&mut cursor, SpillCodec::Zstd, 4).expect_err("truncated frame must error") + } else { + d.read_zstd_frames(&mut cursor, 4).expect_err("truncated frame must error") + }; + assert_eq!( + err.kind(), + io::ErrorKind::UnexpectedEof, + "a short frame must be UnexpectedEof, not a silent short read; got: {err}", + ); + assert!( + err.to_string().contains("truncated zstd spill frame"), + "the error must name the truncation; got: {err}", + ); + } + + /// A zstd frame with no recorded content size must be rejected, not + /// decompressed into an unsized buffer. + /// + /// The reader sizes its scratch from the frame header, so a frame that omits + /// the content size has nothing to size from. Both spill writers compress + /// each block in one `bulk::compress` call, which always records it — so a + /// size-less frame did not come from fgumi, and guessing at its size is the + /// wrong response. `stream::encode_all` produces exactly that shape (verified + /// against zstd 0.13: bulk records the size, the streaming APIs do not), so + /// this is a real frame rather than a hand-forged header. + #[test] + fn a_zstd_frame_without_a_content_size_is_rejected() { + let payload = vec![9u8; 5000]; + let sizeless = zstd::stream::encode_all(&payload[..], 1).expect("stream-compress"); + assert!( + zstd::zstd_safe::get_frame_content_size(&sizeless).expect("readable header").is_none(), + "precondition: the streaming API must omit the content size", + ); + + let mut d = SpillBlockDecompressor::new(); + let err = d + .decompress_one(SpillCodec::Zstd, &sizeless) + .expect_err("a size-less frame must be rejected"); + assert!( + err.to_string().contains("carries no content size"), + "the error must say the size is missing; got: {err}", + ); + } + + /// A single zstd block larger than `SCRATCH_CAP` must survive the round + /// trip. Before the reader sized its scratch to the frame, such a block + /// wrote successfully and then failed on read — a spill file that could be + /// produced but not consumed. Refusing the write was never the answer: a + /// long-read BAM record alone can exceed 256 KiB. + /// + /// `compress_block` does now bound what it accepts, but at + /// `MAX_SPILL_BLOCK_LEN` (64 MiB), which is the reader's ceiling too — the + /// two agree by construction. This block is 768 KiB, above `SCRATCH_CAP` + /// and far below that bound, so it exercises scratch growth specifically + /// and not the writer's limit (which + /// `compress_block_accepts_exactly_what_the_reader_can_decompress` covers). + #[test] + fn oversized_zstd_block_round_trips_through_the_reader() { + let raw: Vec = (0..(SCRATCH_CAP * 3)).map(|i| u8::try_from(i % 251).unwrap()).collect(); + assert!(raw.len() > SCRATCH_CAP, "precondition: larger than the fixed scratch"); + + let mut compressor = + crate::spill_block::SpillBlockCompressor::new(SpillCodec::Zstd, 3).expect("compressor"); + let framed = compressor.compress_block(&raw).expect("compress accepts an oversized block"); + // `compress_block` prefixes the zstd frame with a u32 length; the + // block-level entry point takes the bare frame. + let frame = &framed[4..]; + + let mut decompressor = SpillBlockDecompressor::new(); + let out = decompressor + .decompress_one(SpillCodec::Zstd, frame) + .expect("reader must grow its scratch rather than fail"); + assert_eq!(out, raw, "oversized block round-trips byte-for-byte"); + + // The grown buffer is reused, so a normal-sized block still works after. + let small = vec![7u8; 1024]; + let small_framed = compressor.compress_block(&small).expect("compress"); + let small_frame = &small_framed[4..]; + let small_out = + decompressor.decompress_one(SpillCodec::Zstd, small_frame).expect("decompress"); + assert_eq!(small_out, small, "a later small block is unaffected by the growth"); + } + + /// A spill chunk written with codec X, read back through + /// `SpillBlockDecompressor`, must yield the exact `[len][record]` byte + /// stream that was written — for BOTH codecs. This is the round-trip the + /// streaming merge depends on. + fn roundtrip(codec: SpillCodec) { + use crate::keys::{RawCoordinateKey, RawSortKey}; + // Build framed records whose total size (~440 KB) far exceeds the + // writer's ~64 KB block cap, so the chunk is written as MANY blocks / + // frames. That is what actually exercises the per-block decompress and + // the downstream record reassembly across block boundaries — the whole + // reason `read_blocks` returns per-block `Vec`s. Each record is + // stamped with its index at both ends so a misaligned reassembly is + // caught, not just a stream of zeros. + let records: Vec> = (0usize..80) + .map(|i| { + let size = 3000 + (i % 13) * 500; // 3000..9000 bytes + let mut r = vec![0u8; size]; + let stamp = u8::try_from(i % 251).expect("i % 251 fits u8"); + r[0] = stamp; + let last = r.len() - 1; + r[last] = stamp; + r + }) + .collect(); + + let dir = tempfile::tempdir().expect("tempdir"); + let path = dir.path().join("chunk.spill"); + let pool = Arc::new(SortWorkerPool::new(1, 1, 6, codec)); + { + let mut w = PooledChunkWriter::::new(Arc::clone(&pool), &path, codec) + .expect("writer"); + for rec in &records { + let key = RawCoordinateKey::extract_from_record(rec); + w.write_record(&key, rec).expect("write"); + } + w.start_finish().expect("start_finish").wait().expect("finish"); + } + + // Re-open + position past the codec magic exactly like + // slots_for_chunk_files does. + let mut file = std::fs::File::open(&path).expect("open"); + let mut magic = [0u8; 4]; + let filled = crate::external::read_exact_or_eof(&mut file, &mut magic).expect("magic"); + let detected = if filled { + SpillCodec::from_magic(&magic).unwrap_or(SpillCodec::Bgzf) + } else { + SpillCodec::Bgzf + }; + assert_eq!(detected, codec, "codec must be detected from the file magic"); + if matches!(detected, SpillCodec::Bgzf) { + use std::io::Seek; + file.seek(std::io::SeekFrom::Start(0)).expect("seek"); + } + + let mut reader = std::io::BufReader::new(file); + let mut dec = SpillBlockDecompressor::new(); + let mut all = Vec::new(); + loop { + let blocks = dec.read_blocks(&mut reader, detected, 4).expect("read_blocks"); + let got = blocks.len(); + for b in blocks { + all.extend_from_slice(&b); + } + if got < 4 { + break; + } + } + + // The decompressed stream is `[u32 LE len][record]` per record, in order. + let mut cursor = Cursor::new(&all); + let mut read_back: Vec> = Vec::new(); + let mut len_buf = [0u8; 4]; + while Read::read(&mut cursor, &mut len_buf).map(|n| n == 4).unwrap_or(false) { + let len = u32::from_le_bytes(len_buf) as usize; + let mut rec = vec![0u8; len]; + cursor.read_exact(&mut rec).expect("record body"); + read_back.push(rec); + } + assert_eq!(read_back, records, "round-trip mismatch for {codec:?}"); + } + + #[test] + fn roundtrip_bgzf() { + roundtrip(SpillCodec::Bgzf); + } + + #[test] + fn roundtrip_zstd() { + roundtrip(SpillCodec::Zstd); + } + + /// The split `read_raw` + `decompress_one` path (used by the block-parallel + /// `SortSpillDecompress`) must yield the exact same per-block bytes as the + /// inline `read_blocks` path. We read the same chunk twice and compare. + fn read_raw_matches_read_blocks(codec: SpillCodec) { + use crate::keys::{RawCoordinateKey, RawSortKey}; + + let records: Vec> = (0usize..80) + .map(|i| { + let size = 3000 + (i % 13) * 500; + let mut r = vec![0u8; size]; + let stamp = u8::try_from(i % 251).expect("i % 251 fits u8"); + r[0] = stamp; + let last = r.len() - 1; + r[last] = stamp; + r + }) + .collect(); + + let dir = tempfile::tempdir().expect("tempdir"); + let path = dir.path().join("chunk.spill"); + let pool = Arc::new(SortWorkerPool::new(1, 1, 6, codec)); + { + let mut w = PooledChunkWriter::::new(Arc::clone(&pool), &path, codec) + .expect("writer"); + for rec in &records { + let key = RawCoordinateKey::extract_from_record(rec); + w.write_record(&key, rec).expect("write"); + } + w.start_finish().expect("start_finish").wait().expect("finish"); + } + + // Helper: open the chunk positioned past any codec magic. + let open = || { + let mut file = std::fs::File::open(&path).expect("open"); + let mut magic = [0u8; 4]; + let filled = crate::external::read_exact_or_eof(&mut file, &mut magic).expect("magic"); + let detected = if filled { + SpillCodec::from_magic(&magic).unwrap_or(SpillCodec::Bgzf) + } else { + SpillCodec::Bgzf + }; + if matches!(detected, SpillCodec::Bgzf) { + use std::io::Seek; + file.seek(std::io::SeekFrom::Start(0)).expect("seek"); + } + (std::io::BufReader::new(file), detected) + }; + + // Path A: inline read_blocks. + let (mut reader_a, detected) = open(); + let mut dec_a = SpillBlockDecompressor::new(); + let mut blocks_a: Vec> = Vec::new(); + loop { + let blocks = dec_a.read_blocks(&mut reader_a, detected, 4).expect("read_blocks"); + let got = blocks.len(); + blocks_a.extend(blocks); + if got < 4 { + break; + } + } + + // Path B: read_raw + decompress_one. + let (mut reader_b, _) = open(); + let mut dec_b = SpillBlockDecompressor::new(); + let mut blocks_b: Vec> = Vec::new(); + loop { + let raws = dec_b.read_raw(&mut reader_b, detected, 4).expect("read_raw"); + let got = raws.len(); + for raw in &raws { + blocks_b.push(dec_b.decompress_one(detected, raw).expect("decompress_one")); + } + if got < 4 { + break; + } + } + + assert_eq!( + blocks_a, blocks_b, + "split read_raw path must match inline read_blocks ({codec:?})" + ); + } + + #[test] + fn read_raw_matches_read_blocks_bgzf() { + read_raw_matches_read_blocks(SpillCodec::Bgzf); + } + + #[test] + fn read_raw_matches_read_blocks_zstd() { + read_raw_matches_read_blocks(SpillCodec::Zstd); + } +} diff --git a/crates/fgumi-sort/src/spill_reader.rs b/crates/fgumi-sort/src/spill_reader.rs deleted file mode 100644 index a0afc864d..000000000 --- a/crates/fgumi-sort/src/spill_reader.rs +++ /dev/null @@ -1,1142 +0,0 @@ -//! Scattered reads for the sort's input and spill files. -//! -//! Both phases read a file through one blocking sequential `read()` at a time: -//! phase 1 behind the exclusive `ReadInputBlocks` step, phase 2 behind each -//! spill file's reader mutex. That is queue depth 1, and a device serves one -//! stream far worse than several -- measured on the production volume, 358 MB/s -//! against 1177 MB/s for four concurrent streams, with buffer size making no -//! difference at all (2, 8 and 32 MiB all landed within 0.1%). A bigger single -//! request is not more concurrency. -//! -//! Measured consequence on `1kg-wgs-HG00096` (t16, 44 spill runs): the merge -//! read 52.6 GB of spill back at **313 MB/s**, costing 168.3s of worker-busy -//! time, and phase 1's read span was 141.5s. Raising `read_ahead_kb` to 4096 -//! fixes both, but it is root-only and system-wide, so a tool cannot rely on -//! it. `posix_fadvise` was measured as the per-file substitute and does not -//! work -- `SEQUENTIAL` 359 MB/s and `WILLNEED` 348-350 MB/s against 360 with -//! no hint, `strace` confirming the calls succeed and the kernel simply -//! declines. That leaves doing the concurrency ourselves. -//! -//! # Why bytes and not blocks -//! -//! BGZF blocks and zstd frames are variable-length, so block *N*'s offset is -//! unknowable without scanning from block 0 -- which is why framing is serial. -//! It can stay that way: framing is ~4% of the reader's time. So a fill fetches -//! fixed-size *byte* slices at known offsets and this type serves them in file -//! order through a plain [`Read`]. The framers above are generic over `Read` -//! and do not change. -//! -//! # Why the slices run on the pool -//! -//! The obvious implementation gives each reader its own threads. That was -//! measured, and it works -- and it is wrong, because `--threads N` then buys -//! the user N pool workers plus an ingest thread plus however many readers are -//! live. On a 16 vCPU host the 4-stream configuration ran 21 threads and -//! doubled involuntary context switches (1.39M to 2.81M). It won 31% of wall -//! clock only because that box had five idle cores to lend. -//! -//! So a fill is offered to the pool as [`FetchSlice`] work any worker may -//! steal, and the thread budget stays exactly `--threads` + the ingest thread. -//! The filling thread reads the first slice itself and then **reclaims** every -//! slice no worker has started, rather than waiting on it: submitting and -//! blocking would deadlock the moment every worker is a filler waiting for -//! slices nobody is left to run. Reclaim makes the worst case "exactly the -//! single sequential read we do today", which is a slowdown and never a wedge. - -use std::collections::VecDeque; -use std::fs::File; -use std::io::{self, BufReader, Read}; -use std::os::unix::fs::FileExt; -use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; -use std::sync::{Arc, Condvar, Mutex}; - -use crossbeam_queue::ArrayQueue; - -/// Bytes one fill fetches, across all its slices. -/// -/// Twice the 2 MiB buffer the sequential path uses, so that at four streams a -/// slice is still 1 MiB: the read ladder measured four 1 MiB streams at -/// 975 MB/s against 1081 MB/s at 4 MiB, so slices much below a megabyte start -/// giving the gain back. -pub(crate) const FILL_BYTES: usize = 4 * 1024 * 1024; - -/// Most streams the ramp will ever reach. -/// -/// Eight already matched four on the storage that needs concurrency at all -/// (1080 MB/s against 1081 on EBS gp3), and past this each extra stream is -/// pool work that buys nothing. -const MAX_STREAMS: usize = 8; - -/// Fills measured at one stream before choosing a stream count. -/// -/// Long enough that one slow fill cannot move the decision, short enough that -/// the probe costs 32 MiB of reading at whatever the device does unaided. -const AUTO_PROBE_FILLS: usize = 8; - -/// Single-stream throughput at or above which concurrency has nothing to buy. -/// -/// Derived from the consumer, not from the device: the ingest thread reads -/// 43.1 GB in ~53s, so it can absorb about 810 MB/s, and a reader that already -/// beats that is not what the sort is waiting on. This leaves ~1.5x headroom -/// over that floor. -/// -/// The two devices measured sit either side of it by a wide margin -- EBS gp3 -/// sustains 358 MB/s on one stream and wants four; a local instance-store SSD -/// sustains 2214 MB/s and wants one, where forcing eight measured 1.8% -/// *slower*. Anything from roughly 900 MB/s to 2 GB/s separates them. -const AUTO_TARGET_BYTES_PER_SEC: f64 = 1_200_000_000.0; - -/// Streams to use, from what one stream measured. -/// -/// `ceil(target / measured)`: enough streams to reach a rate the consumer -/// cannot outrun, and no more. Measuring the *device* rather than the pipeline -/// is what makes this work -- an earlier version compared fetch time against -/// the reader's own elapsed time, which is a high fraction on every device -/// because a reader mostly reads, and it duly ramped both gp3 and a local SSD -/// to the cap. -fn streams_for_measured_rate(bytes: u64, nanos: u64) -> usize { - if bytes == 0 || nanos == 0 { - return 1; - } - #[allow(clippy::cast_precision_loss)] - let rate = bytes as f64 * 1e9 / nanos as f64; - let wanted = (AUTO_TARGET_BYTES_PER_SEC / rate).ceil(); - if !wanted.is_finite() || wanted <= 1.0 { - return 1; - } - #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] - let wanted = wanted as usize; - wanted.min(MAX_STREAMS) -} - -/// Smallest slice worth offering to another worker. -/// -/// Below this the per-request overhead and the fill's own barrier cost more -/// than the concurrency buys, so a small fill uses fewer slices than the stream -/// count allows rather than splitting into pinpricks. -const MIN_SLICE_BYTES: usize = 512 * 1024; - -/// One slice of a fill: read a byte range into the buffer waiting for it. -/// -/// Cloned into the pool's queue as an `Arc` and also kept by the filling -/// thread, so whichever gets to it first runs it and the other no-ops. -pub(crate) struct FetchSlice { - file: Arc, - offset: u64, - index: usize, - state: Arc, -} - -/// The parts of a fill its slices report back into. -struct FillState { - progress: Mutex, - completed: Condvar, -} - -/// Slice buffers and how many are still outstanding. -struct FillProgress { - slots: Vec, - remaining: usize, - error: Option, -} - -/// A slice's buffer, and who has it. -/// -/// The transition out of `Pending` under the mutex *is* the claim: exactly one -/// thread can take the buffer, so a worker and the filling thread racing for -/// the same slice cannot both read it, and no separate flag is needed. -enum SlotState { - /// Sized and waiting for someone to read into it. - Pending(Vec), - /// Claimed; a thread is in `read_exact_at` right now. - InFlight, - /// Read, and holding its bytes. - Done(Vec), -} - -/// Fill slices offered to the pool, with the census of what it took. -pub(crate) struct FetchQueue { - jobs: ArrayQueue>, - offered: AtomicU64, - taken: AtomicU64, - /// The stream count one uncontended probe measured, or 0 while none has. - /// - /// Shared because the thing being measured is the *device*, and a merge - /// reads K spill files at once: each reader would otherwise measure its own - /// share of a contended device, conclude it needs more streams, and add - /// contention. Measured on EBS gp3 -- phase 1 alone saw 335 MB/s and chose - /// four, while the 44 spill readers each saw 224-229 MB/s and chose six or - /// seven, which took the merge from 119.3s to 129.1s. - chosen_streams: AtomicUsize, -} - -/// A [`Read`] over a file whose buffer is filled by concurrent positional reads. -pub(crate) struct ScatterReader { - file: Arc, - len: u64, - offset: u64, - ready: VecDeque>, - /// How far the consumer has read into `ready`'s front slice. A cursor - /// rather than draining the front: the framers read a few bytes at a time, - /// and shifting a megabyte down by thirteen bytes per call would cost more - /// than the reads this saves. - front_pos: usize, - free: Vec>, - /// Where slices are offered. `None` reads every fill on this thread, which - /// is what the tests and any pool-less caller get. - fetch: Option>, - streams: usize, - /// Where to book fill time, when this reader is the one the phase-1 - /// reader report describes. A fill is exactly what `TimedReader` measures - /// on the sequential arm -- the fetch, as opposed to the framing above it. - stats: Option>, - /// Fills whose slices have been offered but whose bytes are not collected, - /// oldest first and contiguous in the file. - pending: VecDeque, - /// How many fills to keep in flight ahead of the consumer, so the device - /// works while the framer frames. Zero for the merge's spill readers: the - /// merge already sits at its consumer-serial floor, so all lookahead would - /// buy there is `K * depth * FILL_BYTES` of extra buffers. - lookahead: usize, - /// Whether to grow `streams` from what the fills measure. See - /// [`ramped_streams`]. - auto: bool, - /// Bytes and nanoseconds the probe has measured at one stream so far. - probe_bytes: u64, - probe_nanos: u64, - probe_fills: usize, -} - -/// A fill in flight: its slices, where they report, and where it began. -/// -/// Its length is not stored because it is always `offset - start`: only the -/// most recently issued fill is ever pending, and `issue` advances `offset` by -/// exactly that fill's size. -struct PendingFill { - slices: Vec>, - state: Arc, - start: u64, -} - -/// How one file's bytes are read. -pub(crate) enum SpillSource { - /// One sequential stream -- what both phases did before, and what - /// `--read-streams 1` still does. - Sequential(BufReader), - /// Several positional reads at once, run by pool workers. - Scattered(ScatterReader), -} - -impl FetchQueue { - /// A queue sized for `workers` workers. - /// - /// Capacity only bounds how many slices can be *waiting*; a push that finds - /// it full is not an error, because the filling thread reclaims anything the - /// pool did not take. - pub(crate) fn new(workers: usize) -> Arc { - Arc::new(Self { - jobs: ArrayQueue::new(workers.max(1) * 8), - offered: AtomicU64::new(0), - taken: AtomicU64::new(0), - chosen_streams: AtomicUsize::new(0), - }) - } - - /// Whether any slice is waiting. Drives the step's eligibility. - pub(crate) fn is_empty(&self) -> bool { - self.jobs.is_empty() - } - - /// Run one offered slice. `false` means there was nothing left to do -- - /// either the queue was empty, or the filling thread had already reclaimed - /// the slice this worker popped. - pub(crate) fn run_one(&self) -> bool { - let Some(slice) = self.jobs.pop() else { - return false; - }; - if slice.run() { - self.taken.fetch_add(1, Ordering::Relaxed); - return true; - } - false - } - - /// `(slices offered, slices a worker actually ran)`. - /// - /// The census for "did this engage at all". Zero taken means every fill was - /// read by the thread that wanted it, which is the pre-change behaviour and - /// which no output check would ever notice. - pub(crate) fn census(&self) -> (u64, u64) { - (self.offered.load(Ordering::Relaxed), self.taken.load(Ordering::Relaxed)) - } - - /// The stream count a probe has already settled on, if any. - fn settled_streams(&self) -> Option { - match self.chosen_streams.load(Ordering::Relaxed) { - 0 => None, - n => Some(n), - } - } - - /// Publish what one uncontended probe measured, if nobody has yet. - fn settle_streams(&self, streams: usize) { - let _ = - self.chosen_streams.compare_exchange(0, streams, Ordering::Relaxed, Ordering::Relaxed); - } - - /// Offer a slice, ignoring a full queue -- the filler reclaims either way. - fn offer(&self, slice: &Arc) { - if self.jobs.push(Arc::clone(slice)).is_ok() { - self.offered.fetch_add(1, Ordering::Relaxed); - } - } -} - -impl FetchSlice { - /// Read this slice, unless someone else already claimed it. - /// - /// Returns whether this call did the read, which is what lets a worker tell - /// "I helped" from "the filler beat me to it". - fn run(&self) -> bool { - let mut buf = { - let mut progress = self.state.progress.lock().expect("fill state poisoned"); - match std::mem::replace(&mut progress.slots[self.index], SlotState::InFlight) { - SlotState::Pending(buf) => buf, - claimed => { - progress.slots[self.index] = claimed; - return false; - } - } - }; - let result = self.file.read_exact_at(&mut buf, self.offset); - let mut progress = self.state.progress.lock().expect("fill state poisoned"); - match result { - Ok(()) => progress.slots[self.index] = SlotState::Done(buf), - Err(e) => { - if progress.error.is_none() { - progress.error = Some(e); - } - } - } - progress.remaining -= 1; - let last = progress.remaining == 0; - drop(progress); - if last { - self.state.completed.notify_all(); - } - true - } -} - -impl ScatterReader { - /// Read `file` from `offset`, offering slices to `fetch`. - /// - /// # Errors - /// - /// Returns an error if the file's length cannot be determined. - pub(crate) fn new( - file: File, - offset: u64, - streams: usize, - fetch: Option>, - ) -> io::Result { - let len = file.metadata()?.len(); - Ok(Self { - file: Arc::new(file), - len, - offset, - ready: VecDeque::new(), - front_pos: 0, - free: Vec::new(), - fetch, - streams: streams.max(1), - stats: None, - pending: VecDeque::new(), - lookahead: 0, - auto: false, - probe_bytes: 0, - probe_nanos: 0, - probe_fills: 0, - }) - } - - /// Grow the stream count from what the fills measure, starting at one. - /// - /// Costs nothing before it has measured anything -- one stream is exactly - /// the pre-change path -- and only grows when the consumer is demonstrably - /// waiting on fills. See [`ramped_streams`]. - #[must_use] - pub(crate) fn auto_tuned(mut self) -> Self { - self.auto = true; - self.streams = 1; - self - } - - /// Keep `fills` in flight ahead of the consumer. - /// - /// Without this a fill is issued only once the previous one is exhausted, - /// so the device idles for as long as framing and decompression take. - /// Measured on the phase-1 read span: 69.4s at depth 0, 66.9s at depth 1, - /// and 62.3s for a reader that prefetched continuously with its own - /// threads -- so depth is worth something well past one. - /// - /// Costs `(fills + 1) * FILL_BYTES` of buffers for this reader. - #[must_use] - pub(crate) fn looking_ahead(mut self, fills: usize) -> Self { - self.lookahead = fills; - self - } - - /// Book this reader's fills into `stats`, so framing and fetching stay - /// separable in the phase-1 reader report. - #[must_use] - pub(crate) fn timed(mut self, stats: Arc) -> Self { - self.stats = Some(stats); - self - } - - /// Build a reader for a [`crate::external::ReadStreams`] setting, tuning - /// itself when asked. - /// - /// # Errors - /// - /// Returns an error if the file's length cannot be determined. - pub(crate) fn for_streams( - file: File, - offset: u64, - streams: crate::external::ReadStreams, - fetch: Option>, - ) -> io::Result { - let reader = Self::new(file, offset, streams.initial(), fetch)?; - Ok(if streams.is_auto() { reader.auto_tuned() } else { reader }) - } - - /// A buffer of exactly `len` bytes, recycled where possible. - /// - /// Sized rather than cleared and refilled: `read_exact_at` overwrites every - /// byte, so zero-filling a recycled buffer of the right length is pure cost - /// over the gigabytes a sort reads. - fn take_buf(&mut self, len: usize) -> Vec { - let mut buf = self.free.pop().unwrap_or_default(); - if buf.len() != len { - buf.resize(len, 0); - } - buf - } - - /// Offer the next [`FILL_BYTES`] as slices and advance the fetch offset. - /// - /// Returns `None` at EOF. Nothing is read here beyond what a worker happens - /// to pick up: the bytes are claimed in [`Self::collect`], which is what - /// lets a fill be in flight while the caller does something else. - fn issue(&mut self) -> Option { - let remaining = self.len.saturating_sub(self.offset); - if remaining == 0 { - return None; - } - let want = usize::try_from(remaining.min(FILL_BYTES as u64)).expect("a fill fits usize"); - // Never split below `MIN_SLICE_BYTES`: a fill near EOF is small, and - // cutting it into pinpricks costs more than the concurrency returns. - let slices = self.streams.min(want.div_ceil(MIN_SLICE_BYTES)).max(1); - let slice_len = want.div_ceil(slices); - let start = self.offset; - - let bufs: Vec> = (0..slices) - .map(|index| self.take_buf(slice_len.min(want - index * slice_len))) - .collect(); - let state = Arc::new(FillState { - progress: Mutex::new(FillProgress { - slots: bufs.into_iter().map(SlotState::Pending).collect(), - remaining: slices, - error: None, - }), - completed: Condvar::new(), - }); - let slices: Vec> = (0..slices) - .map(|index| { - Arc::new(FetchSlice { - file: Arc::clone(&self.file), - offset: start + (index * slice_len) as u64, - index, - state: Arc::clone(&state), - }) - }) - .collect(); - - // Slice 0 stays for the collecting thread, which has nothing better to - // do than read it; the rest go where a worker can reach them. - if let Some(queue) = &self.fetch { - for slice in &slices[1..] { - queue.offer(slice); - } - } - self.offset += want as u64; - Some(PendingFill { slices, state, start }) - } - - /// Claim whatever the pool has not taken, wait for the rest, and publish the - /// fill's bytes in file order. - fn collect(&mut self, pending: &PendingFill) -> io::Result<()> { - // Walk every slice: the first is this thread's own share, and the rest - // are reclaimed if no worker has started them. `run` no-ops on a slice - // someone else took, so this is both "do my share" and "take back what - // nobody wanted" in one pass -- and it is why a full queue, an empty - // pool, or a pool of workers all busy filling cannot wedge here. - for slice in &pending.slices { - slice.run(); - } - - let mut progress = pending.state.progress.lock().expect("fill state poisoned"); - // Only slices a worker claimed before we got to them remain, and that - // worker is inside `read_exact_at` rather than waiting on anything, so - // this wait is bounded by one disk read. - while progress.remaining > 0 { - progress = pending.state.completed.wait(progress).expect("fill state poisoned"); - } - if let Some(e) = progress.error.take() { - drop(progress); - // `issue` moved the fetch offset before the bytes were claimed, so - // put it back: a caller that reads again must retry this range, not - // skip it. Anything issued after it covers bytes past the failure - // and is dropped for the same reason. - self.offset = pending.start; - self.pending.clear(); - return Err(e); - } - for slot in &mut progress.slots { - match std::mem::replace(slot, SlotState::InFlight) { - SlotState::Done(buf) => self.ready.push_back(buf), - _ => unreachable!("a fill with no error has every slice done"), - } - } - Ok(()) - } - - /// Fold one single-stream fill into the probe, and decide once it has seen - /// enough. Decides exactly once; after that the reader stops measuring. - fn probe_fill(&mut self, bytes: u64, nanos: u64) { - // Self-guarding rather than relying on the caller to stop: deciding - // twice is exactly how the previous design went wrong, and a probe that - // cannot be re-armed cannot repeat it. Later fills also run at the - // chosen stream count, so they no longer measure one stream and would - // answer a different question. - if !self.auto { - return; - } - // Adopt an answer someone else already measured rather than measuring a - // contended device. In a merge this is every spill reader but the first. - if let Some(settled) = self.fetch.as_ref().and_then(|q| q.settled_streams()) { - self.streams = settled; - self.auto = false; - return; - } - self.probe_bytes = self.probe_bytes.saturating_add(bytes); - self.probe_nanos = self.probe_nanos.saturating_add(nanos); - self.probe_fills += 1; - if self.probe_fills < AUTO_PROBE_FILLS { - return; - } - let chosen = streams_for_measured_rate(self.probe_bytes, self.probe_nanos); - #[allow(clippy::cast_precision_loss)] - let mbps = self.probe_bytes as f64 * 1e3 / self.probe_nanos.max(1) as f64; - log::debug!("read streams: one stream measured {mbps:.0} MB/s, using {chosen}"); - if let Some(queue) = &self.fetch { - queue.settle_streams(chosen); - } - self.streams = chosen; - self.auto = false; - } - - /// Make bytes available in `ready`, leaving it empty only at EOF. - fn fill(&mut self) -> io::Result<()> { - if self.pending.is_empty() { - let Some(first) = self.issue() else { - return Ok(()); - }; - self.pending.push_back(first); - } - let pending = self.pending.pop_front().expect("just ensured non-empty"); - let started = self.auto.then(std::time::Instant::now); - let fill_start = pending.start; - self.collect(&pending)?; - if let Some(t0) = started { - self.probe_fill( - self.offset.saturating_sub(fill_start), - u64::try_from(t0.elapsed().as_nanos()).unwrap_or(u64::MAX), - ); - } - // Topped up after collecting rather than before, so the fills issued - // here overlap the framing of what just landed rather than the wait for - // it. At depth 0 this loop does nothing and the reader is demand-driven. - while self.pending.len() < self.lookahead { - let Some(next) = self.issue() else { - break; - }; - self.pending.push_back(next); - } - Ok(()) - } -} - -impl Read for ScatterReader { - fn read(&mut self, out: &mut [u8]) -> io::Result { - if out.is_empty() { - return Ok(0); - } - if self.ready.is_empty() { - let before = self.offset; - let sched_before = - self.stats.as_ref().and_then(|_| crate::phase1_stats::thread_schedstat()); - let started = std::time::Instant::now(); - let outcome = self.fill(); - if let Some(stats) = &self.stats { - let elapsed = u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX); - if let (Some(a), Some(b)) = (sched_before, crate::phase1_stats::thread_schedstat()) - { - stats.record_refill_sched( - b.0.saturating_sub(a.0), - b.1.saturating_sub(a.1), - b.2.saturating_sub(a.2), - ); - } - // A failed fill still cost time; crediting it zero bytes keeps - // the throughput figure honest. - let fetched = usize::try_from(self.offset - before).unwrap_or(0); - stats.record_refill(elapsed, fetched); - } - outcome?; - if self.ready.is_empty() { - return Ok(0); - } - } - let front = self.ready.front().expect("just filled"); - let n = out.len().min(front.len() - self.front_pos); - out[..n].copy_from_slice(&front[self.front_pos..self.front_pos + n]); - self.front_pos += n; - if self.front_pos == front.len() { - let done = self.ready.pop_front().expect("just borrowed"); - self.free.push(done); - self.front_pos = 0; - } - Ok(n) - } -} - -#[cfg(test)] -impl ScatterReader { - /// Offset of the next byte the consumer will read. - fn position(&self) -> u64 { - let unconsumed: usize = - self.ready.iter().map(Vec::len).sum::().saturating_sub(self.front_pos); - // Pending fills are contiguous and end at the fetch offset, so the - // oldest one's start is where the in-flight region begins. - let in_flight = self.pending.front().map_or(0, |p| self.offset - p.start); - self.offset.saturating_sub(unconsumed as u64).saturating_sub(in_flight) - } -} - -#[cfg(test)] -impl SpillSource { - /// Offset of the next byte the consumer will read. - /// - /// Test-only: production never asks. The codec-detection tests assert that - /// a zstd spill starts past `ZSPILL_MAGIC` and a BGZF one at byte 0, and - /// that has to hold on whichever arm the stream count selects. - pub(crate) fn position(&mut self) -> io::Result { - match self { - Self::Sequential(reader) => std::io::Seek::stream_position(reader), - Self::Scattered(reader) => Ok(reader.position()), - } - } -} - -impl Read for SpillSource { - fn read(&mut self, out: &mut [u8]) -> io::Result { - match self { - Self::Sequential(reader) => reader.read(out), - Self::Scattered(reader) => reader.read(out), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::io::Write; - - /// A temp file of `len` pseudo-random bytes, plus the bytes themselves. - fn fixture(len: usize) -> (tempfile::NamedTempFile, Vec) { - let mut state = 0x2545_F491_4F6C_DD1Du64; - let bytes: Vec = (0..len) - .map(|_| { - state ^= state << 13; - state ^= state >> 7; - state ^= state << 17; - u8::try_from((state >> 24) & 0xff).expect("masked to a byte") - }) - .collect(); - let mut f = tempfile::NamedTempFile::new().expect("temp file"); - f.write_all(&bytes).expect("write"); - f.flush().expect("flush"); - (f, bytes) - } - - fn read_all(reader: &mut impl Read, buf_size: usize) -> Vec { - let mut out = Vec::new(); - let mut buf = vec![0u8; buf_size]; - loop { - let n = reader.read(&mut buf).expect("read"); - if n == 0 { - break; - } - out.extend_from_slice(&buf[..n]); - } - out - } - - /// A reader whose slices are offered to a queue **nobody drains**, which is - /// the reclaim path: every test below therefore exercises it by default. - fn scattered(path: &std::path::Path, offset: u64, streams: usize) -> ScatterReader { - let file = File::open(path).expect("open"); - ScatterReader::new(file, offset, streams, Some(FetchQueue::new(4))).expect("reader") - } - - #[test] - fn test_read_streams_parses_what_a_user_would_type() { - use crate::external::ReadStreams; - use std::str::FromStr; - assert_eq!(ReadStreams::from_str("auto"), Ok(ReadStreams::Auto)); - assert_eq!(ReadStreams::from_str("AUTO"), Ok(ReadStreams::Auto), "case-insensitive"); - assert_eq!(ReadStreams::from_str("1"), Ok(ReadStreams::Fixed(1))); - assert_eq!(ReadStreams::from_str("4"), Ok(ReadStreams::Fixed(4))); - // Zero streams cannot read anything, and silently promoting it to one - // would hide a typo behind a plausible-looking run. - assert!(ReadStreams::from_str("0").is_err(), "zero is rejected, not rounded up"); - assert!(ReadStreams::from_str("four").is_err()); - assert!(ReadStreams::from_str("-1").is_err()); - assert!(ReadStreams::from_str("").is_err()); - } - - #[test] - fn test_read_streams_round_trips_through_its_display() { - // clap prints the default in `--help` via `Display`, so a value it - // shows has to be one it accepts. - use crate::external::ReadStreams; - use std::str::FromStr; - for value in [ReadStreams::Auto, ReadStreams::Fixed(1), ReadStreams::Fixed(8)] { - assert_eq!(ReadStreams::from_str(&value.to_string()), Ok(value)); - } - } - - /// Bytes and nanos for a device sustaining `mbps` on one stream. - fn probe_of(mbps: u64) -> (u64, u64) { - (mbps * 1_000_000, 1_000_000_000) - } - - #[test] - fn test_the_probe_picks_the_stream_count_each_measured_device_wanted() { - // The two devices actually measured, at the rate one stream sustained - // on each. gp3 gained 28% from four streams; the instance-store SSD - // lost 1.8% from being pushed past one. - let (bytes, nanos) = probe_of(358); - assert_eq!(streams_for_measured_rate(bytes, nanos), 4, "EBS gp3 at 358 MB/s wants four"); - let (bytes, nanos) = probe_of(2214); - assert_eq!( - streams_for_measured_rate(bytes, nanos), - 1, - "a local SSD at 2214 MB/s wants one" - ); - } - - #[test] - fn test_the_probe_scales_between_those_two_points() { - // "Smarter than 1 or 4" is the point: the count is computed from the - // measurement, so storage between the two measured devices gets a - // count between the two answers. - assert_eq!(streams_for_measured_rate(probe_of(700).0, probe_of(700).1), 2); - assert_eq!(streams_for_measured_rate(probe_of(1199).0, probe_of(1199).1), 2); - assert_eq!(streams_for_measured_rate(probe_of(1200).0, probe_of(1200).1), 1); - assert_eq!(streams_for_measured_rate(probe_of(5000).0, probe_of(5000).1), 1); - } - - #[test] - fn test_the_probe_is_capped_however_slow_the_device() { - // A very slow mount would ask for dozens of streams; past the cap they - // are pool work that buys nothing. - assert_eq!(streams_for_measured_rate(probe_of(10).0, probe_of(10).1), MAX_STREAMS); - assert_eq!(streams_for_measured_rate(1, u64::MAX), MAX_STREAMS); - } - - #[test] - fn test_the_probe_refuses_to_divide_by_a_degenerate_measurement() { - // A zero-length or zero-byte probe would produce an infinity or a NaN, - // and a NaN comparison is false, which would silently pin the reader at - // whatever it happened to be. - assert_eq!(streams_for_measured_rate(0, 1_000), 1); - assert_eq!(streams_for_measured_rate(1_000, 0), 1); - assert_eq!(streams_for_measured_rate(0, 0), 1); - } - - #[test] - fn test_the_probe_decides_once_and_then_stops_measuring() { - // Deciding repeatedly is what the previous design did, and it ramped - // every device to the cap. One decision on a device-rate measurement - // has no such failure mode -- but only if it really does stop. - let (file, _) = fixture(1024); - let handle = File::open(file.path()).expect("open"); - let mut reader = ScatterReader::new(handle, 0, 1, Some(FetchQueue::new(4))) - .expect("reader") - .auto_tuned(); - for _ in 0..AUTO_PROBE_FILLS { - reader.probe_fill(FILL_BYTES as u64, 12_000_000); - } - assert!(!reader.auto, "the probe must disarm itself"); - assert_eq!(reader.streams, 4, "4 MiB in 12ms is ~350 MB/s, which wants four"); - let settled = reader.streams; - for _ in 0..AUTO_PROBE_FILLS * 4 { - reader.probe_fill(FILL_BYTES as u64, 1); - } - assert_eq!(reader.streams, settled, "a disarmed probe must not revisit its answer"); - } - - #[test] - fn test_readers_sharing_a_queue_share_one_measurement() { - // What is being measured is the device, and a merge reads K spill files - // at once. Measured on EBS gp3: phase 1 alone saw 335 MB/s and chose - // four, while 44 concurrent spill readers each saw 224-229 MB/s of a - // device they were contending for and chose six or seven -- which took - // the merge from 119.3s to 129.1s. A second reader must adopt, not - // re-measure. - let (file, _) = fixture(1024); - let queue = FetchQueue::new(4); - let reader_with = || { - ScatterReader::new( - File::open(file.path()).expect("open"), - 0, - 1, - Some(Arc::clone(&queue)), - ) - .expect("reader") - .auto_tuned() - }; - - let mut first = reader_with(); - for _ in 0..AUTO_PROBE_FILLS { - first.probe_fill(FILL_BYTES as u64, 12_000_000); - } - assert_eq!(first.streams, 4, "the uncontended probe chooses four"); - - // The second reader is handed a *slow* measurement, as a contended one - // would be. It must ignore it and take the settled answer. - let mut second = reader_with(); - second.probe_fill(FILL_BYTES as u64, 900_000_000); - assert_eq!(second.streams, 4, "adopted, not re-measured"); - assert!(!second.auto, "and disarmed on the spot"); - assert_eq!(second.probe_fills, 0, "without accumulating a probe of its own"); - } - - #[test] - fn test_an_incomplete_probe_decides_nothing() { - // Deciding early would let one slow fill choose the count for the run. - let (file, _) = fixture(1024); - let handle = File::open(file.path()).expect("open"); - let mut reader = ScatterReader::new(handle, 0, 1, Some(FetchQueue::new(4))) - .expect("reader") - .auto_tuned(); - for _ in 0..AUTO_PROBE_FILLS - 1 { - reader.probe_fill(FILL_BYTES as u64, 900_000_000); - } - assert_eq!(reader.streams, 1, "a probe one fill short must not decide"); - assert!(reader.auto, "and must still be armed"); - } - - #[test] - fn test_delivers_the_file_verbatim_at_every_stream_count() { - // The contract that matters: splitting a fill must be invisible. A slice - // delivered out of order corrupts every BGZF block or zstd frame the - // sort then tries to parse. - let (file, expected) = fixture(FILL_BYTES * 2 + 98_765); - for streams in [1usize, 2, 4, 8] { - let got = read_all(&mut scattered(file.path(), 0, streams), 64 * 1024); - assert_eq!(got.len(), expected.len(), "length at {streams} streams"); - assert!(got == expected, "bytes differ at {streams} streams"); - } - } - - #[test] - fn test_a_fill_completes_when_no_worker_ever_takes_a_slice() { - // Deadlock-freedom, stated directly. Every worker could be a filler - // waiting on slices nobody is left to run, so a fill that only ever - // *waits* for the pool wedges the sort. It must reclaim instead. - let (file, expected) = fixture(FILL_BYTES * 2); - let queue = FetchQueue::new(4); - let file_handle = File::open(file.path()).expect("open"); - let mut reader = - ScatterReader::new(file_handle, 0, 4, Some(Arc::clone(&queue))).expect("reader"); - - let (tx, rx) = std::sync::mpsc::channel(); - std::thread::spawn(move || drop(tx.send(read_all(&mut reader, 64 * 1024)))); - let got = rx - .recv_timeout(std::time::Duration::from_secs(30)) - .expect("the fill waited on a worker that never came"); - assert!(got == expected, "reclaimed slices must still deliver the file"); - let (offered, taken) = queue.census(); - assert!(offered > 0, "slices were offered"); - assert_eq!(taken, 0, "no worker was running, so none can have been taken"); - } - - #[test] - fn test_exactly_one_caller_ever_runs_a_slice() { - // The claim is a state transition under the fill's mutex, and this is - // what it has to guarantee: a worker and the reclaiming filler racing - // for the same slice must not both read it (a wasted read, and two - // buffers where one is expected) nor both skip it (a fill that never - // completes). Deterministic -- it asserts a count, not who won. - let (file, expected) = fixture(64 * 1024); - let state = Arc::new(FillState { - progress: Mutex::new(FillProgress { - slots: vec![SlotState::Pending(vec![0u8; expected.len()])], - remaining: 1, - error: None, - }), - completed: Condvar::new(), - }); - let slice = Arc::new(FetchSlice { - file: Arc::new(File::open(file.path()).expect("open")), - offset: 0, - index: 0, - state: Arc::clone(&state), - }); - - let racers: Vec<_> = (0..8) - .map(|_| { - let slice = Arc::clone(&slice); - std::thread::spawn(move || usize::from(slice.run())) - }) - .collect(); - let winners: usize = racers.into_iter().map(|h| h.join().expect("racer")).sum(); - assert_eq!(winners, 1, "exactly one caller may read a given slice"); - - let mut progress = state.progress.lock().expect("state"); - assert_eq!(progress.remaining, 0, "the single read completed the fill"); - match std::mem::replace(&mut progress.slots[0], SlotState::InFlight) { - SlotState::Done(buf) => assert!(buf == expected, "the winner read the right bytes"), - _ => panic!("the slice should be done"), - } - } - - #[test] - fn test_a_worker_draining_the_queue_never_corrupts_the_bytes() { - // A live pool worker racing the filler for every slice of every fill. - // Deliberately asserts bytes only: whether the worker wins any given - // race depends on how long slice 0 takes, which is a page-cache hit - // here and a disk read in production. `test_exactly_one_caller...` - // covers the race itself; this covers the result. - let (file, expected) = fixture(FILL_BYTES * 6); - let queue = FetchQueue::new(4); - let stop = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let worker = { - let (queue, stop) = (Arc::clone(&queue), Arc::clone(&stop)); - std::thread::spawn(move || { - while !stop.load(Ordering::Acquire) { - if !queue.run_one() { - std::thread::yield_now(); - } - } - }) - }; - let file_handle = File::open(file.path()).expect("open"); - let mut reader = - ScatterReader::new(file_handle, 0, 4, Some(Arc::clone(&queue))).expect("reader"); - let got = read_all(&mut reader, 64 * 1024); - stop.store(true, Ordering::Release); - worker.join().expect("worker"); - - assert!(got == expected, "bytes must match however the slices were shared"); - } - - /// A reader that keeps one fill in flight, like the phase-1 input reader. - fn ahead(path: &std::path::Path, streams: usize) -> ScatterReader { - let file = File::open(path).expect("open"); - ScatterReader::new(file, 0, streams, Some(FetchQueue::new(4))) - .expect("reader") - .looking_ahead(1) - } - - #[test] - fn test_lookahead_delivers_the_file_verbatim_at_every_stream_count() { - // Prefetching the next fill must not change a byte, and it is the case - // most likely to: two fills are alive at once, so a buffer recycled or - // an offset advanced at the wrong moment corrupts the seam between them. - let (file, expected) = fixture(FILL_BYTES * 3 + 12_345); - for streams in [2usize, 4, 8] { - let got = read_all(&mut ahead(file.path(), streams), 64 * 1024); - assert_eq!(got.len(), expected.len(), "length at {streams} streams"); - assert!(got == expected, "bytes differ at {streams} streams"); - } - } - - #[test] - fn test_lookahead_keeps_the_next_fill_in_flight_while_the_consumer_reads() { - // The whole point: the device should be working on the next fill while - // the framer is still consuming this one. Without it the disk idles for - // exactly as long as framing takes, which is what the demand-driven - // reader measured as 18s of unrecovered read span. - let (file, _) = fixture(FILL_BYTES * 3); - let mut reader = ahead(file.path(), 4); - assert!(reader.pending.is_empty(), "nothing is in flight before the first read"); - assert!(reader.read(&mut [0u8; 64]).expect("first read") > 0, "the fixture is not empty"); - assert_eq!(reader.pending.len(), 1, "the next fill should already be issued"); - } - - #[test] - fn test_lookahead_keeps_the_depth_it_was_given_in_flight() { - // Depth is the knob that decides how much of the framing time the fetch - // can hide behind, so it has to actually take effect. - let (file, _) = fixture(FILL_BYTES * 6); - let handle = File::open(file.path()).expect("open"); - let mut reader = ScatterReader::new(handle, 0, 4, Some(FetchQueue::new(4))) - .expect("reader") - .looking_ahead(3); - assert!(reader.read(&mut [0u8; 64]).expect("read") > 0, "the fixture is not empty"); - assert_eq!(reader.pending.len(), 3, "three fills should be in flight"); - } - - #[test] - fn test_lookahead_does_not_run_ahead_of_the_end_of_the_file() { - // The last fill has no successor. Issuing one anyway would read past EOF - // and turn a clean finish into an error. - let (file, expected) = fixture(FILL_BYTES + 16); - let mut reader = ahead(file.path(), 4); - assert!(read_all(&mut reader, 1024) == expected); - assert!(reader.pending.is_empty(), "nothing may be in flight at EOF"); - } - - #[test] - fn test_a_failed_fill_rewinds_so_the_next_read_retries_the_same_bytes() { - // A fill is issued before it is collected, so the fetch offset moves - // first. If a failure left it moved, the retry would silently skip the - // range that failed -- losing records rather than reporting an error. - let (file, _) = fixture(FILL_BYTES * 3); - let mut reader = ahead(file.path(), 4); - let before = reader.position(); - file.as_file().set_len(0).expect("truncate"); - assert!(reader.read(&mut [0u8; 4096]).is_err(), "the read reports the failure"); - assert_eq!(reader.position(), before, "a failed fill must not advance the reader"); - assert!(reader.read(&mut [0u8; 4096]).is_err(), "and it keeps reporting it"); - } - - #[test] - fn test_position_ignores_bytes_that_are_only_in_flight() { - // `position` is where the consumer is, not how far ahead the fetch has - // run. Counting an in-flight fill would make it jump forward by a whole - // `FILL_BYTES` the moment lookahead was enabled. - let (file, _) = fixture(FILL_BYTES * 3); - let mut reader = ahead(file.path(), 4); - let consumed = reader.read(&mut [0u8; 100]).expect("read"); - assert_eq!(consumed, 100, "a 100-byte read from a full fill returns 100"); - assert_eq!(reader.position(), 100, "only consumed bytes count"); - } - - #[test] - fn test_reading_starts_at_the_requested_offset() { - // A zstd spill's reader is positioned past `ZSPILL_MAGIC` rather than at - // byte 0. Ignoring the offset would feed the frame parser four bytes of - // magic it does not expect. - let (file, expected) = fixture(FILL_BYTES + 4096); - let got = read_all(&mut scattered(file.path(), 4, 4), 32 * 1024); - assert!(got == expected[4..], "did not start at byte 4"); - } - - #[test] - fn test_a_failed_fill_keeps_failing_instead_of_hanging() { - // Fills are demand-driven, so truncating the file after the reader has - // measured it makes the next fill read past EOF deterministically. A - // caller may call `read` again after an error; it must not park. - let (file, _) = fixture(FILL_BYTES * 2); - let mut reader = scattered(file.path(), 0, 4); - file.as_file().set_len(0).expect("truncate"); - assert!(reader.read(&mut [0u8; 4096]).is_err(), "the first read reports the failure"); - assert!(reader.read(&mut [0u8; 4096]).is_err(), "and so does every read after it"); - } - - #[test] - fn test_an_empty_file_reads_as_eof_rather_than_hanging() { - let (file, _) = fixture(0); - assert!(read_all(&mut scattered(file.path(), 0, 4), 4096).is_empty()); - } - - #[test] - fn test_tiny_reads_reassemble_across_slice_boundaries() { - // The framers read a few bytes at a time, so slice hand-offs land - // mid-call-sequence. Getting the advance wrong duplicates or skips bytes - // exactly at the seam between two slices of one fill. - let (file, expected) = fixture(FILL_BYTES + 7); - assert!(read_all(&mut scattered(file.path(), 0, 3), 13) == expected); - } - - #[test] - fn test_one_stream_offers_nothing() { - // `--read-streams 1` must stay exactly as cheap as it is today: no - // slicing, no offer, no barrier -- just the read the caller wanted. - let (file, expected) = fixture(FILL_BYTES * 2); - let queue = FetchQueue::new(4); - let mut reader = ScatterReader::new( - File::open(file.path()).expect("open"), - 0, - 1, - Some(Arc::clone(&queue)), - ) - .expect("reader"); - assert!(read_all(&mut reader, 64 * 1024) == expected); - assert_eq!(queue.census().0, 0, "a single stream offered work to the pool"); - } - - #[test] - fn test_buffers_recycled_through_many_fills_deliver_the_file_verbatim() { - // Seven fills' worth, so every fill after the first reuses buffers the - // consumer handed back. Reusing one at a stale length corrupts the seam - // between one slice and the next, and only once recycling happens. - let (file, expected) = fixture(FILL_BYTES * 6 + 3); - assert!(read_all(&mut scattered(file.path(), 0, 2), 128 * 1024) == expected); - } - - /// Threads in this process, from the kernel's own accounting. - /// - /// Linux-only, which is where CI and every benchmark host run. `nextest` - /// gives each test its own process, so this is a clean count rather than a - /// number shared with a test harness. - #[cfg(target_os = "linux")] - fn thread_count() -> usize { - std::fs::read_dir("/proc/self/task").expect("procfs").count() - } - - #[cfg(target_os = "linux")] - #[test] - fn test_scattered_reading_spawns_no_threads_of_its_own() { - // The whole reason this reader offers slices to the pool instead of - // owning threads: `--threads N` must buy N workers and the ingest - // thread, full stop. The previous design spawned N more per reader, and - // nothing in this suite noticed -- it took a 16 vCPU host running 21 - // threads and double the involuntary context switches to surface it. - let (file, expected) = fixture(FILL_BYTES * 3); - let before = thread_count(); - let mut reader = scattered(file.path(), 0, 8); - assert!(read_all(&mut reader, 64 * 1024) == expected, "and it still reads correctly"); - assert_eq!(thread_count(), before, "the reader spawned threads outside the pool"); - } - - #[test] - fn test_a_scattered_read_matches_the_sequential_one_it_replaces() { - // The two `SpillSource` arms are chosen by a flag, so they have to be - // interchangeable byte for byte -- otherwise the flag changes results, - // not just speed. - let (file, _) = fixture(FILL_BYTES * 2 + 1234); - let mut sequential = SpillSource::Sequential(BufReader::with_capacity( - 2 * 1024 * 1024, - File::open(file.path()).expect("open"), - )); - let mut scatter = SpillSource::Scattered(scattered(file.path(), 0, 4)); - assert!(read_all(&mut sequential, 8192) == read_all(&mut scatter, 8192)); - } -} diff --git a/crates/fgumi-sort/src/sync_spill_writer.rs b/crates/fgumi-sort/src/sync_spill_writer.rs new file mode 100644 index 000000000..2669dbde2 --- /dev/null +++ b/crates/fgumi-sort/src/sync_spill_writer.rs @@ -0,0 +1,790 @@ +//! Synchronous keyed-chunk spill writer (inline compression, no worker pool). +//! +//! The P6 Phase-1 decomposition runs the spill **compress** as a `Parallel` +//! `CompressSpill` step on the framework work-stealing pool, replacing the async +//! [`PooledChunkWriter`] → `SortWorkerPool` path so the sort no longer needs a +//! second, private thread pool. Each `CompressSpill` `try_run` already executes +//! on a framework worker, so the per-chunk write compresses **inline** on that +//! thread. +//! +//! The on-disk format is byte-compatible with [`PooledChunkWriter`], so +//! `SortSpillDecompress` (the Phase-2 streaming reader) and +//! `GenericKeyedChunkReader` (the merge reader) read these files unchanged: +//! +//! - **bgzf** (any `compression`, including `0`): framed BGZF blocks (header + +//! deflate + footer) via [`InlineBgzfCompressor`] then a trailing `BGZF_EOF`, +//! exactly like `PooledChunkWriter`'s bgzf path. Level 0 still produces *framed* +//! (stored) BGZF blocks — NOT raw bytes — because the bgzf reader requires the +//! `0x1f 0x8b` block magic; an unframed raw spill would fail the reader's +//! magic check. +//! - **zstd** (the production default): [`ZSPILL_MAGIC`] then +//! `[u32 LE frame-len][zstd frame]` per ≤`BGZF_MAX_BLOCK_SIZE` raw block, with +//! no trailing marker — mirroring `PooledChunkWriter`'s zstd path and the +//! `ZspillStreamReader` format. +//! +//! In both cases block/frame *boundaries* may differ from the pooled writer, but +//! the decompressed byte stream (and therefore every record read back) is +//! identical, because the reader streams across boundaries. +//! +//! [`PooledChunkWriter`]: crate::pooled_chunk_writer::PooledChunkWriter + +use std::io::{BufWriter, Write}; +use std::marker::PhantomData; +use std::path::{Path, PathBuf}; + +use anyhow::Result; +use fgumi_bgzf::writer::InlineBgzfCompressor; +use fgumi_bgzf::{BGZF_EOF, BGZF_MAX_BLOCK_SIZE}; +use fgumi_raw_bam::RawRecord; +use tempfile::NamedTempFile; +use zstd::bulk::Compressor as ZstdCompressor; + +use crate::codec::{SpillCodec, ZSPILL_MAGIC}; +use crate::keys::RawSortKey; + +/// Open a staging temp file in the same directory as `path`. +/// +/// The spill is written to this temp file and only renamed into place by +/// [`persist_spill`] once `finish()` succeeds, so a crashed or errored write +/// never leaves a partial file at `path` (a `ZspillStreamReader` would otherwise +/// accept a partial file that happens to end on a frame boundary as a valid, +/// shorter spill and silently drop the trailing records). The temp file shares +/// `path`'s directory so the rename is same-filesystem and therefore atomic, and +/// `NamedTempFile` removes it on drop if `finish()` is never reached. +fn stage_spill(path: &Path) -> Result<(BufWriter, PathBuf)> { + let dir = match path.parent() { + Some(parent) if !parent.as_os_str().is_empty() => parent, + _ => Path::new("."), + }; + let temp = NamedTempFile::new_in(dir)?; + Ok((BufWriter::with_capacity(256 * 1024, temp), path.to_path_buf())) +} + +/// Atomically publish a finished spill temp file at `final_path`. +/// +/// Uses `persist_noclobber`, so it fails closed if a file already exists at +/// `final_path` rather than truncating a prior spill — the same "never reuse an +/// existing path" guarantee `SpillWrite::open_file`'s `create_new` provides, +/// preserved now that the destination is opened by rename rather than directly. +fn persist_spill(writer: BufWriter, final_path: &Path) -> Result<()> { + let temp = writer.into_inner().map_err(|e| anyhow::anyhow!("flushing spill temp file: {e}"))?; + temp.persist_noclobber(final_path).map_err(|e| { + anyhow::anyhow!("publishing spill file {}: {}", final_path.display(), e.error) + })?; + Ok(()) +} + +/// Compress and write a fully-sorted in-memory chunk to `path` as a keyed spill +/// file, inline on the calling thread (no `SortWorkerPool`). +/// +/// This is the single-chunk entry point the P6 `CompressSpill` step calls per +/// sorted chunk: it creates a `SyncSpillWriter` for `codec`/`compression`, +/// writes every `(key, record)` pair in `records` order, and closes the file. +/// The on-disk format is byte-compatible with `PooledChunkWriter`, so +/// [`open_spill_slot`](crate::open_spill_slot) → `SortSpillDecompress` / the +/// merge reader consume it unchanged. +/// +/// # Errors +/// +/// Returns an error if the file cannot be created, the zstd compressor cannot be +/// initialized, or any record write/flush fails. +pub fn write_sorted_chunk( + path: &Path, + codec: SpillCodec, + compression: u32, + records: &[(K, RawRecord)], +) -> Result<()> { + let mut writer = SyncSpillWriter::::create(path, codec, compression)?; + for (key, record) in records { + writer.write_record(key, record.as_ref())?; + } + writer.finish() +} + +/// Write an already-sorted [`InMemoryChunk`](crate::InMemoryChunk) to a spill +/// file — the zero-copy analogue of [`write_sorted_chunk`]. The chunk's records +/// share an `Arc` backing store, so this iterates by index and +/// writes each record's bytes directly (no owned `RawRecord`s), avoiding the +/// per-record copy the buffer chain previously paid to materialise +/// `Vec<(K, RawRecord)>`. +/// +/// # Errors +/// +/// Propagates I/O / compression errors from the underlying writer. +pub fn write_sorted_chunk_inmem( + path: &Path, + codec: SpillCodec, + compression: u32, + chunk: &crate::InMemoryChunk, +) -> Result<()> { + let mut writer = SyncSpillWriter::::create(path, codec, compression)?; + for i in 0..chunk.len() { + writer.write_record(chunk.key_at(i), chunk.record_bytes(i))?; + } + writer.finish() +} + +/// Synchronous keyed-chunk spill writer that compresses inline on the calling +/// (framework-worker) thread — no `SortWorkerPool`. +pub(crate) enum SyncSpillWriter { + /// bgzf: framed BGZF blocks (any level, including a stored level-0 block) + + /// trailing `BGZF_EOF`. + Bgzf(BgzfSpillWriter), + /// zstd: inline-compressed `[u32 len][frame]` blocks after `ZSPILL_MAGIC`. + Zstd(ZstdSpillWriter), +} + +impl SyncSpillWriter { + /// Create a writer for `path` using `codec` at `compression` level. + /// + /// For zstd, `compression` is the zstd level (must be ≥ 1 — level 0 is + /// rejected up front by `SortOptions::validate`, since zstd has no + /// uncompressed mode). For bgzf, `compression == 0` writes *framed* stored + /// (uncompressed) BGZF blocks and `> 0` writes deflate-compressed BGZF blocks + /// at that level — in both cases valid, reader-consumable BGZF. + /// + /// # Errors + /// + /// Returns an error if the output file cannot be created or the zstd + /// compressor cannot be initialized. + pub(crate) fn create(path: &Path, codec: SpillCodec, compression: u32) -> Result { + match codec { + SpillCodec::Bgzf => Ok(Self::Bgzf(BgzfSpillWriter::create(path, compression)?)), + SpillCodec::Zstd => Ok(Self::Zstd(ZstdSpillWriter::create(path, compression)?)), + } + } + + /// Write one keyed record in the spill frame format. + /// + /// # Errors + /// + /// Returns an error if key serialization or the underlying write fails. + pub(crate) fn write_record(&mut self, key: &K, record: &[u8]) -> Result<()> { + match self { + Self::Bgzf(w) => w.write_record(key, record), + Self::Zstd(w) => w.write_record(key, record), + } + } + + /// Flush and close the chunk file. + /// + /// # Errors + /// + /// Returns an error if a final flush/compress fails. + pub(crate) fn finish(self) -> Result<()> { + match self { + Self::Bgzf(w) => w.finish(), + Self::Zstd(w) => w.finish(), + } + } +} + +/// Inline BGZF spill writer: framed BGZF blocks (via [`InlineBgzfCompressor`]) +/// followed by a trailing `BGZF_EOF`, matching `PooledChunkWriter`'s bgzf output. +/// Records are framed `[key.write_to() if !EMBEDDED][u32 LE record-len][record]` +/// into the compressor's byte stream, identical to the zstd arm. +pub(crate) struct BgzfSpillWriter { + writer: BufWriter, + /// Final destination the temp file is renamed to by `finish()`. + final_path: PathBuf, + compressor: InlineBgzfCompressor, + /// Reused staging buffer for serialized keys (non-embedded keys only). + /// + /// Kept on the writer rather than allocated per record: this is the + /// queryname and template-coordinate spill path, so a per-record `Vec` + /// costs one malloc/free for every record in the sort. + key_scratch: Vec, + _marker: PhantomData, +} + +impl BgzfSpillWriter { + fn create(path: &Path, level: u32) -> Result { + // Stage into a temp file and rename into place only after `finish()` + // succeeds (see `stage_spill` / `persist_spill`): a partial write never + // reaches `path`, and the final rename still fails closed on a stale + // path, matching `SpillWrite::open_file`. + let (writer, final_path) = stage_spill(path)?; + Ok(Self { + writer, + final_path, + compressor: InlineBgzfCompressor::new(level), + key_scratch: Vec::new(), + _marker: PhantomData, + }) + } + + /// Drain any completed (full-block) compressed output to disk, bounding the + /// compressor's retained-block memory during a large chunk write. + fn drain_blocks(&mut self) -> Result<()> { + for block in self.compressor.take_blocks() { + self.writer.write_all(&block.data)?; + } + Ok(()) + } + + fn write_record(&mut self, key: &K, record: &[u8]) -> Result<()> { + // Validate the record length BEFORE any compressor write, so an oversized + // record fails loud without leaving partial (key) bytes in the stream. + let record_len = u32::try_from(record.len()) + .map_err(|_| anyhow::anyhow!("BAM record too large ({} bytes)", record.len()))?; + if !K::EMBEDDED_IN_RECORD { + // `compressor` and `key_scratch` are disjoint fields, so the scratch + // can be filled and written without moving it out of `self`. + self.key_scratch.clear(); + key.write_to(&mut self.key_scratch)?; + self.compressor.write_all(&self.key_scratch)?; + } + self.compressor.write_all(&record_len.to_le_bytes())?; + // Drain between block-sized chunks so retained compressed-block memory + // stays bounded even for a very large (long-read) record, rather than + // scaling with the record size. + for chunk in record.chunks(BGZF_MAX_BLOCK_SIZE) { + self.compressor.write_all(chunk)?; + self.drain_blocks()?; + } + // Catch the key/length bytes when `record` is empty (the loop is a no-op). + self.drain_blocks() + } + + fn finish(mut self) -> Result<()> { + self.compressor.flush()?; + self.drain_blocks()?; + // BGZF stream terminator (the empty-block EOF marker), as the pooled + // writer and noodles bgzf writer both emit. + self.writer.write_all(&BGZF_EOF)?; + self.writer.flush()?; + persist_spill(self.writer, &self.final_path) + } +} + +/// Inline zstd spill writer: `ZSPILL_MAGIC` then length-prefixed zstd frames, +/// one frame per ≤`BGZF_MAX_BLOCK_SIZE` raw block. +pub(crate) struct ZstdSpillWriter { + writer: BufWriter, + /// Final destination the temp file is renamed to by `finish()`. + final_path: PathBuf, + compressor: ZstdCompressor<'static>, + /// Raw (uncompressed) staging block; flushed as a zstd frame at + /// `BGZF_MAX_BLOCK_SIZE`. + block: Vec, + /// Reused destination for each compressed zstd frame, so `flush_block` does + /// not allocate a fresh `Vec` per frame (`compress` returns an owned one). + frame_buf: Vec, + /// Reused staging buffer for serialized keys — see + /// [`BgzfSpillWriter::key_scratch`]. + key_scratch: Vec, + _marker: PhantomData, +} + +impl ZstdSpillWriter { + fn create(path: &Path, level: u32) -> Result { + // See `BgzfSpillWriter::create`: staged temp file, atomic rename on + // finish, fail-closed (no clobber) on the destination. + let (mut writer, final_path) = stage_spill(path)?; + writer.write_all(&ZSPILL_MAGIC)?; + #[allow(clippy::cast_possible_wrap)] + let compressor = ZstdCompressor::new(level as i32) + .map_err(|e| anyhow::anyhow!("zstd compressor init (level {level}): {e}"))?; + Ok(Self { + writer, + final_path, + compressor, + block: Vec::with_capacity(BGZF_MAX_BLOCK_SIZE + 1024), + frame_buf: Vec::new(), + key_scratch: Vec::new(), + _marker: PhantomData, + }) + } + + /// Compress the staged block as one zstd frame and write `[u32 len][frame]`. + /// No-op when the block is empty (no empty frames in the stream). + fn flush_block(&mut self) -> Result<()> { + if self.block.is_empty() { + return Ok(()); + } + // Reuse `frame_buf` across frames rather than letting `compress` return a + // freshly-allocated `Vec` for every non-empty spill frame. + self.frame_buf.clear(); + self.frame_buf.reserve(zstd::zstd_safe::compress_bound(self.block.len())); + self.compressor + .compress_to_buffer(&self.block, &mut self.frame_buf) + .map_err(|e| anyhow::anyhow!("zstd compress: {e}"))?; + let frame_len = u32::try_from(self.frame_buf.len()) + .map_err(|_| anyhow::anyhow!("zstd frame larger than 4 GiB cannot fit a u32 prefix"))?; + self.writer.write_all(&frame_len.to_le_bytes())?; + self.writer.write_all(&self.frame_buf)?; + self.block.clear(); + Ok(()) + } + + /// Append `data` to the staging block, flushing a frame whenever the block + /// fills. A record (or key) larger than one block therefore spans frames — + /// safe because the reader streams across frame boundaries. + fn append(&mut self, mut data: &[u8]) -> Result<()> { + while !data.is_empty() { + let space = BGZF_MAX_BLOCK_SIZE.saturating_sub(self.block.len()); + let n = data.len().min(space); + self.block.extend_from_slice(&data[..n]); + data = &data[n..]; + if self.block.len() >= BGZF_MAX_BLOCK_SIZE { + self.flush_block()?; + } + } + Ok(()) + } + + fn write_record(&mut self, key: &K, record: &[u8]) -> Result<()> { + // Validate the record length BEFORE any append, so an oversized record + // fails loud without leaving partial (key) bytes in the staging block. + let record_len = u32::try_from(record.len()) + .map_err(|_| anyhow::anyhow!("BAM record too large ({} bytes)", record.len()))?; + if !K::EMBEDDED_IN_RECORD { + // `append` takes `&mut self`, so unlike the bgzf arm the scratch + // cannot stay borrowed from `self` across the call — move it out and + // put it back, which reuses the allocation across records all the + // same. An error here abandons the whole spill, so not restoring the + // buffer on that path costs nothing. + let mut key_bytes = std::mem::take(&mut self.key_scratch); + key_bytes.clear(); + key.write_to(&mut key_bytes)?; + self.append(&key_bytes)?; + self.key_scratch = key_bytes; + } + self.append(&record_len.to_le_bytes())?; + self.append(record)?; + Ok(()) + } + + fn finish(mut self) -> Result<()> { + self.flush_block()?; + self.writer.flush()?; + persist_spill(self.writer, &self.final_path) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::external::GenericKeyedChunkReader; + use crate::inline::TemplateKey; + use crate::pooled_chunk_writer::PooledChunkWriter; + use crate::worker_pool::SortWorkerPool; + use std::sync::Arc; + use tempfile::TempDir; + + #[allow(clippy::cast_possible_truncation)] + fn make_key(i: u64) -> TemplateKey { + TemplateKey::new( + i as i32, + i as i32, + false, + i32::MAX, + i32::MAX, + false, + 0, + 0, + (0, false), + i, + false, + ) + } + + #[allow(clippy::cast_possible_truncation)] + fn sample_records(n: u64) -> Vec<(TemplateKey, Vec)> { + (0..n).map(|i| (make_key(i), vec![(i % 256) as u8; 200 + (i as usize % 50)])).collect() + } + + fn read_back(path: &Path) -> Vec<(TemplateKey, Vec)> { + let mut reader = + GenericKeyedChunkReader::::open(path, None).expect("open reader"); + let mut buf = Vec::new(); + let mut out = Vec::new(); + while let Some(key) = reader.next_record(&mut buf).expect("read record") { + out.push((key, buf.clone())); + } + out + } + + /// A minimal BAM fixed-block body with the coordinate fields set: `ref_id` at + /// 0..4, `pos` at 4..8, `flag` at 14..16 — the three + /// `RawCoordinateKey::extract_from_record` reads. 32 bytes is the fixed block, + /// plus a tail byte to make each record distinguishable. + fn coordinate_body(tid: i32, pos: i32, tail: u8) -> Vec { + let mut b = vec![0u8; 32]; + b[0..4].copy_from_slice(&tid.to_le_bytes()); + b[4..8].copy_from_slice(&pos.to_le_bytes()); + b[14..16].copy_from_slice(&0u16.to_le_bytes()); // flags: mapped, forward + b.push(tail); + b + } + + fn read_back_coordinate(path: &Path) -> Vec> { + let mut reader = GenericKeyedChunkReader::::open(path, None) + .expect("open reader"); + let mut buf = Vec::new(); + let mut out = Vec::new(); + while reader.next_record(&mut buf).expect("read record").is_some() { + out.push(buf.clone()); + } + out + } + + /// The embedded-key writer path, on both codecs. + /// + /// Every other writer test uses `TemplateKey`, whose `EMBEDDED_IN_RECORD` is + /// `false` — so they all take the `if !K::EMBEDDED_IN_RECORD` branch and write + /// a key prefix. The *skip* side of that branch, which the coordinate and + /// queryname sorts use, was never exercised at this layer: those records carry + /// their key inside the body and the spill frame is `[u32 len][record]` with no + /// prefix at all. A writer that emitted a prefix anyway would desync the + /// reader, and no existing test would notice. + #[rstest::rstest] + #[case::zstd(SpillCodec::Zstd)] + #[case::bgzf(SpillCodec::Bgzf)] + fn sync_writers_round_trip_embedded_keys(#[case] codec: SpillCodec) { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("embedded.keyed"); + // Ascending positions so the spill is a sorted run, as production writes. + let bodies: Vec> = + (0..16i32).map(|i| coordinate_body(0, i * 10, u8::try_from(i).unwrap())).collect(); + + let mut w = + SyncSpillWriter::::create(&path, codec, 1).unwrap(); + for b in &bodies { + let key = crate::keys::RawCoordinateKey::extract_from_record(b); + w.write_record(&key, b).unwrap(); + } + w.finish().unwrap(); + + assert_eq!( + read_back_coordinate(&path), + bodies, + "embedded-key records must round-trip with no key prefix ({codec:?})", + ); + } + + /// The sync zstd writer round-trips every record back through the same reader + /// the production merge/Phase-2 path uses, and the file carries `ZSPILL_MAGIC`. + #[test] + fn sync_zstd_round_trips() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("sync_zstd.keyed"); + let records = sample_records(300); + + let mut w = SyncSpillWriter::::create(&path, SpillCodec::Zstd, 3).unwrap(); + for (k, r) in &records { + w.write_record(k, r).unwrap(); + } + w.finish().unwrap(); + + let bytes = std::fs::read(&path).unwrap(); + assert_eq!(&bytes[..ZSPILL_MAGIC.len()], &ZSPILL_MAGIC[..], "missing ZSPILL_MAGIC"); + assert_eq!(read_back(&path), records, "sync zstd round-trip mismatch"); + } + + /// Cross-check: the sync zstd writer and the pooled (async) zstd writer + /// produce files that read back to the *same* records — proving the formats + /// are interchangeable for `SortSpillDecompress` / the merge reader. + #[test] + fn sync_zstd_matches_pooled_zstd_read_back() { + let dir = TempDir::new().unwrap(); + let records = sample_records(500); + + let sync_path = dir.path().join("sync.keyed"); + let mut w = + SyncSpillWriter::::create(&sync_path, SpillCodec::Zstd, 3).unwrap(); + for (k, r) in &records { + w.write_record(k, r).unwrap(); + } + w.finish().unwrap(); + + let pooled_path = dir.path().join("pooled.keyed"); + let pool = Arc::new(SortWorkerPool::new(2, 3, 6, SpillCodec::Zstd)); + { + let mut pw = PooledChunkWriter::::new( + Arc::clone(&pool), + &pooled_path, + SpillCodec::Zstd, + ) + .unwrap(); + for (k, r) in &records { + pw.write_record(k, r).unwrap(); + } + pw.finish().unwrap(); + } + if let Ok(p) = Arc::try_unwrap(pool) { + p.shutdown(); + } + + // Anchor one side to the original `records` first: a comparison of only + // sync-vs-pooled would still pass if BOTH writers dropped every record + // (or the reader returned nothing for both). + let sync_back = read_back(&sync_path); + assert_eq!(sync_back, records, "sync zstd writer must round-trip the original records"); + assert_eq!(sync_back, read_back(&pooled_path), "sync vs pooled read-back differ"); + } + + /// A zero-length record body must round-trip through the writer, on BOTH + /// codecs, without desynchronising the records around it. + /// + /// This is the sibling of the kernel-level + /// `zero_length_record_body_frames_and_keeps_the_stream_aligned`: that one + /// pins `frame_keyed_record_into`, this one pins the streaming writers. + /// + /// The empty body is the degenerate shape of both arms' inner loops — + /// `record.chunks(BGZF_MAX_BLOCK_SIZE)` and the zstd `append` both iterate + /// zero times — so the record reaches the stream as `[key][0u32]` and + /// nothing else, entirely via the surrounding flush paths. The empty record + /// sits between two non-empty ones so a misframed zero length surfaces as + /// its neighbours decoding wrong, rather than as one silently-absent record. + /// + /// Note on scope: the trailing `drain_blocks()` in the bgzf arm is NOT what + /// makes this pass. Removing it keeps every test green, because `finish()` + /// flushes the compressor before draining — that call bounds retained + /// block memory between records, which is what its own comment claims, and + /// is not load-bearing for correctness. + #[rstest::rstest] + #[case::zstd(SpillCodec::Zstd)] + #[case::bgzf(SpillCodec::Bgzf)] + fn sync_writers_round_trip_a_zero_length_record_body(#[case] codec: SpillCodec) { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("empty-body.keyed"); + let records: Vec<(TemplateKey, Vec)> = vec![ + (make_key(1), vec![0xAA; 64]), + (make_key(2), Vec::new()), + (make_key(3), vec![0xBB; 96]), + ]; + + let mut w = SyncSpillWriter::::create(&path, codec, 1).unwrap(); + for (k, r) in &records { + w.write_record(k, r).unwrap(); + } + w.finish().unwrap(); + + assert_eq!( + read_back(&path), + records, + "a zero-length body must round-trip and keep its neighbours aligned ({codec:?})", + ); + } + + /// Records larger than one block must span blocks and still round-trip — + /// on BOTH codecs. + /// + /// The two arms interleave differently and neither can vouch for the other: + /// zstd fills a staging block and flushes a frame when it is full, while + /// bgzf pushes `record.chunks(BGZF_MAX_BLOCK_SIZE)` into the compressor and + /// drains completed blocks between chunks. That chunk/drain interleaving is + /// exactly where a boundary defect would live, and it was untested — the + /// case ran zstd only. + #[rstest::rstest] + #[case::zstd(SpillCodec::Zstd)] + #[case::bgzf(SpillCodec::Bgzf)] + fn sync_writers_span_blocks_for_large_records(#[case] codec: SpillCodec) { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("large.keyed"); + // Two records each ~3x the block size — forces multi-block spanning. + // The `+ 17` keeps the tail from landing on a block boundary. + let big = vec![0xABu8; BGZF_MAX_BLOCK_SIZE * 3 + 17]; + let records = vec![(make_key(1), big.clone()), (make_key(2), big)]; + + let mut w = SyncSpillWriter::::create(&path, codec, 1).unwrap(); + for (k, r) in &records { + w.write_record(k, r).unwrap(); + } + w.finish().unwrap(); + + assert_eq!( + read_back(&path), + records, + "large-record block spanning round-trip mismatch for {codec:?}", + ); + } + + /// Uncompressed bgzf (level 0) must write *framed* (stored) BGZF blocks — + /// starting with the `0x1f 0x8b` gzip magic — NOT raw bytes, or the bgzf + /// reader's magic check rejects the spill. Regression for the + /// `--temp-compression 0 --temp-codec bgzf` path. Also cross-checks that the + /// records read back match the pooled writer at level 0. + #[test] + fn sync_bgzf_level0_is_framed_and_matches_pooled() { + let dir = TempDir::new().unwrap(); + let records = sample_records(300); + + let sync_path = dir.path().join("sync_l0.keyed"); + let mut w = + SyncSpillWriter::::create(&sync_path, SpillCodec::Bgzf, 0).unwrap(); + for (k, r) in &records { + w.write_record(k, r).unwrap(); + } + w.finish().unwrap(); + + let bytes = std::fs::read(&sync_path).unwrap(); + assert_eq!( + &bytes[..2], + &[0x1f, 0x8b], + "level-0 bgzf spill must be framed BGZF (gzip magic), not raw" + ); + assert_eq!(read_back(&sync_path), records, "level-0 bgzf round-trip mismatch"); + + // Cross-check vs the pooled writer at level 0 (its pool is built with + // temp_compression = 0). + let pooled_path = dir.path().join("pooled_l0.keyed"); + let pool = Arc::new(SortWorkerPool::new(2, 0, 6, SpillCodec::Bgzf)); + { + let mut pw = PooledChunkWriter::::new( + Arc::clone(&pool), + &pooled_path, + SpillCodec::Bgzf, + ) + .unwrap(); + for (k, r) in &records { + pw.write_record(k, r).unwrap(); + } + pw.finish().unwrap(); + } + if let Ok(p) = Arc::try_unwrap(pool) { + p.shutdown(); + } + assert_eq!( + read_back(&sync_path), + read_back(&pooled_path), + "level-0 bgzf: sync vs pooled read-back differ" + ); + } + + /// `write_sorted_chunk_inmem` must produce a byte-identical spill file to + /// `write_sorted_chunk` for the same records. + /// + /// It is a separate public entry point that walks the chunk by index + /// (`key_at(i)` / `record_bytes(i)`) instead of iterating owned pairs, and + /// nothing covered it. An off-by-one or a key/record mismatch in that + /// indexing still yields a spill file that parses cleanly, so the merge + /// would emit wrong records rather than fail — the parity assertion is what + /// makes that visible. + #[test] + fn write_sorted_chunk_inmem_matches_write_sorted_chunk() { + let dir = TempDir::new().unwrap(); + let records = sample_records(400); + + let keyed: Vec<(TemplateKey, RawRecord)> = + records.iter().map(|(k, r)| (*k, RawRecord::from(r.clone()))).collect(); + let owned: Vec<(TemplateKey, Vec)> = + records.iter().map(|(k, r)| (*k, r.clone())).collect(); + let chunk = crate::InMemoryChunk::from_owned_records(owned); + + let from_pairs = dir.path().join("pairs.keyed"); + let from_chunk = dir.path().join("chunk.keyed"); + crate::write_sorted_chunk(&from_pairs, SpillCodec::Zstd, 3, &keyed).unwrap(); + crate::write_sorted_chunk_inmem(&from_chunk, SpillCodec::Zstd, 3, &chunk).unwrap(); + + assert_eq!(read_back(&from_chunk), records, "in-mem writer round-trip mismatch"); + assert_eq!( + std::fs::read(&from_chunk).unwrap(), + std::fs::read(&from_pairs).unwrap(), + "the two writers must produce byte-identical spill files", + ); + } + + /// `write_sorted_chunk` (the per-chunk entry the `CompressSpill` step calls) + /// round-trips every record through the merge reader, and `open_spill_slot` + /// opens the result with the requested `file_id` and the detected codec. + #[test] + fn write_sorted_chunk_round_trips_and_open_spill_slot_sets_file_id() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("chunk_0007.keyed"); + let records = sample_records(400); + // `write_sorted_chunk` takes `(K, RawRecord)`; `read_back` yields + // `(K, Vec)`. Build the keyed form once and compare read-back vs the + // original `Vec` payloads. + let keyed: Vec<(TemplateKey, RawRecord)> = + records.iter().map(|(k, r)| (*k, RawRecord::from(r.clone()))).collect(); + + crate::write_sorted_chunk(&path, SpillCodec::Zstd, 3, &keyed).unwrap(); + assert_eq!(read_back(&path), records, "write_sorted_chunk round-trip mismatch"); + + let slot = crate::open_spill_slot(&path, 7).expect("open spill slot"); + assert_eq!(slot.file_id, 7, "open_spill_slot must honor the requested file_id"); + assert_eq!(slot.codec, SpillCodec::Zstd, "codec must be detected from the file magic"); + + // The bgzf arm: codec detection must fall back to Bgzf (no ZSPILL magic). + let bgzf_path = dir.path().join("chunk_0008.keyed"); + crate::write_sorted_chunk(&bgzf_path, SpillCodec::Bgzf, 1, &keyed).unwrap(); + let bgzf_slot = crate::open_spill_slot(&bgzf_path, 8).expect("open bgzf spill slot"); + assert_eq!(bgzf_slot.file_id, 8); + assert_eq!(bgzf_slot.codec, SpillCodec::Bgzf); + } + + /// The bgzf arm frames records itself and drives `InlineBgzfCompressor` + /// directly; confirm the unified type round-trips there too. + #[test] + fn sync_bgzf_round_trips() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("sync_bgzf.keyed"); + let records = sample_records(200); + + let mut w = SyncSpillWriter::::create(&path, SpillCodec::Bgzf, 1).unwrap(); + for (k, r) in &records { + w.write_record(k, r).unwrap(); + } + w.finish().unwrap(); + + assert_eq!(read_back(&path), records, "sync bgzf round-trip mismatch"); + } + + /// A spill file is published atomically: it appears at its final path only + /// after `finish()`, a writer dropped without `finish()` leaves nothing + /// there, and `finish()` refuses to clobber an existing spill (preserving + /// the fail-closed guarantee `create_new` gave), on BOTH codecs. + /// + /// The partial-file hazard this guards is real: `ZspillStreamReader` treats + /// EOF after a complete frame as a clean end, so a partial file ending on a + /// frame boundary would otherwise read back as a valid, shorter spill. + #[rstest::rstest] + #[case::zstd(SpillCodec::Zstd)] + #[case::bgzf(SpillCodec::Bgzf)] + fn spill_is_published_atomically_and_never_clobbers(#[case] codec: SpillCodec) { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("chunk_atomic.keyed"); + + // A writer dropped before `finish()` must leave no file at the final + // path (the temp file is removed on drop). + { + let mut w = SyncSpillWriter::::create(&path, codec, 1).unwrap(); + w.write_record(&make_key(1), &[0xAB; 64]).unwrap(); + // dropped here without `finish()` + } + assert!( + !path.exists(), + "an unfinished spill must not appear at its final path ({codec:?})" + ); + + // `finish()` publishes the file with exactly the written records. + let records = sample_records(24); + let mut w = SyncSpillWriter::::create(&path, codec, 1).unwrap(); + for (k, r) in &records { + w.write_record(k, r).unwrap(); + } + w.finish().unwrap(); + assert!(path.exists(), "a finished spill must appear at its final path ({codec:?})"); + assert_eq!( + read_back(&path), + records, + "published spill must hold the written records ({codec:?})" + ); + + // A second writer targeting the same path fails closed at `finish()` + // rather than truncating the existing spill. + let mut w2 = SyncSpillWriter::::create(&path, codec, 1).unwrap(); + w2.write_record(&make_key(99), &[0xCD; 32]).unwrap(); + assert!( + w2.finish().is_err(), + "finishing onto an existing spill must fail closed, not clobber ({codec:?})", + ); + assert_eq!( + read_back(&path), + records, + "the pre-existing spill must be left intact ({codec:?})" + ); + } +} diff --git a/crates/fgumi-sort/src/template_arena.rs b/crates/fgumi-sort/src/template_arena.rs new file mode 100644 index 000000000..f7cd00d26 --- /dev/null +++ b/crates/fgumi-sort/src/template_arena.rs @@ -0,0 +1,455 @@ +//! Template-coordinate arena-front sort: build a sorted +//! [`InMemoryChunk`] from record bodies already resident in a shared +//! arena, WITHOUT copying record bytes — the template analogue of the coordinate +//! [`coordinate_chunk_from_refs`](crate::ref_sort::coordinate_chunk_from_refs). +//! +//! [`TemplateArenaAccumulator`] encapsulates everything the template key needs +//! that the arena-front pipeline step (`FindBoundariesAndSort` in +//! `fgumi-pipeline-io`) does not have: the [`LibraryLookup`], cell-barcode tag + +//! hasher, the `--key-types` narrowed-lane selection, and the per-record +//! dropped-lane verification. The pipeline step calls [`push`](TemplateArenaAccumulator::push) +//! per record (arena offset + len) and [`seal`](TemplateArenaAccumulator::seal) +//! per run, exactly mirroring the owned [`TemplateChunkSorter`](crate::TemplateChunkSorter) +//! so the output is byte-for-byte identical — only the record bodies stay in the +//! shared arena instead of being copied into owned `RawRecord`s. + +use std::io; +use std::path::Path; +use std::sync::Arc; + +use anyhow::Result; +use noodles::sam::Header; + +use crate::arena_pool::PooledSegmentedBuf; +use crate::external::{ + KeyTypesSpec, LibraryLookup, TemplateKeyVariant, cb_hasher, dropped_lane_error, + extract_template_key_inline, select_template_variant, verify_dropped_lanes, +}; +use crate::inline::{ + CbKey32, InMemoryChunk, TemplateKey, TemplateKey24, TemplateKey40, TemplateLaneKey, + TemplateRecordRef, TertKey32, parallel_radix_sort_template_refs, radix_sort_template_refs, +}; +use crate::ref_sort::PARALLEL_SORT_THRESHOLD; +use crate::{SpillCodec, frame_keyed_record_into, write_sorted_chunk_inmem}; +use fgumi_raw_bam::{RawRecordView, SamTag}; + +/// Variant-carrying erased template residual chunk: one arm per `--key-types` +/// narrowed lane. Lets `MemoryChunkErased::TemplateCoordinate` hold whichever +/// lane variant the sort chose, so template-coordinate rides its natural narrow +/// key end-to-end through merge and spill — exactly like every other sort order +/// rides its own `K` — instead of being pinned to the full 40-byte +/// [`TemplateKey`]. Narrow-lane order equals full-key order (every dropped lane +/// is verified constant on the ingest path), so merging and spilling the narrow +/// key is byte-identical to the full key. The [`K40`](Self::K40) arm is the full +/// key (all lanes), used by the legacy owned path and the full variant. +pub enum TemplateMemChunk { + /// 24-byte core-only lane (neither cb nor tertiary optional word present). + K24(InMemoryChunk), + /// 32-byte lane whose optional word carries `cb_hash`. + Cb32(InMemoryChunk), + /// 32-byte lane whose optional word carries the tertiary (library<<48 | mi). + Tert32(InMemoryChunk), + /// Full 40-byte key (all lanes) — the legacy owned path and full variant. + K40(InMemoryChunk), +} + +/// Run `$body` against the inner `InMemoryChunk` of whichever variant +/// `$self` holds, with the chunk bound to `$chunk`. `$body` is monomorphized per +/// arm, so type inference resolves `K` (e.g. `c.key_at`) from the matched chunk +/// type. Mirrors `chunk_sorter::with_template_buffer!`. +/// +/// Every per-variant dispatch on this enum goes through here: writing the +/// four-arm match out by hand in each method makes a fifth lane variant a +/// six-site edit, and a copy-paste that dispatches the wrong arm still compiles. +macro_rules! with_template_chunk { + ($self:expr, $chunk:ident => $body:expr) => { + match $self { + TemplateMemChunk::K24($chunk) => $body, + TemplateMemChunk::Cb32($chunk) => $body, + TemplateMemChunk::Tert32($chunk) => $body, + TemplateMemChunk::K40($chunk) => $body, + } + }; +} + +impl TemplateMemChunk { + /// Number of records in the chunk. + #[must_use] + pub fn len(&self) -> usize { + with_template_chunk!(self, c => c.len()) + } + + /// `true` iff the chunk holds zero records. + #[must_use] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Total record-payload bytes (sum of record lengths; excludes keys and + /// index overhead). Used for byte-budget accounting at the chunk boundary. + #[must_use] + pub fn payload_bytes(&self) -> usize { + with_template_chunk!(self, c => c.payload_bytes()) + } + + /// Borrow the `i`th record's raw BAM body bytes, in this chunk's sorted order. + /// + /// # Panics + /// + /// Panics if `i >= self.len()`. + #[must_use] + pub fn record_bytes(&self, i: usize) -> &[u8] { + with_template_chunk!(self, c => c.record_bytes(i)) + } + + /// Frame the `i`th record into `out` in the spill layout + /// `[key][u32 LE len][record]`, using this chunk's narrow-lane key. + /// + /// Encapsulates the per-variant key dispatch so the pipeline-io spill + /// serializer stays variant-agnostic. + /// + /// # Errors + /// + /// Returns an error if writing to `out` fails. + pub fn frame_record_into(&self, i: usize, out: &mut Vec) -> io::Result<()> { + with_template_chunk!(self, c => frame_keyed_record_into(out, c.key_at(i), c.record_bytes(i))) + } + + /// Write the whole chunk to a spill file at `path` via + /// [`write_sorted_chunk_inmem`], keyed by this chunk's narrow lane. + /// + /// Encapsulates the per-variant key dispatch so the pipeline-io compress + /// step stays variant-agnostic. + /// + /// # Errors + /// + /// Returns an error if the spill write fails. + pub fn write_spill(&self, path: &Path, codec: SpillCodec, compression: u32) -> Result<()> { + with_template_chunk!(self, c => write_sorted_chunk_inmem(path, codec, compression, c)) + } +} + +/// Build a sorted template chunk from narrow-lane refs pointing into `arena`. +/// +/// Featherweight seal (the template analogue of the coordinate +/// [`coordinate_chunk_from_refs`](crate::ref_sort::coordinate_chunk_from_refs)): +/// sorts `refs` on their cached narrow lane key (stable radix — parallel for +/// large multi-threaded runs, matching the owned path's `par_sort`) and copies +/// that narrow key straight into the returned [`InMemoryChunk`]. NO full-key +/// re-extraction, NO arena body access — the key is already resident in each +/// ref (computed once at `push`). No record bytes are copied: the records +/// reference their bodies in `arena` at `(offset, len)`. Narrow-lane order +/// equals full-key order (every dropped lane is verified constant on the ingest +/// path), so the chunk is correctly ordered for the downstream `MergeDriver`, +/// and spilling/merging the narrow key is byte-identical to the full key. +/// +/// The parallel-vs-serial decision mirrors the coordinate front +/// ([`sort_coordinate_refs`](crate::ref_sort)): it uses the parallel radix only +/// when `sort_threads > 1` AND the run exceeds the shared +/// `PARALLEL_SORT_THRESHOLD`, so a small run does not pay the parallel radix's +/// partition/coordination overhead. Both paths produce byte-identical output +/// (the parallel template radix is stability-tested against the serial one). +#[must_use] +pub fn template_chunk_from_arena_refs( + arena: Arc, + mut refs: Vec>, + sort_threads: usize, +) -> InMemoryChunk { + if sort_threads > 1 && refs.len() >= PARALLEL_SORT_THRESHOLD { + parallel_radix_sort_template_refs(&mut refs); + } else { + radix_sort_template_refs(&mut refs); + } + let records: Vec<(K, u64, u32)> = refs.into_iter().map(|r| (r.key, r.offset, r.len)).collect(); + InMemoryChunk::from_parts(arena, records) +} + +/// Narrow-lane ref accumulator, one arm per `--key-types` variant (mirrors the +/// owned `TemplateBuffer`'s variant set). Each holds `TemplateRecordRef`s +/// pointing at record bodies in the shared inflate arena — no owned byte storage. +enum ArenaRefs { + K24(Vec>), + Cb32(Vec>), + Tert32(Vec>), + K40(Vec>), +} + +/// [`with_template_chunk!`]'s analogue for [`ArenaRefs`]: run `$body` against +/// the inner `Vec>` of whichever variant `$self` holds. +macro_rules! with_arena_refs { + ($self:expr, $refs:ident => $body:expr) => { + match $self { + ArenaRefs::K24($refs) => $body, + ArenaRefs::Cb32($refs) => $body, + ArenaRefs::Tert32($refs) => $body, + ArenaRefs::K40($refs) => $body, + } + }; +} + +impl ArenaRefs { + fn for_variant(v: TemplateKeyVariant) -> Self { + match (v.cb, v.tertiary) { + (false, false) => Self::K24(Vec::new()), + (true, false) => Self::Cb32(Vec::new()), + (false, true) => Self::Tert32(Vec::new()), + (true, true) => Self::K40(Vec::new()), + } + } + + #[inline] + fn push(&mut self, full: &TemplateKey, offset: u64, len: u32) { + // The `key` field type resolves `K` per arm, so `from_full` narrows to the + // matching lane width (identical to the owned `TemplateRecordBuffer::push`). + with_arena_refs!(self, v => v.push(TemplateRecordRef { + key: TemplateLaneKey::from_full(full), + offset, + len, + padding: 0, + })); + } + + fn reserve(&mut self, n: usize) { + with_arena_refs!(self, v => v.reserve(n)); + } +} + +/// State that exists only once the first record has been seen: the chosen +/// narrowed-key variant, the first record's full key (the dropped-lane verify +/// baseline), and the accumulated refs. Set together on the first push and +/// persisted across runs (spills) — matching the owned `TemplateChunkSorter`, +/// whose variant/baseline are chosen once and never re-selected. +struct AccState { + first_key: TemplateKey, + variant: TemplateKeyVariant, + refs: ArenaRefs, +} + +/// Arena-front template-coordinate accumulator: the template analogue of the +/// owned [`TemplateChunkSorter`](crate::TemplateChunkSorter), accumulating +/// arena-pointing refs instead of copying records. Produces byte-identical +/// output (same provisioning, same dropped-lane rejection, same sorted order). +pub struct TemplateArenaAccumulator { + lib_lookup: LibraryLookup, + cell_tag: Option, + cb_hasher: ahash::RandomState, + key_types: KeyTypesSpec, + header_library_varies: bool, + /// Variant + baseline + refs; `None` until the first record provisions it. + state: Option, + /// Reserve hint received before the first record (variant unknown), applied + /// when the ref buffer is provisioned. + pending_reserve: usize, + /// Bounded rayon pool (sized to `sort_threads`) that [`seal`](Self::seal) + /// installs the per-run radix + full-key gather into, so they run on exactly + /// `sort_threads` threads instead of the GLOBAL rayon pool. Without this the + /// parallel radix and `par_iter` gather fan out over every core and + /// oversubscribe the pipeline's own worker pool during a spill (the pipeline + /// runs the sort on one worker while the others inflate/compress). Mirrors + /// the owned `TemplateChunkSorter`'s `rayon_pool.install(...)`. + /// + /// **Shared across worker copies, and it has to be.** `sort_threads` is only + /// known at the first [`seal`](Self::seal), which happens *after* + /// [`fresh`](Self::fresh) has already produced the copies — so a plain + /// `Option`, or even an `Option>`, gives every + /// worker its own pool and puts `workers × sort_threads` threads on the box. + /// That is precisely the oversubscription this field exists to prevent, just + /// reintroduced one level up. `Arc>` lets the copies be made + /// first and the single pool be built once, by whichever worker seals first. + sort_pool: Arc>, +} + +impl TemplateArenaAccumulator { + /// Build an accumulator from the BAM `header`, the sort's cell-barcode tag, + /// and the `--key-types` spec, deriving the library lookup and CB hasher + /// exactly as + /// [`RawExternalSorter::into_template_chunk_sorter`](crate::RawExternalSorter::into_template_chunk_sorter). + #[must_use] + pub fn from_header(header: &Header, cell_tag: Option, key_types: KeyTypesSpec) -> Self { + let lib_lookup = LibraryLookup::from_header(header); + let header_library_varies = lib_lookup.distinct_header_ordinals() > 1; + Self { + lib_lookup, + cell_tag, + cb_hasher: cb_hasher(), + key_types, + header_library_varies, + state: None, + pending_reserve: 0, + sort_pool: Arc::new(std::sync::OnceLock::new()), + } + } + + /// Reserve capacity for approximately `est_records` refs for the current run. + pub fn reserve(&mut self, est_records: usize) { + if let Some(state) = self.state.as_mut() { + state.refs.reserve(est_records); + } else { + self.pending_reserve = self.pending_reserve.max(est_records); + } + } + + /// Extract the template key from `body` (the record's BAM body, `block_size` + /// prefix excluded, at arena offset `body_off`, length `len`), provision the + /// narrowed-lane variant on the first record, verify the dropped lanes on + /// every subsequent record, and accumulate a ref into the arena. + /// + /// # Errors + /// + /// Returns an error if a record carries a dropped-lane value (CB / MI / + /// library) absent from the first record — the same rejection the owned + /// `TemplateChunkSorter::push` performs. + pub fn push(&mut self, body: &[u8], body_off: u64, len: u32) -> Result<()> { + let key = + extract_template_key_inline(body, &self.lib_lookup, self.cell_tag, &self.cb_hasher); + if let Some(state) = self.state.as_mut() { + if let Some(violation) = verify_dropped_lanes(&state.first_key, &key, state.variant) { + let name = RawRecordView::new(body).read_name(); + return Err(dropped_lane_error(&String::from_utf8_lossy(name), violation)); + } + state.refs.push(&key, body_off, len); + } else { + let variant = + select_template_variant(Some(&key), self.key_types, self.header_library_varies); + let mut refs = ArenaRefs::for_variant(variant); + if self.pending_reserve > 0 { + refs.reserve(self.pending_reserve); + } + refs.push(&key, body_off, len); + self.state = Some(AccState { first_key: key, variant, refs }); + } + Ok(()) + } + + /// Sort the refs accumulated for the current run and drain them into an + /// arena-backed [`TemplateMemChunk`] (zero body copies). The chosen variant + + /// baseline are RETAINED for subsequent runs (spills), matching the owned + /// sorter. Empty if nothing was pushed. + /// + /// Featherweight seal: each arm emits the variant-matching narrow + /// [`TemplateMemChunk`] lane using the key already resident in the ref — no + /// per-record re-extraction and no body access — so the downstream + /// merge/spill ride the narrow key lane chosen for this run. + /// + /// # `sort_threads` is per-sort, not per-call + /// + /// The signature takes it per call, which reads as though each call sizes its + /// own execution. It does not: the value builds the shared bounded pool on + /// the FIRST seal across the whole worker fan-out, and every later call — + /// this accumulator's next run, or another worker copy — runs on that pool at + /// its original width, silently ignoring a different value. Pass the same + /// number every time; see `sort_pool`. + /// + /// # Panics + /// + /// Panics if the bounded sort rayon pool cannot be built (an infrastructure + /// failure, e.g. the OS refuses the `sort_threads` worker threads). + #[must_use] + pub fn seal( + &mut self, + arena: Arc, + sort_threads: usize, + ) -> TemplateMemChunk { + // Drain this run's refs, leaving a fresh empty accumulator of the SAME + // variant so the next run re-uses the once-chosen variant + baseline. + // (Scoped so the `self.state` borrow ends before we touch `sort_pool`.) + let refs = { + let Some(state) = self.state.as_mut() else { + return TemplateMemChunk::K40(InMemoryChunk::default()); + }; + std::mem::replace(&mut state.refs, ArenaRefs::for_variant(state.variant)) + }; + let pool = self.bounded_sort_pool(sort_threads); + // Featherweight seal: each arm produces the variant-matching narrow chunk + // (the key is already in the ref — no re-extraction, no body access), and + // tags it with the chosen variant so the downstream merge/spill ride the + // narrow lane. The variant is global, so every run's chunk shares one arm. + pool.install(move || match refs { + ArenaRefs::K24(r) => { + TemplateMemChunk::K24(template_chunk_from_arena_refs(arena, r, sort_threads)) + } + ArenaRefs::Cb32(r) => { + TemplateMemChunk::Cb32(template_chunk_from_arena_refs(arena, r, sort_threads)) + } + ArenaRefs::Tert32(r) => { + TemplateMemChunk::Tert32(template_chunk_from_arena_refs(arena, r, sort_threads)) + } + ArenaRefs::K40(r) => { + TemplateMemChunk::K40(template_chunk_from_arena_refs(arena, r, sort_threads)) + } + }) + } + + /// The shared bounded pool, built on first use at `sort_threads` threads. + /// + /// Running the radix + gather inside `pool.install` bounds BOTH to + /// `sort_threads` (via the pool's `current_num_threads`), so they do not + /// oversubscribe the pipeline's worker pool on a spill. Built once for the + /// whole fan-out — see the `sort_pool` field. + /// + /// **First call wins.** `get_or_init` builds the pool once, so a later call + /// with a different `sort_threads` silently gets the existing pool at the + /// original width. That is deliberate — one pool for the fan-out is the whole + /// point — but it means `sort_threads` is a property of the *sort*, not of an + /// individual call. + fn bounded_sort_pool(&self, sort_threads: usize) -> &rayon::ThreadPool { + self.sort_pool.get_or_init(|| { + rayon::ThreadPoolBuilder::new() + .num_threads(sort_threads.max(1)) + .thread_name(|i| format!("tmpl-sort-{i}")) + .build() + .expect("build bounded template-sort rayon pool") + }) + } + + /// The shared bounded pool, for tests that assert copies share one (the + /// sharing is otherwise unobservable — `seal` returns early before touching + /// the pool when the accumulator was never provisioned). + #[cfg(test)] + pub fn sort_pool_for_test(&mut self, sort_threads: usize) -> &rayon::ThreadPool { + self.bounded_sort_pool(sort_threads) + } + + /// A worker copy: same configuration and provisioning, empty refs. + /// + /// **Provisioning is per-sort, not per-worker, and this carries it across.** + /// The chosen lane variant and the dropped-lane baseline (`first_key`) are + /// selected once from the first record and must then be identical for every + /// worker in the sort. Resetting them here instead would let each `Auto` + /// worker select a lane from *its own* first record, so one sort could emit + /// chunks of two different key widths; would give each worker a different + /// baseline, so a lane constant within every worker but varying across them + /// would pass `verify_dropped_lanes`; and would seal a worker that received + /// no records as [`TemplateMemChunk::K40`] no matter what the others chose. + /// + /// The refs are NOT carried — each worker accumulates its own — so a copy + /// starts empty but already knows which lane it is filling. + /// + /// # Preconditions + /// + /// Call this only **after** the parent has been provisioned (i.e. after its + /// first [`push`](Self::push)). Cloning a worker from an unprovisioned + /// parent yields `state: None`, and that worker will provision itself + /// independently — the exact divergence above. There is no production + /// caller yet; the pipeline step that fans workers out arrives with + /// `fgumi-pipeline-io`, and it owns honouring this. + #[must_use] + pub fn fresh(&self) -> Self { + Self { + lib_lookup: self.lib_lookup.clone(), + cell_tag: self.cell_tag, + cb_hasher: self.cb_hasher.clone(), + key_types: self.key_types, + header_library_varies: self.header_library_varies, + // Carry variant + baseline, not refs: same lane, own accumulation. + state: self.state.as_ref().map(|s| AccState { + first_key: s.first_key, + variant: s.variant, + refs: ArenaRefs::for_variant(s.variant), + }), + pending_reserve: 0, + // Share the holder, not a new one: see the field's doc. + sort_pool: Arc::clone(&self.sort_pool), + } + } +} 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/crates/fgumi-sort/src/worker_pool.rs b/crates/fgumi-sort/src/worker_pool.rs index 2c607c992..63b2809e5 100644 --- a/crates/fgumi-sort/src/worker_pool.rs +++ b/crates/fgumi-sort/src/worker_pool.rs @@ -138,38 +138,17 @@ pub(crate) fn read_length_prefix( Ok(Some(frame_len)) } -/// Cap on the uncompressed size of a zstd spill frame. -/// -/// Derived from the writer's frame size rather than fixed, because a buffer -/// smaller than the largest frame the writer can emit fails to decompress -/// *every* frame at that size -- a total-failure mode, not a slow one. -/// -/// The slack is **proportional** (4x), not a fixed addend, and that is -/// load-bearing rather than tidiness. This buffer is per-worker scratch touched -/// once per decompressed frame -- millions of times in a spill-heavy merge -- so -/// its size is a cache-locality parameter, not just an allocation. A fixed -/// `+ 4 MiB` of slack measured **249.5s against a 199.2s baseline (+25%)** at the -/// default frame size, because it took 8 workers from 2 MiB of hot scratch to -/// 32.5 MiB while peak RSS barely moved. At 4x it evaluates to exactly the 256 -/// KiB this constant held before it was derived, so the default path is -/// unchanged and larger frames scale with it. -/// -/// The 256 KiB floor matters because `BGZF_MAX_BLOCK_SIZE` is 65,280 -- not -/// 65,536 -- so a bare 4x lands 1 KiB *under* the tuned value rather than on it. -#[must_use] -pub(crate) const fn zstd_decomp_cap() -> usize { - let scaled = 4 * crate::bgzf_io::SPILL_FRAME_BYTES; - if scaled > 256 * 1024 { scaled } else { 256 * 1024 } -} - -/// Hard cap on the `u32 LE` length prefix of any zstd spill frame. -/// -/// Scaled from the writer's frame size so the guard can never reject a frame -/// this build is capable of emitting, while still refusing to allocate gigabytes -/// on a corrupt prefix. Compressed frames are smaller than their input in -/// practice, so the doubling is pure slack. -pub(crate) const MAX_ZSTD_FRAME_BYTES: usize = - 2 * 1024 * 1024 + 2 * crate::bgzf_io::SPILL_FRAME_BYTES; +/// Cap on uncompressed size of a zstd spill frame. Production frames are +/// bounded by the staging buffer (`BGZF_MAX_BLOCK_SIZE` + padding ~= 68 KB); +/// this leaves slack but stays small enough that per-frame allocations don't +/// dominate the merge phase when there are many tens of thousands of frames. +pub(crate) const ZSTD_FRAME_DECOMP_CAP: usize = 256 * 1024; + +/// Hard cap on the `u32 LE` length prefix of any zstd spill frame. Frames are +/// produced one per ~64 KiB of input by `compress_job`; even +/// pathological expansion can't reach this. Beyond it, we treat the value as +/// corruption rather than allocate gigabytes. +pub(crate) const MAX_ZSTD_FRAME_BYTES: usize = 2 * 1024 * 1024; /// Maximum zstd compression level recognized by the `zstd` crate. const ZSTD_MAX_CLEVEL: u32 = 22; @@ -265,27 +244,11 @@ pub enum SortStep { Compress = 2, /// Read+decompress one unit of work for some Phase 2 spill file (work-stealing). Phase2FileWork = 3, - /// Extract the sort keys for one batch of already-ingested records. - /// - /// Phase 1 slack work: the ingest thread defers key extraction rather than - /// paying 120 ns/record for it serially, and any worker can pick a batch up. - /// It is scheduled *last* everywhere on purpose — a batch that waits costs - /// nothing until the chunk barrier, whereas displacing `DecompressInput` - /// would starve the very thread this step exists to unblock. - ExtractKeys = 4, - /// Read one byte slice of some reader's in-flight fill. - /// - /// Scheduled *first* everywhere, the mirror image of `ExtractKeys`: a - /// pending slice is what a thread holding an exclusive reader is blocked - /// on, so every moment it waits is a moment nothing downstream is fed. A - /// deferred key batch costs nothing until the chunk barrier; a deferred - /// slice costs the pipeline immediately. - FetchBytes = 5, } impl SortStep { /// Number of distinct sort steps. - pub const COUNT: usize = 6; + pub const COUNT: usize = 4; /// Short label for display. #[must_use] @@ -295,8 +258,6 @@ impl SortStep { Self::DecompressInput => "DecInp", Self::Compress => "Cmprs", Self::Phase2FileWork => "P2File", - Self::ExtractKeys => "ExtKey", - Self::FetchBytes => "Fetch", } } } @@ -323,16 +284,6 @@ pub(crate) struct SortPipelineStats { pub per_thread_step_counts: Box<[[AtomicU64; SortStep::COUNT]; SORT_MAX_THREADS]>, /// Per-thread idle time in nanoseconds (time in backoff/yield). pub per_thread_idle_ns: Box<[AtomicU64; SORT_MAX_THREADS]>, - /// Nanoseconds each worker spent inside a successful step. - /// - /// The existing per-thread array counts steps but not their duration, so a - /// worker doing few expensive steps was indistinguishable from one doing - /// many cheap ones. - pub per_thread_busy_ns: Box<[AtomicU64; SORT_MAX_THREADS]>, - /// Times each worker was the *first* to claim the block the consumer was - /// parked on. Says whether critical-path service is spread across the pool - /// or concentrated on a few workers. - pub per_thread_awaited_claims: Box<[AtomicU64; SORT_MAX_THREADS]>, /// Number of worker threads (for display bounds). pub num_threads: usize, @@ -347,8 +298,6 @@ impl SortPipelineStats { step_count: std::array::from_fn(|_| AtomicU64::new(0)), per_thread_step_counts: new_sort_2d_array(), per_thread_idle_ns: new_sort_1d_array(), - per_thread_busy_ns: new_sort_1d_array(), - per_thread_awaited_claims: new_sort_1d_array(), num_threads, } } @@ -360,14 +309,6 @@ impl SortPipelineStats { self.step_count[step_idx].fetch_add(1, Ordering::Relaxed); if thread_id < SORT_MAX_THREADS { self.per_thread_step_counts[thread_id][step_idx].fetch_add(1, Ordering::Relaxed); - self.per_thread_busy_ns[thread_id].fetch_add(elapsed_ns, Ordering::Relaxed); - } - } - - /// Credit a worker with first-claiming the consumer's awaited block. - pub fn record_awaited_claim(&self, thread_id: usize) { - if thread_id < SORT_MAX_THREADS { - self.per_thread_awaited_claims[thread_id].fetch_add(1, Ordering::Relaxed); } } @@ -390,8 +331,6 @@ impl SortPipelineStats { SortStep::DecompressInput, SortStep::Compress, SortStep::Phase2FileWork, - SortStep::ExtractKeys, - SortStep::FetchBytes, ]; for &step in &all_steps { @@ -556,17 +495,6 @@ pub(crate) struct PermitPool { /// finally surfaced. This never produced wrong output or a hang; it just /// wasted work on the failure path. closed: AtomicBool, - /// Writing one compressed block to the file. - /// - /// Lives here rather than on the pool because the I/O writer thread already - /// holds this `Arc`, and one pool per writer keeps the output writer's stats - /// from being mixed with a spill writer's. - pub(crate) write_dur: crate::merge_trace::DurationHistogram, - /// How long a block sat in the writer's reorder map waiting for an earlier - /// serial to arrive. - pub(crate) write_reorder_wait: crate::merge_trace::DurationHistogram, - /// Blocks held in that map, sampled on each arrival. Read "us" as "blocks". - pub(crate) write_reorder_depth: crate::merge_trace::DurationHistogram, /// Nanoseconds producers spent blocked waiting for a permit, and how many /// waits that was. /// @@ -590,9 +518,6 @@ impl PermitPool { tx: std::sync::Mutex::new(Some(tx)), rx, closed: AtomicBool::new(false), - write_dur: crate::merge_trace::DurationHistogram::default(), - write_reorder_wait: crate::merge_trace::DurationHistogram::default(), - write_reorder_depth: crate::merge_trace::DurationHistogram::default(), blocked_nanos: AtomicU64::new(0), blocked_waits: AtomicU64::new(0), } @@ -639,27 +564,6 @@ impl PermitPool { (secs, self.blocked_waits.load(Ordering::Relaxed)) } - /// Writer-side distributions: per-block write, reorder wait, reorder depth. - /// - /// Read from the pool rather than the writer's staging buffer because the - /// pool outlives `PooledBamWriter::finish`, which consumes the staging: the - /// output drain happens inside `finish`, so a snapshot taken through the - /// still-live writer would omit every block written and every reorder wait - /// incurred during the drain. - pub(crate) fn writer_stats( - &self, - ) -> ( - crate::merge_trace::HistogramReport, - crate::merge_trace::HistogramReport, - crate::merge_trace::HistogramReport, - ) { - ( - self.write_dur.snapshot(), - self.write_reorder_wait.snapshot(), - self.write_reorder_depth.snapshot(), - ) - } - /// Release a permit back to the pool after a block has been written to disk. /// /// A release after `close` is a no-op: the sender is already gone. @@ -710,37 +614,6 @@ pub(crate) const PHASE2_RAW_CAP: usize = 8; /// `ZSTD_FRAME_DECOMP_CAP` (256 KB) for zstd frames. This is a soft cap — the /// "always accept the next-expected serial" rule lets it transiently exceed by /// up to ~`num_workers` blocks per file. -/// **Filling this cap was measured, and it is a large regression. Do not retry.** -/// -/// The pool sits ~0.84 blocks deep on the awaited file while permitted 8, and -/// `decomp-capped` is 2% at t8 -- an obvious invitation to make each claim serve -/// several blocks from the file it already walked to, amortizing the 77.2 -/// fruitless file visits every claim pays. Measured on `1kg-wgs-HG00096`, -/// template-coordinate, 89 spill runs, paired controls in one session: -/// -/// | arm | grab depth | wasted visits/block | merge wall | util | -/// | --- | --- | --- | --- | --- | -/// | t8, one block per claim | 1.00 | 77.6 | 193.8s | 91% | -/// | t8, up to 8 per claim | 6.07 (p50 8) | **10.6** | **263.7s (+36%)** | 66% | -/// | t16, one block per claim | 1.00 | 41.5 | 169.5s | 52% | -/// | t16, up to 8 per claim | 3.16 (p50 2) | 12.6 | 171.8s (+1.4%) | 51% | -/// -/// The target metric moved 7.3x and the sort got 36% slower. **This cap is per -/// file, so depth bought on file X does nothing for a consumer parked on file -/// Y** -- and it consumes the pool capacity that would otherwise have issued Y's -/// disk read. `fully-buffered` scan skips rose 10x (8.6M -> 83.6M): workers -/// scan, find every buffer at cap, and sleep, while the awaited file sits -/// `raw-empty` 98% of the times it is passed over (33% before). Utilization -/// *falls* to 66% -- idle workers and a consumer waiting 158s for one. -/// -/// At t16 the grab could not even deepen (3.16, p50 2) because `decomp-capped` -/// is already 66% there. Raising the cap instead is also measured negative: 8 to -/// 32 cost 1.3%. And `blocks ready on the awaited file at resume` -- the -/// invariant no intervention in this campaign has moved -- stayed at 0.81. -/// -/// The lesson generalizes past this knob: per-file depth is the wrong currency -/// when the consumer needs one specific file's next serial. Breadth is what -/// keeps that file served. pub(crate) const PHASE2_DECOMP_CAP: usize = 8; /// Number of raw blocks to read from disk per `ReadRawBlocks` call. @@ -793,34 +666,13 @@ pub(crate) const NO_AWAITED_SOURCE: usize = usize::MAX; /// | --- | --- | --- | --- | /// | (none: 4 blocks) | 4 | 326.5s | 54% | /// | 1 MiB | 101 | 212.4s | 83% | -/// | 2 MiB | 201 | 197.5s | 89% | -/// | **4 MiB** | 402 | **194.5s** | 90% | -/// -/// Peak RSS was flat across that sweep (4793-4828 MB against a 4794 MB -/// baseline), because the deep allowance is scoped to one file and fewer stalls -/// mean less transient buffering elsewhere. -/// -/// Re-measured later with paired controls in one session, which moved the -/// default here from 2 MiB to 4 MiB and bounded it from above: -/// -/// | threads | 2 MiB | 4 MiB | 32 MiB | -/// | --- | --- | --- | --- | -/// | t8 (89 spill runs) | 199.1s (mean of 5) | **193.5 / 193.8s** | 199.3s, **RSS +58%** | -/// | t16 (44 spill runs) | 178.9s | **169.6s** | not run | +/// | **2 MiB** | 201 | **197.5s** | 89% | +/// | 4 MiB | 402 | 194.5s | 90% | /// -/// So -2.8% at t8 and -5.2% at t16. **Do not raise it further.** At 32 MiB the -/// merge is *slower* than at 4 MiB and peak RSS goes 4793 -> 7562 MB, because -/// `raw-lock` contention climbs with batch depth (29% -> 52% -> 38% of -/// awaited-file skips): more workers collide on one file's `raw_blocks` mutex. -/// A wall-clock-only reading would have called that arm harmless. -/// -/// It helps t16 *more* than t8, and the mechanism is worth knowing because it is -/// not a scheduling improvement. Critical-path worker discovery lag halves -/// (98.5s -> 59.8s) and the share of finds arriving after 320us drops 23% -> -/// 16%: deeper read-ahead does not make recruiting a worker faster, it just -/// needs fewer recruitments. So it pays most where each recruitment is most -/// expensive, which is the regime with the *most* idle capacity. See -/// [`SharedPipelineState::wake_one_worker`] for why recruitment is slow there. +/// 2 MiB is the knee; 4 MiB buys 3s for twice the read-ahead. Peak RSS was flat +/// across the whole sweep (4793-4828 MB against a 4794 MB baseline), because +/// the deep allowance is scoped to one file and fewer stalls mean less +/// transient buffering elsewhere. /// /// The floor read-ahead is working against is total worker busy / threads: /// ~1400 worker-seconds over 8 threads is ~175s, and an uncorrelated merge of @@ -829,13 +681,13 @@ pub(crate) const NO_AWAITED_SOURCE: usize = usize::MAX; /// compression is 68% of the worker total -- or shortening the serial limits /// read-ahead does not touch: the merge consumer's per-record cost, and the /// per-block decompress chain on the file the merge is blocked on. -pub(crate) const PHASE2_STARVING_READ_TARGET_BYTES: u64 = 4 << 20; +pub(crate) const PHASE2_STARVING_READ_TARGET_BYTES: u64 = 2 << 20; /// [`PHASE2_STARVING_READ_TARGET_BYTES`] as a `usize`, for the byte budgets that /// bound the deep FIFO and a single deep read (both of which count `usize` /// bytes). The const assert keeps the two spellings in step, so the target has a /// single source of truth without an `as` cast the pedantic lints reject. -pub(crate) const PHASE2_STARVING_READ_TARGET_BYTES_USIZE: usize = 4 << 20; +pub(crate) const PHASE2_STARVING_READ_TARGET_BYTES_USIZE: usize = 2 << 20; const _: () = assert!(PHASE2_STARVING_READ_TARGET_BYTES_USIZE as u64 == PHASE2_STARVING_READ_TARGET_BYTES); @@ -899,51 +751,24 @@ fn awaited_allowance_for(mean_block_bytes: u64) -> (usize, usize) { /// The file index a Phase 2 scan starts at. /// -/// Three candidates, in order of how well each predicts what the merge needs -/// next: -/// -/// 1. The **awaited source** when starving -- the file the consumer is actually -/// parked on. This is measured demand. -/// 2. The **frontier** when starving -- the lowest source that has not delivered -/// everything. A *proxy* for demand, correct only while sources drain in index -/// order, which is what an input already in the requested order produces. -/// 3. The worker's own round-robin cursor, which spreads workers over the file -/// set and is what keeps a genuinely interleaved merge saturated. -/// -/// The awaited source outranks the frontier because the proxy is wrong exactly -/// where it matters: on a partially-correlated input the merge parks on a source -/// that is not the lowest undrained one, and pointing workers at the frontier -/// sends them to a file nobody is waiting for. Measured on an 89-way merge, 71% -/// of consumer park at 16 threads was spent waiting for *any* worker to claim -/// the needed block while 94% of the other files sat at their buffer cap. +/// The frontier -- the lowest source that has not delivered everything -- when +/// it is starving, and the worker's own round-robin cursor otherwise. Gated on +/// starving rather than applied always: a frontier that is merely *active* is +/// the common case in an interleaved merge, and sending every worker to file 0 +/// there would undo the spread that makes that case fast. /// -/// Both redirects are gated on starving rather than applied always. For the -/// frontier, a merely *active* frontier is the common interleaved case and -/// sending every worker to file 0 would undo the spread. For the awaited source -/// the reason is sharper: it is never cleared between parks, so an ungated -/// version would point every worker at one index permanently, trading a spread -/// problem for a herd problem. +/// `frontier` is only meaningful while it indexes a live file; past the end it +/// names no source and the cursor stands. /// -/// Neither index is meaningful past the end of the file set, where it names no -/// source and the next candidate stands. -/// -/// Pure so every branch of the gate is testable without racing a live pool, +/// Pure so both branches of the gate are testable without racing a live pool, /// following [`classify_scan`](crate::merge_stalls::classify_scan). fn phase2_scan_start( - awaited: usize, - awaited_starving: bool, frontier: usize, num_files: usize, frontier_starving: bool, cursor: usize, ) -> usize { - if awaited < num_files && awaited_starving { - awaited - } else if frontier < num_files && frontier_starving { - frontier - } else { - cursor - } + if frontier < num_files && frontier_starving { frontier } else { cursor } } /// The `(raw_cap, read_batch)` a file may read at. @@ -969,10 +794,8 @@ fn phase2_scan_start( /// /// [`NO_AWAITED_SOURCE`] means the consumer has not parked yet, and matches no /// index. -fn phase2_deserves_deep_read(index: usize, frontier: usize, awaited: usize, next: usize) -> bool { - index == frontier - || (awaited != NO_AWAITED_SOURCE && index == awaited) - || (next != NO_AWAITED_SOURCE && index == next) +fn phase2_deserves_deep_read(index: usize, frontier: usize, awaited: usize) -> bool { + index == frontier || (awaited != NO_AWAITED_SOURCE && index == awaited) } fn phase2_read_allowance(is_frontier: bool) -> (usize, usize) { @@ -1050,7 +873,7 @@ pub(crate) struct TimedBlock { /// Reader state for a single spill file. Locked when reading from disk. pub(crate) struct Phase2Reader { - pub(crate) inner: crate::spill_reader::SpillSource, + pub(crate) inner: BufReader, pub(crate) next_serial: u64, pub(crate) eof: bool, } @@ -1061,6 +884,34 @@ pub(crate) struct Phase2Reader { /// file. The locks here are deliberately fine-grained so different workers can /// be reading, decompressing, and the main thread can be popping records all /// concurrently as long as they touch different sub-states. +/// +/// # Two Phase-2 merge implementations (production vs. the retained oracle) +/// +/// **In this tree THIS pool path is the production sort.** `fgumi sort` and +/// `fgumi merge` both construct a `RawExternalSorter`, so every real sort runs +/// through here. The source branch's note said the opposite — that the buffer +/// chain over `merge_slots.rs` had taken over and this path was reduced to a +/// library entry point and parity oracle — but that end state needs the +/// typed-step `SortMerge` consumer, which lands with `fgumi-pipeline-io` in a +/// later phase. Read the paragraph below as describing why the two Phase-2 +/// implementations differ, not as a claim about which one runs. +/// +/// The two implementations differ deliberately, and that difference is +/// load-bearing for the oracle. This path keeps the `raw_blocks` FIFO, +/// `decomp_in_flight`, the reorder buffer, and the gap-filler because its merge +/// consumer is a parked OS thread (`external.rs::advance_to_next_block`, immune +/// to the v4 framework-Skip deadlock) and its single-*reader* / +/// multi-*decompressor* topology means completion order ≠ pop order. The +/// slimmer `merge_slots.rs` needs only the *inline* subset of that machinery: +/// its `decompressed` queue is a plain `VecDeque` with no gap-filler, because +/// its inline consumer reads AND decompresses in one op so blocks decompress +/// strictly in read order. It still declares `in_flight`, `reader_eof`, and +/// `reorder` — inert on that inline path, live on its block-parallel one — so +/// the difference is the gap-filler and the FIFO discipline, not the fields. +/// Do NOT delete the gap-filler or the reorder buffer +/// here to "match" `merge_slots` — it would reintroduce a real deadlock and/or +/// out-of-order merge output. (History: commits `9d6d7e9` / `9c39dea`, PRs +/// #389 / #395, `docs/design/sort-phase2-unification-deferral.md`.) pub(crate) struct Phase2FileState { /// Disk reader. Held only while popping bytes from disk. pub(crate) reader: Mutex, @@ -1129,7 +980,7 @@ pub(crate) struct Phase2FileState { } impl Phase2FileState { - pub(crate) fn new(reader: crate::spill_reader::SpillSource, codec: SpillCodec) -> Self { + pub(crate) fn new(reader: BufReader, codec: SpillCodec) -> Self { Self { reader: Mutex::new(Phase2Reader { inner: reader, next_serial: 0, eof: false }), reader_eof: AtomicBool::new(false), @@ -1377,12 +1228,6 @@ pub struct SortWorkerPool { pub buffer_pool: BufferPool, num_workers: usize, pub(crate) spill_codec: SpillCodec, - /// How the merge reads each spill file. `Fixed(1)` keeps the sequential - /// `BufReader` a merge has always used. - pub(crate) read_streams: crate::external::ReadStreams, - /// Whether the `--sort-stats` diagnostics are enabled for this run (run-scoped copy of - /// [`crate::RawExternalSorter::sort_stats`]); consulted by this pool's `stat!` emitters. - sort_stats: bool, } /// Shared state visible to all workers and the main thread. @@ -1458,14 +1303,6 @@ pub(crate) struct SharedPipelineState { /// `do_shutdown` checks join results and sets this flag so the main thread /// does not park forever waiting for work that will never arrive. pub(crate) worker_panicked: Arc, - /// What Phase 1's serial ingest thread waited for. Owned here because the - /// input stream records into it and the sorter reports from it, and neither - /// owns the other. - pub(crate) phase1_ingest: Arc, - /// What the exclusively-owned input reader spends its time on. Shared with - /// the `TimedReader` wrapped around the input file, which is the only place - /// the refill syscalls are visible. - pub(crate) reader_stats: Arc, /// Next serial for input block reading (atomic increment for ordering). input_read_serial: AtomicU64, /// Raw input blocks: `ReadInputBlocks` → `DecompressInput`. @@ -1509,23 +1346,6 @@ pub(crate) struct SharedPipelineState { /// Compress jobs: main thread → workers (`ArrayQueue`, non-blocking push). pub(crate) compress_queue: Arc>, - /// Deferred key-extraction batches: ingest thread → workers. - /// - /// Bounded like every other queue here, and a full queue is *not* an error: - /// [`SortWorkerPool::submit_key_job`] hands the batch back and the ingest - /// thread runs it inline. That fallback is what makes the ingest thread's - /// barrier deadlock-free — it never waits on a batch no worker can reach. - pub(crate) key_jobs: Arc>>, - - /// Byte slices offered by whichever reader is filling: reader → workers. - /// - /// Shared by both phases because the job is identical -- read a range at an - /// offset into a waiting buffer. A full queue is not an error and neither is - /// an idle pool: the offering thread reclaims anything nobody started, which - /// is what keeps a fill from waiting on workers that are all themselves - /// filling. See [`crate::spill_reader`]. - pub(crate) fetch_jobs: Arc, - /// Number of workers (for `low_water` threshold in backpressure). num_workers: usize, @@ -1543,77 +1363,9 @@ pub(crate) struct SharedPipelineState { /// slot that is still empty simply cannot be woken yet, which is harmless /// because a worker that has not reached its loop is not sleeping either. worker_threads: Vec>, - /// Per-stage latency distributions and wasted-scan accounting. - pub(crate) stage_latency: crate::merge_phases::StageLatency, - /// Wakes issued by [`Self::wake_one_worker`]. - /// - /// The denominator for the targeting figures above it: a wake that lands on - /// an already-running worker does nothing, so the hit rate is only meaningful - /// against how many wakes were issued at all. - pub(crate) wakes_issued: AtomicU64, - /// When the first worker claimed a block for the awaited source during the - /// consumer's current park, or 0 if none has. Reset by the consumer at park. - pub(crate) awaited_claim_nanos: AtomicU64, - /// When the consumer was last woken, or 0 if not yet during this park. - pub(crate) awaited_publish_nanos: AtomicU64, - /// Where consumer park time goes, split into additive stages. - pub(crate) park_attribution: crate::merge_stalls::ParkAttribution, /// Rotates the target of [`Self::wake_one_worker`] so repeated wakes spread /// across the pool rather than always hitting worker 0. wake_cursor: AtomicUsize, - /// Whether each worker is currently inside `park_timeout`. - /// - /// Written by the worker around its own wait and read by the wake path, so a - /// stale read is possible and harmless -- it can only misclassify a wake in - /// the counters below, never lose one. Set before parking and cleared after, - /// so the window that reads "parked" is a superset of the real one. - worker_parked: Vec, - /// Wakes whose rotating target was not parked, so the unpark was a no-op. - /// - /// The consumer then waits for some *other* worker's backoff to expire, which - /// is bounded by [`MAX_BACKOFF_US`] rather than by wake cost. That is the - /// difference between recruitment costing microseconds and costing a - /// millisecond, and it is invisible in every other counter. - pub(crate) wakes_on_running_worker: AtomicU64, - /// The subset of those where a parked worker *did* exist and was passed over. - /// - /// The recoverable half: targeting a parked worker would have delivered these. - /// A wake with nobody parked at all is not the targeting rule's fault. - pub(crate) wakes_recoverable: AtomicU64, - /// Consumer parks split by what the pool looked like at the instant of the - /// park: a sleeper was available, everyone was busy compressing, or everyone - /// was busy merging. See [`crate::merge_stalls::ParkSupply`]. - pub(crate) park_supply: crate::merge_stalls::ParkSupplyCensus, - /// Parks avoided because the consumer decompressed the block itself. - /// - /// The validity gate for that path: if this stays near zero the consumer never - /// found an unclaimed block and the change is inert, whatever the clock says. - pub(crate) consumer_self_served: AtomicU64, - /// The source the merge expects to consume *after* the current one, or - /// [`NO_AWAITED_SOURCE`]. - /// - /// `phase2_awaited_source` is reactive -- it can only be set once the consumer - /// has already stalled. This is predictive, and the distinction is the whole - /// point: the merge starves at run transitions, where only 2,475 of 167,624 - /// source switches starve but each costs ~20ms against an 8.4ms read latency, - /// because the read had not been *started* when the consumer arrived. Depth - /// cannot fix that -- doubling the read-ahead cap left the starved share of - /// park time at 99% -- but ~300 blocks of advance notice can. - pub(crate) phase2_next_source: AtomicUsize, - /// Predictions actually published to the pool. - /// - /// The validity gate for the predictive read-ahead path. `runner_up()` returns - /// `None` whenever fewer than two sources are active, and a version that - /// returned `None` always would look exactly like a null wall-clock result -- - /// the deep-read path just never fires and nothing says so. A timing gain is - /// only attributable to prediction if this is non-zero. - pub(crate) phase2_predictions: AtomicU64, - /// Set while the merge consumer is parked on a specific spill file. - /// - /// `phase2_awaited_source` cannot serve this purpose: it is never cleared - /// between parks, so keying on it would rank Phase 2 work first permanently - /// and undo the compress-first policy entirely. - pub(crate) consumer_parked: AtomicBool, /// Read batches taken at the deep frontier allowance, and the blocks they /// returned; and the same for batches taken at the uniform allowance. @@ -1659,11 +1411,6 @@ impl SharedPipelineState { fn new(num_workers: usize, main_thread_handle: std::thread::Thread) -> Self { let data_queue_cap = num_workers * 8; let compress_queue_cap = num_workers * 4; - // Deeper than the compress queue: batches are pure CPU with no - // downstream, so a backlog costs only memory, and a shallow queue would - // push the ingest thread onto the inline fallback exactly when the pool - // is busiest — reintroducing the serial cost this step removes. - let key_job_queue_cap = num_workers * 16; Self { merge_phases: crate::merge_phases::MergePhaseCounters::default(), @@ -1683,8 +1430,6 @@ impl SharedPipelineState { decompression_error: Arc::new(AtomicBool::new(false)), chunk_read_error: Arc::new(AtomicBool::new(false)), worker_panicked: Arc::new(AtomicBool::new(false)), - phase1_ingest: Arc::new(crate::phase1_stats::Phase1IngestStats::default()), - reader_stats: Arc::new(crate::phase1_stats::ReaderStats::default()), input_read_serial: AtomicU64::new(0), raw_input_blocks: Arc::new(ArrayQueue::new(data_queue_cap)), decompressed_input: Arc::new(ArrayQueue::new(data_queue_cap)), @@ -1697,32 +1442,17 @@ impl SharedPipelineState { total_sources: AtomicU64::new(0), compress_queue: Arc::new(ArrayQueue::new(compress_queue_cap)), - key_jobs: Arc::new(ArrayQueue::new(key_job_queue_cap)), - fetch_jobs: crate::spill_reader::FetchQueue::new(num_workers), num_workers, main_thread_handle, worker_threads: (0..num_workers).map(|_| std::sync::OnceLock::new()).collect(), wake_cursor: AtomicUsize::new(0), - worker_parked: (0..num_workers).map(|_| AtomicBool::new(false)).collect(), - wakes_on_running_worker: AtomicU64::new(0), - wakes_recoverable: AtomicU64::new(0), - park_supply: crate::merge_stalls::ParkSupplyCensus::default(), - phase2_next_source: AtomicUsize::new(NO_AWAITED_SOURCE), - phase2_predictions: AtomicU64::new(0), - consumer_self_served: AtomicU64::new(0), - consumer_parked: AtomicBool::new(false), deep_read_batches: AtomicU64::new(0), deep_read_blocks: AtomicU64::new(0), shallow_read_batches: AtomicU64::new(0), shallow_read_blocks: AtomicU64::new(0), phase2_lowest_active: AtomicUsize::new(0), phase2_awaited_source: AtomicUsize::new(NO_AWAITED_SOURCE), - stage_latency: crate::merge_phases::StageLatency::default(), - wakes_issued: AtomicU64::new(0), - awaited_claim_nanos: AtomicU64::new(0), - awaited_publish_nanos: AtomicU64::new(0), - park_attribution: crate::merge_stalls::ParkAttribution::default(), awaited_skips: std::array::from_fn(|_| AtomicU64::new(0)), phase2_read_bytes: AtomicU64::new(0), phase2_read_blocks: AtomicU64::new(0), @@ -1736,26 +1466,6 @@ impl SharedPipelineState { } } - /// Wake the merge consumer, stamping when it became runnable. - /// - /// Every consumer wake goes through here so the park decomposition cannot - /// silently miss one of the six unpark sites. The stamp is the moment the - /// consumer *could* run, so the gap to it actually resuming is its own wake - /// latency rather than fetch time. - /// - /// Only the first wake in a park is stamped (compare-exchange from 0): a - /// park ends on the first wake, and later ones belong to the next park. - pub(crate) fn wake_consumer(&self) { - let now = self.now_nanos(); - let _ = self.awaited_publish_nanos.compare_exchange( - 0, - now, - Ordering::Relaxed, - Ordering::Relaxed, - ); - self.main_thread_handle.unpark(); - } - /// Wake one idle worker, rotating which one. /// /// The consumer calls this the instant a reorder buffer drains. Without it @@ -1773,63 +1483,16 @@ impl SharedPipelineState { if self.worker_threads.is_empty() { return; } - self.wakes_issued.fetch_add(1, Ordering::Relaxed); - let cursor = self.wake_cursor.fetch_add(1, Ordering::Relaxed); - let limit = self.active_worker_limit.load(Ordering::Acquire); - let idx = - Self::wake_target_preferring_parked(cursor, limit, self.worker_threads.len(), |i| { - self.worker_parked[i].load(Ordering::Relaxed) - }); - // The same accounting now doubles as this preference's validity gate: if - // it works, "hit an already-running worker" must fall and "recoverable" - // must go to roughly zero, because a passed-over sleeper is exactly what - // the preference removes. A `recoverable` share that stays high means the - // parked flags are going stale between the scan and the unpark. - if !self.worker_parked[idx].load(Ordering::Relaxed) { - self.wakes_on_running_worker.fetch_add(1, Ordering::Relaxed); - let width = limit.min(self.worker_parked.len()); - if crate::merge_stalls::first_parked_from(cursor, width, |i| { - self.worker_parked[i].load(Ordering::Relaxed) - }) - .is_some() - { - self.wakes_recoverable.fetch_add(1, Ordering::Relaxed); - } - } + let idx = Self::wake_target( + self.wake_cursor.fetch_add(1, Ordering::Relaxed), + self.active_worker_limit.load(Ordering::Acquire), + self.worker_threads.len(), + ); if let Some(handle) = self.worker_threads[idx].get() { handle.unpark(); } } - /// Publish the source the merge expects to consume next, so the pool can start - /// its read before the consumer gets there. - pub(crate) fn set_phase2_next_source(&self, next: Option) { - if next.is_some() { - self.phase2_predictions.fetch_add(1, Ordering::Relaxed); - } - self.phase2_next_source.store(next.unwrap_or(NO_AWAITED_SOURCE), Ordering::Relaxed); - } - - /// Predictions published to the pool. See [`Self::phase2_predictions`]. - pub(crate) fn phase2_predictions(&self) -> u64 { - self.phase2_predictions.load(Ordering::Relaxed) - } - - /// How many workers are parked right now, and whether output compression is - /// queued -- the two facts that say why nobody had already started on the - /// block the consumer is about to wait for. - /// - /// Scans the parked flags rather than keeping a running count: it runs once - /// per consumer park against at most `SORT_MAX_THREADS` relaxed loads, and a - /// counter would need an atomic add on both sides of every worker's wait -- - /// the far hotter path. - pub(crate) fn park_supply_now(&self) -> crate::merge_stalls::ParkSupply { - let limit = self.active_worker_limit.load(Ordering::Acquire); - let width = limit.min(self.worker_parked.len()); - let parked = (0..width).filter(|&i| self.worker_parked[i].load(Ordering::Relaxed)).count(); - crate::merge_stalls::classify_park_supply(parked, self.compress_queue.len()) - } - /// Which worker slot the next wake should target. /// /// Rotates over the *active* workers, not the pool width. A pool sized to a @@ -1847,33 +1510,6 @@ impl SharedPipelineState { cursor % active_limit.clamp(1, pool_width.max(1)) } - /// The worker a wake should target, preferring one that is actually parked. - /// - /// The rotating cursor alone is blind to park state, and at low utilization - /// that is where the merge's idle time comes from: a wake spent on a running - /// worker leaves an available sleeper asleep for up to `MAX_BACKOFF_US`. - /// Measured at t16, parks where a sleeper existed cost 190us against 30us - /// when every worker was genuinely busy -- 391,007 of them, 74.4s of a 92.1s - /// park. At t8, where workers are 91% busy and essentially never sleep, that - /// class is 0.1% of parks and this preference has nothing to act on. - /// - /// Falls back to the rotating target when nobody is parked, which keeps the - /// spread that a genuinely interleaved merge needs. The scan still starts at - /// the cursor, so it rotates too rather than concentrating on low indices. - fn wake_target_preferring_parked( - cursor: usize, - active_limit: usize, - pool_width: usize, - is_parked: F, - ) -> usize - where - F: Fn(usize) -> bool, - { - let width = active_limit.clamp(1, pool_width.max(1)); - crate::merge_stalls::first_parked_from(cursor, width, is_parked) - .unwrap_or_else(|| Self::wake_target(cursor, active_limit, pool_width)) - } - /// Advance the drain frontier past every source that is now fully drained. /// /// Called by the merge consumer when a source reports drained. Walks rather @@ -1911,40 +1547,6 @@ impl SharedPipelineState { compress_has_items: !self.compress_queue.is_empty(), phase: current_phase, - consumer_parked: self.consumer_parked.load(Ordering::Relaxed), - } - } -} - -/// The decompression state one thread needs to turn a spill block into records. -/// -/// Extracted from [`SortWorkerState`] so the merge consumer can run -/// [`SortWorkerPool::decompress_and_publish`] itself instead of parking. A second -/// copy of the codec branch would be the alternative, and a spill format that -/// disagreed between the two paths is exactly the class of bug that passes the -/// standard matrix and fails where spill volume is largest. -pub(crate) struct DecompressorSet { - zstd: ZstdDecompressor<'static>, - /// Scratch for one zstd frame, sized by [`zstd_decomp_cap`]. Allocated lazily: - /// a BGZF-only sort must not pay for it. - zstd_buf: Vec, - bgzf: libdeflater::Decompressor, -} - -impl DecompressorSet { - /// Sized for `codec`: the zstd scratch is pre-allocated only for zstd spills, - /// so a BGZF-only sort does not carry `zstd_decomp_cap()` per thread. The - /// decompress path still resizes on demand, so this is an optimization rather - /// than a precondition. - pub(crate) fn for_codec(codec: SpillCodec) -> Self { - let zstd_buf = match codec { - SpillCodec::Zstd => vec![0u8; zstd_decomp_cap()], - SpillCodec::Bgzf => Vec::new(), - }; - Self { - zstd: ZstdDecompressor::new().expect("zstd decompressor init"), - zstd_buf, - bgzf: libdeflater::Decompressor::new(), } } } @@ -1964,8 +1566,12 @@ struct SortWorkerState { output_compressor: InlineBgzfCompressor, /// Zstd compressor reused across spill frames when `SpillCodec::Zstd`. zstd_compressor: ZstdCompressor<'static>, - /// Decompression state, shared in shape with the merge consumer. - decomp: DecompressorSet, + /// Zstd decompressor reused across Phase 2 frames when `SpillCodec::Zstd`. + zstd_decompressor: ZstdDecompressor<'static>, + /// Scratch buffer reused across zstd frame decompressions to avoid + /// allocating a fresh Vec for every frame on the merge hot path. + zstd_decompress_buf: Vec, + decompressor: libdeflater::Decompressor, /// Phase 2 file scan cursor — starts at `worker_id` and advances on success /// for cache locality and reduced lock contention. Workers no longer own a /// fixed subset of files; any worker can do work on any file. @@ -1986,23 +1592,6 @@ struct SortWorkerState { /// cleared) by the next iteration that finds work, which is what makes that /// wait "productive" — see [`crate::merge_stalls::WakeLatencyStats`]. last_wait: Option, - /// Whether the step this iteration completed produced for the file the merge - /// consumer was parked on. - /// - /// Set and read at the two Phase-2 claim sites, where the file index is in - /// hand: it gates the `stamp_awaited_claim` call, so only a worker that - /// served the awaited file can credit the first claim. The - /// productive-wait path reads `won_awaited_claim` instead, not this field. - /// Discovery lag is only paid in wall clock on the awaited file — a - /// late wake for any other file is absorbed by idle capacity — so without - /// this the aggregate lag cannot be told apart from free lag. Cleared at the - /// top of every iteration so a stale `true` cannot survive into a step that - /// served a different file. - served_awaited: bool, - /// Whether this worker won the race to first-claim the consumer's awaited - /// block this iteration. Credited per thread by the main loop, which is - /// where the stats handle is in scope. - won_awaited_claim: bool, } impl SortWorkerState { @@ -2074,15 +1663,6 @@ struct SortBackpressureState { // Shared compress_has_items: bool, phase: u8, - /// Whether the merge consumer is parked right now, waiting on one specific - /// spill file. - /// - /// Output compression is throughput work any worker can do at any time; the - /// awaited block is the only thing that can unblock the merge. Measured at - /// t16 before this existed: only 10.8% of recruitments served the consumer as - /// the woken worker's first action, the rest completing a mean 3.5 other - /// steps first, ~2 of them output compress at 188us each. - consumer_parked: bool, } /// Backpressure-driven priority selection — the sort pipeline's equivalent @@ -2099,29 +1679,15 @@ fn get_sort_priorities(bp: &SortBackpressureState) -> &'static [SortStep] { match bp.phase { phase::PHASE1 => { if bp.input_eof && !bp.compress_has_items && bp.decompressed_input_done { - // Input fully decompressed and no compress work — only leftover - // key batches remain, and eligibility drops the step when there - // are none. - &[SortStep::FetchBytes, SortStep::ExtractKeys] + // Input fully decompressed and no compress work — nothing productive to do + &[] } else if bp.compress_has_items && !bp.decompressed_input_low { // Spill compression is the bottleneck (13.7s at t4). Drain compress // while decompressed blocks are plentiful for the main thread. - &[ - SortStep::FetchBytes, - SortStep::Compress, - SortStep::DecompressInput, - SortStep::ReadInputBlocks, - SortStep::ExtractKeys, - ] + &[SortStep::Compress, SortStep::DecompressInput, SortStep::ReadInputBlocks] } else { // Default/starving: feed the main thread first, compress if available - &[ - SortStep::FetchBytes, - SortStep::DecompressInput, - SortStep::ReadInputBlocks, - SortStep::Compress, - SortStep::ExtractKeys, - ] + &[SortStep::DecompressInput, SortStep::ReadInputBlocks, SortStep::Compress] } } phase::PHASE2 => { @@ -2129,37 +1695,16 @@ fn get_sort_priorities(bp: &SortBackpressureState) -> &'static [SortStep] { // worker grabs whatever has work. We never gate on `all_chunks_eof` // — even after disk reads finish, decompression and parser drain // continue until all per-file reorder buffers empty. - if bp.consumer_parked { - // The merge is blocked on one specific file. Output compression - // is throughput work any worker can do at any time; this block is - // the only thing that can unblock the consumer, so it goes first - // even though the compress queue is the writer-side bottleneck. - &[ - SortStep::FetchBytes, - SortStep::Phase2FileWork, - SortStep::Compress, - SortStep::ExtractKeys, - ] - } else if bp.compress_has_items { + if bp.compress_has_items { // Drain output compression while we can; it's the writer-side bottleneck. - &[ - SortStep::FetchBytes, - SortStep::Compress, - SortStep::Phase2FileWork, - SortStep::ExtractKeys, - ] + &[SortStep::Compress, SortStep::Phase2FileWork] } else { - &[ - SortStep::FetchBytes, - SortStep::Phase2FileWork, - SortStep::Compress, - SortStep::ExtractKeys, - ] + &[SortStep::Phase2FileWork, SortStep::Compress] } } // Legacy/transition: compress only (drain any remaining jobs, of either // kind — each carries its own `CompressTarget`). - _ => &[SortStep::FetchBytes, SortStep::Compress, SortStep::ExtractKeys], + _ => &[SortStep::Compress], } } @@ -2168,37 +1713,6 @@ fn get_sort_priorities(bp: &SortBackpressureState) -> &'static [SortStep] { // ============================================================================ /// Minimum backoff duration in microseconds. -/// Where the consumer's wakes went, and what recruited the worker that served -/// the critical path. -/// -/// **These numbers do not explain the merge's idle time. Measured, both regimes.** -/// A wake aimed at an already-running worker does nothing, which looks like a -/// defect worth fixing until both thread counts are measured: -/// -/// | | t8 (fast) | t16 (slow) | -/// | --- | --- | --- | -/// | hit an already-running worker | **100%** | 88% | -/// | a parked worker was available | **0%** | 27% | -/// | consumer's wait for a worker | **10us** | 72us | -/// -/// t8 scores worse on every count and waits 7x less, because at 91% utilization -/// there is nobody parked to hit and a running worker reaches the awaited file -/// almost at once. So "wakes land on busy workers" is what a *healthy* pool looks -/// like, and targeting parked workers is not the fix for t16. -/// -/// Kept because they are cheap and because rediscovering this costs a machine -/// day. What does separate the regimes is -/// [`crate::merge_stalls::ParkSupply`]. -#[derive(Debug, Clone, Copy)] -pub struct WakeAccounting { - /// Wakes issued. - pub issued: u64, - /// Wakes whose target was not parked. - pub on_running: u64, - /// The subset of `on_running` where a parked worker existed and was skipped. - pub recoverable: u64, -} - pub(crate) const MIN_BACKOFF_US: u64 = 10; /// Maximum backoff duration in microseconds (1ms). /// @@ -2220,7 +1734,6 @@ pub(crate) const MAX_BACKOFF_US: u64 = 1000; /// A pending unpark token makes the next `park_timeout` return at once. That is /// a spurious early wake, which costs one scan and is harmless — the caller /// loops. -/// fn idle_wait_with_jitter(backoff_us: u64, worker_id: usize, iter: u64) { if backoff_us <= MIN_BACKOFF_US { std::thread::yield_now(); @@ -2354,7 +1867,7 @@ impl Drop for WorkerPanicGuard { } // Reached only when the worker loop unwound. self.shared.worker_panicked.store(true, Ordering::Release); - self.shared.wake_consumer(); + self.shared.main_thread_handle.unpark(); } } @@ -2367,15 +1880,12 @@ impl SortWorkerPool { /// - `temp_compression`: BGZF level for Phase 1 spill writes (typically 1 for speed). /// - `output_compression`: BGZF level for Phase 2 merge output (typically 6 for size). /// - `spill_codec`: codec used for spill chunks (BGZF or Zstd). Output is always BGZF. - /// - `sort_stats`: whether `--sort-stats` diagnostics are on for this run (the trailing - /// `bool` at call sites; tests pass `false`). #[must_use] pub fn new( num_workers: usize, temp_compression: u32, output_compression: u32, spill_codec: SpillCodec, - sort_stats: bool, ) -> Self { let buffer_pool = BufferPool::new(num_workers * 4); let stats = PoolStats::default(); @@ -2391,19 +1901,24 @@ impl SortWorkerPool { thread::spawn(move || { let zstd_level = i32::try_from(temp_compression.clamp(1, ZSTD_MAX_CLEVEL)).expect("clamped"); + let zstd_decompress_buf = if matches!(spill_codec, SpillCodec::Zstd) { + vec![0u8; ZSTD_FRAME_DECOMP_CAP] + } else { + Vec::new() + }; let mut worker = SortWorkerState { worker_id, compressor: InlineBgzfCompressor::new(temp_compression), output_compressor: InlineBgzfCompressor::new(output_compression), zstd_compressor: ZstdCompressor::new(zstd_level) .expect("zstd compressor init"), - decomp: DecompressorSet::for_codec(spill_codec), + zstd_decompressor: ZstdDecompressor::new().expect("zstd decompressor init"), + zstd_decompress_buf, + decompressor: libdeflater::Decompressor::new(), phase2_file_cursor: worker_id, held_raw_input_blocks: Vec::new(), held_decompressed_input: None, backoff_us: MIN_BACKOFF_US, - served_awaited: false, - won_awaited_claim: false, idle_iter: 0, last_wait: None, }; @@ -2429,34 +1944,9 @@ impl SortWorkerPool { buffer_pool, num_workers, spill_codec, - read_streams: crate::external::ReadStreams::Fixed(1), - sort_stats, } } - /// Whether the `--sort-stats` diagnostics are enabled for the run that built this pool. - pub(crate) fn sort_stats(&self) -> bool { - self.sort_stats - } - - /// Park for the current backoff, publish the parked state, and return the - /// elapsed park in nanoseconds. - /// - /// `worker_parked` is set before the wait and cleared after, so the window - /// the wake path reads as "parked" is a superset of the real one: a wake is - /// never withheld from a worker that is about to park. The reverse error -- - /// reading "running" for a worker that has just parked -- would lose a wake, - /// so the asymmetry is deliberate. - fn park_and_measure(shared: &SharedPipelineState, worker: &mut SortWorkerState) -> u64 { - let idle_start = Instant::now(); - shared.worker_parked[worker.worker_id].store(true, Ordering::Relaxed); - idle_wait_with_jitter(worker.backoff_us, worker.worker_id, worker.idle_iter); - shared.worker_parked[worker.worker_id].store(false, Ordering::Relaxed); - worker.idle_iter = worker.idle_iter.wrapping_add(1); - worker.backoff_us = (worker.backoff_us * 2).min(MAX_BACKOFF_US); - Self::nanos_u64(idle_start.elapsed()) - } - // ======================================================================== // Worker Loop — modeled on generic_worker_loop (base.rs:4379) // ======================================================================== @@ -2519,8 +2009,6 @@ impl SortWorkerPool { } let mut did_work = false; - worker.served_awaited = false; - worker.won_awaited_claim = false; // 3. Try to advance ALL held items first (deadlock prevention) did_work |= Self::try_advance_all_held(shared, worker); @@ -2602,13 +2090,14 @@ impl SortWorkerPool { .wake_latency .record_productive_sleep(wake_phase(current_phase), waited_ns); } - if worker.won_awaited_claim { - pstats.record_awaited_claim(worker.worker_id); - } worker.backoff_us = MIN_BACKOFF_US; } else { let slept_us = worker.backoff_us; - let idle_ns = Self::park_and_measure(shared, worker); + let idle_start = Instant::now(); + idle_wait_with_jitter(slept_us, worker.worker_id, worker.idle_iter); + worker.idle_iter = worker.idle_iter.wrapping_add(1); + worker.backoff_us = (worker.backoff_us * 2).min(MAX_BACKOFF_US); + let idle_ns = Self::nanos_u64(idle_start.elapsed()); pstats.record_idle(worker.worker_id, idle_ns); shared.wake_latency.record_sleep(wake_phase(current_phase), slept_us, idle_ns); worker.last_wait = Some(PendingWait { phase: current_phase, nanos: idle_ns }); @@ -2624,12 +2113,6 @@ impl SortWorkerPool { phase::PHASE1 => { shared.decompressed_input_done.load(Ordering::Acquire) && shared.compress_queue.is_empty() - // Deferred key batches outlive the input: the last chunk's - // barrier runs after the input is fully decompressed, so - // without this the pool parks precisely when the ingest - // thread is waiting on it and the final chunk's extraction - // collapses back to serial. - && shared.key_jobs.is_empty() } phase::PHASE2 => { if !shared.all_chunks_eof.load(Ordering::Acquire) { @@ -2674,19 +2157,6 @@ impl SortWorkerPool { } } - /// Whether a worker owning `owned_step` may also take key-extraction batches. - /// - /// The worker that exclusively owns [`SortStep::ReadInputBlocks`] may not. - /// A batch is ~0.5 ms of pure CPU, and for that long its worker is not - /// reading — and nobody else can read, because the step is exclusive. On the - /// production cell, letting the reader's owner take batches cost the reader - /// 358 → 345 MB/s and 125.2s → 129.8s, against a phase whose whole remaining - /// cost *is* the reader. Every other worker is free to extract, so this - /// gives up 1/16th of the extraction capacity to protect the critical path. - fn worker_may_extract_keys(owned_step: Option) -> bool { - owned_step != Some(SortStep::ReadInputBlocks) - } - /// Whether a step is exclusive (requires ownership). /// /// Only `ReadInputBlocks` is exclusive — it reads from a shared input file @@ -2739,21 +2209,6 @@ impl SortWorkerPool { } SortStep::Compress => !shared.compress_queue.is_empty(), SortStep::Phase2FileWork => current_phase == phase::PHASE2, - // Deliberately not gated on the phase: a batch queued late in - // Phase 1 must still be reachable while the ingest thread drains at - // the chunk barrier, which can coincide with the phase flip. - // Not gated on the phase either: both phases fill through the - // same queue, and a slice offered as Phase 1 drains must stay - // reachable across the flip. - SortStep::FetchBytes => !shared.fetch_jobs.is_empty(), - SortStep::ExtractKeys => { - !shared.key_jobs.is_empty() - && Self::worker_may_extract_keys(Self::exclusive_step_for( - worker.worker_id, - shared, - current_phase, - )) - } } } @@ -2768,32 +2223,9 @@ impl SortWorkerPool { SortStep::DecompressInput => Self::try_decompress_input(shared, worker), SortStep::Compress => Self::try_compress(shared, worker), SortStep::Phase2FileWork => Self::try_phase2_file_work(shared, worker), - SortStep::ExtractKeys => Self::try_extract_keys(shared), - SortStep::FetchBytes => Self::try_fetch_bytes(shared), } } - /// Read one offered byte slice, if any is still unclaimed. - /// - /// `InputEmpty` covers both "nothing queued" and "the offering thread - /// reclaimed this one first" -- neither did work, and reporting progress - /// for a no-op would keep a worker spinning on a queue of stale slices. - fn try_fetch_bytes(shared: &SharedPipelineState) -> StepResult { - if shared.fetch_jobs.run_one() { StepResult::Success } else { StepResult::InputEmpty } - } - - /// Run one deferred key-extraction batch, if any is queued. - /// - /// Takes no worker state: a batch owns everything it needs and publishes its - /// own result, so there is nothing to hold and nothing to advance. - fn try_extract_keys(shared: &SharedPipelineState) -> StepResult { - let Some(job) = shared.key_jobs.pop() else { - return StepResult::InputEmpty; - }; - job.run(); - StepResult::Success - } - // ======================================================================== // Held-item advancement (deadlock prevention) // ======================================================================== @@ -2818,7 +2250,7 @@ impl SortWorkerPool { let pushed = try_advance_held(&shared.decompressed_input, &mut worker.held_decompressed_input); if pushed { - shared.wake_consumer(); + shared.main_thread_handle.unpark(); advanced = true; // This may have been the last block. Increment `input_blocks_queued` now // that the block is actually in the queue (not held), then re-check the @@ -2833,7 +2265,7 @@ impl SortWorkerPool { && !shared.decompressed_input_done.load(Ordering::Acquire) { shared.decompressed_input_done.store(true, Ordering::Release); - shared.wake_consumer(); + shared.main_thread_handle.unpark(); } } } @@ -2866,21 +2298,14 @@ impl SortWorkerPool { return StepResult::InputEmpty; // No input file set }; - // Read a batch of raw BGZF blocks. - // - // Timed as one region because that is the only honest boundary: the - // 2 MiB refill happens *inside* this call, on whichever block exhausts - // the buffer. Splitting disk from framing needs a timer below the - // buffer, which is what the `TimedReader` on the input file is for -- - // `ReaderReport::framing_secs` is this span minus that one. - let framing_started = Instant::now(); + // Read a batch of raw BGZF blocks let blocks = match read_raw_blocks(reader.as_mut(), INPUT_READ_BATCH_SIZE) { Ok(b) => b, Err(e) => { log::error!("I/O error reading input BAM: {e}"); shared.input_read_error.store(true, Ordering::Release); shared.input_eof.store(true, Ordering::Release); - shared.wake_consumer(); + shared.main_thread_handle.unpark(); return StepResult::InputEmpty; } }; @@ -2890,13 +2315,6 @@ impl SortWorkerPool { return StepResult::InputEmpty; } - // Recorded only past the empty check, so the partition covers exactly the - // calls `record_step` charges to `ReadInputBlocks` -- that step is timed - // on `StepResult::Success` alone, and the EOF-detecting call returns - // `InputEmpty`. Counting it here would put time in the parts that is not - // in the total and drive the residual negative for no reason. - shared.reader_stats.record_batch(Self::nanos_u64(framing_started.elapsed()), blocks.len()); - // Reserve this batch's serial range while STILL HOLDING the input lock. // // Assigning serials after releasing the lock is racy: a worker preempted @@ -2919,14 +2337,12 @@ impl SortWorkerPool { // Drop the lock before pushing to queue drop(guard); - let dispatch_started = Instant::now(); dispatch_reserved_blocks( base_serial, blocks, &shared.raw_input_blocks, &mut worker.held_raw_input_blocks, ); - shared.reader_stats.record_dispatch(Self::nanos_u64(dispatch_started.elapsed())); StepResult::Success } @@ -2954,18 +2370,18 @@ impl SortWorkerPool { let total = shared.input_read_serial.load(Ordering::Acquire); if queued >= total { shared.decompressed_input_done.store(true, Ordering::Release); - shared.wake_consumer(); + shared.main_thread_handle.unpark(); } } return StepResult::InputEmpty; }; - let data = match decompress_block(&block, &mut worker.decomp.bgzf) { + let data = match decompress_block(&block, &mut worker.decompressor) { Ok(d) => d, Err(e) => { log::error!("BGZF decompression error (input block serial {serial}): {e}"); shared.decompression_error.store(true, Ordering::Release); - shared.wake_consumer(); + shared.main_thread_handle.unpark(); return StepResult::InputEmpty; } }; @@ -2976,13 +2392,13 @@ impl SortWorkerPool { // Try to push to decompressed_input ArrayQueue let pushed = match shared.decompressed_input.push((serial, data)) { Ok(()) => { - shared.wake_consumer(); + shared.main_thread_handle.unpark(); true } Err(item) => { worker.held_decompressed_input = Some(item); // Queue full — wake main thread to drain so we can push next time - shared.wake_consumer(); + shared.main_thread_handle.unpark(); false } }; @@ -2997,7 +2413,7 @@ impl SortWorkerPool { let total = shared.input_read_serial.load(Ordering::Acquire); if input_eof && raw_empty && queued >= total { shared.decompressed_input_done.store(true, Ordering::Release); - shared.wake_consumer(); + shared.main_thread_handle.unpark(); } } @@ -3070,18 +2486,7 @@ impl SortWorkerPool { // that makes that case fast. let frontier = shared.phase2_lowest_active.load(Ordering::Relaxed); let frontier_starving = files.get(frontier).is_some_and(Phase2FileState::is_starving); - // The file the consumer is parked on, which the frontier only coincides - // with when sources drain in index order. - let awaited = shared.phase2_awaited_source.load(Ordering::Relaxed); - let awaited_starving = files.get(awaited).is_some_and(Phase2FileState::is_starving); - let start = phase2_scan_start( - awaited, - awaited_starving, - frontier, - n, - frontier_starving, - worker.phase2_file_cursor, - ); + let start = phase2_scan_start(frontier, n, frontier_starving, worker.phase2_file_cursor); for offset in 0..n { let i = (start + offset) % n; @@ -3107,24 +2512,9 @@ impl SortWorkerPool { let pop_skip = match Self::try_pop_raw_for_decompress(file) { Err(skip) => skip, Ok(entry) => { - // `offset` files gave nothing before this one did. The scan - // tally is published only on fruitless scans, so without this - // the walk to a *successful* claim went uncounted. - shared.stage_latency.record_claim(offset as u64); worker.phase2_file_cursor = (i + 1) % n; - worker.served_awaited = - shared.phase2_awaited_source.load(Ordering::Relaxed) == i; - if worker.served_awaited { - worker.won_awaited_claim = Self::stamp_awaited_claim(shared); - } Self::note_claim(shared, file, &entry); - return Self::decompress_and_publish( - shared, - &mut worker.decomp, - file, - i, - entry, - ); + return Self::decompress_and_publish(shared, worker, file, i, entry); } }; @@ -3198,11 +2588,10 @@ impl SortWorkerPool { // entry-count allowances and an unbounded (`usize::MAX`) byte budget, // so they behave exactly as before. let awaited = shared.phase2_awaited_source.load(Ordering::Relaxed); - let next_source = shared.phase2_next_source.load(Ordering::Relaxed); let (raw_cap, read_batch, fifo_byte_budget, read_byte_budget) = if i == frontier { let (cap, batch) = phase2_read_allowance(true); (cap, batch, usize::MAX, usize::MAX) - } else if phase2_deserves_deep_read(i, frontier, awaited, next_source) { + } else if phase2_deserves_deep_read(i, frontier, awaited) { let read_blocks = shared.phase2_read_blocks.load(Ordering::Relaxed); let mean_block_bytes = if read_blocks == 0 { 0 @@ -3252,7 +2641,7 @@ impl SortWorkerPool { // dispatched on the codec but the failure is handled once. The two // arms previously carried identical error handling differing only // in one word of the message, which had to be kept in step by hand. - let read = shared.stage_latency.read.time(|| match file.codec { + let read = match file.codec { SpillCodec::Bgzf => shared .merge_phases .read @@ -3261,7 +2650,7 @@ impl SortWorkerPool { SpillCodec::Zstd => shared.merge_phases.read.time(|| { read_raw_zstd_frames(&mut reader_guard.inner, read_batch, read_byte_budget) }), - }); + }; let raw_bytes: Vec> = match read { Ok(bytes) => bytes, Err(e) => { @@ -3277,8 +2666,7 @@ impl SortWorkerPool { // allowance selection uses, so an awaited non-frontier source -- // which reads at the deep allowance via `awaited_allowance_for` -- // is counted as deep, not shallow. - let (batches, blocks) = if phase2_deserves_deep_read(i, frontier, awaited, next_source) - { + let (batches, blocks) = if phase2_deserves_deep_read(i, frontier, awaited) { (&shared.deep_read_batches, &shared.deep_read_blocks) } else { (&shared.shallow_read_batches, &shared.shallow_read_blocks) @@ -3344,16 +2732,7 @@ impl SortWorkerPool { shared.refill.read_lag.record(enqueued_nanos.saturating_sub(emptied)); } - shared.stage_latency.record_claim(offset as u64); worker.phase2_file_cursor = (i + 1) % n; - // Reading for the awaited file is on the critical path as much as - // decompressing for it: `raw-empty` is 47-70% of the reasons the pool - // passes that file over, so a late wake that ends in a read is a late - // wake the consumer paid for. - worker.served_awaited = shared.phase2_awaited_source.load(Ordering::Relaxed) == i; - if worker.served_awaited { - worker.won_awaited_claim = Self::stamp_awaited_claim(shared); - } return StepResult::Success; } @@ -3407,7 +2786,7 @@ impl SortWorkerPool { /// responsible for releasing. fn decompress_and_publish( shared: &SharedPipelineState, - decomp: &mut DecompressorSet, + worker: &mut SortWorkerState, file: &Phase2FileState, source_idx: usize, entry: RawEntry, @@ -3416,12 +2795,11 @@ impl SortWorkerPool { let data = match file.codec { SpillCodec::Bgzf => { let raw_block = RawBgzfBlock { data: raw_bytes }; - match shared.stage_latency.decompress.time(|| { - shared - .merge_phases - .decompress - .time(|| decompress_block(&raw_block, &mut decomp.bgzf)) - }) { + match shared + .merge_phases + .decompress + .time(|| decompress_block(&raw_block, &mut worker.decompressor)) + { Ok(d) => d, Err(e) => { log::error!( @@ -3429,7 +2807,7 @@ impl SortWorkerPool { ); shared.decompression_error.store(true, Ordering::Release); file.decomp_in_flight.fetch_sub(1, Ordering::AcqRel); - shared.wake_consumer(); + shared.main_thread_handle.unpark(); return StepResult::Success; } } @@ -3437,28 +2815,27 @@ impl SortWorkerPool { SpillCodec::Zstd => { // Allocate the scratch buffer lazily so BGZF-only sorts don't // pay 256 KiB × num_workers of dead memory. - if decomp.zstd_buf.len() < zstd_decomp_cap() { - decomp.zstd_buf.resize(zstd_decomp_cap(), 0); + if worker.zstd_decompress_buf.len() < ZSTD_FRAME_DECOMP_CAP { + worker.zstd_decompress_buf.resize(ZSTD_FRAME_DECOMP_CAP, 0); } - match shared.stage_latency.decompress.time(|| { - shared - .merge_phases - .decompress - .time(|| decomp.zstd.decompress_to_buffer(&raw_bytes, &mut decomp.zstd_buf)) + match shared.merge_phases.decompress.time(|| { + worker + .zstd_decompressor + .decompress_to_buffer(&raw_bytes, &mut worker.zstd_decompress_buf) }) { // Copy the `n` decompressed bytes (≤ one staging-buffer's // worth, typically ~65 KB) into a fresh Vec for the // consumer. The scratch buffer keeps its 256 KiB capacity so // the next frame on this worker reuses it without // reallocating. - Ok(n) => decomp.zstd_buf[..n].to_vec(), + Ok(n) => worker.zstd_decompress_buf[..n].to_vec(), Err(e) => { log::error!( "zstd decompression error (chunk source {source_idx} serial {serial}): {e}" ); shared.decompression_error.store(true, Ordering::Release); file.decomp_in_flight.fetch_sub(1, Ordering::AcqRel); - shared.wake_consumer(); + shared.main_thread_handle.unpark(); return StepResult::Success; } } @@ -3502,25 +2879,11 @@ impl SortWorkerPool { // Wake the consumer either because new data is available or because // the last in-flight decompression for this file just completed and // the file is now fully drained. - shared.wake_consumer(); + shared.main_thread_handle.unpark(); } StepResult::Success } - /// Stamp the first claim of the consumer's awaited block, and credit the - /// worker that made it. - /// - /// Compare-exchange from 0 so only the first claim in a park is recorded -- - /// the consumer waits for whichever worker arrives first, so a sum over all - /// of them would overcount exactly the way the worker-side lag total does. - fn stamp_awaited_claim(shared: &SharedPipelineState) -> bool { - let now = shared.now_nanos(); - shared - .awaited_claim_nanos - .compare_exchange(0, now, Ordering::Relaxed, Ordering::Relaxed) - .is_ok() - } - /// Try to pop a raw block from `file` for decompression, applying /// deadlock-free admission control against the file's reorder buffer. /// @@ -3607,12 +2970,8 @@ impl SortWorkerPool { }; // Only the compression is timed; the handoff below can block on a full // result channel, which is writer backpressure rather than CPU work. - let latency = match target { - CompressTarget::Spill => &shared.stage_latency.spill_compress, - CompressTarget::Output => &shared.stage_latency.output_compress, - }; - let compressed = latency - .time(|| counter.time(|| Self::compress_job(&job, bgzf_compressor, zstd_compressor))); + let compressed = + counter.time(|| Self::compress_job(&job, bgzf_compressor, zstd_compressor)); Self::deliver_compress_result(shared, job, compressed); StepResult::Success } @@ -3647,7 +3006,7 @@ impl SortWorkerPool { ) -> StepResult { file.mark_reader_eof(&mut reader_guard); drop(reader_guard); - shared.wake_consumer(); + shared.main_thread_handle.unpark(); Self::maybe_mark_all_eof(shared); worker.phase2_file_cursor = (source + 1) % num_sources; StepResult::Success @@ -3659,7 +3018,7 @@ impl SortWorkerPool { let total = shared.total_sources.load(Ordering::Acquire); if total > 0 && eof_count >= total { shared.all_chunks_eof.store(true, Ordering::Release); - shared.wake_consumer(); + shared.main_thread_handle.unpark(); } } @@ -3693,44 +3052,6 @@ impl SortWorkerPool { self.shared.merge_phases.snapshot() } - /// Worker seconds spent feeding Phase 1's ingest thread: reading raw input - /// blocks off disk and decompressing them. - /// - /// Deliberately excludes `Compress`, which serves spill *and* output and so - /// cannot be attributed to the ingest span; the spill half is reported on its - /// own through [`crate::merge_phases`]. - pub(crate) fn phase1_input_busy_secs(&self) -> f64 { - let ns = self.pipeline_stats.step_ns[SortStep::ReadInputBlocks as usize] - .load(Ordering::Relaxed) - + self.pipeline_stats.step_ns[SortStep::DecompressInput as usize] - .load(Ordering::Relaxed); - #[allow(clippy::cast_precision_loss, reason = "nanosecond totals stay far below 2^52")] - { - ns as f64 / 1_000_000_000.0 - } - } - - /// Counters for what Phase 1's ingest thread waited on. - pub(crate) fn phase1_ingest_stats(&self) -> Arc { - Arc::clone(&self.shared.phase1_ingest) - } - - /// Counters for the input reader, shared with the `TimedReader` that has to - /// be built before the pool is handed the file. - pub(crate) fn reader_stats(&self) -> Arc { - Arc::clone(&self.shared.reader_stats) - } - - /// The input reader's partition, checked against the exact `ReadInputBlocks` - /// busy total it should add up to. - pub(crate) fn phase1_reader_report(&self) -> crate::phase1_stats::ReaderReport { - let ns = - self.pipeline_stats.step_ns[SortStep::ReadInputBlocks as usize].load(Ordering::Relaxed); - #[allow(clippy::cast_precision_loss, reason = "nanosecond totals stay far below 2^52")] - let step_secs = ns as f64 / 1_000_000_000.0; - self.shared.reader_stats.snapshot(step_secs) - } - /// The pool's shared state, for the merge consumer's own instrumentation. /// /// The consumer runs on the main thread and is not a pool worker, but it is @@ -3762,105 +3083,6 @@ impl SortWorkerPool { (mean, cap, batch) } - /// Where the merge consumer's park time went, split into additive stages. - pub(crate) fn park_attribution_report(&self) -> crate::merge_stalls::ParkAttributionReport { - self.shared.park_attribution.snapshot() - } - - /// Per-worker busy time, idle time and critical-path claims, worker 0 first. - /// - /// Truncated to the workers this pool actually started, so a 16-thread merge - /// does not print 32 rows of zeros. - pub(crate) fn per_thread_report(&self) -> Vec<(u64, u64, u64)> { - let stats = &self.pipeline_stats; - (0..self.num_workers.min(SORT_MAX_THREADS)) - .map(|w| { - ( - stats.per_thread_busy_ns[w].load(Ordering::Relaxed), - stats.per_thread_idle_ns[w].load(Ordering::Relaxed), - stats.per_thread_awaited_claims[w].load(Ordering::Relaxed), - ) - }) - .collect() - } - - /// Per-stage latency distributions and wasted-scan accounting. - pub(crate) fn stage_latency(&self) -> &crate::merge_phases::StageLatency { - &self.shared.stage_latency - } - - /// Publish the source the merge expects to consume next, so the pool can start - /// that file's read before the consumer arrives. - pub(crate) fn set_phase2_next_source(&self, next: Option) { - self.shared.set_phase2_next_source(next); - } - - /// Consumer parks split by what the pool looked like when the park began. - pub(crate) fn park_supply_report(&self) -> crate::merge_stalls::ParkSupplyReport { - self.shared.park_supply.snapshot() - } - - /// Decompress one already-read block for `source_idx` on the calling thread, - /// returning whether it published anything. - /// - /// For the merge consumer to call instead of parking. The block it needs is - /// sitting in the raw FIFO, read but unclaimed, in 12% of parks -- and parking - /// there trades ~53us of work for a ~190us wait on a worker that has to be - /// woken first. - /// - /// The deeper reason is that the consumer's exposed latency is **per block, - /// not per park**: raising `PHASE2_DECOMP_CAP` from 8 to 128 cut parks from - /// 978,325 to 400,487 and left total park time at 89.6s against 89.7s, while - /// wall clock got monotonically worse. A cost that survives a 16x change in - /// buffer depth is not a buffering problem and not a signalling one -- eight - /// scheduling interventions moved their own metrics and left the clock alone. - /// What is left is the producer-to-consumer handoff itself, and the only way - /// to remove a handoff is for one thread to do both halves. - /// - /// Goes through the same [`Self::try_pop_raw_for_decompress`] and - /// [`Self::decompress_and_publish`] as a worker, so the in-flight counter and - /// the reorder-buffer protocol cannot drift between the two paths. Publishing - /// rather than consuming the block directly is deliberate: the consumer then - /// picks it up through its ordinary path, and nothing about serial ordering - /// gets a second implementation. - /// - /// Safe to call while parked-pending: the consumer holds no locks, so it - /// cannot participate in a cycle with a worker. - pub(crate) fn consumer_decompress_one( - shared: &SharedPipelineState, - decomp: &mut DecompressorSet, - file: &Phase2FileState, - source_idx: usize, - ) -> bool { - let Ok(entry) = Self::try_pop_raw_for_decompress(file) else { return false }; - let _ = Self::decompress_and_publish(shared, decomp, file, source_idx, entry); - true - } - - /// Parks the consumer avoided by decompressing a block itself. - pub(crate) fn consumer_self_served(&self) -> u64 { - self.shared.consumer_self_served.load(Ordering::Relaxed) - } - - /// Predictions published to the pool by the merge's loser tree. - pub(crate) fn phase2_predictions(&self) -> u64 { - self.shared.phase2_predictions() - } - - /// Where the consumer's wakes landed. - pub(crate) fn wake_accounting(&self) -> WakeAccounting { - WakeAccounting { - issued: self.shared.wakes_issued.load(Ordering::Relaxed), - on_running: self.shared.wakes_on_running_worker.load(Ordering::Relaxed), - recoverable: self.shared.wakes_recoverable.load(Ordering::Relaxed), - } - } - - /// Wakes issued by the consumer to cut a worker's backoff short. - pub(crate) fn wakes_issued(&self) -> u64 { - self.shared.wakes_issued.load(Ordering::Relaxed) - } - /// Why worker scans passed over the file the consumer was parked on, /// as `[raw-lock, raw-empty, decomp-lock, decomp-capped]`. pub(crate) fn awaited_skip_counts(&self) -> [u64; 4] { @@ -4053,7 +3275,6 @@ impl SortWorkerPool { /// /// Panics if the `phase2_files` rwlock is poisoned. pub fn set_phase2_files(&self, files: &[std::path::PathBuf]) -> anyhow::Result<()> { - let scatter = Arc::clone(&self.shared.fetch_jobs); let total_sources = files.len(); self.shared.total_sources.store(total_sources as u64, Ordering::Release); @@ -4068,12 +3289,6 @@ impl SortWorkerPool { // forward, so neither is self-correcting. self.shared.phase2_lowest_active.store(0, Ordering::Release); self.shared.phase2_awaited_source.store(NO_AWAITED_SOURCE, Ordering::Release); - // The read-ahead prediction indexes the file vector being replaced, so a - // pool that merges twice would expose the prior merge's `runner_up` - // through `phase2_next_source` while the new merge seeds its sources -- - // handing deep-read priority to an unrelated file until the new merge - // publishes its own prediction. Reset it to the sentinel here. - self.shared.phase2_next_source.store(NO_AWAITED_SOURCE, Ordering::Release); // The refill allowance is derived from `phase2_read_bytes / // phase2_read_blocks`, so these must describe only the set being // installed: mean spill-block size moves with the codec, the @@ -4125,26 +3340,7 @@ impl SortWorkerPool { file.seek(SeekFrom::Start(body_start)).map_err(|e| { anyhow::anyhow!("Failed to seek chunk file {}: {e}", path.display()) })?; - let reader = if self.read_streams.is_sequential() { - crate::spill_reader::SpillSource::Sequential(BufReader::with_capacity( - 2 * 1024 * 1024, - file, - )) - } else { - // Positional reads need no file position, so the seek above is - // irrelevant here -- `body_start` is passed explicitly instead. - crate::spill_reader::SpillSource::Scattered( - crate::spill_reader::ScatterReader::for_streams( - file, - body_start, - self.read_streams, - Some(Arc::clone(&scatter)), - ) - .map_err(|e| { - anyhow::anyhow!("Failed to size chunk file {}: {e}", path.display()) - })?, - ) - }; + let reader = BufReader::with_capacity(2 * 1024 * 1024, file); states.push(Phase2FileState::new(reader, codec)); } @@ -4153,11 +3349,6 @@ impl SortWorkerPool { Ok(()) } - /// The queue readers offer byte slices to. See [`crate::spill_reader`]. - pub(crate) fn fetch_queue(&self) -> Arc { - Arc::clone(&self.shared.fetch_jobs) - } - /// Clear the Phase 2 file vector. Call this after Phase 2 finishes (and /// before any subsequent Phase 1) so the file descriptors are released. /// @@ -4191,51 +3382,6 @@ impl SortWorkerPool { } } - /// Offer a deferred key-extraction batch to the pool. - /// - /// Returns the batch unrun in `Err` when the queue is full or the pool has - /// shut down; the caller must run it inline rather than drop it. Unlike - /// [`submit_compress`](Self::submit_compress) this does **not** spin waiting - /// for room: the submitter is the ingest thread, and blocking it to hand off - /// work whose whole purpose is to unblock it would be self-defeating. - /// - /// # Errors - /// - /// Returns the batch when it could not be queued. - pub(crate) fn submit_key_job( - &self, - job: Box, - ) -> Result<(), Box> { - if self.shared.phase.load(Ordering::Acquire) == phase::SHUTDOWN { - return Err(job); // Workers have exited; no one will pop the queue - } - self.shared.key_jobs.push(job)?; - self.shared.wake_one_worker(); - Ok(()) - } - - /// Whether any worker has published a panic. - pub(crate) fn worker_panicked(&self) -> bool { - self.shared.worker_panicked.load(Ordering::Acquire) - } - - /// Run one queued key-extraction batch on the calling thread, if any is - /// queued. Returns whether a batch ran. - /// - /// The ingest thread's self-help path: at the chunk barrier it drains the - /// queue itself rather than waiting on workers that may already have parked - /// for the phase. Without it the barrier could wait on a batch nobody is - /// scheduled to run. - pub(crate) fn run_one_key_job(&self) -> bool { - match self.shared.key_jobs.pop() { - Some(job) => { - job.run(); - true - } - None => false, - } - } - /// Create a new result channel pair for compress results. /// /// The result channel stays as `crossbeam_channel::bounded()` because the @@ -4261,19 +3407,6 @@ impl SortWorkerPool { /// Internal shutdown: signal workers and join them. Safe to call multiple times /// (idempotent via `Option::take`). Called by both `shutdown` and `Drop`. fn do_shutdown(&mut self) { - // The census for "did scattered reading engage at all". Zero taken means - // every fill was read by the thread that wanted it -- the pre-change - // behaviour, and something no output check would ever notice. Reported - // here because the queue spans both phases and this is the one point - // every sort passes through. Gated on the workers still being present so - // it prints once: `do_shutdown` runs from both `shutdown` and `Drop`, - // and only the join below is idempotent on its own. - if self.workers.is_some() { - let (offered, taken) = self.shared.fetch_jobs.census(); - if offered > 0 { - log::info!("Byte fetch: {offered} slices offered, {taken} run by workers"); - } - } self.shared.phase.store(phase::SHUTDOWN, Ordering::Release); if let Some(workers) = self.workers.take() { for w in workers { @@ -4281,7 +3414,7 @@ impl SortWorkerPool { // Worker panicked — set flag and wake main thread so it doesn't // park forever waiting for work that will never arrive. self.shared.worker_panicked.store(true, Ordering::Release); - self.shared.wake_consumer(); + self.shared.main_thread_handle.unpark(); } } } @@ -4676,7 +3809,7 @@ mod tests { fn test_compress_target_decides_level_regardless_of_phase(#[case] phase_under_test: u8) { // Compressible: 8 KiB of one byte shrinks to almost nothing at level 9. let data = vec![b'A'; 8192]; - let pool = SortWorkerPool::new(2, 0, 9, crate::codec::SpillCodec::Bgzf, false); + let pool = SortWorkerPool::new(2, 0, 9, crate::codec::SpillCodec::Bgzf); pool.set_phase(phase_under_test); let compressed_len = |target: CompressTarget| { @@ -4748,7 +3881,7 @@ mod tests { #[test] fn test_pool_compress_roundtrip() { - let pool = SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf, false); + let pool = SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf); let (result_tx, result_rx) = pool.compress_result_channel(); // Submit a compress job @@ -4773,103 +3906,9 @@ mod tests { pool.shutdown(); } - #[test] - fn test_the_pool_runs_every_submitted_key_job() { - // Deferred key extraction is only a win if the pool actually drains the - // batches; a batch the pool never picks up would stall the ingest - // thread's barrier instead. Submits more batches than the queue holds, - // so the full-queue path is exercised too. - use crate::phase1_keys::KeyExtractionJob; - - struct CountingBatch(std::sync::mpsc::Sender, usize); - impl KeyExtractionJob for CountingBatch { - fn run(self: Box) { - // The receiver outlives every batch here; a closed channel would - // just mean the test already finished. - let _ = self.0.send(self.1); - } - } - - let pool = SortWorkerPool::new(4, 1, 6, crate::codec::SpillCodec::Bgzf, false); - pool.set_phase(phase::PHASE1); - let (tx, rx) = std::sync::mpsc::channel(); - - let num_jobs = 64usize; - for i in 0..num_jobs { - // A full queue hands the batch back rather than dropping it; the - // ingest thread runs those inline, so the test does the same. - if let Err(job) = pool.submit_key_job(Box::new(CountingBatch(tx.clone(), i))) { - job.run(); - } - } - drop(tx); - - let mut seen: Vec = rx.iter().collect(); - seen.sort_unstable(); - assert_eq!(seen, (0..num_jobs).collect::>(), "every batch must run exactly once"); - } - - #[test] - fn test_the_worker_that_owns_the_reader_does_not_take_key_batches() { - // Measured on the production cell: adding key extraction to the pool - // slowed the input reader from 358 to 345 MB/s and its step from 125.2s - // to 129.8s. The reader is one exclusively-owned worker, and a key batch - // is ~0.5 ms during which that worker is not reading. Everyone else may - // take batches; the reader's owner may not. - assert!( - !SortWorkerPool::worker_may_extract_keys(Some(SortStep::ReadInputBlocks)), - "the reader's owner must stay on the reader" - ); - assert!( - SortWorkerPool::worker_may_extract_keys(None), - "a worker owning no exclusive step is free to extract" - ); - assert!( - SortWorkerPool::worker_may_extract_keys(Some(SortStep::Compress)), - "owning some other step does not bar extraction" - ); - } - - #[test] - fn test_workers_stay_awake_for_key_batches_after_the_input_is_drained() { - // The final chunk's barrier lands exactly here: the input is read and - // decompressed, so Phase 1 looks finished, but a chunk's worth of key - // batches is still queued. If the phase counts as complete the workers - // all park and the ingest thread drains them one at a time on its own — - // correct output, but the last chunk's extraction becomes fully serial, - // which is the cost this step exists to remove. - use crate::phase1_keys::KeyExtractionJob; - - struct Ping(std::sync::mpsc::Sender<()>); - impl KeyExtractionJob for Ping { - fn run(self: Box) { - let _ = self.0.send(()); - } - } - - let pool = SortWorkerPool::new(4, 1, 6, crate::codec::SpillCodec::Bgzf, false); - pool.set_phase(phase::PHASE1); - // Everything the phase-completion check looks at is now satisfied. - pool.shared.decompressed_input_done.store(true, Ordering::Release); - pool.shared.input_eof.store(true, Ordering::Release); - - let (tx, rx) = std::sync::mpsc::channel(); - let submitted = 8usize; - for _ in 0..submitted { - pool.submit_key_job(Box::new(Ping(tx.clone()))).ok().expect("queue has room"); - } - drop(tx); - - for i in 0..submitted { - rx.recv_timeout(std::time::Duration::from_secs(10)).unwrap_or_else(|_| { - panic!("workers parked with key batches still queued; only {i} of {submitted} ran") - }); - } - } - #[test] fn test_pool_many_jobs() { - let pool = SortWorkerPool::new(4, 1, 6, crate::codec::SpillCodec::Bgzf, false); + let pool = SortWorkerPool::new(4, 1, 6, crate::codec::SpillCodec::Bgzf); let (result_tx, result_rx) = pool.compress_result_channel(); let num_jobs = 100usize; @@ -4906,7 +3945,7 @@ mod tests { #[test] fn set_active_workers_clamps() { - let pool = SortWorkerPool::new(4, 1, 6, crate::codec::SpillCodec::Bgzf, false); + let pool = SortWorkerPool::new(4, 1, 6, crate::codec::SpillCodec::Bgzf); // A fresh pool is fully active, so anything dividing by the active count // before a cap is applied still gets the pool width. assert_eq!(pool.active_workers(), 4, "a pool starts with every worker active"); @@ -4926,7 +3965,7 @@ mod tests { /// utilization and flip the merge verdict from CPU-bound to I/O-bound. #[test] fn active_workers_reports_the_phase2_cap_not_the_pool_width() { - let pool = SortWorkerPool::new(8, 1, 6, crate::codec::SpillCodec::Bgzf, false); + let pool = SortWorkerPool::new(8, 1, 6, crate::codec::SpillCodec::Bgzf); pool.set_active_workers(8); pool.begin_phase2(3); assert_eq!(pool.num_workers(), 8, "the pool is still eight threads wide"); @@ -4964,63 +4003,7 @@ mod tests { #[case] cursor: usize, #[case] expected: usize, ) { - assert_eq!( - phase2_scan_start( - NO_AWAITED_SOURCE, - false, - frontier, - num_files, - frontier_starving, - cursor - ), - expected - ); - } - - /// A starving *awaited* source outranks the frontier, because it is measured - /// demand rather than a proxy for it. - /// - /// The frontier is only the file the merge reaches next when sources drain in - /// index order. On a partially-correlated input they do not, which is the - /// assumption that cost this merge 126s of read-ahead going to the wrong - /// file. The consumer knows which file it is parked on; this is what lets the - /// pool act on it. - /// - /// Gated on starving for the same reason the frontier is: `awaited` is never - /// cleared between parks, so an ungated version would send every worker to - /// one index permanently and trade a spread problem for a herd problem. - #[rstest] - // Measured demand beats the proxy, and beats the cursor. - #[case::starving_awaited_outranks_a_starving_frontier(5, true, 0, 8, true, 2, 5)] - #[case::starving_awaited_outranks_an_active_frontier(5, true, 0, 8, false, 2, 5)] - #[case::starving_awaited_beats_the_cursor(6, true, 9, 8, false, 2, 6)] - // An awaited file with data to give must not pull workers off the spread. - #[case::active_awaited_defers_to_a_starving_frontier(5, false, 3, 8, true, 2, 3)] - #[case::active_awaited_defers_to_the_cursor(5, false, 0, 8, false, 2, 2)] - // Unset, or past the end, names no live source. - #[case::unset_awaited_defers(NO_AWAITED_SOURCE, true, 3, 8, true, 2, 3)] - #[case::awaited_past_the_end_defers(8, true, 3, 8, true, 2, 3)] - #[case::no_files_defers(0, true, 0, 0, true, 0, 0)] - fn phase2_scan_prefers_a_starving_awaited_source( - #[case] awaited: usize, - #[case] awaited_starving: bool, - #[case] frontier: usize, - #[case] num_files: usize, - #[case] frontier_starving: bool, - #[case] cursor: usize, - #[case] expected: usize, - ) { - assert_eq!( - phase2_scan_start( - awaited, - awaited_starving, - frontier, - num_files, - frontier_starving, - cursor - ), - expected - ); + assert_eq!(phase2_scan_start(frontier, num_files, frontier_starving, cursor), expected); } /// Only the drain frontier gets the deep allowance. @@ -5140,24 +4123,17 @@ mod tests { /// shallow allowance: the pool sat 1.9 decompressions deep of a tracked 8 /// while the consumer waited 73% of the merge loop. #[rstest] - #[case::frontier_only(0, usize::MAX, usize::MAX, 0, true)] - #[case::awaited_only(0, 5, usize::MAX, 5, true)] - #[case::neither(0, 5, usize::MAX, 3, false)] - #[case::no_awaited_source_yet(0, usize::MAX, usize::MAX, 7, false)] - // The predicted next source reads deep too: reactive signals fire only after - // the consumer has already stalled, and a run transition costs ~20ms because - // the read had not been started when it arrived. - #[case::predicted_next_source(0, usize::MAX, 9, 9, true)] - #[case::predicted_next_is_not_the_candidate(0, usize::MAX, 9, 4, false)] - #[case::no_prediction_yet(0, usize::MAX, usize::MAX, 4, false)] + #[case::frontier_only(0, usize::MAX, 0, true)] + #[case::awaited_only(0, 5, 5, true)] + #[case::neither(0, 5, 3, false)] + #[case::no_awaited_source_yet(0, usize::MAX, 7, false)] fn the_blocked_file_reads_deep_even_when_it_is_not_the_frontier( #[case] frontier: usize, #[case] awaited: usize, - #[case] next: usize, #[case] candidate: usize, #[case] expect_deep: bool, ) { - let deep = phase2_deserves_deep_read(candidate, frontier, awaited, next); + let deep = phase2_deserves_deep_read(candidate, frontier, awaited); assert_eq!(deep, expect_deep); let expected = if expect_deep { (PHASE2_STARVING_RAW_CAP, PHASE2_STARVING_READ_BATCH) @@ -5307,84 +4283,6 @@ mod tests { /// over the pool width instead sends most wakes to workers idled by the /// Phase 2 cap, which re-park without looking at the starving file — and /// the workers that could have refilled it wait out their full backoff. - /// A published prediction must be counted, and "no prediction" must not be. - /// - /// Without this counter a `runner_up()` that returned `None` on every call - /// would be indistinguishable from a null result: the deep-read path simply - /// never fires, wall clock lands wherever it lands, and nothing in the log - /// says the mechanism was inert. The measured -2.7% is only attributable to - /// prediction if predictions were actually published. - #[test] - fn test_published_predictions_are_counted_and_absent_ones_are_not() { - let shared = SharedPipelineState::new(2, std::thread::current()); - assert_eq!(shared.phase2_predictions(), 0, "nothing published yet"); - - shared.set_phase2_next_source(Some(3)); - shared.set_phase2_next_source(Some(1)); - assert_eq!(shared.phase2_predictions(), 2, "each published source counts once"); - - shared.set_phase2_next_source(None); - assert_eq!(shared.phase2_predictions(), 2, "clearing the prediction is not a prediction"); - } - - /// A wake spent on a running worker while a sleeper sits idle is the merge's - /// largest single idle cost at t16 -- 74.4s of a 92.1s park. The preference - /// must actually find the sleeper. - #[test] - fn test_wake_prefers_a_parked_worker_over_the_rotating_target() { - let parked = [false, false, true, false]; - let at = |i: usize| parked[i]; - assert_eq!( - SharedPipelineState::wake_target_preferring_parked(0, 4, 4, at), - 2, - "the rotating target is 0 and running; worker 2 is asleep and gets the wake" - ); - } - - /// Falling back to the rotation is what preserves the spread an interleaved - /// merge needs, so a saturated pool must behave exactly as before. - #[test] - fn test_wake_falls_back_to_rotation_when_nobody_is_parked() { - for cursor in 0..8 { - assert_eq!( - SharedPipelineState::wake_target_preferring_parked(cursor, 4, 4, |_| false), - SharedPipelineState::wake_target(cursor, 4, 4), - "with nobody parked the target must be unchanged at cursor {cursor}" - ); - } - } - - #[test] - fn test_wake_keeps_the_rotating_target_when_it_is_itself_parked() { - let parked = [true, true, true, true]; - assert_eq!(SharedPipelineState::wake_target_preferring_parked(2, 4, 4, |i| parked[i]), 2); - } - - /// Capped workers will not take Phase 2 work, so waking one is the same lost - /// wake by another route -- the preference must stay inside the active window. - #[test] - fn test_wake_preference_ignores_workers_outside_the_active_limit() { - let parked = [false, false, true, true]; - let at = |i: usize| parked[i]; - assert_eq!( - SharedPipelineState::wake_target_preferring_parked(0, 2, 4, at), - SharedPipelineState::wake_target(0, 2, 4), - "workers 2 and 3 are parked but capped out, so the rotation stands" - ); - } - - /// The scan starts at the cursor, so successive wakes with several sleepers - /// spread rather than piling onto the lowest index. - #[test] - fn test_wake_preference_rotates_across_several_sleepers() { - let parked = [true, false, true, false]; - let at = |i: usize| parked[i]; - let picks: Vec = (0..4) - .map(|c| SharedPipelineState::wake_target_preferring_parked(c, 4, 4, at)) - .collect(); - assert_eq!(picks, vec![0, 2, 2, 0], "each wake starts scanning from its own cursor"); - } - #[rstest] // A pool 8 wide capped to 3 must never select a worker above the cap, at // any point in the rotation. @@ -5485,7 +4383,7 @@ mod tests { (pool, counts) }; - let pool = SortWorkerPool::new(6, 1, 6, crate::codec::SpillCodec::Bgzf, false); + let pool = SortWorkerPool::new(6, 1, 6, crate::codec::SpillCodec::Bgzf); // Batch 1: capped at 2 — only workers 0..2 may run. pool.set_active_workers(2); @@ -5521,7 +4419,7 @@ mod tests { #[test] fn test_pool_stats() { - let pool = SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf, false); + let pool = SortWorkerPool::new(2, 1, 6, crate::codec::SpillCodec::Bgzf); let (c_tx, c_rx) = pool.compress_result_channel(); // Submit one compress job @@ -5604,11 +4502,10 @@ mod tests { input_eof: false, decompressed_input_done: false, compress_has_items: false, - consumer_parked: false, phase: phase::PHASE1, }; let priorities = get_sort_priorities(&bp); - assert!(rank(priorities, SortStep::DecompressInput) < rank(priorities, SortStep::Compress)); + assert_eq!(priorities[0], SortStep::DecompressInput); } #[test] @@ -5618,33 +4515,22 @@ mod tests { input_eof: false, decompressed_input_done: false, compress_has_items: true, - consumer_parked: false, phase: phase::PHASE1, }; let priorities = get_sort_priorities(&bp); - assert!(rank(priorities, SortStep::Compress) < rank(priorities, SortStep::DecompressInput)); + assert_eq!(priorities[0], SortStep::Compress); } #[test] - fn test_sort_priorities_phase1_all_done_offers_no_input_or_compress_work() { - // Once the input is read and decompressed and nothing is queued to - // compress, no worker should be offered work that produces more of it — - // that is what makes them back off instead of spinning. `ExtractKeys` is - // exempt and may appear: it self-gates on a non-empty batch queue, and a - // batch queued just before EOF must still be reachable while the ingest - // thread drains at the barrier. + fn test_sort_priorities_phase1_all_done_returns_empty() { let bp = SortBackpressureState { decompressed_input_low: false, input_eof: true, decompressed_input_done: true, compress_has_items: false, - consumer_parked: false, phase: phase::PHASE1, }; - let priorities = get_sort_priorities(&bp); - for step in [SortStep::ReadInputBlocks, SortStep::DecompressInput, SortStep::Compress] { - assert!(!priorities.contains(&step), "{step:?} must not be offered once input is done"); - } + assert!(get_sort_priorities(&bp).is_empty()); } #[test] @@ -5654,11 +4540,10 @@ mod tests { input_eof: false, decompressed_input_done: false, compress_has_items: false, - consumer_parked: false, phase: phase::PHASE2, }; let priorities = get_sort_priorities(&bp); - assert!(rank(priorities, SortStep::Phase2FileWork) < rank(priorities, SortStep::Compress)); + assert_eq!(priorities[0], SortStep::Phase2FileWork); } #[test] @@ -5668,11 +4553,10 @@ mod tests { input_eof: false, decompressed_input_done: false, compress_has_items: true, - consumer_parked: false, phase: phase::PHASE2, }; let priorities = get_sort_priorities(&bp); - assert!(rank(priorities, SortStep::Compress) < rank(priorities, SortStep::DecompressInput)); + assert_eq!(priorities[0], SortStep::Compress); } #[test] @@ -5684,135 +4568,29 @@ mod tests { input_eof: false, decompressed_input_done: false, compress_has_items: false, - consumer_parked: false, phase: phase::PHASE2, }; let priorities = get_sort_priorities(&bp); - assert!(rank(priorities, SortStep::Phase2FileWork) < rank(priorities, SortStep::Compress)); - } - - /// A pending byte slice is what some thread holding an exclusive reader is - /// blocked on, and that thread reclaims the slice if nobody takes it -- so - /// the whole benefit lives in the short window between offer and reclaim. - /// Anything scheduled above `FetchBytes` spends that window, which is why - /// it leads every list. This is the mirror image of `ExtractKeys`, which - /// goes last everywhere because a deferred key batch costs nothing until - /// the chunk barrier. - #[test] - fn test_fetch_bytes_outranks_every_other_step_in_every_state() { - for phase in [phase::PHASE1, phase::PHASE2, phase::SHUTDOWN] { - for &(low, eof, done, compress, parked) in &[ - (false, false, false, false, false), - (true, false, false, true, false), - (false, true, true, false, false), - (false, false, false, true, true), - (true, true, true, true, true), - ] { - let bp = SortBackpressureState { - decompressed_input_low: low, - input_eof: eof, - decompressed_input_done: done, - compress_has_items: compress, - consumer_parked: parked, - phase, - }; - assert_eq!( - get_sort_priorities(&bp)[0], - SortStep::FetchBytes, - "phase {phase} state {low}{eof}{done}{compress}{parked}" - ); - } - } - } - - /// Where `step` sits in a priority list, for tests that mean "A outranks B" - /// rather than "A is literally first". Index-based assertions broke the - /// moment `FetchBytes` was inserted above them, though nothing they were - /// asserting had changed. - fn rank(priorities: &[SortStep], step: SortStep) -> usize { - priorities.iter().position(|&s| s == step).unwrap_or(usize::MAX) - } - - /// Output compression is throughput work any worker can do at any time. The - /// awaited block is the only thing that can unblock the merge, so a parked - /// consumer must outrank compression. - #[test] - fn test_sort_priorities_phase2_parked_consumer_outranks_compression() { - let bp = SortBackpressureState { - decompressed_input_low: false, - input_eof: false, - decompressed_input_done: false, - compress_has_items: true, - consumer_parked: true, - phase: phase::PHASE2, - }; - let p = get_sort_priorities(&bp); - assert!(rank(p, SortStep::Phase2FileWork) < rank(p, SortStep::Compress)); + assert_eq!(priorities[0], SortStep::Phase2FileWork); } - /// With the consumer running, compress-first is preserved exactly: it is the - /// writer-side bottleneck and draining it is what keeps output moving. #[test] - fn test_sort_priorities_phase2_keeps_compress_first_when_consumer_runs() { + fn test_sort_priorities_legacy_returns_compress_only() { let bp = SortBackpressureState { decompressed_input_low: false, input_eof: false, decompressed_input_done: false, compress_has_items: true, - consumer_parked: false, - phase: phase::PHASE2, - }; - let p = get_sort_priorities(&bp); - assert!(rank(p, SortStep::Compress) < rank(p, SortStep::Phase2FileWork)); - } - - /// Both steps stay reachable either way -- dropping one would starve it - /// outright rather than deprioritize it, and output must keep draining even - /// while the consumer is blocked. - #[rstest] - #[case::consumer_parked(true)] - #[case::consumer_running(false)] - fn test_sort_priorities_phase2_always_offers_both_steps(#[case] consumer_parked: bool) { - let bp = SortBackpressureState { - decompressed_input_low: false, - input_eof: false, - decompressed_input_done: false, - compress_has_items: true, - consumer_parked, - phase: phase::PHASE2, - }; - let p = get_sort_priorities(&bp); - assert!(p.contains(&SortStep::Compress), "compression must stay reachable"); - assert!(p.contains(&SortStep::Phase2FileWork), "file work must stay reachable"); - } - - #[test] - fn test_sort_priorities_legacy_drains_compress_and_no_phased_work() { - // The legacy/transition phase exists to drain queued jobs, so compression - // must be reachable and the phase-specific steps must not be. As above, - // `ExtractKeys` self-gates and is exempt. - let bp = SortBackpressureState { - decompressed_input_low: false, - input_eof: false, - decompressed_input_done: false, - compress_has_items: true, - consumer_parked: false, phase: phase::LEGACY, }; let priorities = get_sort_priorities(&bp); - assert!( - rank(priorities, SortStep::Compress) < rank(priorities, SortStep::Phase2FileWork), - "compression drains first" - ); - for step in [SortStep::ReadInputBlocks, SortStep::DecompressInput, SortStep::Phase2FileWork] - { - assert!(!priorities.contains(&step), "{step:?} belongs to a real phase"); - } + assert_eq!(priorities.len(), 1); + assert_eq!(priorities[0], SortStep::Compress); } #[test] fn test_worker_pool_num_workers() { - let pool = SortWorkerPool::new(3, 1, 6, crate::codec::SpillCodec::Bgzf, false); + let pool = SortWorkerPool::new(3, 1, 6, crate::codec::SpillCodec::Bgzf); assert_eq!(pool.num_workers(), 3); pool.shutdown(); } @@ -5827,7 +4605,7 @@ mod tests { fn empty_phase2_file() -> Phase2FileState { let tmp = tempfile::tempfile().expect("failed to create tempfile"); let reader = BufReader::with_capacity(1024, tmp); - Phase2FileState::new(crate::spill_reader::SpillSource::Sequential(reader), SpillCodec::Bgzf) + Phase2FileState::new(reader, SpillCodec::Bgzf) } /// Build a tiny placeholder raw entry whose contents we never decode. @@ -6330,12 +5108,13 @@ mod tests { // for BGZF (whose decoder reads the gzip header itself). // ======================================================================== - /// Snapshot the position of the next byte the per-file reader will serve. - /// Both `SpillSource` arms account for whatever they have already buffered, - /// so on a reader nothing has read from this is where it will start. + /// Snapshot the file position by locking the per-file reader. `BufReader`'s + /// `stream_position` accounts for any buffered bytes — for a fresh + /// `BufReader` whose buffer hasn't been filled this equals the underlying + /// `File`'s seek position. fn phase2_file_position(state: &Phase2FileState) -> u64 { let mut guard = state.reader.lock().expect("reader lock"); - guard.inner.position().expect("position") + guard.inner.stream_position().expect("stream_position") } #[test] @@ -6348,26 +5127,17 @@ mod tests { file.write_all(&[0xAA, 0xBB, 0xCC]).expect("write body"); drop(file); - // All three shapes: the sequential reader, a pinned scattered one, and - // the auto-tuning one that is now the default. - for read_streams in [ - crate::external::ReadStreams::Fixed(1), - crate::external::ReadStreams::Fixed(4), - crate::external::ReadStreams::Auto, - ] { - let mut pool = SortWorkerPool::new(1, 1, 6, SpillCodec::Bgzf, false); - pool.read_streams = read_streams; - pool.set_phase2_files(std::slice::from_ref(&path)).expect("set_phase2_files"); - let files = pool.phase2_files(); - assert_eq!(files.len(), 1); - assert_eq!(files[0].codec, SpillCodec::Zstd, "ZSPILL_MAGIC must select zstd codec"); - assert_eq!( - phase2_file_position(&files[0]), - ZSPILL_MAGIC.len() as u64, - "zstd reader must start past the 4-byte magic at {read_streams:?}" - ); - pool.shutdown(); - } + let pool = SortWorkerPool::new(1, 1, 6, SpillCodec::Bgzf); + pool.set_phase2_files(std::slice::from_ref(&path)).expect("set_phase2_files"); + let files = pool.phase2_files(); + assert_eq!(files.len(), 1); + assert_eq!(files[0].codec, SpillCodec::Zstd, "ZSPILL_MAGIC must select zstd codec"); + assert_eq!( + phase2_file_position(&files[0]), + ZSPILL_MAGIC.len() as u64, + "zstd reader must be positioned past the 4-byte magic" + ); + pool.shutdown(); } #[test] @@ -6381,26 +5151,17 @@ mod tests { file.write_all(&[0x1f, 0x8b, 0x00, 0x00, 0x55, 0x66]).expect("write magic"); drop(file); - // All three shapes: the sequential reader, a pinned scattered one, and - // the auto-tuning one that is now the default. - for read_streams in [ - crate::external::ReadStreams::Fixed(1), - crate::external::ReadStreams::Fixed(4), - crate::external::ReadStreams::Auto, - ] { - let mut pool = SortWorkerPool::new(1, 1, 6, SpillCodec::Zstd, false); - pool.read_streams = read_streams; - pool.set_phase2_files(std::slice::from_ref(&path)).expect("set_phase2_files"); - let files = pool.phase2_files(); - assert_eq!(files.len(), 1); - assert_eq!(files[0].codec, SpillCodec::Bgzf, "BGZF magic must select bgzf codec"); - assert_eq!( - phase2_file_position(&files[0]), - 0, - "bgzf reader must start at byte 0 at {read_streams:?}, for the header" - ); - pool.shutdown(); - } + let pool = SortWorkerPool::new(1, 1, 6, SpillCodec::Zstd); + pool.set_phase2_files(std::slice::from_ref(&path)).expect("set_phase2_files"); + let files = pool.phase2_files(); + assert_eq!(files.len(), 1); + assert_eq!(files[0].codec, SpillCodec::Bgzf, "BGZF magic must select bgzf codec"); + assert_eq!( + phase2_file_position(&files[0]), + 0, + "bgzf reader must be rewound to byte 0 so the decoder sees the header" + ); + pool.shutdown(); } #[test] @@ -6409,7 +5170,7 @@ mod tests { let path = dir.path().join("empty.spill"); std::fs::File::create(&path).expect("create empty"); - let pool = SortWorkerPool::new(1, 1, 6, SpillCodec::Zstd, false); + let pool = SortWorkerPool::new(1, 1, 6, SpillCodec::Zstd); pool.set_phase2_files(std::slice::from_ref(&path)).expect("set_phase2_files"); let files = pool.phase2_files(); assert_eq!(files.len(), 1); @@ -6440,7 +5201,7 @@ mod tests { path }; - let pool = SortWorkerPool::new(1, 1, 6, SpillCodec::Bgzf, false); + let pool = SortWorkerPool::new(1, 1, 6, SpillCodec::Bgzf); // First merge: one source, drained to completion. pool.set_phase2_files(std::slice::from_ref(&spill("first.spill"))) @@ -6467,52 +5228,6 @@ mod tests { pool.shutdown(); } - /// A pool that merges twice must not carry the first merge's read-ahead - /// prediction into the second. - /// - /// `phase2_next_source` indexes the file vector `set_phase2_files` replaces - /// and hands the named source deep-read priority. A value left from a - /// completed merge would give that priority to an unrelated file while the - /// new merge is still seeding its sources, before it has published a - /// prediction of its own -- so a fresh source set must reset it to the - /// sentinel. - #[test] - fn test_set_phase2_files_clears_the_read_ahead_prediction() { - use std::io::Write; - let dir = tempfile::tempdir().expect("tempdir"); - let spill = |name: &str| { - let path = dir.path().join(name); - let mut file = std::fs::File::create(&path).expect("create"); - file.write_all(&[0x1f, 0x8b, 0x00, 0x00]).expect("write magic"); - path - }; - - let pool = SortWorkerPool::new(1, 1, 6, SpillCodec::Bgzf, false); - - // First merge: install a source and publish a prediction, as the merge - // loop does when the winning run changes. - pool.set_phase2_files(std::slice::from_ref(&spill("first.spill"))) - .expect("set_phase2_files"); - pool.set_phase2_next_source(Some(0)); - assert_eq!( - pool.shared.phase2_next_source.load(Ordering::Relaxed), - 0, - "guard: the first merge's prediction is published" - ); - - // Second merge: a fresh set must start with no prediction so no - // unrelated file inherits the previous merge's deep-read priority. - pool.set_phase2_files(&[spill("second-a.spill"), spill("second-b.spill")]) - .expect("set_phase2_files"); - assert_eq!( - pool.shared.phase2_next_source.load(Ordering::Relaxed), - NO_AWAITED_SOURCE, - "a fresh source set must clear the read-ahead prediction" - ); - - pool.shutdown(); - } - /// A pool that merges twice must size the second merge's refill allowance /// from the second merge's own blocks. /// @@ -6534,7 +5249,7 @@ mod tests { path }; - let pool = SortWorkerPool::new(1, 1, 6, SpillCodec::Bgzf, false); + let pool = SortWorkerPool::new(1, 1, 6, SpillCodec::Bgzf); pool.set_phase2_files(std::slice::from_ref(&spill("first.spill"))) .expect("set_phase2_files"); diff --git a/crates/fgumi-sort/src/zspill_stream.rs b/crates/fgumi-sort/src/zspill_stream.rs index 953a9d4bc..32b8fc73e 100644 --- a/crates/fgumi-sort/src/zspill_stream.rs +++ b/crates/fgumi-sort/src/zspill_stream.rs @@ -8,7 +8,7 @@ //! //! This module presents a [`Read`]-compatible view over the decompressed //! content of such a file, so the legacy -//! [`GenericKeyedChunkReader`](crate::external::GenericKeyedChunkReader) record +//! `GenericKeyedChunkReader` record //! parser can consume zstd-compressed spills without further changes. The //! worker-pool reader takes the faster, per-frame parallel-decompress path; //! this stream reader is the fallback used by the consolidation merge and the @@ -17,27 +17,14 @@ use crate::codec::ZSPILL_MAGIC; use std::io::{self, Read}; -/// Cap on the uncompressed size of a single zstd spill frame. -/// -/// If a frame decompresses to more than this, +/// Cap on the uncompressed size of a single zstd spill frame. Production +/// frames are bounded by the producer's staging buffer +/// (`BGZF_MAX_BLOCK_SIZE` + padding ~= 68 KB), so 256 KiB leaves comfortable +/// slack. If a frame ever decompresses to more bytes than this, /// `zstd::bulk::Decompressor::decompress_to_buffer` surfaces a clear error -/// rather than silently truncating. -/// -/// **Derived** from the pool's cap rather than kept in step with it by comment. -/// This reader serves the consolidation path while `worker_pool` serves the -/// merge, so a cap that is smaller than the writer's frame size fails only the -/// consolidating configurations -- the ones the standard matrix does not -/// exercise. `test_frame_caps_admit_the_largest_frame_the_writer_emits` pins -/// them together. -const FRAME_DECOMP_CAP: usize = crate::worker_pool::zstd_decomp_cap(); - -/// The frame cap this reader enforces, exposed so it can be pinned against the -/// writer's frame size rather than trusted to stay in step. -#[cfg(test)] -#[must_use] -pub(crate) const fn frame_decomp_cap() -> usize { - FRAME_DECOMP_CAP -} +/// rather than silently truncating. Reuses `worker_pool::ZSTD_FRAME_DECOMP_CAP` +/// so producer and consumer share one source and cannot drift. +const FRAME_DECOMP_CAP: usize = crate::worker_pool::ZSTD_FRAME_DECOMP_CAP; /// Streaming decompressor for "ZSP1" spill files. pub struct ZspillStreamReader { diff --git a/crates/fgumi-sort/tests/loom_merge_slots.rs b/crates/fgumi-sort/tests/loom_merge_slots.rs new file mode 100644 index 000000000..a20a07931 --- /dev/null +++ b/crates/fgumi-sort/tests/loom_merge_slots.rs @@ -0,0 +1,358 @@ +//! Loom model-check of the block-parallel Phase-2 decompress protocol +//! (`file_granularity == false`), driving the **real** `SortMergeSlot`. +//! +//! # What this verifies (and what it does NOT) +//! +//! Under `--cfg loom`, `merge_slots.rs` swaps `std::sync` -> `loom::sync`, so +//! the `SortMergeSlot` constructed here uses loom's atomics and mutexes. loom +//! explores the thread interleavings and the memory reorderings the C11 model +//! permits (preemption-bounded — see "Preemption-bounded exploration" below), +//! running the REAL slot methods each time: +//! +//! * [`SortMergeSlot::bp_commit_read`] — the publish order (reserve +//! `in_flight` before setting `reader_eof`) that the original silent- +//! truncation bug lived in. Production +//! (`SortSpillDecompress::try_fill_block_parallel_slot`) calls the SAME +//! method, so the model and the code cannot drift. +//! * [`SortMergeSlot::bp_insert_drain_finalize`] / +//! [`SortMergeSlot::bp_drain_and_finalize`] and the private +//! `drain_locked_and_finalize` (the finalize predicate `!queue_eof && +//! reader_eof && in_flight == 0 && reorder.is_empty()`, the lock order +//! reorder -> decompressed, and the real `in_flight.fetch_sub(AcqRel)`). +//! * The slot's REAL `reader` mutex serializes reads, and its REAL `reorder` +//! buffer ([`fgumi_bam_io::reorder::ReorderBuffer`]) reassembles them. +//! +//! Only two things are *not* the production code, both sound and documented: +//! +//! * **The "read" is simulated.** loom cannot model real file I/O, so a +//! worker computes `(start_seq, got, hit_eof)` from a test-held total under +//! the real `reader` lock instead of calling `read_raw`; "decompression" +//! yields the seq number as the payload. The slot's accounting/finalize +//! methods that run on the result are 100% production code. +//! * **Blocking `reader.lock()` instead of production's `try_lock()`.** The +//! property under test depends only on reads being *serialized* (so +//! `reader_eof` can never become visible while an unreserved block still +//! exists); a blocking lock preserves exactly that while collapsing the +//! `try_lock`-miss-retry fan-out that would otherwise explode loom's state +//! space (and a spin-retry is a loom anti-pattern). It is a sound +//! over-approximation of the serialization invariant. +//! +//! ## Modeling choices (documented for honesty) +//! +//! * **Window/FIFO admission disabled.** `bp_reorder_admits` backpressure is a +//! *bounded-memory* property, separately covered by the `merge_slots` unit +//! test `bp_reorder_window_is_bounded_under_straggler`. Workers here always +//! admit, keeping the model focused on the EOF/truncation invariants. +//! * **Tiny sizes (2-4 blocks, 2-3 workers).** loom's state space is +//! super-exponential; these sizes still exercise every out-of-order +//! insert/finalize interleaving the per-slot protocol can hit (more blocks +//! add only more of the same kind, not a new kind). +//! * **One pass per producer.** `n_workers == reads_needed(N, batch)`, so +//! every block is read by exactly one producer and every producer +//! terminates. The *consumer* does spin: [`consume_until_drained`] breaks +//! only on `is_drained()`, which requires `queue_eof`. So a protocol that +//! never finalizes does NOT reach [`run_model`]'s assertions — the consumer +//! never returns from its poll loop, `consumer.join()` never completes, and +//! the `queue_eof` assertion below is unreachable. It surfaces instead as a +//! loom `max_branches` panic from the spinning consumer, which names the +//! harness rather than the protocol. Read that failure as "`queue_eof` was +//! never finalized", not as a model that needs a bigger branch budget. +//! * **A concurrent merge consumer** ([`consume_until_drained`]) polls the +//! FIFO and STOPS at `is_drained()`, mirroring `SortMerge`. This is what +//! makes a *premature* `queue_eof` observable as truncation (without it the +//! straggler is appended after the join and the bug hides — verified: the +//! `bp_commit_read` order swap is caught only with the consumer present). +//! * **Preemption-bounded exploration.** The consumer's poll loop adds a +//! scheduling point per iteration; combined with 2-3 producers the fully +//! exhaustive state space is minutes-long. Every model is therefore explored +//! under a preemption bound (`Some(k)`) — a recognized technique: essentially +//! all real concurrency bugs (including the truncation race this guards) +//! manifest with ≤2-3 preemptions. The bound is verified to still catch the +//! `bp_commit_read` order swap. +//! +//! Run with: +//! ```text +//! RUSTFLAGS="--cfg loom" cargo test -p fgumi-sort --test loom_merge_slots --release +//! ``` +//! +//! # Complementary coverage and its residual +//! +//! This model is one leg of the block-parallel hardening; the others are the +//! `merge_slots` unit tests (bounded-memory window), the +//! `fgumi-pipeline-io` granularity/proptest/soak-matrix tests (the real +//! end-to-end pipeline over real spill files), and a `ThreadSanitizer` pass. +//! +//! **Sanitizer residual (recorded for honesty):** the `ThreadSanitizer` run +//! exercised the real pipeline on **arm64 only**, and the C decompression codecs +//! (`zstd` via the `zstd` crate, `libdeflate` via `libdeflater`) are +//! **uninstrumented** — the sanitizer only sees the Rust side, so a data race +//! *inside* a C codec would be missed. This is acceptable because the codecs are +//! pure per-block transforms with no shared mutable state across threads (each +//! worker decompresses its own block into its own buffer); the cross-thread +//! protocol the sanitizer and loom actually need to clear is the Rust-side slot +//! accounting, which is fully instrumented here. + +#![cfg(loom)] +#![deny(unsafe_code)] +// Block counts/seqs in this model are tiny (≤ a handful) and always fit a +// usize; the casts below are between u64 model seqs and usize counts. +#![allow(clippy::cast_possible_truncation)] + +use std::fs::File; +use std::io::BufReader; + +use fgumi_sort::{SortMergeSlot, SpillCodec}; +use loom::sync::Arc; +use loom::sync::atomic::Ordering; + +/// A throwaway reader for the slot. The block-parallel slot methods never touch +/// `reader.inner`; the model serializes reads on the `reader` mutex and computes +/// the read result arithmetically, so an empty file is all the struct needs. +fn empty_reader() -> BufReader { + BufReader::new(tempfile::tempfile().expect("create tempfile")) +} + +/// One worker's body: mirrors a single `try_run` of +/// `SortSpillDecompress::try_fill_block_parallel_slot`, but with the file read +/// simulated (see module docs). Reads up to `block_batch` blocks under the REAL +/// `reader` lock, publishes the accounting via the REAL +/// [`SortMergeSlot::bp_commit_read`], "decompresses" outside the lock, then +/// inserts/drains/finalizes via the REAL [`SortMergeSlot::bp_insert_drain_finalize`]. +/// A worker that finds the reader already at EOF falls through to the REAL +/// Phase-B drain-only [`SortMergeSlot::bp_drain_and_finalize`]. +fn worker_one_pass(slot: &SortMergeSlot, block_batch: u64, total_blocks: u64) { + if slot.queue_eof.load(Ordering::Acquire) { + return; + } + let mut did_phase_a = false; + if !slot.reader_eof.load(Ordering::Acquire) { + // Blocking lock (sound over-approximation of production's `try_lock`; + // see module docs) — serializes reads on the REAL reader mutex. + let mut reader = slot.reader.lock().unwrap(); + // Re-check under the lock: another worker may have hit EOF. + if !slot.reader_eof.load(Ordering::Acquire) { + let start_seq = reader.next_seq; + let remaining = total_blocks - start_seq; + let got = block_batch.min(remaining); + let hit_eof = got < block_batch; + // Stamp the read range and commit the accounting BEFORE releasing + // the lock, via the real publish-order method. + reader.next_seq += got; + slot.bp_commit_read(got as usize, hit_eof); + drop(reader); + + // "Decompress" outside the reader lock: the payload is the seq as + // 8 little-endian bytes, so the drained FIFO can be checked for + // in-order, no-loss delivery. + let blocks: Vec> = + (start_seq..start_seq + got).map(|s| s.to_le_bytes().to_vec()).collect(); + slot.bp_insert_drain_finalize(start_seq, blocks, got as usize); + did_phase_a = true; + } + } + if !did_phase_a { + slot.bp_drain_and_finalize(); + } +} + +/// Number of reader-lock acquisitions (= worker passes) needed to read every +/// block and then observe the clean EOF: `ceil((N + 1) / batch)`. The `+ 1` +/// accounts for the read that returns fewer than `batch` blocks (possibly +/// empty), which is what sets `reader_eof`. +fn reads_needed(total_blocks: u64, block_batch: u64) -> usize { + ((total_blocks + 1).div_ceil(block_batch)) as usize +} + +/// Decode an 8-byte little-endian seq payload back to its sequence number. +fn seq_of(block: &[u8]) -> u64 { + let mut buf = [0u8; 8]; + buf.copy_from_slice(&block[..8]); + u64::from_le_bytes(buf) +} + +/// The merge consumer, mirroring `SortMerge`/`slot_try_load_block`: pop every +/// available block, then STOP the instant the slot looks cleanly drained +/// (`is_drained()` == `queue_eof && FIFO empty && !error`). Returns the seqs it +/// collected, in pop (delivery) order. +/// +/// This stop condition is what makes a *premature* `queue_eof` observable: if +/// the slot finalizes EOF while a block is still outstanding (the truncation the +/// publish-order protocol prevents), the consumer sees an empty FIFO + EOF and +/// quits early, so `run_model`'s completeness check fails. Without a consumer +/// the straggler would still be appended after the join and the bug would hide. +/// +/// `yield_now` between polls is the loom scheduling point. Every producer runs +/// exactly one pass, so against a protocol that finalizes `queue_eof` this wait +/// terminates. Against one that does not, it spins forever — see the module +/// header's "One pass per producer" note for why that surfaces as a loom +/// `max_branches` panic rather than as the assertions in [`run_model`]. +fn consume_until_drained(slot: &SortMergeSlot) -> Vec { + let mut collected = Vec::new(); + loop { + loop { + let popped = slot.decompressed.lock().unwrap().pop_front(); + match popped { + Some(b) => collected.push(seq_of(&b)), + None => break, + } + } + if slot.is_drained() { + break; + } + loom::thread::yield_now(); + } + collected +} + +/// Drive one worker pass per required read over `total_blocks` blocks with +/// `block_batch` blocks per read PLUS a concurrent merge consumer, under loom, +/// and assert the no-loss / in-order / clean-EOF invariants for every +/// interleaving against the REAL slot state. +fn run_model(total_blocks: u64, block_batch: u64) { + let slot = Arc::new(SortMergeSlot::new(0, empty_reader(), SpillCodec::Bgzf)); + let n_workers = reads_needed(total_blocks, block_batch); + + let mut handles: Vec<_> = (0..n_workers) + .map(|_| { + let slot = Arc::clone(&slot); + loom::thread::spawn(move || worker_one_pass(&slot, block_batch, total_blocks)) + }) + .collect(); + let consumer = { + let slot = Arc::clone(&slot); + loom::thread::spawn(move || consume_until_drained(&slot)) + }; + for h in handles.drain(..) { + h.join().unwrap(); + } + let delivered = consumer.join().unwrap(); + + // Post-conditions: clean EOF reached, no error, nothing left in flight or + // buffered, and the consumer collected every block exactly once, in read + // order, BEFORE it observed the clean EOF (a premature `queue_eof` truncates + // `delivered`). + assert!(slot.queue_eof.load(Ordering::Acquire), "slot never reached queue_eof"); + assert!(!slot.has_error(), "spurious decomp_error"); + assert_eq!(slot.in_flight.load(Ordering::Acquire), 0, "blocks left in flight at EOF"); + assert!(slot.reorder.lock().unwrap().is_empty(), "reorder buffer not drained at EOF"); + assert!(slot.decompressed.lock().unwrap().is_empty(), "FIFO not fully consumed at EOF"); + + let expected: Vec = (0..total_blocks).collect(); + assert_eq!( + delivered, expected, + "blocks lost, duplicated, reordered, or truncated by early EOF" + ); +} + +/// Run `f` under loom with at most `preemption_bound` preemptions. See the +/// module-level "Preemption-bounded exploration" note for why every model is +/// bounded rather than exhaustive. +fn check_model(preemption_bound: usize, f: F) { + let mut builder = loom::model::Builder::new(); + builder.preemption_bound = Some(preemption_bound); + builder.check(f); +} + +/// Three blocks, batch 2 ⇒ 2 producer passes: the second read is SHORT (carries +/// seq 2 AND sets `reader_eof`) while the first read's blocks (seq 0,1) may +/// still be in flight — the exact `bp_eof_with_straggler` shape. +#[test] +fn loom_three_blocks_batch2() { + check_model(3, || run_model(3, 2)); +} + +/// Two blocks, batch 2 ⇒ 2 producer passes: the second read is the EMPTY +/// EOF-detecting read (got == 0) that sets `reader_eof` while the first read's +/// blocks (seq 0,1) may still be in flight. Exercises the `count == 0` finalize +/// path of `bp_insert_drain_finalize`. +#[test] +fn loom_two_blocks_batch2_empty_eof_read() { + check_model(3, || run_model(2, 2)); +} + +/// Four blocks, batch 3 ⇒ 2 producer passes (read seq 0,1,2; short read seq 3 +/// sets EOF). A larger in-flight batch straggling behind the EOF read. +#[test] +fn loom_four_blocks_batch3() { + check_model(3, || run_model(4, 3)); +} + +/// Two blocks, batch 1 ⇒ 3 producer passes (read seq 0, read seq 1, empty EOF +/// read) plus the consumer — four concurrent threads. Covers the three-way race +/// between two in-flight blocks and the EOF-setter. Bounded tighter (the +/// four-thread × nested-mutex state space is the largest here). +#[test] +fn loom_two_blocks_batch1_three_workers_bounded() { + check_model(2, || run_model(2, 1)); +} + +// ── decomp-error-beats-clean-EOF (#399) ────────────────────────────────────── + +/// The consumer that observes `queue_eof` under the `decompressed` mutex +/// (`is_drained`) must also observe `decomp_error` whenever the producer took +/// the error path — i.e. a failed slot can never be mistaken for a clean EOF. +/// +/// This drives the REAL slot: the producer mirrors +/// `SortSpillDecompress::mark_slot_failed` (store `decomp_error` then +/// `queue_eof`, both under the `decompressed` mutex), and the consumer is the +/// real [`SortMergeSlot::is_drained`] / [`SortMergeSlot::has_error`] pair. The +/// mutex release-acquire makes the two flags jointly visible regardless of which +/// the consumer reads first. +#[test] +fn loom_decomp_error_beats_clean_eof() { + loom::model(|| { + let slot = Arc::new(SortMergeSlot::new(0, empty_reader(), SpillCodec::Bgzf)); + + // Producer: error path. Mirrors `mark_slot_failed` — set decomp_error + // THEN queue_eof, both under the `decompressed` mutex. + let producer = { + let slot = Arc::clone(&slot); + loom::thread::spawn(move || { + let _g = slot.decompressed.lock().unwrap(); + slot.decomp_error.store(true, Ordering::Release); + slot.queue_eof.store(true, Ordering::Release); + }) + }; + + // Consumer: the real `is_drained()` / `has_error()`. + // + // The assertion is `!drained`, flatly, not `!(drained && !errored)`. + // The producer's only path is the error path, so there is no + // interleaving in which a clean drain is legal: before the producer + // runs, `queue_eof` is unset and `is_drained()` is false; after it, + // `decomp_error` is set and `is_drained()` must still be false. + // + // The composite form is what the old assertion used, and it could not + // fail: with the `decomp_error` guard present `drained` is always + // false, and with it deleted `errored` is true, so `drained && !errored` + // is unsatisfiable either way — the test stayed green against the exact + // bug it names. Asserting `!drained` alone restores the discrimination; + // deleting the guard from `is_drained` now fails this model. + // + // `has_error()` is still called, before `is_drained()` in one variant + // and after in the other, so the documented order-independence of the + // two reads is what the model actually exercises. + let consumer = { + let slot = Arc::clone(&slot); + loom::thread::spawn(move || { + let errored = slot.has_error(); + let drained = slot.is_drained(); + assert!(!drained, "errored slot reported a clean drain"); + assert!(!drained || errored, "clean EOF hid a decomp error"); + }) + }; + + let consumer_reversed = { + let slot = Arc::clone(&slot); + loom::thread::spawn(move || { + let drained = slot.is_drained(); + let _errored = slot.has_error(); + assert!(!drained, "errored slot reported a clean drain (reversed read order)"); + }) + }; + + producer.join().unwrap(); + consumer.join().unwrap(); + consumer_reversed.join().unwrap(); + }); +} diff --git a/crates/fgumi-umi/src/assigner.rs b/crates/fgumi-umi/src/assigner.rs index 84f8fd4a9..dd6d1c2c9 100644 --- a/crates/fgumi-umi/src/assigner.rs +++ b/crates/fgumi-umi/src/assigner.rs @@ -565,7 +565,7 @@ pub const BOTTOM_STRAND_DUPLEX: &str = "/B"; /// /// Determines how reads are grouped based on their UMI sequences. Each strategy makes /// different tradeoffs between speed, error tolerance, and grouping behavior. -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] #[cfg_attr(feature = "cli", derive(clap::ValueEnum))] pub enum Strategy { /// Only reads with identical UMI sequences are grouped together diff --git a/crates/xtask/Cargo.toml b/crates/xtask/Cargo.toml index edc0e408e..7e0936d58 100644 --- a/crates/xtask/Cargo.toml +++ b/crates/xtask/Cargo.toml @@ -18,7 +18,7 @@ simulate = ["fgumi/simulate"] anyhow = { workspace = true } clap = { workspace = true } syn = { version = "3", features = ["full", "extra-traits"] } -quote = "1" +quote = { workspace = true } fgumi = { path = "../.." } [dev-dependencies] diff --git a/docs/design/sort-phase2-unification-deferral.md b/docs/design/sort-phase2-unification-deferral.md new file mode 100644 index 000000000..4aae37766 --- /dev/null +++ b/docs/design/sort-phase2-unification-deferral.md @@ -0,0 +1,106 @@ +# Sort Phase-2 merge: two drivers, unification deferred + +## Status: DEFERRED (not a feat-runall blocker) + +fgumi's external sort has **two** Phase-2 (merge) drivers, and they exist by +design: + +- **Pool / standalone** (`crates/fgumi-sort/src/worker_pool.rs`, + `external.rs::merge_chunks_generic`) — drives standalone file-to-file + `fgumi sort`. Work-stealing pool reads + decompresses spill blocks in + parallel; the merge consumer is a **parked OS thread** + (`advance_to_next_block`). Keeps `--write-index`, `--verify`, zstd spill. +- **Streaming / runall** (`crates/fgumi-sort/src/merge_slots.rs`, + `external.rs::MergeDriver`) — drives the fused in-pipeline `runall` sort. A + cooperative `MergeDriverDyn::try_step` state machine that yields + (`Stalled`/`WouldBlock`) + instead of blocking, so it composes with the typed-step pipeline framework. + +> **Which of these runs in this tree.** Only the pool driver. Both halves of the +> streaming driver now exist — `merge_slots.rs::SortMergeSlot` and +> `external.rs::MergeDriver`, the latter arriving with the arena engine — but +> nothing in production drives them: `MergeDriver::from_slots` and +> `open_spill_slot` are reachable only from tests until the typed-step +> `SortMerge` consumer lands with `fgumi-pipeline-io`. `fgumi sort` and +> `fgumi merge` both go through `RawExternalSorter`, i.e. the pool path. +> +> So read the streaming bullet — and every claim below about the cooperative +> consumer, including the deadlock analysis — as describing the end state this +> landing builds toward. The analysis is about that consumer, not about +> `SortMergeSlot`, which on the inline path is a passive handoff slot rather than +> a driver (on the block-parallel path it owns reorder admission and EOF +> finalization — see below). + +Unifying these onto a single driver is a deliberate **future** refactor with its +own design + bench gate + deadlock re-analysis. It is intentionally NOT done in +the feat-runall line. + +## What is already shared (so the duplication is smaller than it looks) + +- `LoserTree` (`loser_tree.rs`) — the actual k-way merge / tie-break + primitive. Both drivers use the one implementation. +- `RawSortKey` comparators and the `keys.rs` key types. +- `ReorderBuffer` (`fgumi-bam-io`) — also used by Phase-1 read-ahead, so it is a + shared utility, not pool-specific duplication. + +What genuinely differs is the **driver** (the outer drive loop and the +read/decompress topology), not the merge core. + +## Why the streaming path dropped the pool's machinery + +The pre-v4 streaming slot also carried a `raw_blocks` FIFO, a `decomp_in_flight` +counter, a reorder buffer, and a cap+gap-filler admission rule. That design +**deadlocked at production scale**: the framework's drain protocol could Skip +workers during a transient "all slots at cap simultaneously" window. v4 collapsed +read-and-decompress into one inline op per worker per slot, which made a plain +bounded FIFO correct **for that inline path** (blocks decompress strictly in read +order) and removed the deadlock surface. + +The reorder buffer and the in-flight counter are **back** on `SortMergeSlot` +(`reorder`, `in_flight`) for the later **block-parallel** path, where several +workers decompress one file's blocks concurrently and results complete out of +order. There `queue_eof` finalizes only when `reader_eof && in_flight == 0 && +reorder.is_empty()`. The gap-filler escape is still gone. + +The pool path keeps all of that machinery and does not hit **that specific +deadlock** — the v4 cooperative-drain Skip described above — because (a) its +consumer is a parked OS thread the engine cannot Skip, and (b) it is +single-*reader* but **multi-*decompressor*** — workers decompress in parallel, +so completion order ≠ pop order, and the reorder buffer + in-flight counter are +a *correctness* dependency (they re-order out-of-order completions and close the +`is_drained` truncation race), not redundancy. + +That is a claim about one analyzed failure mode, not a proof of general deadlock +freedom: avoiding the cooperative `Skip` path is a property of the parked-thread +consumer, and no one has shown the pool path deadlock-free in general. + +## Prerequisites before unification is even attemptable + +1. The streaming path gains a true file-to-file drive (BAM source + + bgzf-decompress + parse → SortAndSpill → SortSpillDecompress → SortMerge → + serialize → WriteBgzfFile). +2. `--write-index` (`.bai`) wired onto/alongside that chain (the pool path has + it via `IndexBamFinalizeHook`, PR #393; the streaming step does not). +3. `--verify` ported (today a separate read-and-check path, designed to stay + legacy). +4. A consumer for the file-to-file case that is not exposed to the v4 drain + deadlock — the streaming consumer is the cooperative `try_step` model that hit + it in production. "Not exposed to *this* deadlock" is the bar; general + deadlock freedom is not something either driver has been shown to have. + +## Explicit risk + +Routing standalone sort through the streaming path **re-opens the v4 +cooperative-consumer deadlock** that standalone sort is not currently exposed to +(parked OS thread). Until all four prerequisites land, it is a net capability +LOSS + added risk and buys no shipped feature. + +## Safe-to-pursue-independently sub-step (does NOT require the above) + +Factor the shared **per-winner merge-emit core**: both drivers already share +`LoserTree`, so extract a `MergeSource` trait (`try_next_record`) + a single +`merge_emit(tree, sources, emit_fn)` core that the pool's blocking loop and the +streaming state machine both call one-winner-at-a-time. Keep the +blocking-vs-cooperative **outer** drive separate — that is the deadlock-relevant +axis and must NOT be merged. MED risk (hot-loop perf; needs a sort bench +matrix), but carries no deadlock surface. diff --git a/scripts/publish-crates.sh b/scripts/publish-crates.sh index 1a414f3da..9da54aed9 100755 --- a/scripts/publish-crates.sh +++ b/scripts/publish-crates.sh @@ -23,12 +23,17 @@ CRATES=( fgumi-bgzf fgumi-simd-fastq fgumi-tag + fgumi-fmt + fgumi-cli-macros + fgumi-cli-common + fgumi-pipeline-core fgumi-raw-bam fgumi-bam-io fgumi-umi fgumi-sam fgumi-metrics fgumi-sort + fgumi-pipeline-io fgumi-consensus fgumi ) diff --git a/src/lib/commands/clip.rs b/src/lib/commands/clip.rs index 2c1ef7084..7c508d36e 100644 --- a/src/lib/commands/clip.rs +++ b/src/lib/commands/clip.rs @@ -242,7 +242,7 @@ impl ClipParams { /// the template — including secondary/supplementary alignments — before the primary pair is /// clipped, matching fgbio `ClipBam`'s `template.allReads.foreach(upgradeAllClipping)` /// (`ClipBam.scala:123`). Doing this template-wide pre-pass here (rather than per-primary - /// inside [`clip_pair`]/[`clip_fragment`]) is what lets supplementary reads' clipping be + /// inside `clip_pair`/`clip_fragment`) is what lets supplementary reads' clipping be /// upgraded too, and keeps both threading paths in lockstep since both go through this method. /// /// This is the single shared implementation used by both the single-threaded and diff --git a/src/lib/commands/codec.rs b/src/lib/commands/codec.rs index 6601a0c27..a5cdef499 100644 --- a/src/lib/commands/codec.rs +++ b/src/lib/commands/codec.rs @@ -309,6 +309,97 @@ pub struct Codec { pub queue_memory: QueueMemoryOptions, } +// ───────────────────────────────────────────────────────────────────────────── +// CodecOptions — the stage's tuning knobs, projected out of the CLI struct +// ───────────────────────────────────────────────────────────────────────────── + +/// Codec-stage tuning, independent of how the values were supplied. +/// +/// See [`crate::commands::zipper::ZipperOptions`] for why this is a plain +/// struct rather than a flattened `clap::Args`. Note that the consensus-calling +/// knobs are held **flat** here even though [`Codec`] nests them behind +/// `#[command(flatten)]` sub-structs: the chain builder wants one bag per stage, +/// not a re-run of the CLI's grouping. +#[derive(Debug, Clone)] +pub struct CodecOptions { + /// Pre-UMI error rate (phred). + pub error_rate_pre_umi: u8, + /// Post-UMI error rate (phred). + pub error_rate_post_umi: u8, + /// Minimum input base quality. + pub min_input_base_quality: u8, + /// Emit per-base consensus tags. + pub output_per_base_tags: bool, + /// Trim consensus reads. + pub trim: bool, + /// Minimum consensus base quality. + pub min_consensus_base_quality: u8, + /// How to resolve a near-tie between the two most likely consensus bases. + pub tie_rule: fgumi_consensus::TieRule, + /// Minimum reads per consensus. + pub min_reads: usize, + /// Cap on reads per consensus. + pub max_reads: Option, + /// Minimum duplex overlap length. + pub min_duplex_length: usize, + /// Reproduce fgbio's legacy (pre-fgumi#761) overlap window for dovetailed + /// FR pairs. Off by default; see [`Codec::legacy_overlap_window`]. + pub legacy_overlap_window: bool, + /// Quality cap for single-strand positions. + pub single_strand_qual: Option, + /// Quality cap for outer bases. + pub outer_bases_qual: Option, + /// How many bases count as outer. + pub outer_bases_length: usize, + /// Maximum duplex disagreement rate. + pub max_duplex_disagreement_rate: f64, + /// Maximum duplex disagreements. + pub max_duplex_disagreements: Option, + /// Let fully-unmapped primary templates through the pre-group filter. + /// + /// Carried as the whole flattened sub-struct, like `io` / `rejects_opts` / + /// `read_group`, rather than as a bare `bool`. + pub allow_unmapped: AllowUnmappedOptions, + /// Input/output paths and reader mode. + pub io: BamIoOptions, + /// Optional rejects output. + pub rejects_opts: RejectsOptions, + /// Optional stats output. + pub stats_opts: StatsOptions, + /// Read-group identity for emitted reads. + pub read_group: ReadGroupOptions, +} + +impl Codec { + /// Project the parsed CLI flags into [`CodecOptions`]. + #[must_use] + pub fn to_codec_options(&self) -> CodecOptions { + CodecOptions { + error_rate_pre_umi: self.consensus.error_rate_pre_umi, + error_rate_post_umi: self.consensus.error_rate_post_umi, + min_input_base_quality: self.consensus.min_input_base_quality, + output_per_base_tags: self.consensus.output_per_base_tags, + trim: self.consensus.trim, + min_consensus_base_quality: self.consensus.min_consensus_base_quality, + tie_rule: self.consensus.tie_rule.into(), + min_reads: self.min_reads, + max_reads: self.max_reads, + min_duplex_length: self.min_duplex_length, + legacy_overlap_window: self.legacy_overlap_window, + single_strand_qual: self.single_strand_qual, + outer_bases_qual: self.outer_bases_qual, + outer_bases_length: self.outer_bases_length, + max_duplex_disagreement_rate: self.max_duplex_disagreement_rate, + max_duplex_disagreements: self.max_duplex_disagreements, + allow_unmapped: self.allow_unmapped.clone(), + io: self.io.clone(), + rejects_opts: self.rejects_opts.clone(), + stats_opts: self.stats_opts.clone(), + read_group: self.read_group.clone(), + } + } +} + /// Decide how the single-thread codec loop should handle a typed /// [`CodecConsensusError`]: silently swallow recoverable duplex-disagreement /// rejects (so the loop continues) and surface any other variant as a fatal @@ -889,6 +980,128 @@ impl Codec { #[cfg(test)] mod tests { use super::*; + + /// Every tuning flag must survive the projection into [`CodecOptions`]. + /// See the simplex counterpart for why this parses rather than constructs. + /// + /// Every value here is deliberately **non-default**, and every field of + /// `CodecOptions` is asserted. Both halves matter: an assertion that + /// compares a default against a default passes even when the projection + /// ignores the parsed field entirely, and a field left unasserted can be + /// dropped from the projection without any test noticing. + #[test] + fn to_codec_options_carries_every_tuning_flag() { + let cmd = Codec::try_parse_from([ + "codec", + "-i", + "in.bam", + "-o", + "out.bam", + "--error-rate-pre-umi", + "42", + "--error-rate-post-umi", + "37", + "--min-input-base-quality", + "16", + "--output-per-base-tags=false", + "--trim=true", + "--min-consensus-base-quality", + "23", + "--tie-rule", + "ulp-relative", + "--min-reads", + "4", + "--max-reads", + "88", + "--min-duplex-length", + "9", + "--legacy-overlap-window", + "--single-strand-qual", + "12", + "--outer-bases-qual", + "15", + "--outer-bases-length", + "7", + "--max-duplex-disagreement-rate", + "0.25", + "--max-duplex-disagreements", + "6", + "--rejects", + "rej.bam", + "--stats", + "stats.txt", + "--read-group-id", + "Z", + "--read-name-prefix", + "pfx", + "--allow-unmapped=true", + ]) + .expect("parses"); + + let opts = cmd.to_codec_options(); + + assert_eq!(opts.error_rate_pre_umi, 42); + assert_eq!(opts.error_rate_post_umi, 37); + assert_eq!(opts.min_input_base_quality, 16); + assert!(!opts.output_per_base_tags, "an explicit false must not be lost"); + assert!(opts.trim); + assert_eq!(opts.min_consensus_base_quality, 23); + assert_eq!( + opts.tie_rule, + fgumi_consensus::TieRule::UlpRelative, + "--tie-rule must reach the projection" + ); + assert_eq!(opts.min_reads, 4); + assert_eq!(opts.max_reads, Some(88)); + assert_eq!(opts.min_duplex_length, 9); + assert!(opts.legacy_overlap_window, "--legacy-overlap-window must reach the projection"); + assert_eq!(opts.single_strand_qual, Some(12)); + assert_eq!(opts.outer_bases_qual, Some(15)); + assert_eq!(opts.outer_bases_length, 7); + assert!((opts.max_duplex_disagreement_rate - 0.25).abs() < f64::EPSILON); + assert_eq!(opts.max_duplex_disagreements, Some(6)); + assert!(opts.allow_unmapped.enabled, "--allow-unmapped must reach the projection"); + // The flattened sub-structs must come across whole, not field by field. + assert_eq!(opts.io.input, std::path::PathBuf::from("in.bam")); + assert_eq!(opts.io.output, std::path::PathBuf::from("out.bam")); + assert_eq!(opts.rejects_opts.rejects, Some(std::path::PathBuf::from("rej.bam"))); + assert_eq!(opts.stats_opts.stats, Some(std::path::PathBuf::from("stats.txt"))); + assert_eq!(opts.read_group.read_group_id, "Z"); + assert_eq!(opts.read_group.read_name_prefix, Some("pfx".to_string())); + } + + /// The projection must carry defaults faithfully too — a field hard-coded to + /// the value the non-default test happens to pass would slip through it. + #[test] + fn to_codec_options_carries_defaults() { + let cmd = + Codec::try_parse_from(["codec", "-i", "in.bam", "-o", "out.bam"]).expect("parses"); + + let opts = cmd.to_codec_options(); + + assert_eq!(opts.error_rate_pre_umi, 45); + assert_eq!(opts.error_rate_post_umi, 40); + assert_eq!(opts.min_input_base_quality, 10); + assert!(opts.output_per_base_tags); + assert!(!opts.trim); + assert_eq!(opts.min_consensus_base_quality, 2); + assert_eq!(opts.tie_rule, fgumi_consensus::TieRule::FgbioCompat); + assert_eq!(opts.min_reads, 1); + assert_eq!(opts.max_reads, None); + assert_eq!(opts.min_duplex_length, 1); + assert!(!opts.legacy_overlap_window, "the default must be preserved, not hard-coded true"); + assert_eq!(opts.single_strand_qual, None); + assert_eq!(opts.outer_bases_qual, None); + assert_eq!(opts.outer_bases_length, 5); + assert!((opts.max_duplex_disagreement_rate - 1.0).abs() < f64::EPSILON); + assert_eq!(opts.max_duplex_disagreements, None); + assert!(!opts.allow_unmapped.enabled); + assert_eq!(opts.rejects_opts.rejects, None); + assert_eq!(opts.stats_opts.stats, None); + assert_eq!(opts.read_group.read_group_id, "A"); + assert_eq!(opts.read_group.read_name_prefix, None); + } + use noodles::sam::alignment::io::Write as AlignmentWrite; use rstest::rstest; use std::path::PathBuf; diff --git a/src/lib/commands/common.rs b/src/lib/commands/common.rs index d30a4b79f..f2a18da1d 100644 --- a/src/lib/commands/common.rs +++ b/src/lib/commands/common.rs @@ -963,7 +963,7 @@ pub struct ThreadingOptions { /// Options for output compression. /// /// Controls BGZF compression level for BAM output files. -#[derive(Debug, Clone, Default, Args)] +#[derive(Debug, Clone, Args)] pub struct CompressionOptions { /// Compression level for output BAM (0-12). /// @@ -975,6 +975,15 @@ pub struct CompressionOptions { pub compression_level: u32, } +impl Default for CompressionOptions { + /// Mirrors the clap `default_value_t = 1` so programmatic and default-constructed callers + /// emit level-1 compression rather than the `u32` default of `0`, which would silently + /// write uncompressed BGZF. + fn default() -> Self { + Self { compression_level: 1 } + } +} + /// Option controlling whether unmapped reads are processed by a consensus caller. /// /// Shared by the `simplex`, `duplex`, and `codec` commands so the `--allow-unmapped` flag @@ -1311,14 +1320,27 @@ fn resolve_memory_budget_with_total( ByteSize(margin as u64), ); } - log::debug!( - "Auto memory: {} of {} ({}/thread × {} threads, reserve {})", - ByteSize(budget as u64), - ByteSize(total as u64), - ByteSize((budget / threads) as u64), - threads, - ByteSize(margin as u64), - ); + // Only the per-thread arm allocates `budget` as `threads` independent slices; + // otherwise it is one shared pool and `budget / threads` would misdescribe it. + // Mirrors the framing `log_memory_config` already uses below. + if per_thread { + log::debug!( + "Auto memory: {} of {} ({}/thread × {} threads, reserve {})", + ByteSize(budget as u64), + ByteSize(total as u64), + ByteSize((budget / threads) as u64), + threads, + ByteSize(margin as u64), + ); + } else { + log::debug!( + "Auto memory: {} of {} (shared across {} threads, reserve {})", + ByteSize(budget as u64), + ByteSize(total as u64), + threads, + ByteSize(margin as u64), + ); + } budget } }; @@ -1884,6 +1906,15 @@ mod tests { assert!(status.success(), "mkfifo {} failed: {status}", path.display()); } + /// A default-constructed `CompressionOptions` must agree with the clap + /// `default_value_t = 1`. The derived `Default` would yield `0`, which is a valid but + /// *uncompressed* BGZF level, so programmatic callers would silently write uncompressed + /// BAM while the same command run from the CLI wrote level 1. + #[test] + fn compression_options_default_matches_the_cli_default() { + assert_eq!(CompressionOptions::default().compression_level, 1); + } + /// `index_threshold_log_message` reports the threshold that is actually in effect, /// which is not always the raw `--index-threshold` value. `Edit` floors a numeric /// flag at its own measured crossover and indexes only at one mismatch; `Adjacency` diff --git a/src/lib/commands/compare/bams.rs b/src/lib/commands/compare/bams.rs index b945d6557..49d6cc941 100644 --- a/src/lib/commands/compare/bams.rs +++ b/src/lib/commands/compare/bams.rs @@ -151,7 +151,7 @@ impl CommandPreset { /// `clip` rewrites CIGAR/SEQ/QUAL in the clipped span and, as a consequence, regenerates /// the alignment tags (`NM`/`UQ`/`MD`) and repairs mate-pair metadata — but neither ever /// recomputes the depth-tag values. All five compare positionally under - /// [`ContentPredicate::Exact`](super::engines::content::ContentPredicate::Exact): + /// [`ContentPredicate::Exact`]: /// because fgumi clamps the depth tags to fgbio's `Short` ceiling at the source, the /// tags are bit-identical and an exact comparison is both sound and complete — no /// consensus-specific predicate is needed. Since `Exact` holds every core SAM field @@ -167,7 +167,7 @@ impl CommandPreset { /// ([`super::engines::molecule_join::molecule_join_compare`]): molecules are matched by /// an MI-invariant canonical id (no re-sort — both inputs must already be grouped), /// and each matched pair is checked for record membership, content under - /// [`ContentPredicate::ExactMinusMi`](super::engines::content::ContentPredicate::ExactMinusMi) + /// [`ContentPredicate::ExactMinusMi`] /// (everything except the MI tag), and duplex `/A`/`/B` strand-partition equivalence — /// the predicate excludes MI precisely because the canonical-id matching, not the /// content predicate, is what verifies MI equivalence. @@ -196,7 +196,7 @@ impl CommandPreset { /// The [`ContentPredicate`] to use for this preset's content comparison. See /// [`Self::resolve`] for the full resolution table and rationale. Note that for `Group` /// (resolved mode `CompareMode::Grouping`), this value is *not* consulted by the - /// molecule-join engine — [`engines::molecule_join::molecule_join_compare`] hardcodes + /// molecule-join engine — [`super::engines::molecule_join::molecule_join_compare`] hardcodes /// [`ContentPredicate::ExactMinusMi`] internally and is not configurable via /// `--command`/`--mode`. `Group` never reaches `execute_content`, but this method is /// still live for it (its resolved value is exercised only by the @@ -997,7 +997,7 @@ impl CompareBams { /// /// This is the sole `CompareMode::Grouping` path (and what `--command group` uses). Both /// inputs are cut into per-molecule runs and matched across the two files by an - /// MI-invariant canonical id (see [`molecule_join_compare`]) — no re-sort, so both inputs + /// MI-invariant canonical id (see [`super::engines::molecule_join::molecule_join_compare`]) — no re-sort, so both inputs /// must already be grouped (same-MI reads consecutive, as `fgumi group`/`fgbio group` /// output is). This makes the comparison inherently order-independent at the molecule /// level — `--ignore-order` is therefore a no-op under `--mode grouping`. Each matched @@ -1009,7 +1009,7 @@ impl CompareBams { /// /// # Errors /// - /// Returns an error if either input cannot be read (see [`molecule_join_compare`]), or + /// Returns an error if either input cannot be read (see [`super::engines::molecule_join::molecule_join_compare`]), or /// [`super::CompareMismatch`] if the two BAMs are found to differ (non-zero exit via the /// `Command` trait). fn execute_grouping(&self, input1: OpenedInput, input2: OpenedInput) -> Result { diff --git a/src/lib/commands/compare/engines/mod.rs b/src/lib/commands/compare/engines/mod.rs index a2b18d870..8f12b4110 100644 --- a/src/lib/commands/compare/engines/mod.rs +++ b/src/lib/commands/compare/engines/mod.rs @@ -1,7 +1,7 @@ //! Comparison engines for `fgumi compare bams`. //! //! An "engine" pairs records from the two input streams (see -//! [`positional`](self::positional), added alongside the positional-alignment work) and +//! [`positional`], added alongside the positional-alignment work) and //! decides whether each pair is content-equal via a [`content::ContentPredicate`]. //! Splitting pairing (positional) from equality (content) keeps the pairing logic honest: //! it can never quietly resync past a mismatch just because two *unrelated* records diff --git a/src/lib/commands/compare/engines/positional.rs b/src/lib/commands/compare/engines/positional.rs index 0bd35730d..7d416670c 100644 --- a/src/lib/commands/compare/engines/positional.rs +++ b/src/lib/commands/compare/engines/positional.rs @@ -15,8 +15,8 @@ //! coincidence, not evidence of parity, so it must never be allowed to mask //! the desync. //! - If the keys agree, the pair is compared under a -//! [`ContentPredicate`](super::content::ContentPredicate) via -//! [`content_diffs`](super::content::content_diffs). +//! [`ContentPredicate`] via +//! [`content_diffs`]. //! //! A record-count mismatch (one file longer than the other) is reported as a //! presence difference, independent of and in addition to any key mismatch. @@ -53,7 +53,7 @@ pub struct PositionalOutcome { /// pairing stopped at the first desync rather than attempting to resync. pub key_mismatch_at: Option, /// `true` if the two inputs' `@HD`/`@SQ`/`@RG` headers disagreed on a field - /// [`compare_headers`](super::header::compare_headers) considers significant (`@PG`/`@CO` + /// [`compare_headers`] considers significant (`@PG`/`@CO` /// are normalized and never contribute here). pub header_mismatch: bool, /// Human-readable diff strings (header, key mismatch, content diffs, and/or the @@ -85,7 +85,7 @@ impl PositionalOutcome { /// always reflect the true totals even after pairing has stopped. /// /// Also compares the two inputs' headers via -/// [`compare_headers`](super::header::compare_headers) (`@HD`/`@SQ`/`@RG`, normalizing away +/// [`compare_headers`] (`@HD`/`@SQ`/`@RG`, normalizing away /// `@PG`/`@CO`); a significant divergence is folded into /// `PositionalOutcome::header_mismatch`/`PositionalOutcome::is_match` alongside the /// record-level findings. diff --git a/src/lib/commands/compare/engines/sort_verify.rs b/src/lib/commands/compare/engines/sort_verify.rs index 303463398..c37e43a00 100644 --- a/src/lib/commands/compare/engines/sort_verify.rs +++ b/src/lib/commands/compare/engines/sort_verify.rs @@ -106,7 +106,7 @@ pub struct SortVerifyOutcome { /// compared (no resync), though both streams are still drained for accurate counts. pub presence_mismatch: bool, /// `true` if the two inputs' `@HD`/`@SQ`/`@RG` headers disagreed on a field - /// [`compare_headers`](super::header::compare_headers) considers significant (`@PG`/`@CO` + /// [`compare_headers`] considers significant (`@PG`/`@CO` /// are normalized and never contribute here). Note that `@HD` `SO`/`GO`/`SS` agreement is /// already implied by `detect_sort_order` succeeding on both inputs with a matching /// [`SortOrder`] (checked before this field is ever populated) *only when both writers use @@ -1072,7 +1072,7 @@ where /// by maximal equal-core-sort-key run (see the module docs). /// /// Never re-sorts either input. Also compares the two inputs' headers via -/// [`compare_headers`](super::header::compare_headers) (`@HD`/`@SQ`/`@RG`, normalizing away +/// [`compare_headers`] (`@HD`/`@SQ`/`@RG`, normalizing away /// `@PG`/`@CO`); a significant divergence is folded into /// [`SortVerifyOutcome::header_mismatch`]/[`SortVerifyOutcome::is_match`] alongside the /// run-comparison findings. `max_diffs` caps the number of entries collected in diff --git a/src/lib/commands/correct.rs b/src/lib/commands/correct.rs index 55d874338..1956c56b5 100644 --- a/src/lib/commands/correct.rs +++ b/src/lib/commands/correct.rs @@ -335,8 +335,65 @@ pub struct CorrectUmis { pub queue_memory: QueueMemoryOptions, } +// ───────────────────────────────────────────────────────────────────────────── +// CorrectOptions — the stage's tuning knobs, projected out of the CLI struct +// ───────────────────────────────────────────────────────────────────────────── + +/// CorrectUmis-stage tuning, independent of how the values were supplied. +/// +/// See [`crate::commands::zipper::ZipperOptions`] for why this is a plain +/// struct rather than a flattened `clap::Args`. Note that the rejects path is +/// held **flat** here even though [`CorrectUmis`] nests it behind a +/// `#[command(flatten)]` sub-struct: the chain builder wants one bag per stage, +/// not a re-run of the CLI's grouping. +#[derive(Debug, Clone)] +pub struct CorrectOptions { + /// Optional metrics output. + pub metrics: Option, + /// Which SAM tag is corrected: `RX`/`OX` for UMIs, `BC`/`ob` for barcodes. + pub target: Target, + /// Maximum mismatches when matching a UMI. + pub max_mismatches: usize, + /// Minimum distance to the runner-up UMI. + pub min_distance_diff: usize, + /// Expected UMI sequences. + pub umis: Vec, + /// Files holding expected UMI sequences. + pub umi_files: Vec, + /// Skip storing the original UMI. + pub dont_store_original_umis: bool, + /// UMI match cache size. + pub cache_size: usize, + /// Minimum corrected fraction before failing. + pub min_corrected: Option, + /// Also match the reverse complement. + pub revcomp: bool, + /// Optional rejects output path. + pub rejects_path: Option, +} + +impl CorrectUmis { + /// Project the parsed CLI flags into [`CorrectOptions`]. + #[must_use] + pub fn to_correct_options(&self) -> CorrectOptions { + CorrectOptions { + metrics: self.metrics.clone(), + target: self.target, + max_mismatches: self.max_mismatches, + min_distance_diff: self.min_distance_diff, + umis: self.umis.clone(), + umi_files: self.umi_files.clone(), + dont_store_original_umis: self.dont_store_original_umis, + cache_size: self.cache_size, + min_corrected: self.min_corrected, + revcomp: self.revcomp, + rejects_path: self.rejects_opts.rejects.clone(), + } + } +} + #[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] -enum RejectionReason { +pub(crate) enum RejectionReason { WrongLength, Mismatched, #[default] @@ -348,9 +405,9 @@ enum RejectionReason { /// This struct holds the result of correcting a UMI once for an entire template, /// which can then be applied to all records in that template. #[derive(Debug)] -struct TemplateCorrection { +pub(crate) struct TemplateCorrection { /// Whether the UMI matched successfully. - matched: bool, + pub(crate) matched: bool, /// The corrected UMI string (if matched). corrected_umi: Option, /// The original UMI string. @@ -360,9 +417,9 @@ struct TemplateCorrection { /// Whether there were actual mismatches (not just revcomp). has_mismatches: bool, /// Match details for metrics. - matches: Vec, + pub(crate) matches: Vec, /// Rejection reason if not matched. - rejection_reason: RejectionReason, + pub(crate) rejection_reason: RejectionReason, } // ============================================================================ @@ -400,17 +457,37 @@ impl MemoryEstimate for CorrectProcessedBatch { /// Metrics collected from UMI correction processing, aggregated post-pipeline. #[derive(Default)] -struct CollectedCorrectMetrics { +pub(crate) struct CollectedCorrectMetrics { /// Total templates processed. - templates_processed: u64, + pub(crate) templates_processed: u64, /// Records with missing UMI tag. - missing_umis: u64, + pub(crate) missing_umis: u64, /// Records with wrong UMI length. - wrong_length: u64, + pub(crate) wrong_length: u64, /// Records that didn't match any fixed UMI. - mismatched: u64, + pub(crate) mismatched: u64, /// Per-UMI match counts (for metrics file). - umi_matches: AHashMap, + pub(crate) umi_matches: AHashMap, +} + +// Slot aggregation for the typed-step correction path +// (`pipeline::steps::correct`). Per-UMI crediting is NOT re-derived here -- +// the step calls `CorrectUmis::credit_umi_metrics` directly, so both paths +// share one definition of fgbio's per-segment accounting. +impl CollectedCorrectMetrics { + /// Drain `other` into `self`, summing counts. + /// + /// Used to aggregate the per-thread accumulator slots once the pipeline has + /// drained. `other` is left empty. + pub(crate) fn merge_into(&mut self, other: &mut CollectedCorrectMetrics) { + self.templates_processed += std::mem::take(&mut other.templates_processed); + self.missing_umis += std::mem::take(&mut other.missing_umis); + self.wrong_length += std::mem::take(&mut other.wrong_length); + self.mismatched += std::mem::take(&mut other.mismatched); + for (umi, counts) in other.umi_matches.drain() { + merge_umi_counts(&mut self.umi_matches, umi, &counts); + } + } } impl Command for CorrectUmis { @@ -640,7 +717,7 @@ impl CorrectUmis { /// Compute UMI correction for a template (called once per template). #[allow(clippy::too_many_arguments)] - fn compute_template_correction( + pub(crate) fn compute_template_correction( umi: &str, umi_length: usize, revcomp: bool, @@ -745,7 +822,7 @@ impl CorrectUmis { /// /// `num_records` scales fgbio's per-record accounting to fgumi's /// per-template batching (both R1 and R2 of a template share one UMI). - fn credit_umi_metrics( + pub(crate) fn credit_umi_metrics( matches: &[UmiMatch], num_records: u64, unmatched_umi: &str, @@ -783,7 +860,7 @@ impl CorrectUmis { /// - Records have different UMIs /// - Some records have UMIs and others don't /// - UMI tag has non-string type - fn extract_and_validate_template_umi_raw( + pub(crate) fn extract_and_validate_template_umi_raw( raw_records: &[RawRecord], umi_tag: [u8; 2], ) -> anyhow::Result> { @@ -849,7 +926,7 @@ impl CorrectUmis { } /// Apply UMI correction to a raw BAM record. - fn apply_correction_to_raw( + pub(crate) fn apply_correction_to_raw( record: &mut RawRecord, correction: &TemplateCorrection, umi_tag: [u8; 2], @@ -1701,6 +1778,99 @@ pub fn find_umi_pairs_within_distance( #[cfg(test)] mod tests { use super::*; + + /// Every tuning flag must survive the projection into [`CorrectOptions`]. + /// + /// Every value here is deliberately **non-default**, and every field of + /// `CorrectOptions` is asserted. Both halves matter: an assertion that + /// compares a default against a default passes even when the projection + /// ignores the parsed field entirely, and a field left unasserted can be + /// dropped from the projection without any test noticing. + /// + /// `rejects_path` is the one field that changes shape: it is read out of the + /// flattened `RejectsOptions`. + #[test] + fn to_correct_options_carries_every_tuning_flag() { + let cmd = CorrectUmis::try_parse_from([ + "correct", + "-i", + "in.bam", + "-o", + "out.bam", + "-u", + "ACGT", + "-u", + "TTTT", + "-U", + "umis.txt", + "--target", + "barcode", + "--max-mismatches", + "5", + "--min-distance", + "3", + "--cache-size", + "4096", + "--min-corrected", + "0.75", + "--revcomp=true", + "--dont-store-original=true", + "--rejects", + "rej.bam", + "--metrics", + "m.txt", + ]) + .expect("parses"); + + let opts = cmd.to_correct_options(); + + assert_eq!(opts.metrics, Some(std::path::PathBuf::from("m.txt"))); + assert_eq!(opts.target, Target::Barcode, "--target must reach the projection"); + assert_eq!(opts.max_mismatches, 5); + assert_eq!(opts.min_distance_diff, 3); + assert_eq!(opts.umis, vec!["ACGT".to_string(), "TTTT".to_string()]); + assert_eq!(opts.umi_files, vec![std::path::PathBuf::from("umis.txt")]); + assert!(opts.dont_store_original_umis); + assert_eq!(opts.cache_size, 4096); + assert_eq!(opts.min_corrected, Some(0.75)); + assert!(opts.revcomp); + assert_eq!( + opts.rejects_path, + Some(std::path::PathBuf::from("rej.bam")), + "rejects_path must be read from the flattened RejectsOptions", + ); + } + + /// The projection must carry defaults faithfully too — a field hard-coded to + /// the value the non-default test happens to pass would slip through it. + #[test] + fn to_correct_options_carries_defaults() { + let cmd = CorrectUmis::try_parse_from([ + "correct", + "-i", + "in.bam", + "-o", + "out.bam", + "-u", + "ACGT", + "--min-distance", + "1", + ]) + .expect("parses"); + + let opts = cmd.to_correct_options(); + + assert_eq!(opts.metrics, None); + assert_eq!(opts.target, Target::Umi); + assert_eq!(opts.max_mismatches, 2); + assert!(opts.umi_files.is_empty()); + assert!(!opts.dont_store_original_umis); + assert_eq!(opts.cache_size, 100_000); + assert_eq!(opts.min_corrected, None); + assert!(!opts.revcomp); + assert_eq!(opts.rejects_path, None); + } + use noodles::sam; use noodles::sam::alignment::io::Write as SamWrite; use noodles::sam::alignment::record_buf::RecordBuf; diff --git a/src/lib/commands/dedup.rs b/src/lib/commands/dedup.rs index 456a75def..4f67f82bf 100644 --- a/src/lib/commands/dedup.rs +++ b/src/lib/commands/dedup.rs @@ -306,8 +306,8 @@ struct LadderLibraryState { duplicate_templates: u64, /// Next cumulative `templates_seen` value that triggers a snapshot row. next_threshold: u64, - /// `templates_seen` at the last emitted row, used by [`Self::finish`] (via - /// [`DuplicationLadderRecorder::finish`]) to avoid a duplicate final row + /// `templates_seen` at the last emitted row, used by + /// [`DuplicationLadderRecorder::finish`] to avoid a duplicate final row /// when the true total already landed exactly on an interval crossing, and /// to size each snapshot's window (`templates_seen - last_emitted_at`). last_emitted_at: u64, diff --git a/src/lib/commands/duplex.rs b/src/lib/commands/duplex.rs index 1832e819d..5c1742d19 100644 --- a/src/lib/commands/duplex.rs +++ b/src/lib/commands/duplex.rs @@ -238,6 +238,86 @@ pub struct Duplex { pub reference: Option, } +// ───────────────────────────────────────────────────────────────────────────── +// DuplexOptions — the stage's tuning knobs, projected out of the CLI struct +// ───────────────────────────────────────────────────────────────────────────── + +/// Duplex-stage tuning, independent of how the values were supplied. +/// +/// See [`crate::commands::zipper::ZipperOptions`] for why this is a plain +/// struct rather than a flattened `clap::Args`. Note that the consensus-calling +/// knobs are held **flat** here even though [`Duplex`] nests them behind +/// `#[command(flatten)]` sub-structs: the chain builder wants one bag per stage, +/// not a re-run of the CLI's grouping. +#[derive(Debug, Clone)] +pub struct DuplexOptions { + /// Pre-UMI error rate (phred). + pub error_rate_pre_umi: u8, + /// Post-UMI error rate (phred). + pub error_rate_post_umi: u8, + /// Minimum input base quality. + pub min_input_base_quality: u8, + /// Emit per-base consensus tags. + pub output_per_base_tags: bool, + /// Trim consensus reads. + pub trim: bool, + /// Minimum consensus base quality. + pub min_consensus_base_quality: u8, + /// How to resolve a near-tie between the two most likely consensus bases. + pub tie_rule: fgumi_consensus::TieRule, + /// Call overlapping bases jointly. + pub consensus_call_overlapping_bases: bool, + /// Minimum reads per consensus, per tier. + pub min_reads: Vec, + /// Cap on reads per strand. + pub max_reads_per_strand: Option, + /// Let fully-unmapped primary templates through the pre-group filter. + /// + /// Carried as the whole flattened sub-struct, like `io` / `rejects_opts` / + /// `read_group`, rather than as a bare `bool`. + pub allow_unmapped: AllowUnmappedOptions, + /// Input/output paths and reader mode. + pub io: BamIoOptions, + /// Optional rejects output. + pub rejects_opts: RejectsOptions, + /// Optional stats output. + pub stats_opts: StatsOptions, + /// Read-group identity for emitted reads. + pub read_group: ReadGroupOptions, + /// Resolved methylation calling mode (`Disabled` when the flag is unset). + pub methylation_mode: fgumi_consensus::MethylationMode, + /// Reference FASTA for methylation-aware modes. + pub reference: Option, +} + +impl Duplex { + /// Project the parsed CLI flags into [`DuplexOptions`]. + #[must_use] + pub fn to_duplex_options(&self) -> DuplexOptions { + DuplexOptions { + error_rate_pre_umi: self.consensus.error_rate_pre_umi, + error_rate_post_umi: self.consensus.error_rate_post_umi, + min_input_base_quality: self.consensus.min_input_base_quality, + output_per_base_tags: self.consensus.output_per_base_tags, + trim: self.consensus.trim, + min_consensus_base_quality: self.consensus.min_consensus_base_quality, + tie_rule: self.consensus.tie_rule.into(), + consensus_call_overlapping_bases: self.overlapping.consensus_call_overlapping_bases, + min_reads: self.min_reads.clone(), + max_reads_per_strand: self.max_reads_per_strand, + allow_unmapped: self.allow_unmapped.clone(), + io: self.io.clone(), + rejects_opts: self.rejects_opts.clone(), + stats_opts: self.stats_opts.clone(), + read_group: self.read_group.clone(), + methylation_mode: crate::commands::common::resolve_methylation_mode( + self.methylation_mode, + ), + reference: self.reference.clone(), + } + } +} + impl Command for Duplex { /// Executes the duplex consensus calling pipeline. /// @@ -980,6 +1060,110 @@ fn has_both_strands_raw(records: &[RawRecord]) -> bool { #[cfg(test)] mod tests { use super::*; + + /// Every tuning flag must survive the projection into [`DuplexOptions`]. + /// See the simplex counterpart for why this parses rather than constructs. + #[test] + fn to_duplex_options_carries_every_tuning_flag() { + let cmd = Duplex::try_parse_from([ + "duplex", + "-i", + "in.bam", + "-o", + "out.bam", + "--error-rate-pre-umi", + "41", + "--error-rate-post-umi", + "36", + "--min-input-base-quality", + "18", + "--output-per-base-tags=false", + "--trim=true", + "--min-consensus-base-quality", + "21", + "--tie-rule", + "ulp-relative", + "--consensus-call-overlapping-bases=false", + "--min-reads", + "3,2,1", + "--max-reads-per-strand", + "55", + "--rejects", + "rej.bam", + "--stats", + "stats.txt", + "--read-group-id", + "Z", + "--read-name-prefix", + "pfx", + "--methylation-mode", + "em-seq", + "--ref", + "ref.fa", + "--allow-unmapped=true", + ]) + .expect("parses"); + + let opts = cmd.to_duplex_options(); + + assert_eq!(opts.error_rate_pre_umi, 41); + assert_eq!(opts.error_rate_post_umi, 36); + assert_eq!(opts.min_input_base_quality, 18); + assert!(!opts.output_per_base_tags, "an explicit false must not be lost"); + assert!(opts.trim); + assert_eq!(opts.min_consensus_base_quality, 21); + assert_eq!( + opts.tie_rule, + fgumi_consensus::TieRule::UlpRelative, + "--tie-rule must reach the projection" + ); + assert!(!opts.consensus_call_overlapping_bases); + assert_eq!(opts.min_reads, vec![3, 2, 1]); + assert_eq!(opts.max_reads_per_strand, Some(55)); + assert!(opts.allow_unmapped.enabled, "--allow-unmapped must reach the projection"); + assert_eq!( + opts.methylation_mode, + fgumi_consensus::MethylationMode::EmSeq, + "--methylation-mode must reach the projection", + ); + assert_eq!(opts.reference, Some(std::path::PathBuf::from("ref.fa"))); + // The flattened sub-structs must come across whole, not field by field. + assert_eq!(opts.io.input, std::path::PathBuf::from("in.bam")); + assert_eq!(opts.io.output, std::path::PathBuf::from("out.bam")); + assert_eq!(opts.rejects_opts.rejects, Some(std::path::PathBuf::from("rej.bam"))); + assert_eq!(opts.stats_opts.stats, Some(std::path::PathBuf::from("stats.txt"))); + assert_eq!(opts.read_group.read_group_id, "Z"); + assert_eq!(opts.read_group.read_name_prefix, Some("pfx".to_string())); + } + + /// The projection must carry defaults faithfully too — a field hard-coded to + /// the value the non-default test happens to pass would slip through it. + #[test] + fn to_duplex_options_carries_defaults() { + let cmd = + Duplex::try_parse_from(["duplex", "-i", "in.bam", "-o", "out.bam"]).expect("parses"); + + let opts = cmd.to_duplex_options(); + + assert_eq!(opts.error_rate_pre_umi, 45); + assert_eq!(opts.error_rate_post_umi, 40); + assert_eq!(opts.min_input_base_quality, 10); + assert!(opts.output_per_base_tags); + assert!(!opts.trim); + assert_eq!(opts.min_consensus_base_quality, 2); + assert_eq!(opts.tie_rule, fgumi_consensus::TieRule::FgbioCompat); + assert!(opts.consensus_call_overlapping_bases); + assert_eq!(opts.min_reads, vec![1]); + assert_eq!(opts.max_reads_per_strand, None); + assert!(!opts.allow_unmapped.enabled); + assert_eq!(opts.methylation_mode, fgumi_consensus::MethylationMode::Disabled); + assert_eq!(opts.reference, None); + assert_eq!(opts.rejects_opts.rejects, None); + assert_eq!(opts.stats_opts.stats, None); + assert_eq!(opts.read_group.read_group_id, "A"); + assert_eq!(opts.read_group.read_name_prefix, None); + } + use anyhow::Result; use fgumi_bam_io::{create_bam_reader, create_bam_writer}; use fgumi_raw_bam::{ diff --git a/src/lib/commands/extract.rs b/src/lib/commands/extract.rs index 3cc843cf4..80c93b542 100644 --- a/src/lib/commands/extract.rs +++ b/src/lib/commands/extract.rs @@ -916,7 +916,7 @@ impl Extract { /// . /// /// An old-style Casava (<1.8) `/1` / `/2` read-number suffix is stripped from - /// the returned name (see [`strip_read_suffix`](crate::fastq_parse::strip_read_suffix)) + /// the returned name (see [`strip_read_suffix`]) /// so both mates of a pair share an identical QNAME, matching fgbio's `FastqSource`. /// Stripping happens before UMI extraction so a read-number digit never leaks into the UMI. /// diff --git a/src/lib/commands/filter.rs b/src/lib/commands/filter.rs index 6af84d4ea..303ca80c8 100644 --- a/src/lib/commands/filter.rs +++ b/src/lib/commands/filter.rs @@ -218,6 +218,80 @@ pub struct Filter { pub queue_memory: QueueMemoryOptions, } +// ───────────────────────────────────────────────────────────────────────────── +// FilterOptions — the stage's tuning knobs, projected out of the CLI struct +// ───────────────────────────────────────────────────────────────────────────── + +/// Filter-stage tuning, independent of how the values were supplied. +/// +/// See [`crate::commands::zipper::ZipperOptions`] for why this is a plain +/// struct rather than a flattened `clap::Args`: the chain builder only reads +/// these values, so moving the fields off [`Filter`] would rewrite this module +/// and its tests for no gain here. +#[derive(Debug, Clone)] +#[allow(clippy::struct_excessive_bools)] +pub struct FilterOptions { + /// Reference FASTA, required by methylation-aware filters. + pub reference: Option, + /// Minimum reads supporting a consensus, per depth tier. + pub min_reads: Vec, + /// Maximum per-read error rate, per depth tier. + pub max_read_error_rate: Vec, + /// Maximum per-base error rate, per depth tier. + pub max_base_error_rate: Vec, + /// Minimum consensus base quality. + pub min_base_quality: Option, + /// Minimum mean base quality across the read. + pub min_mean_base_quality: Option, + /// Maximum fraction of no-called bases. + pub max_no_call_fraction: f64, + /// Reverse per-base tags on negative-strand reads. + pub reverse_per_base_tags: bool, + /// Filter whole templates rather than individual reads. + pub filter_by_template: bool, + /// Optional path for rejected records. + pub rejects: Option, + /// Optional path for filter statistics. + pub stats: Option, + /// Require both single-strand consensuses to agree. + pub require_single_strand_agreement: bool, + /// Minimum methylation depth, per tier. + pub min_methylation_depth: Vec, + /// Require both strands to agree on methylation. + pub require_strand_methylation_agreement: bool, + /// Minimum bisulfite conversion fraction. + pub min_conversion_fraction: Option, + /// Resolved methylation calling mode (`Disabled` when the flag is unset). + pub methylation_mode: fgumi_consensus::MethylationMode, +} + +impl Filter { + /// Project the parsed CLI flags into [`FilterOptions`]. + #[must_use] + pub fn to_filter_options(&self) -> FilterOptions { + FilterOptions { + reference: self.reference.clone(), + min_reads: self.min_reads.clone(), + max_read_error_rate: self.max_read_error_rate.clone(), + max_base_error_rate: self.max_base_error_rate.clone(), + min_base_quality: self.min_base_quality, + min_mean_base_quality: self.min_mean_base_quality, + max_no_call_fraction: self.max_no_call_fraction, + reverse_per_base_tags: self.reverse_per_base_tags, + filter_by_template: self.filter_by_template, + rejects: self.rejects.clone(), + stats: self.stats.clone(), + require_single_strand_agreement: self.require_single_strand_agreement, + min_methylation_depth: self.min_methylation_depth.clone(), + require_strand_methylation_agreement: self.require_strand_methylation_agreement, + min_conversion_fraction: self.min_conversion_fraction, + methylation_mode: crate::commands::common::resolve_methylation_mode( + self.methylation_mode, + ), + } + } +} + // ============================================================================ // 7-Step Pipeline Types // ============================================================================ @@ -1188,6 +1262,104 @@ fn progress_heartbeat_total(before: u64, records: u64) -> Option { #[allow(clippy::float_cmp)] mod tests { use super::*; + + /// Every tuning flag must survive the projection into [`FilterOptions`]. + /// + /// Driven through `try_parse_from` rather than a struct literal: a literal + /// would still compile if a flag were renamed or unwired, whereas parsing + /// pins the whole path from command line to option struct. Non-default + /// values throughout, so a field read from the wrong source fails rather + /// than coincidentally matching its default. + #[test] + fn to_filter_options_carries_every_tuning_flag() { + let cmd = Filter::try_parse_from([ + "filter", + "-i", + "in.bam", + "-o", + "out.bam", + "--ref", + "ref.fa", + "--min-reads", + "3,2,1", + "--max-read-error-rate", + "0.06,0.07", + "--max-base-error-rate", + "0.05", + "--min-base-quality", + "13", + "--min-mean-base-quality", + "22.5", + "--max-no-call-fraction", + "0.3", + "--reverse-per-base-tags=true", + "--filter-by-template=false", + "--rejects", + "rej.bam", + "--stats", + "stats.txt", + "--require-single-strand-agreement=true", + "--min-methylation-depth", + "4,5", + "--require-strand-methylation-agreement=true", + "--min-conversion-fraction", + "0.9", + "--methylation-mode", + "em-seq", + ]) + .expect("parses"); + + let opts = cmd.to_filter_options(); + + assert_eq!(opts.reference, Some(std::path::PathBuf::from("ref.fa"))); + assert_eq!(opts.min_reads, vec![3, 2, 1]); + assert_eq!(opts.max_base_error_rate, vec![0.05]); + assert_eq!(opts.min_base_quality, Some(13)); + assert_eq!(opts.min_mean_base_quality, Some(22.5)); + assert!((opts.max_no_call_fraction - 0.3).abs() < f64::EPSILON); + assert!(opts.reverse_per_base_tags); + assert!(!opts.filter_by_template, "an explicit false must not be lost"); + assert_eq!(opts.rejects, Some(std::path::PathBuf::from("rej.bam"))); + assert_eq!(opts.stats, Some(std::path::PathBuf::from("stats.txt"))); + assert!(opts.require_single_strand_agreement); + assert_eq!(opts.min_methylation_depth, vec![4, 5]); + assert!(opts.require_strand_methylation_agreement); + assert_eq!(opts.min_conversion_fraction, Some(0.9)); + assert_eq!(opts.max_read_error_rate, vec![0.06, 0.07]); + assert_eq!( + opts.methylation_mode, + fgumi_consensus::MethylationMode::EmSeq, + "--methylation-mode must reach the projection", + ); + } + + /// The projection must carry defaults faithfully too — a field hard-coded to + /// the value the non-default test happens to pass would slip through it. + #[test] + fn to_filter_options_carries_defaults() { + let cmd = + Filter::try_parse_from(["filter", "-i", "in.bam", "-o", "out.bam"]).expect("parses"); + + let opts = cmd.to_filter_options(); + + assert_eq!(opts.reference, None); + assert!(opts.min_reads.is_empty()); + assert_eq!(opts.max_read_error_rate, vec![0.025]); + assert_eq!(opts.max_base_error_rate, vec![0.1]); + assert_eq!(opts.min_base_quality, None); + assert_eq!(opts.min_mean_base_quality, None); + assert!((opts.max_no_call_fraction - 0.2).abs() < f64::EPSILON); + assert!(!opts.reverse_per_base_tags); + assert!(opts.filter_by_template); + assert_eq!(opts.rejects, None); + assert_eq!(opts.stats, None); + assert!(!opts.require_single_strand_agreement); + assert!(opts.min_methylation_depth.is_empty()); + assert!(!opts.require_strand_methylation_agreement); + assert_eq!(opts.min_conversion_fraction, None); + assert_eq!(opts.methylation_mode, fgumi_consensus::MethylationMode::Disabled); + } + use crate::sam::SamTag; use fgumi_raw_bam::{RawRecord, SamBuilder as RawSamBuilder, aux_data_slice, flags}; use noodles::sam::alignment::record_buf::RecordBuf; diff --git a/src/lib/commands/group.rs b/src/lib/commands/group.rs index d10fcccb7..f51001968 100644 --- a/src/lib/commands/group.rs +++ b/src/lib/commands/group.rs @@ -721,6 +721,103 @@ pub struct GroupReadsByUmi { pub memory_report_interval: u64, } +// ───────────────────────────────────────────────────────────────────────────── +// GroupOptions — the stage's tuning knobs, projected out of the CLI struct +// ───────────────────────────────────────────────────────────────────────────── + +/// Group-stage tuning, independent of how the values were supplied. +/// +/// See [`crate::commands::zipper::ZipperOptions`] for why this is a plain +/// struct rather than a flattened `clap::Args`. +/// +/// Two fields hold *resolved* values rather than raw flags, because grouping +/// cannot be configured from the raw ones alone: `min_map_q` applies the +/// default that [`GroupReadsByUmi`] would otherwise apply at run time, and +/// `effective_strategy` / `effective_edits` carry the `--no-umi` and +/// identity-implies-zero-edits rules. Both come from the same methods +/// `execute` uses, so the command and the chain builder cannot drift apart. +#[derive(Debug, Clone)] +pub struct GroupOptions { + /// Minimum mapping quality for mapped reads, with the default applied. + pub min_map_q: u8, + /// Include non-PF reads. + pub include_non_pf_reads: bool, + /// Allow fully unmapped templates. + pub allow_unmapped: bool, + /// The strategy as requested on the command line. + pub strategy: Strategy, + /// The edit distance as requested on the command line. + pub edits: u32, + /// Minimum UMI length to accept. + pub min_umi_length: Option, + /// When to build the N-gram/BK-tree index instead of scanning linearly. + /// + /// This is [`fgumi_umi::IndexThreshold`] rather than a bare count: the + /// flag also accepts `always` / `never`, which a number cannot express. + pub index_threshold: IndexThreshold, + /// Skip UMI-based grouping entirely. + pub no_umi: bool, + /// Template-count floor for handing a position group to a parallel assigner. + pub parallel_group_min_templates: Option, + /// The strategy actually used, after applying `--no-umi`. + pub effective_strategy: Strategy, + /// The edit distance actually used, after applying `--no-umi` and the + /// identity-implies-zero rule. + pub effective_edits: u32, + /// Optional family-size histogram output. + pub family_size_histogram: Option, + /// Optional grouping-metrics output. + pub grouping_metrics: Option, + /// Optional output prefix for the full set of metrics files. + pub metrics_prefix: Option, +} + +impl GroupReadsByUmi { + /// The minimum mapping quality to apply, defaulting when the flag is absent. + #[must_use] + pub fn resolved_min_map_q(&self) -> u8 { + self.min_map_q.unwrap_or(1) + } + + /// Resolve the strategy and edit distance grouping will actually use. + /// + /// `--no-umi` forces identity grouping, and identity grouping requires an + /// edit distance of zero; both rules live here so `execute` and the chain + /// builder cannot disagree about what was configured. The caller is + /// responsible for rejecting `--no-umi` with `--strategy paired` and for + /// logging the override — this method only computes. + #[must_use] + pub fn resolve_strategy_and_edits(&self) -> (Strategy, u32) { + if self.no_umi { + return (Strategy::Identity, 0); + } + let edits = if matches!(self.strategy, Strategy::Identity) { 0 } else { self.edits }; + (self.strategy, edits) + } + + /// Project the parsed CLI flags into [`GroupOptions`]. + #[must_use] + pub fn to_group_options(&self) -> GroupOptions { + let (effective_strategy, effective_edits) = self.resolve_strategy_and_edits(); + GroupOptions { + min_map_q: self.resolved_min_map_q(), + include_non_pf_reads: self.include_non_pf_reads, + allow_unmapped: self.allow_unmapped, + strategy: self.strategy, + edits: self.edits, + min_umi_length: self.min_umi_length, + index_threshold: self.index_threshold, + no_umi: self.no_umi, + parallel_group_min_templates: self.parallel_group_min_templates.clone(), + effective_strategy, + effective_edits, + family_size_histogram: self.family_size_histogram.clone(), + grouping_metrics: self.grouping_metrics.clone(), + metrics_prefix: self.metrics.clone(), + } + } +} + /// Build [`UmiGroupingMetrics`] from filter metrics and family size counts. /// /// Shared by both the pipeline and single-threaded execution paths. @@ -802,23 +899,10 @@ impl Command for GroupReadsByUmi { } // Handle --no-umi mode: force identity strategy - let (effective_strategy, no_umi_edits_override) = if self.no_umi { - if !matches!(self.strategy, Strategy::Identity) { - info!("--no-umi mode: overriding strategy to identity"); - } - (Strategy::Identity, true) - } else { - (self.strategy, false) - }; - - // Identity strategy requires edits=0, others use the configured value - // Also force edits=0 in no-umi mode - let effective_edits = - if no_umi_edits_override || matches!(effective_strategy, Strategy::Identity) { - 0 - } else { - self.edits - }; + if self.no_umi && !matches!(self.strategy, Strategy::Identity) { + info!("--no-umi mode: overriding strategy to identity"); + } + let (effective_strategy, effective_edits) = self.resolve_strategy_and_edits(); // `--index-threshold always` asserts that indexing will happen; reject it when // the resolved strategy/edits can never index rather than ignoring the flag. @@ -832,7 +916,7 @@ impl Command for GroupReadsByUmi { self.io.validate()?; // Set minimum mapping quality - let min_mapq: u8 = self.min_map_q.unwrap_or(1); + let min_mapq: u8 = self.resolved_min_map_q(); // Initialize tracking infrastructure let timer = OperationTimer::new("Grouping reads by UMI"); @@ -1801,6 +1885,161 @@ fn with_extension(prefix: &Path, suffix: &str) -> PathBuf { #[cfg(test)] mod tests { use super::*; + + /// The `--no-umi` and identity-implies-zero-edits rules, pinned as a table. + /// + /// These moved out of `execute` into [`GroupReadsByUmi::resolve_strategy_and_edits`] + /// so the command and the chain builder read them from one place; this is + /// the regression test for that move. Identity forces zero edits because a + /// non-zero edit distance would silently change which reads group together, + /// and `--no-umi` forces identity because there is no UMI to compare. + #[rstest] + #[case::adjacency_passes_through(Strategy::Adjacency, 2, false, Strategy::Adjacency, 2)] + #[case::edit_passes_through(Strategy::Edit, 3, false, Strategy::Edit, 3)] + #[case::paired_passes_through(Strategy::Paired, 1, false, Strategy::Paired, 1)] + #[case::identity_forces_zero_edits(Strategy::Identity, 2, false, Strategy::Identity, 0)] + #[case::no_umi_forces_identity_and_zero(Strategy::Adjacency, 2, true, Strategy::Identity, 0)] + #[case::no_umi_over_identity(Strategy::Identity, 0, true, Strategy::Identity, 0)] + fn resolve_strategy_and_edits_applies_the_no_umi_and_identity_rules( + #[case] strategy: Strategy, + #[case] edits: u32, + #[case] no_umi: bool, + #[case] expected_strategy: Strategy, + #[case] expected_edits: u32, + ) { + let mut cmd = GroupReadsByUmi::try_parse_from([ + "group", + "-i", + "in.bam", + "-o", + "out.bam", + "-s", + "adjacency", + ]) + .expect("parses"); + cmd.strategy = strategy; + cmd.edits = edits; + cmd.no_umi = no_umi; + + assert_eq!(cmd.resolve_strategy_and_edits(), (expected_strategy, expected_edits)); + } + + /// `--min-map-q` is optional on the command line but not optional for + /// grouping, so the projection applies the same default `execute` does. + #[rstest] + #[case::absent_defaults_to_one(None, 1)] + #[case::zero_is_honored_not_treated_as_absent(Some(0), 0)] + #[case::explicit_value(Some(30), 30)] + fn resolved_min_map_q_applies_the_default(#[case] flag: Option, #[case] expected: u8) { + let mut cmd = GroupReadsByUmi::try_parse_from([ + "group", + "-i", + "in.bam", + "-o", + "out.bam", + "-s", + "adjacency", + ]) + .expect("parses"); + cmd.min_map_q = flag; + + assert_eq!(cmd.resolved_min_map_q(), expected); + } + + /// Every tuning flag must survive the projection into [`GroupOptions`]. + /// + /// Parsed rather than constructed so a renamed or unwired flag fails here. + /// `metrics_prefix` is the one renamed field — it is fed by `--metrics` — + /// and `effective_*` are the resolved ones, so both are asserted explicitly. + #[test] + fn to_group_options_carries_every_tuning_flag() { + let cmd = GroupReadsByUmi::try_parse_from([ + "group", + "-i", + "in.bam", + "-o", + "out.bam", + "-s", + "adjacency", + "-e", + "2", + "-m", + "30", + "--min-umi-length", + "8", + "--index-threshold", + "250", + "--family-size-histogram", + "fs.txt", + "--grouping-metrics", + "gm.txt", + "--metrics", + "prefix", + "--include-non-pf-reads=true", + "--allow-unmapped=true", + "--parallel-group-min-templates", + "500", + ]) + .expect("parses"); + + let opts = cmd.to_group_options(); + + assert_eq!(opts.min_map_q, 30); + assert!(opts.include_non_pf_reads); + assert!(opts.allow_unmapped); + assert_eq!(opts.strategy, Strategy::Adjacency); + assert_eq!(opts.edits, 2); + assert_eq!(opts.min_umi_length, Some(8)); + assert_eq!(opts.index_threshold, IndexThreshold::MinUmis(250)); + assert_eq!(opts.parallel_group_min_templates, Some(ParallelMinTemplates::Fixed(500))); + assert!(!opts.no_umi); + assert_eq!(opts.family_size_histogram, Some(std::path::PathBuf::from("fs.txt"))); + assert_eq!(opts.grouping_metrics, Some(std::path::PathBuf::from("gm.txt"))); + assert_eq!( + opts.metrics_prefix, + Some(std::path::PathBuf::from("prefix")), + "metrics_prefix is fed by --metrics, not by a flag of its own", + ); + + // Not overridden here, so the resolved pair matches what was requested. + assert_eq!(opts.effective_strategy, Strategy::Adjacency); + assert_eq!(opts.effective_edits, 2); + } + + /// The projection must carry defaults faithfully too — a field hard-coded to + /// the value the non-default test happens to pass would slip through it. + /// + /// `min_map_q` is the interesting one: the CLI holds `Option` and the + /// projection resolves the absent case to 1, so this pins the resolution + /// rather than the raw flag. + #[test] + fn to_group_options_carries_defaults() { + let cmd = GroupReadsByUmi::try_parse_from([ + "group", + "-i", + "in.bam", + "-o", + "out.bam", + "-s", + "adjacency", + ]) + .expect("parses"); + + let opts = cmd.to_group_options(); + + assert_eq!(opts.min_map_q, 1, "an absent --min-map-q resolves to 1"); + assert!(!opts.include_non_pf_reads); + assert!(!opts.allow_unmapped); + assert_eq!(opts.edits, 1); + assert_eq!(opts.min_umi_length, None); + assert_eq!(opts.index_threshold, IndexThreshold::MinUmis(100)); + assert_eq!(opts.parallel_group_min_templates, None); + assert!(!opts.no_umi); + assert_eq!(opts.family_size_histogram, None); + assert_eq!(opts.grouping_metrics, None); + assert_eq!(opts.metrics_prefix, None); + } + use crate::assigner::{IdentityUmiAssigner, PairedUmiAssigner, Strategy}; use crate::metrics::TemplateFilterReason; use bstr::BString; diff --git a/src/lib/commands/review.rs b/src/lib/commands/review.rs index 9069d59e8..7c09549de 100644 --- a/src/lib/commands/review.rs +++ b/src/lib/commands/review.rs @@ -721,7 +721,7 @@ impl Review { None } - /// Extracts the leading numeric value from a typed VCF sample [`Value`] (used to + /// Extracts the leading numeric value from a typed VCF sample `Value` (used to /// read `AF`). Scalars convert directly; arrays take the first present element; /// strings parse their first comma-separated token (fgbio does `AF.toDouble`). fn value_as_f64( diff --git a/src/lib/commands/simplex.rs b/src/lib/commands/simplex.rs index 6d02fb5a3..65a1d8905 100644 --- a/src/lib/commands/simplex.rs +++ b/src/lib/commands/simplex.rs @@ -249,6 +249,86 @@ pub struct Simplex { pub reference: Option, } +// ───────────────────────────────────────────────────────────────────────────── +// SimplexOptions — the stage's tuning knobs, projected out of the CLI struct +// ───────────────────────────────────────────────────────────────────────────── + +/// Simplex-stage tuning, independent of how the values were supplied. +/// +/// See [`crate::commands::zipper::ZipperOptions`] for why this is a plain +/// struct rather than a flattened `clap::Args`. Note that the consensus-calling +/// knobs are held **flat** here even though [`Simplex`] nests them behind +/// `#[command(flatten)]` sub-structs: the chain builder wants one bag per stage, +/// not a re-run of the CLI's grouping. +#[derive(Debug, Clone)] +pub struct SimplexOptions { + /// Pre-UMI error rate (phred). + pub error_rate_pre_umi: u8, + /// Post-UMI error rate (phred). + pub error_rate_post_umi: u8, + /// Minimum input base quality. + pub min_input_base_quality: u8, + /// Emit per-base consensus tags. + pub output_per_base_tags: bool, + /// Trim consensus reads. + pub trim: bool, + /// Minimum consensus base quality. + pub min_consensus_base_quality: u8, + /// How to resolve a near-tie between the two most likely consensus bases. + pub tie_rule: fgumi_consensus::TieRule, + /// Call overlapping bases jointly. + pub consensus_call_overlapping_bases: bool, + /// Minimum reads per consensus. + pub min_reads: usize, + /// Cap on reads per consensus. + pub max_reads: Option, + /// Let fully-unmapped primary templates through the pre-group filter. + /// + /// Carried as the whole flattened sub-struct, like `io` / `rejects_opts` / + /// `read_group`, rather than as a bare `bool`. + pub allow_unmapped: AllowUnmappedOptions, + /// Input/output paths and reader mode. + pub io: BamIoOptions, + /// Optional rejects output. + pub rejects_opts: RejectsOptions, + /// Optional stats output. + pub stats_opts: StatsOptions, + /// Read-group identity for emitted reads. + pub read_group: ReadGroupOptions, + /// Resolved methylation calling mode (`Disabled` when the flag is unset). + pub methylation_mode: fgumi_consensus::MethylationMode, + /// Reference FASTA for methylation-aware modes. + pub reference: Option, +} + +impl Simplex { + /// Project the parsed CLI flags into [`SimplexOptions`]. + #[must_use] + pub fn to_simplex_options(&self) -> SimplexOptions { + SimplexOptions { + error_rate_pre_umi: self.consensus.error_rate_pre_umi, + error_rate_post_umi: self.consensus.error_rate_post_umi, + min_input_base_quality: self.consensus.min_input_base_quality, + output_per_base_tags: self.consensus.output_per_base_tags, + trim: self.consensus.trim, + min_consensus_base_quality: self.consensus.min_consensus_base_quality, + tie_rule: self.consensus.tie_rule.into(), + consensus_call_overlapping_bases: self.overlapping.consensus_call_overlapping_bases, + min_reads: self.min_reads, + max_reads: self.max_reads, + allow_unmapped: self.allow_unmapped.clone(), + io: self.io.clone(), + rejects_opts: self.rejects_opts.clone(), + stats_opts: self.stats_opts.clone(), + read_group: self.read_group.clone(), + methylation_mode: crate::commands::common::resolve_methylation_mode( + self.methylation_mode, + ), + reference: self.reference.clone(), + } + } +} + impl Command for Simplex { fn execute(&self, command_line: &str) -> Result<()> { // Start timing @@ -828,6 +908,122 @@ impl Simplex { #[cfg(test)] mod tests { use super::*; + + /// Every tuning flag must survive the projection into [`SimplexOptions`]. + /// + /// Driven through `try_parse_from` rather than a struct literal: a literal + /// would still compile if a flag were renamed or unwired, whereas parsing + /// pins the whole path from command line to option struct. Non-default + /// values throughout, so a field read from the wrong source fails rather + /// than coincidentally matching its default. + #[test] + fn to_simplex_options_carries_every_tuning_flag() { + let cmd = Simplex::try_parse_from([ + "simplex", + "-i", + "in.bam", + "-o", + "out.bam", + "--error-rate-pre-umi", + "40", + "--error-rate-post-umi", + "35", + "--min-input-base-quality", + "17", + "--output-per-base-tags=false", + "--trim=true", + "--min-consensus-base-quality", + "19", + "--tie-rule", + "ulp-relative", + "--consensus-call-overlapping-bases=false", + "--min-reads", + "3", + "--max-reads", + "77", + "--rejects", + "rej.bam", + "--stats", + "stats.txt", + "--read-group-id", + "Z", + "--read-name-prefix", + "pfx", + "--methylation-mode", + "em-seq", + "--ref", + "ref.fa", + "--allow-unmapped=true", + ]) + .expect("parses"); + + let opts = cmd.to_simplex_options(); + + assert_eq!(opts.error_rate_pre_umi, 40); + assert_eq!(opts.error_rate_post_umi, 35); + assert_eq!(opts.min_input_base_quality, 17); + assert!(!opts.output_per_base_tags, "an explicit false must not be lost"); + assert!(opts.trim); + assert_eq!(opts.min_consensus_base_quality, 19); + assert_eq!( + opts.tie_rule, + fgumi_consensus::TieRule::UlpRelative, + "--tie-rule must reach the projection" + ); + assert!(!opts.consensus_call_overlapping_bases); + assert_eq!(opts.min_reads, 3); + assert_eq!(opts.max_reads, Some(77)); + assert!(opts.allow_unmapped.enabled, "--allow-unmapped must reach the projection"); + assert_eq!(opts.reference, Some(std::path::PathBuf::from("ref.fa"))); + assert_eq!( + opts.methylation_mode, + fgumi_consensus::MethylationMode::EmSeq, + "--methylation-mode must reach the projection", + ); + // The flattened sub-structs must come across whole, not field by field. + assert_eq!(opts.io.input, std::path::PathBuf::from("in.bam")); + assert_eq!(opts.io.output, std::path::PathBuf::from("out.bam")); + assert_eq!(opts.read_group.read_group_id, "Z"); + assert_eq!(opts.read_group.read_name_prefix, Some("pfx".to_string())); + assert_eq!(opts.rejects_opts.rejects, Some(std::path::PathBuf::from("rej.bam"))); + assert_eq!(opts.stats_opts.stats, Some(std::path::PathBuf::from("stats.txt"))); + } + + /// The projection must carry defaults faithfully too — a field hard-coded to + /// the value the non-default test happens to pass would slip through it. + #[test] + fn to_simplex_options_carries_defaults() { + let cmd = Simplex::try_parse_from([ + "simplex", + "-i", + "in.bam", + "-o", + "out.bam", + "--min-reads", + "1", + ]) + .expect("parses"); + + let opts = cmd.to_simplex_options(); + + assert_eq!(opts.error_rate_pre_umi, 45); + assert_eq!(opts.error_rate_post_umi, 40); + assert_eq!(opts.min_input_base_quality, 10); + assert!(opts.output_per_base_tags); + assert!(!opts.trim); + assert_eq!(opts.min_consensus_base_quality, 2); + assert_eq!(opts.tie_rule, fgumi_consensus::TieRule::FgbioCompat); + assert!(opts.consensus_call_overlapping_bases); + assert_eq!(opts.max_reads, None); + assert!(!opts.allow_unmapped.enabled); + assert_eq!(opts.methylation_mode, fgumi_consensus::MethylationMode::Disabled); + assert_eq!(opts.reference, None); + assert_eq!(opts.rejects_opts.rejects, None); + assert_eq!(opts.stats_opts.stats, None); + assert_eq!(opts.read_group.read_group_id, "A"); + assert_eq!(opts.read_group.read_name_prefix, None); + } + use crate::metrics::consensus::ConsensusKvMetric; use noodles::sam::alignment::record::data::field::Tag; use noodles::sam::alignment::record_buf::RecordBuf; diff --git a/src/lib/commands/simulate/common.rs b/src/lib/commands/simulate/common.rs index e6937fb94..038744bdd 100644 --- a/src/lib/commands/simulate/common.rs +++ b/src/lib/commands/simulate/common.rs @@ -434,7 +434,7 @@ impl ReferenceGenome { Some(subseq.to_vec()) } - /// Build a BAM [`Header`] with `@SQ` lines for every loaded contig. + /// Build a BAM `Header` with `@SQ` lines for every loaded contig. pub(super) fn build_bam_header(&self) -> noodles::sam::header::Header { use bstr::BString; use noodles::sam::header::Header; diff --git a/src/lib/commands/zipper.rs b/src/lib/commands/zipper.rs index 2625cdcc7..c62a6cef5 100644 --- a/src/lib/commands/zipper.rs +++ b/src/lib/commands/zipper.rs @@ -234,6 +234,62 @@ pub struct Zipper { pub restore_unconverted_bases: bool, } +// ───────────────────────────────────────────────────────────────────────────── +// ZipperOptions — the stage's tuning knobs, projected out of the CLI struct +// ───────────────────────────────────────────────────────────────────────────── + +/// Zipper-stage tuning, independent of how the values were supplied. +/// +/// The chain builder needs these knobs without knowing where they came from: +/// the standalone `fgumi zipper` command fills them from its own flags, and a +/// fused pipeline fills them from its own. Holding them in a plain struct — +/// rather than having the builder reach into [`Zipper`] — is what lets a chain +/// be constructed with no CLI struct behind it at all. +/// +/// This deliberately does **not** derive `clap::Args`. Flattening it into +/// [`Zipper`] would move the fields off that struct and rewrite every +/// `self.` reference in this module and its tests, which buys the chain +/// builder nothing — it only ever reads the values. That refactor belongs with +/// the fused command that actually needs prefixed `--zipper::` flags, so +/// the CLI surface here is untouched. +#[derive(Debug, Clone)] +pub struct ZipperOptions { + /// Tags to remove from mapped reads before copying unmapped tags. + pub tags_to_remove: Vec, + /// Tags to reverse for reads mapped to the negative strand. + pub tags_to_reverse: Vec, + /// Tags to reverse complement for reads mapped to the negative strand. + pub tags_to_revcomp: Vec, + /// Buffer size for the template channel. + pub buffer: usize, + /// Accepted for backward compatibility; has no effect. Carried so the + /// projection stays total — see [`Zipper::bwa_chunk_size`]. + pub bwa_chunk_size: u64, + /// Drop unmapped-BAM reads absent from the aligned BAM. + pub exclude_missing_reads: bool, + /// Skip adding `tc` tags to secondary/supplementary reads. + pub skip_tc_tags: bool, + /// Restore unconverted bases in EM-seq consensus reads. + pub restore_unconverted_bases: bool, +} + +impl Zipper { + /// Project the parsed CLI flags into [`ZipperOptions`]. + #[must_use] + pub fn to_zipper_options(&self) -> ZipperOptions { + ZipperOptions { + tags_to_remove: self.tags_to_remove.clone(), + tags_to_reverse: self.tags_to_reverse.clone(), + tags_to_revcomp: self.tags_to_revcomp.clone(), + buffer: self.buffer, + bwa_chunk_size: self.bwa_chunk_size, + exclude_missing_reads: self.exclude_missing_reads, + skip_tc_tags: self.skip_tc_tags, + restore_unconverted_bases: self.restore_unconverted_bases, + } + } +} + /// Builds the output BAM header from unmapped and mapped headers /// /// Merges information from both input headers: @@ -3928,6 +3984,71 @@ mod tests { assert_eq!(cmd.skip_tc_tags, expected); } + /// Every tuning flag must survive the projection into [`ZipperOptions`]. + /// + /// Driven through `try_parse_from` rather than a struct literal on purpose: + /// a struct literal would still compile if a flag were renamed or unwired, + /// whereas parsing pins the whole path from command line to option struct. + /// Non-default values throughout, so a field copied from the wrong source — + /// or left at its default — fails rather than coincidentally matching. + #[test] + fn to_zipper_options_carries_every_tuning_flag() { + let cmd = Zipper::try_parse_from([ + "zipper", + "-u", + "u.bam", + "-r", + "ref.fa", + "-o", + "out.bam", + "--tags-to-remove", + "RX,MI", + "--tags-to-reverse", + "QX", + "--tags-to-revcomp", + "OX,ZA", + "--buffer", + "1234", + "--bwa-chunk-size", + "4242", + "--exclude-missing-reads=true", + "--skip-tc-tags=true", + "--restore-unconverted-bases=true", + ]) + .expect("failed to parse Zipper arguments"); + + let opts = cmd.to_zipper_options(); + + assert_eq!(opts.tags_to_remove, vec!["RX".to_string(), "MI".to_string()]); + assert_eq!(opts.tags_to_reverse, vec!["QX".to_string()]); + assert_eq!(opts.tags_to_revcomp, vec!["OX".to_string(), "ZA".to_string()]); + assert_eq!(opts.buffer, 1234); + assert_eq!(opts.bwa_chunk_size, 4242); + assert!(opts.exclude_missing_reads); + assert!(opts.skip_tc_tags); + assert!(opts.restore_unconverted_bases); + } + + /// The projection must also carry defaults faithfully — a field hard-coded + /// to `true`/`0` would pass the all-non-default test above. + #[test] + fn to_zipper_options_carries_defaults() { + let cmd = + Zipper::try_parse_from(["zipper", "-u", "u.bam", "-r", "ref.fa", "-o", "out.bam"]) + .expect("failed to parse Zipper arguments"); + + let opts = cmd.to_zipper_options(); + + assert!(opts.tags_to_remove.is_empty()); + assert!(opts.tags_to_reverse.is_empty()); + assert!(opts.tags_to_revcomp.is_empty()); + assert_eq!(opts.buffer, 50_000); + assert_eq!(opts.bwa_chunk_size, 150_000_000); + assert!(!opts.exclude_missing_reads); + assert!(!opts.skip_tc_tags); + assert!(!opts.restore_unconverted_bases); + } + /// The rendered long help for `--skip-tc-tags`. fn skip_tc_tags_long_help() -> String { let command = Zipper::command(); diff --git a/src/lib/fastq_parse.rs b/src/lib/fastq_parse.rs index a4cd5880c..be3f319e8 100644 --- a/src/lib/fastq_parse.rs +++ b/src/lib/fastq_parse.rs @@ -171,6 +171,15 @@ impl MemoryEstimate for FastqRecord { } } +/// The typed-step pipeline byte-bounds its queues on this. Delegates to +/// [`MemoryEstimate::estimate_heap_size`] so the queue budget and the memory +/// estimate are a single source of truth. +impl crate::pipeline::core::item::HeapSize for FastqRecord { + fn heap_size(&self) -> usize { + MemoryEstimate::estimate_heap_size(self) + } +} + /// Result of parsing a FASTQ record. #[derive(Debug)] enum FastqParseResult { diff --git a/src/lib/grouper.rs b/src/lib/grouper.rs index dbbe05e57..4ae168cb4 100644 --- a/src/lib/grouper.rs +++ b/src/lib/grouper.rs @@ -223,7 +223,7 @@ impl Grouper for TemplateGrouper { // records into the same template. for decoded in records { let name_hash = decoded.key.name_hash; - let raw = decoded.data; + let raw = decoded.into_raw_bytes(); let read_name = fgumi_raw_bam::read_name(&raw); let same_template = match (self.current_name_hash, self.current_name.as_deref()) { (Some(h), Some(name)) => h == name_hash && name == read_name, @@ -430,7 +430,7 @@ impl RecordPositionGrouper { fn validate_mc_tag(decoded: &DecodedRecord) -> io::Result<()> { use fgumi_raw_bam::RawRecordView; - let raw = &decoded.data; + let raw = decoded.record(); let flg = RawRecordView::new(raw).flags(); let is_paired = (flg & fgumi_raw_bam::flags::PAIRED) != 0; let is_secondary = (flg & fgumi_raw_bam::flags::SECONDARY) != 0; @@ -489,8 +489,8 @@ impl RecordPositionGrouper { // Hash match is a fast pre-check; confirm QNAME bytes to guard // against hash collisions merging unrelated templates. last.key.name_hash == decoded.key.name_hash - && fgumi_raw_bam::read_name(&last.data) - == fgumi_raw_bam::read_name(&decoded.data) + && fgumi_raw_bam::read_name(last.raw_bytes()) + == fgumi_raw_bam::read_name(decoded.raw_bytes()) }) => { // Different position but same template (name_hash + QNAME match with @@ -601,7 +601,7 @@ fn group_by_name_and_build( // borrow (and only copy it out when a new group starts) so records within // a group don't each allocate a QNAME `Vec` — the borrow must end before // `extract` consumes `decoded`. - let read_name = fgumi_raw_bam::read_name(&decoded.data); + let read_name = fgumi_raw_bam::read_name(decoded.raw_bytes()); let same = current_name_hash == Some(name_hash) && current_name.as_deref() == Some(read_name); let new_name = if same { None } else { Some(read_name.to_vec()) }; @@ -669,6 +669,15 @@ impl MemoryEstimate for FastqTemplate { } } +/// The typed-step pipeline byte-bounds its queues on this. Delegates to +/// [`MemoryEstimate::estimate_heap_size`] so the queue budget and the memory +/// estimate are a single source of truth. +impl crate::pipeline::core::item::HeapSize for FastqTemplate { + fn heap_size(&self) -> usize { + MemoryEstimate::estimate_heap_size(self) + } +} + /// Groups FASTQ records from multiple synchronized input streams into templates. /// /// This grouper expects decompressed bytes from multiple FASTQ files, @@ -1552,8 +1561,11 @@ mod tests { .finish() .expect("finish must succeed with a group key set") .expect("the buffered position must be emitted"); - let names: Vec> = - group.records.iter().map(|r| fgumi_raw_bam::read_name(&r.data).to_vec()).collect(); + let names: Vec> = group + .records + .iter() + .map(|r| fgumi_raw_bam::read_name(r.raw_bytes()).to_vec()) + .collect(); assert_eq!( names, vec![b"readA".to_vec(), b"readB".to_vec(), b"readC".to_vec()], diff --git a/src/lib/mi_group.rs b/src/lib/mi_group.rs index 8314a28d4..d5f9ad72c 100644 --- a/src/lib/mi_group.rs +++ b/src/lib/mi_group.rs @@ -107,6 +107,111 @@ type MiTransformFn = Box String + Send + Sync>; /// Type alias for raw-byte record filter function. type RecordFilterFn = Box bool + Send + Sync>; +/// Borrowed MI tag transformation, shared by `MiGrouper` and `MiGroupIterator` +/// (whose stored transforms differ only in `Send + Sync` bounds). +pub(crate) type MiTransform<'a> = Option<&'a dyn Fn(&[u8]) -> String>; + +/// The grouping key for a run of consecutive records, stored as the raw +/// comparison bytes so that boundary detection stays allocation-free on the +/// common (no-transform) path: a record is compared against the current +/// group by byte-equality of its tag value(s) rather than by building an +/// owned `String` per record. The display label handed to `MiGroup` / the +/// iterator output is materialized from these bytes only once per group. +pub(crate) struct MiKey { + /// MI tag bytes, already transformed if a transform is configured. + mi: Vec, + /// Cell tag bytes, present (`Some`) only when cell-barcode grouping is + /// enabled. `Some(empty)` when the tag is configured but absent on the + /// record, mirroring the legacy `"MI\t"` composite key. + cell: Option>, +} + +impl MiKey { + /// Locate the MI (and, when cell grouping is enabled, the cell) tag value(s) + /// for a record. When a cell tag is configured both are found in a **single** + /// aux-data walk; otherwise only the MI tag is looked up. Returns `None` when + /// the record has no MI tag (the record is then skipped). The boolean tracks + /// whether cell grouping is enabled so callers can distinguish a configured + /// cell tag that is absent on the record (`Some(b"")`) from no cell tag at all. + #[inline] + fn locate_values( + bam: &[u8], + tag: [u8; 2], + cell_tag: Option<[u8; 2]>, + ) -> Option<(&[u8], Option<&[u8]>)> { + match cell_tag { + Some(ct) => { + // Single aux-block walk for both the MI and the cell tag. + let (mi, cell) = fgumi_raw_bam::find_two_string_tags_in_record(bam, tag, ct); + Some((mi?, Some(cell.unwrap_or(b"")))) + } + None => Some((fgumi_raw_bam::find_string_tag_in_record(bam, tag)?, None)), + } + } + + /// Extract the grouping key from raw BAM bytes, allocating the owned key + /// bytes. Used when a new group begins. Returns `None` when the record has + /// no MI tag (the record is then skipped). + pub(crate) fn from_record( + bam: &[u8], + tag: [u8; 2], + cell_tag: Option<[u8; 2]>, + transform: MiTransform, + ) -> Option { + let (value, cell_value) = Self::locate_values(bam, tag, cell_tag)?; + let mi = match transform { + Some(transform) => transform(value).into_bytes(), + None => value.to_vec(), + }; + let cell = cell_value.map(<[u8]>::to_vec); + Some(Self { mi, cell }) + } + + /// Test whether `bam` belongs to the same MI group as this key. Returns + /// `None` when the record has no MI tag (the record is then skipped). + /// Allocation-free on the no-transform path: the record's tag bytes are + /// compared against the stored key directly instead of building an owned key. + pub(crate) fn matches_record( + &self, + bam: &[u8], + tag: [u8; 2], + cell_tag: Option<[u8; 2]>, + transform: MiTransform, + ) -> Option { + let (value, cell_value) = Self::locate_values(bam, tag, cell_tag)?; + let mi_eq = match transform { + Some(transform) => transform(value).as_bytes() == self.mi.as_slice(), + None => value == self.mi.as_slice(), + }; + let cell_eq = match (cell_value, &self.cell) { + (Some(have), Some(want)) => have == want.as_slice(), + // `cell_tag` is fixed for a grouper's lifetime, so the stored key's + // cell presence always matches the configuration; the mixed arms are + // unreachable. + (None, None) => true, + (Some(_), None) | (None, Some(_)) => { + debug_assert!( + false, + "cell-tag presence is fixed for the grouper's lifetime, so the stored key's \ + cell presence always matches the record's; this mixed arm is unreachable" + ); + false + } + }; + Some(mi_eq && cell_eq) + } + + /// Build the display label. With a cell tag this is `"MI\tCELL"`, + /// matching the legacy composite key; otherwise it is just the MI value. + pub(crate) fn label(&self) -> String { + let mi = String::from_utf8_lossy(&self.mi); + match &self.cell { + Some(cell) => format!("{mi}\t{}", String::from_utf8_lossy(cell)), + None => mi.into_owned(), + } + } +} + /// A Grouper that groups raw-byte BAM records by MI tag. /// /// Records arrive as raw BAM bytes (see [`DecodedRecord::from_raw_bytes`]). MI tags are diff --git a/src/lib/mod.rs b/src/lib/mod.rs index 6188e2e6a..0246d172a 100644 --- a/src/lib/mod.rs +++ b/src/lib/mod.rs @@ -142,6 +142,7 @@ pub mod logging; pub mod metrics; pub mod mi_group; pub mod per_thread_accumulator; +pub mod pipeline; pub use fgumi_consensus::phred; pub mod read_info; pub mod read_structure; diff --git a/src/lib/pipeline/mod.rs b/src/lib/pipeline/mod.rs new file mode 100644 index 000000000..4fe625895 --- /dev/null +++ b/src/lib/pipeline/mod.rs @@ -0,0 +1,27 @@ +//! The typed-step pipeline framework for `--threads N` mode. +//! +//! A pipeline is a chain of [`core::step::Step`]s, each declaring a +//! [`core::step::StepKind`] (`Serial` / `Parallel` / `Exclusive`); the engine +//! runs all `N` worker threads through a round-robin dispatch loop with +//! `Affinity`-based pinning for I/O sources/sinks. This keeps `--threads N` a +//! strict thread cap with no separate I/O thread pools. +//! +//! # Module structure +//! +//! - [`core`]: typed-step execution engine (worker loop, queues, drain +//! protocol, reorder buffers). +//! - [`steps`]: the concrete `Step` implementations (decompress, boundaries, +//! parse, serialize, compress, write, …). +//! +//! A `chains` module (declarative chain construction — `build_for`, +//! `ChainBuilder`, per-command step factories) lands in a follow-up; until it +//! does, nothing in `src/lib/commands` routes through this tree and every +//! command still runs on `unified_pipeline`. + +/// The typed-step execution engine, extracted into the `fgumi-pipeline-core` +/// crate so its lightweight dependency graph (no `noodles`-bam / sort / +/// consensus) compiles fast in isolation. Re-exported here so every +/// `crate::pipeline::core::…` path resolves — which is what lets the ported +/// step sources compile unmodified. +pub use fgumi_pipeline_core as core; +pub mod steps; diff --git a/src/lib/pipeline/steps/bgzf/compress.rs b/src/lib/pipeline/steps/bgzf/compress.rs new file mode 100644 index 000000000..ebfdba2b9 --- /dev/null +++ b/src/lib/pipeline/steps/bgzf/compress.rs @@ -0,0 +1,313 @@ +//! `BgzfCompress` mid-step. `Parallel + ByItemOrdinal`. Compresses each +//! `DecompressedBlock` into one logical `BgzfBlock` whose `bytes` field +//! contains the concatenation of however many physical 64-KiB BGZF blocks +//! `libdeflater` produces internally. +//! +//! ## 1:1 input-to-output mapping (Phase 3 design) +//! +//! Each `try_run` consumes exactly one input and emits exactly one output +//! with `batch_serial = input.batch_serial`. This preserves the +//! consecutive-ordinal invariant required by the downstream +//! `ReorderStage` without any sub-serial encoding. +//! +//! Internally, `InlineBgzfCompressor` may produce multiple physical BGZF +//! blocks when the input exceeds `BGZF_MAX_BLOCK_SIZE = 64 KiB`. We +//! concatenate those physical blocks into a single `BgzfBlock`'s bytes — +//! a valid BGZF stream is just a concatenation of independent blocks, so +//! the downstream writer can emit them verbatim. + +use std::io; + +use fgumi_bgzf::InlineBgzfCompressor; + +use crate::pipeline::core::Unpushed; +use crate::pipeline::core::held::HeldSlot; +use crate::pipeline::core::outputs::OrderedBytesSingle; +use crate::pipeline::core::queues::QueueSpec; +use crate::pipeline::core::reorder::BranchOrdering; +use crate::pipeline::core::step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}; +use crate::pipeline::steps::types::{BgzfBlock, DecompressedBlock}; + +/// Per-worker BGZF compressed-output scratch capacity. The compressed +/// output of `BgzfCompress` is the concatenation of physical 64 KiB BGZF +/// blocks; a single batch typically produces 1-4 such blocks. Sized to +/// the typical case so the same mimalloc size class is hit every call. +/// Matches legacy `bam.rs:1561` `SERIALIZATION_BUFFER_CAPACITY` × 4. +const COMPRESS_SCRATCH_CAPACITY: usize = 256 * 1024; + +/// `Parallel + ByItemOrdinal` BGZF compressor. +pub struct BgzfCompress { + /// Per-worker compressor (each worker holds its own clone). + compressor: InlineBgzfCompressor, + compression_level: u32, + /// Per-worker output scratch buffer; `mem::replace`d on each emit so + /// the freshly-allocated replacement is always the same size class. + /// See `BgzfDecompress::output_scratch` for rationale. + output_scratch: Vec, + held: HeldSlot>, + output_byte_limit: u64, +} + +impl BgzfCompress { + #[must_use] + pub fn new(compression_level: u32, output_byte_limit: u64) -> Self { + Self { + compressor: InlineBgzfCompressor::new(compression_level), + compression_level, + output_scratch: Vec::with_capacity(COMPRESS_SCRATCH_CAPACITY), + held: HeldSlot::new(), + output_byte_limit, + } + } +} + +impl Clone for BgzfCompress { + fn clone(&self) -> Self { + Self { + compressor: InlineBgzfCompressor::new(self.compression_level), + compression_level: self.compression_level, + output_scratch: Vec::with_capacity(COMPRESS_SCRATCH_CAPACITY), + held: HeldSlot::new(), + output_byte_limit: self.output_byte_limit, + } + } +} + +impl BgzfCompress { + /// Drain the compressor's physical BGZF blocks and assemble the single + /// output `BgzfBlock` byte buffer. + /// + /// INVARIANT: each block's `serial` is `InlineBgzfCompressor`'s monotonic + /// per-worker counter; it has no relation to the pipeline-wide + /// `batch_serial`. We deliberately drop it (read only `child.data`) and + /// inherit the parent's `batch_serial` for the emitted `BgzfBlock`. Future + /// code must NOT propagate `child.serial` into the framework's ordinal stream. + /// + /// The output is the concatenation of the physical block bytes (a valid BGZF + /// stream is just independent blocks back to back). When there is exactly one + /// physical block — the common case, since a batch usually fits in one 64 KiB + /// block — its buffer is moved out directly (no scratch allocation and no + /// concatenating memcpy); this is byte-identical to concatenating a single + /// block. For the multi-block case the bytes are concatenated into the + /// per-worker scratch (then `mem::replace`d out so the replacement stays in a + /// fixed mimalloc size class), and each drained child buffer is handed back to + /// the compressor's pool so the next compression reuses it. + fn assemble_output_bytes(&mut self) -> Vec { + let mut children = self.compressor.take_blocks(); + match children.len() { + // No physical blocks — reachable from the all-clean rejects path, + // where an empty `DecompressedBlock` produces no BGZF output. Return + // an empty `Vec` rather than `mem::replace`-ing out the scratch: + // after a prior multi-block batch the scratch retains a + // `COMPRESS_SCRATCH_CAPACITY` allocation, and handing that out as an + // "empty" output would inflate queue memory/backpressure on the + // ordinal-preserving fast path. + 0 => Vec::new(), + 1 => children.pop().expect("len == 1").data, + _ => { + for child in children { + self.output_scratch.extend_from_slice(&child.data); + self.compressor.recycle_buffer(child.data); + } + std::mem::replace( + &mut self.output_scratch, + Vec::with_capacity(COMPRESS_SCRATCH_CAPACITY), + ) + } + } + } +} + +impl Step for BgzfCompress { + type Input = DecompressedBlock; + type Outputs = OrderedBytesSingle; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "BgzfCompress", + 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 let Some(unpushed) = self.held.take() { + match ctx.outputs.retry(unpushed) { + Ok(()) => {} + Err(again) => { + self.held.put(again); + return Ok(StepOutcome::Contention); + } + } + } + + let Some(parent) = ctx.input.pop() else { + // No input this call. If upstream is drained, every item has been + // processed (held output was flushed by the Contention preamble + // above) and this step will never push again — report Finished. + // For a Parallel step only the last clone to finish closes the + // shared output (gated by the StepDrainCounter in the driver). + if ctx.input.is_drained() { + return Ok(StepOutcome::Finished); + } + return Ok(StepOutcome::NoProgress); + }; + let DecompressedBlock { batch_serial, bytes } = parent; + let uncompressed_size = u32::try_from(bytes.len()).unwrap_or(u32::MAX); + + // Compress and harvest physical BGZF blocks. + self.compressor.write_all(&bytes)?; + self.compressor.flush()?; + let bytes = self.assemble_output_bytes(); + + let out = BgzfBlock { batch_serial, bytes, uncompressed_size }; + match ctx.outputs.push(out) { + Ok(()) => Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held.put(unpushed); + Ok(StepOutcome::Progress) + } + } + } + + fn new_worker_copy(&self) -> Self { + self.clone() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn profile_advertises_parallel_byordinal() { + let s = BgzfCompress::new(1, 1024); + let p = s.profile(); + assert_eq!(p.name, "BgzfCompress"); + assert_eq!(p.kind, StepKind::Parallel); + assert_eq!(p.branch_ordering, vec![BranchOrdering::ByItemOrdinal]); + } + + #[test] + fn clone_constructs_fresh_compressor() { + let s = BgzfCompress::new(1, 1024); + let _cloned = s.clone(); + } + + #[test] + fn small_input_produces_single_concatenated_output() { + let mut s = BgzfCompress::new(1, 1024); + s.compressor.write_all(b"hello bgzf").unwrap(); + s.compressor.flush().unwrap(); + let blocks = s.compressor.take_blocks(); + assert!(!blocks.is_empty(), "expected at least one BGZF block"); + // Each block is a valid BGZF block (starts with gzip magic 0x1f 0x8b). + for b in &blocks { + assert_eq!(&b.data[..2], &[0x1f, 0x8b], "BGZF magic"); + } + } + + /// Reference assembly: the straightforward concatenation the move/recycle + /// fast path must reproduce byte-for-byte. + fn reference_concat(level: u32, input: &[u8]) -> Vec { + let mut c = InlineBgzfCompressor::new(level); + c.write_all(input).unwrap(); + c.flush().unwrap(); + let mut out = Vec::new(); + for block in c.take_blocks() { + out.extend_from_slice(&block.data); + } + out + } + + /// Drive `assemble_output_bytes` on a fresh step over `input`. + fn assemble(level: u32, input: &[u8]) -> Vec { + let mut s = BgzfCompress::new(level, 1024); + s.compressor.write_all(input).unwrap(); + s.compressor.flush().unwrap(); + s.assemble_output_bytes() + } + + #[test] + fn single_block_fast_path_is_byte_identical() { + // Small input → one physical block → move fast path. + let input = b"the quick brown fox jumps over the lazy dog"; + let mut s = BgzfCompress::new(1, 1024); + s.compressor.write_all(input).unwrap(); + s.compressor.flush().unwrap(); + assert_eq!(s.compressor.take_blocks().len(), 1, "input should fit one block"); + + let assembled = assemble(1, input); + assert_eq!(assembled, reference_concat(1, input)); + assert_eq!(&assembled[..2], &[0x1f, 0x8b], "valid BGZF magic"); + } + + #[test] + fn multi_block_concat_path_is_byte_identical() { + // Input larger than one 64 KiB block forces the multi-block concat + + // recycle path; output must still match the plain concatenation. + let mut input = Vec::with_capacity(200 * 1024); + for i in 0..(200 * 1024u32) { + input.push((i % 251) as u8); + } + let mut s = BgzfCompress::new(1, 1024); + s.compressor.write_all(&input).unwrap(); + s.compressor.flush().unwrap(); + assert!(s.compressor.take_blocks().len() > 1, "input should span >1 block"); + + let assembled = assemble(1, &input); + assert_eq!(assembled, reference_concat(1, &input)); + } + + #[test] + fn recycle_buffer_repopulates_pool_for_reuse() { + // After the multi-block path recycles buffers, a subsequent compression + // reuses them, and output remains byte-identical to a fresh compressor. + let big: Vec = (0..(200 * 1024u32)).map(|i| (i % 251) as u8).collect(); + let mut s = BgzfCompress::new(1, 1024); + + s.compressor.write_all(&big).unwrap(); + s.compressor.flush().unwrap(); + let first = s.assemble_output_bytes(); + assert_eq!(first, reference_concat(1, &big)); + + // Second block through the same (now pool-primed) compressor. + s.compressor.write_all(&big).unwrap(); + s.compressor.flush().unwrap(); + let second = s.assemble_output_bytes(); + assert_eq!(second, reference_concat(1, &big), "recycled buffers must not corrupt output"); + } + + #[test] + fn empty_batch_returns_unallocated_vec() { + // An all-clean rejects batch produces no physical BGZF blocks. The + // assembled output must be a truly empty `Vec`, not the per-worker + // scratch — otherwise an "empty" output carries a retained + // `COMPRESS_SCRATCH_CAPACITY` allocation and inflates queue memory. + let mut s = BgzfCompress::new(1, 1024); + + // Fresh compressor, nothing written → no blocks. + let fresh = s.assemble_output_bytes(); + assert!(fresh.is_empty()); + assert_eq!(fresh.capacity(), 0, "empty output must not carry a scratch allocation"); + + // After a real multi-block batch (which routes through the scratch and + // replaces it with a `COMPRESS_SCRATCH_CAPACITY` Vec), a following empty + // batch must still return an unallocated Vec. + let big: Vec = (0..(200 * 1024u32)).map(|i| (i % 251) as u8).collect(); + s.compressor.write_all(&big).unwrap(); + s.compressor.flush().unwrap(); + let multi = s.assemble_output_bytes(); + assert!(!multi.is_empty(), "multi-block batch should produce output"); + + let after = s.assemble_output_bytes(); // no new input → empty batch + assert!(after.is_empty()); + assert_eq!( + after.capacity(), + 0, + "empty batch after a multi-block batch must not retain scratch" + ); + } +} diff --git a/src/lib/pipeline/steps/bgzf/decompress.rs b/src/lib/pipeline/steps/bgzf/decompress.rs new file mode 100644 index 000000000..178aa19ee --- /dev/null +++ b/src/lib/pipeline/steps/bgzf/decompress.rs @@ -0,0 +1,159 @@ +//! `BgzfDecompress` mid-step. `Parallel + ByItemOrdinal`. Decompresses +//! incoming `BgzfBlock`s into `DecompressedBlock`s using `libdeflater`'s +//! decompressor (one per worker via `Clone`). + +use std::io; + +use fgumi_bgzf::reader::decompress_block_slice_into; +use libdeflater::Decompressor; + +use crate::pipeline::core::Unpushed; +use crate::pipeline::core::held::HeldSlot; +use crate::pipeline::core::outputs::OrderedBytesSingle; +use crate::pipeline::core::queues::QueueSpec; +use crate::pipeline::core::reorder::BranchOrdering; +use crate::pipeline::core::step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}; +use crate::pipeline::steps::types::{BgzfBlock, DecompressedBlock}; + +/// Per-worker decompression scratch capacity. Sized to the upper end +/// of typical BGZF block decompressed sizes (256 KiB ≈ 4 blocks × 64 KiB). +/// Mirrors legacy `pipeline/bam.rs:1560` `DECOMPRESSION_BUFFER_CAPACITY`. +/// +/// Why a fixed size: when each `try_run` calls `Vec::with_capacity(N)` with +/// a *variable* `N` (= `block.uncompressed_size`), every allocation hits a +/// different mimalloc size class and defeats the thread-local cache. A +/// fixed capacity always pulls from the same size class so the freed +/// buffer that just left for downstream gets reused on the next call. +const DECOMPRESS_SCRATCH_CAPACITY: usize = 256 * 1024; + +/// `Parallel + ByItemOrdinal` decompressor. Each worker holds its own +/// `libdeflater::Decompressor` (`Clone` constructs a fresh one) and a +/// fixed-capacity output scratch buffer that's emitted via `mem::replace` +/// to avoid variable-size allocations on the hot path. +pub struct BgzfDecompress { + decompressor: Decompressor, + /// Per-worker output scratch. `try_run` decompresses into this + /// buffer, then swaps it out with a fresh fixed-capacity Vec via + /// `mem::replace`. Keeps allocations on a single mimalloc size + /// class so the thread-local cache can recycle them. Mirrors legacy + /// `WorkerState::decompression_buffer`. + output_scratch: Vec, + held: HeldSlot>, + output_byte_limit: u64, +} + +impl BgzfDecompress { + #[must_use] + pub fn new(output_byte_limit: u64) -> Self { + Self { + decompressor: Decompressor::new(), + output_scratch: Vec::with_capacity(DECOMPRESS_SCRATCH_CAPACITY), + held: HeldSlot::new(), + output_byte_limit, + } + } +} + +impl Clone for BgzfDecompress { + fn clone(&self) -> Self { + Self { + decompressor: Decompressor::new(), + output_scratch: Vec::with_capacity(DECOMPRESS_SCRATCH_CAPACITY), + held: HeldSlot::new(), + output_byte_limit: self.output_byte_limit, + } + } +} + +impl Step for BgzfDecompress { + type Input = BgzfBlock; + type Outputs = OrderedBytesSingle; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "BgzfDecompress", + 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 let Some(unpushed) = self.held.take() { + match ctx.outputs.retry(unpushed) { + Ok(()) => {} + Err(again) => { + self.held.put(again); + // `Contention` (not `NoProgress`) keeps the worker + // alive for retry — `NoProgress` would let the + // framework mark this worker `Skip` if input is also + // drained, silently dropping the held item. + return Ok(StepOutcome::Contention); + } + } + } + + let Some(block) = ctx.input.pop() else { + // No input this call. If upstream is drained, every item has been + // processed (held output was flushed by the Contention preamble + // above) and this step will never push again — report Finished. + // For a Parallel step only the last clone to finish closes the + // shared output (gated by the StepDrainCounter in the driver). + if ctx.input.is_drained() { + return Ok(StepOutcome::Finished); + } + return Ok(StepOutcome::NoProgress); + }; + + // Decompress into the per-worker scratch buffer. After the + // decompressor fills it, `mem::replace` swaps it out (the filled + // buffer goes into the emitted block; a fresh fixed-capacity + // buffer takes its place). Single mimalloc size class keeps the + // thread-local cache hot. + decompress_block_slice_into( + &block.bytes, + &mut self.decompressor, + &mut self.output_scratch, + )?; + let bytes = std::mem::replace( + &mut self.output_scratch, + Vec::with_capacity(DECOMPRESS_SCRATCH_CAPACITY), + ); + + let decompressed = DecompressedBlock { batch_serial: block.batch_serial, bytes }; + + match ctx.outputs.push(decompressed) { + Ok(()) => Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held.put(unpushed); + Ok(StepOutcome::Progress) + } + } + } + + fn new_worker_copy(&self) -> Self { + self.clone() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn profile_advertises_parallel_byordinal() { + let s = BgzfDecompress::new(1024); + let p = s.profile(); + assert_eq!(p.name, "BgzfDecompress"); + assert_eq!(p.kind, StepKind::Parallel); + assert!(!p.sticky); + assert_eq!(p.branch_ordering, vec![BranchOrdering::ByItemOrdinal]); + } + + #[test] + fn clone_constructs_fresh_decompressor() { + let s = BgzfDecompress::new(1024); + let _cloned = s.clone(); + } +} diff --git a/src/lib/pipeline/steps/bgzf/mod.rs b/src/lib/pipeline/steps/bgzf/mod.rs new file mode 100644 index 000000000..83b0c9a8a --- /dev/null +++ b/src/lib/pipeline/steps/bgzf/mod.rs @@ -0,0 +1,4 @@ +//! BGZF compress/decompress steps. + +pub mod compress; +pub mod decompress; diff --git a/src/lib/pipeline/steps/boundaries/bam.rs b/src/lib/pipeline/steps/boundaries/bam.rs new file mode 100644 index 000000000..e472f0030 --- /dev/null +++ b/src/lib/pipeline/steps/boundaries/bam.rs @@ -0,0 +1,177 @@ +//! `FindBamBoundaries` mid-step. `Serial + ByItemOrdinal`. Strips the BAM +//! header (on first input) and finds record boundaries within decompressed +//! BGZF block data, emitting `DecompressedBlock`s whose bytes contain only +//! complete records (4-byte `block_size` prefix + record body, repeated). +//! +//! Wraps [`super::state::BoundaryState`] — the boundary-finding state +//! machine handles header skipping, cross-block record carryover, and +//! validation. Reusing it ensures the new framework's boundary semantics +//! match the legacy pipeline's exactly. + +use std::io; + +use super::state::BoundaryState; +use crate::pipeline::core::Unpushed; +use crate::pipeline::core::held::HeldSlot; +use crate::pipeline::core::outputs::OrderedBytesSingle; +use crate::pipeline::core::queues::QueueSpec; +use crate::pipeline::core::reorder::BranchOrdering; +use crate::pipeline::core::step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}; +use crate::pipeline::steps::types::DecompressedBlock; + +/// Max inputs processed per `try_run` invocation. Amortizes the Serial +/// mutex acquisition; matches legacy `bam.rs:2050+` (`MAX_BATCHES_PER_LOCK`). +const MAX_BATCHES_PER_LOCK: usize = 8; + +/// `Serial + ByItemOrdinal` boundary finder. Holds `BoundaryState` (which +/// owns the cross-block carryover buffer + header-skip flag). +pub struct FindBamBoundaries { + state: BoundaryState, + /// Pending output batch when we found boundaries but the push was + /// rejected. Held across retries until pushed. + held: HeldSlot>, + /// Self-managed output ordinal. Incremented only when a new + /// `DecompressedBlock` is emitted. Using the input's `batch_serial` + /// directly would skip ordinals when the boundary state absorbs an + /// input without producing output (header bytes, mid-record carryover); + /// the downstream `ReorderStage` requires consecutive `0, 1, 2, …`. + next_output_serial: u64, + /// Set once the final-flush path has called the (non-idempotent) + /// `state.finish()`; guards against a second call across the multi-pass + /// completion drain. + finalized: bool, + output_byte_limit: u64, +} + +impl FindBamBoundaries { + /// Construct expecting the first input batch to begin with the BAM + /// header (matches `BAM_MAGIC`). The header bytes are skipped on the + /// first call; subsequent calls just find record boundaries. + #[must_use] + pub fn new(output_byte_limit: u64) -> Self { + Self { + state: BoundaryState::new(), + held: HeldSlot::new(), + next_output_serial: 0, + finalized: false, + output_byte_limit, + } + } + + /// Construct expecting the input stream to be already past the BAM + /// header. Useful for runall-spliced sub-pipelines. + #[must_use] + pub fn new_no_header(output_byte_limit: u64) -> Self { + Self { + state: BoundaryState::new_no_header(), + held: HeldSlot::new(), + next_output_serial: 0, + finalized: false, + output_byte_limit, + } + } +} + +impl Step for FindBamBoundaries { + type Input = DecompressedBlock; + type Outputs = OrderedBytesSingle; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "FindBamBoundaries", + kind: StepKind::Serial, + 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 let Some(unpushed) = self.held.take() { + match ctx.outputs.retry(unpushed) { + Ok(()) => {} + Err(again) => { + self.held.put(again); + return Ok(StepOutcome::Contention); + } + } + } + + // Process up to `MAX_BATCHES_PER_LOCK` inputs per `try_run` to + // amortize the Serial mutex acquisition. Each iteration pops one + // input, runs find_boundaries, and pushes at most one output — + // header-only or carryover-only inputs are fully absorbed and produce + // no output (the loop `continue`s, and `did_work` still records the + // consumed input as Progress). We stop early on push rejection (held + // the unpushed; subsequent iterations would contend on backpressure) + // or on input exhaustion. + let mut did_work = false; + for _ in 0..MAX_BATCHES_PER_LOCK { + let Some(block) = ctx.input.pop() else { break }; + did_work = true; + + let boundary_batch = self.state.find_boundaries(&block.bytes)?; + if boundary_batch.buffer.is_empty() { + // Input fully absorbed (header/leftover); try next input. + continue; + } + + let serial = self.next_output_serial; + self.next_output_serial += 1; + let out = DecompressedBlock { batch_serial: serial, bytes: boundary_batch.buffer }; + match ctx.outputs.push(out) { + Ok(()) => {} + Err(unpushed) => { + self.held.put(unpushed); + // Hold off on more inputs — the held item must clear + // before we accept new work. + return Ok(StepOutcome::Progress); + } + } + } + + if did_work { + return Ok(StepOutcome::Progress); + } + + // No input this call. If upstream is drained, flush the boundary + // state's final partial-record buffer once (guarded by `finalized` — + // `state.finish()` is not idempotent), emit it, and report `Finished` + // once nothing remains. `held` is empty here (step 1 returned + // `Contention` otherwise); a bounced final push is parked in `held` + // for the held-drain at the top of the next pass. + if ctx.input.is_drained() { + if !self.finalized { + self.finalized = true; + if let Some(boundary_batch) = self.state.finish()? + && !boundary_batch.buffer.is_empty() + { + let serial = self.next_output_serial; + self.next_output_serial += 1; + let out = + DecompressedBlock { batch_serial: serial, bytes: boundary_batch.buffer }; + if let Err(unpushed) = ctx.outputs.push(out) { + self.held.put(unpushed); + } + return Ok(StepOutcome::Progress); + } + } + return Ok(StepOutcome::Finished); + } + Ok(StepOutcome::NoProgress) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn profile_advertises_serial_byordinal() { + let s = FindBamBoundaries::new(1024); + let p = s.profile(); + assert_eq!(p.name, "FindBamBoundaries"); + assert_eq!(p.kind, StepKind::Serial); + assert_eq!(p.branch_ordering, vec![BranchOrdering::ByItemOrdinal]); + } +} diff --git a/src/lib/pipeline/steps/boundaries/mod.rs b/src/lib/pipeline/steps/boundaries/mod.rs new file mode 100644 index 000000000..8465bfd93 --- /dev/null +++ b/src/lib/pipeline/steps/boundaries/mod.rs @@ -0,0 +1,4 @@ +//! Per-format boundary-finding steps. + +pub mod bam; +pub mod state; diff --git a/src/lib/pipeline/steps/boundaries/state.rs b/src/lib/pipeline/steps/boundaries/state.rs new file mode 100644 index 000000000..491649f27 --- /dev/null +++ b/src/lib/pipeline/steps/boundaries/state.rs @@ -0,0 +1,876 @@ +//! `BoundaryState` / `BoundaryBatch`: the BAM record-boundary state machine. +//! +//! 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; + +/// Upper bound on bytes the scanner will carry forward waiting for one record +/// (or the BAM header) to complete. +/// +/// This is a **corruption backstop, not a biological limit**. The scanner is a +/// push scanner: it is handed decompressed blocks and cannot ask for "the rest +/// of this record" the way a pull reader (htslib's `bam_read1`) can, so a +/// corrupt length field makes every block look like "not enough data yet" and +/// the carry grows for the rest of the stream. That still terminates — +/// [`BoundaryState::finish`] reports the shortfall at EOF — but only after +/// buffering the remaining input, so a large corrupt file exhausts memory +/// before it can print the diagnostic it was about to print. +/// +/// The value sits in the gap between what BAM permits and what BAM files +/// contain. `block_size` is a `u32`, so the format allows a ~4 GiB record, while +/// the largest records seen in practice — ONT ultra-long reads carrying +/// methylation tags — run a few MB. 256 MiB is roughly 25-50x above that and +/// still bounds memory hard. +/// +/// It is deliberately NOT tied to the pipeline's queue budget: byte-bounded +/// queues always admit at least one item regardless of size, so a legitimate +/// multi-MB record flows through a 4 MiB per-step budget, and bounding the carry +/// there would reject valid BAMs. +const MAX_CARRY_BYTES: usize = 256 * 1024 * 1024; + +/// 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, +} + +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, + } + } + + /// Parse BAM header and return the number of bytes consumed. + /// + /// `Ok(None)` means more data is needed. `Err` means the header's own length + /// fields are corrupt — their sum overflows `usize`, so no amount of further + /// data could satisfy them. + /// + /// # Errors + /// + /// Returns `InvalidData` when `l_text` or a reference's `l_name` overflows + /// the running offset. Unreachable on 64-bit targets (both are `u32` widened + /// to `usize`); present because a length field taken from the stream must be + /// folded in with checked arithmetic rather than trusted to fit. + fn parse_header_size(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 + if &data[0..4] != fgumi_raw_bam::BAM_MAGIC { + // Not a valid BAM file, but let's not error here + // Just return 0 so records start immediately + return Ok(Some(0)); + } + + let l_text = u32::from_le_bytes([data[4], data[5], data[6], data[7]]) as usize; + let Some(mut offset) = 8usize.checked_add(l_text) else { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "FindBamBoundaries: BAM header l_text={l_text} overflows the header offset" + ), + )); + }; + + 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; + // l_name + name + l_ref, folded in with checked arithmetic for the + // same reason as `l_text` above. + let Some(next) = offset.checked_add(8).and_then(|o| o.checked_add(l_name)) else { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "FindBamBoundaries: BAM header l_name={l_name} overflows \ + the header offset at {offset}" + ), + )); + }; + offset = next; + + if data.len() < offset { + return Ok(None); + } + } + + Ok(Some(offset)) + } + + /// 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 `InvalidData` if the bytes carried forward waiting for the header + /// or for one record to complete exceed this module's `MAX_CARRY_BYTES` + /// corruption backstop. A well-formed stream always completes a record + /// within a bounded carry, so exceeding it means a length field is corrupt. + /// + /// # 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 { + // Step 1: Combine leftover with new data into reusable work_buffer + // This avoids allocating a new Vec on every call + 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 { + if let Some(header_size) = Self::parse_header_size(&self.work_buffer)? { + cursor = header_size; + self.header_skipped = true; + } else { + // Not enough data to parse header yet. Bound the carry: a corrupt + // `l_text` / `l_name` makes this branch unreachable-to-satisfy, so + // without the check every remaining block accumulates here. + if self.work_buffer.len() > MAX_CARRY_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "FindBamBoundaries: BAM header still unresolved after {} byte(s) \ + (limit {MAX_CARRY_BYTES}); the header length fields are corrupt", + self.work_buffer.len(), + ), + )); + } + // Save as leftover and return an empty batch. + std::mem::swap(&mut self.leftover, &mut self.work_buffer); + return Ok(BoundaryBatch { buffer: Vec::new(), offsets: vec![0] }); + } + } + + // Step 3: Scan for record boundaries (FAST - just read integers) + let start_cursor = cursor; + // Pre-size from the previous block's record count so the per-record + // pushes below don't trigger repeated Vec regrowth (adjacent blocks + // hold near-identical record counts). First offset is 0 (relative to + // the start of records). + 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; + + // Checked: `block_size` comes from the stream. Unreachable on + // 64-bit, but the framing rule is checked-before-use. + let Some(record_end) = cursor.checked_add(4).and_then(|c| c.checked_add(block_size)) + else { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "FindBamBoundaries: record end overflows \ + (cursor={cursor}, block_size={block_size})" + ), + )); + }; + if record_end > self.work_buffer.len() { + break; // Incomplete record - becomes leftover + } + + cursor = record_end; + // Offset is relative to start of records (after header) + offsets.push(cursor - start_cursor); + } + + // Remember this block's offset count to pre-size the next call. + self.prev_offsets_len = offsets.len(); + + // Step 4: Save leftover for next block (reuse allocation) + // Split work_buffer: [0..start_cursor | start_cursor..cursor | cursor..] + // header (discard) | records (output) | leftover + // Bound the carry before retaining it: a corrupt `block_size` keeps the + // scan breaking at the same record forever, so without this the leftover + // grows to the size of the remaining input. `finish` would still catch + // the shortfall at EOF, but only after the memory was already spent. + let carry_len = self.work_buffer.len() - cursor; + if carry_len > MAX_CARRY_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "FindBamBoundaries: incomplete BAM record has carried {carry_len} byte(s) \ + (limit {MAX_CARRY_BYTES}); its block_size prefix is corrupt", + ), + )); + } + self.leftover.clear(); + self.leftover.extend_from_slice(&self.work_buffer[cursor..]); + + // Extract the records buffer - this allocation is unavoidable as we return ownership + let buffer = self.work_buffer[start_cursor..cursor].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 }) + } + + /// 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; + + // Checked for the same reason as the scan loop above. + let Some(record_end) = cursor.checked_add(4).and_then(|c| c.checked_add(block_size)) + else { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "FindBamBoundaries: record end overflows at EOF \ + (cursor={cursor}, block_size={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; + + /// Frame `payload` as one BAM record: `block_size: u32 LE` + body. The + /// scanner never inspects the body, so an opaque payload is sufficient — + /// and keeps each case's record lengths readable in the table. + fn record(payload: &[u8]) -> Vec { + let mut framed = Vec::with_capacity(4 + payload.len()); + framed.extend_from_slice( + &u32::try_from(payload.len()).expect("payload fits u32").to_le_bytes(), + ); + framed.extend_from_slice(payload); + framed + } + + /// A well-formed binary BAM header: magic, `l_text` + text, `n_ref`, then + /// `(l_name, name\0, l_ref)` per reference. This is exactly the prefix + /// `parse_header_size` walks, so building it honestly is what makes the + /// header-skip assertions meaningful. + fn bam_header(text: &str, refs: &[(&str, u32)]) -> Vec { + let mut header = Vec::new(); + header.extend_from_slice(fgumi_raw_bam::BAM_MAGIC); + header.extend_from_slice(&u32::try_from(text.len()).expect("text fits u32").to_le_bytes()); + header.extend_from_slice(text.as_bytes()); + header.extend_from_slice(&u32::try_from(refs.len()).expect("n_ref fits u32").to_le_bytes()); + for (name, ref_len) in refs { + let mut name_bytes = name.as_bytes().to_vec(); + name_bytes.push(0); // names are NUL-terminated, and l_name counts the NUL + header.extend_from_slice( + &u32::try_from(name_bytes.len()).expect("l_name fits u32").to_le_bytes(), + ); + header.extend_from_slice(&name_bytes); + header.extend_from_slice(&ref_len.to_le_bytes()); + } + header + } + + /// Concatenate framed records, as one decompressed block would hold them. + fn records(payloads: &[&[u8]]) -> Vec { + payloads.iter().flat_map(|p| record(p)).collect() + } + + // ======================================================================== + // Header skipping + // ======================================================================== + + /// The header must be consumed and excluded from the emitted buffer, with + /// only record bytes surviving. Parameterized over header shapes because + /// `parse_header_size`'s arithmetic differs per section: text length and + /// per-reference `(l_name, name, l_ref)` walking are separate loops, and a + /// bug in either would pass the other's case. + #[rstest] + #[case::no_text_no_refs("", &[])] + #[case::text_only("@HD\tVN:1.6\n", &[])] + #[case::one_ref("@HD\tVN:1.6\n", &[("chr1", 1000)])] + #[case::several_refs("@HD\tVN:1.6\n", &[("chr1", 1000), ("chr2", 2000), ("chrM", 16569)])] + #[case::empty_text_with_refs("", &[("chr1", 248_956_422)])] + fn find_boundaries_skips_the_header_and_emits_only_record_bytes( + #[case] text: &str, + #[case] refs: &[(&str, u32)], + ) { + let payloads: [&[u8]; 2] = [&[1u8; 8], &[2u8; 16]]; + let mut block = bam_header(text, refs); + block.extend_from_slice(&records(&payloads)); + + let mut state = BoundaryState::new(); + let batch = state.find_boundaries(&block).expect("well-formed block scans"); + + assert_eq!( + batch.buffer, + records(&payloads), + "header bytes must not reach the output buffer" + ); + assert_eq!(batch.offsets, vec![0, 12, 32], "offsets are relative to the first record"); + } + + /// A `BoundaryState::new_no_header()` treats byte 0 as the first record — + /// this is the runall-spliced sub-pipeline case, where an upstream stage + /// already consumed the header. + #[test] + fn new_no_header_treats_the_first_byte_as_a_record_boundary() { + let payloads: [&[u8]; 2] = [&[7u8; 4], &[9u8; 4]]; + let block = records(&payloads); + + let mut state = BoundaryState::new_no_header(); + let batch = state.find_boundaries(&block).expect("record-aligned block scans"); + + assert_eq!(batch.buffer, block, "every byte is a record byte when the header is skipped"); + assert_eq!(batch.offsets, vec![0, 8, 16]); + } + + /// `Default` must agree with `new()` — i.e. it expects a header. Asserted + /// behaviorally (the header is consumed) rather than by comparing private + /// fields, so the test survives a field-level refactor. + #[test] + fn default_expects_a_header_like_new() { + let payloads: [&[u8]; 1] = [&[1u8; 8]]; + let mut block = bam_header("", &[]); + block.extend_from_slice(&records(&payloads)); + + let batch = BoundaryState::default().find_boundaries(&block).expect("scans"); + assert_eq!(batch.buffer, records(&payloads)); + } + + /// A header split across blocks must be buffered, not mis-parsed: the first + /// call cannot determine the header size, so it emits an empty batch and + /// carries every byte forward. Parameterized over cut points that land in + /// each distinct section of `parse_header_size`'s walk, since each has its + /// own "need more data" early return. + #[rstest] + #[case::mid_magic(2)] + #[case::after_magic_before_l_text(4)] + #[case::mid_text(10)] + #[case::before_n_ref(15)] + #[case::mid_ref_name(24)] + fn a_header_split_across_blocks_is_buffered_until_complete(#[case] cut: usize) { + let payloads: [&[u8]; 1] = [&[3u8; 8]]; + let header = bam_header("@HD\tVN:1.6\n", &[("chr1", 1000)]); + assert!( + cut < header.len(), + "cut must fall inside the header for this case to mean anything" + ); + let mut full = header.clone(); + full.extend_from_slice(&records(&payloads)); + + let mut state = BoundaryState::new(); + let first = state.find_boundaries(&full[..cut]).expect("partial header is not an error"); + assert!(first.buffer.is_empty(), "no records can be emitted before the header is parsed"); + assert_eq!(first.offsets, vec![0]); + + let second = state.find_boundaries(&full[cut..]).expect("rest of the header completes it"); + assert_eq!(second.buffer, records(&payloads), "records resume once the header is consumed"); + assert_eq!(second.offsets, vec![0, 12]); + } + + /// Data that does not start with `BAM\1` is treated as record bytes from + /// offset 0 rather than erroring — the documented "not a valid BAM file, + /// but let's not error here" behavior. Worth pinning because it is + /// surprising: a caller who feeds the wrong stream gets garbage records, + /// not a diagnostic. + #[test] + fn a_non_bam_magic_prefix_is_scanned_as_records_from_offset_zero() { + let payloads: [&[u8]; 1] = [&[5u8; 8]]; + let block = records(&payloads); + assert_ne!(&block[0..4], fgumi_raw_bam::BAM_MAGIC, "fixture must not look like a header"); + + let mut state = BoundaryState::new(); + let batch = state.find_boundaries(&block).expect("scans"); + + assert_eq!(batch.buffer, block); + assert_eq!(batch.offsets, vec![0, 12]); + } + + /// Fewer than 8 bytes cannot even hold magic + `l_text`, so the scanner + /// must ask for more data instead of reading past the end. + #[test] + fn fewer_than_eight_bytes_cannot_resolve_the_header() { + let mut state = BoundaryState::new(); + let batch = state.find_boundaries(b"BAM\x01").expect("short input is not an error"); + assert!(batch.buffer.is_empty()); + assert_eq!(batch.offsets, vec![0]); + } + + // ======================================================================== + // Record scanning and cross-block carryover + // ======================================================================== + + /// The core carryover contract: a record straddling two blocks is held back + /// on the first call and emitted whole on the second, never split. The + /// concatenated output across both calls must equal the input records + /// exactly — that is the invariant downstream parsing depends on. + #[rstest] + #[case::split_mid_body(18)] + #[case::split_on_a_record_start(12)] + #[case::split_mid_length_prefix(14)] + #[case::split_one_byte_in(1)] + fn a_record_spanning_two_blocks_is_emitted_whole(#[case] cut: usize) { + let payloads: [&[u8]; 3] = [&[1u8; 8], &[2u8; 8], &[3u8; 8]]; + let all = records(&payloads); + assert!(cut < all.len()); + + let mut state = BoundaryState::new_no_header(); + let first = state.find_boundaries(&all[..cut]).expect("first half scans"); + let second = state.find_boundaries(&all[cut..]).expect("second half scans"); + + let mut seen = first.buffer.clone(); + seen.extend_from_slice(&second.buffer); + assert_eq!(seen, all, "no byte may be dropped or duplicated across the split"); + assert!(state.finish().expect("nothing left over").is_none()); + } + + /// Offsets must always start at 0, end at `buffer.len()`, and step by each + /// record's framed size — this is the contract that lets a consumer slice + /// records with `offsets.windows(2)`. + #[test] + fn offsets_bracket_every_record_in_the_emitted_buffer() { + let payloads: [&[u8]; 3] = [&[1u8; 4], &[2u8; 20], &[3u8; 0]]; + let block = records(&payloads); + + let mut state = BoundaryState::new_no_header(); + let batch = state.find_boundaries(&block).expect("scans"); + + assert_eq!(batch.offsets, vec![0, 8, 32, 36]); + assert_eq!(*batch.offsets.last().expect("non-empty"), batch.buffer.len()); + for (i, window) in batch.offsets.windows(2).enumerate() { + let body = &batch.buffer[window[0] + 4..window[1]]; + assert_eq!(body, payloads[i], "record {i} must round-trip through its offsets"); + } + } + + /// An empty block is a legal no-op: nothing to scan, nothing carried. + #[test] + fn an_empty_block_yields_an_empty_batch() { + let mut state = BoundaryState::new_no_header(); + let batch = state.find_boundaries(&[]).expect("empty input is not an error"); + assert!(batch.buffer.is_empty()); + assert_eq!(batch.offsets, vec![0]); + } + + /// Feeding many blocks in sequence must give the same records as feeding + /// one concatenated block. This is what exercises the `prev_offsets_len` + /// pre-sizing path across calls, and it is the property that actually + /// matters: block framing must not be observable in the output. + #[test] + fn streaming_many_blocks_matches_scanning_one_concatenated_block() { + let payloads: Vec> = (0..32u8).map(|i| vec![i; usize::from(i) + 1]).collect(); + let refs: Vec<&[u8]> = payloads.iter().map(Vec::as_slice).collect(); + let all = records(&refs); + + let one_shot = { + let mut state = BoundaryState::new_no_header(); + let batch = state.find_boundaries(&all).expect("scans"); + assert!(state.finish().expect("nothing left over").is_none()); + batch.buffer + }; + + let streamed = { + let mut state = BoundaryState::new_no_header(); + let mut out = Vec::new(); + for chunk in all.chunks(7) { + out.extend_from_slice(&state.find_boundaries(chunk).expect("scans").buffer); + } + assert!(state.finish().expect("nothing left over").is_none()); + out + }; + + assert_eq!(streamed, one_shot, "block framing must not be observable downstream"); + assert_eq!(one_shot, all); + } + + // ======================================================================== + // finish(): EOF validation + // ======================================================================== + + /// With no carryover there is nothing to flush. + #[test] + fn finish_returns_none_when_no_bytes_were_carried_over() { + let mut state = BoundaryState::new_no_header(); + state.find_boundaries(&records(&[&[1u8; 8]])).expect("scans"); + assert!(state.finish().expect("no leftover").is_none()); + } + + /// `finish` flushes carried-over bytes when they happen to form whole + /// records. That is narrower than it sounds, and the narrowness is the + /// point: `find_boundaries` consumes *every* complete record it can see, so + /// after a successful scan the carryover is by construction an INCOMPLETE + /// record — making `None` or an error the normal outcomes at EOF. + /// + /// The one reachable path to a flushed batch is a stream that ended before + /// the header could even be resolved (< 8 bytes, so `parse_header_size` + /// never ran) whose bytes nonetheless frame cleanly. Pinning it keeps the + /// flush path honest and documents why it is nearly dead code. + #[test] + fn finish_flushes_carried_over_bytes_that_form_whole_records() { + // Seven bytes: too short for magic + `l_text`, so the whole block is + // carried over unparsed — but a clean `block_size = 3` + 3-byte body. + let block = record(&[7u8; 3]); + assert_eq!(block.len(), 7); + + let mut state = BoundaryState::new(); + let batch = state.find_boundaries(&block).expect("short input is not an error"); + assert!(batch.buffer.is_empty(), "the header was never resolved, so nothing is emitted"); + + let flushed = + state.finish().expect("carryover frames cleanly").expect("a batch is flushed"); + assert_eq!(flushed.buffer, block); + assert_eq!(flushed.offsets, vec![0, 7]); + } + + /// After a successful scan the carryover is always a partial record, so the + /// normal EOF outcome is `None` — the complete records were already emitted + /// by `find_boundaries` itself. + #[test] + fn finish_returns_none_after_a_scan_that_consumed_every_record() { + let payloads: [&[u8]; 2] = [&[1u8; 8], &[2u8; 8]]; + let all = records(&payloads); + + let mut state = BoundaryState::new_no_header(); + let batch = state.find_boundaries(&all[..12]).expect("scans"); + assert_eq!(batch.buffer, record(&[1u8; 8]), "the first record is emitted immediately"); + + assert!( + state.finish().expect("no carryover").is_none(), + "a scan that ended on a record boundary leaves nothing to flush" + ); + } + + /// A truncated record at EOF is corruption and must surface as + /// `UnexpectedEof` rather than silently dropping bytes. Two distinct + /// truncations reach two different error sites: a declared `block_size` + /// that outruns the buffer, and trailing bytes too short to even hold the + /// 4-byte length prefix. + #[rstest] + #[case::body_shorter_than_declared_block_size(&[0x20, 0, 0, 0, 1, 2, 3])] + #[case::three_trailing_bytes(&[1, 2, 3])] + #[case::one_trailing_byte(&[9])] + fn finish_rejects_an_incomplete_record_at_eof(#[case] tail: &[u8]) { + let mut block = record(&[1u8; 8]); + block.extend_from_slice(tail); + + let mut state = BoundaryState::new_no_header(); + let batch = state.find_boundaries(&block).expect("the complete record scans"); + assert_eq!(batch.buffer, record(&[1u8; 8])); + + let err = state.finish().expect_err("a truncated tail must not be dropped"); + assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof); + } + + /// `finish` takes the leftover, so a second call reports nothing remaining + /// rather than re-emitting the same bytes. `FindBamBoundaries` guards this + /// with its own `finalized` flag, but the state machine must not depend on + /// that guard for correctness of the buffer it returns. + #[test] + fn finish_does_not_re_emit_the_same_leftover_twice() { + // Same short-input setup as the flush test above — the only shape that + // leaves whole records parked in carryover. + let block = record(&[7u8; 3]); + let mut state = BoundaryState::new(); + state.find_boundaries(&block).expect("short input buffers"); + + assert!(state.finish().expect("first flush").is_some()); + assert!(state.finish().expect("second flush").is_none(), "leftover must be consumed once"); + } + + // ======================================================================== + // Bounded carryover on corrupt framing + // ======================================================================== + + /// A corrupt `block_size` must be rejected once the carry passes + /// [`MAX_CARRY_BYTES`], rather than buffering the rest of the stream. + /// + /// Without the bound this terminates *correctly* — `finish` catches the + /// shortfall at EOF — but only after carrying every remaining byte, so peak + /// memory tracks input size and a large corrupt file OOMs before it can + /// print the diagnostic it was about to print. + #[test] + fn a_corrupt_block_size_is_rejected_once_the_carry_passes_the_bound() { + let mut state = BoundaryState::new_no_header(); + // Declare a body far larger than any real record, and never supply it. + let mut first = + u32::try_from(MAX_CARRY_BYTES * 4).expect("fits u32").to_le_bytes().to_vec(); + first.extend_from_slice(&[7u8; 64]); + state.find_boundaries(&first).expect("first block just starts the carry"); + + // Feed blocks until the bound trips. It must trip well before the + // declared size is reached. + let block = vec![9u8; 1024 * 1024]; + let mut fed = first.len(); + let err = loop { + match state.find_boundaries(&block) { + Ok(_) => { + fed += block.len(); + assert!( + fed <= MAX_CARRY_BYTES + block.len(), + "carry ran past the bound without erroring ({fed} bytes fed)", + ); + } + Err(e) => break e, + } + }; + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + } + + /// A header whose declared `l_text` never resolves must be bounded the same + /// way — it is the other path that carries the whole stream forward. + #[test] + fn a_header_that_never_resolves_is_rejected_once_the_carry_passes_the_bound() { + let mut header = fgumi_raw_bam::BAM_MAGIC.to_vec(); + header.extend_from_slice(&u32::MAX.to_le_bytes()); // absurd l_text + let mut state = BoundaryState::new(); + state.find_boundaries(&header).expect("first block just starts the carry"); + + let block = vec![0u8; 1024 * 1024]; + let mut fed = header.len(); + let err = loop { + match state.find_boundaries(&block) { + Ok(_) => { + fed += block.len(); + assert!( + fed <= MAX_CARRY_BYTES + block.len(), + "header carry ran past the bound without erroring ({fed} bytes fed)", + ); + } + Err(e) => break e, + } + }; + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + } + + /// The bound must not reject legitimately large records. An ONT ultra-long + /// read with methylation tags can reach a few MB, and such a record spans + /// many BGZF blocks — so it is carried across dozens of calls before it + /// completes. That is exactly the shape the bound must let through, which is + /// why it sits far above any real record rather than at the queue budget. + #[test] + fn a_multi_megabyte_record_spanning_many_blocks_is_still_accepted() { + const BODY: usize = 8 * 1024 * 1024; // ~8 MB — above any real BAM record + // Compile-time: if the bound is ever lowered under this fixture, the + // test must fail to build rather than silently stop proving anything. + const { assert!(BODY < MAX_CARRY_BYTES) }; + + let mut stream = u32::try_from(BODY).expect("fits u32").to_le_bytes().to_vec(); + stream.extend_from_slice(&vec![3u8; BODY]); + + let mut state = BoundaryState::new_no_header(); + let mut emitted = Vec::new(); + for chunk in stream.chunks(64 * 1024) { + emitted.extend_from_slice( + &state.find_boundaries(chunk).expect("a large but legal record must scan").buffer, + ); + } + assert!(state.finish().expect("nothing left over").is_none()); + assert_eq!(emitted, stream, "the whole record must survive the carry intact"); + } + + /// A header that never completes before EOF leaves the partial header in + /// `leftover`; `finish` then reports it as corruption rather than emitting + /// header bytes as if they were records. + #[test] + fn finish_reports_a_header_that_never_completed_as_corruption() { + let header = bam_header("@HD\tVN:1.6\n", &[("chr1", 1000)]); + let mut state = BoundaryState::new(); + let batch = state.find_boundaries(&header[..12]).expect("partial header buffers"); + assert!(batch.buffer.is_empty()); + + let err = state.finish().expect_err("a partial header is not a valid record stream"); + assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof); + } +} diff --git a/src/lib/pipeline/steps/chain_tests.rs b/src/lib/pipeline/steps/chain_tests.rs new file mode 100644 index 000000000..09f0ee88a --- /dev/null +++ b/src/lib/pipeline/steps/chain_tests.rs @@ -0,0 +1,2202 @@ +//! Cross-step integration tests for the step subset in this port. +//! +//! Every step in `steps/` is a `try_run` state machine whose interesting +//! behavior — held-item retry, drain detection, `Finished` reporting — only +//! occurs when the framework drives it. Unit tests over a step's helpers +//! therefore leave the step body itself unexercised, so these tests assemble +//! real chains and run them through [`Pipeline::run`]: +//! +//! ```text +//! ReplaySource → BgzfDecompress → FindBamBoundaries → DecodeRecords → sink +//! ReplaySource → BgzfDecompress → FindBamBoundaries → ParseBamRecords → sink +//! ReplaySource → BgzfDecompress → FindBamBoundaries → BgzfCompress → sink +//! ReplaySource → TemplatesToRecordBatch → sink +//! ReplaySource → GroupByMi → sink +//! ReplaySource → Process2Ordered → sink A + sink B +//! ``` +//! +//! `ReplaySource` is a test-local source because the real `source/` steps +//! are not part of this port; the block bytes it replays are built by the same +//! BGZF compressor the production writer uses, so the decode path sees exactly +//! what it would in production. +//! +//! Each chain runs at 1 and 4 threads. That is not redundancy: at 1 thread the +//! `Parallel` steps have a single clone and drain trivially, while at 4 the +//! last-worker barrier (only the final clone may close a shared output) is +//! actually exercised. A drain bug that wedges the pipeline shows up as a +//! hang in the 4-thread case only. +//! +//! The full production chain (through `GroupBam`, serialization, and the file +//! sink) is covered by the integration tests that arrive with the remaining +//! step subtrees. + +use std::collections::VecDeque; +use std::io; +use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering}; +use std::sync::{Arc, Mutex}; + +use fgumi_bam_io::GroupKeyConfig; +use fgumi_raw_bam::{RawRecord, SamBuilder}; +use rstest::rstest; + +use crate::grouper::ProcessedPositionGroup; +use crate::pipeline::core::Unpushed; +use crate::pipeline::core::builder::{Pipeline, PipelineConfig}; +use crate::pipeline::core::held::HeldSlot; +use crate::pipeline::core::item::{HeapSize, Ordered}; +use crate::pipeline::core::outputs::OrderedBytesSingle; +use crate::pipeline::core::queues::QueueSpec; +use crate::pipeline::core::reorder::BranchOrdering; +use crate::pipeline::core::step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}; +use crate::pipeline::steps::bgzf::compress::BgzfCompress; +use crate::pipeline::steps::bgzf::decompress::BgzfDecompress; +use crate::pipeline::steps::boundaries::bam::FindBamBoundaries; +use crate::pipeline::steps::group::position::BatchedProcessedPositionGroups; +use crate::pipeline::steps::parse::bam::ParseBamRecords; +use crate::pipeline::steps::parse::decode::{DecodeFromRecords, DecodeRecords}; +use crate::pipeline::steps::types::{BgzfBlock, DecodedRecordBatch, RecordBatch}; +use crate::template::Template; + +/// Per-edge byte budget. Deliberately small relative to the multi-record +/// fixtures so the byte-bounded edges bind mid-run rather than swallowing the +/// whole stream in one go. `the_edge_budget_binds_on_the_multi_record_fixtures` +/// pins that this constant actually has teeth — a future edit that raised it +/// past the fixture volume would silently turn every chain below into an +/// unbounded-queue test. +const EDGE_LIMIT_BYTES: u64 = 16 * 1024; + +// ============================================================================ +// Fixtures +// ============================================================================ + +/// A minimal binary BAM header: magic, empty text, and `n_ref` references. +/// This is the prefix `FindBamBoundaries` must consume and discard. +fn minimal_binary_bam_header(n_ref: u32) -> Vec { + let mut header = Vec::new(); + header.extend_from_slice(fgumi_raw_bam::BAM_MAGIC); + 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 ("r" + NUL) + header.extend_from_slice(b"r\0"); + header.extend_from_slice(&1_000_000u32.to_le_bytes()); // l_ref + } + header +} + +/// `n` mapped, paired records with distinct read names and positions — enough +/// field variety that a group key computed from them is not trivially constant. +fn test_records(n: usize) -> Vec { + (0..n) + .map(|i| { + let mut b = SamBuilder::new(); + b.read_name(format!("read{i}").as_bytes()) + .flags(fgumi_raw_bam::flags::PAIRED | fgumi_raw_bam::flags::FIRST_SEGMENT) + .ref_id(0) + .pos(i32::try_from(100 + i * 10).expect("pos fits i32")) + .cigar_ops(&[4u32 << 4]) // 4M + .sequence(b"ACGT") + .qualities(&[30u8; 4]); + b.build() + }) + .collect() +} + +/// Serialize a header + records into a BAM byte stream, then cut it into BGZF +/// blocks of at most `payload_bytes` uncompressed bytes each. +/// +/// Cutting at a fixed byte size (rather than on record boundaries) is the +/// point: records land straddling block edges, which is exactly the carryover +/// case `FindBamBoundaries` exists to handle. +fn bgzf_blocks_for(records: &[RawRecord], payload_bytes: usize) -> Vec { + let mut stream = minimal_binary_bam_header(1); + 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::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"); + BgzfBlock { + batch_serial: i as u64, + bytes: blocks.remove(0).data, + uncompressed_size: u32::try_from(payload.len()).expect("payload fits u32"), + } + }) + .collect() +} + +/// A four-record template for one queryname: primary R1 and R2 plus an R1 +/// supplementary and an R2 secondary. +/// +/// The split alignments are the point — a flatten written against +/// [`Template::r1`]/[`Template::r2`] would silently drop them, and every record +/// here carries a distinct position, sequence, and length so any dropped, +/// duplicated, or reordered record changes the emitted byte stream. +fn split_alignment_template(qname: &[u8]) -> Template { + use fgumi_raw_bam::flags::{FIRST_SEGMENT, LAST_SEGMENT, PAIRED, SECONDARY, SUPPLEMENTARY}; + + let record = |flags: u16, pos: i32, seq: &[u8], quals: &[u8]| { + let mut b = SamBuilder::new(); + b.read_name(qname) + .flags(flags) + .ref_id(0) + .pos(pos) + // (length << 4) | op, where op 0 is `M`. + .cigar_ops(&[u32::try_from(seq.len()).expect("seq length fits u32") << 4]) + .sequence(seq) + .qualities(quals); + b.build() + }; + + Template::from_records(vec![ + record(PAIRED | FIRST_SEGMENT, 100, b"ACGT", &[30u8; 4]), + record(PAIRED | LAST_SEGMENT, 200, b"TTGCA", &[31u8; 5]), + record(PAIRED | FIRST_SEGMENT | SUPPLEMENTARY, 300, b"GGGCCC", &[32u8; 6]), + record(PAIRED | LAST_SEGMENT | SECONDARY, 400, b"TATATATA", &[33u8; 8]), + ]) + .expect("split-alignment template") +} + +/// A minimal record carrying `MI:Z:` — the only field `GroupByMi` reads. +fn mi_record(mi: &str) -> RawRecord { + let mut b = SamBuilder::new(); + b.read_name(b"read") + .flags(0) + .sequence(b"ACGT") + .qualities(&[30u8; 4]) + .add_string_tag(*crate::sam::SamTag::MI, mi.as_bytes()); + b.build() +} + +/// A minimal record with no `MI` tag at all — dropped by `GroupByMi` and +/// counted against `skipped_no_mi`. +fn record_without_mi() -> RawRecord { + let mut b = SamBuilder::new(); + b.read_name(b"read").flags(0).sequence(b"ACGT").qualities(&[30u8; 4]); + b.build() +} + +/// A minimal record carrying `MI:i:` — an integer where a string is +/// required, so it groups nothing and is counted against +/// `skipped_non_string_mi`. +fn record_with_integer_mi(value: i32) -> RawRecord { + let mut b = SamBuilder::new(); + b.read_name(b"read") + .flags(0) + .sequence(b"ACGT") + .qualities(&[30u8; 4]) + .add_int_tag(*crate::sam::SamTag::MI, value); + b.build() +} + +/// Wrap raw records in a [`DecodedRecordBatch`]. The [`GroupKey`] is left at +/// its default because `GroupByMi` re-reads the MI tag from the record bytes +/// and never consults the pre-computed key. +fn decoded_batch(batch_serial: u64, records: Vec) -> DecodedRecordBatch { + DecodedRecordBatch::new( + batch_serial, + records + .into_iter() + .map(|raw| { + fgumi_bam_io::DecodedRecord::from_raw_bytes(raw, fgumi_bam_io::GroupKey::default()) + }) + .collect(), + ) +} + +// ============================================================================ +// Test-local source and sink steps +// ============================================================================ + +/// Replays pre-built items into the chain. `Exclusive` because it owns mutable +/// cursor state and must not be cloned per worker. +/// +/// Holds a rejected item and retries it on a later tick rather than dropping +/// it — dropping would punch a hole in the ordinal sequence, which the +/// downstream `ByItemOrdinal` reorder stages would then wait on forever. +struct ReplaySource { + items: VecDeque, + held: HeldSlot>, +} + +impl ReplaySource { + fn new(items: Vec) -> Self { + Self { items: items.into(), held: HeldSlot::new() } + } +} + +impl Step for ReplaySource { + type Input = (); + type Outputs = OrderedBytesSingle; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "ReplaySource", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: EDGE_LIMIT_BYTES }], + 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::Contention); + } + } + } + let Some(item) = self.items.pop_front() else { + return Ok(StepOutcome::Finished); + }; + if let Err(unpushed) = ctx.outputs.push(item) { + self.held.put(unpushed); + } + Ok(StepOutcome::Progress) + } +} + +/// Terminal sink that accumulates every item it receives, in arrival order. +/// `Exclusive` so arrival order is the chain's output order with no +/// sink-side interleaving to reason about. +struct CollectSink { + collected: Arc>>, +} + +impl Step for CollectSink +where + T: Send + HeapSize + 'static, +{ + type Input = T; + type Outputs = (); + + fn profile(&self) -> StepProfile { + StepProfile { + name: "CollectSink", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + match ctx.input.pop() { + Some(item) => { + self.collected.lock().expect("sink mutex not poisoned").push(item); + Ok(StepOutcome::Progress) + } + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } +} + +/// Like [`CollectSink`], but ignores its input for the first `remaining_stalls` +/// calls. +/// +/// Stalling the terminal step is what lets an upstream byte-bounded output +/// queue actually fill up. With a sink that drains on every tick the queue +/// returns to zero bytes between pushes, so `cur >= limit` never holds at push +/// time and the producer's held-slot retry path is never taken — at one thread +/// it is structurally unreachable, and at four it depends on scheduling. That +/// makes a "this exercises backpressure" claim untrue exactly when the test +/// looks like it passes. +struct StallingCollectSink { + collected: Arc>>, + remaining_stalls: usize, +} + +impl Step for StallingCollectSink +where + T: Send + HeapSize + 'static, +{ + type Input = T; + type Outputs = (); + + fn profile(&self) -> StepProfile { + StepProfile { + name: "StallingCollectSink", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if self.remaining_stalls > 0 { + self.remaining_stalls -= 1; + return Ok(StepOutcome::NoProgress); + } + match ctx.input.pop() { + Some(item) => { + self.collected.lock().expect("sink mutex not poisoned").push(item); + Ok(StepOutcome::Progress) + } + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } +} + +// ============================================================================ +// Chains +// ============================================================================ + +/// `ReplaySource → BgzfDecompress → FindBamBoundaries → DecodeRecords`. +/// +/// Every record in the input must arrive downstream exactly once, in order, +/// with its bytes intact — across BGZF block boundaries, byte backpressure, +/// and (at 4 threads) parallel decompress/decode workers. +#[rstest] +fn decode_chain_round_trips_every_record_in_order( + #[values(1, 4)] threads: usize, + #[values(1, 64, 500)] n_records: usize, +) { + let records = test_records(n_records); + let blocks = bgzf_blocks_for(&records, 4096); + let collected: Arc>> = Arc::new(Mutex::new(Vec::new())); + + let sink_handle = Arc::clone(&collected); + let builder = Pipeline::builder(); + builder + .chain(ReplaySource::new(blocks)) + .chain(BgzfDecompress::new(EDGE_LIMIT_BYTES)) + .chain(FindBamBoundaries::new(EDGE_LIMIT_BYTES)) + .chain(DecodeRecords::new(GroupKeyConfig::default(), EDGE_LIMIT_BYTES)) + .chain(CollectSink { collected: sink_handle }) + .into_sink_marker(); + let pipeline = builder.build().expect("chain builds"); + pipeline.run(PipelineConfig { threads, ..Default::default() }).expect("chain runs"); + + let batches = collected.lock().expect("mutex not poisoned"); + let seen: Vec<&[u8]> = batches + .iter() + .flat_map(|batch| batch.records().iter().map(fgumi_bam_io::DecodedRecord::raw_bytes)) + .collect(); + let expected: Vec<&[u8]> = records.iter().map(AsRef::as_ref).collect(); + assert_eq!(seen, expected, "records must survive the chain intact and in order"); +} + +/// Same chain but terminating in `ParseBamRecords`, the zero-alloc sibling of +/// `DecodeRecords` that emits shared-buffer `RecordBatch`es instead of owned +/// `DecodedRecord`s. It must yield the same records in the same order. +#[rstest] +fn parse_chain_round_trips_every_record_in_order( + #[values(1, 4)] threads: usize, + #[values(1, 64, 500)] n_records: usize, +) { + let records = test_records(n_records); + let blocks = bgzf_blocks_for(&records, 4096); + let collected: Arc>> = Arc::new(Mutex::new(Vec::new())); + + let sink_handle = Arc::clone(&collected); + let builder = Pipeline::builder(); + builder + .chain(ReplaySource::new(blocks)) + .chain(BgzfDecompress::new(EDGE_LIMIT_BYTES)) + .chain(FindBamBoundaries::new(EDGE_LIMIT_BYTES)) + .chain(ParseBamRecords::new(EDGE_LIMIT_BYTES)) + .chain(CollectSink { collected: sink_handle }) + .into_sink_marker(); + let pipeline = builder.build().expect("chain builds"); + pipeline.run(PipelineConfig { threads, ..Default::default() }).expect("chain runs"); + + let batches = collected.lock().expect("mutex not poisoned"); + let seen: Vec> = + batches.iter().flat_map(|batch| batch.iter_record_bytes().map(<[u8]>::to_vec)).collect(); + let expected: Vec> = records.iter().map(|r| r.as_ref().to_vec()).collect(); + assert_eq!(seen, expected, "records must survive the parse chain intact and in order"); +} + +/// `… → FindBamBoundaries → BgzfCompress` re-compresses the record stream. +/// Decompressing the sink's output must reproduce the record bytes exactly — +/// proving the compress step preserves both content and batch ordering. +#[rstest] +fn compress_chain_output_decompresses_back_to_the_record_stream(#[values(1, 4)] threads: usize) { + let records = test_records(200); + let blocks = bgzf_blocks_for(&records, 4096); + let collected: Arc>> = Arc::new(Mutex::new(Vec::new())); + + let sink_handle = Arc::clone(&collected); + let builder = Pipeline::builder(); + builder + .chain(ReplaySource::new(blocks)) + .chain(BgzfDecompress::new(EDGE_LIMIT_BYTES)) + .chain(FindBamBoundaries::new(EDGE_LIMIT_BYTES)) + .chain(BgzfCompress::new(1, EDGE_LIMIT_BYTES)) + .chain(CollectSink { collected: sink_handle }) + .into_sink_marker(); + let pipeline = builder.build().expect("chain builds"); + pipeline.run(PipelineConfig { threads, ..Default::default() }).expect("chain runs"); + + // Concatenate the emitted BGZF blocks and inflate the whole stream. + let compressed: Vec = collected + .lock() + .expect("mutex not poisoned") + .iter() + .flat_map(|b| b.bytes.clone()) + .collect(); + // Read in bounded batches: `read_raw_blocks` pre-allocates `max_blocks`, + // so it must be given a sane cap rather than "everything". + let mut cursor = io::Cursor::new(compressed); + let mut decompressor = libdeflater::Decompressor::new(); + let mut inflated = Vec::new(); + loop { + let raw_blocks = fgumi_bgzf::reader::read_raw_blocks(&mut cursor, 64) + .expect("output is a valid BGZF stream"); + if raw_blocks.is_empty() { + break; + } + for block in &raw_blocks { + inflated.extend_from_slice( + &fgumi_bgzf::reader::decompress_block(block, &mut decompressor) + .expect("each emitted block inflates"), + ); + } + } + + // The compressed stream is the header-stripped record stream: framed + // `block_size` + body per record, in input order. + let mut expected = Vec::new(); + for record in &records { + let bytes = record.as_ref(); + expected.extend_from_slice(&u32::try_from(bytes.len()).expect("fits u32").to_le_bytes()); + expected.extend_from_slice(bytes); + } + assert_eq!(inflated, expected, "compress must preserve the record stream byte-for-byte"); +} + +/// An input with a header but no records must still complete cleanly and +/// produce nothing — the empty-BAM case, where every step sees drain before it +/// ever sees an item. +#[rstest] +fn a_header_only_input_completes_with_no_records(#[values(1, 4)] threads: usize) { + let blocks = bgzf_blocks_for(&[], 4096); + let collected: Arc>> = Arc::new(Mutex::new(Vec::new())); + + let sink_handle = Arc::clone(&collected); + let builder = Pipeline::builder(); + builder + .chain(ReplaySource::new(blocks)) + .chain(BgzfDecompress::new(EDGE_LIMIT_BYTES)) + .chain(FindBamBoundaries::new(EDGE_LIMIT_BYTES)) + .chain(DecodeRecords::new(GroupKeyConfig::default(), EDGE_LIMIT_BYTES)) + .chain(CollectSink { collected: sink_handle }) + .into_sink_marker(); + let pipeline = builder.build().expect("chain builds"); + pipeline.run(PipelineConfig { threads, ..Default::default() }).expect("chain runs"); + + let batches = collected.lock().expect("mutex not poisoned"); + let total: usize = batches.iter().map(|b| b.records().len()).sum(); + assert_eq!(total, 0, "a header-only BAM yields no records"); +} + +/// A single BGZF block holding the entire stream exercises the path where +/// `FindBamBoundaries` never carries anything over — the whole header and every +/// record arrive together in one input. +#[rstest] +fn a_single_block_input_needs_no_carryover(#[values(1, 4)] threads: usize) { + let records = test_records(50); + // One block big enough for the whole stream. + let blocks = bgzf_blocks_for(&records, 64 * 1024); + assert_eq!(blocks.len(), 1, "fixture must fit in a single BGZF block"); + + let collected: Arc>> = Arc::new(Mutex::new(Vec::new())); + let sink_handle = Arc::clone(&collected); + let builder = Pipeline::builder(); + builder + .chain(ReplaySource::new(blocks)) + .chain(BgzfDecompress::new(EDGE_LIMIT_BYTES)) + .chain(FindBamBoundaries::new(EDGE_LIMIT_BYTES)) + .chain(DecodeRecords::new(GroupKeyConfig::default(), EDGE_LIMIT_BYTES)) + .chain(CollectSink { collected: sink_handle }) + .into_sink_marker(); + let pipeline = builder.build().expect("chain builds"); + pipeline.run(PipelineConfig { threads, ..Default::default() }).expect("chain runs"); + + let batches = collected.lock().expect("mutex not poisoned"); + let total: usize = batches.iter().map(|b| b.records().len()).sum(); + assert_eq!(total, records.len()); +} + +/// `FindBamBoundaries::new_no_header` is the runall-spliced entry point: the +/// stream is already past the header, so byte 0 is a record boundary. Feeding +/// it a header-stripped stream must yield every record. +#[rstest] +fn the_no_header_boundary_variant_consumes_a_header_stripped_stream( + #[values(1, 4)] threads: usize, +) { + let records = test_records(64); + // Build blocks over the record stream ONLY — no BAM header prefix. + let mut stream = Vec::new(); + for record in &records { + let bytes = record.as_ref(); + stream.extend_from_slice(&u32::try_from(bytes.len()).expect("fits u32").to_le_bytes()); + stream.extend_from_slice(bytes); + } + let blocks: Vec = stream + .chunks(1024) + .enumerate() + .map(|(i, payload)| { + let mut compressor = fgumi_bgzf::InlineBgzfCompressor::new(1); + compressor.write_all(payload).expect("compress"); + compressor.flush().expect("flush"); + let mut blocks = compressor.take_blocks(); + assert_eq!(blocks.len(), 1); + BgzfBlock { + batch_serial: i as u64, + bytes: blocks.remove(0).data, + uncompressed_size: u32::try_from(payload.len()).expect("fits u32"), + } + }) + .collect(); + + let collected: Arc>> = Arc::new(Mutex::new(Vec::new())); + let sink_handle = Arc::clone(&collected); + let builder = Pipeline::builder(); + builder + .chain(ReplaySource::new(blocks)) + .chain(BgzfDecompress::new(EDGE_LIMIT_BYTES)) + .chain(FindBamBoundaries::new_no_header(EDGE_LIMIT_BYTES)) + .chain(DecodeRecords::new(GroupKeyConfig::default(), EDGE_LIMIT_BYTES)) + .chain(CollectSink { collected: sink_handle }) + .into_sink_marker(); + let pipeline = builder.build().expect("chain builds"); + pipeline.run(PipelineConfig { threads, ..Default::default() }).expect("chain runs"); + + let batches = collected.lock().expect("mutex not poisoned"); + let seen: Vec<&[u8]> = batches + .iter() + .flat_map(|batch| batch.records().iter().map(fgumi_bam_io::DecodedRecord::raw_bytes)) + .collect(); + let expected: Vec<&[u8]> = records.iter().map(AsRef::as_ref).collect(); + assert_eq!(seen, expected, "no-header mode must treat byte 0 as a record boundary"); +} + +/// `… → ParseBamRecords → DecodeFromRecords` is the runall fusion shape: parse +/// once into shared-buffer batches, then re-attach group keys downstream +/// instead of re-serializing. It must yield the same records, in the same +/// order, as the fused `DecodeRecords` chain above. +#[rstest] +fn decode_from_records_chain_matches_the_fused_decode_chain(#[values(1, 4)] threads: usize) { + let records = test_records(120); + + let collect_via = |use_split: bool| -> Vec> { + let blocks = bgzf_blocks_for(&records, 4096); + let collected: Arc>> = Arc::new(Mutex::new(Vec::new())); + let sink_handle = Arc::clone(&collected); + let builder = Pipeline::builder(); + let head = builder + .chain(ReplaySource::new(blocks)) + .chain(BgzfDecompress::new(EDGE_LIMIT_BYTES)) + .chain(FindBamBoundaries::new(EDGE_LIMIT_BYTES)); + if use_split { + head.chain(ParseBamRecords::new(EDGE_LIMIT_BYTES)) + .chain(DecodeFromRecords::new(GroupKeyConfig::default(), EDGE_LIMIT_BYTES)) + .chain(CollectSink { collected: sink_handle }) + .into_sink_marker(); + } else { + head.chain(DecodeRecords::new(GroupKeyConfig::default(), EDGE_LIMIT_BYTES)) + .chain(CollectSink { collected: sink_handle }) + .into_sink_marker(); + } + let pipeline = builder.build().expect("chain builds"); + pipeline.run(PipelineConfig { threads, ..Default::default() }).expect("chain runs"); + + let batches = collected.lock().expect("mutex not poisoned"); + batches + .iter() + .flat_map(|batch| batch.records().iter().map(|r| r.raw_bytes().to_vec())) + .collect() + }; + + let expected: Vec> = records.iter().map(|r| r.as_ref().to_vec()).collect(); + assert_eq!(collect_via(false), expected, "fused decode chain must round-trip records"); + assert_eq!(collect_via(true), expected, "split parse+decode chain must agree with it"); +} + +/// `ReplaySource → ParseSamChunk` covers the SAM ingest path, whose records +/// converge on the same `DecodedRecordBatch` the BAM chain produces. Driving it +/// through the framework (rather than calling `parse_sam_chunk_into_decoded` +/// directly) is what exercises the step's drain and held-item handling. +#[rstest] +fn sam_chain_decodes_every_line_in_order(#[values(1, 4)] threads: usize) { + use crate::pipeline::steps::parse::sam::ParseSamChunk; + use crate::pipeline::steps::types::SamChunk; + + const HEADER_TEXT: &str = "@HD\tVN:1.6\tSO:unsorted\n@SQ\tSN:chr1\tLN:100000\n"; + let header = { + let mut reader = noodles::sam::io::Reader::new(HEADER_TEXT.as_bytes()); + Arc::new(reader.read_header().expect("header parses")) + }; + + // One chunk per record keeps the ordinal stream long enough that the + // byte-bounded edge rejects pushes and the retry path is taken. + let n_records = 200usize; + let chunks: Vec = (0..n_records) + .map(|i| { + let line = format!("read{i}\t0\tchr1\t{}\t60\t4M\t*\t0\t0\tACGT\tIIII\n", 100 + i); + let bytes = line.into_bytes(); + let line_offsets = vec![0u32, u32::try_from(bytes.len()).expect("fits u32")]; + SamChunk { batch_serial: i as u64, bytes, line_offsets } + }) + .collect(); + + let key_config = + GroupKeyConfig::new_raw_no_cell(fgumi_bam_io::LibraryIndex::from_header(&header)); + let collected: Arc>> = Arc::new(Mutex::new(Vec::new())); + let sink_handle = Arc::clone(&collected); + let builder = Pipeline::builder(); + builder + .chain(ReplaySource::new(chunks)) + .chain(ParseSamChunk::new(Arc::clone(&header), key_config, EDGE_LIMIT_BYTES)) + .chain(CollectSink { collected: sink_handle }) + .into_sink_marker(); + let pipeline = builder.build().expect("chain builds"); + pipeline.run(PipelineConfig { threads, ..Default::default() }).expect("chain runs"); + + let batches = collected.lock().expect("mutex not poisoned"); + let names: Vec> = batches + .iter() + .flat_map(DecodedRecordBatch::records) + .map(|record| { + let body = record.raw_bytes(); + let l_read_name = body[8] as usize; + body[32..32 + l_read_name - 1].to_vec() + }) + .collect(); + let expected: Vec> = (0..n_records).map(|i| format!("read{i}").into_bytes()).collect(); + assert_eq!(names, expected, "SAM lines must decode in order, one record per line"); +} + +/// A stream that ends mid-record must fail the run, not silently drop the +/// truncated tail. `FindBamBoundaries` carries the partial record forward, and +/// at drain its EOF validation rejects it — so the error surfaces through +/// `Pipeline::run` rather than showing up as quietly missing records. +#[rstest] +fn a_truncated_record_stream_fails_the_run(#[values(1, 4)] threads: usize) { + let records = test_records(32); + let mut stream = minimal_binary_bam_header(1); + for record in &records { + let bytes = record.as_ref(); + stream.extend_from_slice(&u32::try_from(bytes.len()).expect("fits u32").to_le_bytes()); + stream.extend_from_slice(bytes); + } + // Chop the final record in half — a `block_size` prefix promising bytes + // that never arrive. + stream.truncate(stream.len() - 20); + + let blocks: Vec = stream + .chunks(1024) + .enumerate() + .map(|(i, payload)| { + let mut compressor = fgumi_bgzf::InlineBgzfCompressor::new(1); + compressor.write_all(payload).expect("compress"); + compressor.flush().expect("flush"); + let mut blocks = compressor.take_blocks(); + assert_eq!(blocks.len(), 1); + BgzfBlock { + batch_serial: i as u64, + bytes: blocks.remove(0).data, + uncompressed_size: u32::try_from(payload.len()).expect("fits u32"), + } + }) + .collect(); + + let collected: Arc>> = Arc::new(Mutex::new(Vec::new())); + let builder = Pipeline::builder(); + builder + .chain(ReplaySource::new(blocks)) + .chain(BgzfDecompress::new(EDGE_LIMIT_BYTES)) + .chain(FindBamBoundaries::new(EDGE_LIMIT_BYTES)) + .chain(DecodeRecords::new(GroupKeyConfig::default(), EDGE_LIMIT_BYTES)) + .chain(CollectSink { collected: Arc::clone(&collected) }) + .into_sink_marker(); + let pipeline = builder.build().expect("chain builds"); + + let err = pipeline + .run(PipelineConfig { threads, ..Default::default() }) + .expect_err("a truncated record stream must fail the run"); + match err { + crate::pipeline::core::signal::PipelineError::Io { step, source } => { + assert_eq!(step, "FindBamBoundaries", "the boundary scanner owns EOF validation"); + assert_eq!(source.kind(), io::ErrorKind::UnexpectedEof); + } + other => panic!("expected an I/O failure from the boundary step, got {other:?}"), + } +} + +/// `ReplaySource → TemplatesToRecordBatch` is the AAM-fusion adapter: it turns +/// the queryname-template view back into the flat-record view the sort ingest +/// consumes. Driving it through `Pipeline::run` (rather than re-running its +/// flattening loop in a unit test) is what exercises `try_run` itself — the +/// held-slot retry, the drain path, and the last-clone `Finished` gate that a +/// `Parallel` step only reaches under the framework. +/// +/// Every record of every template must arrive downstream exactly once, +/// byte-for-byte, in template order, in a batch carrying the input batch's +/// serial — *including* the secondary and supplementary alignments. +#[rstest] +fn templates_to_records_flattens_every_record_including_split_alignments( + #[values(1, 4)] threads: usize, +) { + use crate::pipeline::steps::templates_to_records::TemplatesToRecordBatch; + use crate::pipeline::steps::types::BamTemplateBatch; + + const N_BATCHES: u64 = 200; + const TEMPLATES_PER_BATCH: usize = 2; + const RECORDS_PER_TEMPLATE: usize = 4; + + // Record each template's records in `records` order — the order the step is + // contracted to emit. Reading that off the template rather than hard-coding + // r1/r2/supplementary/secondary keeps this a test of the flattening step, + // not of `Template::from_records`' internal layout. + let mut expected: Vec> = Vec::new(); + let mut batches: Vec = Vec::new(); + for serial in 0..N_BATCHES { + let templates: Vec