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..eb81fbc8 100644 --- a/modkit-core/src/dmr/pairwise.rs +++ b/modkit-core/src/dmr/pairwise.rs @@ -183,7 +183,11 @@ 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())?; + 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); @@ -247,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(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); } @@ -271,7 +286,7 @@ pub(super) fn run_pairwise_dmr( record(s), {message}, stopping" ); }); - err = Some(e); + err = Some(e.into()); break 'rcv_loop; } _ => {} @@ -294,7 +309,7 @@ pub(super) fn run_pairwise_dmr( } }); batch_failures.inc(1u64); - err = Some(error); + err = Some(error.into()); break 'rcv_loop; } } @@ -302,9 +317,75 @@ pub(super) fn run_pairwise_dmr( pb.finish_and_clear(); - if let Some(e) = err { - Err(e.into()) - } else { - Ok((success_count, region_error_counts)) + finish_pairwise_output(err, writer.as_mut())?; + Ok((success_count, region_error_counts)) +} + +fn finish_pairwise_output( + first_error: Option, + writer: &mut dyn std::io::Write, +) -> anyhow::Result<()> { + let flush_result = writer.flush().map_err(anyhow::Error::from); + match first_error { + Some(error) => Err(error), + None => flush_result, + } +} + +#[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 6e0f6e62..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( - 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) => { - if let Err(e) = segmenter.add(&scores) { - self.multi_progress.suspend(|| { - error!("segmentation error, {e}"); - }) + 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( + 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; } } @@ -347,16 +386,15 @@ 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()?; - - if let Some(e) = err { - return Err(e.into()); - } + finish_single_site_outputs( + err, + run_final_chunk, + segmenter.as_mut(), + writer.as_mut(), + &self.multi_progress, + )?; if !error_counts.is_empty() { self.multi_progress.suspend(|| { @@ -984,6 +1022,55 @@ 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<()> { + let result = segmenter.add(scores); + if let Err(e) = &result { + multi_progress.suspend(|| error!("segmentation error, {e}")); + } + result?; + Ok(()) +} + +fn finish_segmenter( + segmenter: &mut dyn DmrSegmenter, + multi_progress: &MultiProgress, +) -> anyhow::Result<()> { + let result = segmenter.run_current_chunk(); + if let Err(e) = &result { + multi_progress.suspend(|| error!("segmentation error, {e}")); + } + let clean_up_result = segmenter.clean_up(); + result?; + 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<()> { + let segmenter_result = if run_final_chunk { + finish_segmenter(segmenter, multi_progress) + } else { + 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, + } +} + #[derive(new)] struct DummySegmenter {} @@ -1174,6 +1261,7 @@ impl DmrSegmenter for HmmDmrSegmenter { "HMM segmenter finished, wrote {} segments", self.segments_written.position() ); + self.writer.flush()?; Ok(()) } } @@ -1345,3 +1433,217 @@ fn path_to_region_labels( agg } } + +#[cfg(test)] +mod segmenter_error_tests { + use super::{ + 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"; + + #[derive(Default)] + struct StubSegmenter { + 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 { + 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<()> { + self.events.lock().unwrap().push("final chunk"); + 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.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 { + 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); + } + + #[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 66658e15..91c42786 100644 --- a/modkit-core/src/dmr/subcommands.rs +++ b/modkit-core/src/dmr/subcommands.rs @@ -38,6 +38,66 @@ 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<()>, +{ + 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()); + } + 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 @@ -1207,7 +1267,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 +1354,7 @@ impl EntryDmrIsoform { } } } + writer.flush()?; Ok(()) } @@ -1311,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(); @@ -1347,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 || { @@ -1406,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); @@ -1422,11 +1496,13 @@ 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"); + drop(records_rx); + drop(empties_tx); + finish_threaded_dmr_output( + output_error, + || join_threaded_dmr_pipeline(source_thread, handles, aggregator), + writer.as_mut(), + )?; multi_progress.suspend(|| { info!("finished, processed {} genes", pb.position()); @@ -1451,11 +1527,11 @@ impl EntryDmrIsoform { Box::new(BufWriter::new(fh)) } }; - writer.write( - GeneIsoformDmrRecord::::header( + write_dmr_header( + writer.as_mut(), + &GeneIsoformDmrRecord::::header( self.emit_full_results, - ) - .as_bytes(), + ), )?; Ok(writer) } @@ -1698,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(); @@ -1736,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 || { @@ -1796,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, @@ -1826,53 +1912,65 @@ 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", + )); + } + } } - 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"); + drop(records_rx); + drop(empties_tx); + finish_threaded_dmr_output( + output_error, + || join_threaded_dmr_pipeline(source_thread, handles, aggregator), + writer.as_mut(), + )?; multi_progress.suspend(|| { info!("finished, {} errors", errs.len()); @@ -1903,12 +2001,12 @@ impl EntryGeneTx { Box::new(BufWriter::new(fh)) } }; - writer.write( - GeneIsoformDmrRecord::::header( + write_dmr_header( + writer.as_mut(), + &GeneIsoformDmrRecord::::header( single_mod_code, self.emit_full_results, - ) - .as_bytes(), + ), )?; Ok(writer) } @@ -1938,3 +2036,145 @@ impl EntryGeneTx { } } } + +#[cfg(test)] +mod output_finalization_tests { + 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 { + 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<()> { + 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_write: false, + 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"]); + } + + #[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"]); + } +} 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 e1b16e93..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 { @@ -197,19 +219,22 @@ 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) } + + fn finish(&mut self) -> anyhow::Result<()> { + self.inner.flush() + } } pub fn bedmethyl_header() -> String { @@ -256,7 +281,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 +317,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(()) @@ -319,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 { @@ -383,7 +413,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; } @@ -391,6 +421,11 @@ impl PileupWriter for BedMethylWriter { } Ok(rows_written) } + + fn finish(&mut self) -> anyhow::Result<()> { + self.buf_writer.flush()?; + Ok(()) + } } pub struct MultipleMotifBedmethylWriter { @@ -410,7 +445,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 +453,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()); @@ -561,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 { @@ -712,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 { @@ -720,7 +764,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() } } @@ -732,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 }) } } @@ -756,7 +805,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 +838,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 +848,13 @@ 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) + } + + fn finish(&mut self) -> AnyhowResult<()> { + self.flush()?; + Ok(()) } } @@ -878,9 +930,14 @@ 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) } + + fn finish(&mut self) -> AnyhowResult<()> { + self.flush()?; + Ok(()) + } } #[derive(new)] @@ -1142,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; @@ -1149,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 { @@ -1157,37 +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(blob.as_bytes()).map(|_x| ())? - } - 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(blob.as_bytes()).map(|_x| ())?, - 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(); } } @@ -1419,57 +1582,311 @@ 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) } + + 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::{ + 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}; + use std::sync::{Arc, Mutex}; + + #[derive(Clone)] + struct ShortWriter { + bytes: Arc>>, + max_write: usize, + fail_flush: bool, + flush_count: Arc, + } + + 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, + flush_count: Arc::new(AtomicUsize::new(0)), + } + } + + fn failing_flush(max_write: usize) -> Self { + Self { fail_flush: true, ..Self::new(max_write) } + } + + fn bytes(&self) -> Vec { + self.bytes.lock().unwrap().clone() + } + + fn flush_count(&self) -> usize { + self.flush_count.load(Ordering::SeqCst) + } + } + + 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<()> { + self.flush_count.fetch_add(1, Ordering::SeqCst); + 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()); + } + + #[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); + } + + #[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); + } }