From e17db4fedea9aa10dbeb0ddd1d1015f35710f256 Mon Sep 17 00:00:00 2001 From: SuhasSrinivasan <32346517+SuhasSrinivasan@users.noreply.github.com> Date: Wed, 5 Aug 2026 15:27:18 -0700 Subject: [PATCH 1/7] test: cover short text writes and flush failures --- modkit-core/src/writers.rs | 129 +++++++++++++++++++++++++++++++++++++ 1 file changed, 129 insertions(+) diff --git a/modkit-core/src/writers.rs b/modkit-core/src/writers.rs index e1b16e93..4ee922a5 100644 --- a/modkit-core/src/writers.rs +++ b/modkit-core/src/writers.rs @@ -1473,3 +1473,132 @@ where Ok(total_rows as u64) } } + +#[cfg(test)] +mod writer_tests { + use super::{ + BedMethylWriter, BedMethylWriter2, PileupWriter, RecordingWriter, + TsvWriter, + }; + use crate::mod_base_code::ModCodeRepr; + use crate::pileup::{ModBasePileup2, PileupFeatureCounts2}; + use indicatif::ProgressBar; + use std::io::{self, BufWriter, Write}; + use std::panic::{catch_unwind, AssertUnwindSafe}; + use std::sync::{Arc, Mutex}; + + #[derive(Clone)] + struct ShortWriter { + bytes: Arc>>, + max_write: usize, + fail_flush: bool, + } + + impl ShortWriter { + fn new(max_write: usize) -> Self { + assert!(max_write > 0); + Self { + bytes: Arc::new(Mutex::new(Vec::new())), + max_write, + fail_flush: false, + } + } + + fn failing_flush(max_write: usize) -> Self { + Self { fail_flush: true, ..Self::new(max_write) } + } + + fn bytes(&self) -> Vec { + self.bytes.lock().unwrap().clone() + } + } + + impl Write for ShortWriter { + fn write(&mut self, buf: &[u8]) -> io::Result { + let count = buf.len().min(self.max_write); + self.bytes.lock().unwrap().extend_from_slice(&buf[..count]); + Ok(count) + } + + fn flush(&mut self) -> io::Result<()> { + if self.fail_flush { + Err(io::Error::new(io::ErrorKind::Other, "flush failed")) + } else { + Ok(()) + } + } + } + + fn feature_counts() -> PileupFeatureCounts2 { + PileupFeatureCounts2::new( + 7, + '+', + 4, + ModCodeRepr::Code('m'), + 3, + 1, + 0, + 0, + 0, + 0, + 0, + 0, + ) + } + + fn pileup() -> ModBasePileup2 { + ModBasePileup2 { + chrom_name: "chr1".to_string(), + position_feature_counts: vec![feature_counts()], + interval_width: 1, + stride: 1, + failed_records: 0, + phased_feature_counts: [Vec::new(), Vec::new()], + } + } + + #[test] + fn tsv_writer_retries_short_writes_without_changing_bytes() { + let sink = ShortWriter::new(2); + let mut writer = TsvWriter { writer: sink.clone() }; + let expected = b"alpha\tbeta\n"; + + assert_eq!(writer.write(expected).unwrap(), expected.len()); + assert_eq!(sink.bytes(), expected); + } + + #[test] + fn bedmethyl_writer_retries_short_writes_without_changing_row() { + let sink = ShortWriter::new(3); + let buf_writer = BufWriter::with_capacity(1, sink.clone()); + let mut writer = BedMethylWriter { buf_writer, tabs_and_spaces: false }; + + assert_eq!(PileupWriter::write(&mut writer, pileup(), &[]).unwrap(), 1); + writer.buf_writer.flush().unwrap(); + assert_eq!( + sink.bytes(), + b"chr1\t7\t8\tm\t4\t+\t7\t8\t255,0,0\t4\t25.00\t1\t3\t0\t0\t0\t0\t0\n" + ); + } + + #[test] + fn bedmethyl_writer_returns_flush_error_instead_of_panicking() { + let sink = ShortWriter::failing_flush(usize::MAX); + let (return_mem, _returned) = crossbeam_channel::unbounded(); + let mut writer = BedMethylWriter2 { + buff: std::io::Cursor::new(vec![0u8; 1 << 20]), + inner: RecordingWriter { inner: sink, pb: ProgressBar::hidden() }, + return_mem, + bedrmod_spec: false, + }; + + let result = catch_unwind(AssertUnwindSafe(|| { + PileupWriter::write(&mut writer, pileup(), &[]) + })); + assert!( + result.is_ok(), + "writer panicked instead of returning the error" + ); + assert!(result.unwrap().is_err()); + } +} From 039d5cad35dad26dcb813721d89144cf996c50fc Mon Sep 17 00:00:00 2001 From: SuhasSrinivasan <32346517+SuhasSrinivasan@users.noreply.github.com> Date: Wed, 5 Aug 2026 15:37:34 -0700 Subject: [PATCH 2/7] fix: make text output writes complete and fallible --- modkit-core/src/bedmethyl_util/subcommands.rs | 10 +- modkit-core/src/dmr/isoform/mod.rs | 4 +- modkit-core/src/dmr/pairwise.rs | 6 +- modkit-core/src/dmr/single_site.rs | 6 +- modkit-core/src/dmr/subcommands.rs | 9 +- modkit-core/src/writers.rs | 99 ++++++++++--------- 6 files changed, 77 insertions(+), 57 deletions(-) diff --git a/modkit-core/src/bedmethyl_util/subcommands.rs b/modkit-core/src/bedmethyl_util/subcommands.rs index 96cf1b56..b1132b01 100644 --- a/modkit-core/src/bedmethyl_util/subcommands.rs +++ b/modkit-core/src/bedmethyl_util/subcommands.rs @@ -306,7 +306,7 @@ impl EntryMergeBedMethyl { } }; if self.with_header { - writer.write(bedmethyl_header().as_bytes())?; + writer.write_all(bedmethyl_header().as_bytes())?; } let readers = self @@ -444,7 +444,7 @@ impl EntryMergeBedMethyl { .collect::>() }); for row in rows { - writer.write(row.as_bytes())?; + writer.write_all(row.as_bytes())?; rows_written.inc(1); } } @@ -456,6 +456,7 @@ impl EntryMergeBedMethyl { } } + writer.flush()?; Ok(()) } } @@ -745,7 +746,7 @@ impl EntryMapToGenome { p @ _ => Box::new(BufWriter::new(File::create(p)?)), }; if self.header { - writer.write(bedmethyl_header().as_bytes())?; + writer.write_all(bedmethyl_header().as_bytes())?; } reader.fetch(tid, 0, tm.transcript_len)?; @@ -777,7 +778,7 @@ impl EntryMapToGenome { bml.chrom = tm.chrom.clone(); bml.interval = Iv { start: genome_start, stop: genome_stop, val: () }; - writer.write(bml.to_line().as_bytes())?; + writer.write_all(bml.to_line().as_bytes())?; processed_records.inc(1); } @@ -789,6 +790,7 @@ impl EntryMapToGenome { ); }); + writer.flush()?; Ok(()) } } diff --git a/modkit-core/src/dmr/isoform/mod.rs b/modkit-core/src/dmr/isoform/mod.rs index e5bc52cd..1791ce50 100644 --- a/modkit-core/src/dmr/isoform/mod.rs +++ b/modkit-core/src/dmr/isoform/mod.rs @@ -1294,7 +1294,7 @@ impl GeneIsoformDmr { self.gene.gene_name.as_ref(), emit_full_results, ); - writer.write(row.as_bytes())?; + writer.write_all(row.as_bytes())?; records_written = records_written.saturating_add(1); } } @@ -2059,7 +2059,7 @@ impl GeneTxDmr { single_mod_code, emit_full_results, ); - writer.write(row.as_bytes())?; + writer.write_all(row.as_bytes())?; records_written = records_written.saturating_add(1); } } diff --git a/modkit-core/src/dmr/pairwise.rs b/modkit-core/src/dmr/pairwise.rs index 21d2e239..bcc6b1c5 100644 --- a/modkit-core/src/dmr/pairwise.rs +++ b/modkit-core/src/dmr/pairwise.rs @@ -183,7 +183,8 @@ pub(super) fn run_pairwise_dmr( multi_progress: MultiProgress, ) -> anyhow::Result<(usize, FxHashMap)> { if header { - writer.write(ModificationCounts::header(a_name, b_name).as_bytes())?; + writer + .write_all(ModificationCounts::header(a_name, b_name).as_bytes())?; } let (snd, rcv) = crossbeam_channel::bounded(1000); @@ -254,7 +255,7 @@ pub(super) fn run_pairwise_dmr( for result in results { match result { Ok(counts) => { - writer.write(counts.to_row()?.as_bytes())?; + writer.write_all(counts.to_row()?.as_bytes())?; success_count += 1; pb.inc(1); } @@ -305,6 +306,7 @@ pub(super) fn run_pairwise_dmr( if let Some(e) = err { Err(e.into()) } else { + writer.flush()?; Ok((success_count, region_error_counts)) } } diff --git a/modkit-core/src/dmr/single_site.rs b/modkit-core/src/dmr/single_site.rs index 6e0f6e62..6c87fa8a 100644 --- a/modkit-core/src/dmr/single_site.rs +++ b/modkit-core/src/dmr/single_site.rs @@ -155,7 +155,7 @@ impl SingleSiteDmrAnalysis { } if self.header { - writer.write( + writer.write_all( SingleSiteDmrScore::header(multiple_samples, matched_samples) .as_bytes(), )?; @@ -307,7 +307,7 @@ impl SingleSiteDmrAnalysis { for result in results { match result { Ok(scores) => { - writer.write( + writer.write_all( scores .to_row( multiple_samples, @@ -370,6 +370,7 @@ impl SingleSiteDmrAnalysis { success_count, failure_counter.position(), ); + writer.flush()?; Ok(()) } } @@ -1174,6 +1175,7 @@ impl DmrSegmenter for HmmDmrSegmenter { "HMM segmenter finished, wrote {} segments", self.segments_written.position() ); + self.writer.flush()?; Ok(()) } } diff --git a/modkit-core/src/dmr/subcommands.rs b/modkit-core/src/dmr/subcommands.rs index 66658e15..306f4544 100644 --- a/modkit-core/src/dmr/subcommands.rs +++ b/modkit-core/src/dmr/subcommands.rs @@ -1207,7 +1207,7 @@ impl EntryDmrIsoform { self.emit_full_results, ) }) { - writer.write(row.as_bytes())?; + writer.write_all(row.as_bytes())?; } if let Some(plot_dir) = self.plot.as_ref() { if !plot_dir.exists() { @@ -1294,6 +1294,7 @@ impl EntryDmrIsoform { } } } + writer.flush()?; Ok(()) } @@ -1432,6 +1433,7 @@ impl EntryDmrIsoform { info!("finished, processed {} genes", pb.position()); }); + writer.flush()?; Ok(()) } @@ -1451,7 +1453,7 @@ impl EntryDmrIsoform { Box::new(BufWriter::new(fh)) } }; - writer.write( + writer.write_all( GeneIsoformDmrRecord::::header( self.emit_full_results, ) @@ -1883,6 +1885,7 @@ impl EntryGeneTx { multi_progress.suspend(|| info!("{err_table}")); } + writer.flush()?; Ok(()) } @@ -1903,7 +1906,7 @@ impl EntryGeneTx { Box::new(BufWriter::new(fh)) } }; - writer.write( + writer.write_all( GeneIsoformDmrRecord::::header( single_mod_code, self.emit_full_results, diff --git a/modkit-core/src/writers.rs b/modkit-core/src/writers.rs index 4ee922a5..2474886d 100644 --- a/modkit-core/src/writers.rs +++ b/modkit-core/src/writers.rs @@ -197,16 +197,15 @@ impl PileupWriter for BedMethylWriter2 { &mut self.buff, pfc, self.bedrmod_spec, - ) - .unwrap(); + )?; let pos = self.buff.position() as usize; if pos >= 1 << 20 { - self.inner.write(&self.buff.get_ref()[..pos]).unwrap(); + self.inner.write(&self.buff.get_ref()[..pos])?; self.buff.set_position(0); } } let pos = self.buff.position() as usize; - self.inner.write(&self.buff.get_ref()[..pos]).unwrap(); + self.inner.write(&self.buff.get_ref()[..pos])?; let _ = self.return_mem.send(item); Ok(n_rows) } @@ -256,7 +255,7 @@ impl BedMethylWriter { with_header: bool, ) -> anyhow::Result { if with_header { - buf_writer.write(Self::header().as_bytes())?; + buf_writer.write_all(Self::header().as_bytes())?; } Ok(Self { buf_writer, tabs_and_spaces }) @@ -292,7 +291,7 @@ impl BedMethylWriter { let pos = buff.position() as usize; writer - .write(&buff.get_ref()[..pos]) + .write_all(&buff.get_ref()[..pos]) .with_context(|| "failed to write row")?; Ok(()) @@ -383,7 +382,7 @@ impl PileupWriter for BedMethylWriter { pattern.n_nocall, ); self.buf_writer - .write(row.as_bytes()) + .write_all(row.as_bytes()) .with_context(|| "failed to write row")?; rows_written += 1; } @@ -410,7 +409,7 @@ impl MultipleMotifBedmethylWriter> { ) -> anyhow::Result { let mut writer = BufWriter::new(stdout()); if with_header { - writer.write(bedmethyl_header().as_bytes())?; + writer.write_all(bedmethyl_header().as_bytes())?; } else if bed_rmod_args.enabled() { let modified_bases_options = modified_bases_options.ok_or_else(|| { @@ -418,7 +417,7 @@ impl MultipleMotifBedmethylWriter> { })?; let bedrmod_header = bed_rmod_args.header(&header, modified_bases_options)?; - writer.write(bedrmod_header.as_bytes())?; + writer.write_all(bedrmod_header.as_bytes())?; } let write_pb = multi_progress.add(get_ticker_with_rate()); @@ -720,7 +719,12 @@ pub struct TsvWriter { impl TsvWriter { pub fn write(&mut self, raw: &[u8]) -> std::io::Result { - self.writer.write(raw) + self.writer.write_all(raw)?; + Ok(raw.len()) + } + + pub fn flush(&mut self) -> std::io::Result<()> { + self.writer.flush() } } @@ -756,7 +760,7 @@ impl TsvWriter> { let fh = File::create(path)?; let mut buf_writer = BufWriter::new(fh); if let Some(header) = header { - buf_writer.write(format!("{header}\n").as_bytes())?; + buf_writer.write_all(format!("{header}\n").as_bytes())?; } Ok(Self { writer: buf_writer }) } @@ -789,8 +793,8 @@ impl TsvWriter> { .unwrap() .from_writer(out_fh); if let Some(header) = header { - writer.write(header.as_bytes())?; - writer.write(&['\n' as u8])?; + writer.write_all(header.as_bytes())?; + writer.write_all(&['\n' as u8])?; } Ok(Self { writer }) @@ -799,10 +803,8 @@ impl TsvWriter> { impl OutWriter for TsvWriter { fn write(&mut self, item: String) -> anyhow::Result { - self.writer - .write(item.as_bytes()) - .map(|b| b as u64) - .map_err(|e| anyhow!("{e}")) + self.writer.write_all(item.as_bytes())?; + Ok(item.len() as u64) } } @@ -878,7 +880,7 @@ impl<'a, W: Write> OutWriter> for TsvWriter { item.total_reads_used )); - self.writer.write(report.as_bytes())?; + self.writer.write_all(report.as_bytes())?; Ok(1) } } @@ -1174,14 +1176,12 @@ impl OutWriter for MultiTableWriter { ); tab.to_csv_writer(csv_writer)?; match HtmlRenderer::new("Counts", 800, 800).render(&counts_chart) { - Ok(blob) => { - counts_plot_fh.write(blob.as_bytes()).map(|_x| ())? - } + Ok(blob) => counts_plot_fh.write_all(blob.as_bytes())?, Err(e) => debug!("failed to render counts plot, {e:?}"), } match HtmlRenderer::new("Proportions", 800, 800).render(&prop_chart) { - Ok(blob) => prop_plot_fh.write(blob.as_bytes()).map(|_x| ())?, + Ok(blob) => prop_plot_fh.write_all(blob.as_bytes())?, Err(e) => debug!("failed to render proportions plot, {e:?}"), } } @@ -1419,55 +1419,66 @@ where let total_rows = combined_counts.len() + hp1.len() + hp2.len(); // TODO: make the "buff"s part of the object. - std::thread::scope(|scope| { - let hp1_handle = scope.spawn(|| { + std::thread::scope(|scope| -> anyhow::Result<()> { + let hp1_handle = scope.spawn(|| -> anyhow::Result<()> { let mut buff = Cursor::new(vec![0u8; 1 << 20]); for pfc in hp1.iter().filter(|x| x.is_valid()) { - format_feature_counts2(chrom_name, &mut buff, pfc, false) - .unwrap(); + format_feature_counts2(chrom_name, &mut buff, pfc, false)?; let pos = buff.position() as usize; if pos >= 1 << 20 { - self.hp1_writer.write(&buff.get_ref()[..pos]).unwrap(); + self.hp1_writer.write(&buff.get_ref()[..pos])?; buff.set_position(0); } } let pos = buff.position() as usize; - self.hp1_writer.write(&buff.get_ref()[..pos]).unwrap(); + self.hp1_writer.write(&buff.get_ref()[..pos])?; + Ok(()) }); - let hp2_handle = scope.spawn(|| { + let hp2_handle = scope.spawn(|| -> anyhow::Result<()> { let mut buff = Cursor::new(vec![0u8; 1 << 20]); for pfc in hp2.iter().filter(|x| x.is_valid()) { - format_feature_counts2(chrom_name, &mut buff, pfc, false) - .unwrap(); + format_feature_counts2(chrom_name, &mut buff, pfc, false)?; let pos = buff.position() as usize; if pos >= 1 << 20 { - self.hp2_writer.write(&buff.get_ref()[..pos]).unwrap(); + self.hp2_writer.write(&buff.get_ref()[..pos])?; buff.set_position(0); } } let pos = buff.position() as usize; - self.hp2_writer.write(&buff.get_ref()[..pos]).unwrap(); + self.hp2_writer.write(&buff.get_ref()[..pos])?; + Ok(()) }); - let combined_handle = scope.spawn(|| { + let combined_handle = scope.spawn(|| -> anyhow::Result<()> { let mut buff = Cursor::new(vec![0u8; 1 << 20]); for pfc in combined_counts.iter().filter(|x| x.is_valid()) { - format_feature_counts2(&chrom_name, &mut buff, pfc, false) - .unwrap(); + format_feature_counts2(&chrom_name, &mut buff, pfc, false)?; let pos = buff.position() as usize; if pos >= 1 << 20 { - self.combined_writer - .write(&buff.get_ref()[..pos]) - .unwrap(); + self.combined_writer.write(&buff.get_ref()[..pos])?; buff.set_position(0); } } let pos = buff.position() as usize; - self.combined_writer.write(&buff.get_ref()[..pos]).unwrap(); + self.combined_writer.write(&buff.get_ref()[..pos])?; + Ok(()) }); - let _ = hp1_handle.join().unwrap(); - let _ = hp2_handle.join().unwrap(); - let _ = combined_handle.join().unwrap(); - }); + let hp1_result = hp1_handle + .join() + .map_err(|_| anyhow!("hp1 writer thread panicked")) + .and_then(|result| result); + let hp2_result = hp2_handle + .join() + .map_err(|_| anyhow!("hp2 writer thread panicked")) + .and_then(|result| result); + let combined_result = combined_handle + .join() + .map_err(|_| anyhow!("combined writer thread panicked")) + .and_then(|result| result); + hp1_result?; + hp2_result?; + combined_result?; + Ok(()) + })?; let _ = self.return_mem.send(item); Ok(total_rows as u64) From d8e15ffcb3f4b05edbfa2062f3179ca7360efd73 Mon Sep 17 00:00:00 2001 From: SuhasSrinivasan <32346517+SuhasSrinivasan@users.noreply.github.com> Date: Wed, 5 Aug 2026 15:40:45 -0700 Subject: [PATCH 3/7] test: expose swallowed DMR segmenter errors --- modkit-core/src/dmr/single_site.rs | 131 +++++++++++++++++++++++++++-- 1 file changed, 122 insertions(+), 9 deletions(-) diff --git a/modkit-core/src/dmr/single_site.rs b/modkit-core/src/dmr/single_site.rs index 6c87fa8a..90ffd8a5 100644 --- a/modkit-core/src/dmr/single_site.rs +++ b/modkit-core/src/dmr/single_site.rs @@ -298,11 +298,11 @@ impl SingleSiteDmrAnalysis { break 'rcv_loop; } Ok(scores) => { - if let Err(e) = segmenter.add(&scores) { - self.multi_progress.suspend(|| { - error!("segmentation error, {e}"); - }) - } + add_scores_to_segmenter( + segmenter.as_mut(), + &scores, + &self.multi_progress, + )?; for (chrom, results) in scores { for result in results { match result { @@ -347,12 +347,9 @@ impl SingleSiteDmrAnalysis { } } - if let Err(e) = segmenter.run_current_chunk() { - self.multi_progress.suspend(|| error!("segmentation error, {e}")); - } success_counter.finish_and_clear(); failure_counter.finish_and_clear(); - segmenter.clean_up()?; + finish_segmenter(segmenter.as_mut(), &self.multi_progress)?; if let Some(e) = err { return Err(e.into()); @@ -985,6 +982,27 @@ trait DmrSegmenter { fn clean_up(&mut self) -> anyhow::Result<()>; } +fn add_scores_to_segmenter( + segmenter: &mut dyn DmrSegmenter, + scores: &[ChromToSingleScores], + multi_progress: &MultiProgress, +) -> anyhow::Result<()> { + if let Err(e) = segmenter.add(scores) { + multi_progress.suspend(|| error!("segmentation error, {e}")); + } + Ok(()) +} + +fn finish_segmenter( + segmenter: &mut dyn DmrSegmenter, + multi_progress: &MultiProgress, +) -> anyhow::Result<()> { + if let Err(e) = segmenter.run_current_chunk() { + multi_progress.suspend(|| error!("segmentation error, {e}")); + } + segmenter.clean_up() +} + #[derive(new)] struct DummySegmenter {} @@ -1347,3 +1365,98 @@ fn path_to_region_labels( agg } } + +#[cfg(test)] +mod segmenter_error_tests { + use super::{ + add_scores_to_segmenter, finish_segmenter, ChromToSingleScores, + DmrSegmenter, + }; + use anyhow::anyhow; + use indicatif::MultiProgress; + + const ADDED_ROW: &[u8] = b"chr1\t10\t11\tSAME\n"; + const FINAL_ROW: &[u8] = b"chr1\t20\t21\tDIFF\n"; + + #[derive(Default)] + struct StubSegmenter { + output: Vec, + add_error: Option<&'static str>, + final_error: Option<&'static str>, + cleaned_up: bool, + } + + impl DmrSegmenter for StubSegmenter { + fn add( + &mut self, + _dmr_scores: &[ChromToSingleScores], + ) -> anyhow::Result<()> { + if let Some(message) = self.add_error { + return Err(anyhow!(message)); + } + self.output.extend_from_slice(ADDED_ROW); + Ok(()) + } + + fn run_current_chunk(&mut self) -> anyhow::Result<()> { + if let Some(message) = self.final_error { + return Err(anyhow!(message)); + } + self.output.extend_from_slice(FINAL_ROW); + Ok(()) + } + + fn clean_up(&mut self) -> anyhow::Result<()> { + self.cleaned_up = true; + Ok(()) + } + } + + #[test] + fn segmenter_add_error_is_returned_exactly() { + let mut segmenter = StubSegmenter { + add_error: Some("stub segmenter add failed"), + ..StubSegmenter::default() + }; + + let error = + add_scores_to_segmenter(&mut segmenter, &[], &MultiProgress::new()) + .expect_err("add failure must be returned"); + + assert_eq!(error.to_string(), "stub segmenter add failed"); + assert!(segmenter.output.is_empty()); + assert!(!segmenter.cleaned_up); + } + + #[test] + fn segmenter_final_chunk_error_is_returned_exactly() { + let mut segmenter = StubSegmenter { + final_error: Some("stub segmenter final chunk failed"), + ..StubSegmenter::default() + }; + add_scores_to_segmenter(&mut segmenter, &[], &MultiProgress::new()) + .unwrap(); + + let error = finish_segmenter(&mut segmenter, &MultiProgress::new()) + .expect_err("final chunk failure must be returned"); + + assert_eq!(error.to_string(), "stub segmenter final chunk failed"); + assert_eq!(segmenter.output, ADDED_ROW); + assert!(!segmenter.cleaned_up); + } + + #[test] + fn successful_segmenter_bytes_and_lifecycle_are_unchanged() { + let mut segmenter = StubSegmenter::default(); + let progress = MultiProgress::new(); + + add_scores_to_segmenter(&mut segmenter, &[], &progress).unwrap(); + finish_segmenter(&mut segmenter, &progress).unwrap(); + + let mut expected = Vec::new(); + expected.extend_from_slice(ADDED_ROW); + expected.extend_from_slice(FINAL_ROW); + assert_eq!(segmenter.output, expected); + assert!(segmenter.cleaned_up); + } +} From c4ed9b8a127a133340ea6e79ab27958b8b5c38f7 Mon Sep 17 00:00:00 2001 From: SuhasSrinivasan <32346517+SuhasSrinivasan@users.noreply.github.com> Date: Wed, 5 Aug 2026 18:53:38 -0700 Subject: [PATCH 4/7] test: require fallible text output finalization --- modkit-core/src/writers.rs | 60 ++++++++++++++++++++++++++++++++++++-- 1 file changed, 58 insertions(+), 2 deletions(-) diff --git a/modkit-core/src/writers.rs b/modkit-core/src/writers.rs index 2474886d..0814dc86 100644 --- a/modkit-core/src/writers.rs +++ b/modkit-core/src/writers.rs @@ -1488,14 +1488,15 @@ where #[cfg(test)] mod writer_tests { use super::{ - BedMethylWriter, BedMethylWriter2, PileupWriter, RecordingWriter, - TsvWriter, + BedMethylWriter, BedMethylWriter2, OutWriter, PileupWriter, + RecordingWriter, TsvWriter, }; use crate::mod_base_code::ModCodeRepr; use crate::pileup::{ModBasePileup2, PileupFeatureCounts2}; use indicatif::ProgressBar; use std::io::{self, BufWriter, Write}; use std::panic::{catch_unwind, AssertUnwindSafe}; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; #[derive(Clone)] @@ -1503,6 +1504,7 @@ mod writer_tests { bytes: Arc>>, max_write: usize, fail_flush: bool, + flush_count: Arc, } impl ShortWriter { @@ -1512,6 +1514,7 @@ mod writer_tests { bytes: Arc::new(Mutex::new(Vec::new())), max_write, fail_flush: false, + flush_count: Arc::new(AtomicUsize::new(0)), } } @@ -1522,6 +1525,10 @@ mod writer_tests { fn bytes(&self) -> Vec { self.bytes.lock().unwrap().clone() } + + fn flush_count(&self) -> usize { + self.flush_count.load(Ordering::SeqCst) + } } impl Write for ShortWriter { @@ -1532,6 +1539,7 @@ mod writer_tests { } fn flush(&mut self) -> io::Result<()> { + self.flush_count.fetch_add(1, Ordering::SeqCst); if self.fail_flush { Err(io::Error::new(io::ErrorKind::Other, "flush failed")) } else { @@ -1612,4 +1620,52 @@ mod writer_tests { ); assert!(result.unwrap().is_err()); } + + #[test] + fn bedmethyl_writer_finish_propagates_buffered_flush_error() { + let sink = ShortWriter::failing_flush(usize::MAX); + let buf_writer = BufWriter::with_capacity(1024, sink); + let mut writer = + BedMethylWriter { buf_writer, tabs_and_spaces: false }; + PileupWriter::write(&mut writer, pileup(), &[]).unwrap(); + + let result = catch_unwind(AssertUnwindSafe(|| { + as PileupWriter< + ModBasePileup2, + >>::finish(&mut writer) + })); + assert!(result.is_ok(), "final flush panicked"); + let error = result.unwrap().expect_err("final flush should fail"); + assert!(error.to_string().contains("flush failed")); + } + + #[test] + fn out_writer_finish_propagates_tsv_flush_error() { + let sink = ShortWriter::failing_flush(usize::MAX); + let mut writer = TsvWriter { writer: sink }; + OutWriter::write(&mut writer, "row\n".to_string()).unwrap(); + + let result = catch_unwind(AssertUnwindSafe(|| { + as OutWriter>::finish(&mut writer) + })); + assert!(result.is_ok(), "final flush panicked"); + let error = result.unwrap().expect_err("final flush should fail"); + assert!(error.to_string().contains("flush failed")); + } + + #[test] + fn recording_writer_does_not_retry_a_surfaced_flush_error_on_drop() { + let sink = ShortWriter::failing_flush(usize::MAX); + let observed = sink.clone(); + let mut writer = RecordingWriter { + inner: sink, + pb: ProgressBar::hidden(), + }; + + let error = writer.write(b"row\n").expect_err("flush should fail"); + assert!(error.to_string().contains("flush failed")); + drop(writer); + assert_eq!(observed.flush_count(), 1); + } + } From 84af71f9cb2c9d9a9bbb9f2b77307cce50f702aa Mon Sep 17 00:00:00 2001 From: SuhasSrinivasan <32346517+SuhasSrinivasan@users.noreply.github.com> Date: Wed, 5 Aug 2026 19:20:33 -0700 Subject: [PATCH 5/7] fix: propagate text output finalization errors --- modkit-core/src/dmr/single_site.rs | 13 +- modkit-core/src/extract/subcommand.rs | 76 ++++-- modkit-core/src/extract/writer.rs | 11 + modkit-core/src/modbam_util/subcommands.rs | 32 ++- modkit-core/src/pileup/subcommand.rs | 53 +++- modkit-core/src/writers.rs | 287 ++++++++++++++++++--- 6 files changed, 396 insertions(+), 76 deletions(-) diff --git a/modkit-core/src/dmr/single_site.rs b/modkit-core/src/dmr/single_site.rs index 90ffd8a5..acef3fa6 100644 --- a/modkit-core/src/dmr/single_site.rs +++ b/modkit-core/src/dmr/single_site.rs @@ -987,9 +987,11 @@ fn add_scores_to_segmenter( scores: &[ChromToSingleScores], multi_progress: &MultiProgress, ) -> anyhow::Result<()> { - if let Err(e) = segmenter.add(scores) { + let result = segmenter.add(scores); + if let Err(e) = &result { multi_progress.suspend(|| error!("segmentation error, {e}")); } + result?; Ok(()) } @@ -997,10 +999,13 @@ fn finish_segmenter( segmenter: &mut dyn DmrSegmenter, multi_progress: &MultiProgress, ) -> anyhow::Result<()> { - if let Err(e) = segmenter.run_current_chunk() { + let result = segmenter.run_current_chunk(); + if let Err(e) = &result { multi_progress.suspend(|| error!("segmentation error, {e}")); } - segmenter.clean_up() + let clean_up_result = segmenter.clean_up(); + result?; + clean_up_result } #[derive(new)] @@ -1442,7 +1447,7 @@ mod segmenter_error_tests { assert_eq!(error.to_string(), "stub segmenter final chunk failed"); assert_eq!(segmenter.output, ADDED_ROW); - assert!(!segmenter.cleaned_up); + assert!(segmenter.cleaned_up); } #[test] diff --git a/modkit-core/src/extract/subcommand.rs b/modkit-core/src/extract/subcommand.rs index 42b4d6ed..b81379fe 100644 --- a/modkit-core/src/extract/subcommand.rs +++ b/modkit-core/src/extract/subcommand.rs @@ -42,7 +42,7 @@ use crate::record_processor::WithRecords; use crate::sample_probs::calc_per_base_thresholds_from_indexed_hts_file; use crate::threshold_mod_caller::MultipleThresholdModCaller; use crate::util::{format_errors_table, get_ticker, Region, KMER_SIZE}; -use crate::writers::TsvWriter; +use crate::writers::{finish_with_first_error, TsvWriter}; #[derive(Subcommand)] pub enum ExtractMods { @@ -285,7 +285,7 @@ impl EntryExtractFull { let mut writer: Box> = match self.input_args.out_path.as_str() { "stdout" | "-" => { - let tsv_writer = TsvWriter::new_stdout(output_header); + let tsv_writer = TsvWriter::new_stdout(output_header)?; let writer = TsvWriterWithContigNames::new( tsv_writer, tid_to_name, @@ -326,30 +326,47 @@ impl EntryExtractFull { } }; + let mut output_error = None; for result in rcv { match result { Ok(mod_profile) => { n_used.inc(mod_profile.num_reads() as u64); n_failed.inc(mod_profile.num_fails as u64); n_skipped.inc(mod_profile.num_skips as u64); - match writer - .write(mod_profile, motif_position_lookup.as_ref()) - { - Ok(n) => n_rows.inc(n), - Err(e) => { - error!("failed to write {}", e.to_string()); + if output_error.is_none() { + match writer + .write(mod_profile, motif_position_lookup.as_ref()) + { + Ok(n) => n_rows.inc(n), + Err(error) => { + output_error = + Some(error.context( + "failed to write extract output", + )); + } } } } - Err(e) => { + Err(error) => { debug!( "failed to calculate read-level mod probs, {}", - e.to_string() + error.to_string() ); + if output_error.is_none() { + output_error = Some(error.context( + "failed to calculate read-level mod probabilities", + )); + } } } } + let output_result = finish_with_first_error( + output_error, + || writer.finish(), + "failed to flush extract output", + ); + n_failed.finish_and_clear(); n_skipped.finish_and_clear(); n_used.finish_and_clear(); @@ -361,7 +378,7 @@ impl EntryExtractFull { n_skipped.position(), n_failed.position() ); - Ok(()) + output_result } } @@ -725,7 +742,7 @@ impl EntryExtractCalls { let mut writer: Box> = match self.input_args.out_path.as_str() { "stdout" | "-" => { - let tsv_writer = TsvWriter::new_stdout(output_header); + let tsv_writer = TsvWriter::new_stdout(output_header)?; let writer = TsvWriterWithContigNames::new_with_caller( tsv_writer, tid_to_name, @@ -836,30 +853,47 @@ impl EntryExtractCalls { ); }); + let mut output_error = None; for result in rcv { match result { Ok(mod_profile) => { n_used.inc(mod_profile.num_reads() as u64); n_failed.inc(mod_profile.num_fails as u64); n_skipped.inc(mod_profile.num_skips as u64); - match writer - .write(mod_profile, motif_position_lookup.as_ref()) - { - Ok(n) => n_rows.inc(n), - Err(e) => { - error!("failed to write {}", e.to_string()); + if output_error.is_none() { + match writer + .write(mod_profile, motif_position_lookup.as_ref()) + { + Ok(n) => n_rows.inc(n), + Err(error) => { + output_error = + Some(error.context( + "failed to write extract output", + )); + } } } } - Err(e) => { + Err(error) => { debug!( "failed to calculate read-level mod probs, {}", - e.to_string() + error.to_string() ); + if output_error.is_none() { + output_error = Some(error.context( + "failed to calculate read-level mod probabilities", + )); + } } } } + let output_result = finish_with_first_error( + output_error, + || writer.finish(), + "failed to flush extract output", + ); + n_failed.finish_and_clear(); n_skipped.finish_and_clear(); n_used.finish_and_clear(); @@ -871,7 +905,7 @@ impl EntryExtractCalls { n_skipped.position(), n_failed.position() ); - Ok(()) + output_result } } diff --git a/modkit-core/src/extract/writer.rs b/modkit-core/src/extract/writer.rs index 237827ba..a6bcd879 100644 --- a/modkit-core/src/extract/writer.rs +++ b/modkit-core/src/extract/writer.rs @@ -182,6 +182,7 @@ pub(crate) trait OutwriterWithMemory { motif_position_lookup: Option<&MotifPositionLookup>, ) -> anyhow::Result; fn num_reads(&self) -> usize; + fn finish(&mut self) -> anyhow::Result<()>; } pub struct TsvWriterWithContigNames { @@ -251,6 +252,11 @@ impl OutwriterWithMemory fn num_reads(&self) -> usize { self.number_of_written_reads } + + fn finish(&mut self) -> anyhow::Result<()> { + self.tsv_writer.flush()?; + Ok(()) + } } impl TsvWriterWithContigNames { @@ -311,6 +317,11 @@ impl OutwriterWithMemory fn num_reads(&self) -> usize { self.number_of_written_reads } + + fn finish(&mut self) -> anyhow::Result<()> { + self.tsv_writer.flush()?; + Ok(()) + } } pub(super) trait CanWriteReadModStatsRecords { diff --git a/modkit-core/src/modbam_util/subcommands.rs b/modkit-core/src/modbam_util/subcommands.rs index 509b270b..67c6f4e7 100644 --- a/modkit-core/src/modbam_util/subcommands.rs +++ b/modkit-core/src/modbam_util/subcommands.rs @@ -57,7 +57,8 @@ use crate::util::{ get_ticker, ReferenceRecord, Region, DEFAULT_NUM_READS, }; use crate::writers::{ - MultiTableWriter, OutWriter, SampledProbs, TableWriter, TsvWriter, + finish_with_first_error, MultiTableWriter, OutWriter, SampledProbs, + TableWriter, TsvWriter, }; #[derive(Subcommand)] @@ -1555,12 +1556,19 @@ impl SampleModBaseProbs { sampled_probs.check_path(p, self.force)?; Box::new(MultiTableWriter::new(p.clone())) } else { - Box::new(TsvWriter::new_stdout(None)) + Box::new(TsvWriter::new_stdout(None)?) }; - writer.write(sampled_probs)?; - - Ok(()) + let output_error = writer + .write(sampled_probs) + .map(|_| ()) + .context("failed to write sampled probabilities output") + .err(); + finish_with_first_error( + output_error, + || writer.finish(), + "failed to flush sampled probabilities output", + ) } } @@ -2015,12 +2023,20 @@ impl ModSummarize { )?; let mut writer: Box> = if self.tsv_format { - Box::new(TsvWriter::new_stdout(None)) + Box::new(TsvWriter::new_stdout(None)?) } else { Box::new(TableWriter::new()) }; - writer.write(mod_summary)?; - Ok(()) + let output_error = writer + .write(mod_summary) + .map(|_| ()) + .context("failed to write summary output") + .err(); + finish_with_first_error( + output_error, + || writer.finish(), + "failed to flush summary output", + ) } } diff --git a/modkit-core/src/pileup/subcommand.rs b/modkit-core/src/pileup/subcommand.rs index 4503f99e..1e0a4934 100644 --- a/modkit-core/src/pileup/subcommand.rs +++ b/modkit-core/src/pileup/subcommand.rs @@ -53,8 +53,8 @@ use crate::util::{ get_ticker, reader_is_bam, reader_is_cram, Region, }; use crate::writers::{ - BedMethylWriter, BedMethylWriter2, MultipleMotifBedmethylWriter, - PhasedBedMethylWriter, PileupWriter, + finish_with_first_error, BedMethylWriter, BedMethylWriter2, + MultipleMotifBedmethylWriter, PhasedBedMethylWriter, PileupWriter, }; #[derive(Args)] @@ -1623,20 +1623,37 @@ impl ModBamPileup { drop(records_tx); }); - for result in records_rx.into_iter() { + let mut output_error = None; + for result in records_rx.iter() { match result { Ok(mod_base_pileup) => { tid_progress.inc(mod_base_pileup.interval_width as u64); erred_reads.inc(mod_base_pileup.failed_records as u64); - let rows_written = - writer.write(mod_base_pileup, &motif_labels)?; - write_progress.inc(rows_written); + match writer.write(mod_base_pileup, &motif_labels) { + Ok(rows_written) => write_progress.inc(rows_written), + Err(error) => { + output_error = Some( + error.context("failed to write pileup output"), + ); + break; + } + } } Err(message) => { debug!("unexpected error {message}"); } } } + // If writing failed, disconnect the collector before joining the + // pipeline so workers do not remain blocked sending further records. + drop(records_rx); + let output_result = finish_with_first_error( + output_error, + || writer.finish(), + "failed to flush pileup output", + ); + drop(writer); + drop(empties_tx); let rows_processed = write_progress.position(); let n_failed_reads = erred_reads.position(); @@ -1660,7 +1677,7 @@ impl ModBamPileup { } aggregator.join().expect("aggregator theread paniced"); - Ok(()) + output_result } } @@ -2343,20 +2360,36 @@ impl DuplexModBamPileup { tid_progress.finish_and_clear(); }); + let mut output_error = None; for result in rx.into_iter() { match result { Ok(mod_base_pileup) => { processed_reads .inc(mod_base_pileup.processed_records as u64); skipped_reads.inc(mod_base_pileup.skipped_records as u64); - let rows_written = writer.write(mod_base_pileup, &[])?; - write_progress.inc(rows_written); + if output_error.is_none() { + match writer.write(mod_base_pileup, &[]) { + Ok(rows_written) => { + write_progress.inc(rows_written) + } + Err(error) => { + output_error = Some(error.context( + "failed to write duplex pileup output", + )); + } + } + } } Err(message) => { debug!("> unexpected error {message}"); } } } + let output_result = finish_with_first_error( + output_error, + || writer.finish(), + "failed to flush duplex pileup output", + ); let rows_processed = write_progress.position(); let n_skipped_reads = skipped_reads.position(); let n_skipped_message = if n_skipped_reads == 0 { @@ -2372,6 +2405,6 @@ impl DuplexModBamPileup { "Done, processed {rows_processed} rows. Processed \ ~{n_processed_reads} reads and skipped {n_skipped_message}." ); - Ok(()) + output_result } } diff --git a/modkit-core/src/writers.rs b/modkit-core/src/writers.rs index 0814dc86..7baef34e 100644 --- a/modkit-core/src/writers.rs +++ b/modkit-core/src/writers.rs @@ -45,10 +45,32 @@ pub trait PileupWriter { item: T, motif_labels: &[String], ) -> anyhow::Result; + + /// Flush all command-visible output before reporting success. + fn finish(&mut self) -> anyhow::Result<()>; } pub trait OutWriter { fn write(&mut self, item: T) -> AnyhowResult; + fn finish(&mut self) -> AnyhowResult<()>; +} + +/// Always attempt output finalization, but do not let a later flush error +/// replace an error that was already surfaced while producing or writing +/// output. +pub(crate) fn finish_with_first_error( + first_error: Option, + finish: F, + context: &'static str, +) -> AnyhowResult<()> +where + F: FnOnce() -> AnyhowResult<()>, +{ + let finish_result = finish().context(context); + match first_error { + Some(error) => Err(error), + None => finish_result, + } } pub struct BedMethylWriter { @@ -209,6 +231,10 @@ impl PileupWriter for BedMethylWriter2 { let _ = self.return_mem.send(item); Ok(n_rows) } + + fn finish(&mut self) -> anyhow::Result<()> { + self.inner.flush() + } } pub fn bedmethyl_header() -> String { @@ -318,6 +344,11 @@ impl PileupWriter for BedMethylWriter { std::thread::spawn(|| drop(item)); Ok(rows_written) } + + fn finish(&mut self) -> anyhow::Result<()> { + self.buf_writer.flush()?; + Ok(()) + } } impl PileupWriter for BedMethylWriter { @@ -390,6 +421,11 @@ impl PileupWriter for BedMethylWriter { } Ok(rows_written) } + + fn finish(&mut self) -> anyhow::Result<()> { + self.buf_writer.flush()?; + Ok(()) + } } pub struct MultipleMotifBedmethylWriter { @@ -560,6 +596,10 @@ impl PileupWriter let _ = self.return_mem.send(item); Ok(rows_written) } + + fn finish(&mut self) -> anyhow::Result<()> { + self.writer.flush() + } } pub struct TableWriter { @@ -711,6 +751,11 @@ impl<'a, W: Write> OutWriter> for TableWriter { report_emitted += emitted; Ok(report_emitted as u64) } + + fn finish(&mut self) -> AnyhowResult<()> { + self.writer.flush()?; + Ok(()) + } } pub struct TsvWriter { @@ -736,13 +781,13 @@ impl TsvWriter> { } impl TsvWriter> { - pub fn new_stdout(header: Option) -> Self { - let out = BufWriter::new(std::io::stdout()); + pub fn new_stdout(header: Option) -> anyhow::Result { + let mut out = BufWriter::new(std::io::stdout()); if let Some(header) = header { - println!("{header}"); + out.write_all(format!("{header}\n").as_bytes())?; } - Self { writer: out } + Ok(Self { writer: out }) } } @@ -806,6 +851,11 @@ impl OutWriter for TsvWriter { self.writer.write_all(item.as_bytes())?; Ok(item.len() as u64) } + + fn finish(&mut self) -> AnyhowResult<()> { + self.flush()?; + Ok(()) + } } impl<'a, W: Write> OutWriter> for TsvWriter { @@ -883,6 +933,11 @@ impl<'a, W: Write> OutWriter> for TsvWriter { self.writer.write_all(report.as_bytes())?; Ok(1) } + + fn finish(&mut self) -> AnyhowResult<()> { + self.flush()?; + Ok(()) + } } #[derive(new)] @@ -1144,6 +1199,84 @@ impl ProbHistogram { } } +fn remember_first_error( + first_error: &mut Option, + result: AnyhowResult<()>, +) { + if let Err(error) = result { + if first_error.is_none() { + *first_error = Some(error); + } + } +} + +fn write_probability_artifacts( + table: &Table, + probabilities_writer: P, + counts_html: Option<&str>, + counts_writer: C, + proportions_html: Option<&str>, + proportions_writer: R, +) -> AnyhowResult<()> +where + P: Write, + C: Write, + R: Write, +{ + let mut first_error = None; + + let csv_writer = csv::WriterBuilder::new() + .has_headers(true) + .delimiter(b'\t') + .from_writer(probabilities_writer); + let probabilities_result = table + .to_csv_writer(csv_writer) + .map_err(anyhow::Error::from) + .and_then(|writer| { + writer + .into_inner() + .map(|_| ()) + .map_err(|error| anyhow::Error::new(error.into_error())) + }) + .context("failed to finalize probabilities table"); + remember_first_error(&mut first_error, probabilities_result); + + let mut counts_writer = BufWriter::new(counts_writer); + if let Some(blob) = counts_html { + remember_first_error( + &mut first_error, + counts_writer + .write_all(blob.as_bytes()) + .context("failed to write counts plot"), + ); + } + remember_first_error( + &mut first_error, + counts_writer.flush().context("failed to finalize counts plot"), + ); + + let mut proportions_writer = BufWriter::new(proportions_writer); + if let Some(blob) = proportions_html { + remember_first_error( + &mut first_error, + proportions_writer + .write_all(blob.as_bytes()) + .context("failed to write proportions plot"), + ); + } + remember_first_error( + &mut first_error, + proportions_writer + .flush() + .context("failed to finalize proportions plot"), + ); + + match first_error { + Some(error) => Err(error), + None => Ok(()), + } +} + impl OutWriter for MultiTableWriter { fn write(&mut self, item: SampledProbs) -> AnyhowResult { let mut rows_written = 0u64; @@ -1151,7 +1284,19 @@ impl OutWriter for MultiTableWriter { let threshold_fn = self.out_dir.join(item.get_thresholds_filename()); let mut fh = File::create(threshold_fn)?; - let n_written = thresh_table.print(&mut fh)?; + let threshold_write_result = thresh_table.print(&mut fh); + let threshold_finish_result = fh.flush(); + let n_written = match threshold_write_result { + Ok(n_written) => { + threshold_finish_result + .context("failed to finalize thresholds table")?; + n_written + } + Err(error) => { + let _ = threshold_finish_result; + return Err(error.into()); + } + }; rows_written += n_written as u64; if let Some(histograms) = &item.histograms { @@ -1159,35 +1304,49 @@ impl OutWriter for MultiTableWriter { SampledProbs::get_probabilities_filenames(item.prefix.as_ref()); let probs_table_fh = File::create(self.out_dir.join(probs_table_fn))?; - let mut counts_plot_fh = BufWriter::new(File::create( - self.out_dir.join(counts_plot_fn), - )?); - let mut prop_plot_fh = - BufWriter::new(File::create(self.out_dir.join(prop_plot_fn))?); - - let csv_writer = csv::WriterBuilder::new() - .has_headers(true) - .delimiter('\t' as u8) - .from_writer(probs_table_fh); + let counts_plot_fh = + File::create(self.out_dir.join(counts_plot_fn))?; + let prop_plot_fh = File::create(self.out_dir.join(prop_plot_fn))?; let (tab, counts_chart, prop_chart) = histograms.get_artifacts( &item.primary_base_colors, &item.mod_base_colors, ); - tab.to_csv_writer(csv_writer)?; - match HtmlRenderer::new("Counts", 800, 800).render(&counts_chart) { - Ok(blob) => counts_plot_fh.write_all(blob.as_bytes())?, - Err(e) => debug!("failed to render counts plot, {e:?}"), - } - match HtmlRenderer::new("Proportions", 800, 800).render(&prop_chart) + let counts_html = match HtmlRenderer::new("Counts", 800, 800) + .render(&counts_chart) { - Ok(blob) => prop_plot_fh.write_all(blob.as_bytes())?, - Err(e) => debug!("failed to render proportions plot, {e:?}"), - } + Ok(blob) => Some(blob), + Err(e) => { + debug!("failed to render counts plot, {e:?}"); + None + } + }; + let proportions_html = + match HtmlRenderer::new("Proportions", 800, 800) + .render(&prop_chart) + { + Ok(blob) => Some(blob), + Err(e) => { + debug!("failed to render proportions plot, {e:?}"); + None + } + }; + write_probability_artifacts( + &tab, + probs_table_fh, + counts_html.as_deref(), + counts_plot_fh, + proportions_html.as_deref(), + prop_plot_fh, + )?; } Ok(rows_written) } + + fn finish(&mut self) -> AnyhowResult<()> { + Ok(()) + } } impl OutWriter for TsvWriter> { @@ -1198,6 +1357,11 @@ impl OutWriter for TsvWriter> { rows_written += n_written as u64; Ok(rows_written) } + + fn finish(&mut self) -> AnyhowResult<()> { + self.flush()?; + Ok(()) + } } #[inline] @@ -1283,7 +1447,6 @@ impl RecordingWriter { impl Drop for RecordingWriter { fn drop(&mut self) { self.pb.finish_and_clear(); - let _ = self.inner.flush(); } } @@ -1483,17 +1646,29 @@ where Ok(total_rows as u64) } + + fn finish(&mut self) -> anyhow::Result<()> { + // Attempt every output even if an earlier flush failed so no phased + // writer is left with command-visible bytes still buffered. + let hp1_result = self.hp1_writer.flush(); + let hp2_result = self.hp2_writer.flush(); + let combined_result = self.combined_writer.flush(); + hp1_result?; + hp2_result?; + combined_result + } } #[cfg(test)] mod writer_tests { use super::{ - BedMethylWriter, BedMethylWriter2, OutWriter, PileupWriter, - RecordingWriter, TsvWriter, + finish_with_first_error, write_probability_artifacts, BedMethylWriter, + BedMethylWriter2, OutWriter, PileupWriter, RecordingWriter, TsvWriter, }; use crate::mod_base_code::ModCodeRepr; use crate::pileup::{ModBasePileup2, PileupFeatureCounts2}; use indicatif::ProgressBar; + use prettytable::Table; use std::io::{self, BufWriter, Write}; use std::panic::{catch_unwind, AssertUnwindSafe}; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -1625,8 +1800,7 @@ mod writer_tests { fn bedmethyl_writer_finish_propagates_buffered_flush_error() { let sink = ShortWriter::failing_flush(usize::MAX); let buf_writer = BufWriter::with_capacity(1024, sink); - let mut writer = - BedMethylWriter { buf_writer, tabs_and_spaces: false }; + let mut writer = BedMethylWriter { buf_writer, tabs_and_spaces: false }; PileupWriter::write(&mut writer, pileup(), &[]).unwrap(); let result = catch_unwind(AssertUnwindSafe(|| { @@ -1657,10 +1831,8 @@ mod writer_tests { fn recording_writer_does_not_retry_a_surfaced_flush_error_on_drop() { let sink = ShortWriter::failing_flush(usize::MAX); let observed = sink.clone(); - let mut writer = RecordingWriter { - inner: sink, - pb: ProgressBar::hidden(), - }; + let mut writer = + RecordingWriter { inner: sink, pb: ProgressBar::hidden() }; let error = writer.write(b"row\n").expect_err("flush should fail"); assert!(error.to_string().contains("flush failed")); @@ -1668,4 +1840,53 @@ mod writer_tests { assert_eq!(observed.flush_count(), 1); } + #[test] + fn finish_is_attempted_without_replacing_an_earlier_output_error() { + let mut finish_attempts = 0; + let result = finish_with_first_error( + Some(anyhow::anyhow!("write failed first")), + || { + finish_attempts += 1; + anyhow::bail!("flush failed later") + }, + "failed to finalize output", + ); + + assert_eq!(finish_attempts, 1); + assert_eq!(result.unwrap_err().to_string(), "write failed first"); + } + + #[test] + fn finish_error_is_returned_when_there_is_no_earlier_error() { + let result = finish_with_first_error( + None, + || anyhow::bail!("flush failed"), + "failed to finalize output", + ); + + let error = result.unwrap_err(); + assert_eq!(error.to_string(), "failed to finalize output"); + assert!(format!("{error:#}").contains("flush failed")); + } + + #[test] + fn probability_artifacts_all_finalize_and_keep_the_first_error() { + let probabilities = ShortWriter::failing_flush(usize::MAX); + let counts = ShortWriter::failing_flush(usize::MAX); + let proportions = ShortWriter::failing_flush(usize::MAX); + let error = write_probability_artifacts( + &Table::new(), + probabilities.clone(), + Some("counts"), + counts.clone(), + Some("proportions"), + proportions.clone(), + ) + .expect_err("artifact finalization should fail"); + + assert_eq!(error.to_string(), "failed to finalize probabilities table"); + assert!(probabilities.flush_count() >= 1); + assert_eq!(counts.flush_count(), 1); + assert_eq!(proportions.flush_count(), 1); + } } From 8bb45f9a1f321f58dcaf2e0619bb70039ad78992 Mon Sep 17 00:00:00 2001 From: SuhasSrinivasan <32346517+SuhasSrinivasan@users.noreply.github.com> Date: Wed, 5 Aug 2026 19:58:55 -0700 Subject: [PATCH 6/7] test: expose DMR finalization ordering --- modkit-core/src/dmr/pairwise.rs | 75 +++++++++++++- modkit-core/src/dmr/single_site.rs | 155 +++++++++++++++++++++++++++-- modkit-core/src/dmr/subcommands.rs | 138 ++++++++++++++++++++++--- 3 files changed, 343 insertions(+), 25 deletions(-) diff --git a/modkit-core/src/dmr/pairwise.rs b/modkit-core/src/dmr/pairwise.rs index bcc6b1c5..ccae279b 100644 --- a/modkit-core/src/dmr/pairwise.rs +++ b/modkit-core/src/dmr/pairwise.rs @@ -303,10 +303,75 @@ pub(super) fn run_pairwise_dmr( pb.finish_and_clear(); - if let Some(e) = err { - Err(e.into()) - } else { - writer.flush()?; - Ok((success_count, region_error_counts)) + finish_pairwise_output(err.map(anyhow::Error::from), writer.as_mut())?; + Ok((success_count, region_error_counts)) +} + +fn finish_pairwise_output( + first_error: Option, + writer: &mut dyn std::io::Write, +) -> anyhow::Result<()> { + if let Some(error) = first_error { + return Err(error); + } + writer.flush()?; + Ok(()) +} + +#[cfg(test)] +mod output_finalization_tests { + use super::finish_pairwise_output; + use anyhow::anyhow; + use std::io::{self, Write}; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + struct FlushWriter { + flushes: Arc, + fail_flush: bool, + } + + impl Write for FlushWriter { + fn write(&mut self, buf: &[u8]) -> io::Result { + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + self.flushes.fetch_add(1, Ordering::SeqCst); + if self.fail_flush { + Err(io::Error::other("pairwise flush failed later")) + } else { + Ok(()) + } + } + } + + #[test] + fn earlier_pairwise_error_is_retained_and_flush_is_attempted() { + let flushes = Arc::new(AtomicUsize::new(0)); + let mut writer = + FlushWriter { flushes: flushes.clone(), fail_flush: true }; + + let error = finish_pairwise_output( + Some(anyhow!("pairwise processing failed first")), + &mut writer, + ) + .expect_err("the first error must be returned"); + + assert_eq!(error.to_string(), "pairwise processing failed first"); + assert_eq!(flushes.load(Ordering::SeqCst), 1); + } + + #[test] + fn pairwise_flush_error_is_returned_when_it_is_first() { + let flushes = Arc::new(AtomicUsize::new(0)); + let mut writer = + FlushWriter { flushes: flushes.clone(), fail_flush: true }; + + let error = finish_pairwise_output(None, &mut writer) + .expect_err("flush failure must be returned"); + + assert_eq!(error.to_string(), "pairwise flush failed later"); + assert_eq!(flushes.load(Ordering::SeqCst), 1); } } diff --git a/modkit-core/src/dmr/single_site.rs b/modkit-core/src/dmr/single_site.rs index acef3fa6..1e707abe 100644 --- a/modkit-core/src/dmr/single_site.rs +++ b/modkit-core/src/dmr/single_site.rs @@ -349,11 +349,13 @@ impl SingleSiteDmrAnalysis { success_counter.finish_and_clear(); failure_counter.finish_and_clear(); - finish_segmenter(segmenter.as_mut(), &self.multi_progress)?; - - if let Some(e) = err { - return Err(e.into()); - } + finish_single_site_outputs( + err.map(anyhow::Error::from), + true, + segmenter.as_mut(), + writer.as_mut(), + &self.multi_progress, + )?; if !error_counts.is_empty() { self.multi_progress.suspend(|| { @@ -367,7 +369,6 @@ impl SingleSiteDmrAnalysis { success_count, failure_counter.position(), ); - writer.flush()?; Ok(()) } } @@ -1008,6 +1009,25 @@ fn finish_segmenter( clean_up_result } +fn finish_single_site_outputs( + first_error: Option, + run_final_chunk: bool, + segmenter: &mut dyn DmrSegmenter, + writer: &mut dyn Write, + multi_progress: &MultiProgress, +) -> anyhow::Result<()> { + if run_final_chunk { + finish_segmenter(segmenter, multi_progress)?; + } else { + segmenter.clean_up()?; + } + if let Some(error) = first_error { + return Err(error); + } + writer.flush()?; + Ok(()) +} + #[derive(new)] struct DummySegmenter {} @@ -1374,11 +1394,13 @@ fn path_to_region_labels( #[cfg(test)] mod segmenter_error_tests { use super::{ - add_scores_to_segmenter, finish_segmenter, ChromToSingleScores, - DmrSegmenter, + add_scores_to_segmenter, finish_segmenter, finish_single_site_outputs, + ChromToSingleScores, DmrSegmenter, }; use anyhow::anyhow; use indicatif::MultiProgress; + use std::io::{self, Write}; + use std::sync::{Arc, Mutex}; const ADDED_ROW: &[u8] = b"chr1\t10\t11\tSAME\n"; const FINAL_ROW: &[u8] = b"chr1\t20\t21\tDIFF\n"; @@ -1388,7 +1410,9 @@ mod segmenter_error_tests { output: Vec, add_error: Option<&'static str>, final_error: Option<&'static str>, + cleanup_error: Option<&'static str>, cleaned_up: bool, + events: Arc>>, } impl DmrSegmenter for StubSegmenter { @@ -1404,6 +1428,7 @@ mod segmenter_error_tests { } fn run_current_chunk(&mut self) -> anyhow::Result<()> { + self.events.lock().unwrap().push("final chunk"); if let Some(message) = self.final_error { return Err(anyhow!(message)); } @@ -1412,11 +1437,60 @@ mod segmenter_error_tests { } fn clean_up(&mut self) -> anyhow::Result<()> { + self.events.lock().unwrap().push("segmenter cleanup"); self.cleaned_up = true; + if let Some(message) = self.cleanup_error { + return Err(anyhow!(message)); + } Ok(()) } } + struct TrackingWriter { + events: Arc>>, + flush_error: Option<&'static str>, + } + + impl Write for TrackingWriter { + fn write(&mut self, buf: &[u8]) -> io::Result { + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + self.events.lock().unwrap().push("main writer flush"); + match self.flush_error { + Some(message) => Err(io::Error::other(message)), + None => Ok(()), + } + } + } + + fn lifecycle( + first_error: Option, + run_final_chunk: bool, + final_error: Option<&'static str>, + cleanup_error: Option<&'static str>, + flush_error: Option<&'static str>, + ) -> (anyhow::Result<()>, Vec<&'static str>) { + let events = Arc::new(Mutex::new(Vec::new())); + let mut segmenter = StubSegmenter { + final_error, + cleanup_error, + events: events.clone(), + ..StubSegmenter::default() + }; + let mut writer = TrackingWriter { events: events.clone(), flush_error }; + let result = finish_single_site_outputs( + first_error, + run_final_chunk, + &mut segmenter, + &mut writer, + &MultiProgress::new(), + ); + let observed = events.lock().unwrap().clone(); + (result, observed) + } + #[test] fn segmenter_add_error_is_returned_exactly() { let mut segmenter = StubSegmenter { @@ -1464,4 +1538,69 @@ mod segmenter_error_tests { assert_eq!(segmenter.output, expected); assert!(segmenter.cleaned_up); } + + #[test] + fn add_error_skips_chunk_retry_but_attempts_cleanup_and_main_flush() { + let (result, events) = lifecycle( + Some(anyhow!("segmenter add failed")), + false, + Some("must not retry final chunk"), + None, + Some("main flush failed"), + ); + + assert_eq!(result.unwrap_err().to_string(), "segmenter add failed"); + assert_eq!(events, ["segmenter cleanup", "main writer flush"]); + } + + #[test] + fn batch_error_runs_pending_chunk_and_retains_first_error() { + let (result, events) = lifecycle( + Some(anyhow!("batch failed first")), + true, + Some("final chunk failed later"), + Some("cleanup failed later"), + Some("main flush failed later"), + ); + + assert_eq!(result.unwrap_err().to_string(), "batch failed first"); + assert_eq!( + events, + ["final chunk", "segmenter cleanup", "main writer flush"] + ); + } + + #[test] + fn main_write_error_still_attempts_all_finalization() { + let (result, events) = lifecycle( + Some(anyhow!("main write failed first")), + true, + None, + Some("cleanup failed later"), + Some("main flush failed later"), + ); + + assert_eq!(result.unwrap_err().to_string(), "main write failed first"); + assert_eq!( + events, + ["final chunk", "segmenter cleanup", "main writer flush"] + ); + } + + #[test] + fn final_chunk_error_wins_but_cleanup_and_main_flush_are_attempted() { + let (result, events) = lifecycle( + None, + true, + Some("final chunk failed first"), + Some("cleanup failed later"), + Some("main flush failed later"), + ); + + assert_eq!(result.unwrap_err().to_string(), "final chunk failed first"); + assert_eq!( + events, + ["final chunk", "segmenter cleanup", "main writer flush"] + ); + } } diff --git a/modkit-core/src/dmr/subcommands.rs b/modkit-core/src/dmr/subcommands.rs index 306f4544..f4de988e 100644 --- a/modkit-core/src/dmr/subcommands.rs +++ b/modkit-core/src/dmr/subcommands.rs @@ -38,6 +38,22 @@ use crate::util::{ }; use modkit_logging::init_logging; +fn finish_threaded_dmr_output( + first_error: Option, + join_pipeline: F, + writer: &mut dyn Write, +) -> anyhow::Result<()> +where + F: FnOnce() -> anyhow::Result<()>, +{ + if let Some(error) = first_error { + return Err(error); + } + join_pipeline()?; + writer.flush()?; + Ok(()) +} + #[derive(Subcommand)] pub enum BedMethylDmr { /// Compare regions in a pair of samples (for example, tumor and normal or @@ -1423,17 +1439,23 @@ impl EntryDmrIsoform { } } } - source_thread.join().expect("source thread paniced"); - for (i, worker_thread) in handles.into_iter().enumerate() { - worker_thread.join().expect(&format!("worker {i} paniced")); - } - aggregator.join().expect("aggregator theread paniced"); + finish_threaded_dmr_output( + None, + || { + source_thread.join().expect("source thread paniced"); + for (i, worker_thread) in handles.into_iter().enumerate() { + worker_thread.join().expect(&format!("worker {i} paniced")); + } + aggregator.join().expect("aggregator theread paniced"); + Ok(()) + }, + writer.as_mut(), + )?; multi_progress.suspend(|| { info!("finished, processed {} genes", pb.position()); }); - writer.flush()?; Ok(()) } @@ -1870,11 +1892,18 @@ impl EntryGeneTx { std::fs::write(fp, svg)?; } - source_thread.join().expect("source thread paniced"); - for (i, worker_thread) in handles.into_iter().enumerate() { - worker_thread.join().expect(&format!("worker {i} paniced")); - } - aggregator.join().expect("aggregator theread paniced"); + finish_threaded_dmr_output( + None, + || { + source_thread.join().expect("source thread paniced"); + for (i, worker_thread) in handles.into_iter().enumerate() { + worker_thread.join().expect(&format!("worker {i} paniced")); + } + aggregator.join().expect("aggregator theread paniced"); + Ok(()) + }, + writer.as_mut(), + )?; multi_progress.suspend(|| { info!("finished, {} errors", errs.len()); @@ -1885,7 +1914,6 @@ impl EntryGeneTx { multi_progress.suspend(|| info!("{err_table}")); } - writer.flush()?; Ok(()) } @@ -1941,3 +1969,89 @@ impl EntryGeneTx { } } } + +#[cfg(test)] +mod output_finalization_tests { + use super::finish_threaded_dmr_output; + use anyhow::anyhow; + use std::io::{self, Write}; + use std::sync::{Arc, Mutex}; + + struct EventWriter { + events: Arc>>, + fail_flush: bool, + } + + impl Write for EventWriter { + fn write(&mut self, buf: &[u8]) -> io::Result { + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + self.events.lock().unwrap().push("writer flush"); + if self.fail_flush { + Err(io::Error::other("writer flush failed later")) + } else { + Ok(()) + } + } + } + + fn lifecycle( + first_error: Option, + join_error: Option<&'static str>, + fail_flush: bool, + ) -> (anyhow::Result<()>, Vec<&'static str>) { + let events = Arc::new(Mutex::new(Vec::new())); + let join_events = events.clone(); + let mut writer = EventWriter { events: events.clone(), fail_flush }; + let result = finish_threaded_dmr_output( + first_error, + move || { + join_events.lock().unwrap().push("pipeline joins"); + match join_error { + Some(message) => Err(anyhow!(message)), + None => Ok(()), + } + }, + &mut writer, + ); + let observed = events.lock().unwrap().clone(); + (result, observed) + } + + #[test] + fn all_gene_write_error_still_joins_and_flushes_preserving_first_error() { + let (result, events) = lifecycle( + Some(anyhow!("all-gene write failed first")), + Some("pipeline join failed later"), + true, + ); + + assert_eq!( + result.unwrap_err().to_string(), + "all-gene write failed first" + ); + assert_eq!(events, ["pipeline joins", "writer flush"]); + } + + #[test] + fn gene_transcript_join_error_still_flushes_and_is_retained() { + let (result, events) = + lifecycle(None, Some("pipeline join failed first"), true); + + assert_eq!( + result.unwrap_err().to_string(), + "pipeline join failed first" + ); + assert_eq!(events, ["pipeline joins", "writer flush"]); + } + + #[test] + fn successful_threaded_dmr_finalization_order_is_join_then_flush() { + let (result, events) = lifecycle(None, None, false); + + result.unwrap(); + assert_eq!(events, ["pipeline joins", "writer flush"]); + } +} From 096b04636d77e0d8cf8dd74021f953df2819004b Mon Sep 17 00:00:00 2001 From: SuhasSrinivasan <32346517+SuhasSrinivasan@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:13:02 -0700 Subject: [PATCH 7/7] fix: finalize DMR outputs after failures --- modkit-core/src/dmr/pairwise.rs | 36 ++-- modkit-core/src/dmr/single_site.rs | 101 +++++++--- modkit-core/src/dmr/subcommands.rs | 283 +++++++++++++++++++++-------- 3 files changed, 300 insertions(+), 120 deletions(-) diff --git a/modkit-core/src/dmr/pairwise.rs b/modkit-core/src/dmr/pairwise.rs index ccae279b..eb81fbc8 100644 --- a/modkit-core/src/dmr/pairwise.rs +++ b/modkit-core/src/dmr/pairwise.rs @@ -183,8 +183,11 @@ pub(super) fn run_pairwise_dmr( multi_progress: MultiProgress, ) -> anyhow::Result<(usize, FxHashMap)> { if header { - writer - .write_all(ModificationCounts::header(a_name, b_name).as_bytes())?; + if let Err(error) = writer + .write_all(ModificationCounts::header(a_name, b_name).as_bytes()) + { + finish_pairwise_output(Some(error.into()), writer.as_mut())?; + } } let (snd, rcv) = crossbeam_channel::bounded(1000); @@ -248,14 +251,25 @@ pub(super) fn run_pairwise_dmr( let mut success_count = 0; let mut region_error_counts = FxHashMap::::default(); - let mut err: Option = None; + let mut err: Option = None; 'rcv_loop: for batch_result in rcv { match batch_result { BatchResult::Results(results) => { for result in results { match result { Ok(counts) => { - writer.write_all(counts.to_row()?.as_bytes())?; + let row = match counts.to_row() { + Ok(row) => row, + Err(error) => { + err = Some(error); + break 'rcv_loop; + } + }; + if let Err(error) = writer.write_all(row.as_bytes()) + { + err = Some(error.into()); + break 'rcv_loop; + } success_count += 1; pb.inc(1); } @@ -272,7 +286,7 @@ pub(super) fn run_pairwise_dmr( record(s), {message}, stopping" ); }); - err = Some(e); + err = Some(e.into()); break 'rcv_loop; } _ => {} @@ -295,7 +309,7 @@ pub(super) fn run_pairwise_dmr( } }); batch_failures.inc(1u64); - err = Some(error); + err = Some(error.into()); break 'rcv_loop; } } @@ -303,7 +317,7 @@ pub(super) fn run_pairwise_dmr( pb.finish_and_clear(); - finish_pairwise_output(err.map(anyhow::Error::from), writer.as_mut())?; + finish_pairwise_output(err, writer.as_mut())?; Ok((success_count, region_error_counts)) } @@ -311,11 +325,11 @@ fn finish_pairwise_output( first_error: Option, writer: &mut dyn std::io::Write, ) -> anyhow::Result<()> { - if let Some(error) = first_error { - return Err(error); + let flush_result = writer.flush().map_err(anyhow::Error::from); + match first_error { + Some(error) => Err(error), + None => flush_result, } - writer.flush()?; - Ok(()) } #[cfg(test)] diff --git a/modkit-core/src/dmr/single_site.rs b/modkit-core/src/dmr/single_site.rs index 1e707abe..f52ab57f 100644 --- a/modkit-core/src/dmr/single_site.rs +++ b/modkit-core/src/dmr/single_site.rs @@ -30,7 +30,7 @@ use crate::util::{ format_errors_table, get_subroutine_progress_bar, get_ticker, Region, Strand, StrandRule, }; -use crate::writers::TsvWriter; +use crate::writers::{finish_with_first_error, TsvWriter}; pub(super) struct SingleSiteDmrAnalysis { sample_index: Arc, @@ -154,16 +154,9 @@ impl SingleSiteDmrAnalysis { info!("running with replicates, but not matched samples"); } - if self.header { - writer.write_all( - SingleSiteDmrScore::header(multiple_samples, matched_samples) - .as_bytes(), - )?; - } - let mut segmenter: Box = if let Some(segmentation_fp) = &self.segmentation_fp { - Box::new(HmmDmrSegmenter::new( + match HmmDmrSegmenter::new( segmentation_fp, max_gap_size, dmr_prior, @@ -175,11 +168,35 @@ impl SingleSiteDmrAnalysis { decay_distance, &self.multi_progress, self.header, - )?) + ) { + Ok(segmenter) => Box::new(segmenter), + Err(error) => { + return finish_with_first_error( + Some(error), + || writer.flush().map_err(anyhow::Error::from), + "failed to flush single-site DMR output", + ) + } + } } else { Box::new(DummySegmenter::new()) }; + if self.header { + if let Err(error) = writer.write_all( + SingleSiteDmrScore::header(multiple_samples, matched_samples) + .as_bytes(), + ) { + return finish_single_site_outputs( + Some(error.into()), + true, + segmenter.as_mut(), + writer.as_mut(), + &self.multi_progress, + ); + } + } + let (scores_snd, scores_rcv) = crossbeam::channel::bounded(1000); let processed_batches = self.multi_progress.add(get_ticker()); let failure_counter = self.multi_progress.add(get_ticker()); @@ -189,12 +206,23 @@ impl SingleSiteDmrAnalysis { failure_counter.set_message("sites failed"); success_counter.set_message("sites processed successfully"); - let batch_iter = SingleSiteBatches::new( + let batch_iter = match SingleSiteBatches::new( self.sample_index.clone(), self.genome_positions.clone(), self.batch_size, self.interval_size, - )?; + ) { + Ok(batch_iter) => batch_iter, + Err(error) => { + return finish_single_site_outputs( + Some(error), + true, + segmenter.as_mut(), + writer.as_mut(), + &self.multi_progress, + ) + } + }; let sample_index = self.sample_index.clone(); let pmap_estimator = self.pmap_estimator.clone(); @@ -278,7 +306,8 @@ impl SingleSiteDmrAnalysis { let mut success_count = 0usize; let mut error_counts = FxHashMap::::default(); - let mut err: Option = None; + let mut err: Option = None; + let mut run_final_chunk = true; 'rcv_loop: for batch_result in scores_rcv { match batch_result { Err(e) => { @@ -294,20 +323,27 @@ impl SingleSiteDmrAnalysis { }); } } - err = Some(e); + err = Some(e.into()); break 'rcv_loop; } Ok(scores) => { - add_scores_to_segmenter( + if let Err(error) = add_scores_to_segmenter( segmenter.as_mut(), &scores, &self.multi_progress, - )?; + ) { + err = Some(error); + // `add` may have failed after partially emitting a + // chunk. Retrying it during finalization could + // duplicate rows that were already written. + run_final_chunk = false; + break 'rcv_loop; + } for (chrom, results) in scores { for result in results { match result { Ok(scores) => { - writer.write_all( + if let Err(error) = writer.write_all( scores .to_row( multiple_samples, @@ -315,7 +351,10 @@ impl SingleSiteDmrAnalysis { &chrom, ) .as_bytes(), - )?; + ) { + err = Some(error.into()); + break 'rcv_loop; + } success_counter.inc(1); success_count += 1; } @@ -336,7 +375,7 @@ impl SingleSiteDmrAnalysis { {message}, stopping" ); }); - err = Some(e); + err = Some(e.into()); break 'rcv_loop; } } @@ -350,8 +389,8 @@ impl SingleSiteDmrAnalysis { success_counter.finish_and_clear(); failure_counter.finish_and_clear(); finish_single_site_outputs( - err.map(anyhow::Error::from), - true, + err, + run_final_chunk, segmenter.as_mut(), writer.as_mut(), &self.multi_progress, @@ -1016,16 +1055,20 @@ fn finish_single_site_outputs( writer: &mut dyn Write, multi_progress: &MultiProgress, ) -> anyhow::Result<()> { - if run_final_chunk { - finish_segmenter(segmenter, multi_progress)?; + let segmenter_result = if run_final_chunk { + finish_segmenter(segmenter, multi_progress) } else { - segmenter.clean_up()?; - } - if let Some(error) = first_error { - return Err(error); + segmenter.clean_up() + }; + let finalization_result = finish_with_first_error( + segmenter_result.err(), + || writer.flush().map_err(anyhow::Error::from), + "failed to flush single-site DMR output", + ); + match first_error { + Some(error) => Err(error), + None => finalization_result, } - writer.flush()?; - Ok(()) } #[derive(new)] diff --git a/modkit-core/src/dmr/subcommands.rs b/modkit-core/src/dmr/subcommands.rs index f4de988e..91c42786 100644 --- a/modkit-core/src/dmr/subcommands.rs +++ b/modkit-core/src/dmr/subcommands.rs @@ -46,14 +46,58 @@ fn finish_threaded_dmr_output( where F: FnOnce() -> anyhow::Result<()>, { - if let Some(error) = first_error { - return Err(error); + let join_result = join_pipeline(); + let flush_result = writer.flush().map_err(anyhow::Error::from); + match first_error { + Some(error) => Err(error), + None => { + join_result?; + flush_result + } + } +} + +fn write_dmr_header( + writer: &mut dyn Write, + header: &str, +) -> anyhow::Result<()> { + if let Err(error) = writer.write_all(header.as_bytes()) { + // The write failure happened first, but still try to flush any bytes + // the writer accepted before failing. + let _ = writer.flush(); + return Err(error.into()); } - join_pipeline()?; - writer.flush()?; Ok(()) } +fn join_threaded_dmr_pipeline( + source_thread: std::thread::JoinHandle<()>, + worker_threads: Vec>, + aggregator: std::thread::JoinHandle<()>, +) -> anyhow::Result<()> { + // Join every owned thread before selecting an error. This prevents one + // panic from detaching the rest of the pipeline during output cleanup. + let source_result = + source_thread.join().map_err(|_| anyhow!("source thread panicked")); + let worker_results = worker_threads + .into_iter() + .enumerate() + .map(|(i, worker_thread)| { + worker_thread + .join() + .map_err(|_| anyhow!("worker thread {i} panicked")) + }) + .collect::>>(); + let aggregator_result = + aggregator.join().map_err(|_| anyhow!("aggregator thread panicked")); + + source_result?; + for worker_result in worker_results { + worker_result?; + } + aggregator_result +} + #[derive(Subcommand)] pub enum BedMethylDmr { /// Compare regions in a pair of samples (for example, tumor and normal or @@ -1328,7 +1372,6 @@ impl EntryDmrIsoform { sorted_gene_common_coords.len() ); }); - let mut writer = self.get_writer(&multi_progress)?; let transcript_models = Arc::new(transcript_models); let (empties_tx, empties_rx) = crossbeam_channel::unbounded(); @@ -1364,6 +1407,8 @@ impl EntryDmrIsoform { multi_progress .suspend(|| info!("workers staged, starting processing..")); + let mut writer = self.get_writer(&multi_progress)?; + let source_thread = std::thread::spawn({ let results_handle = results_tx.clone(); move || { @@ -1423,15 +1468,27 @@ impl EntryDmrIsoform { drop(records_tx); }); let mut errs = HashMap::new(); - for result in records_rx { + let mut output_error = None; + for result in records_rx.iter() { match result { Ok(mut gene_isoform_dmr) => { - let records_written = gene_isoform_dmr - .write(&mut writer, self.emit_full_results)?; - gene_isoform_dmr.clear(); - let _ = empties_tx.send(gene_isoform_dmr); - records_counter.inc(records_written as u64); - pb.inc(1); + match gene_isoform_dmr + .write(&mut writer, self.emit_full_results) + { + Ok(records_written) => { + gene_isoform_dmr.clear(); + let _ = empties_tx.send(gene_isoform_dmr); + records_counter.inc(records_written as u64); + pb.inc(1); + } + Err(error) => { + output_error = + Some(error.context( + "failed to write isoform DMR output", + )); + break; + } + } } Err(e) => { let c = errs.entry(e.to_string()).or_insert(0usize); @@ -1439,16 +1496,11 @@ impl EntryDmrIsoform { } } } + drop(records_rx); + drop(empties_tx); finish_threaded_dmr_output( - None, - || { - source_thread.join().expect("source thread paniced"); - for (i, worker_thread) in handles.into_iter().enumerate() { - worker_thread.join().expect(&format!("worker {i} paniced")); - } - aggregator.join().expect("aggregator theread paniced"); - Ok(()) - }, + output_error, + || join_threaded_dmr_pipeline(source_thread, handles, aggregator), writer.as_mut(), )?; @@ -1475,11 +1527,11 @@ impl EntryDmrIsoform { Box::new(BufWriter::new(fh)) } }; - writer.write_all( - GeneIsoformDmrRecord::::header( + write_dmr_header( + writer.as_mut(), + &GeneIsoformDmrRecord::::header( self.emit_full_results, - ) - .as_bytes(), + ), )?; Ok(writer) } @@ -1722,7 +1774,7 @@ impl EntryGeneTx { &mut sorted_by_gene_common_coordinates, )?; - let mut writer = self.get_writer(single_mod_code, &multi_progress)?; + let mut gene_labels = self.get_gene_labels(&multi_progress)?; let transcript_models = Arc::new(transcript_models); let (empties_tx, empties_rx) = crossbeam_channel::unbounded(); @@ -1760,6 +1812,8 @@ impl EntryGeneTx { multi_progress .suspend(|| info!("workers staged, starting processing..")); + let mut writer = self.get_writer(single_mod_code, &multi_progress)?; + let source_thread = std::thread::spawn({ let results_handle = results_tx.clone(); move || { @@ -1820,15 +1874,23 @@ impl EntryGeneTx { let mut errs = FxHashMap::default(); let mut plot_points = Vec::with_capacity(n_genes * self.top_k); - let mut gene_labels = self.get_gene_labels(&multi_progress)?; - for result in records_rx { + let mut output_error = None; + for result in records_rx.iter() { match result { Ok(mut gene_tx_dmr) => { - let records_written = gene_tx_dmr.write( + let records_written = match gene_tx_dmr.write( &mut writer, single_mod_code, self.emit_full_results, - )?; + ) { + Ok(records_written) => records_written, + Err(error) => { + output_error = Some(error.context( + "failed to write gene-transcript DMR output", + )); + break; + } + }; if self.plot.is_some() { let points = gene_tx_dmr.topk_records( self.top_k, @@ -1850,58 +1912,63 @@ impl EntryGeneTx { } } - if let Some(fp) = self.plot.as_ref() { - multi_progress.suspend(|| { - info!("plotting {} points to {fp:?}", plot_points.len()) - }); - if let Some(label_top_k_genes) = self.label_top_k_genes { - plot_points.sort_by(|a, b| { - b.neg_log_pvalue - .partial_cmp(&a.neg_log_pvalue) - .unwrap_or(std::cmp::Ordering::Equal) + if output_error.is_none() { + if let Some(fp) = self.plot.as_ref() { + multi_progress.suspend(|| { + info!("plotting {} points to {fp:?}", plot_points.len()) }); - for pp in plot_points.iter() { - let gene_label = - pp.gene_name.clone().unwrap_or_else(|| pp.gene.clone()); - gene_labels.insert(gene_label); - if gene_labels.len() >= label_top_k_genes { - break; + if let Some(label_top_k_genes) = self.label_top_k_genes { + plot_points.sort_by(|a, b| { + b.neg_log_pvalue + .partial_cmp(&a.neg_log_pvalue) + .unwrap_or(std::cmp::Ordering::Equal) + }); + for pp in plot_points.iter() { + let gene_label = pp + .gene_name + .clone() + .unwrap_or_else(|| pp.gene.clone()); + gene_labels.insert(gene_label); + if gene_labels.len() >= label_top_k_genes { + break; + } } - } - for pp in plot_points.iter_mut() { - let gene_label = pp.gene_name.as_ref().unwrap_or(&pp.gene); - if gene_labels.contains(gene_label) { - pp.label_point = true; + for pp in plot_points.iter_mut() { + let gene_label = + pp.gene_name.as_ref().unwrap_or(&pp.gene); + if gene_labels.contains(gene_label) { + pp.label_point = true; + } } } - } - multi_progress.suspend(|| { - let sorted_by = if self.sort_by_effect_size { - "effect size" - } else { - "p-value" - }; - info!( - "plotting the top {} points from each gene, sorted by \ + multi_progress.suspend(|| { + let sorted_by = if self.sort_by_effect_size { + "effect size" + } else { + "p-value" + }; + info!( + "plotting the top {} points from each gene, sorted by \ {sorted_by}", - self.top_k - ); - }); - let svg = volcano_svg(&plot_points, self.plot_title.as_ref()); - std::fs::write(fp, svg)?; + self.top_k + ); + }); + let svg = volcano_svg(&plot_points, self.plot_title.as_ref()); + if let Err(error) = std::fs::write(fp, svg) { + output_error = + Some(anyhow::Error::from(error).context( + "failed to write gene-transcript DMR plot", + )); + } + } } + drop(records_rx); + drop(empties_tx); finish_threaded_dmr_output( - None, - || { - source_thread.join().expect("source thread paniced"); - for (i, worker_thread) in handles.into_iter().enumerate() { - worker_thread.join().expect(&format!("worker {i} paniced")); - } - aggregator.join().expect("aggregator theread paniced"); - Ok(()) - }, + output_error, + || join_threaded_dmr_pipeline(source_thread, handles, aggregator), writer.as_mut(), )?; @@ -1934,12 +2001,12 @@ impl EntryGeneTx { Box::new(BufWriter::new(fh)) } }; - writer.write_all( - GeneIsoformDmrRecord::::header( + write_dmr_header( + writer.as_mut(), + &GeneIsoformDmrRecord::::header( single_mod_code, self.emit_full_results, - ) - .as_bytes(), + ), )?; Ok(writer) } @@ -1972,19 +2039,28 @@ impl EntryGeneTx { #[cfg(test)] mod output_finalization_tests { - use super::finish_threaded_dmr_output; + use super::{ + finish_threaded_dmr_output, join_threaded_dmr_pipeline, + write_dmr_header, + }; use anyhow::anyhow; use std::io::{self, Write}; use std::sync::{Arc, Mutex}; struct EventWriter { events: Arc>>, + fail_write: bool, fail_flush: bool, } impl Write for EventWriter { fn write(&mut self, buf: &[u8]) -> io::Result { - Ok(buf.len()) + self.events.lock().unwrap().push("writer write"); + if self.fail_write { + Err(io::Error::other("writer write failed first")) + } else { + Ok(buf.len()) + } } fn flush(&mut self) -> io::Result<()> { @@ -2004,7 +2080,11 @@ mod output_finalization_tests { ) -> (anyhow::Result<()>, Vec<&'static str>) { let events = Arc::new(Mutex::new(Vec::new())); let join_events = events.clone(); - let mut writer = EventWriter { events: events.clone(), fail_flush }; + let mut writer = EventWriter { + events: events.clone(), + fail_write: false, + fail_flush, + }; let result = finish_threaded_dmr_output( first_error, move || { @@ -2054,4 +2134,47 @@ mod output_finalization_tests { result.unwrap(); assert_eq!(events, ["pipeline joins", "writer flush"]); } + + #[test] + fn header_write_error_is_retained_and_flush_is_attempted() { + let events = Arc::new(Mutex::new(Vec::new())); + let mut writer = EventWriter { + events: events.clone(), + fail_write: true, + fail_flush: true, + }; + + let error = write_dmr_header(&mut writer, "header\n") + .expect_err("header write failure must be returned"); + + assert_eq!(error.to_string(), "writer write failed first"); + assert_eq!(*events.lock().unwrap(), ["writer write", "writer flush"]); + } + + #[test] + fn pipeline_panics_are_fallible_and_every_thread_is_joined() { + let events = Arc::new(Mutex::new(Vec::new())); + let source_events = events.clone(); + let source = std::thread::spawn(move || { + source_events.lock().unwrap().push("source"); + panic!("source test panic"); + }); + let worker_events = events.clone(); + let worker = std::thread::spawn(move || { + worker_events.lock().unwrap().push("worker"); + }); + let aggregator_events = events.clone(); + let aggregator = std::thread::spawn(move || { + aggregator_events.lock().unwrap().push("aggregator"); + }); + + let error = + join_threaded_dmr_pipeline(source, vec![worker], aggregator) + .expect_err("thread panic must be returned as an error"); + + assert_eq!(error.to_string(), "source thread panicked"); + let mut observed = events.lock().unwrap().clone(); + observed.sort_unstable(); + assert_eq!(observed, ["aggregator", "source", "worker"]); + } }