diff --git a/CHANGELOG.md b/CHANGELOG.md index b6e4893..7cf53c4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,18 @@ Sections commonly used: Features, Bug fixes, Other changes. ## [Unreleased] +### Other changes + +- BAM output now compresses and encodes on the rayon pool sized by + `--runThreadN`, instead of deflating one BGZF block at a time on the + writing thread. Output is unchanged: the multithreaded BGZF writer + stages into the same buffer size, emits through the same frame writer + and appends the same EOF block, and record encoding is split into + sub-ranges whose concatenation in input order is the serial stream. + Pinned by `multithreaded_bgzf_is_byte_identical_to_the_serial_writer`, + `parallel_encoding_is_byte_identical_at_every_worker_count` and + `bam_file_is_byte_identical_to_the_serial_path`. + ### Features - **STARsolo single-cell quantification (`--soloType`)** — the 10x diff --git a/src/io/bam.rs b/src/io/bam.rs index 1d2a143..7385d5d 100644 --- a/src/io/bam.rs +++ b/src/io/bam.rs @@ -45,20 +45,128 @@ fn bgzf_compression(level: i32) -> bgzf::io::writer::CompressionLevel { } /// Create a BGZF writer with the given STAR compression level. -fn make_bgzf_writer(inner: W, compression: i32) -> bgzf::io::Writer { - bgzf::io::writer::Builder::default() +/// +/// Deflate runs on the global rayon pool, which `run()` sizes from +/// `--runThreadN`, while a dedicated thread writes the finished frames in +/// order. Block boundaries, frame layout and the EOF marker are the same as +/// the single-threaded `bgzf::io::Writer`: both stage into a buffer of +/// `MAX_BUF_SIZE`, both emit through `write_frame`, and libdeflate is +/// deterministic at a fixed level. The bytes are therefore identical to what +/// the serial writer produced, which +/// `multithreaded_bgzf_is_byte_identical_to_the_serial_writer` pins. +/// +/// **The returned writer must not be written to from a rayon worker thread.** +/// It hands each finished block to the pool via `rayon::spawn` and blocks on a +/// channel bounded by the worker count while a dedicated thread drains it. A +/// caller that is itself occupying a worker can therefore fill the channel and +/// then wait for a compression task that has no free worker to run on: a real +/// deadlock, not a slowdown, and it is reachable at any pool size. rustar's +/// writes come from the pipeline's dispatcher thread (`run_batch_pipeline` +/// calls `consume` on its own thread, never on the pool), which is why this is +/// safe here; `write_batch` carries a `debug_assert` so a future call site that +/// breaks the rule fails loudly in tests instead of hanging in production. +fn make_bgzf_writer( + inner: W, + compression: i32, +) -> bgzf::io::MultithreadedWriter { + bgzf::io::multithreaded_writer::Builder::default() .set_compression_level(bgzf_compression(compression)) .build_from_writer(inner) } +/// Below this many records in a batch, encoding on the calling thread beats +/// paying for the fan-out. +const PARALLEL_ENCODE_MIN_RECORDS: usize = 512; + +/// Encode records to the raw BAM byte stream: per record, a little-endian +/// `block_size` followed by the encoded record. +/// +/// This is what `bam::io::Writer::write_alignment_record` writes, and it is +/// produced here by that same call against an in-memory writer, so the bytes +/// cannot drift from the noodles encoder. +fn encode_records_into( + header: &sam::Header, + records: &[RecordBuf], + out: Vec, +) -> Result, Error> { + let mut writer = bam::io::Writer::from(out); + for record in records { + writer.write_alignment_record(header, record)?; + } + Ok(writer.into_inner()) +} + +/// How many records the sorted writers encode at a time on `finish()`. +/// +/// The whole sorted set is already resident as `RecordBuf`s; encoding it in one +/// pass would hold the entire encoded stream alongside them. Encoding in slices +/// keeps the extra buffer bounded while still handing the pool enough records +/// per slice to be worth the fan-out. +const ENCODE_SLICE_RECORDS: usize = 65_536; + +/// Guard the rule documented on `make_bgzf_writer`: writing to the +/// multithreaded BGZF writer from a rayon worker can deadlock. Cheap enough to +/// run on every batch, and only in debug builds. +fn assert_not_on_rayon_worker() { + debug_assert!( + rayon::current_thread_index().is_none(), + "BGZF writes must come from outside the rayon pool: writing from a worker \ + can block the pool on its own compression tasks and deadlock" + ); +} + +/// Encode `records` in bounded slices and write each to `writer`. +fn write_records_chunked( + writer: &mut W, + header: &sam::Header, + records: &[RecordBuf], +) -> Result<(), Error> { + for slice in records.chunks(ENCODE_SLICE_RECORDS) { + let bytes = encode_batch(header, slice)?; + writer.write_all(&bytes)?; + } + Ok(()) +} + +/// Encode a batch to raw BAM bytes, splitting the work across the rayon pool. +/// +/// `bam::io::Writer` carries no state between records other than a scratch +/// buffer it clears at the top of every `write_alignment_record`, so encoding +/// a sub-range against a fresh in-memory writer yields exactly the bytes that +/// sub-range would have contributed to the serial stream. Concatenating the +/// sub-ranges in input order therefore reproduces the serial stream byte for +/// byte, which is why this is output-neutral rather than merely equivalent. +fn encode_batch(header: &sam::Header, batch: &[RecordBuf]) -> Result, Error> { + let threads = rayon::current_num_threads(); + if threads < 2 || batch.len() < PARALLEL_ENCODE_MIN_RECORDS { + return encode_records_into(header, batch, Vec::new()); + } + + use rayon::prelude::*; + let chunk_len = batch.len().div_ceil(threads).max(64); + // `par_chunks` is an indexed parallel iterator, so `collect` restores input + // order regardless of which worker finished first. + let chunks: Vec> = batch + .par_chunks(chunk_len) + .map(|chunk| encode_records_into(header, chunk, Vec::new())) + .collect::>()?; + + let mut out = Vec::with_capacity(chunks.iter().map(Vec::len).sum()); + for chunk in &chunks { + out.extend_from_slice(chunk); + } + Ok(out) +} + /// BAM file writer (streaming, unsorted) /// /// This writer streams BAM records directly to disk as they're generated, /// without buffering or sorting. The output is BGZF-compressed but unsorted. /// Users can sort the output with `samtools sort` if needed. pub struct BamWriter { - writer: bam::io::Writer>>, + writer: bgzf::io::MultithreadedWriter>, header: sam::Header, + finished: bool, } /// BAM file writer that collects all records in memory, sorts by coordinate, @@ -87,10 +195,13 @@ impl BamWriter { compression: i32, ) -> Result { let buf_writer = BufWriter::new(File::create(output_path)?); - let mut bgzf = make_bgzf_writer(buf_writer, compression); - write_bam_header_lenient(&mut bgzf, &header, None)?; - let writer = bam::io::Writer::from(bgzf); - Ok(Self { writer, header }) + let mut writer = make_bgzf_writer(buf_writer, compression); + write_bam_header_lenient(&mut writer, &header, None)?; + Ok(Self { + writer, + header, + finished: false, + }) } /// Create a new BAM writer with header from genome index. @@ -127,15 +238,20 @@ impl BamWriter { /// # Arguments /// * `batch` - Slice of records to write pub fn write_batch(&mut self, batch: &[RecordBuf]) -> Result<(), Error> { - for record in batch { - self.writer.write_alignment_record(&self.header, record)?; - } + assert_not_on_rayon_worker(); + let bytes = encode_batch(&self.header, batch)?; + self.writer.write_all(&bytes)?; Ok(()) } /// Flush and close BAM file pub fn finish(&mut self) -> Result<(), Error> { - self.writer.finish(&self.header)?; + // `MultithreadedWriter::finish` shuts the workers down and panics if + // called twice; `Drop` calls it too when it has not run yet. + if !self.finished { + self.writer.finish()?; + self.finished = true; + } log::info!("BAM file written successfully"); Ok(()) } @@ -198,13 +314,10 @@ impl SortedBamWriter { }); let buf_writer = BufWriter::new(File::create(&self.output_path)?); - let mut bgzf = make_bgzf_writer(buf_writer, self.compression); - write_bam_header_lenient(&mut bgzf, &self.header, Some("coordinate"))?; - let mut bam_writer = bam::io::Writer::from(bgzf); - for record in &self.records { - bam_writer.write_alignment_record(&self.header, record)?; - } - bam_writer.finish(&self.header)?; + let mut writer = make_bgzf_writer(buf_writer, self.compression); + write_bam_header_lenient(&mut writer, &self.header, Some("coordinate"))?; + write_records_chunked(&mut writer, &self.header, &self.records)?; + writer.finish()?; log::info!("Sorted BAM written ({} records)", self.records.len()); Ok(()) } @@ -218,15 +331,13 @@ impl SortedBamWriter { _ => (usize::MAX, 0), }); - let stdout = std::io::stdout(); - let buf_writer = BufWriter::new(stdout.lock()); - let mut bgzf = make_bgzf_writer(buf_writer, self.compression); - write_bam_header_lenient(&mut bgzf, &self.header, Some("coordinate"))?; - let mut bam_writer = bam::io::Writer::from(bgzf); - for record in &self.records { - bam_writer.write_alignment_record(&self.header, record)?; - } - bam_writer.finish(&self.header)?; + // `std::io::stdout()` rather than a lock guard: the BGZF writer owns + // its sink on a worker thread, so the sink has to be `'static`. + let buf_writer = BufWriter::new(std::io::stdout()); + let mut writer = make_bgzf_writer(buf_writer, self.compression); + write_bam_header_lenient(&mut writer, &self.header, Some("coordinate"))?; + write_records_chunked(&mut writer, &self.header, &self.records)?; + writer.finish()?; log::info!( "Sorted BAM written to stdout ({} records)", self.records.len() @@ -371,31 +482,38 @@ fn render_sam_text_lenient(header: &sam::Header, sort_order: Option<&str>) -> Ve /// Streaming unsorted BAM writer that writes to stdout. pub struct BamStdoutWriter { - writer: bam::io::Writer>>, + writer: bgzf::io::MultithreadedWriter>, header: sam::Header, + finished: bool, } impl BamStdoutWriter { pub fn create(genome: &crate::genome::Genome, params: &Parameters) -> Result { let header = crate::io::sam::build_sam_header(genome, params)?; - let mut bgzf = make_bgzf_writer( + let mut writer = make_bgzf_writer( BufWriter::new(std::io::stdout()), params.out_bam_compression, ); - write_bam_header_lenient(&mut bgzf, &header, None)?; - let writer = bam::io::Writer::from(bgzf); - Ok(Self { writer, header }) + write_bam_header_lenient(&mut writer, &header, None)?; + Ok(Self { + writer, + header, + finished: false, + }) } pub fn write_batch(&mut self, batch: &[RecordBuf]) -> Result<(), Error> { - for record in batch { - self.writer.write_alignment_record(&self.header, record)?; - } + assert_not_on_rayon_worker(); + let bytes = encode_batch(&self.header, batch)?; + self.writer.write_all(&bytes)?; Ok(()) } pub fn finish(&mut self) -> Result<(), Error> { - self.writer.finish(&self.header)?; + if !self.finished { + self.writer.finish()?; + self.finished = true; + } Ok(()) } } @@ -661,6 +779,156 @@ mod tests { ); } + /// The pre-change write path: one serial `bgzf::io::Writer`, records fed + /// one at a time through `bam::io::Writer`. Every neutrality test below + /// compares against these bytes, so what is being asserted is not "the two + /// new halves agree with each other" but "the output did not move". + fn serial_reference_bytes(header: &sam::Header, records: &[RecordBuf]) -> Vec { + let mut bgzf = bgzf::io::writer::Builder::default() + .set_compression_level(bgzf_compression(1)) + .build_from_writer(Vec::new()); + write_bam_header_lenient(&mut bgzf, header, None).unwrap(); + let mut writer = bam::io::Writer::from(bgzf); + for record in records { + writer.write_alignment_record(header, record).unwrap(); + } + writer.try_finish().unwrap(); + writer.into_inner().into_inner() + } + + fn unmapped_records(params: &Parameters, n: usize) -> Vec { + (0..n) + .map(|i| { + crate::io::sam::SamWriter::build_unmapped_record( + &format!("read{i}"), + &[0, 1, 2, 3, 3, 2, 1, 0], + &[30; 8], + params, + crate::stats::UnmappedReason::Other, + ) + .unwrap() + }) + .collect() + } + + /// The multithreaded BGZF writer stages into the same `MAX_BUF_SIZE` + /// buffer, emits through the same `write_frame`, and appends the same EOF + /// block as the serial writer, so switching to it must not move a byte. + /// Enough records to span several BGZF blocks, so block boundaries are + /// actually exercised rather than a single short block. + #[test] + fn multithreaded_bgzf_is_byte_identical_to_the_serial_writer() { + let params = default_params(); + let header = sam::Header::default(); + let records = unmapped_records(¶ms, 4000); + let expected = serial_reference_bytes(&header, &records); + + let mut writer = make_bgzf_writer(Vec::new(), 1); + write_bam_header_lenient(&mut writer, &header, None).unwrap(); + for record in &records { + let bytes = + encode_records_into(&header, std::slice::from_ref(record), Vec::new()).unwrap(); + writer.write_all(&bytes).unwrap(); + } + let got = writer.finish().unwrap(); + + assert_eq!(got, expected); + } + + /// Splitting the batch across the pool is only sound if concatenating the + /// sub-ranges in input order reproduces the serial stream. Checked at + /// several pool sizes, including one worker (which takes the serial branch) + /// and a size that does not divide the batch evenly. + #[test] + fn parallel_encoding_is_byte_identical_at_every_worker_count() { + let params = default_params(); + let header = sam::Header::default(); + let records = unmapped_records(¶ms, 3000); + let expected = encode_records_into(&header, &records, Vec::new()).unwrap(); + + for threads in [1, 3, 4, 7, 16] { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap(); + let got = pool.install(|| encode_batch(&header, &records)).unwrap(); + assert_eq!(got, expected, "encoding diverged at {threads} workers"); + } + } + + /// Batches under the fan-out threshold take the serial branch; that branch + /// has to produce the same bytes as the parallel one, not merely a valid + /// stream. + #[test] + fn a_batch_below_the_fanout_threshold_encodes_the_same_bytes() { + let params = default_params(); + let header = sam::Header::default(); + let records = unmapped_records(¶ms, PARALLEL_ENCODE_MIN_RECORDS - 1); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(8) + .build() + .unwrap(); + let got = pool.install(|| encode_batch(&header, &records)).unwrap(); + assert_eq!( + got, + encode_records_into(&header, &records, Vec::new()).unwrap() + ); + } + + /// The sorted writers encode in bounded slices rather than in one pass, so + /// the slice boundary must not show up in the output. + #[test] + fn chunked_encoding_matches_a_single_pass() { + let params = default_params(); + let header = sam::Header::default(); + let records = unmapped_records(¶ms, 2500); + let expected = encode_records_into(&header, &records, Vec::new()).unwrap(); + + let mut got = Vec::new(); + for slice in records.chunks(700) { + got.extend_from_slice(&encode_batch(&header, slice).unwrap()); + } + assert_eq!(got, expected); + } + + /// End to end through `BamWriter`: the file on disk is byte-identical to + /// what the serial path wrote. + /// + /// Deliberately not wrapped in a `ThreadPool::install`. Doing so puts the + /// caller on a worker and deadlocks, for the reason spelled out on + /// `make_bgzf_writer` — this test found that the hard way. Worker-count + /// invariance of the part that has one is covered by + /// `parallel_encoding_is_byte_identical_at_every_worker_count`; BGZF block + /// boundaries do not depend on the worker count at all, since the staging + /// buffer fills to a fixed size before any block is handed to the pool. + #[test] + fn bam_file_is_byte_identical_to_the_serial_path() { + let genome = create_test_genome(); + let params = default_params(); + let header = crate::io::sam::build_sam_header(&genome, ¶ms).unwrap(); + let records = unmapped_records(¶ms, 2000); + let expected = serial_reference_bytes(&header, &records); + + let temp_file = NamedTempFile::new().unwrap(); + let mut writer = BamWriter::create(temp_file.path(), &genome, ¶ms).unwrap(); + writer.write_batch(&records).unwrap(); + writer.finish().unwrap(); + + assert_eq!(std::fs::read(temp_file.path()).unwrap(), expected); + } + + /// `finish()` shuts the BGZF workers down, and calling it twice would panic + /// inside noodles. `Drop` also calls it. Both paths have to stay safe. + #[test] + fn finishing_twice_is_not_an_error() { + let genome = create_test_genome(); + let params = default_params(); + let temp_file = NamedTempFile::new().unwrap(); + let mut writer = BamWriter::create(temp_file.path(), &genome, ¶ms).unwrap(); + writer.finish().unwrap(); + writer.finish().unwrap(); + } + #[test] fn test_sorted_bam_limit_ram_unlimited() { let genome = create_test_genome(); diff --git a/test/bench_bam_write.sh b/test/bench_bam_write.sh new file mode 100755 index 0000000..5d4532c --- /dev/null +++ b/test/bench_bam_write.sh @@ -0,0 +1,91 @@ +#!/usr/bin/env bash +# Interleaved A/B of BAM writing, before and after the parallel writer. +# +# test/bench_bam_write.sh [threads] [pairs] [baseRef] +# +# Two binaries are built: one from `baseRef` (default origin/main) and one from +# the working tree. They are then alternated pair by pair, with the order +# flipped on even pairs so a machine that drifts (thermal, page cache warming) +# cannot favour whichever runs first. +# +# Each pair times three output modes: +# +# None alignment only, no BAM written +# BAM Unsorted streaming write +# BAM SortedByCoordinate sort then write +# +# `None` is not decoration. Only the difference between a BAM mode and `None` +# is the work this change touches; the total is mostly alignment and hides it. +# A previous measurement on this repo reported a BAM-writer difference that the +# workload could not have resolved, because BAM writing was 1-4% of the run. +# Report the delta, and if the delta is smaller than the run-to-run spread, say +# that the workload cannot settle the question rather than publishing a median. +set -euo pipefail + +GENOME_DIR=${1:?usage: bench_bam_write.sh [threads] [pairs] [baseRef]} +READS=${2:?usage: bench_bam_write.sh [threads] [pairs] [baseRef]} +THREADS=${3:-16} +PAIRS=${4:-6} +BASE_REF=${5:-origin/main} + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +WORK=$(mktemp -d) +BASE_TREE=$(mktemp -d) +trap 'rm -rf "$WORK"; git -C "$ROOT" worktree remove --force "$BASE_TREE" 2>/dev/null || rm -rf "$BASE_TREE"' EXIT + +# Refuse to measure on a busy machine. A single unrelated job saturating the +# cores does not just add noise, it can invert the result, and it is invisible +# in the numbers afterwards: the series simply drifts. +# +# The check is on CPU idle, not on load average. Load average is an exponential +# average over minutes, so it stays high long after the offending job has gone +# and would refuse to measure on a machine that is now perfectly quiet. +cpu_idle() { # percent idle, sampled over one second + top -l 2 -n 0 2>/dev/null | awk '/CPU usage/ {gsub("%","",$(NF-1)); v=$(NF-1)} END {print v+0}' +} +IDLE=$(cpu_idle) +if awk -v i="$IDLE" 'BEGIN{exit !(i < 80)}'; then + echo "CPU is only ${IDLE}% idle; other work is running and the numbers would not mean anything." >&2 + echo "wait for the machine to be idle, or set BENCH_IGNORE_LOAD=1 to override." >&2 + [ "${BENCH_IGNORE_LOAD:-0}" = "1" ] || exit 1 +fi + +echo "building new (working tree)" >&2 +cargo build --release --manifest-path "$ROOT/Cargo.toml" >&2 +cp "$ROOT/target/release/rustar-aligner" "$WORK/new" + +echo "building old ($BASE_REF)" >&2 +rm -rf "$BASE_TREE" +git -C "$ROOT" worktree add --detach "$BASE_TREE" "$BASE_REF" >&2 +cargo build --release --manifest-path "$BASE_TREE/Cargo.toml" >&2 +cp "$BASE_TREE/target/release/rustar-aligner" "$WORK/old" + +run() { # $1 = binary, $2 = outSAMtype words + local dir="$WORK/run" + rm -rf "$dir" + mkdir -p "$dir" + local t0 t1 + t0=$(python3 -c 'import time; print(time.time())') + # shellcheck disable=SC2086 # $2 is deliberately two words for "BAM Unsorted" + "$1" --genomeDir "$GENOME_DIR" --readFilesIn "$READS" --runThreadN "$THREADS" \ + --outSAMtype $2 --outFileNamePrefix "$dir/" >/dev/null 2>&1 + t1=$(python3 -c "import time; print(f'{time.time()-$t0:.2f}')") + echo "$t1" +} + +# One discarded run so the reads and the index are in page cache for pair 1. +run "$WORK/new" "None" >/dev/null + +for mode in "None" "BAM Unsorted" "BAM SortedByCoordinate"; do + echo "=== --outSAMtype $mode ===" + for i in $(seq 1 "$PAIRS"); do + if (( i % 2 )); then + o=$(run "$WORK/old" "$mode") + n=$(run "$WORK/new" "$mode") + else + n=$(run "$WORK/new" "$mode") + o=$(run "$WORK/old" "$mode") + fi + echo "pair$i old=${o}s new=${n}s idle=$(cpu_idle)%" + done +done