From 619efb1205afde40c123adb2dc0099efc9c114ec Mon Sep 17 00:00:00 2001 From: SuhasSrinivasan <32346517+SuhasSrinivasan@users.noreply.github.com> Date: Sun, 9 Aug 2026 19:14:50 -0700 Subject: [PATCH 1/4] Add generic ordered scheduler for repair --- modkit-core/src/lib.rs | 1 + modkit-core/src/ordered_scheduler.rs | 907 +++++++++++++++++++++++++++ 2 files changed, 908 insertions(+) create mode 100644 modkit-core/src/ordered_scheduler.rs diff --git a/modkit-core/src/lib.rs b/modkit-core/src/lib.rs index b64d7709..dee93a29 100644 --- a/modkit-core/src/lib.rs +++ b/modkit-core/src/lib.rs @@ -16,6 +16,7 @@ pub mod mod_base_code; pub mod modbam_util; pub mod monoid; pub mod motifs; +pub(crate) mod ordered_scheduler; pub mod pileup; pub mod position_filter; pub mod read_ids_to_base_mod_probs; diff --git a/modkit-core/src/ordered_scheduler.rs b/modkit-core/src/ordered_scheduler.rs new file mode 100644 index 00000000..2d34257f --- /dev/null +++ b/modkit-core/src/ordered_scheduler.rs @@ -0,0 +1,907 @@ +use std::any::Any; +use std::collections::BTreeMap; +use std::panic::{catch_unwind, AssertUnwindSafe}; +use std::sync::{Arc, Mutex}; +use std::thread::JoinHandle; + +use anyhow::anyhow; +use crossbeam_channel::{bounded, unbounded, Receiver, Sender}; + +#[derive(Clone)] +struct FailureState { + first: Arc>>, + signal: Sender<()>, +} + +impl FailureState { + fn report(&self, error: anyhow::Error) { + let mut first = + self.first.lock().unwrap_or_else(|poisoned| poisoned.into_inner()); + if first.is_none() { + *first = Some(error); + let _ = self.signal.try_send(()); + } + } + + fn has_failed(&self) -> bool { + self.first + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .is_some() + } + + fn take(&self) -> Option { + self.first + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .take() + } +} + +fn panic_message(panic: Box) -> String { + if let Some(message) = panic.downcast_ref::<&str>() { + (*message).to_string() + } else if let Some(message) = panic.downcast_ref::() { + message.clone() + } else { + "unknown panic payload".to_string() + } +} + +fn run_stage( + label: &str, + failure: FailureState, + mut resources: R, + task: F, +) where + F: FnOnce(&mut R) -> anyhow::Result<()>, +{ + // Keep every channel in `resources` alive until the originating failure is + // recorded, so a derived disconnection cannot replace the root cause. + let outcome = catch_unwind(AssertUnwindSafe(|| task(&mut resources))); + match outcome { + Ok(Ok(())) => {} + Ok(Err(error)) => failure.report(anyhow!("{label} failed: {error:#}")), + Err(panic) => failure + .report(anyhow!("{label} panicked: {}", panic_message(panic))), + } +} + +fn join_stage(label: &str, handle: JoinHandle<()>, failure: &FailureState) { + if let Err(panic) = handle.join() { + failure.report(anyhow!( + "{label} panicked outside its scheduler guard: {}", + panic_message(panic) + )); + } +} + +pub(crate) trait OrderedWorker: Send { + fn process(&mut self, job: J, buffer: B) -> anyhow::Result; +} + +pub(crate) fn run_ordered_scheduler( + stage_name: &'static str, + feeder: I, + workers: Vec, + empty_buffers: Receiver, + output_queue_size: usize, + mut consume: C, +) -> anyhow::Result<()> +where + I: Iterator> + Send + 'static, + J: Send + 'static, + B: Send + 'static, + W: OrderedWorker + 'static, + C: FnMut(B) -> anyhow::Result<()>, +{ + if workers.is_empty() { + return Err(anyhow!( + "{stage_name} scheduler requires at least one worker" + )); + } + + let (jobs_tx, jobs_rx) = bounded(workers.len() * 2); + let (results_tx, results_rx) = unbounded(); + let (records_tx, records_rx) = bounded(output_queue_size); + // Dropping the only sender broadcasts cancellation to every blocked + // receiver without consuming a one-shot cancellation message. + let (cancel_tx, cancel_rx) = bounded::<()>(0); + let (failure_signal_tx, failure_signal_rx) = bounded(1); + let failure = FailureState { + first: Arc::new(Mutex::new(None)), + signal: failure_signal_tx, + }; + + let source_failure = failure.clone(); + let source_cancel = cancel_rx.clone(); + let source = std::thread::spawn(move || { + run_stage( + &format!("{stage_name} source"), + source_failure, + (feeder, jobs_tx, empty_buffers, source_cancel), + |resources| { + let (feeder, jobs_tx, empty_buffers, cancel) = resources; + let mut seq = 0usize; + loop { + let job = match feeder.next() { + Some(Ok(job)) => job, + Some(Err(error)) => { + return Err(anyhow!("feeder error: {error:#}")) + } + None => return Ok(()), + }; + let buffer = crossbeam_channel::select_biased! { + recv(cancel) -> _ => return Ok(()), + recv(empty_buffers) -> result => result.map_err(|_| { + anyhow!("buffer-credit channel disconnected") + })?, + }; + let sent = crossbeam_channel::select_biased! { + recv(cancel) -> _ => return Ok(()), + send(jobs_tx, (seq, job, buffer)) -> result => result, + }; + sent.map_err(|_| anyhow!("job channel disconnected"))?; + seq = seq.wrapping_add(1); + } + }, + ); + }); + + let mut worker_handles = Vec::with_capacity(workers.len()); + for (index, worker) in workers.into_iter().enumerate() { + let worker_failure = failure.clone(); + let worker_cancel = cancel_rx.clone(); + let worker_jobs = jobs_rx.clone(); + let worker_results = results_tx.clone(); + worker_handles.push(std::thread::spawn(move || { + let label = format!("{stage_name} worker {index}"); + run_stage( + &label, + worker_failure, + (worker, worker_jobs, worker_results, worker_cancel), + |resources| { + let (worker, jobs, results, cancel) = resources; + loop { + let received = crossbeam_channel::select_biased! { + recv(cancel) -> _ => return Ok(()), + recv(jobs) -> result => result, + }; + let (seq, job, buffer) = match received { + Ok(job) => job, + Err(_) => return Ok(()), + }; + let output = worker.process(job, buffer)?; + let sent = crossbeam_channel::select_biased! { + recv(cancel) -> _ => return Ok(()), + send(results, (seq, output)) -> result => result, + }; + sent.map_err(|_| { + anyhow!("results channel disconnected") + })?; + } + }, + ); + })); + } + drop(results_tx); + drop(jobs_rx); + + let aggregator_failure = failure.clone(); + let aggregator_cancel = cancel_rx; + let aggregator = std::thread::spawn(move || { + run_stage( + &format!("{stage_name} aggregator"), + aggregator_failure, + (results_rx, records_tx, aggregator_cancel), + |resources| { + let (results, records, cancel) = resources; + let mut next_seq = 0usize; + let mut buffer = BTreeMap::new(); + loop { + let received = crossbeam_channel::select_biased! { + recv(cancel) -> _ => return Ok(()), + recv(results) -> result => result, + }; + let (seq, output) = match received { + Ok(result) => result, + Err(_) if buffer.is_empty() => return Ok(()), + Err(_) => { + return Err(anyhow!( + "results channel closed before sequence \ + {next_seq}" + )) + } + }; + if buffer.insert(seq, output).is_some() { + return Err(anyhow!("duplicate result sequence {seq}")); + } + while let Some(output) = buffer.remove(&next_seq) { + let sent = crossbeam_channel::select_biased! { + recv(cancel) -> _ => return Ok(()), + send(records, output) -> result => result, + }; + sent.map_err(|_| { + anyhow!("ordered-output channel disconnected") + })?; + next_seq = next_seq.wrapping_add(1); + } + } + }, + ); + }); + + let mut cancel_tx = Some(cancel_tx); + loop { + crossbeam_channel::select_biased! { + recv(failure_signal_rx) -> _ => { + drop(cancel_tx.take()); + break; + }, + recv(records_rx) -> result => match result { + Ok(output) => { + match catch_unwind(AssertUnwindSafe(|| consume(output))) { + Ok(Ok(())) => {} + Ok(Err(error)) => { + failure.report(anyhow!( + "{stage_name} output consumer failed: {error:#}" + )); + drop(cancel_tx.take()); + break; + } + Err(panic) => { + failure.report(anyhow!( + "{stage_name} output consumer panicked: {}", + panic_message(panic) + )); + drop(cancel_tx.take()); + break; + } + } + } + Err(_) => { + if failure.has_failed() { + drop(cancel_tx.take()); + } + break; + } + }, + } + } + drop(records_rx); + // Cancellation makes every channel-blocked stage wake before joining. + drop(cancel_tx.take()); + + join_stage(&format!("{stage_name} source"), source, &failure); + for (index, worker) in worker_handles.into_iter().enumerate() { + join_stage(&format!("{stage_name} worker {index}"), worker, &failure); + } + join_stage(&format!("{stage_name} aggregator"), aggregator, &failure); + + match failure.take() { + Some(error) => Err(error), + None => Ok(()), + } +} + +#[cfg(test)] +mod tests { + use std::collections::VecDeque; + use std::panic::{catch_unwind, AssertUnwindSafe}; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, Barrier, Mutex}; + use std::thread; + use std::time::Duration; + + use anyhow::{anyhow, bail}; + use crossbeam_channel::{bounded, unbounded, Receiver, Sender}; + + use super::{run_ordered_scheduler, OrderedWorker}; + + struct TestBuffer { + value: usize, + drops: Option>, + } + + impl Drop for TestBuffer { + fn drop(&mut self) { + if let Some(drops) = self.drops.as_ref() { + drops.fetch_add(1, Ordering::SeqCst); + } + } + } + + #[derive(Clone, Copy)] + enum WorkerBehavior { + Succeed, + ErrorAt(usize), + PanicAt(usize), + } + + struct TestWorker { + behavior: WorkerBehavior, + drops: Option>, + } + + impl Drop for TestWorker { + fn drop(&mut self) { + if let Some(drops) = self.drops.as_ref() { + drops.fetch_add(1, Ordering::SeqCst); + } + } + } + + impl OrderedWorker for TestWorker { + fn process( + &mut self, + job: usize, + mut buffer: TestBuffer, + ) -> anyhow::Result { + match self.behavior { + WorkerBehavior::Succeed => {} + WorkerBehavior::ErrorAt(target) if job == target => { + bail!("single worker failure at job {job}") + } + WorkerBehavior::PanicAt(target) if job == target => { + panic!("scripted worker panic at job {job}") + } + WorkerBehavior::ErrorAt(_) | WorkerBehavior::PanicAt(_) => {} + } + buffer.value = job; + Ok(buffer) + } + } + + struct BlockingJoinWorker { + job_two_started: Sender<()>, + wait_for_job_two: Receiver<()>, + release_job_two: Receiver<()>, + failure_ready: Sender<()>, + drops: Arc, + } + + impl Drop for BlockingJoinWorker { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + } + } + + impl OrderedWorker for BlockingJoinWorker { + fn process( + &mut self, + job: usize, + mut buffer: TestBuffer, + ) -> anyhow::Result { + match job { + 0 => { + self.wait_for_job_two + .recv_timeout(Duration::from_secs(1)) + .map_err(|_| anyhow!("job 2 never started"))?; + let _ = self.failure_ready.send(()); + bail!("cancellation trigger at job 0") + } + 2 => { + let _ = self.job_two_started.send(()); + self.release_job_two + .recv_timeout(Duration::from_secs(1)) + .map_err(|_| anyhow!("job 2 was not released"))?; + } + _ => {} + } + buffer.value = job; + Ok(buffer) + } + } + + struct CoordinatedErrorWorker { + ready: Arc, + } + + impl OrderedWorker for CoordinatedErrorWorker { + fn process( + &mut self, + job: usize, + _buffer: TestBuffer, + ) -> anyhow::Result { + self.ready.wait(); + bail!("scripted root failure at job {job}") + } + } + + struct ConsumerJoinWorker { + job_one_started: Sender<()>, + wait_for_job_one: Receiver<()>, + release_job_one: Receiver<()>, + drops: Arc, + } + + impl Drop for ConsumerJoinWorker { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + } + } + + impl OrderedWorker for ConsumerJoinWorker { + fn process( + &mut self, + job: usize, + mut buffer: TestBuffer, + ) -> anyhow::Result { + match job { + 0 => { + self.wait_for_job_one + .recv_timeout(Duration::from_secs(1)) + .map_err(|_| anyhow!("job 1 never started"))?; + } + 1 => { + let _ = self.job_one_started.send(()); + self.release_job_one + .recv_timeout(Duration::from_secs(1)) + .map_err(|_| anyhow!("job 1 was not released"))?; + } + _ => {} + } + buffer.value = job; + Ok(buffer) + } + } + + struct OrderedCompletionWorker { + job_one_completed: Sender<()>, + release_job_zero: Receiver<()>, + completions: Arc>>, + } + + impl OrderedWorker for OrderedCompletionWorker { + fn process( + &mut self, + job: usize, + mut buffer: TestBuffer, + ) -> anyhow::Result { + if job == 0 { + self.release_job_zero + .recv_timeout(Duration::from_secs(1)) + .map_err(|_| anyhow!("job 0 was not released"))?; + } + self.completions.lock().unwrap().push(job); + if job == 1 { + let _ = self.job_one_completed.send(()); + } + buffer.value = job; + Ok(buffer) + } + } + + #[derive(Clone, Copy)] + enum ConsumerFailure { + Error, + Panic, + } + + struct TestFeeder { + items: VecDeque>, + drops: Option>, + } + + impl TestFeeder { + fn jobs(jobs: impl IntoIterator) -> Self { + Self { items: jobs.into_iter().map(Ok).collect(), drops: None } + } + } + + impl Iterator for TestFeeder { + type Item = anyhow::Result; + + fn next(&mut self) -> Option { + self.items.pop_front() + } + } + + impl Drop for TestFeeder { + fn drop(&mut self) { + if let Some(drops) = self.drops.as_ref() { + drops.fetch_add(1, Ordering::SeqCst); + } + } + } + + fn make_buffer_pool( + count: usize, + drops: Option>, + ) -> (Sender, crossbeam_channel::Receiver) { + let (sender, receiver) = unbounded(); + for _ in 0..count { + sender + .send(TestBuffer { value: usize::MAX, drops: drops.clone() }) + .unwrap(); + } + (sender, receiver) + } + + fn recycle_buffer( + sender: &Sender, + buffer: TestBuffer, + ) -> anyhow::Result<()> { + let _ = sender.send(buffer); + Ok(()) + } + + fn run_scenario( + feeder: TestFeeder, + behavior: WorkerBehavior, + worker_count: usize, + buffer_count: usize, + ) -> anyhow::Result<()> { + let workers = (0..worker_count) + .map(|_| TestWorker { behavior, drops: None }) + .collect(); + let (empty_sender, empty_buffers) = + make_buffer_pool(buffer_count, None); + let recycle = empty_sender.clone(); + run_ordered_scheduler( + "pileup", + feeder, + workers, + empty_buffers, + worker_count * 2, + move |buffer| recycle_buffer(&recycle, buffer), + ) + } + + fn run_with_watchdog(timeout: Duration, task: F) -> anyhow::Result<()> + where + F: FnOnce() -> anyhow::Result<()> + Send + 'static, + { + let (sender, receiver) = bounded(1); + thread::spawn(move || { + let result = + catch_unwind(AssertUnwindSafe(task)).map_err(|panic| { + if let Some(message) = panic.downcast_ref::<&str>() { + anyhow!("scheduler panicked: {message}") + } else if let Some(message) = panic.downcast_ref::() + { + anyhow!("scheduler panicked: {message}") + } else { + anyhow!("scheduler panicked") + } + }); + let _ = sender.send(result.and_then(|result| result)); + }); + receiver.recv_timeout(timeout).map_err(|_| { + anyhow!("scheduler watchdog expired after {timeout:?}") + })? + } + + fn assert_consumer_failure_joins_all_threads(mode: ConsumerFailure) { + let worker_drops = Arc::new(AtomicUsize::new(0)); + let feeder_drops = Arc::new(AtomicUsize::new(0)); + let buffer_drops = Arc::new(AtomicUsize::new(0)); + let (job_one_started, wait_for_job_one) = bounded(1); + let (release_job_one, job_one_release) = bounded(1); + let (consumer_called, wait_for_consumer) = bounded(1); + let (finished, wait_for_finish) = bounded(1); + let worker_drop_probe = worker_drops.clone(); + let feeder_drop_probe = feeder_drops.clone(); + let buffer_drop_probe = buffer_drops.clone(); + + thread::spawn(move || { + let feeder = TestFeeder { + items: (0..64).map(Ok).collect(), + drops: Some(feeder_drop_probe), + }; + let workers = (0..2) + .map(|_| ConsumerJoinWorker { + job_one_started: job_one_started.clone(), + wait_for_job_one: wait_for_job_one.clone(), + release_job_one: job_one_release.clone(), + drops: worker_drop_probe.clone(), + }) + .collect(); + let (empty_sender, empty_buffers) = + make_buffer_pool(4, Some(buffer_drop_probe)); + let outcome = catch_unwind(AssertUnwindSafe(|| { + run_ordered_scheduler( + "pileup", + feeder, + workers, + empty_buffers, + 0, + move |_buffer| { + let _ = consumer_called.send(()); + match mode { + ConsumerFailure::Error => { + bail!("scripted consumer failure") + } + ConsumerFailure::Panic => { + panic!("scripted consumer panic") + } + } + }, + ) + })); + drop(empty_sender); + let _ = finished.send(outcome); + }); + + wait_for_consumer + .recv_timeout(Duration::from_secs(1)) + .expect("consumer was never called"); + if wait_for_finish.recv_timeout(Duration::from_millis(50)).is_ok() { + let _ = release_job_one.send(()); + panic!("scheduler returned before joining the blocked worker"); + } + release_job_one.send(()).expect("failed to release blocked worker"); + let outcome = wait_for_finish + .recv_timeout(Duration::from_secs(1)) + .expect("scheduler did not finish after cancellation"); + let result = match outcome { + Ok(result) => result, + Err(_) => panic!("consumer panic escaped the scheduler"), + }; + let error = result.expect_err("consumer failure should be propagated"); + match mode { + ConsumerFailure::Error => { + assert!(error.to_string().contains("scripted consumer failure")) + } + ConsumerFailure::Panic => { + assert!(error.to_string().contains( + "pileup output consumer panicked: scripted consumer panic" + )) + } + } + assert_eq!(worker_drops.load(Ordering::SeqCst), 2); + assert_eq!(feeder_drops.load(Ordering::SeqCst), 1); + assert_eq!(buffer_drops.load(Ordering::SeqCst), 4); + } + + #[test] + fn zero_workers_are_rejected_without_feeding_or_consuming() { + let feeder_calls = Arc::new(AtomicUsize::new(0)); + let consumer_calls = Arc::new(AtomicUsize::new(0)); + let feeder_probe = feeder_calls.clone(); + let consumer_probe = consumer_calls.clone(); + let error = run_with_watchdog(Duration::from_secs(1), move || { + let feeder = std::iter::from_fn(move || { + feeder_probe.fetch_add(1, Ordering::SeqCst); + Some(Ok(0)) + }); + let (_empty_sender, empty_buffers) = unbounded(); + run_ordered_scheduler( + "pileup", + feeder, + Vec::::new(), + empty_buffers, + 0, + move |_buffer: TestBuffer| { + consumer_probe.fetch_add(1, Ordering::SeqCst); + Ok(()) + }, + ) + }) + .expect_err("an empty worker set should be rejected"); + + assert!(error.to_string().contains("at least one worker")); + assert_eq!(feeder_calls.load(Ordering::SeqCst), 0); + assert_eq!(consumer_calls.load(Ordering::SeqCst), 0); + } + + #[test] + fn propagates_single_worker_error() { + let error = run_with_watchdog(Duration::from_secs(1), || { + run_scenario( + TestFeeder::jobs([0]), + WorkerBehavior::ErrorAt(0), + 1, + 2, + ) + }) + .expect_err("worker error should be propagated"); + assert!(error.to_string().contains("single worker failure at job 0")); + } + + #[test] + fn concurrent_worker_errors_preserve_a_scripted_root_cause() { + let error = run_with_watchdog(Duration::from_secs(1), || { + let ready = Arc::new(Barrier::new(2)); + let workers = (0..2) + .map(|_| CoordinatedErrorWorker { ready: ready.clone() }) + .collect(); + let (empty_sender, empty_buffers) = make_buffer_pool(4, None); + let recycle = empty_sender.clone(); + run_ordered_scheduler( + "pileup", + TestFeeder::jobs(0..64), + workers, + empty_buffers, + 4, + move |buffer| recycle_buffer(&recycle, buffer), + ) + }) + .expect_err("a coordinated worker error should be propagated"); + let message = error.to_string(); + assert!( + message.contains("scripted root failure at job 0") + || message.contains("scripted root failure at job 1"), + "unexpected error: {message}" + ); + assert!( + !message.contains("disconnected"), + "derived error won: {message}" + ); + } + + #[test] + fn worker_panic_is_propagated_without_sequence_hole_deadlock() { + let error = run_with_watchdog(Duration::from_secs(1), || { + run_scenario( + TestFeeder::jobs(0..64), + WorkerBehavior::PanicAt(0), + 1, + 2, + ) + }) + .expect_err("worker panic should be propagated"); + assert!(error.to_string().contains("scripted worker panic at job 0")); + } + + #[test] + fn feeder_error_is_propagated() { + let error = run_with_watchdog(Duration::from_secs(1), || { + run_scenario( + TestFeeder { + items: vec![Err(anyhow!("scripted feeder failure"))].into(), + drops: None, + }, + WorkerBehavior::Succeed, + 1, + 2, + ) + }) + .expect_err("feeder error should be propagated"); + assert!(error.to_string().contains("scripted feeder failure")); + } + + #[test] + fn disconnected_buffer_credit_channel_is_propagated() { + let error = run_with_watchdog(Duration::from_secs(1), || { + let (empty_sender, empty_buffers) = unbounded::(); + drop(empty_sender); + run_ordered_scheduler( + "pileup", + TestFeeder::jobs([0]), + vec![TestWorker { + behavior: WorkerBehavior::Succeed, + drops: None, + }], + empty_buffers, + 2, + |_| Ok(()), + ) + }) + .expect_err("buffer-credit channel failure should be propagated"); + assert!(error.to_string().contains("buffer")); + } + + #[test] + fn failure_cancels_blocked_stages_and_joins_all_threads() { + let worker_drops = Arc::new(AtomicUsize::new(0)); + let feeder_drops = Arc::new(AtomicUsize::new(0)); + let buffer_drops = Arc::new(AtomicUsize::new(0)); + let (job_two_started, wait_for_job_two) = bounded(1); + let (release_job_two, job_two_release) = bounded(1); + let (failure_ready, wait_for_failure) = bounded(1); + let (finished, wait_for_finish) = bounded(1); + let worker_drop_probe = worker_drops.clone(); + let feeder_drop_probe = feeder_drops.clone(); + let buffer_drop_probe = buffer_drops.clone(); + + thread::spawn(move || { + let feeder = TestFeeder { + items: (0..64).map(Ok).collect(), + drops: Some(feeder_drop_probe), + }; + let workers = (0..2) + .map(|_| BlockingJoinWorker { + job_two_started: job_two_started.clone(), + wait_for_job_two: wait_for_job_two.clone(), + release_job_two: job_two_release.clone(), + failure_ready: failure_ready.clone(), + drops: worker_drop_probe.clone(), + }) + .collect(); + let (empty_sender, empty_buffers) = + make_buffer_pool(4, Some(buffer_drop_probe)); + let recycle = empty_sender.clone(); + let result = run_ordered_scheduler( + "pileup", + feeder, + workers, + empty_buffers, + 1, + move |buffer| recycle_buffer(&recycle, buffer), + ); + drop(empty_sender); + let _ = finished.send(result); + }); + + wait_for_failure + .recv_timeout(Duration::from_secs(1)) + .expect("scripted worker failure never became ready"); + assert!( + wait_for_finish.recv_timeout(Duration::from_millis(50)).is_err(), + "scheduler returned before joining the blocked worker" + ); + release_job_two.send(()).expect("failed to release blocked worker"); + let error = wait_for_finish + .recv_timeout(Duration::from_secs(1)) + .expect("scheduler did not finish after cancellation") + .expect_err("worker failure should cancel the pipeline"); + assert!(error.to_string().contains("cancellation trigger at job 0")); + assert_eq!(worker_drops.load(Ordering::SeqCst), 2); + assert_eq!(feeder_drops.load(Ordering::SeqCst), 1); + assert_eq!(buffer_drops.load(Ordering::SeqCst), 4); + } + + #[test] + fn consumer_error_cancels_and_joins_all_threads() { + assert_consumer_failure_joins_all_threads(ConsumerFailure::Error); + } + + #[test] + fn consumer_panic_is_converted_and_joins_all_threads() { + assert_consumer_failure_joins_all_threads(ConsumerFailure::Panic); + } + + #[test] + fn out_of_order_completions_are_consumed_in_feeder_order() { + let completions = Arc::new(Mutex::new(Vec::new())); + let consumed = Arc::new(Mutex::new(Vec::new())); + let (job_one_completed, wait_for_job_one) = bounded(1); + let (release_job_zero, job_zero_release) = bounded(1); + let (finished, wait_for_finish) = bounded(1); + let completion_probe = completions.clone(); + let consumed_probe = consumed.clone(); + + thread::spawn(move || { + let workers = (0..2) + .map(|_| OrderedCompletionWorker { + job_one_completed: job_one_completed.clone(), + release_job_zero: job_zero_release.clone(), + completions: completion_probe.clone(), + }) + .collect(); + let (empty_sender, empty_buffers) = make_buffer_pool(2, None); + let recycle = empty_sender.clone(); + let result = run_ordered_scheduler( + "pileup", + TestFeeder::jobs(0..2), + workers, + empty_buffers, + 0, + move |buffer| { + consumed_probe.lock().unwrap().push(buffer.value); + recycle_buffer(&recycle, buffer) + }, + ); + drop(empty_sender); + let _ = finished.send(result); + }); + + wait_for_job_one + .recv_timeout(Duration::from_secs(1)) + .expect("job 1 did not complete"); + assert_eq!(*completions.lock().unwrap(), vec![1]); + assert!(consumed.lock().unwrap().is_empty()); + assert!(wait_for_finish + .recv_timeout(Duration::from_millis(50)) + .is_err()); + release_job_zero.send(()).expect("failed to release job 0"); + wait_for_finish + .recv_timeout(Duration::from_secs(1)) + .expect("scheduler did not finish") + .expect("successful scheduler run should complete"); + + assert_eq!(*completions.lock().unwrap(), vec![1, 0]); + assert_eq!(*consumed.lock().unwrap(), vec![0, 1]); + } +} From 69db1a2069719183f6486002ecd6fe2878801ab7 Mon Sep 17 00:00:00 2001 From: SuhasSrinivasan <32346517+SuhasSrinivasan@users.noreply.github.com> Date: Sun, 9 Aug 2026 19:20:45 -0700 Subject: [PATCH 2/4] Preserve repair query-name and output ordering --- modkit-core/src/repair_tags.rs | 825 +++++++++++++++++++++++++++------ modkit/tests/test_repair.rs | 327 ++++++++++++- 2 files changed, 985 insertions(+), 167 deletions(-) diff --git a/modkit-core/src/repair_tags.rs b/modkit-core/src/repair_tags.rs index 3aa8f2af..ded5201a 100644 --- a/modkit-core/src/repair_tags.rs +++ b/modkit-core/src/repair_tags.rs @@ -1,12 +1,13 @@ +use std::cmp::Ordering; use std::path::PathBuf; use std::sync::Arc; -use anyhow::{anyhow, bail, Context}; +use anyhow::{anyhow, bail}; use clap::Args; +use crossbeam_channel::unbounded; use derive_new::new; use indicatif::{MultiProgress, ProgressBar}; use log::{debug, error, info, warn}; -use rayon::prelude::*; use rust_htslib::bam::record::{Aux, AuxArray}; use rust_htslib::bam::{self, Read}; use rustc_hash::FxHashMap; @@ -17,6 +18,7 @@ use crate::mod_bam::{ format_mm_ml_tag, BaseModProbs, DeltaListConverter, ModBaseInfo, SeqPosBaseModProbs, ML_TAGS, MM_TAGS, MN_TAG, }; +use crate::ordered_scheduler::{run_ordered_scheduler, OrderedWorker}; use crate::util::{ get_forward_sequence_str, get_query_name_string, get_ticker, record_is_not_primary, @@ -53,18 +55,28 @@ impl RepairTags { std::cmp::min(half, 16) }; let threads_per_reader = std::cmp::max(reader_threads / 2, 1); - let pool_threads = - self.threads.checked_sub(reader_threads).unwrap_or(1); + let pool_threads = self.threads.saturating_sub(reader_threads).max(1); debug!( "assigning {threads_per_reader} to each reader and using \ {pool_threads} to process records" ); - let (pair_snd, pair_rcv) = std::sync::mpsc::sync_channel(1000); let mut donor_records = bam::Reader::from_path(&self.donor_bam)?; donor_records.set_threads(threads_per_reader)?; let mut acceptor_records = bam::Reader::from_path(&self.acceptor_bam)?; acceptor_records.set_threads(threads_per_reader)?; + let donor_order = query_name_order(donor_records.header(), "donor")?; + let acceptor_order = + query_name_order(acceptor_records.header(), "acceptor")?; + let merge_order = match donor_order.common_merge_order(acceptor_order) { + Some(order) => order, + None => bail!( + "donor and acceptor BAMs use incompatible query-name sorting: \ + donor is {}, acceptor is {}", + donor_order.label(), + acceptor_order.label() + ), + }; let header = bam::Header::from_template(acceptor_records.header()); let mut writer = bam::Writer::from_path( &self.output_bam, @@ -88,51 +100,19 @@ impl RepairTags { let written_ticker = master_progress.add(get_ticker()).with_message("~records written"); - std::thread::spawn(move || { - let pair_iter = ZipRecordsIter::new( - donor_records.records(), - acceptor_records.records(), - donor_ticker, - acceptor_ticker, - ); - for pair in pair_iter { - match pair_snd.send(pair) { - Ok(_) => {} - Err(e) => { - error!( - "failed to send record pair on channel, {}", - e.to_string() - ); - } - } - } - }); - - let (repair_snd, repair_rcv) = std::sync::mpsc::sync_channel(1000); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(pool_threads) - .build() - .context("failed to make thread pool")?; - std::thread::spawn(move || { - pool.install(|| { - pair_rcv.into_iter().par_bridge().for_each(|record_pair| { - let repaired = repair_record_pair(record_pair); - match repair_snd.send(repaired) { - Ok(_) => repaired_ticker.inc(1), - Err(e) => { - error!( - "failed to send repaired record on channel, {}", - e.to_string() - ); - } - } - }) - }) - }); - + let pair_iter = ZipRecordsIter::new( + donor_records, + acceptor_records, + donor_ticker, + acceptor_ticker, + merge_order, + ); + let workers = + (0..pool_threads).map(|_| RepairWorker).collect::>(); let mut n_repaired = 0usize; let mut n_failed = 0usize; - for res in repair_rcv { + run_repair_scheduler(pair_iter, workers, 1000, |res| { + repaired_ticker.inc(1); match res { Ok(record) => { if let Err(e) = writer.write(&record) { @@ -148,152 +128,486 @@ impl RepairTags { n_failed += 1; } } - } + Ok(()) + })?; info!("finished, repaired {n_repaired} records, {n_failed} failed."); Ok(()) } } +struct RepairSlot { + outcome: Option>, +} + +enum RepairJob { + Matched(RecordPair), + MissingDonor { read_name: Vec }, +} + +struct RepairWorker; + +impl OrderedWorker for RepairWorker { + fn process( + &mut self, + job: RepairJob, + mut slot: RepairSlot, + ) -> anyhow::Result { + debug_assert!(slot.outcome.is_none()); + slot.outcome = Some(match job { + RepairJob::Matched(record_pair) => repair_record_pair(record_pair), + RepairJob::MissingDonor { read_name } => Err(anyhow!( + "record {} failed, no primary donor record", + String::from_utf8_lossy(&read_name) + )), + }); + Ok(slot) + } +} + +fn run_repair_scheduler( + feeder: I, + workers: Vec, + output_queue_size: usize, + mut consume: C, +) -> anyhow::Result<()> +where + I: Iterator> + Send + 'static, + W: OrderedWorker + 'static, + C: FnMut(anyhow::Result) -> anyhow::Result<()>, +{ + let (empty_sender, empty_buffers) = unbounded(); + for _ in 0..(workers.len() * 2) { + empty_sender + .send(RepairSlot { outcome: None }) + .expect("unbounded repair buffer channel should be connected"); + } + let recycle = empty_sender.clone(); + run_ordered_scheduler( + "repair", + feeder, + workers, + empty_buffers, + output_queue_size, + move |mut slot| { + let outcome = slot.outcome.take().ok_or_else(|| { + anyhow!("repair worker returned an empty slot") + })?; + let consumed = consume(outcome); + // The source drops its receiver after it has fed the last job, so + // recycling the final ordered slots may legitimately disconnect. + let _ = recycle.send(slot); + consumed + }, + ) +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum QueryNameOrder { + /// `SS:queryname:natural`, matching samtools' numeric-run comparator. + Natural, + /// SO-only bundled inputs were sorted by an older natural comparator that + /// orders numerically tied runs by leading-zero count. + LegacyNatural, + Lexicographical, +} + +impl QueryNameOrder { + fn common_merge_order(self, other: Self) -> Option { + match (self, other) { + (Self::Natural, Self::Natural) => Some(Self::Natural), + (Self::LegacyNatural, Self::LegacyNatural) => { + Some(Self::LegacyNatural) + } + (Self::Lexicographical, Self::Lexicographical) => { + Some(Self::Lexicographical) + } + _ => None, + } + } + + fn compare(self, left: &[u8], right: &[u8]) -> Ordering { + match self { + Self::Natural => natural_query_name_cmp(left, right), + Self::LegacyNatural => legacy_natural_query_name_cmp(left, right), + Self::Lexicographical => left.cmp(right), + } + } + + fn label(self) -> &'static str { + match self { + Self::Natural => "natural", + Self::LegacyNatural => "natural (SO-only legacy fallback)", + Self::Lexicographical => "lexicographical", + } + } +} + +fn natural_query_name_cmp(left: &[u8], right: &[u8]) -> Ordering { + natural_query_name_cmp_impl(left, right, false) +} + +fn legacy_natural_query_name_cmp(left: &[u8], right: &[u8]) -> Ordering { + natural_query_name_cmp_impl(left, right, true) +} + +fn natural_query_name_cmp_impl( + left: &[u8], + right: &[u8], + tie_break_leading_zeros: bool, +) -> Ordering { + let mut left_idx = 0usize; + let mut right_idx = 0usize; + while left_idx < left.len() && right_idx < right.len() { + let left_byte = left[left_idx]; + let right_byte = right[right_idx]; + if left_byte.is_ascii_digit() && right_byte.is_ascii_digit() { + let left_end = left[left_idx..] + .iter() + .position(|byte| !byte.is_ascii_digit()) + .map_or(left.len(), |offset| left_idx + offset); + let right_end = right[right_idx..] + .iter() + .position(|byte| !byte.is_ascii_digit()) + .map_or(right.len(), |offset| right_idx + offset); + let left_digits = &left[left_idx..left_end]; + let right_digits = &right[right_idx..right_end]; + let left_significant = left_digits + .iter() + .position(|byte| *byte != b'0') + .map_or(&left_digits[left_digits.len()..], |idx| { + &left_digits[idx..] + }); + let right_significant = right_digits + .iter() + .position(|byte| *byte != b'0') + .map_or(&right_digits[right_digits.len()..], |idx| { + &right_digits[idx..] + }); + match left_significant.len().cmp(&right_significant.len()) { + Ordering::Equal => {} + ordering => return ordering, + } + match left_significant.cmp(right_significant) { + Ordering::Equal => {} + ordering => return ordering, + } + if tie_break_leading_zeros { + let left_zero_count = + left_digits.len() - left_significant.len(); + let right_zero_count = + right_digits.len() - right_significant.len(); + match right_zero_count.cmp(&left_zero_count) { + Ordering::Equal => {} + // Older samtools natural sorting placed more leading + // zeros before fewer leading zeros. + ordering => return ordering, + } + } + // Keep numerically equal runs comparator-equivalent. The merge + // groups these names and still matches their exact QNAME bytes. + left_idx = left_end; + right_idx = right_end; + } else { + match left_byte.cmp(&right_byte) { + Ordering::Equal => { + left_idx += 1; + right_idx += 1; + } + ordering => return ordering, + } + } + } + left.len() + .saturating_sub(left_idx) + .cmp(&right.len().saturating_sub(right_idx)) +} + +fn header_tag<'a>(line: &'a [u8], tag: &[u8; 2]) -> Option<&'a [u8]> { + line.split(|byte| *byte == b'\t').skip(1).find_map(|field| { + field.strip_prefix(tag).and_then(|value| value.strip_prefix(b":")) + }) +} + +fn query_name_order( + header: &bam::HeaderView, + input_label: &str, +) -> anyhow::Result { + let hd = header + .as_bytes() + .split(|byte| *byte == b'\n') + .find(|line| line.starts_with(b"@HD\t")) + .ok_or_else(|| anyhow!("{input_label} BAM has no @HD header record"))?; + let sort_order = header_tag(hd, b"SO") + .ok_or_else(|| anyhow!("{input_label} BAM @HD record has no SO tag"))?; + if sort_order != b"queryname" { + bail!( + "{input_label} BAM must be query-name sorted, found SO:{}", + String::from_utf8_lossy(sort_order) + ) + } + let Some(sub_sort) = header_tag(hd, b"SS") else { + warn!( + "{input_label} BAM declares SO:queryname without SS; assuming \ + natural query-name order with the legacy leading-zero tie-break" + ); + return Ok(QueryNameOrder::LegacyNatural); + }; + let mut parts = sub_sort.split(|byte| *byte == b':'); + match (parts.next(), parts.next(), parts.next()) { + (Some(b"queryname"), Some(b"natural"), None) => { + Ok(QueryNameOrder::Natural) + } + (Some(b"queryname"), Some(b"lexicographical"), None) => { + Ok(QueryNameOrder::Lexicographical) + } + _ => bail!( + "{input_label} BAM uses unsupported query-name sub-sort SS:{}", + String::from_utf8_lossy(sub_sort) + ), + } +} + #[derive(new)] struct RecordPair { donor: Arc, acceptor: bam::Record, } -struct ZipRecordsIter<'a, T: Read> { - donor_records: bam::Records<'a, T>, - acceptor_records: bam::Records<'a, T>, - cur_donor_record: Option>, +struct DonorGroup { + representative: Vec, + by_exact_name: FxHashMap, Arc>, +} + +struct ZipRecordsIter { + donor_records: D, + acceptor_records: A, + pending_donor: Option, + donor_group: Option, cur_acceptor_record: Option, + last_donor_name: Option>, + last_acceptor_name: Option>, donor_ticker: ProgressBar, acceptor_ticker: ProgressBar, + order: QueryNameOrder, + terminated: bool, } -impl<'a, T: Read> ZipRecordsIter<'a, T> { +impl ZipRecordsIter { fn new( - donor: bam::Records<'a, T>, - acceptor: bam::Records<'a, T>, + donor: D, + acceptor: A, donor_ticker: ProgressBar, acceptor_ticker: ProgressBar, + order: QueryNameOrder, ) -> Self { Self { donor_records: donor, acceptor_records: acceptor, - cur_donor_record: None, + pending_donor: None, + donor_group: None, cur_acceptor_record: None, + last_donor_name: None, + last_acceptor_name: None, donor_ticker, acceptor_ticker, + order, + terminated: false, } } -} -fn get_next_record( - records: &mut bam::Records, - donor: bool, -) -> Option { - loop { - match records.next() { - Some(Ok(record)) => { - if record_is_not_primary(&record) && donor { - continue; - } else { - break Some(record); - } - } - Some(Err(e)) => { - let label = if donor { "donor" } else { "acceptor" }; - warn!( - "failed to parse record from {label} BAM, {}", - e.to_string() - ); - continue; + fn validate_monotonic( + order: QueryNameOrder, + previous: Option<&[u8]>, + current: &[u8], + input_label: &str, + ) -> anyhow::Result<()> { + if let Some(previous) = previous { + if order.compare(previous, current) == Ordering::Greater { + bail!( + "{input_label} BAM is not {} query-name sorted: {} \ + appears after {}", + order.label(), + String::from_utf8_lossy(current), + String::from_utf8_lossy(previous) + ) } - None => break None, } + Ok(()) } -} -impl<'a, T: Read> ZipRecordsIter<'a, T> { - fn advance_donor_record(&mut self) { - match self.cur_donor_record { - Some(_) => { - return; - } - None => { - self.cur_donor_record = - get_next_record(&mut self.donor_records, true) - .map(|rec| Arc::new(rec)) + fn read_primary_donor(&mut self) -> anyhow::Result> { + loop { + let Some(record) = + get_next_record(&mut self.donor_records, "donor") + else { + return Ok(None); + }; + let name = record.qname().to_vec(); + Self::validate_monotonic( + self.order, + self.last_donor_name.as_deref(), + &name, + "donor", + )?; + self.last_donor_name = Some(name); + self.donor_ticker.inc(1); + if !record_is_not_primary(&record) { + return Ok(Some(record)); } } } - fn advance_acceptor_record(&mut self) { - match self.cur_acceptor_record { - Some(_) => { - return; - } - None => { - self.cur_acceptor_record = - get_next_record(&mut self.acceptor_records, false); + + fn load_donor_group(&mut self) -> anyhow::Result> { + let first = match self.pending_donor.take() { + Some(record) => record, + None => match self.read_primary_donor()? { + Some(record) => record, + None => return Ok(None), + }, + }; + let representative = first.qname().to_vec(); + let mut by_exact_name = FxHashMap::default(); + by_exact_name.insert(representative.clone(), Arc::new(first)); + while let Some(record) = self.read_primary_donor()? { + let name = record.qname().to_vec(); + match self.order.compare(&representative, &name) { + Ordering::Equal => { + if by_exact_name + .insert(name.clone(), Arc::new(record)) + .is_some() + { + bail!( + "donor BAM has multiple primary records for \ + query name {}", + String::from_utf8_lossy(&name) + ) + } + } + Ordering::Less => { + self.pending_donor = Some(record); + break; + } + Ordering::Greater => unreachable!( + "donor monotonicity is checked before grouping" + ), } } + Ok(Some(DonorGroup { representative, by_exact_name })) } -} -impl<'a, T: Read> Iterator for ZipRecordsIter<'a, T> { - type Item = RecordPair; + fn load_acceptor(&mut self) -> anyhow::Result { + if self.cur_acceptor_record.is_some() { + return Ok(true); + } + let Some(record) = + get_next_record(&mut self.acceptor_records, "acceptor") + else { + return Ok(false); + }; + let name = record.qname().to_vec(); + Self::validate_monotonic( + self.order, + self.last_acceptor_name.as_deref(), + &name, + "acceptor", + )?; + self.last_acceptor_name = Some(name); + self.cur_acceptor_record = Some(record); + Ok(true) + } - fn next(&mut self) -> Option { - // advance a to next record + fn consume_acceptor(&mut self) -> bam::Record { + self.acceptor_ticker.inc(1); + self.cur_acceptor_record + .take() + .expect("acceptor is loaded before it is consumed") + } + + fn next_job(&mut self) -> anyhow::Result> { + if !self.load_acceptor()? { + debug!("exhausted acceptor BAM reader, finished."); + return Ok(None); + } + let acceptor_name = self + .cur_acceptor_record + .as_ref() + .expect("acceptor was loaded") + .qname() + .to_vec(); loop { - self.advance_donor_record(); - self.advance_acceptor_record(); - - return match ( - self.cur_donor_record.as_ref(), - self.cur_acceptor_record.as_ref(), - ) { - (Some(donor), Some(acceptor)) => { - match donor.qname().eq(acceptor.qname()) { - true => { - // unwrap are safe because of the above match, - // advances acceptor on next - // call to .next - let acceptor_record = std::mem::replace( - &mut self.cur_acceptor_record, - None, - ) - .unwrap(); - self.acceptor_ticker.inc(1); - return Some(RecordPair::new( - donor.clone(), - acceptor_record, - )); - } - false => { - // advance donor record in attempt to find this - // acceptor - // todo consider logging? - let _ = std::mem::replace( - &mut self.cur_donor_record, - None, - ); - self.donor_ticker.inc(1); - continue; - } - } + if self.donor_group.is_none() { + self.donor_group = self.load_donor_group()?; + } + let Some(group) = self.donor_group.as_ref() else { + self.consume_acceptor(); + return Ok(Some(RepairJob::MissingDonor { + read_name: acceptor_name, + })); + }; + match self.order.compare(&group.representative, &acceptor_name) { + Ordering::Less => { + self.donor_group = None; } - (None, Some(_)) => { - // no more donors, but still some acceptors.. error case - error!("ran out of donor records"); - None + Ordering::Greater => { + self.consume_acceptor(); + return Ok(Some(RepairJob::MissingDonor { + read_name: acceptor_name, + })); } - (Some(_), None) => { - debug!("exhausted acceptor BAM reader, finished."); - None + Ordering::Equal => { + let donor = + group.by_exact_name.get(&acceptor_name).cloned(); + let acceptor = self.consume_acceptor(); + return Ok(Some(match donor { + Some(donor) => { + RepairJob::Matched(RecordPair::new(donor, acceptor)) + } + None => { + RepairJob::MissingDonor { read_name: acceptor_name } + } + })); } - (None, None) => None, - }; + } + } + } +} + +fn get_next_record( + records: &mut T, + input_label: &str, +) -> Option { + loop { + let mut record = bam::Record::new(); + match records.read(&mut record) { + Some(Ok(())) => return Some(record), + Some(Err(error)) => { + // Commit 2 makes malformed input fatal and contextual. Keep + // the current skip-and-warn policy in this ordering-only unit. + warn!("failed to parse record from {input_label} BAM, {error}"); + } + None => return None, + } + } +} + +impl Iterator for ZipRecordsIter { + type Item = anyhow::Result; + + fn next(&mut self) -> Option { + if self.terminated { + return None; + } + match self.next_job() { + Ok(Some(job)) => Some(Ok(job)), + Ok(None) => { + self.terminated = true; + None + } + Err(error) => { + self.terminated = true; + Some(Err(error)) + } } } } @@ -392,3 +706,218 @@ fn repair_record_pair(record_pair: RecordPair) -> anyhow::Result { Ok(repaired_record) } } +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + use std::thread; + use std::time::Duration; + + use crossbeam_channel::{bounded, Receiver, Sender}; + use rust_htslib::bam; + + use crate::ordered_scheduler::OrderedWorker; + + use super::{ + query_name_order, run_repair_scheduler, QueryNameOrder, RecordPair, + RepairJob, RepairSlot, + }; + + fn make_record(name: &[u8], sequence: &str) -> bam::Record { + let mut record = bam::Record::new(); + record.set(name, None, sequence.as_bytes(), &vec![255; sequence.len()]); + record + } + + struct GatedRepairWorker { + job_one_completed: Sender<()>, + release_job_zero: Receiver<()>, + completions: Arc>>, + } + + impl OrderedWorker for GatedRepairWorker { + fn process( + &mut self, + job: RepairJob, + mut slot: RepairSlot, + ) -> anyhow::Result { + let RepairJob::Matched(record_pair) = job else { + panic!("ordering fixture only uses matched jobs") + }; + let job = match record_pair.acceptor.qname() { + b"job0" => 0, + b"job1" => 1, + name => panic!("unexpected job name: {:?}", name), + }; + if job == 0 { + self.release_job_zero + .recv_timeout(Duration::from_secs(2)) + .expect("job 0 was not released"); + } + self.completions.lock().unwrap().push(job); + if job == 1 { + let _ = self.job_one_completed.send(()); + } + slot.outcome = Some(Ok(record_pair.acceptor)); + Ok(slot) + } + } + + fn matched_job(name: &[u8]) -> anyhow::Result { + Ok(RepairJob::Matched(RecordPair::new( + Arc::new(make_record(name, "ACG")), + make_record(name, "ACG"), + ))) + } + + #[test] + fn repair_scheduler_preserves_acceptor_order_after_reordered_completion() { + let completions = Arc::new(Mutex::new(Vec::new())); + let consumed = Arc::new(Mutex::new(Vec::new())); + let (job_one_completed, wait_for_job_one) = bounded(1); + let (release_job_zero, wait_for_job_zero_release) = bounded(1); + let workers = (0..2) + .map(|_| GatedRepairWorker { + job_one_completed: job_one_completed.clone(), + release_job_zero: wait_for_job_zero_release.clone(), + completions: completions.clone(), + }) + .collect(); + let consumed_probe = consumed.clone(); + let (finished, wait_for_finish) = bounded(1); + let handle = thread::spawn(move || { + let result = run_repair_scheduler( + vec![matched_job(b"job0"), matched_job(b"job1")].into_iter(), + workers, + 0, + move |record| { + consumed_probe + .lock() + .unwrap() + .push(record?.qname().to_vec()); + Ok(()) + }, + ); + let _ = finished.send(result); + }); + + wait_for_job_one + .recv_timeout(Duration::from_secs(2)) + .expect("job 1 did not complete while job 0 was blocked"); + assert_eq!(*completions.lock().unwrap(), vec![1]); + assert!( + consumed.lock().unwrap().is_empty(), + "later repair was consumed before sequence 0" + ); + release_job_zero.send(()).unwrap(); + wait_for_finish + .recv_timeout(Duration::from_secs(2)) + .expect("repair scheduler watchdog expired") + .unwrap(); + handle.join().expect("repair scheduler thread panicked"); + + assert_eq!(*completions.lock().unwrap(), vec![1, 0]); + assert_eq!( + *consumed.lock().unwrap(), + vec![b"job0".to_vec(), b"job1".to_vec()] + ); + } + + #[test] + fn query_name_comparators_cover_numeric_and_equivalent_names() { + assert_eq!( + QueryNameOrder::Natural + .common_merge_order(QueryNameOrder::LegacyNatural), + None + ); + assert_eq!( + QueryNameOrder::LegacyNatural + .common_merge_order(QueryNameOrder::Natural), + None + ); + assert_eq!( + QueryNameOrder::LegacyNatural + .common_merge_order(QueryNameOrder::Lexicographical), + None + ); + assert_eq!( + QueryNameOrder::Lexicographical + .common_merge_order(QueryNameOrder::LegacyNatural), + None + ); + assert_eq!( + QueryNameOrder::Natural + .common_merge_order(QueryNameOrder::Lexicographical), + None + ); + assert_eq!( + QueryNameOrder::Lexicographical + .common_merge_order(QueryNameOrder::Natural), + None + ); + assert_eq!( + QueryNameOrder::Natural.compare(b"read1", b"read2"), + std::cmp::Ordering::Less + ); + assert_eq!( + QueryNameOrder::Natural.compare(b"read2", b"read10"), + std::cmp::Ordering::Less + ); + assert_eq!( + QueryNameOrder::Natural.compare(b"read1", b"read01"), + std::cmp::Ordering::Equal + ); + assert_eq!( + QueryNameOrder::LegacyNatural.compare(b"read01", b"read1"), + std::cmp::Ordering::Less + ); + assert_eq!( + QueryNameOrder::Natural.compare( + b"read999999999999999999999999", + b"read1000000000000000000000000", + ), + std::cmp::Ordering::Less + ); + assert_eq!( + QueryNameOrder::Lexicographical.compare(b"read10", b"read2"), + std::cmp::Ordering::Less + ); + } + + #[test] + fn query_name_sort_is_read_from_header_and_must_be_supported() { + let natural = bam::HeaderView::from_bytes( + b"@HD\tVN:1.6\tSO:queryname\tSS:queryname:natural\n", + ); + let lexical = bam::HeaderView::from_bytes( + b"@HD\tVN:1.6\tSO:queryname\tSS:queryname:lexicographical\n", + ); + let fallback = + bam::HeaderView::from_bytes(b"@HD\tVN:1.5\tSO:queryname\n"); + let coordinate = + bam::HeaderView::from_bytes(b"@HD\tVN:1.6\tSO:coordinate\n"); + let malformed = bam::HeaderView::from_bytes( + b"@HD\tVN:1.6\tSO:queryname\tSS:queryname:natural:extra\n", + ); + + assert_eq!( + query_name_order(&natural, "test").unwrap(), + QueryNameOrder::Natural + ); + assert_eq!( + query_name_order(&lexical, "test").unwrap(), + QueryNameOrder::Lexicographical + ); + assert_eq!( + query_name_order(&fallback, "test").unwrap(), + QueryNameOrder::LegacyNatural + ); + assert!(query_name_order(&coordinate, "test") + .unwrap_err() + .to_string() + .contains("must be query-name sorted")); + assert!(query_name_order(&malformed, "test") + .unwrap_err() + .to_string() + .contains("unsupported query-name sub-sort")); + } +} diff --git a/modkit/tests/test_repair.rs b/modkit/tests/test_repair.rs index 6fb548a1..7d259c1d 100644 --- a/modkit/tests/test_repair.rs +++ b/modkit/tests/test_repair.rs @@ -1,7 +1,9 @@ -use std::collections::HashMap; +use std::path::Path; use mod_kit::mod_bam::MN_TAG; use rust_htslib::bam; +use rust_htslib::bam::header::HeaderRecord; +use rust_htslib::bam::record::{Aux, AuxArray}; use rust_htslib::bam::Read; use crate::common::{run_modkit, run_simple_summary}; @@ -15,7 +17,8 @@ fn test_repair_help() { #[test] fn test_repair_regression() { - let out_bam = std::env::temp_dir().join("test_repair_regression.bam"); + let temp_dir = tempfile::tempdir().unwrap(); + let out_bam = temp_dir.path().join("test_repair_regression.bam"); let donor_fp = "../tests/resources/donor_read_sort.bam"; let acceptor_fp = "../tests/resources/trimmed_read_sort.mapped.bam"; run_modkit(&[ @@ -36,41 +39,327 @@ fn test_repair_regression() { "../tests/resources/trimmed_read_sort_mods.mapped.bam", ) .unwrap(); - let test_records = test_bam + let mut test_records = test_bam .records() .map(|r| r.unwrap()) .map(|mut record| { - let qname = - record.qname().iter().map(|&b| b as char).collect::(); // todo consider removing this later, added MN tag by default but // the old test data doesn't have it. record.remove_aux(MN_TAG.as_bytes()).expect("should remove MN tag"); - (qname, record) + record }) - .collect::>(); + .collect::>(); - let expected_records = ref_bam + let mut expected_records = + ref_bam.records().map(|r| r.unwrap()).collect::>(); + + // The hand-checked legacy reference predates ordered repair output. Sort + // both complete vectors for the content oracle; the synthetic regression + // below independently checks exact acceptor-relative output order. + test_records.sort_by(|left, right| left.qname().cmp(right.qname())); + expected_records.sort_by(|left, right| left.qname().cmp(right.qname())); + assert_eq!( + test_records.len(), + expected_records.len(), + "repair output cardinality changed" + ); + for (index, (actual, expected)) in + test_records.iter().zip(expected_records.iter()).enumerate() + { + assert_eq!(actual.qname(), expected.qname(), "record {index} missing"); + assert_eq!(actual, expected, "record {index} not the same"); + } +} + +fn query_name_header(sub_sort: Option<&str>) -> bam::Header { + let mut header = bam::Header::new(); + let mut hd = HeaderRecord::new(b"HD"); + hd.push_tag(b"VN", "1.6").push_tag(b"SO", "queryname"); + if let Some(sub_sort) = sub_sort { + hd.push_tag(b"SS", sub_sort); + } + header.push_record(&hd); + header +} + +fn write_repair_input( + path: &Path, + records: &[(&str, u32)], + donor_probabilities: Option<&[u8]>, + sub_sort: Option<&str>, +) { + if let Some(probabilities) = donor_probabilities { + assert_eq!(probabilities.len(), records.len()); + } + let header = query_name_header(sub_sort); + let mut writer = + bam::Writer::from_path(path, &header, bam::Format::Bam).unwrap(); + for (index, (name, marker)) in records.iter().enumerate() { + let sequence = + if donor_probabilities.is_some() { "TACGT" } else { "ACG" }; + let mut record = bam::Record::new(); + record.set( + name.as_bytes(), + None, + sequence.as_bytes(), + &vec![255; sequence.len()], + ); + record.push_aux(b"XI", Aux::U32(*marker)).unwrap(); + if let Some(probabilities) = donor_probabilities { + record.push_aux(b"MM", Aux::String("A+a.,0;")).unwrap(); + let probability = [probabilities[index]]; + let probability_array: AuxArray = (&probability[..]).into(); + record.push_aux(b"ML", Aux::ArrayU8(probability_array)).unwrap(); + } + writer.write(&record).unwrap(); + } +} + +fn repaired_identities(path: &Path) -> Vec<(String, u32, u8)> { + let mut reader = bam::Reader::from_path(path).unwrap(); + reader .records() - .map(|r| r.unwrap()) .map(|record| { - let qname = - record.qname().iter().map(|&b| b as char).collect::(); - (qname, record) + let record = record.unwrap(); + let marker = match record.aux(b"XI").unwrap() { + Aux::U32(marker) => marker, + other => panic!("unexpected marker tag: {other:?}"), + }; + let probability = match record.aux(b"ML").unwrap() { + Aux::ArrayU8(probabilities) => { + let probabilities = + probabilities.iter().collect::>(); + assert_eq!(probabilities.len(), 1); + probabilities[0] + } + other => panic!("unexpected ML tag: {other:?}"), + }; + ( + String::from_utf8_lossy(record.qname()).into_owned(), + marker, + probability, + ) }) - .collect::>(); + .collect() +} + +fn repair_error(donor: &Path, acceptor: &Path, output: &Path) -> String { + let executable = Path::new(env!("CARGO_BIN_EXE_modkit")); + let result = std::process::Command::new(executable) + .args([ + "repair", + "--donor", + donor.to_str().unwrap(), + "--acceptor", + acceptor.to_str().unwrap(), + "--output", + output.to_str().unwrap(), + "--threads", + "2", + ]) + .output() + .unwrap(); + assert!(!result.status.success(), "repair unexpectedly succeeded"); + String::from_utf8_lossy(&result.stderr).into_owned() +} + +#[test] +fn test_repair_preserves_acceptor_order_and_exact_cardinality() { + let temp_dir = tempfile::tempdir().unwrap(); + let donor = temp_dir.path().join("donor.bam"); + let acceptor = temp_dir.path().join("acceptor.bam"); + let output = temp_dir.path().join("repaired.bam"); + write_repair_input( + &donor, + &[ + ("read01", 90), + ("read1", 91), + ("read001", 92), + ("read3", 93), + ("read10", 100), + ], + Some(&[22, 11, 33, 44, 55]), + Some("queryname:natural"), + ); + write_repair_input( + &acceptor, + &[ + ("read1", 0), + ("read1", 1), + ("read01", 2), + ("read001", 7), + ("read0001", 8), + ("read2", 3), + ("read3", 4), + ("read3", 5), + ("read10", 6), + ], + None, + Some("queryname:natural"), + ); - for (q, r) in test_records.iter() { - assert_eq!( - expected_records.get(q).unwrap(), - r, - "record {q} not the same" + run_modkit(&[ + "repair", + "--donor", + donor.to_str().unwrap(), + "--acceptor", + acceptor.to_str().unwrap(), + "--output", + output.to_str().unwrap(), + "--threads", + "2", + ]) + .unwrap(); + + let reader = bam::Reader::from_path(&output).unwrap(); + let header = String::from_utf8_lossy(reader.header().as_bytes()); + assert!(header.contains("SO:queryname")); + assert!(header.contains("SS:queryname:natural")); + drop(reader); + let actual = repaired_identities(&output); + assert_eq!( + actual, + vec![ + ("read1".to_string(), 0, 11), + ("read1".to_string(), 1, 11), + ("read01".to_string(), 2, 22), + ("read001".to_string(), 7, 33), + ("read3".to_string(), 4, 44), + ("read3".to_string(), 5, 44), + ("read10".to_string(), 6, 55), + ] + ); +} + +#[test] +fn test_repair_rejects_mixed_natural_headers_in_both_directions() { + let temp_dir = tempfile::tempdir().unwrap(); + for (index, donor_uses_explicit_natural) in + [true, false].into_iter().enumerate() + { + let donor = temp_dir.path().join(format!("donor-{index}.bam")); + let acceptor = temp_dir.path().join(format!("acceptor-{index}.bam")); + let output = temp_dir.path().join(format!("repaired-{index}.bam")); + if donor_uses_explicit_natural { + write_repair_input( + &donor, + &[("read1a1", 1), ("read01a2", 2)], + Some(&[11, 22]), + Some("queryname:natural"), + ); + write_repair_input( + &acceptor, + &[("read01a2", 2), ("read1a1", 1)], + None, + None, + ); + } else { + write_repair_input( + &donor, + &[("read01a2", 2), ("read1a1", 1)], + Some(&[22, 11]), + None, + ); + write_repair_input( + &acceptor, + &[("read1a1", 1), ("read01a2", 2)], + None, + Some("queryname:natural"), + ); + } + + let expected_diagnostic = if donor_uses_explicit_natural { + "donor and acceptor BAMs use incompatible query-name sorting: \ + donor is natural, acceptor is natural (SO-only legacy fallback)" + } else { + "donor and acceptor BAMs use incompatible query-name sorting: \ + donor is natural (SO-only legacy fallback), acceptor is natural" + }; + let error = repair_error(&donor, &acceptor, &output); + assert!( + error.contains(expected_diagnostic), + "unexpected mixed-sort error: {error}" ); } } +#[test] +fn test_repair_rejects_ambiguous_or_incompatibly_sorted_inputs() { + let temp_dir = tempfile::tempdir().unwrap(); + let donor = temp_dir.path().join("donor.bam"); + let acceptor = temp_dir.path().join("acceptor.bam"); + let output = temp_dir.path().join("repaired.bam"); + write_repair_input( + &donor, + &[("read1", 0), ("read1", 1)], + Some(&[100, 101]), + Some("queryname:natural"), + ); + write_repair_input( + &acceptor, + &[("read1", 0)], + None, + Some("queryname:natural"), + ); + let error = repair_error(&donor, &acceptor, &output); + assert!( + error.contains("multiple primary records for query name read1"), + "unexpected duplicate-donor error: {error}" + ); + + write_repair_input( + &donor, + &[("read2", 0), ("read1", 1)], + Some(&[100, 101]), + Some("queryname:natural"), + ); + let error = repair_error(&donor, &acceptor, &output); + assert!( + error.contains("donor BAM is not natural query-name sorted"), + "unexpected donor-inversion error: {error}" + ); + + write_repair_input( + &donor, + &[("read1", 0), ("read2", 1)], + Some(&[100, 101]), + Some("queryname:natural"), + ); + write_repair_input( + &acceptor, + &[("read2", 0), ("read1", 1)], + None, + Some("queryname:natural"), + ); + let error = repair_error(&donor, &acceptor, &output); + assert!( + error.contains("acceptor BAM is not natural query-name sorted"), + "unexpected acceptor-inversion error: {error}" + ); + + write_repair_input( + &donor, + &[("read1", 0)], + Some(&[100]), + Some("queryname:lexicographical"), + ); + write_repair_input( + &acceptor, + &[("read1", 0)], + None, + Some("queryname:natural"), + ); + let error = repair_error(&donor, &acceptor, &output); + assert!( + error.contains("incompatible query-name sorting"), + "unexpected incompatible-sort error: {error}" + ); +} + #[test] fn test_repair_mn_tag() { - let out_bam = std::env::temp_dir().join("test_repair_mn_tag.bam"); + let temp_dir = tempfile::tempdir().unwrap(); + let out_bam = temp_dir.path().join("test_repair_mn_tag.bam"); let donor_fp = "../tests/resources/donor_read_sort_mn_tag.bam"; let acceptor_fp = "../tests/resources/trimmed_read_sort_mn_tag.mapped.bam"; run_modkit(&[ From d0b47e207182b8cd14a1fe3fe764c8b60c376d9f Mon Sep 17 00:00:00 2001 From: SuhasSrinivasan <32346517+SuhasSrinivasan@users.noreply.github.com> Date: Sun, 9 Aug 2026 19:24:02 -0700 Subject: [PATCH 3/4] Make repair output atomic on fatal errors --- modkit-core/Cargo.toml | 2 +- modkit-core/src/repair_tags.rs | 1011 ++++++++++++++++++++++++++++++-- 2 files changed, 974 insertions(+), 39 deletions(-) diff --git a/modkit-core/Cargo.toml b/modkit-core/Cargo.toml index 881c7d87..a75e67a5 100644 --- a/modkit-core/Cargo.toml +++ b/modkit-core/Cargo.toml @@ -49,6 +49,7 @@ rustc-hash = "1.1.0" rv = "=0.16.0" statrs = "0.16.0" substring = "1.4.5" +tempfile = "3.2" thiserror = "2.0.11" tokio = "1.42.0" tracing = "0.1.41" @@ -59,5 +60,4 @@ serde_json = "1.0.149" [dev-dependencies] assert_approx_eq = "1.1.0" similar-asserts = "1.4.2" -tempfile = "3.2" serde = {version = "1.0.219", features = ["derive"]} diff --git a/modkit-core/src/repair_tags.rs b/modkit-core/src/repair_tags.rs index ded5201a..e8fbfe00 100644 --- a/modkit-core/src/repair_tags.rs +++ b/modkit-core/src/repair_tags.rs @@ -1,16 +1,21 @@ use std::cmp::Ordering; -use std::path::PathBuf; +use std::fs::{self, File}; +use std::io::{Read as IoRead, Seek, SeekFrom}; +#[cfg(unix)] +use std::os::unix::fs::{MetadataExt, PermissionsExt}; +use std::path::{Path, PathBuf}; use std::sync::Arc; -use anyhow::{anyhow, bail}; +use anyhow::{anyhow, bail, Context}; use clap::Args; use crossbeam_channel::unbounded; use derive_new::new; use indicatif::{MultiProgress, ProgressBar}; -use log::{debug, error, info, warn}; +use log::{debug, info, warn}; use rust_htslib::bam::record::{Aux, AuxArray}; use rust_htslib::bam::{self, Read}; use rustc_hash::FxHashMap; +use tempfile::{Builder as TempFileBuilder, NamedTempFile, TempDir}; use modkit_logging::init_logging; @@ -46,10 +51,452 @@ pub struct RepairTags { threads: usize, } +const REPAIR_TEMP_PREFIX: &str = ".modkit-repair-"; +const BGZF_EOF_BLOCK: [u8; 28] = [ + 31, 139, 8, 4, 0, 0, 0, 0, 0, 255, 6, 0, 66, 67, 2, 0, 27, 0, 3, 0, 0, 0, + 0, 0, 0, 0, 0, 0, +]; + +trait RepairSink { + fn write(&mut self, record: &bam::Record) -> anyhow::Result<()>; + fn finish(self, expected_records: usize) -> anyhow::Result<()>; +} + +trait StagedBamValidator { + fn validate( + &self, + file: &mut File, + path: &Path, + expected_identity: StagedFileIdentity, + expected_records: usize, + ) -> anyhow::Result<()>; +} + +struct CompleteBamValidator; + +impl StagedBamValidator for CompleteBamValidator { + fn validate( + &self, + file: &mut File, + path: &Path, + expected_identity: StagedFileIdentity, + expected_records: usize, + ) -> anyhow::Result<()> { + validate_complete_bam_file( + file, + path, + expected_identity, + expected_records, + ) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +struct StagedFileIdentity { + #[cfg(unix)] + device: u64, + #[cfg(unix)] + inode: u64, +} + +impl StagedFileIdentity { + fn from_metadata(metadata: &fs::Metadata) -> Self { + #[cfg(unix)] + { + Self { device: metadata.dev(), inode: metadata.ino() } + } + #[cfg(not(unix))] + { + let _ = metadata; + Self {} + } + } +} + +fn require_staged_path_identity( + file: &File, + path: &Path, + expected: StagedFileIdentity, +) -> anyhow::Result<()> { + let retained_metadata = file.metadata().with_context(|| { + format!( + "failed to inspect retained staged repair BAM descriptor for {}", + path.display() + ) + })?; + if !retained_metadata.file_type().is_file() { + bail!( + "retained staged repair BAM descriptor for {} is not a regular file", + path.display() + ); + } + let path_metadata = fs::symlink_metadata(path).with_context(|| { + format!( + "staged repair BAM pathname {} no longer identifies the retained file", + path.display() + ) + })?; + if !path_metadata.file_type().is_file() { + bail!( + "staged repair BAM pathname {} no longer identifies the retained regular file", + path.display() + ); + } + #[cfg(unix)] + { + let retained_identity = + StagedFileIdentity::from_metadata(&retained_metadata); + let path_identity = StagedFileIdentity::from_metadata(&path_metadata); + if retained_identity != expected || path_identity != expected { + bail!( + "staged repair BAM pathname {} no longer identifies the retained file \ + (expected dev {} ino {}, retained dev {} ino {}, path dev {} ino {})", + path.display(), + expected.device, + expected.inode, + retained_identity.device, + retained_identity.inode, + path_identity.device, + path_identity.inode + ); + } + } + #[cfg(not(unix))] + let _ = expected; + Ok(()) +} + +struct StagedBamOutput { + // Keep the writer and temporary file before the directory so early-drop + // closes htslib, removes the staged file, then removes the private dir. + writer: Option, + temp_file: NamedTempFile, + _staging_dir: TempDir, + destination: PathBuf, + destination_permissions: Option, + staged_identity: StagedFileIdentity, + validator: V, +} + +fn existing_destination_permissions( + destination: &Path, +) -> anyhow::Result> { + match fs::symlink_metadata(destination) { + Ok(metadata) if metadata.file_type().is_symlink() => bail!( + "repair output destination {} must not be a symbolic link", + destination.display() + ), + Ok(metadata) if !metadata.file_type().is_file() => bail!( + "repair output destination {} must be a regular file when it already exists", + destination.display() + ), + Ok(metadata) => Ok(Some(metadata.permissions())), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None), + Err(error) => Err(error).with_context(|| { + format!( + "failed to inspect repair output destination {}", + destination.display() + ) + }), + } +} + +impl StagedBamOutput { + fn new(destination: &Path, header: &bam::Header) -> anyhow::Result { + Self::with_validator(destination, header, CompleteBamValidator) + } +} + +impl StagedBamOutput { + fn with_validator( + destination: &Path, + header: &bam::Header, + validator: V, + ) -> anyhow::Result { + let destination_permissions = + existing_destination_permissions(destination)?; + let destination_dir = destination + .parent() + .filter(|parent| !parent.as_os_str().is_empty()) + .unwrap_or_else(|| Path::new(".")); + + let mut staging_dir_builder = TempFileBuilder::new(); + staging_dir_builder.prefix(REPAIR_TEMP_PREFIX); + #[cfg(unix)] + staging_dir_builder.permissions(fs::Permissions::from_mode(0o700)); + let staging_dir = staging_dir_builder + .tempdir_in(destination_dir) + .with_context(|| { + format!( + "failed to create private repair staging directory in {}", + destination_dir.display() + ) + })?; + #[cfg(unix)] + { + // Builder applies the process umask. A restrictive umask may + // remove owner bits, but group/other access must never be added. + let staging_mode = fs::symlink_metadata(staging_dir.path()) + .with_context(|| { + format!( + "failed to inspect repair staging directory {}", + staging_dir.path().display() + ) + })? + .permissions() + .mode() + & 0o7777; + if staging_mode & 0o077 != 0 { + bail!( + "repair staging directory {} has non-owner mode bits {staging_mode:o}", + staging_dir.path().display() + ); + } + } + + let mut temp_file_builder = TempFileBuilder::new(); + temp_file_builder.prefix("staged-").suffix(".bam"); + #[cfg(unix)] + { + // Match an ordinary File::create rather than tempfile's 0600 + // default; Builder applies the process umask to this mode. + temp_file_builder.permissions(fs::Permissions::from_mode(0o666)); + } + let temp_file = temp_file_builder + .tempfile_in(staging_dir.path()) + .with_context(|| { + format!( + "failed to create staged repair BAM in private directory {}", + staging_dir.path().display() + ) + })?; + let staged_identity = StagedFileIdentity::from_metadata( + &temp_file.as_file().metadata().with_context(|| { + format!( + "failed to inspect retained staged repair BAM {}", + temp_file.path().display() + ) + })?, + ); + let writer = + bam::Writer::from_path(temp_file.path(), header, bam::Format::Bam) + .with_context(|| { + format!( + "failed to open staged repair BAM {}", + temp_file.path().display() + ) + })?; + require_staged_path_identity( + temp_file.as_file(), + temp_file.path(), + staged_identity, + ) + .context("staged repair BAM changed while opening its htslib writer")?; + Ok(Self { + writer: Some(writer), + temp_file, + _staging_dir: staging_dir, + destination: destination.to_path_buf(), + destination_permissions, + staged_identity, + validator, + }) + } +} + +impl RepairSink for StagedBamOutput { + fn write(&mut self, record: &bam::Record) -> anyhow::Result<()> { + self.writer + .as_mut() + .ok_or_else(|| { + anyhow!("staged repair BAM writer is already closed") + })? + .write(record) + .context("failed to write staged repair BAM record") + } + + fn finish(mut self, expected_records: usize) -> anyhow::Result<()> { + let writer = self.writer.take().ok_or_else(|| { + anyhow!("staged repair BAM writer is already closed") + })?; + // rust-htslib 0.46 does not expose hts_close's return status. Drop the + // writer, then independently require a readable BAM, canonical EOF, + // and the expected record count before making it visible. + drop(writer); + let staged_path = self.temp_file.path().to_path_buf(); + require_staged_path_identity( + self.temp_file.as_file(), + &staged_path, + self.staged_identity, + ) + .context("staged repair BAM changed after closing its htslib writer")?; + self.validator + .validate( + self.temp_file.as_file_mut(), + &staged_path, + self.staged_identity, + expected_records, + ) + .context("failed to validate staged repair BAM")?; + require_staged_path_identity( + self.temp_file.as_file(), + &staged_path, + self.staged_identity, + ) + .context("staged repair BAM changed during validation")?; + if let Some(permissions) = self.destination_permissions.take() { + // Restore after all writes because a write may clear special mode + // bits. Any failure still precedes the atomic replacement. + self.temp_file + .as_file() + .set_permissions(permissions) + .with_context(|| { + format!( + "failed to preserve permissions for existing repair output {}", + self.destination.display() + ) + })?; + } + self.temp_file.as_file().sync_all().with_context(|| { + format!( + "failed to sync staged repair BAM {}", + staged_path.display() + ) + })?; + // NamedTempFile::persist still names the source path. These checks + // reject observed substitutions, although a hostile writer of a + // non-sticky parent directory can still race the check and rename. + require_staged_path_identity( + self.temp_file.as_file(), + &staged_path, + self.staged_identity, + ) + .context("staged repair BAM changed before atomic replacement")?; + self.temp_file.persist(&self.destination).map_err(|error| { + anyhow!( + "failed to atomically replace repair output {}: {error}", + self.destination.display() + ) + })?; + Ok(()) + } +} + +#[cfg(test)] +fn validate_complete_bam( + path: &Path, + expected_records: usize, +) -> anyhow::Result<()> { + let mut file = File::open(path).with_context(|| { + format!("failed to open staged repair BAM {}", path.display()) + })?; + let expected_identity = StagedFileIdentity::from_metadata( + &file.metadata().with_context(|| { + format!("failed to inspect staged repair BAM {}", path.display()) + })?, + ); + validate_complete_bam_file( + &mut file, + path, + expected_identity, + expected_records, + ) +} + +fn validate_complete_bam_file( + file: &mut File, + path: &Path, + expected_identity: StagedFileIdentity, + expected_records: usize, +) -> anyhow::Result<()> { + let file_len = file + .metadata() + .with_context(|| { + format!("failed to inspect staged repair BAM {}", path.display()) + })? + .len(); + if file_len < BGZF_EOF_BLOCK.len() as u64 { + bail!( + "staged repair BAM {} is too short to contain a BGZF EOF block", + path.display() + ) + } + file.seek(SeekFrom::End(-(BGZF_EOF_BLOCK.len() as i64))).with_context( + || format!("failed to seek staged repair BAM {}", path.display()), + )?; + let mut eof = [0u8; BGZF_EOF_BLOCK.len()]; + file.read_exact(&mut eof).with_context(|| { + format!("failed to read staged repair BAM EOF at {}", path.display()) + })?; + if eof != BGZF_EOF_BLOCK { + bail!( + "staged repair BAM {} is missing the canonical BGZF EOF block", + path.display() + ) + } + + require_staged_path_identity(file, path, expected_identity) + .context("staged repair BAM changed before full BAM decoding")?; + let mut reader = bam::Reader::from_path(path).with_context(|| { + format!("failed to reopen staged repair BAM {}", path.display()) + })?; + let mut observed_records = 0usize; + for (index, record) in reader.records().enumerate() { + record.with_context(|| { + format!( + "failed to decode staged repair BAM record {} at {}", + index + 1, + path.display() + ) + })?; + observed_records += 1; + } + drop(reader); + require_staged_path_identity(file, path, expected_identity) + .context("staged repair BAM changed during full BAM decoding")?; + if observed_records != expected_records { + bail!( + "staged repair BAM {} contains {observed_records} records, \ + expected {expected_records}", + path.display() + ) + } + Ok(()) +} + +fn open_bam_reader( + path: &Path, + input_label: &str, + threads: usize, +) -> anyhow::Result { + let mut reader = bam::Reader::from_path(path).with_context(|| { + format!("failed to open {input_label} BAM {}", path.display()) + })?; + reader.set_threads(threads).with_context(|| { + format!( + "failed to configure {input_label} BAM reader {}", + path.display() + ) + })?; + Ok(reader) +} + impl RepairTags { pub fn run(&self) -> anyhow::Result<()> { let _handle = init_logging(self.log_filepath.as_ref()); + self.run_with_output_factory(|header| { + StagedBamOutput::new(&self.output_bam, header) + }) + } + fn run_with_output_factory( + &self, + make_output: F, + ) -> anyhow::Result<()> + where + S: RepairSink, + F: FnOnce(&bam::Header) -> anyhow::Result, + { let reader_threads = { let half = self.threads / 2; std::cmp::min(half, 16) @@ -61,10 +508,13 @@ impl RepairTags { {pool_threads} to process records" ); - let mut donor_records = bam::Reader::from_path(&self.donor_bam)?; - donor_records.set_threads(threads_per_reader)?; - let mut acceptor_records = bam::Reader::from_path(&self.acceptor_bam)?; - acceptor_records.set_threads(threads_per_reader)?; + let donor_records = + open_bam_reader(&self.donor_bam, "donor", threads_per_reader)?; + let acceptor_records = open_bam_reader( + &self.acceptor_bam, + "acceptor", + threads_per_reader, + )?; let donor_order = query_name_order(donor_records.header(), "donor")?; let acceptor_order = query_name_order(acceptor_records.header(), "acceptor")?; @@ -78,11 +528,7 @@ impl RepairTags { ), }; let header = bam::Header::from_template(acceptor_records.header()); - let mut writer = bam::Writer::from_path( - &self.output_bam, - &header, - bam::Format::Bam, - )?; + let mut output = make_output(&header)?; info!( "repairing records in {} with base modification information in {}", &self.acceptor_bam.to_str().unwrap_or_else(|| "??"), @@ -115,13 +561,14 @@ impl RepairTags { repaired_ticker.inc(1); match res { Ok(record) => { - if let Err(e) = writer.write(&record) { - error!("failed to write record {}", e.to_string()); - n_failed += 1; - } else { - written_ticker.inc(1); - n_repaired += 1; - } + output.write(&record).with_context(|| { + format!( + "failed to write repaired record {}", + String::from_utf8_lossy(record.qname()) + ) + })?; + written_ticker.inc(1); + n_repaired += 1; } Err(e) => { debug!("record failed to be repaired: {}", e.to_string()); @@ -131,6 +578,7 @@ impl RepairTags { Ok(()) })?; + output.finish(n_repaired)?; info!("finished, repaired {n_repaired} records, {n_failed} failed."); Ok(()) } @@ -387,6 +835,8 @@ struct ZipRecordsIter { cur_acceptor_record: Option, last_donor_name: Option>, last_acceptor_name: Option>, + donor_records_read: usize, + acceptor_records_read: usize, donor_ticker: ProgressBar, acceptor_ticker: ProgressBar, order: QueryNameOrder, @@ -409,6 +859,8 @@ impl ZipRecordsIter { cur_acceptor_record: None, last_donor_name: None, last_acceptor_name: None, + donor_records_read: 0, + acceptor_records_read: 0, donor_ticker, acceptor_ticker, order, @@ -438,8 +890,11 @@ impl ZipRecordsIter { fn read_primary_donor(&mut self) -> anyhow::Result> { loop { - let Some(record) = - get_next_record(&mut self.donor_records, "donor") + let Some(record) = get_next_record( + &mut self.donor_records, + "donor", + &mut self.donor_records_read, + )? else { return Ok(None); }; @@ -500,8 +955,11 @@ impl ZipRecordsIter { if self.cur_acceptor_record.is_some() { return Ok(true); } - let Some(record) = - get_next_record(&mut self.acceptor_records, "acceptor") + let Some(record) = get_next_record( + &mut self.acceptor_records, + "acceptor", + &mut self.acceptor_records_read, + )? else { return Ok(false); }; @@ -576,18 +1034,19 @@ impl ZipRecordsIter { fn get_next_record( records: &mut T, input_label: &str, -) -> Option { - loop { - let mut record = bam::Record::new(); - match records.read(&mut record) { - Some(Ok(())) => return Some(record), - Some(Err(error)) => { - // Commit 2 makes malformed input fatal and contextual. Keep - // the current skip-and-warn policy in this ordering-only unit. - warn!("failed to parse record from {input_label} BAM, {error}"); - } - None => return None, + records_read: &mut usize, +) -> anyhow::Result> { + let mut record = bam::Record::new(); + match records.read(&mut record) { + Some(Ok(())) => { + *records_read += 1; + Ok(Some(record)) } + Some(Err(error)) => bail!( + "failed to decode {input_label} BAM record {}: {error}", + *records_read + 1 + ), + None => Ok(None), } } @@ -706,20 +1165,30 @@ fn repair_record_pair(record_pair: RecordPair) -> anyhow::Result { Ok(repaired_record) } } + #[cfg(test)] mod tests { + use std::env; + use std::fs::{self, File, OpenOptions}; + #[cfg(unix)] + use std::os::unix::fs::{symlink, FileTypeExt, PermissionsExt}; + use std::path::{Path, PathBuf}; + use std::process::{Command, Stdio}; use std::sync::{Arc, Mutex}; use std::thread; - use std::time::Duration; + use std::time::{Duration, Instant}; use crossbeam_channel::{bounded, Receiver, Sender}; - use rust_htslib::bam; + use rust_htslib::bam::header::HeaderRecord; + use rust_htslib::bam::{self, record::Aux}; use crate::ordered_scheduler::OrderedWorker; use super::{ - query_name_order, run_repair_scheduler, QueryNameOrder, RecordPair, - RepairJob, RepairSlot, + query_name_order, run_repair_scheduler, validate_complete_bam, + CompleteBamValidator, QueryNameOrder, RecordPair, RepairJob, + RepairSink, RepairSlot, RepairTags, StagedBamOutput, + StagedBamValidator, StagedFileIdentity, REPAIR_TEMP_PREFIX, }; fn make_record(name: &[u8], sequence: &str) -> bam::Record { @@ -728,6 +1197,472 @@ mod tests { record } + struct FailNthWriteSink { + inner: S, + fail_at: usize, + writes: usize, + } + + impl RepairSink for FailNthWriteSink { + fn write(&mut self, record: &bam::Record) -> anyhow::Result<()> { + self.writes += 1; + if self.writes == self.fail_at { + anyhow::bail!( + "injected repair sink write failure at record {}", + self.fail_at + ) + } + self.inner.write(record) + } + + fn finish(self, expected_records: usize) -> anyhow::Result<()> { + self.inner.finish(expected_records) + } + } + + struct CorruptBeforeValidation; + + impl StagedBamValidator for CorruptBeforeValidation { + fn validate( + &self, + file: &mut File, + path: &Path, + expected_identity: StagedFileIdentity, + expected_records: usize, + ) -> anyhow::Result<()> { + let file_len = file.metadata()?.len(); + file.set_len(file_len.checked_sub(1).expect("BAM is non-empty"))?; + CompleteBamValidator.validate( + file, + path, + expected_identity, + expected_records, + ) + } + } + + #[cfg(unix)] + struct SubstituteBeforeDecode; + + #[cfg(unix)] + impl StagedBamValidator for SubstituteBeforeDecode { + fn validate( + &self, + file: &mut File, + path: &Path, + expected_identity: StagedFileIdentity, + expected_records: usize, + ) -> anyhow::Result<()> { + let replacement = path.with_file_name("replacement.bam"); + let replacement_writer = bam::Writer::from_path( + &replacement, + &bam::Header::new(), + bam::Format::Bam, + )?; + drop(replacement_writer); + fs::rename(&replacement, path)?; + CompleteBamValidator.validate( + file, + path, + expected_identity, + expected_records, + ) + } + } + + const REPAIR_CHILD_CASE_ENV: &str = "MODKIT_REPAIR_CHILD_CASE"; + const REPAIR_CHILD_DIR_ENV: &str = "MODKIT_REPAIR_CHILD_DIR"; + const REPAIR_CHILD_TEST: &str = + "repair_tags::tests::repair_lifecycle_child"; + + fn write_lifecycle_bam(path: &Path, records: usize, donor: bool) { + let mut header = bam::Header::new(); + let mut hd = HeaderRecord::new(b"HD"); + hd.push_tag(b"VN", "1.6") + .push_tag(b"SO", "queryname") + .push_tag(b"SS", "queryname:natural"); + header.push_record(&hd); + let mut writer = + bam::Writer::from_path(path, &header, bam::Format::Bam).unwrap(); + for index in 0..records { + let name = format!("read{index:06}"); + let sequence = if donor { "TACGT" } else { "ACG" }; + let mut record = make_record(name.as_bytes(), sequence); + if donor { + record.push_aux(b"MM", Aux::String("A+a.,0;")).unwrap(); + let probability = [((index % 200) + 1) as u8]; + record + .push_aux(b"ML", Aux::ArrayU8((&probability[..]).into())) + .unwrap(); + } + writer.write(&record).unwrap(); + } + drop(writer); + validate_complete_bam(path, records).unwrap(); + } + + fn truncate_bam_data(path: &Path) { + let file = OpenOptions::new().write(true).open(path).unwrap(); + let file_len = file.metadata().unwrap().len(); + assert!(file_len > 40); + // Remove the canonical 28-byte EOF plus part of the final data block. + file.set_len(file_len - 40).unwrap(); + } + + fn lifecycle_paths(root: &Path) -> (PathBuf, PathBuf, PathBuf) { + ( + root.join("donor.bam"), + root.join("acceptor.bam"), + root.join("output.bam"), + ) + } + + fn run_lifecycle_child(case: &str, root: &Path) { + let mut child = Command::new(env::current_exe().unwrap()) + .args(["--exact", REPAIR_CHILD_TEST, "--nocapture"]) + .env(REPAIR_CHILD_CASE_ENV, case) + .env(REPAIR_CHILD_DIR_ENV, root) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .unwrap(); + let deadline = Instant::now() + Duration::from_secs(10); + loop { + if child.try_wait().unwrap().is_some() { + break; + } + if Instant::now() >= deadline { + let _ = child.kill(); + let output = child.wait_with_output().unwrap(); + panic!( + "repair lifecycle child {case} exceeded watchdog; \ + stdout={} stderr={}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + } + thread::sleep(Duration::from_millis(10)); + } + let output = child.wait_with_output().unwrap(); + assert!( + output.status.success(), + "repair lifecycle child {case} failed; stdout={} stderr={}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + } + + fn staged_repair_entries(root: &Path) -> Vec { + fs::read_dir(root) + .unwrap() + .map(|entry| entry.unwrap().path()) + .filter(|path| { + path.file_name() + .unwrap() + .to_string_lossy() + .starts_with(REPAIR_TEMP_PREFIX) + }) + .collect() + } + + #[test] + fn repair_lifecycle_child() { + let Ok(case) = env::var(REPAIR_CHILD_CASE_ENV) else { + return; + }; + let root = PathBuf::from(env::var_os(REPAIR_CHILD_DIR_ENV).unwrap()); + let (donor_bam, acceptor_bam, output_bam) = lifecycle_paths(&root); + let repair = RepairTags { + donor_bam, + acceptor_bam, + output_bam, + log_filepath: None, + threads: 2, + }; + let result = match case.as_str() { + "nth-write" => repair.run_with_output_factory(|header| { + Ok(FailNthWriteSink { + inner: StagedBamOutput::new(&repair.output_bam, header)?, + fail_at: 2, + writes: 0, + }) + }), + "validation" => repair.run_with_output_factory(|header| { + StagedBamOutput::with_validator( + &repair.output_bam, + header, + CorruptBeforeValidation, + ) + }), + _ => repair.run_with_output_factory(|header| { + StagedBamOutput::new(&repair.output_bam, header) + }), + }; + match case.as_str() { + "success" => result.expect("successful repair should finish"), + "truncated-donor" => { + let error = result.expect_err("truncated donor must fail"); + assert!( + format!("{error:#}") + .contains("failed to decode donor BAM record"), + "unexpected truncated-donor error: {error:#}" + ); + } + "truncated-acceptor" => { + let error = result.expect_err("truncated acceptor must fail"); + assert!( + format!("{error:#}") + .contains("failed to decode acceptor BAM record"), + "unexpected truncated-acceptor error: {error:#}" + ); + } + "nth-write" => { + let error = result.expect_err("injected write must fail"); + assert!( + format!("{error:#}").contains( + "injected repair sink write failure at record 2" + ), + "unexpected injected-write error: {error:#}" + ); + } + "validation" => { + let error = result.expect_err("corrupt staged BAM must fail"); + assert!( + format!("{error:#}") + .contains("missing the canonical BGZF EOF block"), + "unexpected validation error: {error:#}" + ); + } + other => panic!("unknown repair lifecycle child case {other}"), + } + } + + #[test] + fn repair_lifecycle_failures_preserve_outputs_and_clean_staging() { + const SENTINEL: &[u8] = b"existing repair output sentinel"; + for case in [ + "truncated-donor", + "truncated-acceptor", + "nth-write", + "validation", + "success", + ] { + for output_exists in [false, true] { + let temp_dir = tempfile::tempdir().unwrap(); + let (donor, acceptor, output) = + lifecycle_paths(temp_dir.path()); + let records = + if case.starts_with("truncated-") { 2_000 } else { 3 }; + write_lifecycle_bam(&donor, records, true); + write_lifecycle_bam(&acceptor, records, false); + match case { + "truncated-donor" => truncate_bam_data(&donor), + "truncated-acceptor" => truncate_bam_data(&acceptor), + _ => {} + } + if output_exists { + fs::write(&output, SENTINEL).unwrap(); + } + + run_lifecycle_child(case, temp_dir.path()); + + if case == "success" { + validate_complete_bam(&output, records).unwrap(); + assert_ne!(fs::read(&output).unwrap(), SENTINEL); + } else if output_exists { + assert_eq!(fs::read(&output).unwrap(), SENTINEL); + } else { + assert!(!output.exists()); + } + let staged = staged_repair_entries(temp_dir.path()); + assert!( + staged.is_empty(), + "staged repair entries remain: {staged:?}" + ); + } + } + } + + #[cfg(unix)] + #[test] + fn staged_output_matches_file_create_mode_and_preserves_existing_mode() { + let temp_dir = tempfile::tempdir().unwrap(); + // Model a shared, non-sticky destination directory. The staged BAM + // must still live under an owner-only child directory. + fs::set_permissions(temp_dir.path(), fs::Permissions::from_mode(0o777)) + .unwrap(); + let header = bam::Header::new(); + + let control = temp_dir.path().join("ordinary-file-create"); + drop(File::create(&control).unwrap()); + let ordinary_create_mode = + fs::metadata(&control).unwrap().permissions().mode() & 0o7777; + + let new_output = temp_dir.path().join("new-output.bam"); + let staged_output = StagedBamOutput::new(&new_output, &header).unwrap(); + let staging_dir = staged_output._staging_dir.path().to_path_buf(); + assert_eq!(staging_dir.parent(), Some(temp_dir.path())); + assert_eq!( + fs::metadata(&staging_dir).unwrap().permissions().mode() & 0o077, + 0 + ); + assert_eq!( + staged_output.temp_file.path().parent(), + Some(staging_dir.as_path()) + ); + assert_eq!( + staged_output + .temp_file + .as_file() + .metadata() + .unwrap() + .permissions() + .mode() + & 0o7777, + ordinary_create_mode + ); + staged_output.finish(0).unwrap(); + assert!(!staging_dir.exists()); + let new_output_mode = + fs::metadata(&new_output).unwrap().permissions().mode() & 0o7777; + assert_eq!(new_output_mode, ordinary_create_mode); + + let existing_output = temp_dir.path().join("existing-output.bam"); + fs::write(&existing_output, b"sentinel").unwrap(); + let deliberate_mode = 0o404; + fs::set_permissions( + &existing_output, + fs::Permissions::from_mode(deliberate_mode), + ) + .unwrap(); + StagedBamOutput::new(&existing_output, &header) + .unwrap() + .finish(0) + .unwrap(); + let replacement_mode = + fs::metadata(&existing_output).unwrap().permissions().mode() + & 0o7777; + assert_eq!(replacement_mode, deliberate_mode); + assert!(staged_repair_entries(temp_dir.path()).is_empty()); + } + + #[cfg(unix)] + #[test] + fn staged_inode_substitution_is_fatal_and_cleans_private_directory() { + const SENTINEL: &[u8] = b"existing repair output sentinel"; + for output_exists in [false, true] { + let temp_dir = tempfile::tempdir().unwrap(); + let output = temp_dir.path().join("output.bam"); + if output_exists { + fs::write(&output, SENTINEL).unwrap(); + } + let staged_output = StagedBamOutput::with_validator( + &output, + &bam::Header::new(), + SubstituteBeforeDecode, + ) + .unwrap(); + let staging_dir = staged_output._staging_dir.path().to_path_buf(); + + let error = staged_output + .finish(0) + .expect_err("replacement inode must not be persisted"); + assert!( + format!("{error:#}") + .contains("no longer identifies the retained file"), + "unexpected substitution error: {error:#}" + ); + if output_exists { + assert_eq!(fs::read(&output).unwrap(), SENTINEL); + } else { + assert!(!output.exists()); + } + assert!(!staging_dir.exists()); + assert!(staged_repair_entries(temp_dir.path()).is_empty()); + } + } + + #[cfg(unix)] + #[test] + fn staged_output_rejects_symlink_and_nonregular_destinations() { + const SENTINEL: &[u8] = b"symlink target sentinel"; + let temp_dir = tempfile::tempdir().unwrap(); + let header = bam::Header::new(); + + let target = temp_dir.path().join("target.bam"); + fs::write(&target, SENTINEL).unwrap(); + let symlink_output = temp_dir.path().join("symlink-output.bam"); + symlink(&target, &symlink_output).unwrap(); + let symlink_error = match StagedBamOutput::new(&symlink_output, &header) + { + Ok(_) => panic!("symlink destination should be rejected"), + Err(error) => error, + }; + assert!( + format!("{symlink_error:#}") + .contains("must not be a symbolic link"), + "unexpected symlink error: {symlink_error:#}" + ); + assert_eq!(fs::read(&target).unwrap(), SENTINEL); + assert!(fs::symlink_metadata(&symlink_output) + .unwrap() + .file_type() + .is_symlink()); + + let dangling_output = temp_dir.path().join("dangling-output.bam"); + symlink(temp_dir.path().join("missing-target.bam"), &dangling_output) + .unwrap(); + let dangling_error = + match StagedBamOutput::new(&dangling_output, &header) { + Ok(_) => { + panic!("dangling symlink destination should be rejected") + } + Err(error) => error, + }; + assert!( + format!("{dangling_error:#}") + .contains("must not be a symbolic link"), + "unexpected dangling-symlink error: {dangling_error:#}" + ); + assert!(fs::symlink_metadata(&dangling_output) + .unwrap() + .file_type() + .is_symlink()); + + let fifo_output = temp_dir.path().join("fifo-output.bam"); + let fifo_status = + Command::new("mkfifo").arg(&fifo_output).status().unwrap(); + assert!(fifo_status.success(), "failed to create FIFO fixture"); + let fifo_error = match StagedBamOutput::new(&fifo_output, &header) { + Ok(_) => panic!("FIFO destination should be rejected"), + Err(error) => error, + }; + assert!( + format!("{fifo_error:#}") + .contains("must be a regular file when it already exists"), + "unexpected FIFO error: {fifo_error:#}" + ); + assert!(fs::symlink_metadata(&fifo_output) + .unwrap() + .file_type() + .is_fifo()); + + let directory_output = temp_dir.path().join("directory-output.bam"); + fs::create_dir(&directory_output).unwrap(); + let directory_error = + match StagedBamOutput::new(&directory_output, &header) { + Ok(_) => panic!("directory destination should be rejected"), + Err(error) => error, + }; + assert!( + format!("{directory_error:#}") + .contains("must be a regular file when it already exists"), + "unexpected directory error: {directory_error:#}" + ); + assert!(directory_output.is_dir()); + assert!(staged_repair_entries(temp_dir.path()).is_empty()); + } + struct GatedRepairWorker { job_one_completed: Sender<()>, release_job_zero: Receiver<()>, From 23709cf3aff0da3aeec6beb02498629e0920aedf Mon Sep 17 00:00:00 2001 From: SuhasSrinivasan <32346517+SuhasSrinivasan@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:52:58 -0700 Subject: [PATCH 4/4] Reject overlapping repair placements --- modkit-core/src/repair_tags.rs | 121 +++++++++++++++++++++++++++++---- 1 file changed, 108 insertions(+), 13 deletions(-) diff --git a/modkit-core/src/repair_tags.rs b/modkit-core/src/repair_tags.rs index e8fbfe00..e5fd2013 100644 --- a/modkit-core/src/repair_tags.rs +++ b/modkit-core/src/repair_tags.rs @@ -12,6 +12,7 @@ use crossbeam_channel::unbounded; use derive_new::new; use indicatif::{MultiProgress, ProgressBar}; use log::{debug, info, warn}; +use memchr::memmem::Finder; use rust_htslib::bam::record::{Aux, AuxArray}; use rust_htslib::bam::{self, Read}; use rustc_hash::FxHashMap; @@ -1095,19 +1096,18 @@ fn repair_record_pair(record_pair: RecordPair) -> anyhow::Result { })?; if donor_seq.len() < acceptor_seq.len() { - bail!("donor sequence for {read_name} is longer than acceptor sequence") + bail!( + "donor sequence for {read_name} is shorter than acceptor sequence" + ) } - let matches = donor_seq.match_indices(&acceptor_seq); - - let starts = - matches.into_iter().map(|(start, _)| start).collect::>(); - if starts.len() > 1 { - bail!("multiple potential corrections found for {read_name}") - } else if starts.is_empty() { + let finder = Finder::new(acceptor_seq.as_bytes()); + let Some(start) = finder.find(donor_seq.as_bytes()) else { bail!("acceptor sequence is not a substring of the donor sequence") + }; + if finder.find(&donor_seq.as_bytes()[start + 1..]).is_some() { + bail!("multiple potential corrections found for {read_name}") } else { let acceptor_seq_len = acceptor_seq.len(); - let start = *starts.get(0).unwrap(); let end = start + acceptor_seq_len; let mm_style = modbase_info.mm_style; @@ -1185,10 +1185,11 @@ mod tests { use crate::ordered_scheduler::OrderedWorker; use super::{ - query_name_order, run_repair_scheduler, validate_complete_bam, - CompleteBamValidator, QueryNameOrder, RecordPair, RepairJob, - RepairSink, RepairSlot, RepairTags, StagedBamOutput, - StagedBamValidator, StagedFileIdentity, REPAIR_TEMP_PREFIX, + query_name_order, repair_record_pair, run_repair_scheduler, + validate_complete_bam, CompleteBamValidator, QueryNameOrder, + RecordPair, RepairJob, RepairSink, RepairSlot, RepairTags, + StagedBamOutput, StagedBamValidator, StagedFileIdentity, + REPAIR_TEMP_PREFIX, }; fn make_record(name: &[u8], sequence: &str) -> bam::Record { @@ -1197,6 +1198,100 @@ mod tests { record } + fn make_donor(name: &[u8], sequence: &str) -> bam::Record { + let mut record = make_record(name, sequence); + record.push_aux(b"MM", Aux::String("A+a.,0;")).unwrap(); + record.push_aux(b"ML", Aux::ArrayU8((&[200u8][..]).into())).unwrap(); + record + } + + fn repair(donor: &str, acceptor: &str) -> anyhow::Result { + repair_record_pair(RecordPair::new( + Arc::new(make_donor(b"read", donor)), + make_record(b"read", acceptor), + )) + } + + fn repair_reverse( + donor_stored_sequence: &str, + acceptor_stored_sequence: &str, + ) -> anyhow::Result { + let mut donor = make_donor(b"reverse", donor_stored_sequence); + donor.set_reverse(); + let mut acceptor = make_record(b"reverse", acceptor_stored_sequence); + acceptor.set_reverse(); + repair_record_pair(RecordPair::new(Arc::new(donor), acceptor)) + } + + #[test] + fn repair_rejects_overlapping_ambiguous_placements() { + let error = repair("ACACAC", "ACAC").unwrap_err(); + assert!( + error.to_string().contains("multiple potential corrections"), + "unexpected error: {error:#}" + ); + + assert!(repair("ACACGGACAC", "ACAC") + .unwrap_err() + .to_string() + .contains("multiple potential corrections")); + } + + #[test] + fn repair_preserves_unique_and_absent_placement_behavior() { + let repaired = repair("TACACG", "ACAC").unwrap(); + assert!(matches!(repaired.aux(b"MM").unwrap(), Aux::String("A+a.,0;"))); + match repaired.aux(b"ML").unwrap() { + Aux::ArrayU8(qualities) => { + assert_eq!(qualities.iter().collect::>(), vec![200]); + } + other => panic!("unexpected ML tag: {other:?}"), + } + assert!(matches!(repaired.aux(b"MN").unwrap(), Aux::U32(4))); + + assert!(repair("ACACAC", "CGCG") + .unwrap_err() + .to_string() + .contains("not a substring")); + } + + #[test] + fn repair_searches_forward_sequences_for_reverse_records() { + // Stored placement starts at 3, while the forward-sequence placement + // TACGAAA -> ACG starts at 1. + let repaired = repair_reverse("TTTCGTA", "CGT").unwrap(); + assert!(repaired.is_reverse()); + assert!(matches!(repaired.aux(b"MM").unwrap(), Aux::String("A+a.,0;"))); + match repaired.aux(b"ML").unwrap() { + Aux::ArrayU8(qualities) => { + assert_eq!(qualities.iter().collect::>(), vec![200]); + } + other => panic!("unexpected ML tag: {other:?}"), + } + assert!(matches!(repaired.aux(b"MN").unwrap(), Aux::U32(3))); + + assert!(repair_reverse("GTGTGT", "GTGT") + .unwrap_err() + .to_string() + .contains("multiple potential corrections")); + } + + #[test] + fn repair_returns_errors_for_empty_sequences() { + for (donor, acceptor, expected) in [ + ( + "ACAC", + "", + "acceptor sequence for record read failed, empty-read-sequence", + ), + ("", "ACAC", "record read failed, invalid-MM-tag"), + ("", "", "record read failed, invalid-MM-tag"), + ] { + let error = repair(donor, acceptor).unwrap_err(); + assert_eq!(error.to_string(), expected); + } + } + struct FailNthWriteSink { inner: S, fail_at: usize,