diff --git a/Cargo.lock b/Cargo.lock index 82c556de7..536cdc14a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -669,15 +669,18 @@ dependencies = [ "fgoxide", "fgumi-bam-io", "fgumi-bgzf", + "fgumi-cli-common", "fgumi-cli-macros", "fgumi-consensus", "fgumi-dna", "fgumi-metrics", "fgumi-pipeline-core", + "fgumi-pipeline-io", "fgumi-raw-bam", "fgumi-sam", "fgumi-simd-fastq", "fgumi-sort", + "fgumi-sort-cli", "fgumi-umi", "flate2", "fs4", @@ -742,6 +745,21 @@ dependencies = [ "libdeflater", ] +[[package]] +name = "fgumi-cli-common" +version = "0.3.0" +dependencies = [ + "anyhow", + "bytesize", + "clap", + "enum_dispatch", + "log", + "num_cpus", + "rstest", + "sysinfo", + "thiserror 2.0.18", +] + [[package]] name = "fgumi-cli-macros" version = "0.3.0" @@ -804,6 +822,23 @@ dependencies = [ "trybuild", ] +[[package]] +name = "fgumi-pipeline-io" +version = "0.3.0" +dependencies = [ + "anyhow", + "fgumi-bam-io", + "fgumi-bgzf", + "fgumi-pipeline-core", + "fgumi-raw-bam", + "fgumi-sort", + "log", + "noodles", + "parking_lot", + "rstest", + "tempfile", +] + [[package]] name = "fgumi-raw-bam" version = "0.3.0" @@ -880,6 +915,27 @@ dependencies = [ "zstd", ] +[[package]] +name = "fgumi-sort-cli" +version = "0.3.0" +dependencies = [ + "anyhow", + "bytesize", + "clap", + "fgumi-bam-io", + "fgumi-cli-common", + "fgumi-cli-macros", + "fgumi-pipeline-core", + "fgumi-pipeline-io", + "fgumi-sam", + "fgumi-sort", + "log", + "noodles", + "parking_lot", + "rstest", + "tempfile", +] + [[package]] name = "fgumi-tag" version = "0.3.0" diff --git a/Cargo.toml b/Cargo.toml index 974134f83..3612aebc7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [workspace] -members = [".", "crates/fgumi-raw-bam", "crates/fgumi-dna", "crates/fgumi-bgzf", "crates/fgumi-metrics", "crates/fgumi-sam", "crates/fgumi-simd-fastq", "crates/fgumi-tag", "crates/fgumi-umi", "crates/fgumi-consensus", "crates/fgumi-bam-io", "crates/fgumi-sort", "crates/fgumi-pipeline-core", "crates/fgumi-cli-macros", "crates/xtask"] +members = [".", "crates/fgumi-raw-bam", "crates/fgumi-dna", "crates/fgumi-bgzf", "crates/fgumi-metrics", "crates/fgumi-sam", "crates/fgumi-simd-fastq", "crates/fgumi-tag", "crates/fgumi-umi", "crates/fgumi-consensus", "crates/fgumi-bam-io", "crates/fgumi-sort", "crates/fgumi-pipeline-core", "crates/fgumi-cli-macros", "crates/fgumi-cli-common", "crates/fgumi-pipeline-io", "crates/fgumi-sort-cli", "crates/xtask"] resolver = "2" [workspace.package] @@ -21,6 +21,9 @@ fgumi-sam = { version = "0.3.0", path = "crates/fgumi-sam" } fgumi-simd-fastq = { version = "0.3.0", path = "crates/fgumi-simd-fastq" } fgumi-sort = { version = "0.3.0", path = "crates/fgumi-sort" } fgumi-pipeline-core = { version = "0.3.0", path = "crates/fgumi-pipeline-core" } +fgumi-cli-common = { version = "0.3.0", path = "crates/fgumi-cli-common" } +fgumi-pipeline-io = { version = "0.3.0", path = "crates/fgumi-pipeline-io" } +fgumi-sort-cli = { version = "0.3.0", path = "crates/fgumi-sort-cli" } fgumi-tag = { version = "0.3.0", path = "crates/fgumi-tag" } fgumi-umi = { version = "0.3.0", path = "crates/fgumi-umi" } bytemuck = { version = "1.14", features = ["derive"] } @@ -101,6 +104,9 @@ fgumi-consensus = { workspace = true } fgumi-bam-io = { workspace = true } fgumi-sort = { workspace = true } fgumi-pipeline-core = { workspace = true } +fgumi-cli-common = { workspace = true } +fgumi-pipeline-io = { workspace = true } +fgumi-sort-cli = { workspace = true } [target.'cfg(target_os = "macos")'.dependencies] mach2 = "0.6" diff --git a/crates/fgumi-cli-common/Cargo.toml b/crates/fgumi-cli-common/Cargo.toml new file mode 100644 index 000000000..65091f2ff --- /dev/null +++ b/crates/fgumi-cli-common/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "fgumi-cli-common" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +description = "Shared CLI types and helpers for fgumi commands" +repository.workspace = true +license.workspace = true + +[dependencies] +anyhow = "1.0.102" +bytesize = "2.3" +clap = { version = "4", features = ["derive", "string"] } +enum_dispatch = "0.3.13" +log = "0" +num_cpus = "1.16" +sysinfo = { version = "0.38", default-features = false, features = ["system"] } +thiserror = "2" + +[dev-dependencies] +rstest = "0" diff --git a/crates/fgumi-cli-common/src/lib.rs b/crates/fgumi-cli-common/src/lib.rs new file mode 100644 index 000000000..e4c35a0a6 --- /dev/null +++ b/crates/fgumi-cli-common/src/lib.rs @@ -0,0 +1,699 @@ +#![deny(unsafe_code)] +//! Shared CLI types and helpers for fgumi commands. + +// ───────────────────────────────────────────────────────────────────────────── +// Command trait +// ───────────────────────────────────────────────────────────────────────────── + +use anyhow::Result; +use enum_dispatch::enum_dispatch; + +/// Trait implemented by all fgumi CLI commands. +/// +/// Each command provides an `execute` method that runs the command's main logic. +/// The `command_line` parameter contains the full command invocation for @PG records. +#[enum_dispatch] +pub trait Command { + #[allow(clippy::missing_errors_doc)] + fn execute(&self, command_line: &str) -> Result<()>; +} + +// ───────────────────────────────────────────────────────────────────────────── +// Error types +// ───────────────────────────────────────────────────────────────────────────── + +use thiserror::Error; + +/// Result type alias for fgumi operations +pub type FgumiResult = std::result::Result; + +/// Error type for fgumi operations +#[derive(Error, Debug)] +pub enum FgumiError { + /// Invalid parameter value provided + #[error("Invalid parameter '{parameter}': {reason}")] + InvalidParameter { + /// The parameter name + parameter: String, + /// Explanation of why it's invalid + reason: String, + }, + + /// Invalid frequency threshold + #[error("Invalid frequency threshold: {value} (must be between {min} and {max})")] + InvalidFrequency { + /// The invalid frequency value + value: f64, + /// Minimum valid value + min: f64, + /// Maximum valid value + max: f64, + }, + + /// Invalid quality threshold + #[error("Invalid quality threshold: {value} (must be between 0 and {max})")] + InvalidQuality { + /// The invalid quality value + value: u8, + /// Maximum valid value (usually 93 for SAM/BAM) + max: u8, + }, + + /// File format error + #[error("Invalid {file_type} file '{path}': {reason}")] + InvalidFileFormat { + /// Type of file (e.g., "BAM", "FASTQ") + file_type: String, + /// Path to the file + path: String, + /// Explanation of the problem + reason: String, + }, + + /// Required reference sequence not found + #[error("Reference sequence '{ref_name}' not found in header")] + ReferenceNotFound { + /// The reference sequence name + ref_name: String, + }, + + /// Invalid memory size string + #[error("Invalid memory size: {reason}")] + InvalidMemorySize { + /// Explanation of why the value is invalid + reason: String, + }, +} + +// ───────────────────────────────────────────────────────────────────────────── +// System detection +// ───────────────────────────────────────────────────────────────────────────── + +/// Returns the effective total memory available to this process in bytes. +/// +/// Checks cgroup memory limits for container environments, falling back to +/// physical RAM on bare-metal or macOS. Always returns `min(cgroup, physical)`. +#[must_use] +pub fn detect_total_memory() -> usize { + let mut system = sysinfo::System::new(); + system.refresh_memory(); + let physical = system.total_memory(); + let bytes = system.cgroup_limits().map_or(physical, |c| c.total_memory.min(physical)); + usize::try_from(bytes).unwrap_or(usize::MAX) +} + +/// Returns the number of logical CPUs available to this process. +/// +/// Uses `num_cpus::get()`, which attempts to honor cgroup CPU quotas (e.g. +/// `--cpus` in Docker or Kubernetes resource limits) but may fall back to the +/// physical core count — its cgroup v2 quota handling is incomplete (see +/// [num_cpus#122](https://github.com/seanmonstar/num_cpus/issues/122)). Returns +/// at least 1. +#[must_use] +pub fn detect_cpu_count() -> usize { + num_cpus::get().max(1) +} + +// ───────────────────────────────────────────────────────────────────────────── +// Formatting helpers +// ───────────────────────────────────────────────────────────────────────────── + +/// Format an integer with comma separators (e.g. 1234567 → "1,234,567"). +fn format_count(n: u64) -> String { + let s = n.to_string(); + let mut result = String::with_capacity(s.len() + s.len() / 3); + let offset = s.len() % 3; + for (i, c) in s.chars().enumerate() { + // Insert a comma before every group of three digits, but not before + // the first character. `offset` is the length of the leading partial + // group, so commas fall at indices `offset, offset+3, offset+6, …` + // (skipping index 0 when `offset == 0`). + if i >= offset && i > 0 && (i - offset).is_multiple_of(3) { + result.push(','); + } + result.push(c); + } + result +} + +/// Formats a duration in human-readable form. +/// +/// # Examples +/// +/// ``` +/// use fgumi_cli_common::format_duration; +/// use std::time::Duration; +/// +/// assert_eq!(format_duration(Duration::from_secs(45)), "45s"); +/// assert_eq!(format_duration(Duration::from_secs(135)), "2m 15s"); +/// assert_eq!(format_duration(Duration::from_secs(5400)), "1h 30m"); +/// ``` +#[must_use] +pub fn format_duration(duration: std::time::Duration) -> String { + let secs = duration.as_secs(); + if secs < 60 { + format!("{secs}s") + } else if secs < 3600 { + let mins = secs / 60; + let remaining_secs = secs % 60; + if remaining_secs == 0 { format!("{mins}m") } else { format!("{mins}m {remaining_secs}s") } + } else { + let hours = secs / 3600; + let mins = (secs % 3600) / 60; + if mins == 0 { format!("{hours}h") } else { format!("{hours}h {mins}m") } + } +} + +/// Formats a rate (items per second) with appropriate units. +/// +/// # Examples +/// +/// ``` +/// use fgumi_cli_common::format_rate; +/// use std::time::Duration; +/// +/// assert_eq!(format_rate(1000, Duration::from_secs(1)), "1,000 items/s"); +/// assert_eq!(format_rate(600, Duration::from_secs(60)), "10 items/s"); +/// ``` +#[must_use] +#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation, clippy::cast_sign_loss)] +pub fn format_rate(count: u64, duration: std::time::Duration) -> String { + let secs = duration.as_secs_f64(); + if secs < 0.001 { + return format!("{} items/s", format_count(count)); + } + + let rate = count as f64 / secs; + if rate >= 1.0 { + format!("{} items/s", format_count(rate as u64)) + } else { + let items_per_min = count as f64 / (secs / 60.0); + format!("{items_per_min:.1} items/min") + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Operation timer +// ───────────────────────────────────────────────────────────────────────────── + +/// Operation timing and summary helper. +/// +/// Tracks operation timing and provides formatted summary output. +pub struct OperationTimer { + operation: String, + start_time: std::time::Instant, +} + +impl OperationTimer { + /// Creates a new operation timer and logs the start. + #[must_use] + pub fn new(operation: &str) -> Self { + log::info!("{operation} ..."); + Self { operation: operation.to_string(), start_time: std::time::Instant::now() } + } + + /// Logs the completion with item count and rate. + pub fn log_completion(&self, count: u64) { + let duration = self.start_time.elapsed(); + log::info!( + "{} completed: {} in {} ({})", + self.operation, + format_count(count), + format_duration(duration), + format_rate(count, duration) + ); + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Validation helpers +// ───────────────────────────────────────────────────────────────────────────── + +use bytesize::ByteSize; +use std::path::Path; + +/// Validate that a file exists. +/// +/// # Errors +/// +/// Returns [`FgumiError::InvalidFileFormat`] if the file does not exist. +pub fn validate_file_exists>(path: P, description: &str) -> FgumiResult<()> { + let path_ref = path.as_ref(); + if !path_ref.exists() { + return Err(FgumiError::InvalidFileFormat { + file_type: description.to_string(), + path: path_ref.display().to_string(), + reason: "File does not exist".to_string(), + }); + } + Ok(()) +} + +/// Parses a memory size string into bytes. +/// +/// Accepts both plain numbers (interpreted as MiB) and human-readable formats like: +/// - "2GB", "2G" -> 2 gigabytes (decimal: 2,000,000,000) +/// - "1.5GB" -> 1.5 gigabytes +/// - "1024MB", "1024M" -> 1024 megabytes (decimal) +/// - "512MiB" -> 512 mebibytes (binary: 536,870,912) +/// - "768" -> 768 MiB (plain numbers are interpreted as mebibytes) +/// +/// # Errors +/// +/// Returns [`FgumiError::InvalidMemorySize`] if the string cannot be parsed. +pub fn parse_memory_size(size_str: &str) -> FgumiResult { + let trimmed = size_str.trim(); + if trimmed.is_empty() { + return Err(FgumiError::InvalidMemorySize { + reason: "Memory size cannot be empty".to_string(), + }); + } + + if trimmed.starts_with('-') { + return Err(FgumiError::InvalidMemorySize { + reason: format!("Memory size cannot be negative: '{trimmed}'"), + }); + } + + if let Ok(mb_value) = trimmed.parse::() { + if mb_value == 0 { + return Err(FgumiError::InvalidMemorySize { + reason: "Memory size cannot be zero".to_string(), + }); + } + if mb_value > 1_000_000 { + return Err(FgumiError::InvalidMemorySize { + reason: format!( + "Plain number memory size too large: {} MiB. Use human-readable format like '{}GB' instead.", + mb_value, + mb_value / 1000 + ), + }); + } + return mb_value.checked_mul(1024 * 1024).ok_or_else(|| FgumiError::InvalidMemorySize { + reason: format!("Memory size calculation overflow for {mb_value} MiB"), + }); + } + + if trimmed.contains('e') || trimmed.contains('E') { + return Err(FgumiError::InvalidMemorySize { + reason: format!( + "Scientific notation not supported: '{trimmed}'. Use integer values or human-readable formats like '2GB'." + ), + }); + } + + if trimmed.contains('.') && trimmed.chars().all(|c| c.is_ascii_digit() || c == '.') { + return Err(FgumiError::InvalidMemorySize { + reason: format!( + "Plain decimal numbers not supported: '{trimmed}'. Use an integer for MiB (e.g. '768') or a human-readable format (e.g. '1.5GB')." + ), + }); + } + + match trimmed.parse::() { + Ok(size) => { + if size.0 == 0 { + return Err(FgumiError::InvalidMemorySize { + reason: format!("Memory size cannot be zero: '{trimmed}'"), + }); + } + Ok(size.0) + } + Err(_) => Err(FgumiError::InvalidMemorySize { + reason: format!( + "Invalid memory size '{trimmed}'. Valid formats:\n\ + - Plain numbers (interpreted as MiB): '768', '4096'\n\ + - Human-readable (decimal): '2GB', '1024MB'\n\ + - Human-readable (binary): '1GiB', '512MiB'" + ), + }), + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Memory/compression options +// ───────────────────────────────────────────────────────────────────────────── + +/// A memory limit, either auto-detected from the host or a fixed byte count. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MemoryLimit { + /// Detect the (cgroup-aware) host memory and subtract the reserve. + Auto, + /// Use a fixed memory limit in bytes. + Fixed(usize), +} + +impl Default for MemoryLimit { + fn default() -> Self { + Self::Fixed(768 * 1024 * 1024) + } +} + +impl std::fmt::Display for MemoryLimit { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Auto => f.write_str("auto"), + Self::Fixed(bytes) => format_binary_bytes(*bytes, f), + } + } +} + +/// How much memory to reserve for other processes when [`MemoryLimit::Auto`] is used. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum MemoryReserve { + /// Automatic: `min(10 GiB, 50% of host memory)`. + /// Matches the clap `default_value = "auto"` on `SortOptions::memory_reserve`. + #[default] + Auto, + /// Reserve a fixed number of bytes. + Fixed(usize), +} + +impl std::fmt::Display for MemoryReserve { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Auto => f.write_str("auto"), + Self::Fixed(bytes) => format_binary_bytes(*bytes, f), + } + } +} + +/// Format a byte count in the largest binary unit that divides cleanly. +fn format_binary_bytes(bytes: usize, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + const K: usize = 1024; + const M: usize = K * 1024; + const G: usize = M * 1024; + if bytes >= G && bytes.is_multiple_of(G) { + write!(f, "{}GiB", bytes / G) + } else if bytes >= M && bytes.is_multiple_of(M) { + write!(f, "{}MiB", bytes / M) + } else if bytes >= K && bytes.is_multiple_of(K) { + write!(f, "{}KiB", bytes / K) + } else { + write!(f, "{bytes}B") + } +} + +/// The minimum per-thread memory budget (256 MiB). +pub const MIN_MEMORY_PER_THREAD: usize = 256 * 1024 * 1024; + +/// Default auto-reserve cap: 10 GiB. +const AUTO_RESERVE_CAP: usize = 10 * 1024 * 1024 * 1024; + +/// Parse a memory-limit string (e.g. "512M", "1G", "768", "auto"). +/// +/// # Errors +/// +/// Returns an error string if parsing fails. +pub fn parse_memory(s: &str) -> Result { + let s = s.trim(); + if s.eq_ignore_ascii_case("auto") { + return Ok(MemoryLimit::Auto); + } + Ok(MemoryLimit::Fixed(parse_memory_bytes(s, "Memory size")?)) +} + +/// Parse a memory-reserve string (e.g. "10G", "auto"). +/// +/// # Errors +/// +/// Returns an error string if parsing fails. +pub fn parse_memory_reserve(s: &str) -> Result { + let s = s.trim(); + if s.eq_ignore_ascii_case("auto") { + return Ok(MemoryReserve::Auto); + } + Ok(MemoryReserve::Fixed(parse_memory_bytes(s, "Memory reserve")?)) +} + +/// Resolve a [`MemoryReserve`] to a concrete byte count given total host memory. +#[must_use] +pub fn resolve_reserve(reserve: MemoryReserve, total_memory: usize) -> usize { + match reserve { + MemoryReserve::Fixed(bytes) => bytes, + MemoryReserve::Auto => AUTO_RESERVE_CAP.min(total_memory / 2), + } +} + +/// Resolve a memory budget to a concrete byte count. +/// +/// # Errors +/// +/// Returns an error if `threads` is 0 or the multiplication overflows. +pub fn resolve_memory_budget( + limit: MemoryLimit, + reserve: MemoryReserve, + threads: usize, + per_thread: bool, +) -> anyhow::Result { + resolve_memory_budget_with_total(limit, reserve, threads, per_thread, detect_total_memory()) +} + +/// Pure resolver behind [`resolve_memory_budget`], with `total` injected for testability. +fn resolve_memory_budget_with_total( + limit: MemoryLimit, + reserve: MemoryReserve, + threads: usize, + per_thread: bool, + total: usize, +) -> anyhow::Result { + if threads == 0 { + anyhow::bail!("--threads must be at least 1"); + } + + let budget = match limit { + MemoryLimit::Fixed(bytes) => { + if per_thread { + bytes + .checked_mul(threads) + .ok_or_else(|| anyhow::anyhow!("memory limit × {threads} threads overflowed"))? + } else { + bytes + } + } + MemoryLimit::Auto => { + let margin = resolve_reserve(reserve, total); + let available = total.saturating_sub(margin); + let target = if per_thread { + (available / threads) + .max(MIN_MEMORY_PER_THREAD) + .checked_mul(threads) + .ok_or_else(|| anyhow::anyhow!("auto memory budget overflowed"))? + } else { + available.max(MIN_MEMORY_PER_THREAD) + }; + let budget = target.min(available); + if budget < target { + log::warn!( + "Auto memory: capping budget to host-available {} (minimum viable target {} \ + exceeds it after reserve {}); throughput may drop but the run stays within memory", + bytesize::ByteSize(budget as u64), + bytesize::ByteSize(target as u64), + bytesize::ByteSize(margin as u64), + ); + } + log::debug!( + "Auto memory: {} of {} ({}/thread × {} threads, reserve {})", + bytesize::ByteSize(budget as u64), + bytesize::ByteSize(total as u64), + bytesize::ByteSize((budget / threads) as u64), + threads, + bytesize::ByteSize(margin as u64), + ); + budget + } + }; + + if budget > total { + log::warn!( + "Memory budget {} exceeds total host memory {}; this may cause OOM (or, for sort, earlier spill-to-disk)", + bytesize::ByteSize(budget as u64), + bytesize::ByteSize(total as u64), + ); + } + + Ok(budget) +} + +/// Parse a memory size string into `usize` bytes (private helper). +fn parse_memory_bytes(s: &str, label: &str) -> Result { + let bytes = parse_memory_size(s).map_err(|e| e.to_string())?; + usize::try_from(bytes).map_err(|_| format!("{label} too large: {bytes}")) +} + +/// Parses a boolean value from a string, accepting: true/false, yes/no, y/n, t/f +/// (case-insensitive). Matches sopt/fgbio behavior. +/// +/// # Errors +/// +/// Returns an error string if the input is not a recognized boolean. +pub fn parse_bool(s: &str) -> Result { + match s.to_ascii_lowercase().as_str() { + "true" | "t" | "yes" | "y" => Ok(true), + "false" | "f" | "no" | "n" => Ok(false), + _ => Err(format!("Invalid boolean value '{s}'. Expected: true|false|yes|no|y|n|t|f")), + } +} + +/// Options for output compression. +/// +/// Controls BGZF compression level for BAM output files. +#[derive(Debug, Clone, clap::Args)] +pub struct CompressionOptions { + /// Compression level for output BAM (0-12). + /// + /// Level 0 disables compression (uncompressed BGZF blocks). + /// Level 1 is fastest of the compressing levels with larger files; + /// level 12 produces the smallest files but is slowest. + #[arg(long, default_value_t = 1, value_parser = clap::value_parser!(u32).range(0..=12))] + pub compression_level: u32, +} + +impl Default for CompressionOptions { + /// Mirrors the clap `default_value_t = 1` so programmatic/default-constructed + /// callers emit level-1 compression rather than the `u32` default of `0` + /// (uncompressed BGZF). + fn default() -> Self { + Self { compression_level: 1 } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[test] + fn test_detect_total_memory_nonzero() { + let total = detect_total_memory(); + assert!(total > 0, "expected non-zero total memory, got {total}"); + } + + #[test] + fn test_detect_cpu_count_at_least_one() { + assert!(detect_cpu_count() >= 1); + } + + #[test] + fn test_format_count_no_commas() { + assert_eq!(format_count(0), "0"); + assert_eq!(format_count(999), "999"); + } + + #[test] + fn test_format_count_with_commas() { + assert_eq!(format_count(1000), "1,000"); + assert_eq!(format_count(1_234_567), "1,234,567"); + } + + #[test] + fn test_format_duration() { + use std::time::Duration; + assert_eq!(format_duration(Duration::from_secs(45)), "45s"); + assert_eq!(format_duration(Duration::from_secs(135)), "2m 15s"); + assert_eq!(format_duration(Duration::from_secs(5400)), "1h 30m"); + } + + #[test] + fn test_format_rate() { + use std::time::Duration; + assert_eq!(format_rate(1000, Duration::from_secs(1)), "1,000 items/s"); + } + + #[test] + fn test_parse_memory_size_plain() { + assert_eq!(parse_memory_size("768").unwrap(), 768 * 1024 * 1024); + } + + #[test] + fn test_parse_memory_size_human() { + assert_eq!(parse_memory_size("2GB").unwrap(), 2 * 1000 * 1000 * 1000); + assert_eq!(parse_memory_size("512MiB").unwrap(), 512 * 1024 * 1024); + } + + #[test] + fn test_parse_memory_size_errors() { + assert!(parse_memory_size("").is_err()); + assert!(parse_memory_size("-1").is_err()); + assert!(parse_memory_size("0").is_err()); + } + + #[test] + fn test_memory_limit_display() { + assert_eq!(MemoryLimit::Auto.to_string(), "auto"); + assert_eq!(MemoryLimit::Fixed(768 * 1024 * 1024).to_string(), "768MiB"); + assert_eq!(MemoryLimit::Fixed(1024 * 1024 * 1024).to_string(), "1GiB"); + } + + #[test] + fn test_resolve_reserve_auto() { + let total = 32 * 1024 * 1024 * 1024_usize; // 32 GiB + let r = resolve_reserve(MemoryReserve::Auto, total); + assert_eq!(r, AUTO_RESERVE_CAP); // capped at 10 GiB + } + + #[test] + fn test_resolve_memory_budget_fixed() { + let budget = resolve_memory_budget( + MemoryLimit::Fixed(512 * 1024 * 1024), + MemoryReserve::Auto, + 4, + true, + ) + .unwrap(); + assert_eq!(budget, 4 * 512 * 1024 * 1024); + } + + #[test] + fn test_parse_bool() { + assert!(parse_bool("true").unwrap()); + assert!(!parse_bool("false").unwrap()); + assert!(parse_bool("yes").unwrap()); + assert!(!parse_bool("no").unwrap()); + assert!(parse_bool("maybe").is_err()); + } + + #[derive(clap::Parser)] + struct CompressionHarness { + #[command(flatten)] + compression: CompressionOptions, + } + + // In-range values (the 0-12 bounds; 0 = uncompressed) parse. + #[rstest] + #[case(0_u32)] + #[case(1_u32)] + #[case(6_u32)] + #[case(12_u32)] + fn test_compression_level_accepts_in_range(#[case] level: u32) { + use clap::Parser; + + assert_eq!( + CompressionHarness::try_parse_from(["prog", "--compression-level", &level.to_string()]) + .unwrap() + .compression + .compression_level, + level + ); + } + + // The programmatic `Default` must match the clap `default_value_t = 1`; a + // derived `Default` would silently yield level 0 (uncompressed BGZF). + #[test] + fn test_compression_options_default_matches_cli_default() { + assert_eq!(CompressionOptions::default().compression_level, 1); + } + + #[test] + fn test_compression_level_default_and_rejects_out_of_range() { + use clap::Parser; + + // Default is 1 when the flag is omitted. + assert_eq!(CompressionHarness::parse_from(["prog"]).compression.compression_level, 1); + + // Out-of-range values are rejected at parse time rather than silently accepted. + assert!(CompressionHarness::try_parse_from(["prog", "--compression-level", "13"]).is_err()); + assert!(CompressionHarness::try_parse_from(["prog", "--compression-level", "99"]).is_err()); + } +} diff --git a/crates/fgumi-pipeline-core/src/handles.rs b/crates/fgumi-pipeline-core/src/handles.rs index 8a7cba91c..1bb4451f8 100644 --- a/crates/fgumi-pipeline-core/src/handles.rs +++ b/crates/fgumi-pipeline-core/src/handles.rs @@ -1102,7 +1102,12 @@ pub(crate) fn build_single_queues( assert_eq!(specs.len(), 1, "Single::build_queues requires 1 spec"); assert_eq!(ordering.len(), 1, "Single::build_queues requires 1 ordering"); - let branch = build_branch::(specs[0], ordering[0]); + // `Single` bounds `T: HeapSize`, so use the byte-aware build path: it + // honors `QueueSpec::ByteBounded` (documented as supported for `Single` + // outputs without item-carried serials) and delegates every non-byte spec + // straight back to `build_branch::`, so count/unbounded paths are + // unchanged. + let branch = build_branch_byte_aware::(specs[0], ordering[0]); let view = SingleOutputsView { primary: branch.output }; let outputs_view = OutputsViewAny { inner: Box::new(view) }; let queue_set = OutputQueueSet::new(vec![BranchEntry { @@ -1357,18 +1362,23 @@ mod handle_tests { } #[test] - #[should_panic(expected = "ByteBounded requires `T: HeapSize`")] - fn byte_bounded_panics_via_user_facing_build_queues() { - // Regression: user step declaring `QueueSpec::ByteBounded` in its - // `StepProfile::output_queues` reaches the panic via - // `StepOutputs::build_queues` → `build_single_queues` → - // `build_branch::`. The panic message must clearly direct the - // step author to PR 2's byte-aware path. + fn byte_bounded_single_builds_via_user_facing_build_queues() { + // Regression: a user step declaring `QueueSpec::ByteBounded` in its + // `StepProfile::output_queues` with a `Single` output reaches + // `StepOutputs::build_queues` → `build_single_queues`. `Single` + // bounds `T: HeapSize`, so this builds a byte-bounded queue (the + // documented "byte-bounded without item-carried serials" shape) + // rather than panicking. use crate::outputs::{Single, StepOutputs}; - let _ = as StepOutputs>::build_queues( + use crate::step::OutputHandles; + let (mut queue_set, outputs_view) = as StepOutputs>::build_queues( &[QueueSpec::ByteBounded { limit_bytes: 1000 }], &[BranchOrdering::None], ); + let outputs: OutputHandles> = OutputHandles::new(outputs_view); + outputs.push(7).unwrap(); + let input = queue_set.take_typed_input::(0); + assert_eq!(input.pop(), Some(7)); } #[test] diff --git a/crates/fgumi-pipeline-io/Cargo.toml b/crates/fgumi-pipeline-io/Cargo.toml new file mode 100644 index 000000000..92b109d85 --- /dev/null +++ b/crates/fgumi-pipeline-io/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "fgumi-pipeline-io" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +description = "BAM pipeline I/O types and steps for fgumi" +repository.workspace = true +license.workspace = true + +[dependencies] +fgumi-bam-io = { workspace = true } +fgumi-bgzf = { workspace = true } +fgumi-pipeline-core = { workspace = true } +fgumi-raw-bam = { workspace = true, features = ["noodles"] } +fgumi-sort = { workspace = true } +anyhow = "1.0.102" +log = "0" +noodles = { version = "0.111.0", features = ["bam", "sam", "bgzf"] } +parking_lot = "0.12" +tempfile = "3.4" + +[dev-dependencies] +fgumi-bam-io = { workspace = true } +fgumi-raw-bam = { workspace = true, features = ["noodles", "test-utils"] } +tempfile = "3.4" +noodles = { version = "0.111.0", features = ["bam", "sam"] } +fgumi-bgzf = { workspace = true } +rstest = "0" diff --git a/crates/fgumi-pipeline-io/src/lib.rs b/crates/fgumi-pipeline-io/src/lib.rs new file mode 100644 index 000000000..d3ff9cf58 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/lib.rs @@ -0,0 +1,15 @@ +#![deny(unsafe_code)] + +pub mod sink; +pub mod sort; +pub mod source; +pub mod types; + +pub use fgumi_pipeline_core::HeaderHandle; +pub use sink::write_bgzf::WriteBgzfFile; +pub use sort::{SortAndSpill, SortBamFile, SortMerge, SortSpillDecompress}; +pub use source::read_bam::{ + DEFAULT_BLOCKS_PER_BATCH, ReadBgzfBlocks, read_bam, read_bam_auto, read_bam_from_reader, + read_bam_stdin, +}; +pub use types::{BgzfBlock, DecompressedBlock, RecordBatch, RecordBatchBuilder}; diff --git a/crates/fgumi-pipeline-io/src/sink/mod.rs b/crates/fgumi-pipeline-io/src/sink/mod.rs new file mode 100644 index 000000000..201bc6e29 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sink/mod.rs @@ -0,0 +1 @@ +pub mod write_bgzf; diff --git a/crates/fgumi-pipeline-io/src/sink/write_bgzf.rs b/crates/fgumi-pipeline-io/src/sink/write_bgzf.rs new file mode 100644 index 000000000..0a65b8ecd --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sink/write_bgzf.rs @@ -0,0 +1,301 @@ +//! `WriteBgzfFile` sink step. `Serial` + `Affinity::Writer`. Receives +//! pre-compressed `BgzfBlock`s from `BgzfCompress` and writes them +//! directly to disk. + +use std::fs::File; +use std::io::{self, BufWriter, Write}; +use std::path::Path; + +use fgumi_bgzf::{BGZF_EOF, InlineBgzfCompressor}; +use noodles::sam::Header; +use parking_lot::Mutex; + +use crate::types::BgzfBlock; +use fgumi_pipeline_core::{ + header::HeaderHandle, + step::{Affinity, Step, StepCtx, StepKind, StepOutcome, StepProfile}, +}; + +/// `Exclusive + sticky` BAM sink that consumes pre-compressed `BgzfBlock`s. +pub struct WriteBgzfFile { + state: Mutex>, + name: &'static str, +} + +struct WriterState { + out: BufWriter, + pending_header: Option, +} + +struct PendingHeader { + handle: HeaderHandle, + compression_level: u32, +} + +impl WriteBgzfFile { + /// Open `path`, BGZF-compress and write the BAM header bytes, return + /// the sink ready to receive `BgzfBlock`s. + /// + /// # Errors + /// + /// Returns I/O errors from path open or header write. + pub fn new>( + path: P, + header: &Header, + compression_level: u32, + ) -> io::Result { + let file = File::create(path.as_ref())?; + let mut out = BufWriter::with_capacity(256 * 1024, file); + + let mut header_bytes = Vec::new(); + fgumi_bam_io::write_bam_header(&mut header_bytes, header) + .map_err(|e| io::Error::other(format!("write_bam_header: {e}")))?; + + let mut hc = InlineBgzfCompressor::new(compression_level); + hc.write_all(&header_bytes)?; + hc.flush()?; + hc.write_blocks_to(&mut out)?; + + Ok(Self { + state: Mutex::new(Some(WriterState { out, pending_header: None })), + name: "WriteBgzfFile", + }) + } + + /// Open `path` and return the sink with the BAM header write + /// deferred until an upstream step resolves `handle`. + /// + /// # Errors + /// + /// Returns I/O errors from path open. Header-write errors are + /// surfaced from `try_run` once the handle resolves. + pub fn new_with_handle>( + path: P, + handle: HeaderHandle, + compression_level: u32, + ) -> io::Result { + let file = File::create(path.as_ref())?; + let out = BufWriter::with_capacity(256 * 1024, file); + Ok(Self { + state: Mutex::new(Some(WriterState { + out, + pending_header: Some(PendingHeader { handle, compression_level }), + })), + name: "WriteBgzfFile", + }) + } + + fn try_write_pending_header(state: &mut WriterState) -> io::Result { + let Some(pending) = state.pending_header.as_ref() else { + return Ok(true); + }; + let header_clone = match pending.handle.try_get() { + None => return Ok(false), + Some(Err(e)) => return Err(e), + Some(Ok(h)) => h.clone(), + }; + let level = pending.compression_level; + + let mut header_bytes = Vec::new(); + fgumi_bam_io::write_bam_header(&mut header_bytes, &header_clone) + .map_err(|e| io::Error::other(format!("write_bam_header: {e}")))?; + let mut hc = InlineBgzfCompressor::new(level); + hc.write_all(&header_bytes)?; + hc.flush()?; + hc.write_blocks_to(&mut state.out)?; + + state.pending_header = None; + Ok(true) + } +} + +impl Step for WriteBgzfFile { + type Input = BgzfBlock; + type Outputs = (); + + fn profile(&self) -> StepProfile { + StepProfile { + name: self.name, + kind: StepKind::Serial, + sticky: true, + output_queues: vec![], + branch_ordering: vec![], + } + } + + fn affinity(&self) -> Affinity { + Affinity::Writer + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + let mut guard = self.state.lock(); + let Some(state) = guard.as_mut() else { + return Ok(StepOutcome::Finished); + }; + + let header_ready = Self::try_write_pending_header(state)?; + if header_ready { + if let Some(block) = ctx.input.pop() { + state.out.write_all(&block.bytes)?; + return Ok(StepOutcome::Progress); + } + } + + if ctx.input.is_drained() { + if !header_ready { + return Err(io::Error::other( + "WriteBgzfFile: input drained before HeaderHandle was resolved", + )); + } + state.out.write_all(&BGZF_EOF)?; + state.out.flush()?; + let _ = guard.take(); + return Ok(StepOutcome::Finished); + } + Ok(StepOutcome::NoProgress) + } +} + +impl Drop for WriteBgzfFile { + /// Cleanup-only drop. The BGZF EOF marker is written **exclusively** by the + /// drained-finish path in `try_run` (which then takes the state so this drop + /// is a no-op for a normally-finished sink). If state is still present here + /// the sink was dropped before that path ran — i.e. an aborted/partial + /// stream — so we deliberately do **not** append `BGZF_EOF`: stamping the + /// EOF marker onto a truncated BAM would make it look like a complete stream + /// and hide the truncation from downstream readers. We only flush whatever + /// bytes were already buffered so the on-disk file reflects what was written + /// (and stays detectably truncated). A still-pending header means nothing + /// valid was written, so leave the file empty. + fn drop(&mut self) { + let mut guard = self.state.lock(); + if let Some(mut state) = guard.take() { + if state.pending_header.is_some() { + return; + } + let _ = state.out.flush(); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn empty_header() -> Header { + Header::default() + } + + #[test] + fn profile_advertises_serial_writer_sink() { + let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); + let header = empty_header(); + let step = WriteBgzfFile::new(&path, &header, 1).unwrap(); + let profile = step.profile(); + assert_eq!(profile.name, "WriteBgzfFile"); + assert_eq!(profile.kind, StepKind::Serial); + assert!(profile.sticky); + assert_eq!(step.affinity(), Affinity::Writer); + assert_eq!(profile.output_queues.len(), 0); + assert_eq!(profile.branch_ordering.len(), 0); + } + + #[test] + fn header_only_round_trip() { + let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); + let header = empty_header(); + let step = WriteBgzfFile::new(&path, &header, 1).unwrap(); + let mut guard = step.state.lock(); + let mut state = guard.take().expect("state present"); + state.out.write_all(&BGZF_EOF).unwrap(); + state.out.flush().unwrap(); + drop(guard); + + let bytes = std::fs::read(&path).unwrap(); + assert!(bytes.len() >= 28, "BGZF EOF + header should be at least 28 bytes"); + assert_eq!(&bytes[0..2], &[0x1f, 0x8b], "BGZF/gzip magic at start"); + let tail = &bytes[bytes.len() - 28..]; + assert_eq!(tail, &BGZF_EOF, "file ends with BGZF EOF marker"); + } + + #[test] + fn new_with_handle_defers_header_until_resolved() { + let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); + let handle = HeaderHandle::new(); + let step = WriteBgzfFile::new_with_handle(&path, handle.clone(), 1).unwrap(); + + let bytes_before = std::fs::read(&path).unwrap(); + assert_eq!(bytes_before.len(), 0, "no bytes written until header resolves"); + + { + let mut guard = step.state.lock(); + let state = guard.as_mut().expect("state present"); + assert!(state.pending_header.is_some(), "handle still pending"); + let wrote = WriteBgzfFile::try_write_pending_header(state).unwrap(); + assert!(!wrote, "unresolved handle should yield without writing"); + assert!(state.pending_header.is_some(), "still pending after no-op probe"); + } + let bytes_mid = std::fs::read(&path).unwrap(); + assert_eq!(bytes_mid.len(), 0, "still nothing on disk after no-op probe"); + + handle.set(empty_header()).expect("first set"); + { + let mut guard = step.state.lock(); + let state = guard.as_mut().expect("state present"); + let wrote = WriteBgzfFile::try_write_pending_header(state).unwrap(); + assert!(wrote, "resolved handle should write"); + assert!(state.pending_header.is_none(), "pending slot cleared"); + state.out.flush().unwrap(); + } + let bytes_after = std::fs::read(&path).unwrap(); + assert!(bytes_after.len() >= 2, "header BGZF block emitted"); + assert_eq!(&bytes_after[0..2], &[0x1f, 0x8b], "BGZF/gzip magic at start"); + } + + #[test] + fn new_with_handle_propagates_poison() { + let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); + let handle = HeaderHandle::new(); + let step = WriteBgzfFile::new_with_handle(&path, handle.clone(), 1).unwrap(); + + handle.poison(io::Error::new(io::ErrorKind::BrokenPipe, "aligner died")).unwrap(); + let mut guard = step.state.lock(); + let state = guard.as_mut().expect("state present"); + let err = WriteBgzfFile::try_write_pending_header(state).expect_err("poison"); + assert_eq!(err.kind(), io::ErrorKind::BrokenPipe); + assert_eq!(err.to_string(), "aligner died"); + } + + #[test] + fn drop_before_finish_does_not_append_eof_marker() { + // A sink dropped before the drained-finish path in `try_run` (e.g. a + // pipeline abort) must NOT append the BGZF EOF marker. Appending it + // would stamp a "complete stream" signature onto a truncated BAM, + // hiding the truncation from downstream readers. See the `Drop` doc. + let path = tempfile::NamedTempFile::new().unwrap(); + let path_buf = path.path().to_path_buf(); + let header = empty_header(); + // `new` eagerly writes the header (pending_header is None), so the + // only thing standing between this state and a valid EOF marker is + // the `try_run` drained-finish path, which we never reach. + let step = WriteBgzfFile::new(&path_buf, &header, 1).unwrap(); + drop(step); + + let bytes = std::fs::read(&path_buf).unwrap(); + assert!(bytes.len() >= 28, "header bytes should be on disk"); + let tail = &bytes[bytes.len() - 28..]; + assert_ne!(tail, &BGZF_EOF, "aborted output must not end with a valid BGZF EOF marker"); + } + + #[test] + fn drop_with_unresolved_handle_leaves_empty_file() { + let path = tempfile::NamedTempFile::new().unwrap(); + let path_buf = path.path().to_path_buf(); + let handle = HeaderHandle::new(); + let step = WriteBgzfFile::new_with_handle(&path_buf, handle, 1).unwrap(); + drop(step); + + let bytes = std::fs::read(&path_buf).unwrap(); + assert_eq!(bytes.len(), 0, "Drop with unresolved handle must skip EOF — see Drop doc"); + } +} diff --git a/src/lib/pipeline/steps/sort/and_spill.rs b/crates/fgumi-pipeline-io/src/sort/and_spill.rs similarity index 53% rename from src/lib/pipeline/steps/sort/and_spill.rs rename to crates/fgumi-pipeline-io/src/sort/and_spill.rs index 18d147a18..bf2cc38d5 100644 --- a/src/lib/pipeline/steps/sort/and_spill.rs +++ b/crates/fgumi-pipeline-io/src/sort/and_spill.rs @@ -1,36 +1,4 @@ //! `SortAndSpill` — first step of the runall-sort three-step chain. -//! -//! ```text -//! [RecordBatch] → SortAndSpill ──SortPhase1Event──> SortSpillDecompress ──SortPhase2Event──> SortMerge → [RecordBatch] -//! (Serial) (Parallel) (Serial) -//! ``` -//! -//! Drives the streaming-sort engine's ingestion + spill phase -//! (`*SortStream::push_records` and internal spill management). On input -//! drain, finalizes via `*SortStream::into_slot_setup`, converts the -//! resulting `SlotSetup` into `SortPhase1Event`s, and emits them -//! downstream: -//! -//! * One `SpillReady` event per spill file (carries the shared -//! `Arc` constructed by the slot-setup conversion). -//! * One `MemoryChunk` event per residual in-memory chunk (after the -//! par-sort prologue). -//! -//! See `docs/design/sort-step-split.md` for the locked design. -//! -//! ## State machine -//! -//! 1. **`Ingesting`** — the streaming-sort handle is collecting input -//! `RecordBatch`es and managing internal spills via `push_records`. -//! Transitions to `Emitting` when `try_run` observes -//! `ctx.input.is_drained()`. -//! 2. **`Emitting`** — the conversion to `SlotSetup` has completed and -//! we hold a `VecDeque` of events to push downstream -//! plus the spill-directory `Vec` (kept alive for the -//! duration of the merge — slot readers reference files inside). -//! Transitions to `Done` when the queue is empty. -//! 3. **`Done`** — terminal; `try_run` returns `Finished` (every event has -//! been emitted and the step will never push again). use std::collections::VecDeque; use std::io; @@ -44,34 +12,28 @@ use fgumi_sort::{ use noodles::sam::Header; use tempfile::TempDir; -use crate::pipeline::core::Unpushed; -use crate::pipeline::core::held::HeldSlot; -use crate::pipeline::core::outputs::Single; -use crate::pipeline::core::queues::QueueSpec; -use crate::pipeline::core::reorder::BranchOrdering; -use crate::pipeline::core::step::{Affinity, Step, StepCtx, StepKind, StepOutcome, StepProfile}; -use crate::pipeline::steps::sort::protocol::{MemoryChunkErased, SortPhase1Event}; -use crate::pipeline::steps::types::RecordBatch; +use crate::sort::protocol::{MemoryChunkErased, SortPhase1Event}; +use crate::types::RecordBatch; +use fgumi_pipeline_core::{ + Unpushed, + held::HeldSlot, + outputs::Single, + queues::QueueSpec, + reorder::BranchOrdering, + step::{Affinity, Step, StepCtx, StepKind, StepOutcome, StepProfile}, +}; -/// Max input batches consumed per `try_run` invocation in the `Ingesting` -/// state. Bounds how long this worker holds the Sort lock before yielding -/// to the round-robin so other steps make progress. +/// Max input batches consumed per `try_run` invocation in the `Ingesting` state. const MAX_INGEST_BATCHES_PER_LOCK: usize = 8; /// Max events emitted per `try_run` invocation in the `Emitting` state. -/// Keeps the per-call work bounded; events are typically a handful, so -/// this is rarely the bottleneck. const MAX_EVENTS_PER_LOCK: usize = 8; -/// Return type of `SortStream::finalize` and the per-queryname-variant -/// helpers. Factored out because the inner tuple is too wide for clippy's -/// `type_complexity` threshold. +/// Return type of `SortStream::finalize`. type FinalizeResult = (Vec>, Vec, Vec, u64); -/// Wraps the three `*SortStream` variants so `SortAndSpill` can hold one -/// concrete state regardless of sort order. Boxed (see `SortAndSpillState`) -/// because the variants are ~hundreds of bytes. +/// Wraps the three `*SortStream` variants. enum SortStream { Coordinate(CoordinateSortStream), Queryname(QuerynameSortStream), @@ -87,11 +49,6 @@ impl SortStream { } } - /// Finalize the stream: drain pending spill, par-sort residual buffer - /// (matching legacy tie-break), open `Arc`s for each - /// spill file. Returns the slots, K-erased memory chunks (as - /// `Vec`), and the temp-dir RAII handles that - /// must outlive the slots' readers. fn finalize(self) -> Result { match self { Self::Coordinate(s) => { @@ -140,10 +97,6 @@ fn finalize_queryname_natural( } fn build_stream(sorter: RawExternalSorter, header: &Header) -> Result { - // The streaming-sort path uses whatever spill codec the sorter is configured - // for (main's default is zstd, #341). `SortSpillDecompress` detects each - // spill chunk's codec from its file magic (`slot.codec`) and decodes BGZF - // blocks or zstd frames accordingly, so no codec pinning is needed here. match sorter.sort_order() { SortOrder::Coordinate => Ok(SortStream::Coordinate(sorter.into_coordinate_stream(header)?)), SortOrder::Queryname(_) => Ok(SortStream::Queryname(sorter.into_queryname_stream(header)?)), @@ -155,45 +108,25 @@ fn build_stream(sorter: RawExternalSorter, header: &Header) -> Result), - Emitting { - pending_events: VecDeque, - /// RAII for the spill directories. Slots' `BufReader` - /// references point inside these directories; the temp dirs - /// must live until the slots are dropped (after `SortMerge` - /// finishes consuming them). - /// - /// Held inside `SortAndSpill` rather than embedded in each - /// event because the `SortAndSpill` step instance is alive for - /// the entire pipeline run. - _temp_dirs: Vec, - }, + Emitting { pending_events: VecDeque, _temp_dirs: Vec }, Done, } -/// `Serial` step that ingests `RecordBatch`es into a streaming sort -/// engine, manages internal spills, and on input drain emits -/// `SortPhase1Event`s describing the spill files + residual in-memory -/// chunks to the downstream `SortSpillDecompress` step. +/// `Serial` step that ingests `RecordBatch`es into a streaming sort engine. pub struct SortAndSpill { state: SortAndSpillState, held: HeldSlot>, - output_capacity: usize, - /// Running `records_ingested` counter — snapshotted into each - /// emitted event. Stamped onto every spill ready by the streaming - /// sort engine's `total_records` after `into_slot_setup`; for - /// memory chunks we use the same value (records ingested at - /// finalization time equals the final total). - /// - /// V1 (this commit): all events emitted on input drain, so the - /// snapshot is always the final value. V2 (Phase 1/Phase 2 overlap - /// follow-up) will move `SpillReady` emission into `try_run` per - /// spill close and snapshot the counter at each spill. + output_byte_limit: u64, affinity: Affinity, } impl SortAndSpill { - /// Build a `SortAndSpill` from a configured `RawExternalSorter` and - /// the output `Header`. + /// Build a `SortAndSpill` from a configured `RawExternalSorter` and the output `Header`. + /// + /// `output_byte_limit` byte-bounds the output event queue. The emitted + /// `SortPhase1Event::MemoryChunk` variant retains sorted record chunks, so + /// this queue must budget on bytes (`HeapSize`), not event count, to keep + /// retained memory a function of configuration. /// /// # Errors /// @@ -201,18 +134,18 @@ impl SortAndSpill { pub fn from_sorter( sorter: RawExternalSorter, header: &Header, - output_capacity: usize, + output_byte_limit: u64, ) -> Result { let stream = build_stream(sorter, header)?; Ok(Self { state: SortAndSpillState::Ingesting(Box::new(stream)), held: HeldSlot::new(), - output_capacity, + output_byte_limit, affinity: Affinity::None, }) } - /// Override the affinity hint. Not load-bearing for correctness. + /// Override the affinity hint. #[must_use] pub fn with_affinity(mut self, affinity: Affinity) -> Self { self.affinity = affinity; @@ -232,32 +165,9 @@ impl SortAndSpill { } } - /// Consume the `Ingesting` state's `SortStream`, finalize it, and - /// build the `pending_events` queue (plus its companion - /// `_temp_dirs` RAII bundle) that the `Emitting` state will push - /// to `out_chan`. - /// - /// Returns `Ok(None)` for empty inputs — no slots, no memory - /// chunks. Callers leave `self.state == Done` (already set by the - /// caller's `mem::replace`) and return early. - /// - /// Called from `try_run` once `ctx.input.is_drained()` is observed. The - /// block was previously inlined at two call sites (the drain-detected path - /// and a now-deleted `on_input_drained` race fallback); factoring it out - /// kept the two in sync. - /// - /// `caller_label` is interpolated into the error message so - /// future post-mortems can attribute a finalize failure to the - /// correct call site without changing the error surface. - /// /// # Panics /// - /// Panics if the number of spill slots exceeds `u32::MAX`. Each slot - /// corresponds to one on-disk spill file, so reaching `u32::MAX` - /// (~4 billion) slots is physically unreachable in any real run — the - /// process would exhaust file descriptors and disk long before. The - /// `u32` cast is required because `slot_count` is carried in the - /// `AllAnnounced` sentinel as a `u32`. + /// Panics if the number of spill slots exceeds `u32::MAX`. fn finalize_into_pending( stream: SortStream, caller_label: &str, @@ -287,15 +197,8 @@ impl SortAndSpill { }); } if pending_events.is_empty() { - // Empty input — no slots, no memory chunks. Don't emit - // even AllAnnounced; SortMerge falls through to its - // `is_drained` fallback with zero sources. return Ok(None); } - // AllAnnounced sentinel — always emitted LAST so SortMerge - // can transition WaitingForSetup → Merging as soon as the - // count predicate matches. See - // `docs/design/sort-step-split-parity-fix.md` Change 1. pending_events.push_back(SortPhase1Event::AllAnnounced { slot_count, memory_chunk_count, @@ -304,15 +207,9 @@ impl SortAndSpill { Ok(Some((pending_events, temp_dirs))) } - /// `try_run` Phase 1: pop up to `MAX_INGEST_BATCHES_PER_LOCK` - /// batches from upstream and push records into the `SortStream`. - /// Returns `true` if any records were ingested (caller returns - /// `Progress`), `false` if upstream was empty this call. - /// /// # Panics /// - /// Panics if `self.state` is not `Ingesting`. Callers must check - /// the state before invoking. + /// Panics if `self.state` is not `Ingesting`. fn ingest_batches(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { let SortAndSpillState::Ingesting(stream) = &mut self.state else { unreachable!("ingest_batches called outside Ingesting state"); @@ -328,10 +225,6 @@ impl SortAndSpill { Ok(did_work) } - /// If currently in `Ingesting`, finalize the stream and transition - /// to `Emitting`. On empty input, leave state at `Done` (set by - /// the internal `mem::replace`). No-op when already in `Emitting` - /// or `Done`. Called from `try_run` once the input is drained. fn transition_to_emitting(&mut self, caller_label: &str) -> io::Result<()> { if !matches!(&self.state, SortAndSpillState::Ingesting(_)) { return Ok(()); @@ -346,21 +239,12 @@ impl SortAndSpill { { self.state = SortAndSpillState::Emitting { pending_events, _temp_dirs: temp_dirs }; } - // else: empty input — state stays Done. Ok(()) } - /// Cooperative event emit. Push up to `MAX_EVENTS_PER_LOCK` - /// events from the `Emitting` state's `pending_events` queue, - /// stashing the first rejected push in `held`, returning the - /// appropriate `StepOutcome` so the framework can interleave - /// other steps' dispatches. Transitions state to `Done` when - /// the queue is empty; `try_run` then reports `Finished`. - /// /// # Panics /// - /// Panics if `self.state` is not `Emitting`. Callers must check - /// the state before invoking. + /// Panics if `self.state` is not `Emitting`. fn emit_pending_cooperative(&mut self, ctx: &mut StepCtx<'_, Self>) -> StepOutcome { let mut emitted = 0usize; let drained; @@ -394,7 +278,7 @@ impl Step for SortAndSpill { name: "SortAndSpill", kind: StepKind::Serial, sticky: false, - output_queues: vec![QueueSpec::CountBounded { capacity: self.output_capacity }], + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], branch_ordering: vec![BranchOrdering::None], } } @@ -408,10 +292,6 @@ impl Step for SortAndSpill { return Ok(StepOutcome::Contention); } - // Phase 1: ingest from upstream while we're in Ingesting. If - // upstream is drained AND we ingested nothing this call, - // transition Ingesting → Emitting (or directly to Done on - // empty input) and fall through. if matches!(&self.state, SortAndSpillState::Ingesting(_)) { if self.ingest_batches(ctx)? { return Ok(StepOutcome::Progress); @@ -422,11 +302,6 @@ impl Step for SortAndSpill { self.transition_to_emitting("try_run")?; } - // Phase 2: cooperative emit if we're in Emitting. Done reports - // `Finished` — the step has emitted every event and will never push - // again (it only transitions to Emitting on input drain, so reaching - // Done means input is drained and all events are out). Ingesting is - // unreachable — Phase 1 either returned early or transitioned us out. match &self.state { SortAndSpillState::Emitting { .. } => Ok(self.emit_pending_cooperative(ctx)), SortAndSpillState::Done => Ok(StepOutcome::Finished), diff --git a/crates/fgumi-pipeline-io/src/sort/and_spill/tests.rs b/crates/fgumi-pipeline-io/src/sort/and_spill/tests.rs new file mode 100644 index 000000000..1094133c5 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/and_spill/tests.rs @@ -0,0 +1 @@ +// No unit tests for and_spill directly; integration tests live in sort/tests.rs. diff --git a/src/lib/pipeline/steps/sort/merge.rs b/crates/fgumi-pipeline-io/src/sort/merge.rs similarity index 51% rename from src/lib/pipeline/steps/sort/merge.rs rename to crates/fgumi-pipeline-io/src/sort/merge.rs index 4fa181830..74383fca2 100644 --- a/src/lib/pipeline/steps/sort/merge.rs +++ b/crates/fgumi-pipeline-io/src/sort/merge.rs @@ -1,57 +1,4 @@ //! `SortMerge` — third step of the runall-sort three-step chain. -//! -//! ```text -//! [RecordBatch] → SortAndSpill ──SortPhase1Event──> SortSpillDecompress ──SortPhase2Event──> SortMerge → [RecordBatch] -//! (Serial) (Parallel) (Serial) -//! ``` -//! -//! Consumes `SortPhase2Event`s from the parallel decompress stage, drives -//! a non-blocking `MergeDriverDyn` over the slot table + memory chunks, and -//! emits sorted `RecordBatch`es downstream. The decompressed BGZF block -//! bytes flow out-of-band on each `SortMergeSlot`'s bounded queue (pushed by -//! `SortSpillDecompress`, popped by the merge driver) — only the slot -//! handles and counts ride the typed `SortPhase2Event` chain. -//! -//! See `docs/design/sort-step-split.md` for the locked design and -//! `docs/design/sort-merge-nonblocking.md` for the cooperative, resumable -//! merge consumer (issue #330 BUG #3). -//! -//! ## State machine -//! -//! 1. **`WaitingForSetup`** — collect upstream events. `SpillReady` events -//! contribute their `slot` to the slot table (deduplicated by -//! `file_id`); `MemoryChunk` events accumulate as typed memory chunks; -//! `AllAnnounced` carries the final counts. We `max()` -//! `records_ingested_so_far` from every event so `total_records` is -//! order-independent. Transitions to `Merging` once the announced -//! slot/chunk counts have all arrived (or, as a fallback, on input -//! drain). -//! -//! 2. **`Merging`** — build a typed `MergeDriver` via -//! [`fgumi_sort::MergeDriver::from_slots`], type-erase to -//! [`fgumi_sort::MergeDriverDyn`], and pump -//! [`MergeDriverDyn::try_step`]. `try_step` is **non-blocking**: when a -//! slot's decompressed queue is momentarily empty (and not at EOF) it -//! returns [`fgumi_sort::MergeStep::Stalled`], and this step flushes any -//! ready batch and returns `Contention` — yielding the worker so it can -//! round-robin to `SortSpillDecompress` to refill the slot, then resume -//! the merge on a later dispatch. Each produced record is appended to a -//! [`RecordBatchBuilder`]; batches flush downstream when full -//! (`target_batch_count` or `output_byte_limit`). -//! -//! 3. **`Done`** — terminal. -//! -//! ## Why the merge must never block -//! -//! `MergeDriver` runs on whichever framework worker holds `SortMerge`'s -//! Serial-step lock. A blocking wait there would park that worker (and hold -//! the lock) until a sibling worker refilled the slot — which a single-worker -//! configuration can never do, deadlocking the chain (issue #330 BUG #3, -//! review finding M1). The non-blocking `try_step` keeps the merge -//! cooperative: the same worker that stalls can refill the slot itself. Lazy -//! priming inside `from_slots` extends the same property to driver -//! construction, so building the `LoserTree` over not-yet-decompressed slots -//! never blocks either. use std::collections::HashMap; use std::io; @@ -63,36 +10,23 @@ use fgumi_sort::{ RawQuerynameLexKey, SortMergeSlot, SortOrder, TemplateKey, }; -use crate::pipeline::core::Unpushed; -use crate::pipeline::core::held::HeldSlot; -use crate::pipeline::core::outputs::OrderedBytesSingle; -use crate::pipeline::core::queues::QueueSpec; -use crate::pipeline::core::reorder::BranchOrdering; -use crate::pipeline::core::step::{Affinity, Step, StepCtx, StepKind, StepOutcome, StepProfile}; -use crate::pipeline::steps::sort::protocol::{MemoryChunkErased, SortPhase2Event}; -use crate::pipeline::steps::types::{RecordBatch, RecordBatchBuilder}; +use crate::sort::protocol::{MemoryChunkErased, SortPhase2Event}; +use crate::types::{RecordBatch, RecordBatchBuilder}; +use fgumi_pipeline_core::{ + Unpushed, + held::HeldSlot, + outputs::OrderedBytesSingle, + queues::QueueSpec, + reorder::BranchOrdering, + step::{Affinity, Step, StepCtx, StepKind, StepOutcome, StepProfile}, +}; /// Default output batch size: 1024 records per emitted `RecordBatch`. -/// Matches the legacy `Sort::DEFAULT_TARGET_BATCH_COUNT`. pub const DEFAULT_TARGET_BATCH_COUNT: usize = 1024; -/// Max input events drained per `try_run` invocation in the -/// `WaitingForSetup` state. Bounds the per-call work so other Serial -/// steps in the chain make progress. -const MAX_EVENTS_PER_LOCK: usize = 16; - -/// Max output batches emitted per `try_run` invocation in the `Merging` -/// state. Same intent as `MAX_EVENTS_PER_LOCK`. +/// Max output batches emitted per `try_run` invocation in `Merging`. const MAX_DRAIN_BATCHES_PER_LOCK: usize = 8; -/// Type-erased memory-chunk accumulator. Per-variant `Vec`s grow as -/// `MemoryChunk` events arrive in `WaitingForSetup`; on transition to -/// `Merging` exactly one variant is non-empty (matching `sort_order`) -/// and we hand its contents to the typed `MergeDriver::from_slots`. -/// -/// We don't gate which variant fills purely by `sort_order` — the -/// upstream `SortAndSpill` step is single-typed too — but the four -/// fields exist so the dispatch is statically typed. #[derive(Default)] struct MemoryChunksByKind { coordinate: Vec>, @@ -111,7 +45,6 @@ impl MemoryChunksByKind { } } - /// Total non-empty memory chunks across all four variants. fn total_len(&self) -> usize { self.coordinate.len() + self.queryname_lex.len() @@ -120,10 +53,6 @@ impl MemoryChunksByKind { } } -/// Build a typed merge driver from the collected slots + memory chunks -/// for the configured sort order. Construction is infallible and always -/// yields a driver; an all-empty input is reported by the driver's first -/// `try_step` returning [`MergeStep::Done`]. fn build_driver( sort_order: SortOrder, slots: Vec>, @@ -158,46 +87,21 @@ fn build_driver( } } -/// Result of one [`SortMerge::next_batch`] pump. enum NextBatch { - /// A full output batch is ready to push; more may follow immediately. Batch(RecordBatch), - /// The merge driver stalled on a not-yet-ready slot. Carries any - /// partially-filled batch to flush before yielding. Stalled(Option), - /// The merge is exhausted. Carries any final partial batch to flush, plus - /// the driver's total `records_merged()` — read in `next_batch` while the - /// driver is still in scope — for the completion log. Done(Option, u64), } -/// Three-phase state machine for [`SortMerge`]. enum SortMergeState { WaitingForSetup { - /// Slots received via `SpillReady` events. Sorted by `file_id` - /// before passing to `MergeDriver::from_slots` so the - /// `LoserTree` tie-break for equal sort keys is deterministic - /// and matches the legacy chunk-files order. slots: Vec>, - /// `slot.file_id -> index into slots` so deduplication is - /// O(1) per `SpillReady`. slot_index: HashMap, - /// Typed accumulators for the four sort-key variants. memory_chunks: MemoryChunksByKind, - /// Running `max(records_ingested_so_far)` across events. total_records: u64, - /// Expected slot count from `AllAnnounced`. Set once when - /// the sentinel arrives. The transition gate compares this - /// against `slots.len()`; when both match (and - /// `expected_memory_chunk_count` matches), we can transition - /// to `Merging` without waiting for `ctx.input.is_drained()`. expected_slot_count: Option, - /// Expected non-empty memory-chunk count from - /// `AllAnnounced`. Compared against `memory_chunks.total_len()`. expected_memory_chunk_count: Option, }, - /// Active merge state. Boxed because `dyn MergeDriverDyn` is a wide - /// pointer; `RecordBatchBuilder` carries a `Vec` allocation. Merging { driver: Box, builder: RecordBatchBuilder, @@ -206,9 +110,6 @@ enum SortMergeState { Done, } -/// Pop one `SortPhase2Event` and route it into the in-flight setup -/// accumulators. Factored out of `try_run`'s `WaitingForSetup` arm so -/// the latter stays under clippy's `too_many_lines` threshold. fn absorb_phase2_event( event: SortPhase2Event, slots: &mut Vec>, @@ -227,20 +128,14 @@ fn absorb_phase2_event( *total_records = (*total_records).max(records_ingested_so_far); } SortPhase2Event::MemoryChunk { chunk, records_ingested_so_far } => { - let inner = Arc::try_unwrap(chunk).unwrap_or_else(|arc| { - // Fallback: clone. Only hit if a future change adds a - // second `Arc` holder. - match arc.as_ref() { - MemoryChunkErased::Coordinate(v) => MemoryChunkErased::Coordinate(v.clone()), - MemoryChunkErased::QuerynameLex(v) => { - MemoryChunkErased::QuerynameLex(v.clone()) - } - MemoryChunkErased::QuerynameNatural(v) => { - MemoryChunkErased::QuerynameNatural(v.clone()) - } - MemoryChunkErased::TemplateCoordinate(v) => { - MemoryChunkErased::TemplateCoordinate(v.clone()) - } + let inner = Arc::try_unwrap(chunk).unwrap_or_else(|arc| match arc.as_ref() { + MemoryChunkErased::Coordinate(v) => MemoryChunkErased::Coordinate(v.clone()), + MemoryChunkErased::QuerynameLex(v) => MemoryChunkErased::QuerynameLex(v.clone()), + MemoryChunkErased::QuerynameNatural(v) => { + MemoryChunkErased::QuerynameNatural(v.clone()) + } + MemoryChunkErased::TemplateCoordinate(v) => { + MemoryChunkErased::TemplateCoordinate(v.clone()) } }); memory_chunks.push(inner); @@ -265,10 +160,6 @@ fn absorb_phase2_event( } } -/// Returns `true` if the slot set is complete: `AllAnnounced` has -/// arrived and the actual slot + memory-chunk counts match the -/// expected counts. Used to gate the early `WaitingForSetup → Merging` -/// transition without waiting for `ctx.input.is_drained()`. fn slot_set_complete( slots_len: usize, memory_chunks_total_len: usize, @@ -283,11 +174,7 @@ fn slot_set_complete( ) } -/// `Serial + ByItemOrdinal` middle-of-chain merge. Input = -/// `SortPhase2Event`, Output = sorted `RecordBatch`. Drives -/// [`fgumi_sort::MergeDriver`] via pausable `peek` / `advance` so the -/// merge phase runs interleaved with the framework's round-robin -/// dispatch — no dedicated OS thread. +/// `Serial + ByItemOrdinal` middle-of-chain merge. pub struct SortMerge { state: SortMergeState, held: HeldSlot>, @@ -328,17 +215,13 @@ impl SortMerge { } } - /// Override the affinity hint. Not load-bearing for correctness — the - /// pausable driver releases the Serial lock between bounded batches. + /// Override the affinity hint. #[must_use] pub fn with_affinity(mut self, affinity: Affinity) -> Self { self.affinity = affinity; self } - /// Try to deliver a previously-held batch. Returns `true` if the held - /// slot is now empty, `false` if it's still held (caller signals - /// `Contention`). fn flush_held(&mut self, ctx: &mut StepCtx<'_, Self>) -> bool { let Some(unpushed) = self.held.take() else { return true; @@ -352,31 +235,14 @@ impl SortMerge { } } - /// If currently in `WaitingForSetup`, sort the accumulated slots, - /// build the merge driver, and transition to `Merging` (or - /// directly to `Done` for empty input — `build_driver` returns - /// `Ok(None)` when there's nothing to merge). No-op when already - /// in `Merging` or `Done`. Called from `try_run` once the setup - /// predicate is satisfied or the input is drained. - /// - /// `caller_label` is interpolated into the error message on - /// `build_driver` failure so post-mortems can attribute the - /// failure to the correct call site. - /// `WaitingForSetup` Phase 1: pop events from upstream and absorb - /// each into the setup accumulators. `cap` bounds the per-call - /// absorb count; `try_run` passes `Some(MAX_EVENTS_PER_LOCK)` to keep - /// each dispatch bounded. Returns the count of events absorbed so - /// `try_run` can short-circuit on `did_work`. + /// Drains every currently-available input event into the setup state and + /// returns the number absorbed. The drain is intentionally unbounded — the + /// upstream queue is byte-bounded, so memory is gated on the producer side. /// /// # Panics /// - /// Panics if `self.state` is not `WaitingForSetup`. Callers must - /// check the state before invoking. - fn absorb_events_into_setup( - &mut self, - ctx: &mut StepCtx<'_, Self>, - cap: Option, - ) -> usize { + /// Panics if `self.state` is not `WaitingForSetup`. + fn absorb_events_into_setup(&mut self, ctx: &mut StepCtx<'_, Self>) -> usize { let SortMergeState::WaitingForSetup { slots, slot_index, @@ -389,8 +255,7 @@ impl SortMerge { unreachable!("absorb_events_into_setup called outside WaitingForSetup state"); }; let mut absorbed = 0usize; - while cap.is_none_or(|c| absorbed < c) { - let Some(event) = ctx.input.pop() else { break }; + while let Some(event) = ctx.input.pop() { absorb_phase2_event( event, slots, @@ -405,10 +270,6 @@ impl SortMerge { absorbed } - /// True when `self.state` is `WaitingForSetup` AND the - /// `AllAnnounced` predicates match. False in any other state. - /// Used by `try_run` to early-transition `WaitingForSetup -> - /// Merging` without waiting for upstream's drained flag. fn is_ready_to_merge(&self) -> bool { let SortMergeState::WaitingForSetup { slots, @@ -428,17 +289,9 @@ impl SortMerge { ) } - /// Pump the merge driver into the current batch builder via the - /// non-blocking [`MergeDriverDyn::try_step`] until the builder is - /// full (a `Batch` is ready), the driver stalls on a not-yet-ready - /// slot (`Stalled`), or the merge is exhausted (`Done`). For the - /// latter two, any partially-filled builder is flushed alongside so - /// the caller can push it before yielding / finishing. - /// /// # Panics /// - /// Panics if `self.state` is not `Merging`. Callers must check - /// the state before invoking. + /// Panics if `self.state` is not `Merging`. fn next_batch(&mut self) -> io::Result { let target = self.target_batch_count; let byte_limit = self.output_byte_limit; @@ -447,14 +300,11 @@ impl SortMerge { unreachable!("next_batch called outside Merging state"); }; - // Flush the current builder into a `RecordBatch`, advancing the - // ordinal and installing a fresh builder. let flush = |builder: &mut RecordBatchBuilder, next_ordinal: &mut u64| { *next_ordinal += 1; let next_builder = RecordBatchBuilder::with_capacity(*next_ordinal, bytes_cap, target); std::mem::replace(builder, next_builder).build() }; - // Flush only if the builder holds at least one record. let flush_partial = |builder: &mut RecordBatchBuilder, next_ordinal: &mut u64| { if builder.is_empty() { None } else { Some(flush(builder, next_ordinal)) } }; @@ -485,21 +335,11 @@ impl SortMerge { } } - /// Cooperative merge emit. Produce up to `MAX_DRAIN_BATCHES_PER_LOCK` - /// batches via `next_batch`, push each downstream, stash the first - /// rejection in `held` and return `Progress`. Transitions state to - /// `Done` when the driver is exhausted. - /// - /// Called from `try_run` only. - /// /// # Panics /// /// Panics if `self.state` is not `Merging`. fn emit_batches_cooperative(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { let mut delivered = 0usize; - // Push a batch downstream; on backpressure stash it in `held` and - // signal the caller to return `Progress` (so `flush_held` retries it - // next dispatch). Returns `false` if the push was held. loop { match self.next_batch()? { NextBatch::Batch(batch) => { @@ -520,14 +360,6 @@ impl SortMerge { } delivered += 1; } - // A slot's queue is momentarily empty. Yield so this - // worker round-robins to `SortSpillDecompress` to refill - // it, then resume the merge on a later dispatch. Either - // outcome yields: `Progress` (we delivered batches this - // call) triggers the framework's priority-restart at step - // 0, which subsumes the `Contention` retry — so reporting - // `Progress` when `delivered > 0` is the honest, stronger - // signal, and `Contention` covers the made-no-progress case. return Ok(if delivered > 0 { StepOutcome::Progress } else { @@ -538,15 +370,10 @@ impl SortMerge { if let Some(batch) = partial { if let Err(unpushed) = ctx.outputs.push(batch) { self.held.put(unpushed); - // Stay in `Merging`; next dispatch flushes `held`, - // then `next_batch` returns `Done(None)` → `Done`. return Ok(StepOutcome::Progress); } delivered += 1; } - // Surface the merged-record count carried up from `next_batch` - // (where the driver was in scope). The streaming sort's - // analogue of standalone sort's Records-written summary. log::info!("Sort merge complete: {merged} records merged"); self.state = SortMergeState::Done; return Ok(if delivered > 0 { @@ -559,11 +386,6 @@ impl SortMerge { } } - /// If currently `WaitingForSetup`, build the merge driver and move to - /// `Merging`. No-op in any other state. Infallible: `build_driver` always - /// yields a driver, and empty input is reported by the driver's first - /// `try_step` returning `Done` (handled by the `Merging` emit paths), so - /// there is no separate empty-input or error branch. fn transition_to_merging(&mut self) { if !matches!(&self.state, SortMergeState::WaitingForSetup { .. }) { return; @@ -579,11 +401,6 @@ impl SortMerge { else { unreachable!("just matched WaitingForSetup") }; - // Sort slots by `file_id` so `MergeDriver::from_slots` source - // order is deterministic and matches the legacy chunk-files - // order. `SpillReady` events arrive in `SortAndSpill`'s - // emission order (which IS `file_id` order today), but - // defensively sort to avoid relying on that. slots.sort_by_key(|s| s.file_id); let driver = build_driver(self.sort_order, slots, memory_chunks, total_records); let bytes_cap = usize::try_from(self.output_byte_limit).unwrap_or(usize::MAX); @@ -615,15 +432,13 @@ impl Step for SortMerge { return Ok(StepOutcome::Contention); } - // Phase 1: absorb events while in WaitingForSetup. Transition - // to Merging when AllAnnounced predicates match (early- - // transition gate — avoids waiting for upstream's drained flag - // which would require SortSpillDecompress workers to Skip - // themselves first, blocking further decompress work) OR - // when upstream is drained with zero slots/chunks (empty- - // input fall-back). if matches!(&self.state, SortMergeState::WaitingForSetup { .. }) { - let absorbed = self.absorb_events_into_setup(ctx, Some(MAX_EVENTS_PER_LOCK)); + // Drain the input queue unbounded: the setup absorb is cheap (it just + // moves `Arc`s/`Vec`s into the setup state) and the upstream queue is + // already byte-bounded, so backpressure belongs on the producer, not + // on a consumer-side drain cap. A cap here would cycle a full upstream + // queue through repeated partial drains and add producer contention. + let absorbed = self.absorb_events_into_setup(ctx); if !self.is_ready_to_merge() { if absorbed > 0 { return Ok(StepOutcome::Progress); @@ -631,20 +446,38 @@ impl Step for SortMerge { if !ctx.input.is_drained() { return Ok(StepOutcome::NoProgress); } + // Input is drained but setup never completed. Fail closed for any + // setup that saw payload or a (mismatched) `AllAnnounced`, so an + // incomplete setup can never silently merge a partial result. The + // only legitimate drained-but-not-ready case is a wholly empty + // input (no slots, no chunks, no announcement), which merges to an + // empty output. + let SortMergeState::WaitingForSetup { + slots, + memory_chunks, + expected_slot_count, + expected_memory_chunk_count, + .. + } = &self.state + else { + unreachable!("state matched WaitingForSetup above"); + }; + let saw_payload = !slots.is_empty() || memory_chunks.total_len() > 0; + let saw_expectations = + expected_slot_count.is_some() || expected_memory_chunk_count.is_some(); + if saw_payload || saw_expectations { + return Err(io::Error::other(format!( + "SortMerge: setup incomplete at input drain \ + (slots={}, chunks={}, expected_slots={expected_slot_count:?}, \ + expected_chunks={expected_memory_chunk_count:?})", + slots.len(), + memory_chunks.total_len(), + ))); + } } self.transition_to_merging(); } - // Phase 2: cooperative emit if Merging; Done reports `Finished`. - // `Done` means the merge driver is exhausted and every batch has - // landed — completion is keyed on the merge state, NOT on - // `ctx.input.is_drained()` (the input edge drains early, before the - // merge finishes, because the records come from spilled chunks / - // sibling decompress workers; finishing on `is_drained` would - // truncate the sorted output). Empty input reaches `Merging`, then - // `emit_batches_cooperative`'s first `next_batch` returns `Done(None)` - // → state `Done` → `Finished` on the next pass. WaitingForSetup is - // unreachable — Phase 1 either returned early or transitioned us out. match &self.state { SortMergeState::Merging { .. } => self.emit_batches_cooperative(ctx), SortMergeState::Done => Ok(StepOutcome::Finished), diff --git a/crates/fgumi-pipeline-io/src/sort/merge/tests.rs b/crates/fgumi-pipeline-io/src/sort/merge/tests.rs new file mode 100644 index 000000000..837159411 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/merge/tests.rs @@ -0,0 +1 @@ +// No unit tests for merge directly; integration tests live in sort/tests.rs. diff --git a/crates/fgumi-pipeline-io/src/sort/mod.rs b/crates/fgumi-pipeline-io/src/sort/mod.rs new file mode 100644 index 000000000..6a4323d1f --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/mod.rs @@ -0,0 +1,75 @@ +//! Sort typed-steps for the unified pipeline. + +pub mod and_spill; +pub mod merge; +pub mod protocol; +pub mod spill_decompress; + +pub use and_spill::SortAndSpill; +pub use merge::SortMerge; +pub use spill_decompress::SortSpillDecompress; + +use std::io; +use std::path::PathBuf; +use std::sync::Arc; + +use fgumi_sort::RawExternalSorter; +use parking_lot::Mutex; + +use fgumi_pipeline_core::step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}; + +// ───────────────────────────────────────────────────────────────────────────── +// SortBamFile — Exclusive single-step wrapping legacy sort end-to-end. +// ───────────────────────────────────────────────────────────────────────────── + +/// `Exclusive` step that drives a complete legacy +/// `RawExternalSorter::sort(input, output)` call to completion in a +/// single `try_run`. +pub struct SortBamFile { + sorter: Option, + input: PathBuf, + output: PathBuf, + stats_out: Arc>>, +} + +impl SortBamFile { + /// Build a `SortBamFile` step. + #[must_use] + pub fn new( + sorter: RawExternalSorter, + input: PathBuf, + output: PathBuf, + stats_out: Arc>>, + ) -> Self { + Self { sorter: Some(sorter), input, output, stats_out } + } +} + +impl Step for SortBamFile { + type Input = (); + type Outputs = (); + + fn profile(&self) -> StepProfile { + StepProfile { + name: "SortBamFile", + kind: StepKind::Exclusive, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + let Some(sorter) = self.sorter.take() else { + return Ok(StepOutcome::Finished); + }; + let stats = sorter + .sort(&self.input, &self.output) + .map_err(|e| io::Error::other(format!("SortBamFile: sort failed: {e:#}")))?; + *self.stats_out.lock() = Some(stats); + Ok(StepOutcome::Finished) + } +} + +#[cfg(test)] +pub mod tests; diff --git a/src/lib/pipeline/steps/sort/protocol.rs b/crates/fgumi-pipeline-io/src/sort/protocol.rs similarity index 61% rename from src/lib/pipeline/steps/sort/protocol.rs rename to crates/fgumi-pipeline-io/src/sort/protocol.rs index cd5c2107c..7c65c507a 100644 --- a/src/lib/pipeline/steps/sort/protocol.rs +++ b/crates/fgumi-pipeline-io/src/sort/protocol.rs @@ -1,31 +1,5 @@ //! Typed-event protocol between the three sort steps in the runall-sort //! fused chain. -//! -//! ```text -//! [RecordBatch] → SortAndSpill ──SortPhase1Event──> SortSpillDecompress ──SortPhase2Event──> SortMerge → [RecordBatch] -//! (Serial) (Parallel) (Serial) -//! ``` -//! -//! v4 design: Decompressed BGZF block bytes do NOT flow on the typed -//! event chain. They live exclusively in `slot.decompressed` (per-slot -//! bounded queue on `SortMergeSlot`); `SortSpillDecompress` pushes, -//! `SortMerge` pops. The events on the chain carry only: -//! -//! * `SpillReady` — registers a slot (with `Arc`) for -//! `SortSpillDecompress` to operate on AND for `SortMerge` to install -//! in its slot table. -//! * `MemoryChunk` — residual in-memory chunk emitted by `SortAndSpill`'s -//! drained-completion path; nothing for `SortSpillDecompress` to do -//! (forwarded verbatim). -//! * `AllAnnounced` — sentinel emitted LAST by `SortAndSpill`. Carries -//! the final counts so `SortMerge` can transition to `Merging` as -//! soon as it has received the announced slot/chunk count, without -//! waiting for `ctx.input.is_drained()` (which would require upstream -//! workers to Skip themselves first — see -//! `docs/design/sort-step-split-parity-fix.md`). -//! -//! See `docs/design/sort-step-split-parity-fix.md` for the full -//! locked design and rationale. use std::path::PathBuf; use std::sync::Arc; @@ -35,25 +9,13 @@ use fgumi_sort::{ RawCoordinateKey, RawQuerynameKey, RawQuerynameLexKey, SortMergeSlot, TemplateKey, }; -use crate::pipeline::core::item::HeapSize; +use fgumi_pipeline_core::item::HeapSize; -/// Approximate fixed overhead of a `Vec<(K, RawRecord)>` entry (key + record -/// header). Used by `MemoryChunkErased::heap_size` to estimate buffer-bounded -/// queue pressure when a memory chunk flows through the typed chain. -/// -/// Conservatively chosen to match the per-record memory budget used by -/// `RawExternalSorter` for spill triggering (`bytes_per_record = 354` in the -/// sort engine). Slightly over-counting is harmless — under-counting risks -/// queue overshoot. +/// Approximate fixed overhead of a `Vec<(K, RawRecord)>` entry. const PER_MEMORY_RECORD_OVERHEAD: usize = 354; /// In-memory sorted residual chunk produced by `SortAndSpill`, type-erased -/// over the sort-key variant `K` so it can flow through the typed-event -/// chain without `K` leaking into step signatures. -/// -/// `SortMerge` is constructed with `sort_order` known statically; it -/// `match`es on the variant to recover the typed `Vec<(K, RawRecord)>` for -/// `MergeDriver::from_slots`'s `memory_chunks` argument. +/// over the sort-key variant `K`. pub enum MemoryChunkErased { /// Coordinate-sort residual. `K = RawCoordinateKey`. Coordinate(Vec<(RawCoordinateKey, RawRecord)>), @@ -83,10 +45,7 @@ impl MemoryChunkErased { self.len() == 0 } - /// Approximate heap footprint in bytes. Sum of each record's payload - /// length plus a fixed per-entry overhead approximating `K + RawRecord` - /// inline-struct cost. Used by `HeapSize` for byte-bounded queue - /// accounting. + /// Approximate heap footprint in bytes. #[must_use] pub fn approx_heap_bytes(&self) -> usize { let (count, payload): (usize, usize) = match self { @@ -109,67 +68,55 @@ impl MemoryChunkErased { /// Events from `SortAndSpill` → `SortSpillDecompress`. pub enum SortPhase1Event { - /// A spill chunk file has been closed and is ready for Phase 2 - /// decompression. Carries the per-file shared state, the on-disk - /// path (informational), and a snapshot of - /// `records_ingested_so_far` at spill-close. + /// A spill chunk file has been closed and is ready for Phase 2 decompression. SpillReady { slot: Arc, path: PathBuf, records_ingested_so_far: u64 }, /// A par-sorted residual in-memory chunk produced by `SortAndSpill`'s - /// drained-completion path. Forwarded by `SortSpillDecompress` to - /// `SortMerge` (no decompression needed). + /// drained-completion path. MemoryChunk { chunk: Arc, records_ingested_so_far: u64 }, - /// Sentinel emitted as the LAST event by `SortAndSpill`'s drained-completion - /// path (after all `SpillReady` and `MemoryChunk` events). Carries the final - /// counts so `SortMerge` - /// can transition to `Merging` early — without waiting for - /// `ctx.input.is_drained()`. See - /// `docs/design/sort-step-split-parity-fix.md` Change 1. + /// Sentinel emitted as the LAST event by `SortAndSpill`'s drained-completion path. AllAnnounced { slot_count: u32, memory_chunk_count: u32, total_records: u64 }, } impl HeapSize for SortPhase1Event { fn heap_size(&self) -> usize { + // Charge a fixed per-event base so the byte-bounded transport queues + // cannot accept an unbounded count of near-zero-cost control events + // (`SpillReady` with an empty path, `AllAnnounced`). Memory stays a + // function of configuration rather than event count. + let base = std::mem::size_of::(); match self { - Self::SpillReady { path, .. } => path.as_os_str().len(), - Self::MemoryChunk { chunk, .. } => chunk.approx_heap_bytes(), - // Three small numeric fields plus the discriminant. - Self::AllAnnounced { .. } => 0, + Self::SpillReady { path, .. } => base + path.as_os_str().len(), + Self::MemoryChunk { chunk, .. } => base + chunk.approx_heap_bytes(), + Self::AllAnnounced { .. } => base, } } } /// Events from `SortSpillDecompress` → `SortMerge`. -/// -/// Same shape as `SortPhase1Event`; `SortSpillDecompress` forwards all -/// three variants verbatim. (Decompressed Block bytes do NOT flow on -/// this chain — they live in `slot.decompressed`.) pub enum SortPhase2Event { - /// Forwarded `SortPhase1Event::SpillReady` — registers the slot in - /// `SortMerge`'s slot table. + /// Forwarded `SortPhase1Event::SpillReady`. SpillReady { slot: Arc, path: PathBuf, records_ingested_so_far: u64 }, - /// Forwarded `SortPhase1Event::MemoryChunk`. `SortMerge` consumes - /// these directly as `ChunkSource::Memory` sources in - /// `MergeDriver::from_slots`. + /// Forwarded `SortPhase1Event::MemoryChunk`. MemoryChunk { chunk: Arc, records_ingested_so_far: u64 }, - /// Forwarded `SortPhase1Event::AllAnnounced`. Gates `SortMerge`'s - /// early transition from `WaitingForSetup` to `Merging`. + /// Forwarded `SortPhase1Event::AllAnnounced`. AllAnnounced { slot_count: u32, memory_chunk_count: u32, total_records: u64 }, } impl HeapSize for SortPhase2Event { fn heap_size(&self) -> usize { + // See `SortPhase1Event::heap_size`: a fixed per-event base keeps the + // byte-bounded queues from absorbing unbounded control-event counts. + let base = std::mem::size_of::(); match self { - Self::SpillReady { path, .. } => path.as_os_str().len(), - Self::MemoryChunk { chunk, .. } => chunk.approx_heap_bytes(), - Self::AllAnnounced { .. } => 0, + Self::SpillReady { path, .. } => base + path.as_os_str().len(), + Self::MemoryChunk { chunk, .. } => base + chunk.approx_heap_bytes(), + Self::AllAnnounced { .. } => base, } } } impl SortPhase1Event { - /// Running snapshot of records ingested at the moment this event was - /// emitted. `SortMerge` takes `max()` over all popped events to recover - /// the final `total_records` count. + /// Running snapshot of records ingested at the moment this event was emitted. #[must_use] pub fn records_ingested_so_far(&self) -> u64 { match self { @@ -263,18 +210,21 @@ mod tests { } #[test] - fn all_announced_heap_size_is_zero() { + fn all_announced_heap_size_charges_base_cost() { + // Control events carry no heap payload but must still cost a fixed, + // non-zero amount so the byte-bounded queues cannot accept an unbounded + // count of them. let ev1 = SortPhase1Event::AllAnnounced { slot_count: 16, memory_chunk_count: 4, total_records: 1_000_000, }; - assert_eq!(ev1.heap_size(), 0); + assert_eq!(ev1.heap_size(), std::mem::size_of::()); let ev2 = SortPhase2Event::AllAnnounced { slot_count: 16, memory_chunk_count: 4, total_records: 1_000_000, }; - assert_eq!(ev2.heap_size(), 0); + assert_eq!(ev2.heap_size(), std::mem::size_of::()); } } diff --git a/crates/fgumi-pipeline-io/src/sort/spill_decompress.rs b/crates/fgumi-pipeline-io/src/sort/spill_decompress.rs new file mode 100644 index 000000000..4a16d848b --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/spill_decompress.rs @@ -0,0 +1,218 @@ +//! `SortSpillDecompress` — Parallel typed step that reads spill chunk +//! files, decompresses them inline, and pushes the decompressed bytes +//! into per-slot bounded queues on `SortMergeSlot`. + +use std::io; +use std::sync::Arc; + +use fgumi_sort::{PHASE2_DECOMP_CAP, SortMergeSlot, SpillBlockDecompressor}; +use parking_lot::Mutex; + +use crate::sort::protocol::{SortPhase1Event, SortPhase2Event}; +use fgumi_pipeline_core::{ + Unpushed, + held::HeldSlot, + outputs::Single, + queues::QueueSpec, + reorder::BranchOrdering, + step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}, +}; + +/// Max raw blocks read+decompressed by a single `try_fill_some_slot` call. +const MAX_BATCH_PER_CALL: usize = 4; + +struct RegisteredSpill { + slot: Arc, +} + +/// Parallel step that reads + decompresses spill chunk files and +/// pushes results into per-slot queues. Forwards `SortPhase1Event`s +/// verbatim to `SortMerge`. +pub struct SortSpillDecompress { + registry: Arc>>, + block_dec: SpillBlockDecompressor, + held: HeldSlot>, + output_byte_limit: u64, +} + +impl SortSpillDecompress { + /// Construct a fresh step with an empty registry. + /// + /// `output_byte_limit` byte-bounds the forwarded-event output queue. + /// The forwarded `SortPhase2Event::MemoryChunk` variant retains sorted + /// record chunks, so this queue must budget on bytes (`HeapSize`), not + /// event count, to keep retained memory a function of configuration. + #[must_use] + pub fn new(output_byte_limit: u64) -> Self { + Self { + registry: Arc::new(Mutex::new(Vec::new())), + block_dec: SpillBlockDecompressor::new(), + held: HeldSlot::new(), + output_byte_limit, + } + } + + fn flush_held(&mut self, ctx: &mut StepCtx<'_, Self>) -> bool { + let Some(unpushed) = self.held.take() else { + return true; + }; + match ctx.outputs.retry(unpushed) { + Ok(()) => true, + Err(again) => { + self.held.put(again); + false + } + } + } + + fn push_or_hold(&mut self, ctx: &mut StepCtx<'_, Self>, event: SortPhase2Event) -> bool { + match ctx.outputs.push(event) { + Ok(()) => true, + Err(unpushed) => { + self.held.put(unpushed); + false + } + } + } + + fn snapshot_registry(&self) -> Vec> { + let registry = self.registry.lock(); + registry.iter().map(|e| Arc::clone(&e.slot)).collect() + } + + fn try_fill_some_slot(&mut self) -> io::Result { + for slot in self.snapshot_registry() { + if slot.queue_eof.load(std::sync::atomic::Ordering::Acquire) { + continue; + } + + let Ok(mut reader_guard) = slot.reader.try_lock() else { + continue; + }; + + let m = { + let dec = slot.decompressed.lock().expect("decompressed mutex poisoned"); + PHASE2_DECOMP_CAP.saturating_sub(dec.len()) + }; + if m == 0 { + continue; + } + let m = m.min(MAX_BATCH_PER_CALL); + + let decompressed_batch = + match self.block_dec.read_blocks(&mut reader_guard.inner, slot.codec, m) { + Ok(b) => b, + Err(e) => { + { + let _g = slot.decompressed.lock().expect("decompressed mutex poisoned"); + slot.decomp_error.store(true, std::sync::atomic::Ordering::Release); + slot.queue_eof.store(true, std::sync::atomic::Ordering::Release); + } + drop(reader_guard); + return Err(e); + } + }; + let got = decompressed_batch.len(); + let hit_eof = got < m; + + if got == 0 { + { + let _g = slot.decompressed.lock().expect("decompressed mutex poisoned"); + slot.queue_eof.store(true, std::sync::atomic::Ordering::Release); + } + drop(reader_guard); + return Ok(true); + } + + { + let mut dec = slot.decompressed.lock().expect("decompressed mutex poisoned"); + for b in decompressed_batch { + dec.push_back(b); + } + if hit_eof { + slot.queue_eof.store(true, std::sync::atomic::Ordering::Release); + } + } + drop(reader_guard); + return Ok(true); + } + Ok(false) + } +} + +impl Clone for SortSpillDecompress { + fn clone(&self) -> Self { + Self { + registry: Arc::clone(&self.registry), + block_dec: SpillBlockDecompressor::new(), + held: HeldSlot::new(), + output_byte_limit: self.output_byte_limit, + } + } +} + +impl Step for SortSpillDecompress { + type Input = SortPhase1Event; + type Outputs = Single; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "SortSpillDecompress", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], + branch_ordering: vec![BranchOrdering::None], + } + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + // 1. Drain held output first. + if !self.flush_held(ctx) { + return Ok(StepOutcome::Contention); + } + + // 2. Pop one input event, register if SpillReady, forward all. + if let Some(event) = ctx.input.pop() { + let forwarded = match event { + SortPhase1Event::SpillReady { slot, path, records_ingested_so_far } => { + self.registry.lock().push(RegisteredSpill { slot: Arc::clone(&slot) }); + SortPhase2Event::SpillReady { slot, path, records_ingested_so_far } + } + SortPhase1Event::MemoryChunk { chunk, records_ingested_so_far } => { + SortPhase2Event::MemoryChunk { chunk, records_ingested_so_far } + } + SortPhase1Event::AllAnnounced { slot_count, memory_chunk_count, total_records } => { + SortPhase2Event::AllAnnounced { slot_count, memory_chunk_count, total_records } + } + }; + let _ = self.push_or_hold(ctx, forwarded); + return Ok(StepOutcome::Progress); + } + + // 3. Greedy slot-fill step. + if self.try_fill_some_slot()? { + return Ok(StepOutcome::Progress); + } + + // 4. No fill work. + let any_alive = self + .snapshot_registry() + .iter() + .any(|slot| !slot.queue_eof.load(std::sync::atomic::Ordering::Acquire)); + if any_alive { + return Ok(StepOutcome::Contention); + } + + if ctx.input.is_drained() { + return Ok(StepOutcome::Finished); + } + Ok(StepOutcome::NoProgress) + } + + fn new_worker_copy(&self) -> Self { + self.clone() + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/fgumi-pipeline-io/src/sort/spill_decompress/tests.rs b/crates/fgumi-pipeline-io/src/sort/spill_decompress/tests.rs new file mode 100644 index 000000000..a22bfec24 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/sort/spill_decompress/tests.rs @@ -0,0 +1 @@ +// No unit tests for spill_decompress directly; integration tests live in sort/tests.rs. diff --git a/src/lib/pipeline/steps/sort/tests.rs b/crates/fgumi-pipeline-io/src/sort/tests.rs similarity index 67% rename from src/lib/pipeline/steps/sort/tests.rs rename to crates/fgumi-pipeline-io/src/sort/tests.rs index 7f81a3905..1f145a691 100644 --- a/src/lib/pipeline/steps/sort/tests.rs +++ b/crates/fgumi-pipeline-io/src/sort/tests.rs @@ -1,16 +1,6 @@ //! Tests for the runall-sort three-step chain //! (`SortAndSpill` → `SortSpillDecompress` → `SortMerge`) and for the //! single-step `SortBamFile` wrapper. -//! -//! Strategy: build a 5-step pipeline -//! `VecSource → SortAndSpill → SortSpillDecompress → SortMerge → VecSink` -//! and compare emitted record bytes against `RawExternalSorter::sort()` -//! byte-for-byte. The sort engine itself is covered by 360+ tests in -//! `fgumi-sort`; these tests pin the typed-step adapter (push, drain, -//! batch, retry, ordinal-restamp, inter-step event flow). -//! -//! We exercise both the in-memory fast path (`memory_limit` > total) and -//! the multi-spill / k-way merge path (`memory_limit` << total). use std::io; use std::sync::Arc; @@ -24,21 +14,20 @@ use parking_lot::Mutex; use rstest::rstest; use super::*; -use crate::pipeline::core::Unpushed; -use crate::pipeline::core::builder::{Pipeline, PipelineConfig}; -use crate::pipeline::core::held::HeldSlot; -use crate::pipeline::core::outputs::OrderedBytesSingle; -use crate::pipeline::core::queues::QueueSpec; -use crate::pipeline::core::reorder::BranchOrdering; -use crate::pipeline::core::step::{Step, StepCtx, StepProfile}; -use crate::pipeline::steps::types::RecordBatch; +use crate::types::RecordBatch; +use fgumi_pipeline_core::{ + Unpushed, + builder::{Pipeline, PipelineConfig}, + held::HeldSlot, + outputs::OrderedBytesSingle, + queues::QueueSpec, + reorder::BranchOrdering, + step::{Step, StepCtx, StepProfile}, +}; // ── In-memory source / sink test steps ────────────────────────────────────── -/// `Exclusive` source that drains a `Vec` one batch per -/// `try_run` call. Returns `Finished` when empty. Output ordering is -/// `ByOrdinal` so the framework's reorder stage exercises the same path -/// real producers go through. +/// `Exclusive` source that drains a `Vec` one batch per `try_run` call. struct VecSource { batches: Vec, held: HeldSlot>, @@ -47,7 +36,6 @@ struct VecSource { impl VecSource { fn new(mut batches: Vec, output_byte_limit: u64) -> Self { - // Pop from the end, so reverse to preserve user-supplied order. batches.reverse(); Self { batches, held: HeldSlot::new(), output_byte_limit } } @@ -90,11 +78,7 @@ impl Step for VecSource { } } -/// Sink that appends every received batch into a shared `Vec`, -/// for byte-level equivalence checks. `kind` selects the step kind: most tests -/// use `Exclusive`; the `threads == 1` deadlock repro uses `Serial` so the -/// chain has only ONE `Exclusive` step (`VecSource`) and is allowed to run at -/// a single worker. +/// Sink that appends every received batch into a shared `Vec`. struct VecSink { received: Arc>>, kind: StepKind, @@ -128,17 +112,12 @@ impl Step for VecSink { // ── Synthetic-record helpers ──────────────────────────────────────────────── -/// Build `n` synthetic unmapped records with deterministic, sort-relevant -/// variation (read names, tid/pos, paired flags). Thin wrapper over -/// [`synthesize_sized_records`] with a zero-length sequence — each record is -/// the minimal ~50-byte header, so spills are ≤ 1 BGZF block. fn synthesize_records(n: usize, seed: u64) -> (Header, Vec) { synthesize_sized_records(n, seed, 0) } -/// Pack records into batches of `batch_size` records each. fn pack_batches(records: &[RawRecord], batch_size: usize) -> Vec { - use crate::pipeline::steps::types::RecordBatchBuilder; + use crate::types::RecordBatchBuilder; records .chunks(batch_size) .enumerate() @@ -153,12 +132,6 @@ fn pack_batches(records: &[RawRecord], batch_size: usize) -> Vec { .collect() } -/// Drive the pipeline: -/// `VecSource(batches)` → `SortAndSpill` → `SortSpillDecompress` → -/// `SortMerge` → `VecSink(sink_kind)` → collected. -/// Returns the per-record bytes in emitted order. `sink_kind` is -/// `Exclusive` for the parity tests and `Serial` for the `threads == 1` -/// repro (so the chain has a single `Exclusive` step and one worker is legal). fn drive_sort_pipeline( sorter: RawExternalSorter, header: &Header, @@ -171,9 +144,8 @@ fn drive_sort_pipeline( let sort_order = sorter.sort_order(); let source = VecSource::new(batches, output_byte_limit); - let and_spill = SortAndSpill::from_sorter(sorter, header, 32)?; - // 64 events × ~256 KiB/block ≈ 16 MiB queue ceiling for spill blocks. - let decompress = SortSpillDecompress::new(64); + let and_spill = SortAndSpill::from_sorter(sorter, header, output_byte_limit)?; + let decompress = SortSpillDecompress::new(output_byte_limit); let merge = SortMerge::with_target_batch_count(sort_order, output_byte_limit, 256); let sink = VecSink { received: Arc::clone(&received), kind: sink_kind }; @@ -200,8 +172,6 @@ fn drive_sort_pipeline( // ── Reference: RawExternalSorter::sort to bytes ───────────────────────────── -/// Run `RawExternalSorter::sort()` on `records` and return per-record -/// bytes in the resulting BAM. fn sort_via_legacy( sort_order: SortOrder, header: &Header, @@ -254,9 +224,6 @@ fn three_step_chain_in_memory_path_matches_legacy( #[case] sort_order: SortOrder, #[case] threads: usize, ) { - // 5K small records, memory_limit large enough to keep everything - // in memory — exercises the in-memory residual chunk (no spill files) - // path of the new chain. let (header, records) = synthesize_records(5_000, 0x00C0_FFEE); let memory_limit = 256 * 1024 * 1024; @@ -293,9 +260,6 @@ fn three_step_chain_multi_spill_path_matches_legacy( #[case] sort_order: SortOrder, #[case] threads: usize, ) { - // 20K records under a tight 256 KB memory_limit → many spills, - // exercising the spill + decompress + k-way merge path of the new - // chain. This is the hot path that motivated the split. let (header, records) = synthesize_records(20_000, 0xFEED_FACE); let memory_limit = 256 * 1024; @@ -323,13 +287,6 @@ fn three_step_chain_multi_spill_path_matches_legacy( } } -/// Like [`synthesize_records`] but each record carries a `seq_len`-base -/// sequence payload so the record byte size is tunable. Large records make -/// each spill chunk span MANY BGZF blocks (more than `PHASE2_DECOMP_CAP`), -/// which is the production condition the small-record tests above miss: a -/// slot does not reach `queue_eof` on its first cap-fill, so the merge -/// consumer drains it to empty and re-parks on `block_ready` while the -/// producer is still feeding the same slot. fn synthesize_sized_records(n: usize, seed: u64, seq_len: usize) -> (Header, Vec) { let header = Header::default(); let mut state = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15).wrapping_add(1); @@ -354,36 +311,6 @@ fn synthesize_sized_records(n: usize, seed: u64, seq_len: usize) -> (Header, Vec (header, records) } -/// Regression test for issue #330 BUG #3 — the fused-sort-chain deadlock. -/// -/// Production symptom (fgumi-benchmarks, c7g.4xlarge, `--threads 8`, real -/// CODEC BAM): `runall --start-from sort` reads/sorts/spills then wedges -/// forever at the sort→group handoff. Root cause: the merge consumer -/// blocked on `slot.block_ready.wait()` whenever a slot's decompressed queue -/// was momentarily empty and not yet at EOF, parking a framework worker (and, -/// being `Serial`, holding the step lock) with no guarantee a sibling worker -/// would refill the slot. -/// -/// The condition the small-record `*_multi_spill_*` tests miss is **large -/// spill files** — each spill chunk spans far more than `PHASE2_DECOMP_CAP` -/// (8) BGZF blocks, so a slot stays `!queue_eof` while the merge consumer -/// drains it to empty mid-merge. The fix makes the merge non-blocking and -/// resumable (`MergeDriver::try_step` → `SortMerge` returns `Contention` -/// instead of parking); this test asserts the chain completes at every worker -/// count, deterministically (a hang ⇒ the bug is back). -/// -/// `threads == 1` is the M1 case the fix targets (a sole worker that parks in -/// the merge can never round-robin to refill). It runs here with a large -/// output-queue limit on purpose: at `threads == 1` the round-robin's -/// priority-restart separately starves the single worker's downstream drain -/// under output backpressure — a pre-existing, `threads == 1`-only artifact of -/// this minimal harness. It does not affect production: the real fused chain -/// has two `Exclusive` steps (BAM reader + writer) and so always runs at -/// `threads >= 2`, where a separate worker drains the output. Sizing the queue -/// above the total output isolates the merge-liveness property this test -/// guards (the old blocking merge hangs at `threads == 1` regardless of queue -/// size). The non-blocking-merge unit coverage is in `fgumi-sort` -/// (`from_slots_try_step_*`). #[rstest] #[case::t1(1)] #[case::t2(2)] @@ -393,17 +320,9 @@ fn three_step_chain_large_spill_completes(#[case] pipeline_threads: usize) { use std::sync::mpsc; use std::time::Duration; - // ~320-byte records (seq_len=200 → 200 + 100 seq/qual bytes + header) - // under a 1 MiB memory_limit ⇒ ~3K records/spill ⇒ ~1 MiB ⇒ ~15 BGZF - // blocks/spill (> CAP 8). 60K records ⇒ ~20 spill files. The - // many-blocks-per-spill shape is what makes the merge consumer drain a - // slot to empty while that slot is still `!queue_eof` — the production - // condition that triggered the original `block_ready` park. let total_records = 60_000usize; let memory_limit = 1024 * 1024; let sort_order = SortOrder::Coordinate; - // Sort-engine internal threads are independent of the pipeline worker - // count; keep them low so the spill shape is stable across cases. let sorter_threads = 2; let (header, records) = synthesize_sized_records(total_records, 0xBADD_CAFE, 200); @@ -414,24 +333,14 @@ fn three_step_chain_large_spill_completes(#[case] pipeline_threads: usize) { .output_compression(1) .temp_compression(1); - // Reference output (byte-for-byte) so a hang ISN'T the only thing this - // test catches — a merge/refill ordering regression must also fail. let legacy_out = sort_via_legacy(sort_order, &header, &records, memory_limit, sorter_threads) .expect("legacy sort"); - // Output-queue ceiling sized above the ~19 MiB total output so the - // `threads == 1` case isn't confounded by the unrelated single-worker - // output-backpressure starvation (see the test doc). The old blocking - // merge hangs at `threads == 1` regardless of this size. let output_queue_limit = 256 * 1024 * 1024; let (tx, rx) = mpsc::channel(); let worker = std::thread::Builder::new() .name(format!("bug3-repro-t{pipeline_threads}")) .spawn(move || { - // Drive the chain directly at the requested pipeline worker count - // — NOT `threads.max(3)` — with a `Serial` sink so the chain has a - // single `Exclusive` step and `threads == 1` is legal, exercising - // single-worker merge liveness rather than hiding it. let result = drive_sort_pipeline( sorter, &header, @@ -480,13 +389,6 @@ fn three_step_chain_empty_input_drains_cleanly() { // ── SortBamFile (Exclusive single-step) tests ────────────────────────────── -/// Smoke test for the `SortBamFile` step. Constructs a synthetic input -/// BAM, sorts it via `[SortBamFile]` through the pipeline, then sorts -/// the same input via `RawExternalSorter::sort()` directly, and -/// verifies byte-identical output. Since `SortBamFile`'s body is just -/// `sorter.sort(input, output)`, this test mostly verifies the -/// framework wrapping doesn't corrupt anything — the heavy testing is -/// in `fgumi-sort`'s own 360+ tests. #[test] fn sort_bam_file_matches_legacy_sort() { let (header, records) = synthesize_records(2_000, 0xCAFE_F00D); @@ -518,18 +420,21 @@ fn sort_bam_file_matches_legacy_sort() { .threads(1) .output_compression(1) .temp_compression(1); - let step = SortBamFile::new( - sorter, - input.clone(), - pipeline_out.clone(), - Arc::new(parking_lot::Mutex::new(None)), - ); + let stats_slot = Arc::new(parking_lot::Mutex::new(None)); + let step = + SortBamFile::new(sorter, input.clone(), pipeline_out.clone(), Arc::clone(&stats_slot)); let builder = Pipeline::builder(); builder.chain(step).into_sink_marker(); let pipeline = builder.build().expect("Pipeline::build"); pipeline.run(PipelineConfig { threads: 1, ..Default::default() }).expect("Pipeline::run"); + // `SortBamFile::try_run` must publish its stats into the shared slot; the + // `SortFinalizeHook` reads this slot after `Pipeline::run` to log the summary. + let stats = stats_slot.lock().take().expect("SortBamFile should publish SortStats"); + assert_eq!(stats.total_records, records.len() as u64); + assert_eq!(stats.output_records, records.len() as u64); + let legacy_records = read_all_records(&legacy_out); let pipeline_records = read_all_records(&pipeline_out); assert_eq!(legacy_records.len(), pipeline_records.len(), "record count mismatch"); @@ -556,7 +461,111 @@ fn sort_bam_file_profile_is_exclusive() { assert!(p.branch_ordering.is_empty()); } -/// Read every record's bytes from a BAM file. +// ── SortMerge fail-closed regression tests ────────────────────────────────── + +/// `Exclusive` source that drains a `Vec` one event per +/// `try_run`, feeding `SortMerge` directly. Used to drive the merge into a +/// drained-but-incomplete-setup state without standing up the spill machinery. +struct Phase2EventSource { + events: Vec, + held: HeldSlot>, + output_byte_limit: u64, +} + +impl Phase2EventSource { + fn new( + mut events: Vec, + output_byte_limit: u64, + ) -> Self { + events.reverse(); + Self { events, held: HeldSlot::new(), output_byte_limit } + } +} + +impl Step for Phase2EventSource { + type Input = (); + type Outputs = fgumi_pipeline_core::outputs::Single; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "Phase2EventSource", + kind: StepKind::Exclusive, + sticky: true, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], + branch_ordering: vec![BranchOrdering::None], + } + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if let Some(unpushed) = self.held.take() { + match ctx.outputs.retry(unpushed) { + Ok(()) => {} + Err(again) => { + self.held.put(again); + return Ok(StepOutcome::Progress); + } + } + } + let Some(event) = self.events.pop() else { + return Ok(StepOutcome::Finished); + }; + match ctx.outputs.push(event) { + Ok(()) => Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held.put(unpushed); + Ok(StepOutcome::Progress) + } + } + } +} + +fn run_merge_over_events( + events: Vec, +) -> Result>> { + let output_byte_limit = 1 << 20; + let received: Arc>> = Arc::new(Mutex::new(Vec::new())); + let source = Phase2EventSource::new(events, output_byte_limit); + let merge = SortMerge::with_target_batch_count(SortOrder::Coordinate, output_byte_limit, 256); + let sink = VecSink { received: Arc::clone(&received), kind: StepKind::Serial }; + + let builder = Pipeline::builder(); + builder.chain(source).chain(merge).chain(sink).into_sink_marker(); + let pipeline = builder.build()?; + pipeline.run(PipelineConfig { threads: 1, ..Default::default() })?; + + let collected = std::mem::take(&mut *received.lock()); + let mut out = Vec::new(); + for batch in collected { + for bytes in batch.iter_record_bytes() { + out.push(bytes.to_vec()); + } + } + Ok(out) +} + +/// A wholly empty event stream (no payload, no `AllAnnounced`) is the one +/// legitimate drained-but-not-ready case: it merges to an empty output. +#[test] +fn test_sort_merge_empty_input_merges_to_empty() { + let out = run_merge_over_events(Vec::new()).expect("empty merge should succeed"); + assert!(out.is_empty(), "expected no records, got {}", out.len()); +} + +/// An `AllAnnounced` that promises a slot which never arrives leaves the setup +/// incomplete when the input drains; `SortMerge` must fail closed rather than +/// silently merge a partial result. +#[test] +fn test_sort_merge_fails_closed_on_incomplete_setup() { + let events = vec![crate::sort::protocol::SortPhase2Event::AllAnnounced { + slot_count: 1, + memory_chunk_count: 0, + total_records: 0, + }]; + let err = run_merge_over_events(events).expect_err("incomplete setup must error"); + let msg = err.to_string(); + assert!(msg.contains("setup incomplete at input drain"), "unexpected error message: {msg}"); +} + fn read_all_records(path: &std::path::Path) -> Vec> { let (mut reader, _hdr) = fgumi_bam_io::create_raw_bam_reader_with_opts( path, diff --git a/crates/fgumi-pipeline-io/src/source/mod.rs b/crates/fgumi-pipeline-io/src/source/mod.rs new file mode 100644 index 000000000..5d82941cb --- /dev/null +++ b/crates/fgumi-pipeline-io/src/source/mod.rs @@ -0,0 +1 @@ +pub mod read_bam; diff --git a/crates/fgumi-pipeline-io/src/source/read_bam.rs b/crates/fgumi-pipeline-io/src/source/read_bam.rs new file mode 100644 index 000000000..03eefbc14 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/source/read_bam.rs @@ -0,0 +1,288 @@ +//! `ReadBgzfBlocks` source step + `read_bam(path)` convenience helper. +//! +//! Reads raw BGZF blocks from a file (no decompression) and emits them as +//! `BgzfBlock` items with monotonically increasing `batch_serial`. The +//! header bytes are NOT skipped here — they pass through as part of the +//! first block(s); `FindBamBoundaries` strips them downstream. + +use std::collections::VecDeque; +use std::fs::File; +use std::io; +use std::path::Path; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +use fgumi_bam_io::PipelineReaderOpts; +use fgumi_bgzf::reader::read_raw_blocks; +use noodles::sam::Header; +use parking_lot::Mutex; + +use crate::types::BgzfBlock; +use fgumi_pipeline_core::{ + Unpushed, + held::HeldSlot, + outputs::OrderedBytesSingle, + queues::QueueSpec, + reorder::BranchOrdering, + step::{Affinity, Step, StepCtx, StepKind, StepOutcome, StepProfile}, +}; + +/// Legacy default blocks-per-batch. +pub const DEFAULT_BLOCKS_PER_BATCH: usize = 16; + +/// `Exclusive + sticky` source step that reads raw BGZF blocks from a file. +pub struct ReadBgzfBlocks { + reader: Arc>>>, + blocks_per_batch: usize, + next_serial: u64, + pending: VecDeque, + held: HeldSlot>, + output_byte_limit: u64, + finished: Arc, +} + +impl ReadBgzfBlocks { + #[must_use] + pub fn new( + reader: Box, + blocks_per_batch: usize, + output_byte_limit: u64, + ) -> Self { + Self { + reader: Arc::new(Mutex::new(Some(reader))), + blocks_per_batch: blocks_per_batch.max(1), + next_serial: 0, + pending: VecDeque::new(), + held: HeldSlot::new(), + output_byte_limit, + finished: Arc::new(AtomicBool::new(false)), + } + } +} + +impl Step for ReadBgzfBlocks { + type Input = (); + type Outputs = OrderedBytesSingle; + + fn profile(&self) -> StepProfile { + StepProfile { + name: "ReadBgzfBlocks", + kind: StepKind::Serial, + sticky: true, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], + branch_ordering: vec![BranchOrdering::ByItemOrdinal], + } + } + + fn affinity(&self) -> Affinity { + Affinity::Reader + } + + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + // 1. Drain the held slot first. + if let Some(unpushed) = self.held.take() { + match ctx.outputs.retry(unpushed) { + Ok(()) => {} + Err(again) => { + self.held.put(again); + return Ok(StepOutcome::Contention); + } + } + } + + // 2. Drain pending blocks (one per call iteration). + if let Some(block) = self.pending.pop_front() { + match ctx.outputs.push(block) { + Ok(()) => return Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held.put(unpushed); + return Ok(StepOutcome::Progress); + } + } + } + + if self.finished.load(Ordering::Acquire) { + return Ok(StepOutcome::Finished); + } + + // 3. Read up to `blocks_per_batch` raw BGZF blocks. + let raw_blocks = { + let mut guard = self.reader.lock(); + let reader = + guard.as_mut().expect("ReadBgzfBlocks: reader missing — was clone() called?"); + read_raw_blocks(reader.as_mut(), self.blocks_per_batch)? + }; + + if raw_blocks.is_empty() { + self.finished.store(true, Ordering::Release); + return Ok(StepOutcome::Finished); + } + + for raw in raw_blocks { + let serial = self.next_serial; + self.next_serial += 1; + self.pending.push_back(BgzfBlock { + batch_serial: serial, + uncompressed_size: u32::try_from(raw.uncompressed_size()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidData, + format!( + "ReadBgzfBlocks: BGZF uncompressed_size out of range: {}", + raw.uncompressed_size() + ), + ) + })?, + bytes: raw.data, + }); + } + + if let Some(block) = self.pending.pop_front() { + match ctx.outputs.push(block) { + Ok(()) => Ok(StepOutcome::Progress), + Err(unpushed) => { + self.held.put(unpushed); + Ok(StepOutcome::Progress) + } + } + } else { + Ok(StepOutcome::NoProgress) + } + } +} + +/// Build a [`ReadBgzfBlocks`] step from an already-prepared reader + header. +#[must_use] +pub fn read_bam_from_reader( + reader: Box, + header: Header, + blocks_per_batch: usize, + output_byte_limit: u64, +) -> (ReadBgzfBlocks, Header) { + (ReadBgzfBlocks::new(reader, blocks_per_batch, output_byte_limit), header) +} + +/// Convenience helper: open a BAM file, parse its header, return the +/// `(step, header)` pair. +/// +/// # Errors +/// +/// Returns I/O errors from file open or BAM-header parse. +pub fn read_bam>( + path: P, + opts: PipelineReaderOpts, + blocks_per_batch: usize, + output_byte_limit: u64, +) -> io::Result<(ReadBgzfBlocks, Header)> { + let path = path.as_ref(); + let (_, header) = fgumi_bam_io::create_raw_bam_reader_with_opts(path, 1, opts) + .map_err(|e| io::Error::other(format!("create_raw_bam_reader_with_opts: {e}")))?; + + let file = File::open(path)?; + let reader: Box = + Box::new(io::BufReader::with_capacity(2 * 1024 * 1024, file)); + Ok(read_bam_from_reader(reader, header, blocks_per_batch, output_byte_limit)) +} + +/// Stdin counterpart to [`read_bam`]. +/// +/// # Errors +/// +/// Returns I/O errors from stdin read or BAM-header parse. +pub fn read_bam_stdin( + opts: PipelineReaderOpts, + blocks_per_batch: usize, + output_byte_limit: u64, +) -> io::Result<(ReadBgzfBlocks, Header)> { + let (reader, header) = + fgumi_bam_io::create_bam_reader_for_pipeline_with_opts(Path::new("-"), opts).map_err( + |e| io::Error::other(format!("create_bam_reader_for_pipeline_with_opts: {e}")), + )?; + Ok(read_bam_from_reader(reader, header, blocks_per_batch, output_byte_limit)) +} + +/// Path-aware dispatcher: routes to [`read_bam_stdin`] when `path` is a +/// stdin sentinel (`-` or `/dev/stdin`) and to [`read_bam`] otherwise. +/// +/// # Errors +/// +/// Returns I/O errors from file open, stdin read, or BAM-header parse. +pub fn read_bam_auto>( + path: P, + opts: PipelineReaderOpts, + blocks_per_batch: usize, + output_byte_limit: u64, +) -> io::Result<(ReadBgzfBlocks, Header)> { + if fgumi_bam_io::is_stdin_path(path.as_ref()) { + read_bam_stdin(opts, blocks_per_batch, output_byte_limit) + } else { + read_bam(path, opts, blocks_per_batch, output_byte_limit) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn profile_advertises_serial_reader_byordinal() { + let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); + let header = noodles::sam::Header::default(); + let writer = fgumi_bam_io::create_raw_bam_writer(&path, &header, 1, 1).unwrap(); + writer.finish().unwrap(); + + let (step, _hdr) = + read_bam(&path, PipelineReaderOpts::default(), DEFAULT_BLOCKS_PER_BATCH, 1024 * 1024) + .unwrap(); + let profile = step.profile(); + assert_eq!(profile.name, "ReadBgzfBlocks"); + assert_eq!(profile.kind, StepKind::Serial); + assert!(profile.sticky); + assert_eq!(step.affinity(), Affinity::Reader); + assert_eq!(profile.branch_ordering, vec![BranchOrdering::ByItemOrdinal]); + assert!(matches!(profile.output_queues[0], QueueSpec::ByteBounded { .. })); + } + + #[test] + fn read_bam_from_reader_round_trips_bytes() { + const BGZF_EOF_LEN: usize = 28; + + let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); + let header = noodles::sam::Header::default(); + let writer = fgumi_bam_io::create_raw_bam_writer(&path, &header, 1, 1).unwrap(); + writer.finish().unwrap(); + + let on_disk = std::fs::read(&path).unwrap(); + assert!(!on_disk.is_empty(), "BAM file should contain header + EOF block"); + + let cursor = std::io::Cursor::new(on_disk.clone()); + let reader: Box = Box::new(cursor); + let (mut step, _hdr) = + read_bam_from_reader(reader, header, DEFAULT_BLOCKS_PER_BATCH, 1024 * 1024); + + let mut collected = Vec::with_capacity(on_disk.len()); + let mut last_serial: Option = None; + loop { + let raw_blocks = { + let mut guard = step.reader.lock(); + let reader = guard.as_mut().expect("reader present"); + fgumi_bgzf::reader::read_raw_blocks(reader.as_mut(), DEFAULT_BLOCKS_PER_BATCH) + .unwrap() + }; + if raw_blocks.is_empty() { + break; + } + for raw in raw_blocks { + let serial = step.next_serial; + step.next_serial += 1; + if let Some(prev) = last_serial { + assert_eq!(serial, prev + 1, "serials must be monotonic"); + } + last_serial = Some(serial); + collected.extend_from_slice(&raw.data); + } + } + let expected = &on_disk[..on_disk.len() - BGZF_EOF_LEN]; + assert_eq!(collected, expected, "concatenated blocks must equal source bytes minus EOF"); + assert!(last_serial.is_some(), "should have read at least one block"); + } +} diff --git a/crates/fgumi-pipeline-io/src/types.rs b/crates/fgumi-pipeline-io/src/types.rs new file mode 100644 index 000000000..05f683be7 --- /dev/null +++ b/crates/fgumi-pipeline-io/src/types.rs @@ -0,0 +1,302 @@ +//! Concrete data types that flow through the BAM step library. +//! +//! Every flowing type carries an explicit `batch_serial: u64` field and +//! impls both [`HeapSize`] (so byte-bounded queues can budget memory) and +//! [`Ordered`] (so `BranchOrdering::ByItemOrdinal` reorder stages preserve +//! global ordering across multi-step Parallel transforms). + +use fgumi_pipeline_core::{HeapSize, Ordered}; +use fgumi_raw_bam::RawRecord; + +// ───────────────────────────────────────────────────────────────────────────── +// BgzfBlock — raw compressed BGZF block + read-order serial. +// ───────────────────────────────────────────────────────────────────────────── + +/// Raw compressed BGZF block + parsed metadata. Carries read-order serial. +/// +/// Sentinel/EOF blocks have `bytes` containing the 28-byte BGZF EOF marker +/// and `uncompressed_size = 0`. +#[derive(Debug)] +pub struct BgzfBlock { + /// Read-order serial. Set by `ReadBgzfBlocks` based on block read index. + pub batch_serial: u64, + pub bytes: Vec, + /// Decompressed size, parsed from the BGZF block header. + pub uncompressed_size: u32, +} + +impl HeapSize for BgzfBlock { + fn heap_size(&self) -> usize { + // Byte-bounded queues budget on resident heap, so account for the full + // allocation (`capacity`), not just the populated prefix (`len`). + self.bytes.capacity() + } +} + +impl Ordered for BgzfBlock { + fn ordinal(&self) -> u64 { + self.batch_serial + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// DecompressedBlock — raw bytes from a BGZF decompression, record-aligned +// or not. +// ───────────────────────────────────────────────────────────────────────────── + +/// Decompressed bytes from one or more BGZF blocks. Carries a serial for +/// ordering purposes; record alignment is the consumer's responsibility. +#[derive(Debug)] +pub struct DecompressedBlock { + pub batch_serial: u64, + pub bytes: Vec, +} + +impl HeapSize for DecompressedBlock { + fn heap_size(&self) -> usize { + // Account for the full allocation (`capacity`), not just `len` — see + // the `BgzfBlock` impl above. + self.bytes.capacity() + } +} + +impl Ordered for DecompressedBlock { + fn ordinal(&self) -> u64 { + self.batch_serial + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// RecordBatch — parsed BAM records grouped into a batch. +// ───────────────────────────────────────────────────────────────────────────── + +/// A batch of parsed BAM records, stored as a flat backing buffer + per-record +/// `(start, end)` ranges. +#[derive(Debug)] +pub struct RecordBatch { + batch_serial: u64, + /// All record bodies concatenated, in batch order. + backing: Vec, + /// (start, end) byte ranges into `backing`, one per record. + ranges: Vec<(u32, u32)>, +} + +impl RecordBatch { + /// Construct a batch from a pre-parsed backing buffer and `(start, end)` ranges. + #[must_use] + pub fn from_parsed(batch_serial: u64, backing: Vec, ranges: Vec<(u32, u32)>) -> Self { + Self { batch_serial, backing, ranges } + } + + /// Convenience constructor: serializes a slice of `RawRecord`s into the + /// flat representation. + /// + /// # Panics + /// + /// Panics if the concatenated record bodies exceed `u32::MAX` bytes. + #[must_use] + pub fn new(batch_serial: u64, records: &[RawRecord]) -> Self { + let total: usize = records.iter().map(RawRecord::len).sum(); + let mut backing = Vec::with_capacity(total); + let mut ranges = Vec::with_capacity(records.len()); + for rec in records { + let start = u32::try_from(backing.len()).expect("backing fits in u32"); + backing.extend_from_slice(rec.as_ref()); + let end = u32::try_from(backing.len()).expect("backing fits in u32"); + ranges.push((start, end)); + } + Self { batch_serial, backing, ranges } + } + + /// Self-managed ordinal. + #[must_use] + pub fn batch_serial(&self) -> u64 { + self.batch_serial + } + + /// Number of records in the batch. + #[must_use] + pub fn len(&self) -> usize { + self.ranges.len() + } + + /// `true` iff the batch contains zero records. + #[must_use] + pub fn is_empty(&self) -> bool { + self.ranges.is_empty() + } + + /// Total bytes across all record bodies. + #[must_use] + pub fn total_bytes(&self) -> usize { + self.backing.len() + } + + /// Iterate the record bodies as borrowed byte slices into the backing buffer. + pub fn iter_record_bytes(&self) -> impl Iterator + '_ { + let backing = &self.backing[..]; + self.ranges.iter().map(move |&(s, e)| &backing[s as usize..e as usize]) + } +} + +/// Builder for emit-side `RecordBatch` construction. +#[derive(Debug)] +pub struct RecordBatchBuilder { + batch_serial: u64, + backing: Vec, + ranges: Vec<(u32, u32)>, +} + +impl RecordBatchBuilder { + /// Build an empty builder with reserved capacity. + #[must_use] + pub fn with_capacity(batch_serial: u64, bytes_cap: usize, records_cap: usize) -> Self { + Self { + batch_serial, + backing: Vec::with_capacity(bytes_cap), + ranges: Vec::with_capacity(records_cap), + } + } + + /// Append one record's body bytes. + /// + /// # Panics + /// + /// Panics if accumulated bytes would exceed `u32::MAX`. + pub fn push_record_bytes(&mut self, bytes: &[u8]) { + let start = u32::try_from(self.backing.len()).expect("backing fits in u32"); + self.backing.extend_from_slice(bytes); + let end = u32::try_from(self.backing.len()).expect("backing fits in u32"); + self.ranges.push((start, end)); + } + + /// Number of records appended so far. + #[must_use] + pub fn len(&self) -> usize { + self.ranges.len() + } + + /// `true` iff no records have been appended. + #[must_use] + pub fn is_empty(&self) -> bool { + self.ranges.is_empty() + } + + /// Total record bytes appended so far. + #[must_use] + pub fn total_bytes(&self) -> usize { + self.backing.len() + } + + /// Finalize and produce the `RecordBatch`. Consumes the builder. + #[must_use] + pub fn build(self) -> RecordBatch { + RecordBatch { batch_serial: self.batch_serial, backing: self.backing, ranges: self.ranges } + } +} + +impl HeapSize for RecordBatch { + fn heap_size(&self) -> usize { + // Account for the full allocation (`capacity`) of both buffers, not + // just their populated prefixes — see the `BgzfBlock` impl above. + self.backing.capacity() + self.ranges.capacity() * std::mem::size_of::<(u32, u32)>() + } +} + +impl Ordered for RecordBatch { + fn ordinal(&self) -> u64 { + self.batch_serial + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn bgzf_block_heap_size_matches_bytes_capacity() { + let b = BgzfBlock { batch_serial: 0, bytes: vec![0u8; 1024], uncompressed_size: 4096 }; + assert_eq!(b.heap_size(), 1024); + assert_eq!(b.ordinal(), 0); + } + + #[test] + fn decompressed_block_heap_size_matches_bytes_capacity() { + let b = DecompressedBlock { batch_serial: 7, bytes: vec![0u8; 4096] }; + assert_eq!(b.heap_size(), 4096); + assert_eq!(b.ordinal(), 7); + } + + #[test] + fn heap_size_counts_allocated_capacity_not_logical_len() { + // A buffer with spare capacity (e.g. after `with_capacity`) holds more + // resident heap than its `len` — byte-bounded queues must budget the + // full allocation or they undercount and bypass configured limits. + let mut bytes = Vec::with_capacity(8192); + bytes.extend_from_slice(&[0u8; 100]); + assert!(bytes.capacity() >= 8192 && bytes.len() == 100); + let cap = bytes.capacity(); + + let block = BgzfBlock { batch_serial: 0, bytes, uncompressed_size: 0 }; + assert_eq!(block.heap_size(), cap); + + let mut backing = Vec::with_capacity(4096); + backing.extend_from_slice(&[0u8; 10]); + let mut ranges = Vec::with_capacity(64); + ranges.push((0u32, 10u32)); + let backing_cap = backing.capacity(); + let ranges_cap = ranges.capacity(); + let batch = RecordBatch::from_parsed(0, backing, ranges); + assert_eq!(batch.heap_size(), backing_cap + ranges_cap * std::mem::size_of::<(u32, u32)>()); + } + + #[test] + fn record_batch_total_bytes_sums_record_lengths() { + let r1: RawRecord = vec![0u8; 100].into(); + let r2: RawRecord = vec![0u8; 200].into(); + let batch = RecordBatch::new(3, &[r1, r2]); + assert_eq!(batch.len(), 2); + assert_eq!(batch.total_bytes(), 300); + // `Vec::with_capacity` may over-allocate, so assert against the actual + // allocated capacities rather than the requested sizes (see the sibling + // `heap_size_counts_allocated_capacity_not_logical_len` test). + assert_eq!( + batch.heap_size(), + batch.backing.capacity() + batch.ranges.capacity() * std::mem::size_of::<(u32, u32)>() + ); + assert_eq!(batch.ordinal(), 3); + } + + #[test] + fn record_batch_from_parsed_round_trips_ranges() { + let backing = b"AAABBBBCC".to_vec(); + let ranges = vec![(0u32, 3u32), (3u32, 7u32), (7u32, 9u32)]; + let batch = RecordBatch::from_parsed(11, backing, ranges); + let got: Vec<&[u8]> = batch.iter_record_bytes().collect(); + assert_eq!(got, vec![&b"AAA"[..], &b"BBBB"[..], &b"CC"[..]]); + assert_eq!(batch.len(), 3); + assert_eq!(batch.total_bytes(), 9); + assert_eq!(batch.batch_serial(), 11); + } + + #[test] + fn record_batch_builder_collects_records() { + let mut b = RecordBatchBuilder::with_capacity(0, 64, 4); + b.push_record_bytes(&[0u8; 50]); + b.push_record_bytes(&[0u8; 75]); + assert_eq!(b.len(), 2); + assert!(!b.is_empty()); + assert_eq!(b.total_bytes(), 125); + let batch = b.build(); + let got: Vec = batch.iter_record_bytes().map(<[u8]>::len).collect(); + assert_eq!(got, vec![50, 75]); + // `heap_size` budgets allocated capacity (not logical length): the + // builder was seeded with a 64-byte backing buffer but 125 bytes were + // pushed, so `backing` reallocated and its capacity now exceeds 125. + assert_eq!( + batch.heap_size(), + batch.backing.capacity() + batch.ranges.capacity() * std::mem::size_of::<(u32, u32)>() + ); + assert!(batch.heap_size() >= 125 + 2 * std::mem::size_of::<(u32, u32)>()); + } +} diff --git a/crates/fgumi-sort-cli/Cargo.toml b/crates/fgumi-sort-cli/Cargo.toml new file mode 100644 index 000000000..0c129affe --- /dev/null +++ b/crates/fgumi-sort-cli/Cargo.toml @@ -0,0 +1,31 @@ +[package] +name = "fgumi-sort-cli" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +description = "Standalone fgumi sort CLI command (framework-light)" +repository.workspace = true +license.workspace = true + +[dependencies] +fgumi-bam-io = { workspace = true } +fgumi-cli-common = { workspace = true } +fgumi-cli-macros = { workspace = true } +fgumi-pipeline-core = { workspace = true } +fgumi-pipeline-io = { workspace = true } +fgumi-sam = { workspace = true } +fgumi-sort = { workspace = true } +anyhow = "1.0.102" +bytesize = "2.3" +clap = { version = "4", features = ["derive", "string"] } +log = "0" +noodles = { version = "0.111.0", features = ["bam", "sam"] } +parking_lot = "0.12" + +[dev-dependencies] +noodles = { version = "0.111.0", features = ["bam", "sam"] } +rstest = "0" +tempfile = "3.4" + +[lints.clippy] +pedantic = { level = "deny", priority = -1 } diff --git a/crates/fgumi-sort-cli/src/chains.rs b/crates/fgumi-sort-cli/src/chains.rs new file mode 100644 index 000000000..4cb2c8133 --- /dev/null +++ b/crates/fgumi-sort-cli/src/chains.rs @@ -0,0 +1,295 @@ +//! Sort-specific finalize hooks and the `SortBamFile` step factory. +//! +//! These were lifted out of the umbrella's `pipeline::chains::commands::sort` +//! module so that `fgumi sort` can build its pipeline directly via +//! [`fgumi_pipeline_core::Pipeline::builder`] without depending on the +//! umbrella's monolithic `ChainBuilder`. + +use std::path::PathBuf; +use std::sync::Arc; + +use anyhow::Result; +use bytesize::ByteSize; +use fgumi_cli_common::{MemoryLimit, OperationTimer, format_duration}; +use fgumi_pipeline_io::SortBamFile; +use fgumi_sort::{RawExternalSorter, SortOrder}; +use log::info; +use parking_lot::Mutex; + +use crate::sort::{SortOptions, TMP_DIRS_ENV, parse_cell_tag, resolve_tmp_dirs}; +use crate::version; + +// ───────────────────────────────────────────────────────────────────────────── +// SortFinalizeHook +// ───────────────────────────────────────────────────────────────────────────── + +/// Post-pipeline finalize action for sort. Reads `SortStats` out of the +/// shared slot, logs the Summary block, and calls +/// [`OperationTimer::log_completion`]. +pub struct SortFinalizeHook { + /// Shared slot the `SortBamFile` step fills after `sort()` returns. + pub stats_slot: Arc>>, + /// Path of the finished BAM (logged in the Summary block). + pub output_path: PathBuf, + /// Timer started when the sort began (logs wall-time on completion). + pub timer: OperationTimer, +} + +impl SortFinalizeHook { + /// Run the post-pipeline summary logging. + pub fn finalize(self) { + let SortFinalizeHook { stats_slot, output_path, timer } = self; + + let stats = stats_slot.lock().take().unwrap_or_default(); + info!("=== Summary ==="); + info!("Records processed: {}", stats.total_records); + info!("Records written: {}", stats.output_records); + if stats.chunks_written > 0 { + info!("Temporary chunks: {}", stats.chunks_written); + } + info!("Output: {}", output_path.display()); + + timer.log_completion(stats.total_records); + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// IndexBamFinalizeHook +// ───────────────────────────────────────────────────────────────────────────── + +/// Post-pipeline action that builds a BAI index for a finished +/// coordinate-sorted BAM, writing `.bam.bai` next to the BAM. +/// +/// **Invariant:** the BAM at `output_path` must be writer-closed before +/// `finalize` runs — guaranteed by `Pipeline::run` returning (which only +/// happens after every step's writer has dropped). +pub struct IndexBamFinalizeHook { + /// Path of the finished coordinate-sorted BAM to index. + pub output_path: PathBuf, +} + +impl IndexBamFinalizeHook { + /// Build and write the BAI index alongside the BAM. + /// + /// # Errors + /// + /// Returns an error if indexing or writing the BAI sidecar fails. + pub fn finalize(self) -> Result<()> { + use fgumi_bam_io::write_bai_index; + use noodles::bam; + use std::time::Instant; + + let IndexBamFinalizeHook { output_path } = self; + + info!("Indexing BAM: {}", output_path.display()); + let start = Instant::now(); + + let index = bam::fs::index(&output_path) + .map_err(|e| anyhow::anyhow!("Failed to index {}: {e}", output_path.display()))?; + + // BAI convention: sit next to the BAM with `.bam.bai`, e.g. + // `foo.bam` → `foo.bam.bai`. + let index_path = output_path.with_extension("bam.bai"); + write_bai_index(&index_path, &index) + .map_err(|e| anyhow::anyhow!("Failed to write BAI to {}: {e}", index_path.display()))?; + info!("Wrote BAM index: {} ({})", index_path.display(), format_duration(start.elapsed())); + + Ok(()) + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Step factory +// ───────────────────────────────────────────────────────────────────────────── + +/// Parameters for [`build_sort_step`]. Bundles all captures that the +/// factory needs so the call site stays readable. +pub struct SortStepCaptures { + /// The resolved per-stage sort options (memory, tmp dirs, order, …). + pub sort: SortOptions, + /// Input BAM/SAM path. + pub input_path: PathBuf, + /// Output BAM path. + pub output_path: PathBuf, + /// Number of threads for the sort engine's worker pool. + pub num_sorter_threads: usize, + /// Resolved memory budget, already computed by the caller via + /// [`resolve_memory_budget`](fgumi_cli_common::resolve_memory_budget). + /// Passed in (rather than recomputed here) so + /// the configured value the caller logged matches the value the sorter + /// uses, and to avoid a redundant `detect_total_memory()` syscall. + pub effective_memory: usize, + /// Output BGZF compression level. + pub output_compression: u32, + /// Full command-line string for the `@PG` record. + pub command_line: String, + /// Shared slot for post-run stats retrieval. The caller (finalize hook) + /// holds the other end; the step fills it when `RawExternalSorter::sort` + /// returns. + pub stats_slot: Arc>>, +} + +/// Build the `SortBamFile` step, constructing and fully configuring the +/// [`RawExternalSorter`]. +/// +/// # Errors +/// +/// Returns an error if the `MemoryLimit::Auto` initial-capacity calculation +/// overflows, or the cell tag is invalid for the sort order. +pub fn build_sort_step(cap: SortStepCaptures) -> Result { + let sort_order: SortOrder = cap.sort.order.into(); + let cell_tag = parse_cell_tag(cap.sort.order)?; + + let effective_memory = cap.effective_memory; + + let mut sorter = RawExternalSorter::new(sort_order) + .memory_limit(effective_memory) + .threads(cap.num_sorter_threads) + .output_compression(cap.output_compression) + .temp_compression(cap.sort.temp_compression) + .pg_info(version::version_string(), cap.command_line); + + if matches!(cap.sort.max_memory, MemoryLimit::Auto) { + let init = 768_usize + .checked_mul(1024 * 1024) + .and_then(|b| b.checked_mul(cap.num_sorter_threads)) + .ok_or_else(|| anyhow::anyhow!("initial auto buffer size overflowed"))?; + sorter = sorter.initial_capacity(effective_memory.min(init)); + } + + if let Some(ct) = cell_tag { + sorter = sorter.cell_tag(ct); + } + + let env_value = std::env::var(TMP_DIRS_ENV).ok(); + let resolved_tmp_dirs = resolve_tmp_dirs(&cap.sort.tmp_dirs, env_value.as_deref()); + if !resolved_tmp_dirs.is_empty() { + sorter = sorter.temp_dirs(resolved_tmp_dirs); + } + + Ok(SortBamFile::new(sorter, cap.input_path, cap.output_path, cap.stats_slot)) +} + +// ───────────────────────────────────────────────────────────────────────────── +// Logging helper +// ───────────────────────────────────────────────────────────────────────────── + +/// Log the sort startup banner. +pub fn log_sort_start( + sort: &SortOptions, + input_path: &std::path::Path, + output_path: &std::path::Path, + num_sorter_threads: usize, + effective_memory: usize, +) { + let cell_tag = parse_cell_tag(sort.order).unwrap_or(None); + info!("Starting Sort"); + info!("Input: {}", input_path.display()); + info!("Output: {}", output_path.display()); + info!("Sort order: {:?}", sort.order); + if let Some(ct) = cell_tag { + let ct_bytes = *ct; + info!("Cell tag: {}{}", ct_bytes[0] as char, ct_bytes[1] as char); + } + if let MemoryLimit::Fixed(per_thread) = sort.max_memory { + if sort.memory_per_thread { + info!( + "Max memory: {} ({}/thread x {} threads)", + ByteSize(effective_memory as u64), + ByteSize(per_thread as u64), + num_sorter_threads + ); + } else { + info!("Max memory: {} (fixed)", ByteSize(effective_memory as u64)); + } + } + info!("Threads: {num_sorter_threads}"); + info!("Temp compression level: {}", sort.temp_compression); + let env_value = std::env::var(TMP_DIRS_ENV).ok(); + let resolved_tmp_dirs = resolve_tmp_dirs(&sort.tmp_dirs, env_value.as_deref()); + if !resolved_tmp_dirs.is_empty() { + let joined = resolved_tmp_dirs + .iter() + .map(|p| p.display().to_string()) + .collect::>() + .join(", "); + info!("Temp directories: {joined}"); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use fgumi_bam_io::create_raw_bam_writer; + use noodles::sam::Header; + use noodles::sam::header::record::value::{Map, map::ReferenceSequence}; + use std::num::NonZeroUsize; + use tempfile::tempdir; + + /// Build a minimal two-reference, `SO:coordinate` header for indexer + /// round-trip tests. + fn test_header() -> Header { + let mut builder = Header::builder(); + let chr1 = Map::::new(NonZeroUsize::new(1000).expect("non-zero")); + let chr2 = Map::::new(NonZeroUsize::new(1000).expect("non-zero")); + builder = builder.add_reference_sequence(b"chr1", chr1); + builder = builder.add_reference_sequence(b"chr2", chr2); + let header = builder.build(); + fgumi_sort::create_output_header(fgumi_sort::SortOrder::Coordinate, &header) + } + + /// Build the raw BAM bytes for a single 10-base mapped record. + #[allow(clippy::cast_possible_truncation)] + fn record_bytes(ref_id: i32, pos: i32, name: &[u8]) -> Vec { + let name_with_null = name.len() + 1; + let padding = (4 - (name_with_null % 4)) % 4; + let l_read_name = (name_with_null + padding) as u8; + + let mut record = Vec::with_capacity(64); + record.extend_from_slice(&ref_id.to_le_bytes()); + record.extend_from_slice(&pos.to_le_bytes()); + record.push(l_read_name); + record.push(60_u8); // mapq + record.extend_from_slice(&4681_u16.to_le_bytes()); // bin + record.extend_from_slice(&1_u16.to_le_bytes()); // n_cigar_op + record.extend_from_slice(&0_u16.to_le_bytes()); // flag + record.extend_from_slice(&10_u32.to_le_bytes()); // l_seq + record.extend_from_slice(&(-1_i32).to_le_bytes()); // next_ref_id + record.extend_from_slice(&(-1_i32).to_le_bytes()); // next_pos + record.extend_from_slice(&0_i32.to_le_bytes()); // tlen + record.extend_from_slice(name); + record.push(0); + record.extend(std::iter::repeat_n(0_u8, padding)); + let cigar_op: u32 = 10 << 4; // 10M + record.extend_from_slice(&cigar_op.to_le_bytes()); + record.extend_from_slice(&[0x11_u8; 5]); // packed seq (AAAAAAAAAA) + record.extend_from_slice(&[30_u8; 10]); // qualities + record + } + + /// Smoke test: the hook reads a finished coordinate-sorted BAM and + /// emits a parseable BAI alongside it. + #[test] + fn index_bam_finalize_hook_produces_readable_bai() { + let dir = tempdir().expect("tempdir"); + let bam_path = dir.path().join("test.bam"); + + let header = test_header(); + let mut writer = + create_raw_bam_writer(&bam_path, &header, 1, 6).expect("create_raw_bam_writer"); + for (ref_id, pos, name) in + [(0_i32, 100_i32, b"r1".as_ref()), (0, 200, b"r2"), (1, 50, b"r3")] + { + writer.write_raw_record(&record_bytes(ref_id, pos, name)).expect("write"); + } + writer.finish().expect("finish"); + + let hook = IndexBamFinalizeHook { output_path: bam_path.clone() }; + hook.finalize().expect("hook finalize"); + + let bai_path = bam_path.with_extension("bam.bai"); + assert!(bai_path.exists(), "BAI not created at {}", bai_path.display()); + let index = noodles::bam::bai::fs::read(&bai_path).expect("read bai"); + assert!(!index.reference_sequences().is_empty(), "BAI has no reference sequences"); + } +} diff --git a/crates/fgumi-sort-cli/src/lib.rs b/crates/fgumi-sort-cli/src/lib.rs new file mode 100644 index 000000000..48e2ffbf0 --- /dev/null +++ b/crates/fgumi-sort-cli/src/lib.rs @@ -0,0 +1,19 @@ +//! Standalone, framework-light `fgumi sort` command. +//! +//! This crate carries the `fgumi sort` CLI command (`Sort`), its options +//! (`SortOptions`, `SortOrderArg`), and the sort-specific finalize hooks and +//! step factory. `Sort::execute` builds the typed-step pipeline directly via +//! [`fgumi_pipeline_core::Pipeline::builder`] rather than going through the +//! umbrella's monolithic `ChainBuilder`, so the dependency graph stays light +//! (`fgumi-pipeline-core`, `fgumi-pipeline-io`, `fgumi-sort`, `fgumi-bam-io`, +//! `fgumi-sam`, `fgumi-cli-common`, `fgumi-cli-macros`) with no pull on +//! `fgumi-consensus`, `fgumi-umi`, or `fgumi-simd-fastq`. + +#![deny(unsafe_code)] + +pub mod chains; +pub mod sort; +pub mod version; + +pub use fgumi_cli_common::Command; +pub use sort::{Sort, SortOptions, SortOrderArg, TMP_DIRS_ENV, parse_cell_tag, resolve_tmp_dirs}; diff --git a/crates/fgumi-sort-cli/src/sort.rs b/crates/fgumi-sort-cli/src/sort.rs new file mode 100644 index 000000000..dbf5ba364 --- /dev/null +++ b/crates/fgumi-sort-cli/src/sort.rs @@ -0,0 +1,748 @@ +//! Sort BAM files by various orderings. +//! +//! Uses high-performance raw-bytes sorting with radix sort for in-memory +//! chunks and O(1) merge comparisons via pre-computed sort keys. +//! +//! # Sort Orders +//! +//! - **Template-coordinate**: Groups paired-end reads by template position (for `fgumi group`) +//! - **Queryname**: Groups reads by read name (for `fgumi zipper`) +//! - **Coordinate**: Standard genomic coordinate order (for IGV, `fgumi review`) +//! +//! # Performance +//! +//! - 1.9x faster than samtools on template-coordinate sort +//! - Handles BAM files larger than available RAM via spill-to-disk +//! - Uses parallel sorting for in-memory chunks +//! - Configurable temp file compression (--temp-compression) +//! +//! # Verification +//! +//! Use `--verify` to check if a BAM file is correctly sorted without writing output. + +use std::path::PathBuf; +use std::sync::Arc; + +use anyhow::{Result, bail}; +use clap::Parser; +use fgumi_bam_io::create_raw_bam_reader; +use fgumi_cli_common::{ + Command, CompressionOptions, MemoryLimit, MemoryReserve, OperationTimer, parse_bool, + parse_memory, parse_memory_reserve, resolve_memory_budget, validate_file_exists, +}; +use fgumi_pipeline_core::{Pipeline, PipelineConfig}; +use fgumi_sam::SamTag; +use fgumi_sort::{QuerynameComparator, SortOrder, verify_sort_order}; +use log::info; +use parking_lot::Mutex; + +use crate::chains::{ + IndexBamFinalizeHook, SortFinalizeHook, SortStepCaptures, build_sort_step, log_sort_start, +}; + +/// Sort order for BAM files. +/// +/// Queryname sort supports sub-sort specification via `::` syntax: +/// - `queryname` — lexicographic ordering (default, fast) +/// - `queryname::lexicographic` — explicit lexicographic ordering +/// - `queryname::natural` — natural numeric ordering (samtools-compatible) +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum SortOrderArg { + /// Coordinate sort (tid → pos → strand) + Coordinate, + /// Queryname sort with lexicographic ordering (default) + Queryname, + /// Queryname sort with natural numeric ordering + QuerynameNatural, + /// Template-coordinate sort (for UMI grouping) + #[default] + TemplateCoordinate, +} + +impl SortOrderArg { + /// Parse a sort order string, supporting `::` sub-sort syntax for queryname. + /// + /// Valid values: + /// - `coordinate` + /// - `queryname` (default: lexicographic) + /// - `queryname::lexicographic` + /// - `queryname::natural` + /// - `template-coordinate` + /// + /// # Errors + /// + /// Returns an error if the string is not a valid sort order or has an + /// unrecognized sub-sort specifier. + pub fn parse(s: &str) -> Result { + match s { + "coordinate" => Ok(Self::Coordinate), + "queryname" | "queryname::lex" | "queryname::lexicographic" => Ok(Self::Queryname), + "queryname::natural" => Ok(Self::QuerynameNatural), + "template-coordinate" => Ok(Self::TemplateCoordinate), + other => { + if let Some(sub) = other.strip_prefix("queryname::") { + Err(format!( + "unknown queryname sub-sort '{sub}', expected 'lex', 'lexicographic', or 'natural'" + )) + } else { + Err(format!( + "unknown sort order '{other}', expected 'coordinate', 'queryname', \ + 'queryname::lex', 'queryname::lexicographic', 'queryname::natural', \ + or 'template-coordinate'" + )) + } + } + } + } +} + +impl From for SortOrder { + fn from(arg: SortOrderArg) -> Self { + match arg { + SortOrderArg::Coordinate => SortOrder::Coordinate, + SortOrderArg::Queryname => SortOrder::Queryname(QuerynameComparator::Lexicographic), + SortOrderArg::QuerynameNatural => SortOrder::Queryname(QuerynameComparator::Natural), + SortOrderArg::TemplateCoordinate => SortOrder::TemplateCoordinate, + } + } +} + +/// Sort a BAM file. +/// +/// Sorts BAM files using high-performance external merge-sort, supporting +/// multiple sort orders required by the fgumi pipeline. +#[derive(Debug, Parser)] +#[command( + name = "sort", + about = "\x1b[38;5;72m[ALIGNMENT]\x1b[0m \x1b[36mSort BAM file by coordinate, queryname, or template-coordinate\x1b[0m", + long_about = r#" +Sort a BAM file using high-performance external merge-sort. + +This tool provides efficient BAM sorting with support for multiple sort orders: + +SORT ORDERS: + + coordinate Standard genomic coordinate sort (tid → pos → strand). + Use for IGV visualization, variant calling, `fgumi review`. + + queryname Lexicographic read name sort (fast, default sub-sort). + queryname::lex Short alias for lexicographic ordering (same as above). + queryname::lexicographic Explicit lexicographic ordering (same as above). + queryname::natural Natural numeric ordering (samtools-compatible). + Use for `fgumi zipper`, template-level operations. + + template-coordinate Template-level position sort for UMI grouping. + Use for `fgumi group`, `fgumi dedup`, and `fgumi downsample` input. + +PERFORMANCE: + + - 1.9x faster than samtools on template-coordinate sort + - Handles BAM files larger than available RAM via spill-to-disk + - Uses parallel sorting (--threads) for in-memory chunks + - Configurable temp file compression (--temp-compression) + - Default 768M per-thread memory limit (samtools-compatible); pass + `--max-memory auto` to detect system memory (opt-in) + +EXAMPLES: + + # Sort for fgumi group input + fgumi sort -i aligned.bam -o sorted.bam --order template-coordinate + + # Sort by coordinate for IGV + fgumi sort -i input.bam -o sorted.bam --order coordinate + + # Sort by queryname for zipper + fgumi sort -i input.bam -o sorted.bam --order queryname + + # Multi-threaded sort (default 768M per thread) + fgumi sort -i input.bam -o sorted.bam --order template-coordinate --threads 8 + + # Override the per-thread memory limit + fgumi sort -i input.bam -o sorted.bam -m 2GiB --threads 8 + + # Opt in to auto-detected system memory (subtracts --memory-reserve) + fgumi sort -i input.bam -o sorted.bam -m auto --threads 8 + + # Reserve extra memory for bwa mem running in a pipeline + fgumi sort -i input.bam -o sorted.bam --memory-reserve 12GiB --threads 4 + + # Verify a BAM file is correctly sorted + fgumi sort -i sorted.bam --verify --order template-coordinate + + # Spread spill chunks across multiple temp dirs (round-robin, free-space aware) + fgumi sort -i in.bam -o out.bam -T /mnt/ssd1 -T /mnt/ssd2 + + # Same via FGUMI_TMP_DIRS env var (PATH-style list) + FGUMI_TMP_DIRS=/mnt/ssd1:/mnt/ssd2 fgumi sort -i in.bam -o out.bam +"# +)] +#[allow(clippy::struct_excessive_bools)] +pub struct Sort { + /// Input BAM file. + #[arg(short = 'i', long = "input")] + pub input: PathBuf, + + /// Output BAM file (required unless --verify is used). + #[arg(short = 'o', long = "output")] + pub output: Option, + + /// Wrap the input in a userspace async prefetch reader: a background + /// thread reads ahead so disk I/O overlaps sorting/compression. Helps + /// on slow or networked storage. Defaults to off. + #[arg(long = "async-reader", default_value_t = false, hide = true)] + pub async_reader: bool, + + /// Verify the input file is correctly sorted (no output written). + /// + /// Reads records sequentially and checks that each record's sort key + /// is >= the previous record's key. Exits 0 if sorted correctly, + /// non-zero if any records are out of order. + #[arg(long = "verify", default_value = "false", num_args = 0..=1, default_missing_value = "true", action = clap::ArgAction::Set, value_parser = parse_bool)] + pub verify: bool, + + /// Sort order. + /// + /// Queryname sort supports sub-sort specifiers: + /// `queryname` Lexicographic byte ordering (default, fast) + /// `queryname::lexicographic` Explicit lexicographic ordering (alias: `queryname::lex`) + /// `queryname::natural` Natural numeric ordering (samtools-compatible) + #[arg(long = "order", default_value = "template-coordinate", value_parser = SortOrderArg::parse)] + pub order: SortOrderArg, + + /// Per-stage sort tuning knobs (max memory, tmp dirs, etc.). + /// Flattened here so `fgumi sort` exposes them as unprefixed + /// flags (`--max-memory`, `-T`, …) and `fgumi runall` exposes + /// them as prefixed `--sort::*` via `MultiSortOptions`. + #[command(flatten)] + pub options: SortOptions, + + /// Number of threads for parallel operations. + /// + /// Used for parallel sorting of in-memory chunks and + /// multi-threaded BGZF compression. + #[arg(short = '@', short_alias = 't', long = "threads", default_value = "1")] + pub threads: usize, + + /// Compression options for output BAM. + #[command(flatten)] + pub compression: CompressionOptions, + + /// Write BAM index (.bai) alongside output. + /// + /// Only valid for coordinate sort. The index file is written to + /// `.bam.bai` after sort completes, by re-reading the + /// finished BAM (a post-pipeline pass). The BAM itself uses the + /// same multi-threaded BGZF compression as a non-`--write-index` + /// run, so output bytes are identical between the two; only the + /// extra BAI sidecar distinguishes them. + #[arg(long = "write-index", default_value = "false", num_args = 0..=1, default_missing_value = "true", action = clap::ArgAction::Set, value_parser = parse_bool)] + pub write_index: bool, +} + +/// Per-stage sort tuning options, flattened into both the standalone +/// `Sort` command (as bare `--max-memory`, `-T`, … flags) and the +/// `RunAll` command (as prefixed `--sort::max-memory`, `--sort::tmp-dir` +/// flags via the `MultiSortOptions` companion struct generated by +/// `#[multi_options]`). +/// +/// `Default` must match each field's clap `default_value` exactly — +/// `MultiSortOptions`' generated `default_value_t` reads from +/// `SortOptions::default().field`, and clap requires the default +/// shown in `--help` to round-trip through the value-parser. +#[fgumi_cli_macros::multi_options("sort", "Sort Options")] +#[derive(clap::Args, Debug, Clone)] +pub struct SortOptions { + /// Maximum memory for in-memory sorting. + /// + /// Default is "768MiB" per thread (matching samtools' 768 MiB). Pass "auto" + /// to detect system memory and subtract --memory-reserve, leaving room + /// for the OS and co-running processes (e.g. an aligner). Explicit values + /// like "512MiB", "1GiB", "4GiB" are per-thread when --memory-per-thread is + /// enabled (default). Note bare "M"/"G" are decimal (1000ⁿ); "MiB"/"GiB" are + /// binary (1024ⁿ). + /// + /// When the limit is reached, sorted chunks spill to temporary files. + #[arg(short = 'm', long = "max-memory", default_value = "768MiB", value_parser = parse_memory)] + pub max_memory: MemoryLimit, + + /// Memory to reserve for other processes when --max-memory=auto. + /// + /// "auto" (default) reserves min(10 GiB, 50% of system memory). Explicit + /// values like "10G", "8GiB" set a fixed reservation. Set higher when + /// running alongside a memory-intensive aligner (e.g. `bwa mem` with a + /// human genome index uses ~8 GiB). + /// + /// Ignored when --max-memory is set to an explicit value. + #[arg(long = "memory-reserve", default_value = "auto", value_parser = parse_memory_reserve)] + pub memory_reserve: MemoryReserve, + + /// Scale memory limit by thread count (samtools behavior). + /// + /// When enabled (default), --max-memory specifies memory per thread. + /// Total memory = `max_memory` × threads. Disable for fixed total memory. + #[arg(long = "memory-per-thread", default_value = "true", num_args = 0..=1, default_missing_value = "true", action = clap::ArgAction::Set, value_parser = parse_bool)] + pub memory_per_thread: bool, + + /// Temporary directory for intermediate files. Repeatable. + /// + /// Pass `-T ` one or more times to spread spill chunks across multiple + /// directories in free-space-aware round-robin order. Useful when one + /// filesystem is too small or slower than the aggregate of several. + /// + /// If no flags are given and the `FGUMI_TMP_DIRS` environment variable is + /// set, its value is parsed as a `PATH`-style list (colon-separated on + /// Unix, semicolon-separated on Windows) and used instead. + /// + /// If neither is provided, the system default temp directory is used. + /// For best performance, use fast SSDs. + #[arg(short = 'T', long = "tmp-dir", action = clap::ArgAction::Append)] + pub tmp_dirs: Vec, + + /// Compression level for temporary chunk files (0-9). + /// + /// Level 0 disables compression (fastest, uses most disk space). + /// Level 1 (default) provides fast compression with reasonable space savings. + /// Higher levels (up to 9) provide better compression but are slower. + #[arg(long = "temp-compression", default_value = "1", value_parser = clap::value_parser!(u32).range(0..=9))] + pub temp_compression: u32, + + /// Sort order (chain-builder slot). + /// + /// Carried here so the chain builder can read it from the bag without + /// needing a separate out-of-band parameter. Populated by `Sort::execute` + /// from `Sort::order` before constructing the step. + #[arg(skip)] + pub order: SortOrderArg, +} + +impl Default for SortOptions { + fn default() -> Self { + Self { + max_memory: MemoryLimit::default(), + memory_reserve: MemoryReserve::default(), + memory_per_thread: true, + tmp_dirs: Vec::new(), + temp_compression: 1, + order: SortOrderArg::TemplateCoordinate, + } + } +} + +/// Environment variable name for the fallback temp-dir list, parsed as a +/// `PATH`-style list when no `-T/--tmp-dir` flags are passed. +pub const TMP_DIRS_ENV: &str = "FGUMI_TMP_DIRS"; + +/// Resolve the final list of temp directories for a sort run. +/// +/// Precedence: CLI flags (if non-empty) > `FGUMI_TMP_DIRS` env var > empty. +/// Empty strings and whitespace-only entries are filtered out of the env-var +/// value so that `FGUMI_TMP_DIRS=:` or trailing separators don't produce bogus +/// paths. +#[must_use] +pub fn resolve_tmp_dirs(cli: &[PathBuf], env_value: Option<&str>) -> Vec { + if !cli.is_empty() { + return cli.to_vec(); + } + + let Some(value) = env_value else { return Vec::new() }; + if value.is_empty() { + return Vec::new(); + } + + std::env::split_paths(value) + .filter(|p| !p.as_os_str().is_empty()) + .filter(|p| !p.to_string_lossy().trim().is_empty()) + .collect() +} + +/// Parse the cell tag for template-coordinate sort/verify, returning `None` +/// for other sort orders. +/// +/// # Errors +/// +/// Currently infallible, but returns `Result` for forward compatibility with +/// configurable cell tags. +pub fn parse_cell_tag(order: SortOrderArg) -> Result> { + if matches!(order, SortOrderArg::TemplateCoordinate) { Ok(Some(SamTag::CB)) } else { Ok(None) } +} + +impl Command for Sort { + fn execute(&self, command_line: &str) -> Result<()> { + if self.verify && self.output.is_some() { + bail!("--verify cannot be used with --output"); + } + if self.verify && self.write_index { + bail!("--write-index cannot be used with --verify"); + } + + // Validate inputs. Exempt stdin paths (`-` / `/dev/stdin`): the + // streaming sort path reads stdin once (the sort engine's reader + // handles stdin directly), so a file-existence check would spuriously + // reject it. `--verify` is the exception: it re-scans the input, which + // a non-seekable stdin can't satisfy, so reject stdin there up front. + if fgumi_bam_io::is_stdin_path(&self.input) { + if self.verify { + bail!( + "fgumi sort --verify cannot read from stdin (it re-scans the input); \ + provide a file path instead" + ); + } + } else { + validate_file_exists(&self.input, "Input BAM")?; + } + + // Either --output or --verify must be specified + if !self.verify && self.output.is_none() { + bail!("Either --output or --verify must be specified"); + } + + if self.verify { + // --verify is a standalone read-and-check path: no pipeline, just a + // raw record-stream walk that compares each key to the previous. + return self.execute_verify(); + } + + let output = self.output.as_ref().expect("output required for sort mode"); + + // BAI is only defined for coordinate sort. + if self.write_index && !matches!(self.order, SortOrderArg::Coordinate) { + bail!("--write-index is only valid for coordinate sort"); + } + + // Copy order into SortOptions so the step factory can read it. + let mut sort_opts = self.options.clone(); + sort_opts.order = self.order; + + self.execute_sort(sort_opts, output, command_line) + } +} + +impl Sort { + /// Execute the non-verify sort path by building the typed-step pipeline + /// directly (no chains monolith): a single `Exclusive` `SortBamFile` step + /// registered as a source, run on one framework driver thread while the + /// sort engine's own worker pool provides internal concurrency. Finalize + /// actions (summary log, optional BAI index) run after `Pipeline::run`. + fn execute_sort( + &self, + sort_opts: SortOptions, + output: &std::path::Path, + command_line: &str, + ) -> Result<()> { + let input_path = self.input.clone(); + let output_path = output.to_path_buf(); + let num_sorter_threads = self.threads.max(1); + + let effective_memory = resolve_memory_budget( + sort_opts.max_memory, + sort_opts.memory_reserve, + num_sorter_threads, + sort_opts.memory_per_thread, + )?; + + let timer = OperationTimer::new("Sorting BAM"); + log_sort_start(&sort_opts, &input_path, &output_path, num_sorter_threads, effective_memory); + + let stats_slot = Arc::new(Mutex::new(None::)); + + let sort_step = build_sort_step(SortStepCaptures { + sort: sort_opts, + input_path, + output_path: output_path.clone(), + num_sorter_threads, + effective_memory, + output_compression: self.compression.compression_level, + command_line: command_line.to_string(), + stats_slot: Arc::clone(&stats_slot), + })?; + + // SortBamFile has `Input = ()` and `Outputs = ()`, so it registers via + // `append_source` and the builder won't flag unwired branches. A single + // Exclusive step needs only one framework driver thread. + let builder = Pipeline::builder(); + let _ = builder.append_source(sort_step); + let pipeline = + builder.build().map_err(|e| anyhow::anyhow!("Pipeline build failed: {e}"))?; + let config = PipelineConfig { threads: 1, ..Default::default() }; + + // Always drain the summary finalize action, even on the error path, so + // the summary logs. The pipeline-run error takes precedence. + let run_result = pipeline.run(config).map_err(|e| anyhow::anyhow!("Pipeline::run: {e:?}")); + + SortFinalizeHook { stats_slot, output_path: output_path.clone(), timer }.finalize(); + + // Only write the BAM index after a successful sort: a failed run leaves + // the output BAM incomplete, so emitting (or overwriting) `.bam.bai` + // would publish a stale index for a partial file. + run_result?; + if self.write_index { + IndexBamFinalizeHook { output_path }.finalize()?; + } + Ok(()) + } + + /// Parse the cell tag for template-coordinate sort/verify, returning `None` + /// for other sort orders. + fn parse_cell_tag(&self) -> Result> { + parse_cell_tag(self.order) + } + + /// Execute verify mode: read records and check sort order. + fn execute_verify(&self) -> Result<()> { + use fgumi_sort::RawBamRecordReader; + use fgumi_sort::{ + LibraryLookup, RawQuerynameKey, RawQuerynameLexKey, RawSortKey, SortContext, cb_hasher, + extract_coordinate_key_inline, extract_template_key_inline, + }; + use std::cmp::Ordering; + use std::fs::File; + + let cell_tag = self.parse_cell_tag()?; + + let timer = OperationTimer::new("Verifying BAM sort order"); + + info!("Starting Sort Verification"); + info!("Input: {}", self.input.display()); + info!("Expected order: {:?}", self.order); + if let Some(ct) = cell_tag { + let ct_bytes = *ct; + info!("Cell tag: {}{}", ct_bytes[0] as char, ct_bytes[1] as char); + } + + // Get header via the raw-byte reader, then re-open for raw record iteration. + let (_, header) = create_raw_bam_reader(&self.input, 1)?; + + let file = File::open(&self.input)?; + let mut raw_reader = RawBamRecordReader::new(file)?; + raw_reader.skip_header()?; + + let (total_records, violations, first_violation) = match self.order { + SortOrderArg::Coordinate => { + let nref = u32::try_from(header.reference_sequences().len()).unwrap_or(u32::MAX); + verify_sort_order( + raw_reader, + |bam| extract_coordinate_key_inline(bam, nref), + |key, prev| key < prev, + )? + } + SortOrderArg::Queryname => { + let ctx = SortContext::from_header(&header); + verify_sort_order( + raw_reader, + |bam| RawQuerynameLexKey::extract(bam, &ctx), + |key, prev| key < prev, + )? + } + SortOrderArg::QuerynameNatural => { + let ctx = SortContext::from_header(&header); + verify_sort_order( + raw_reader, + |bam| RawQuerynameKey::extract(bam, &ctx), + |key, prev| key < prev, + )? + } + SortOrderArg::TemplateCoordinate => { + let lib_lookup = LibraryLookup::from_header(&header); + let hasher = cb_hasher(); + verify_sort_order( + raw_reader, + |bam| extract_template_key_inline(bam, &lib_lookup, cell_tag, &hasher), + // Use core_cmp to ignore name_hash tie-breaker differences + // so both fgumi- and samtools-sorted files pass. + |key, prev| key.core_cmp(prev) == Ordering::Less, + )? + } + }; + + info!("=== Verification Summary ==="); + info!("Records checked: {total_records}"); + info!("Sort order violations: {violations}"); + + if violations > 0 { + if let Some((record_num, name)) = first_violation { + info!("First violation at record {record_num}: {name}"); + } + timer.log_completion(total_records); + bail!( + "BAM file is NOT correctly sorted by {:?}: {violations} violations found", + self.order + ); + } + + info!("Result: PASS - file is correctly sorted by {:?}", self.order); + timer.log_completion(total_records); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use clap::Parser; + use rstest::rstest; + + #[test] + fn test_resolve_tmp_dirs_empty() { + assert!(resolve_tmp_dirs(&[], None).is_empty()); + assert!(resolve_tmp_dirs(&[], Some("")).is_empty()); + } + + #[test] + fn test_resolve_tmp_dirs_cli_only() { + let cli = vec![PathBuf::from("/tmp/a"), PathBuf::from("/tmp/b")]; + let got = resolve_tmp_dirs(&cli, None); + assert_eq!(got, cli); + } + + #[test] + fn test_resolve_tmp_dirs_env_only() { + #[cfg(unix)] + let env = "/tmp/x:/tmp/y"; + #[cfg(windows)] + let env = "C:/tmp/x;C:/tmp/y"; + + let got = resolve_tmp_dirs(&[], Some(env)); + assert_eq!(got.len(), 2); + assert!(got[0].to_string_lossy().ends_with('x')); + assert!(got[1].to_string_lossy().ends_with('y')); + } + + #[test] + fn test_resolve_tmp_dirs_cli_overrides_env() { + let cli = vec![PathBuf::from("/tmp/cli")]; + #[cfg(unix)] + let env = "/tmp/env1:/tmp/env2"; + #[cfg(windows)] + let env = "C:/tmp/env1;C:/tmp/env2"; + + let got = resolve_tmp_dirs(&cli, Some(env)); + assert_eq!(got, cli, "CLI flags must take precedence over env var"); + } + + #[rstest] + #[case("coordinate", SortOrderArg::Coordinate)] + #[case("queryname", SortOrderArg::Queryname)] + #[case("queryname::lex", SortOrderArg::Queryname)] + #[case("queryname::lexicographic", SortOrderArg::Queryname)] + #[case("queryname::natural", SortOrderArg::QuerynameNatural)] + #[case("template-coordinate", SortOrderArg::TemplateCoordinate)] + fn test_sort_order_parse_valid(#[case] input: &str, #[case] expected: SortOrderArg) { + assert_eq!(SortOrderArg::parse(input).unwrap(), expected); + } + + #[test] + fn test_sort_order_parse_invalid() { + assert!(SortOrderArg::parse("bogus").is_err()); + assert!(SortOrderArg::parse("queryname::bogus").is_err()); + } + + #[test] + fn test_sort_order_to_sort_order() { + assert_eq!(SortOrder::from(SortOrderArg::Coordinate), SortOrder::Coordinate); + assert_eq!( + SortOrder::from(SortOrderArg::TemplateCoordinate), + SortOrder::TemplateCoordinate + ); + } + + #[test] + fn test_parse_cell_tag() { + assert_eq!(parse_cell_tag(SortOrderArg::TemplateCoordinate).unwrap(), Some(SamTag::CB)); + assert_eq!(parse_cell_tag(SortOrderArg::Coordinate).unwrap(), None); + assert_eq!(parse_cell_tag(SortOrderArg::Queryname).unwrap(), None); + } + + #[test] + fn test_sort_options_default_matches_clap() { + // The Sort command must parse with only required args, taking all + // SortOptions defaults. + let sort = Sort::try_parse_from(["sort", "-i", "in.bam", "-o", "out.bam"]).unwrap(); + let defaults = SortOptions::default(); + assert_eq!(sort.options.max_memory, defaults.max_memory); + assert_eq!(sort.options.memory_reserve, defaults.memory_reserve); + assert_eq!(sort.options.memory_per_thread, defaults.memory_per_thread); + assert_eq!(sort.options.temp_compression, defaults.temp_compression); + } + + /// Write a minimal valid coordinate-sorted BAM (one mapped record) to + /// `path`, so a pre-placed output file looks indexable to the BAI hook. + #[allow(clippy::cast_possible_truncation)] + fn write_coordinate_bam(path: &std::path::Path) { + use noodles::sam::Header; + use noodles::sam::header::record::value::{Map, map::ReferenceSequence}; + use std::num::NonZeroUsize; + + let mut builder = Header::builder(); + builder = builder.add_reference_sequence( + b"chr1", + Map::::new(NonZeroUsize::new(1000).expect("non-zero")), + ); + let header = + fgumi_sort::create_output_header(fgumi_sort::SortOrder::Coordinate, &builder.build()); + + // One 10-base mapped record on chr1:100 (BAM record body bytes). + let name = b"r1"; + let name_with_null = name.len() + 1; + let padding = (4 - (name_with_null % 4)) % 4; + let mut record = Vec::with_capacity(64); + record.extend_from_slice(&0_i32.to_le_bytes()); // ref_id + record.extend_from_slice(&100_i32.to_le_bytes()); // pos + record.push((name_with_null + padding) as u8); // l_read_name + record.push(60_u8); // mapq + record.extend_from_slice(&4681_u16.to_le_bytes()); // bin + record.extend_from_slice(&1_u16.to_le_bytes()); // n_cigar_op + record.extend_from_slice(&0_u16.to_le_bytes()); // flag + record.extend_from_slice(&10_u32.to_le_bytes()); // l_seq + record.extend_from_slice(&(-1_i32).to_le_bytes()); // next_ref_id + record.extend_from_slice(&(-1_i32).to_le_bytes()); // next_pos + record.extend_from_slice(&0_i32.to_le_bytes()); // tlen + record.extend_from_slice(name); + record.push(0); + record.extend(std::iter::repeat_n(0_u8, padding)); + record.extend_from_slice(&(10_u32 << 4).to_le_bytes()); // 10M cigar + record.extend_from_slice(&[0x11_u8; 5]); // packed seq + record.extend_from_slice(&[30_u8; 10]); // qualities + + let mut writer = fgumi_bam_io::create_raw_bam_writer(path, &header, 1, 1) + .expect("create_raw_bam_writer"); + writer.write_raw_record(&record).expect("write record"); + writer.finish().expect("finish"); + } + + /// A failed sort must not write (or overwrite) the `.bam.bai`: the index + /// finalize hook only runs once the pipeline run succeeds. + #[test] + fn execute_sort_skips_index_when_sort_fails() { + let dir = tempfile::tempdir().expect("tempdir"); + let missing_input = dir.path().join("does-not-exist.bam"); + let output = dir.path().join("out.bam"); + + // Pre-place a valid coordinate BAM at the output path. The sort opens its + // (missing) input before touching the output, so this file survives — + // making the BAI's presence a clean signal of whether the index hook ran. + write_coordinate_bam(&output); + let bai = output.with_extension("bam.bai"); + assert!(!bai.exists(), "precondition: no index before the failed sort"); + + // A Sort configured to write the index on success. + let sort = Sort::try_parse_from([ + "sort", + "--input", + missing_input.to_str().expect("utf8"), + "--output", + output.to_str().expect("utf8"), + "--order", + "coordinate", + "--write-index", + "true", + ]) + .expect("parse Sort"); + + let mut sort_opts = sort.options.clone(); + sort_opts.order = sort.order; + let result = sort.execute_sort(sort_opts, &output, "fgumi sort test"); + + assert!(result.is_err(), "sort should fail when the input is missing"); + assert!(!bai.exists(), "index hook must not run after a failed sort"); + } +} diff --git a/crates/fgumi-sort-cli/src/version.rs b/crates/fgumi-sort-cli/src/version.rs new file mode 100644 index 000000000..b75825434 --- /dev/null +++ b/crates/fgumi-sort-cli/src/version.rs @@ -0,0 +1,33 @@ +//! Version string for the `@PG` record emitted by `fgumi sort`. +//! +//! This crate is framework-light and standalone, so its built-in default is +//! the Cargo package version. The umbrella `fgumi` binary derives a richer +//! version string (git commit + dirty flag) via `built`; it installs that +//! richer string once at startup through [`set_version_override`] so the +//! `@PG VN` field on sorted BAMs matches the rest of the `fgumi` toolchain. + +use std::sync::OnceLock; + +/// Process-global version override. Set at most once by the umbrella binary's +/// `main` so the `@PG` record carries the git-augmented version. When unset +/// (e.g. a standalone `fgumi-sort-cli` consumer), the Cargo package version is +/// used. +static VERSION_OVERRIDE: OnceLock = OnceLock::new(); + +/// Install the version string used in the sort `@PG` record. +/// +/// Idempotent: only the first call takes effect (subsequent calls are ignored), +/// matching the "set once at startup" contract. Returns `true` if this call +/// installed the value, or `false` if a value was already set. +pub fn set_version_override(version: String) -> bool { + VERSION_OVERRIDE.set(version).is_ok() +} + +/// Returns the version string used in the sort `@PG` record. +/// +/// Prefers the override installed via [`set_version_override`]; otherwise falls +/// back to the Cargo package version. +#[must_use] +pub fn version_string() -> String { + VERSION_OVERRIDE.get().cloned().unwrap_or_else(|| env!("CARGO_PKG_VERSION").to_string()) +} diff --git a/src/lib/commands/command.rs b/src/lib/commands/command.rs index 54396f66c..545e5871e 100644 --- a/src/lib/commands/command.rs +++ b/src/lib/commands/command.rs @@ -1,17 +1,9 @@ //! Command trait definition for CLI commands. //! -//! This module defines the [`Command`] trait that all fgumi CLI commands implement. -//! The trait uses `enum_dispatch` for efficient dynamic dispatch across command variants. - -use anyhow::Result; -use enum_dispatch::enum_dispatch; - -/// Trait implemented by all fgumi CLI commands. -/// -/// Each command provides an `execute` method that runs the command's main logic. -/// The `command_line` parameter contains the full command invocation for @PG records. -#[enum_dispatch] -pub trait Command { - #[allow(clippy::missing_errors_doc)] - fn execute(&self, command_line: &str) -> Result<()>; -} +//! Re-export shim. The [`Command`] trait that all fgumi CLI commands implement +//! now lives in the `fgumi-cli-common` crate. Re-exporting it here keeps +//! `crate::commands::command::Command` resolving AND ensures there is ONE +//! canonical `Command` trait shared with `fgumi-sort-cli` (so the umbrella's +//! `enum_dispatch` `Commands` enum and `fgumi-sort-cli`'s `Sort` impl agree on +//! the same trait). +pub use fgumi_cli_common::Command; diff --git a/src/lib/commands/common.rs b/src/lib/commands/common.rs index f9bb9ba49..1330322b8 100644 --- a/src/lib/commands/common.rs +++ b/src/lib/commands/common.rs @@ -412,19 +412,6 @@ pub struct ThreadingOptions { pub threads: Option, } -/// Options for output compression. -/// -/// Controls BGZF compression level for BAM output files. -#[derive(Debug, Clone, Default, Args)] -pub struct CompressionOptions { - /// Compression level for output BAM (1-12). - /// - /// Level 1 is fastest with larger files. - /// Level 12 produces smallest files but is slowest. - #[arg(long, default_value_t = 1)] - pub compression_level: u32, -} - /// Pipeline debugging/diagnostics options: statistics output, deadlock /// timeout, and deadlock recovery. Flattened into every pipeline command. /// @@ -560,224 +547,14 @@ impl ThreadingOptions { // in lockstep. ////////////////////////////////////////////////////////////////////////////// -/// A memory limit, either auto-detected from the host or a fixed byte count. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum MemoryLimit { - /// Detect the (cgroup-aware) host memory and subtract the reserve. - Auto, - /// Use a fixed memory limit in bytes. - Fixed(usize), -} - -impl Default for MemoryLimit { - /// Matches the clap `default_value = "768MiB"` on `SortOptions::max_memory`. - fn default() -> Self { - Self::Fixed(768 * 1024 * 1024) - } -} - -impl std::fmt::Display for MemoryLimit { - /// Round-trips through [`parse_memory`]. `Fixed(N)` is rendered in the - /// largest binary unit that divides cleanly (`GiB` → `MiB` → `KiB` → `B`) - /// so `--help` shows e.g. `"768MiB"` rather than `"805306368B"`. - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Auto => f.write_str("auto"), - Self::Fixed(bytes) => format_binary_bytes(*bytes, f), - } - } -} - -/// How much memory to reserve for other processes (OS, aligners, etc.) when a -/// memory limit is set to [`MemoryLimit::Auto`]. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum MemoryReserve { - /// Automatic: `min(10 GiB, 50% of host memory)`. - Auto, - /// Reserve a fixed number of bytes. - Fixed(usize), -} - -impl Default for MemoryReserve { - /// Matches the clap `default_value = "auto"` on `SortOptions::memory_reserve`. - fn default() -> Self { - Self::Auto - } -} - -impl std::fmt::Display for MemoryReserve { - /// Round-trips through [`parse_memory_reserve`]. See [`MemoryLimit`]'s - /// `Display` for the formatting rule. - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Auto => f.write_str("auto"), - Self::Fixed(bytes) => format_binary_bytes(*bytes, f), - } - } -} - -/// Format a byte count in the largest binary unit that divides cleanly -/// (`G` → `M` → `K` → `B`). Used by [`MemoryLimit`]'s and [`MemoryReserve`]'s -/// `Display` so that [`parse_memory`] / [`parse_memory_reserve`] round-trip the -/// result. -fn format_binary_bytes(bytes: usize, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - const K: usize = 1024; - const M: usize = K * 1024; - const G: usize = M * 1024; - // Emit binary (Mi/Gi) suffixes: `parse_memory` reads bare `M`/`G` as - // *decimal* (1000ⁿ) and only `MiB`/`GiB` as binary (1024ⁿ), so a bare - // suffix here would not round-trip a binary value. - if bytes >= G && bytes.is_multiple_of(G) { - write!(f, "{}GiB", bytes / G) - } else if bytes >= M && bytes.is_multiple_of(M) { - write!(f, "{}MiB", bytes / M) - } else if bytes >= K && bytes.is_multiple_of(K) { - write!(f, "{}KiB", bytes / K) - } else { - write!(f, "{bytes}B") - } -} - -/// The minimum per-thread memory budget (256 MiB). -pub(crate) const MIN_MEMORY_PER_THREAD: usize = 256 * 1024 * 1024; - -/// Default auto-reserve cap: 10 GiB. -pub(crate) const AUTO_RESERVE_CAP: usize = 10 * 1024 * 1024 * 1024; - -/// Parse a memory size string into `usize` bytes, suitable for use in clap -/// value parsers. -/// -/// Delegates to [`parse_memory_size`] for numeric parsing. Plain numbers are -/// interpreted as MiB (e.g. "768" = 768 MiB). Supports human-readable formats -/// like "2GB", "1GiB", "512MiB". See [`parse_memory_size`] for full details. -fn parse_memory_bytes(s: &str, label: &str) -> Result { - let bytes = parse_memory_size(s).map_err(|e| e.to_string())?; - usize::try_from(bytes).map_err(|_| format!("{label} too large: {bytes}")) -} - -/// Parse a memory-limit string (e.g. "512M", "1G", "768", "auto"). -pub(crate) fn parse_memory(s: &str) -> Result { - let s = s.trim(); - if s.eq_ignore_ascii_case("auto") { - return Ok(MemoryLimit::Auto); - } - Ok(MemoryLimit::Fixed(parse_memory_bytes(s, "Memory size")?)) -} - -/// Parse a memory-reserve string (e.g. "10G", "auto"). -pub(crate) fn parse_memory_reserve(s: &str) -> Result { - let s = s.trim(); - if s.eq_ignore_ascii_case("auto") { - return Ok(MemoryReserve::Auto); - } - Ok(MemoryReserve::Fixed(parse_memory_bytes(s, "Memory reserve")?)) -} - -/// Resolve a [`MemoryReserve`] to a concrete byte count given total host memory. -pub(crate) fn resolve_reserve(reserve: MemoryReserve, total_memory: usize) -> usize { - match reserve { - MemoryReserve::Fixed(bytes) => bytes, - // min(10 GiB, 50% of host memory) - MemoryReserve::Auto => AUTO_RESERVE_CAP.min(total_memory / 2), - } -} - -/// Resolve a memory budget to a concrete byte count. -/// -/// For [`MemoryLimit::Auto`]: detects total host memory (cgroup-aware via -/// [`detect_total_memory`]), subtracts the reserve, and—when `per_thread` is -/// set—targets each thread's share with a 256 MiB floor. The result is then -/// **capped to the available (post-reserve) memory**, so the floor can never -/// push the total past what the host has. The reserve makes the budget shrink -/// to fit the host, which is what lets pipeline commands self-throttle instead -/// of OOM-ing. -/// -/// For [`MemoryLimit::Fixed`]: multiplies by `threads` when `per_thread` is set; -/// the reserve and host size are ignored. -/// -/// Calls [`detect_total_memory`] exactly once (it invokes `sysinfo`, which is -/// not free). -pub(crate) fn resolve_memory_budget( - limit: MemoryLimit, - reserve: MemoryReserve, - threads: usize, - per_thread: bool, -) -> anyhow::Result { - // Call once — detect_total_memory() invokes sysinfo, which is not free. - resolve_memory_budget_with_total(limit, reserve, threads, per_thread, detect_total_memory()) -} - -/// Pure resolver behind [`resolve_memory_budget`], with `total` (host memory) -/// injected so the `Auto` math is unit-testable on simulated small hosts. -fn resolve_memory_budget_with_total( - limit: MemoryLimit, - reserve: MemoryReserve, - threads: usize, - per_thread: bool, - total: usize, -) -> anyhow::Result { - if threads == 0 { - anyhow::bail!("--threads must be at least 1"); - } - - let budget = match limit { - MemoryLimit::Fixed(bytes) => { - if per_thread { - bytes - .checked_mul(threads) - .ok_or_else(|| anyhow::anyhow!("memory limit × {threads} threads overflowed"))? - } else { - bytes - } - } - MemoryLimit::Auto => { - let margin = resolve_reserve(reserve, total); - let available = total.saturating_sub(margin); - // The per-thread floor is a *target*, not a guarantee. On a small - // host (or high thread count) the floor-based budget can exceed what - // is actually available; cap it to `available` so `auto` truly - // self-throttles instead of multiplying the floor past physical - // memory — the exact OOM this feature exists to prevent (#380). - let target = if per_thread { - (available / threads) - .max(MIN_MEMORY_PER_THREAD) - .checked_mul(threads) - .ok_or_else(|| anyhow::anyhow!("auto memory budget overflowed"))? - } else { - available.max(MIN_MEMORY_PER_THREAD) - }; - let budget = target.min(available); - if budget < target { - log::warn!( - "Auto memory: capping budget to host-available {} (minimum viable target {} \ - exceeds it after reserve {}); throughput may drop but the run stays within memory", - ByteSize(budget as u64), - ByteSize(target as u64), - ByteSize(margin as u64), - ); - } - log::debug!( - "Auto memory: {} of {} ({}/thread × {} threads, reserve {})", - ByteSize(budget as u64), - ByteSize(total as u64), - ByteSize((budget / threads) as u64), - threads, - ByteSize(margin as u64), - ); - budget - } - }; - - if budget > total { - log::warn!( - "Memory budget {} exceeds total host memory {}; this may cause OOM (or, for sort, earlier spill-to-disk)", - ByteSize(budget as u64), - ByteSize(total as u64), - ); - } - - Ok(budget) -} +// The memory-limit types, parsers, and resolution logic now live in the +// `fgumi-cli-common` crate. Re-export them so `crate::commands::common::*` +// paths keep resolving AND so there is ONE `MemoryLimit`/`MemoryReserve` +// type shared with `fgumi-sort-cli`'s `SortOptions` (type unification). +pub use fgumi_cli_common::{ + MIN_MEMORY_PER_THREAD, MemoryLimit, MemoryReserve, parse_bool, parse_memory, + parse_memory_reserve, resolve_memory_budget, resolve_reserve, +}; /// Options for pipeline queue memory limits. /// @@ -965,17 +742,18 @@ impl QueueMemoryOptions { } } -/// Parses a boolean value from a string, accepting: true/false, yes/no, y/n, t/f -/// (case-insensitive). Matches sopt/fgbio behavior. -pub(crate) fn parse_bool(s: &str) -> Result { - match s.to_ascii_lowercase().as_str() { - "true" | "t" | "yes" | "y" => Ok(true), - "false" | "f" | "no" | "n" => Ok(false), - _ => Err(format!("Invalid boolean value '{s}'. Expected: true|false|yes|no|y|n|t|f")), - } -} +// `parse_bool` is re-exported from `fgumi_cli_common` above. + +// Output compression options now live in `fgumi-cli-common`; re-export so +// `crate::commands::common::CompressionOptions` keeps resolving and is the +// same type as `fgumi-sort-cli`'s `Sort::compression`. +pub use fgumi_cli_common::CompressionOptions; // Re-export from the library crate for backward compatibility. +// Now only the module's own tests reference `detect_total_memory` (the memory +// resolvers that used it moved to `fgumi-cli-common`), so gate the import to +// test builds to avoid an unused-import warning in the library build. +#[cfg(test)] pub(crate) use crate::system::detect_total_memory; pub use crate::validation::parse_memory_size; @@ -1080,16 +858,6 @@ pub fn run_new_pipeline( mod tests { use super::*; - /// Enable an at-Trace logger so `log::warn!`/`debug!` macros evaluate their - /// arguments — without an enabled logger the `log` crate skips argument - /// evaluation, leaving the formatting expressions inside the memory-budget - /// warn/debug branches unexecuted under test. nextest runs each test in its - /// own process, so `try_init` is local and idempotent. - fn enable_logging() { - let _ = - env_logger::builder().is_test(true).filter_level(log::LevelFilter::Trace).try_init(); - } - #[test] fn test_none_is_single_threaded() { let opts = ThreadingOptions::none(); @@ -1313,61 +1081,12 @@ mod tests { assert!(result <= total, "auto budget {result} exceeded host total {total}"); } - #[test] - fn test_auto_never_oversubscribes_small_host() { - enable_logging(); // exercise the cap-warning and auto-debug log branches - // Simulated 4 GiB host, 16 threads: the 256 MiB/thread floor would want - // 4 GiB before reserve, which cannot fit after the auto reserve. The - // budget must be capped to `available`, never `floor × threads`. - let total = 4 * 1024 * 1024 * 1024; // 4 GiB - let margin = resolve_reserve(MemoryReserve::Auto, total); // min(10 GiB, 2 GiB) = 2 GiB - let available = total - margin; - let budget = resolve_memory_budget_with_total( - MemoryLimit::Auto, - MemoryReserve::Auto, - 16, - true, - total, - ) - .expect("should resolve"); - assert!(budget <= available, "budget {budget} oversubscribed available {available}"); - assert!(budget <= total, "budget {budget} oversubscribed host {total}"); - } - - #[test] - fn test_auto_uses_floor_when_host_is_ample() { - // Simulated 256 GiB host, 4 threads: plenty of room, so the budget is - // the per-thread share and stays under available. - let total = 256 * 1024 * 1024 * 1024; - let margin = resolve_reserve(MemoryReserve::Auto, total); // 10 GiB cap - let available = total - margin; - let budget = resolve_memory_budget_with_total( - MemoryLimit::Auto, - MemoryReserve::Auto, - 4, - true, - total, - ) - .expect("should resolve"); - assert!(budget >= MIN_MEMORY_PER_THREAD * 4, "budget {budget} fell below the floor"); - assert!(budget <= available, "budget {budget} exceeded available {available}"); - } - - #[test] - fn test_fixed_budget_independent_of_host() { - enable_logging(); // exercise the "budget exceeds host total" warn branch - // Fixed limits ignore host size entirely (reserve is irrelevant). - let tiny_host = 512 * 1024 * 1024; - let budget = resolve_memory_budget_with_total( - MemoryLimit::Fixed(2 * 1024 * 1024 * 1024), - MemoryReserve::Auto, - 4, - false, - tiny_host, - ) - .expect("should resolve"); - assert_eq!(budget, 2 * 1024 * 1024 * 1024); - } + // NOTE: `test_auto_never_oversubscribes_small_host`, + // `test_auto_uses_floor_when_host_is_ample`, and + // `test_fixed_budget_independent_of_host` exercised the (now-private) + // `resolve_memory_budget_with_total` helper, which moved to + // `fgumi-cli-common` along with the resolver itself. Their coverage lives + // in that crate's unit tests now. #[test] fn test_queue_memory_auto_reserve_shrinks_budget() { @@ -1467,19 +1186,8 @@ mod tests { assert!(opts.calculate_memory_limit(4).is_err()); } - #[test] - fn test_auto_per_thread_overflow_is_error() { - // A pathological thread count makes the per-thread floor × threads - // overflow; this must surface as an error, not wrap. - let result = resolve_memory_budget_with_total( - MemoryLimit::Auto, - MemoryReserve::Auto, - usize::MAX, - true, - 1024 * 1024 * 1024, - ); - assert!(result.is_err()); - } + // `test_auto_per_thread_overflow_is_error` moved to `fgumi-cli-common` + // with the `resolve_memory_budget_with_total` helper it exercised. #[test] fn test_log_memory_config_exercises_both_branches() { diff --git a/src/lib/commands/sort.rs b/src/lib/commands/sort.rs index cf03de336..b8635a6a9 100644 --- a/src/lib/commands/sort.rs +++ b/src/lib/commands/sort.rs @@ -1,1054 +1,6 @@ -//! Sort BAM files by various orderings. -//! -//! Uses high-performance raw-bytes sorting with radix sort for in-memory -//! chunks and O(1) merge comparisons via pre-computed sort keys. -//! -//! # Sort Orders -//! -//! - **Template-coordinate**: Groups paired-end reads by template position (for `fgumi group`) -//! - **Queryname**: Groups reads by read name (for `fgumi zipper`) -//! - **Coordinate**: Standard genomic coordinate order (for IGV, `fgumi review`) -//! -//! # Performance -//! -//! - 1.9x faster than samtools on template-coordinate sort -//! - Handles BAM files larger than available RAM via spill-to-disk -//! - Uses parallel sorting for in-memory chunks -//! - Configurable temp file compression (--temp-compression) -//! -//! # Verification -//! -//! Use `--verify` to check if a BAM file is correctly sorted without writing output. - -use crate::logging::OperationTimer; -use crate::sam::SamTag; -use crate::validation::validate_file_exists; -use anyhow::{Result, bail}; -use clap::Parser; -use fgumi_bam_io::create_raw_bam_reader; -use fgumi_sort::{QuerynameComparator, SortOrder}; - -use log::info; -use std::path::PathBuf; - -use crate::commands::command::Command; -use crate::commands::common::{ - CompressionOptions, MemoryLimit, MemoryReserve, parse_bool, parse_memory, parse_memory_reserve, +//! Re-export shim. The `fgumi sort` command now lives in the +//! `fgumi-sort-cli` crate. +pub use fgumi_sort_cli::sort::{ + MultiSortOptions, Sort, SortOptions, SortOrderArg, TMP_DIRS_ENV, parse_cell_tag, + resolve_tmp_dirs, }; - -/// Sort order for BAM files. -/// -/// Queryname sort supports sub-sort specification via `::` syntax: -/// - `queryname` — lexicographic ordering (default, fast) -/// - `queryname::lexicographic` — explicit lexicographic ordering -/// - `queryname::natural` — natural numeric ordering (samtools-compatible) -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] -pub enum SortOrderArg { - /// Coordinate sort (tid → pos → strand) - Coordinate, - /// Queryname sort with lexicographic ordering (default) - Queryname, - /// Queryname sort with natural numeric ordering - QuerynameNatural, - /// Template-coordinate sort (for UMI grouping) - #[default] - TemplateCoordinate, -} - -impl SortOrderArg { - /// Parse a sort order string, supporting `::` sub-sort syntax for queryname. - /// - /// Valid values: - /// - `coordinate` - /// - `queryname` (default: lexicographic) - /// - `queryname::lexicographic` - /// - `queryname::natural` - /// - `template-coordinate` - /// - /// # Errors - /// - /// Returns an error if the string is not a valid sort order or has an - /// unrecognized sub-sort specifier. - pub fn parse(s: &str) -> Result { - match s { - "coordinate" => Ok(Self::Coordinate), - "queryname" | "queryname::lex" | "queryname::lexicographic" => Ok(Self::Queryname), - "queryname::natural" => Ok(Self::QuerynameNatural), - "template-coordinate" => Ok(Self::TemplateCoordinate), - other => { - if other.starts_with("queryname::") { - let sub = - other.strip_prefix("queryname::").expect("guarded by starts_with check"); - Err(format!( - "unknown queryname sub-sort '{sub}', expected 'lex', 'lexicographic', or 'natural'" - )) - } else { - Err(format!( - "unknown sort order '{other}', expected 'coordinate', 'queryname', \ - 'queryname::lex', 'queryname::lexicographic', 'queryname::natural', \ - or 'template-coordinate'" - )) - } - } - } - } -} - -impl From for SortOrder { - fn from(arg: SortOrderArg) -> Self { - match arg { - SortOrderArg::Coordinate => SortOrder::Coordinate, - SortOrderArg::Queryname => SortOrder::Queryname(QuerynameComparator::Lexicographic), - SortOrderArg::QuerynameNatural => SortOrder::Queryname(QuerynameComparator::Natural), - SortOrderArg::TemplateCoordinate => SortOrder::TemplateCoordinate, - } - } -} - -/// Sort a BAM file. -/// -/// Sorts BAM files using high-performance external merge-sort, supporting -/// multiple sort orders required by the fgumi pipeline. -#[derive(Debug, Parser)] -#[command( - name = "sort", - about = "\x1b[38;5;72m[ALIGNMENT]\x1b[0m \x1b[36mSort BAM file by coordinate, queryname, or template-coordinate\x1b[0m", - long_about = r#" -Sort a BAM file using high-performance external merge-sort. - -This tool provides efficient BAM sorting with support for multiple sort orders: - -SORT ORDERS: - - coordinate Standard genomic coordinate sort (tid → pos → strand). - Use for IGV visualization, variant calling, `fgumi review`. - - queryname Lexicographic read name sort (fast, default sub-sort). - queryname::lex Short alias for lexicographic ordering (same as above). - queryname::lexicographic Explicit lexicographic ordering (same as above). - queryname::natural Natural numeric ordering (samtools-compatible). - Use for `fgumi zipper`, template-level operations. - - template-coordinate Template-level position sort for UMI grouping. - Use for `fgumi group`, `fgumi dedup`, and `fgumi downsample` input. - -PERFORMANCE: - - - 1.9x faster than samtools on template-coordinate sort - - Handles BAM files larger than available RAM via spill-to-disk - - Uses parallel sorting (--threads) for in-memory chunks - - Configurable temp file compression (--temp-compression) - - Default 768M per-thread memory limit (samtools-compatible); pass - `--max-memory auto` to detect system memory (opt-in) - -EXAMPLES: - - # Sort for fgumi group input - fgumi sort -i aligned.bam -o sorted.bam --order template-coordinate - - # Sort by coordinate for IGV - fgumi sort -i input.bam -o sorted.bam --order coordinate - - # Sort by queryname for zipper - fgumi sort -i input.bam -o sorted.bam --order queryname - - # Multi-threaded sort (default 768M per thread) - fgumi sort -i input.bam -o sorted.bam --order template-coordinate --threads 8 - - # Override the per-thread memory limit - fgumi sort -i input.bam -o sorted.bam -m 2GiB --threads 8 - - # Opt in to auto-detected system memory (subtracts --memory-reserve) - fgumi sort -i input.bam -o sorted.bam -m auto --threads 8 - - # Reserve extra memory for bwa mem running in a pipeline - fgumi sort -i input.bam -o sorted.bam --memory-reserve 12GiB --threads 4 - - # Verify a BAM file is correctly sorted - fgumi sort -i sorted.bam --verify --order template-coordinate - - # Spread spill chunks across multiple temp dirs (round-robin, free-space aware) - fgumi sort -i in.bam -o out.bam -T /mnt/ssd1 -T /mnt/ssd2 - - # Same via FGUMI_TMP_DIRS env var (PATH-style list) - FGUMI_TMP_DIRS=/mnt/ssd1:/mnt/ssd2 fgumi sort -i in.bam -o out.bam -"# -)] -#[allow(clippy::struct_excessive_bools)] -pub struct Sort { - /// Input BAM file. - #[arg(short = 'i', long = "input")] - pub input: PathBuf, - - /// Output BAM file (required unless --verify is used). - #[arg(short = 'o', long = "output")] - pub output: Option, - - /// Wrap the input in a userspace async prefetch reader: a background - /// thread reads ahead so disk I/O overlaps sorting/compression. Helps - /// on slow or networked storage. Defaults to off. - #[arg(long = "async-reader", default_value_t = false, hide = true)] - pub async_reader: bool, - - /// Verify the input file is correctly sorted (no output written). - /// - /// Reads records sequentially and checks that each record's sort key - /// is >= the previous record's key. Exits 0 if sorted correctly, - /// non-zero if any records are out of order. - #[arg(long = "verify", default_value = "false", num_args = 0..=1, default_missing_value = "true", action = clap::ArgAction::Set, value_parser = parse_bool)] - pub verify: bool, - - /// Sort order. - /// - /// Queryname sort supports sub-sort specifiers: - /// `queryname` Lexicographic byte ordering (default, fast) - /// `queryname::lexicographic` Explicit lexicographic ordering (alias: `queryname::lex`) - /// `queryname::natural` Natural numeric ordering (samtools-compatible) - #[arg(long = "order", default_value = "template-coordinate", value_parser = SortOrderArg::parse)] - pub order: SortOrderArg, - - /// Per-stage sort tuning knobs (max memory, tmp dirs, etc.). - /// Flattened here so `fgumi sort` exposes them as unprefixed - /// flags (`--max-memory`, `-T`, …) and `fgumi runall` exposes - /// them as prefixed `--sort::*` via `MultiSortOptions`. - #[command(flatten)] - pub options: SortOptions, - - /// Number of threads for parallel operations. - /// - /// Used for parallel sorting of in-memory chunks and - /// multi-threaded BGZF compression. - #[arg(short = '@', short_alias = 't', long = "threads", default_value = "1")] - pub threads: usize, - - /// Compression options for output BAM. - #[command(flatten)] - pub compression: CompressionOptions, - - /// Write BAM index (.bai) alongside output. - /// - /// Only valid for coordinate sort. The index file is written to - /// `.bam.bai` after sort completes, by re-reading the - /// finished BAM (a post-pipeline pass). The BAM itself uses the - /// same multi-threaded BGZF compression as a non-`--write-index` - /// run, so output bytes are identical between the two; only the - /// extra BAI sidecar distinguishes them. - #[arg(long = "write-index", default_value = "false", num_args = 0..=1, default_missing_value = "true", action = clap::ArgAction::Set, value_parser = parse_bool)] - pub write_index: bool, -} - -/// Per-stage sort tuning options, flattened into both the standalone -/// `Sort` command (as bare `--max-memory`, `-T`, … flags) and the -/// `RunAll` command (as prefixed `--sort::max-memory`, `--sort::tmp-dir` -/// flags via the `MultiSortOptions` companion struct generated by -/// `#[multi_options]`). -/// -/// `Default` must match each field's clap `default_value` exactly — -/// `MultiSortOptions`' generated `default_value_t` reads from -/// `SortOptions::default().field`, and clap requires the default -/// shown in `--help` to round-trip through the value-parser. -#[fgumi_cli_macros::multi_options("sort", "Sort Options")] -#[derive(clap::Args, Debug, Clone)] -pub struct SortOptions { - /// Maximum memory for in-memory sorting. - /// - /// Default is "768MiB" per thread (matching samtools' 768 MiB). Pass "auto" - /// to detect system memory and subtract --memory-reserve, leaving room - /// for the OS and co-running processes (e.g. an aligner). Explicit values - /// like "512MiB", "1GiB", "4GiB" are per-thread when --memory-per-thread is - /// enabled (default). Note bare "M"/"G" are decimal (1000ⁿ); "MiB"/"GiB" are - /// binary (1024ⁿ). - /// - /// When the limit is reached, sorted chunks spill to temporary files. - #[arg(short = 'm', long = "max-memory", default_value = "768MiB", value_parser = parse_memory)] - pub max_memory: MemoryLimit, - - /// Memory to reserve for other processes when --max-memory=auto. - /// - /// "auto" (default) reserves min(10 GiB, 50% of system memory). Explicit - /// values like "10G", "8GiB" set a fixed reservation. Set higher when - /// running alongside a memory-intensive aligner (e.g. `bwa mem` with a - /// human genome index uses ~8 GiB). - /// - /// Ignored when --max-memory is set to an explicit value. - #[arg(long = "memory-reserve", default_value = "auto", value_parser = parse_memory_reserve)] - pub memory_reserve: MemoryReserve, - - /// Scale memory limit by thread count (samtools behavior). - /// - /// When enabled (default), --max-memory specifies memory per thread. - /// Total memory = `max_memory` × threads. Disable for fixed total memory. - #[arg(long = "memory-per-thread", default_value = "true", num_args = 0..=1, default_missing_value = "true", action = clap::ArgAction::Set, value_parser = parse_bool)] - pub memory_per_thread: bool, - - /// Temporary directory for intermediate files. Repeatable. - /// - /// Pass `-T ` one or more times to spread spill chunks across multiple - /// directories in free-space-aware round-robin order. Useful when one - /// filesystem is too small or slower than the aggregate of several. - /// - /// If no flags are given and the `FGUMI_TMP_DIRS` environment variable is - /// set, its value is parsed as a `PATH`-style list (colon-separated on - /// Unix, semicolon-separated on Windows) and used instead. - /// - /// If neither is provided, the system default temp directory is used. - /// For best performance, use fast SSDs. - #[arg(short = 'T', long = "tmp-dir", action = clap::ArgAction::Append)] - pub tmp_dirs: Vec, - - /// Compression level for temporary chunk files (0-9). - /// - /// Level 0 disables compression (fastest, uses most disk space). - /// Level 1 (default) provides fast compression with reasonable space savings. - /// Higher levels (up to 9) provide better compression but are slower. - #[arg(long = "temp-compression", default_value = "1", value_parser = clap::value_parser!(u32).range(0..=9))] - pub temp_compression: u32, - - /// Sort order (chain-builder slot). - /// - /// Carried here so the chain builder can read it from the bag without - /// needing a separate out-of-band parameter. Populated by `Sort::execute` - /// from `Sort::order` before constructing the `ChainSpec`. - #[arg(skip)] - pub order: SortOrderArg, -} - -impl Default for SortOptions { - fn default() -> Self { - Self { - max_memory: MemoryLimit::default(), - memory_reserve: MemoryReserve::default(), - memory_per_thread: true, - tmp_dirs: Vec::new(), - temp_compression: 1, - order: SortOrderArg::TemplateCoordinate, - } - } -} - -/// Environment variable name for the fallback temp-dir list, parsed as a -/// `PATH`-style list when no `-T/--tmp-dir` flags are passed. -pub(crate) const TMP_DIRS_ENV: &str = "FGUMI_TMP_DIRS"; - -/// Resolve the final list of temp directories for a sort run. -/// -/// Precedence: CLI flags (if non-empty) > `FGUMI_TMP_DIRS` env var > empty. -/// Empty strings and whitespace-only entries are filtered out of the env-var -/// value so that `FGUMI_TMP_DIRS=:` or trailing separators don't produce bogus -/// paths. -pub(crate) fn resolve_tmp_dirs(cli: &[PathBuf], env_value: Option<&str>) -> Vec { - if !cli.is_empty() { - return cli.to_vec(); - } - - let Some(value) = env_value else { return Vec::new() }; - if value.is_empty() { - return Vec::new(); - } - - std::env::split_paths(value) - .filter(|p| !p.as_os_str().is_empty()) - .filter(|p| !p.to_string_lossy().trim().is_empty()) - .collect() -} - -/// Parse the cell tag for template-coordinate sort/verify, returning `None` -/// for other sort orders. -pub(crate) fn parse_cell_tag(order: SortOrderArg) -> Result> { - if matches!(order, SortOrderArg::TemplateCoordinate) { Ok(Some(SamTag::CB)) } else { Ok(None) } -} - -use fgumi_sort::verify_sort_order; - -impl Command for Sort { - fn execute(&self, command_line: &str) -> Result<()> { - if self.verify && self.output.is_some() { - bail!("--verify cannot be used with --output"); - } - if self.verify && self.write_index { - bail!("--write-index cannot be used with --verify"); - } - - // Validate inputs. Exempt stdin paths (`-` / `/dev/stdin`): the - // streaming sort path reads stdin once (the sort engine's reader - // handles stdin directly), so a file-existence check would spuriously - // reject it — matching the stdin exemption in group/dedup/correct. - // `--verify` is the exception: it re-scans the input (header probe + - // a fresh record pass), which a non-seekable stdin can't satisfy, so - // reject stdin there up front with a clear message. - if fgumi_bam_io::is_stdin_path(&self.input) { - if self.verify { - bail!( - "fgumi sort --verify cannot read from stdin (it re-scans the input); \ - provide a file path instead" - ); - } - } else { - validate_file_exists(&self.input, "Input BAM")?; - } - - // Either --output or --verify must be specified - if !self.verify && self.output.is_none() { - bail!("Either --output or --verify must be specified"); - } - - if self.verify { - // --verify is a standalone read-and-check path: no - // `ChainSpec`, no `Pipeline`, just a raw record-stream walk - // that compares each key to the previous (see - // `execute_verify`). It is intentionally NOT modeled as a - // `SinkSpec` variant because there is no BAM/BAI output — - // the path emits only log lines + a nonzero exit on - // violations. - return self.execute_verify(); - } - - // Route the non-verify sort path through chains::build_for. - let output = self.output.as_ref().expect("output required for sort mode"); - - // BAI is only defined for coordinate sort. The cross-stage validator - // (Rule 3 in `chains::validate`) catches the spec-level constraint - // "BamWithIndex requires Stage::Sort terminal"; this CLI check - // enforces the file-format-physics constraint "BAI indexes - // coordinate-sorted BGZF offsets specifically, not template-coordinate - // or queryname". Together they cover both axes (chain shape and sort - // order) of the indexability invariant. - if self.write_index && !matches!(self.order, SortOrderArg::Coordinate) { - bail!("--write-index is only valid for coordinate sort"); - } - - // Copy order into SortOptions so the chain builder can read it from - // the bag. `--write-index` no longer rides on SortOptions: it routes - // through the sink variant below so the chain-builder's - // `IndexBamFinalizeHook` (registered by `add_sort`) owns BAI - // generation as a post-pipeline step, decoupled from the BGZF - // compression path. - let mut sort_opts = self.options.clone(); - sort_opts.order = self.order; - - let sink = if self.write_index { - crate::pipeline::chains::SinkSpec::BamWithIndex(output.clone()) - } else { - crate::pipeline::chains::SinkSpec::Bam(output.clone()) - }; - - let spec = crate::pipeline::chains::ChainSpec { - stages: vec![crate::pipeline::chains::Stage::Sort], - source: crate::pipeline::chains::SourceSpec::Bam(self.input.clone()), - sink, - stage_opts: crate::pipeline::chains::StageOptionsBag { - sort: Some(sort_opts), - ..Default::default() - }, - threading: crate::commands::common::ThreadingOptions::new(self.threads), - compression: self.compression.clone(), - scheduler: crate::commands::common::SchedulerOptions::default(), - queue_memory: crate::commands::common::QueueMemoryOptions::default(), - async_reader: self.async_reader, - command_line: command_line.to_string(), - }; - crate::pipeline::chains::build_for(spec)?.run() - } -} - -impl Sort { - /// Parse the cell tag for template-coordinate sort/verify, returning `None` - /// for other sort orders. - fn parse_cell_tag(&self) -> Result> { - parse_cell_tag(self.order) - } - - /// Execute verify mode: read records and check sort order. - fn execute_verify(&self) -> Result<()> { - use fgumi_sort::RawBamRecordReader; - use fgumi_sort::{ - LibraryLookup, RawQuerynameKey, RawQuerynameLexKey, RawSortKey, SortContext, cb_hasher, - extract_coordinate_key_inline, extract_template_key_inline, - }; - use std::cmp::Ordering; - use std::fs::File; - - let cell_tag = self.parse_cell_tag()?; - - let timer = OperationTimer::new("Verifying BAM sort order"); - - info!("Starting Sort Verification"); - info!("Input: {}", self.input.display()); - info!("Expected order: {:?}", self.order); - if let Some(ct) = cell_tag { - let ct_bytes = *ct; - info!("Cell tag: {}{}", ct_bytes[0] as char, ct_bytes[1] as char); - } - - // Get header via the raw-byte reader, then re-open for raw record iteration. - let (_, header) = create_raw_bam_reader(&self.input, 1)?; - - let file = File::open(&self.input)?; - let mut raw_reader = RawBamRecordReader::new(file)?; - raw_reader.skip_header()?; - - let (total_records, violations, first_violation) = match self.order { - SortOrderArg::Coordinate => { - let nref = header.reference_sequences().len() as u32; - verify_sort_order( - raw_reader, - |bam| extract_coordinate_key_inline(bam, nref), - |key, prev| key < prev, - )? - } - SortOrderArg::Queryname => { - let ctx = SortContext::from_header(&header); - verify_sort_order( - raw_reader, - |bam| RawQuerynameLexKey::extract(bam, &ctx), - |key, prev| key < prev, - )? - } - SortOrderArg::QuerynameNatural => { - let ctx = SortContext::from_header(&header); - verify_sort_order( - raw_reader, - |bam| RawQuerynameKey::extract(bam, &ctx), - |key, prev| key < prev, - )? - } - SortOrderArg::TemplateCoordinate => { - let lib_lookup = LibraryLookup::from_header(&header); - let hasher = cb_hasher(); - verify_sort_order( - raw_reader, - |bam| extract_template_key_inline(bam, &lib_lookup, cell_tag, &hasher), - // Use core_cmp to ignore name_hash tie-breaker differences - // This allows both fgumi and samtools sorted files to pass - |key, prev| key.core_cmp(prev) == Ordering::Less, - )? - } - }; - - // Summary - info!("=== Verification Summary ==="); - info!("Records checked: {total_records}"); - info!("Sort order violations: {violations}"); - - if violations > 0 { - if let Some((record_num, name)) = first_violation { - info!("First violation at record {record_num}: {name}"); - } - timer.log_completion(total_records); - bail!( - "BAM file is NOT correctly sorted by {:?}: {violations} violations found", - self.order - ); - } - - info!("Result: PASS - file is correctly sorted by {:?}", self.order); - timer.log_completion(total_records); - Ok(()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - // Memory-budget helpers moved to `commands::common`; import the `pub(crate)` - // items these tests exercise that are not re-exported through `super::*`. - use crate::commands::common::{ - MIN_MEMORY_PER_THREAD, detect_total_memory, resolve_memory_budget, resolve_reserve, - }; - use clap::Parser; - use rstest::rstest; - - // ======================================================================== - // Temp-dir resolution tests - // ======================================================================== - - #[test] - fn test_resolve_tmp_dirs_empty() { - assert!(resolve_tmp_dirs(&[], None).is_empty()); - assert!(resolve_tmp_dirs(&[], Some("")).is_empty()); - } - - #[test] - fn test_resolve_tmp_dirs_cli_only() { - let cli = vec![PathBuf::from("/tmp/a"), PathBuf::from("/tmp/b")]; - let got = resolve_tmp_dirs(&cli, None); - assert_eq!(got, cli); - } - - #[test] - fn test_resolve_tmp_dirs_env_only() { - #[cfg(unix)] - let env = "/tmp/x:/tmp/y"; - #[cfg(windows)] - let env = "C:/tmp/x;C:/tmp/y"; - - let got = resolve_tmp_dirs(&[], Some(env)); - assert_eq!(got.len(), 2); - assert!(got[0].to_string_lossy().ends_with('x')); - assert!(got[1].to_string_lossy().ends_with('y')); - } - - #[test] - fn test_resolve_tmp_dirs_cli_overrides_env() { - let cli = vec![PathBuf::from("/tmp/cli")]; - #[cfg(unix)] - let env = "/tmp/env1:/tmp/env2"; - #[cfg(windows)] - let env = "C:/tmp/env1;C:/tmp/env2"; - - let got = resolve_tmp_dirs(&cli, Some(env)); - assert_eq!(got, cli, "CLI flags must take precedence over env var"); - } - - #[test] - fn test_resolve_tmp_dirs_skips_empty_segments() { - // Trailing separator / empty segment must not produce a bogus empty PathBuf. - #[cfg(unix)] - let env = "/tmp/a::/tmp/b:"; - #[cfg(windows)] - let env = "C:/tmp/a;;C:/tmp/b;"; - - let got = resolve_tmp_dirs(&[], Some(env)); - assert_eq!(got.len(), 2, "empty path segments must be filtered: {got:?}"); - } - - // ======================================================================== - // Clap parsing tests for repeatable -T flag - // ======================================================================== - - #[rstest] - #[case::zero(&[], vec![])] - #[case::single_short(&["-T", "/tmp/a"], vec![PathBuf::from("/tmp/a")])] - #[case::multiple_short( - &["-T", "/tmp/a", "-T", "/tmp/b", "-T", "/tmp/c"], - vec![PathBuf::from("/tmp/a"), PathBuf::from("/tmp/b"), PathBuf::from("/tmp/c")], - )] - #[case::multiple_long( - &["--tmp-dir", "/tmp/a", "--tmp-dir", "/tmp/b"], - vec![PathBuf::from("/tmp/a"), PathBuf::from("/tmp/b")], - )] - fn test_clap_tmp_dir_repeatable(#[case] extra: &[&str], #[case] expected: Vec) { - let base = ["sort", "-i", "in.bam", "-o", "out.bam", "--order", "coordinate"]; - let args: Vec<&str> = base.iter().copied().chain(extra.iter().copied()).collect(); - let sort = Sort::try_parse_from(args).expect("parse should succeed"); - assert_eq!(sort.options.tmp_dirs, expected); - } - - /// Helper to construct a `Sort` struct with a given order. - fn make_sort(order: SortOrderArg) -> Sort { - Sort { - input: PathBuf::from("test.bam"), - output: None, - async_reader: false, - verify: false, - order, - options: SortOptions { - max_memory: MemoryLimit::Fixed(512 * 1024 * 1024), - memory_reserve: MemoryReserve::Auto, - memory_per_thread: true, - tmp_dirs: Vec::new(), - temp_compression: 1, - order, - }, - threads: 1, - compression: CompressionOptions::default(), - write_index: false, - } - } - - #[rstest] - #[case(SortOrderArg::TemplateCoordinate, Some(SamTag::CB))] - #[case(SortOrderArg::Coordinate, None)] - #[case(SortOrderArg::Queryname, None)] - fn test_parse_cell_tag(#[case] order: SortOrderArg, #[case] expected: Option) { - let sort = make_sort(order); - assert_eq!(sort.parse_cell_tag().expect("parse_cell_tag should succeed"), expected); - } - - #[test] - fn test_parse_memory_auto() { - assert_eq!( - parse_memory("auto").expect("parse_memory should succeed for 'auto'"), - MemoryLimit::Auto - ); - assert_eq!( - parse_memory("AUTO").expect("parse_memory should succeed for 'AUTO'"), - MemoryLimit::Auto - ); - assert_eq!( - parse_memory("Auto").expect("parse_memory should succeed for 'Auto'"), - MemoryLimit::Auto - ); - } - - #[test] - fn test_parse_memory_plain_numbers_as_mb() { - // Plain numbers are interpreted as MB (via parse_memory_size) - assert_eq!( - parse_memory("768").expect("parse_memory should succeed for 768"), - MemoryLimit::Fixed(768 * 1024 * 1024) - ); - assert_eq!( - parse_memory("1").expect("parse_memory should succeed for 1"), - MemoryLimit::Fixed(1024 * 1024) - ); - } - - #[test] - fn memory_limit_display_round_trips_through_parse() { - // Display must emit a binary suffix the parser reads as binary so that - // the rendered value re-parses unchanged — the path clap takes for the - // generated runall `--sort::*` `default_value_t`. Previously Display - // emitted "768M" (a binary value with a decimal-parsed suffix), so a - // 768 MiB default re-parsed as 768 MB. - for bytes in [768 * 1024 * 1024, 4 * 1024 * 1024 * 1024, 512 * 1024 * 1024, 256 * 1024] { - let limit = MemoryLimit::Fixed(bytes); - let rendered = limit.to_string(); - assert_eq!( - parse_memory(&rendered).expect("Display output must parse"), - limit, - "MemoryLimit::Fixed({bytes}) displayed as {rendered:?} did not round-trip", - ); - } - } - - #[test] - fn test_parse_memory_human_readable() { - // Suffixed values use ByteSize (decimal: G=1000^3, M=1000^2) - assert_eq!( - parse_memory("512MB").expect("parse_memory should succeed for 512MB"), - MemoryLimit::Fixed(512 * 1000 * 1000) - ); - assert_eq!( - parse_memory("1G").expect("parse_memory should succeed for 1G"), - MemoryLimit::Fixed(1_000_000_000) - ); - assert_eq!( - parse_memory("2GB").expect("parse_memory should succeed for 2GB"), - MemoryLimit::Fixed(2_000_000_000) - ); - // Binary suffixes (GiB, MiB) - assert_eq!( - parse_memory("1GiB").expect("parse_memory should succeed for 1GiB"), - MemoryLimit::Fixed(1024 * 1024 * 1024) - ); - assert_eq!( - parse_memory("512MiB").expect("parse_memory should succeed for 512MiB"), - MemoryLimit::Fixed(512 * 1024 * 1024) - ); - } - - #[test] - fn test_parse_memory_case_insensitive() { - assert_eq!( - parse_memory("512mb").expect("parse_memory should succeed for lowercase 512mb"), - MemoryLimit::Fixed(512 * 1000 * 1000) - ); - assert_eq!( - parse_memory("1gb").expect("parse_memory should succeed for lowercase 1gb"), - MemoryLimit::Fixed(1_000_000_000) - ); - } - - #[test] - fn test_parse_memory_decimal_with_suffix() { - assert_eq!( - parse_memory("1.5GB").expect("parse_memory should succeed for 1.5GB"), - MemoryLimit::Fixed(1_500_000_000) - ); - } - - #[test] - fn test_parse_memory_invalid() { - assert!(parse_memory("").is_err()); - assert!(parse_memory("abc").is_err()); - assert!(parse_memory("-1G").is_err()); - } - - #[test] - fn test_resolve_memory_limit_fixed() { - let fixed = MemoryLimit::Fixed(1024 * 1024 * 1024); // 1 GiB - // Reserve is ignored for fixed limits - let resolved = - resolve_memory_budget(fixed, MemoryReserve::Auto, 4, true).expect("should succeed"); - // Fixed + memory_per_thread: total = 1 GiB * 4 = 4 GiB - assert_eq!(resolved, 4 * 1024 * 1024 * 1024); - } - - #[test] - fn test_resolve_memory_limit_fixed_no_per_thread() { - let fixed = MemoryLimit::Fixed(4 * 1024 * 1024 * 1024); // 4 GiB - let resolved = - resolve_memory_budget(fixed, MemoryReserve::Auto, 4, false).expect("should succeed"); - // Fixed + no per-thread: total = 4 GiB - assert_eq!(resolved, 4 * 1024 * 1024 * 1024); - } - - #[test] - fn test_resolve_memory_limit_auto() { - let total = detect_total_memory(); - - let resolved = resolve_memory_budget(MemoryLimit::Auto, MemoryReserve::Auto, 4, true) - .expect("should succeed"); - // Must be at least the per-thread minimum floor (256 MiB * 4 threads) - let min_expected = MIN_MEMORY_PER_THREAD.saturating_mul(4).min(total); - assert!( - resolved >= min_expected, - "auto resolved to {resolved} bytes, expected at least {min_expected}" - ); - // And not more than effective total memory (cgroup-aware), unless total - // memory is so small that the per-thread floor already exceeds it. - if total >= MIN_MEMORY_PER_THREAD.saturating_mul(4) { - assert!(resolved <= total); - } - } - - #[test] - fn test_resolve_memory_limit_auto_no_per_thread() { - let resolved = resolve_memory_budget(MemoryLimit::Auto, MemoryReserve::Auto, 8, false) - .expect("should succeed"); - // Auto + no per-thread: should be total budget, not divided by threads - // Must be at least the minimum floor (256 MB) - assert!(resolved >= 256 * 1024 * 1024); - } - - #[test] - fn test_resolve_reserve_auto() { - let gib = 1024 * 1024 * 1024; - // 32 GiB system: min(10 GiB, 16 GiB) = 10 GiB - assert_eq!(resolve_reserve(MemoryReserve::Auto, 32 * gib), 10 * gib); - // 16 GiB system: min(10 GiB, 8 GiB) = 8 GiB - assert_eq!(resolve_reserve(MemoryReserve::Auto, 16 * gib), 8 * gib); - // 8 GiB system: min(10 GiB, 4 GiB) = 4 GiB - assert_eq!(resolve_reserve(MemoryReserve::Auto, 8 * gib), 4 * gib); - // 128 GiB system: min(10 GiB, 64 GiB) = 10 GiB - assert_eq!(resolve_reserve(MemoryReserve::Auto, 128 * gib), 10 * gib); - } - - #[test] - fn test_resolve_reserve_fixed() { - let gib = 1024 * 1024 * 1024; - assert_eq!(resolve_reserve(MemoryReserve::Fixed(12 * gib), 64 * gib), 12 * gib); - } - - #[test] - fn test_parse_memory_reserve() { - assert_eq!(parse_memory_reserve("auto").expect("should parse 'auto'"), MemoryReserve::Auto,); - assert_eq!( - parse_memory_reserve("10GiB").expect("should parse '10GiB'"), - MemoryReserve::Fixed(10 * 1024 * 1024 * 1024), - ); - assert_eq!( - parse_memory_reserve("8G").expect("should parse '8G'"), - MemoryReserve::Fixed(8_000_000_000), - ); - } - - #[test] - fn test_resolve_memory_limit_auto_with_fixed_reserve() { - // With a larger fixed reserve, auto should return less memory. - // Use modest reserve sizes (128 MiB vs 512 MiB) to stay well within - // CI runner RAM and avoid the per-thread floor clamping both results. - let large_reserve = resolve_memory_budget( - MemoryLimit::Auto, - MemoryReserve::Fixed(512 * 1024 * 1024), - 4, - true, - ) - .expect("should succeed"); - let small_reserve = resolve_memory_budget( - MemoryLimit::Auto, - MemoryReserve::Fixed(128 * 1024 * 1024), - 4, - true, - ) - .expect("should succeed"); - assert!(large_reserve < small_reserve); - } - - #[test] - fn test_sort_order_conversion() { - assert_eq!(SortOrder::from(SortOrderArg::Coordinate), SortOrder::Coordinate); - assert_eq!( - SortOrder::from(SortOrderArg::Queryname), - SortOrder::Queryname(QuerynameComparator::Lexicographic) - ); - assert_eq!( - SortOrder::from(SortOrderArg::QuerynameNatural), - SortOrder::Queryname(QuerynameComparator::Natural) - ); - assert_eq!( - SortOrder::from(SortOrderArg::TemplateCoordinate), - SortOrder::TemplateCoordinate - ); - } - - // ======================================================================== - // SortOrderArg::parse tests - // ======================================================================== - - #[rstest] - #[case("coordinate", Ok(SortOrderArg::Coordinate))] - #[case("queryname", Ok(SortOrderArg::Queryname))] - #[case("queryname::lexicographic", Ok(SortOrderArg::Queryname))] - #[case("queryname::lex", Ok(SortOrderArg::Queryname))] - #[case("queryname::natural", Ok(SortOrderArg::QuerynameNatural))] - #[case("template-coordinate", Ok(SortOrderArg::TemplateCoordinate))] - #[case("queryname::fast", Err("unknown queryname sub-sort 'fast'"))] - #[case("random", Err("unknown sort order 'random'"))] - #[case("queryname::", Err("unknown queryname sub-sort ''"))] - fn test_parse_sort_order(#[case] input: &str, #[case] expected: Result) { - match expected { - Ok(order) => assert_eq!( - SortOrderArg::parse(input).expect("parse should succeed for valid sort order"), - order - ), - Err(msg) => { - let err = SortOrderArg::parse(input) - .expect_err("parse should fail for invalid sort order"); - assert!(err.contains(msg), "expected error containing {msg:?}, got: {err}"); - } - } - } - - // ======================================================================== - // Header sub-sort tag tests - // ======================================================================== - - #[test] - fn test_queryname_lex_header_has_subsort() { - let order = SortOrder::from(SortOrderArg::Queryname); - assert_eq!(order.header_so_tag(), "queryname"); - assert_eq!(order.header_ss_tag(), Some("lexicographic")); - } - - #[test] - fn test_queryname_natural_header_has_subsort() { - let order = SortOrder::from(SortOrderArg::QuerynameNatural); - assert_eq!(order.header_so_tag(), "queryname"); - assert_eq!(order.header_ss_tag(), Some("natural")); - } - - #[test] - fn test_coordinate_header_no_subsort() { - let order = SortOrder::from(SortOrderArg::Coordinate); - assert_eq!(order.header_so_tag(), "coordinate"); - assert_eq!(order.header_ss_tag(), None); - } - - #[test] - fn test_template_coordinate_header_subsort() { - let order = SortOrder::from(SortOrderArg::TemplateCoordinate); - assert_eq!(order.header_so_tag(), "unsorted"); - assert_eq!(order.header_ss_tag(), Some("template-coordinate")); - } - - #[test] - fn test_verify_conflicts_with_output() { - let sort = Sort { - verify: true, - output: Some(PathBuf::from("out.bam")), - ..make_sort(SortOrderArg::Coordinate) - }; - let err = sort.execute("test").unwrap_err(); - assert!(err.to_string().contains("--verify cannot be used with --output")); - } - - #[test] - fn test_verify_conflicts_with_write_index() { - let sort = Sort { verify: true, write_index: true, ..make_sort(SortOrderArg::Coordinate) }; - let err = sort.execute("test").unwrap_err(); - assert!(err.to_string().contains("--write-index cannot be used with --verify")); - } - - #[test] - fn test_verify_coordinate_fails_on_unsorted() -> Result<()> { - use fgumi_sort::RawBamRecordReader; - use fgumi_sort::extract_coordinate_key_inline; - - // Build BAM with records deliberately out of coordinate order - let mut builder = crate::sam::builder::SamBuilder::new(); - let _ = builder.add_pair().name("a").contig(1).start1(100).build(); - let _ = builder.add_pair().name("b").contig(0).start1(200).build(); - - let dir = tempfile::tempdir()?; - let bam_path = dir.path().join("unsorted.bam"); - builder.write_bam(&bam_path)?; - - let file = std::fs::File::open(&bam_path)?; - let (_, header) = fgumi_bam_io::create_bam_reader(&bam_path, 1)?; - let mut reader = RawBamRecordReader::new(file)?; - reader.skip_header()?; - - let nref = header.reference_sequences().len() as u32; - let (total, violations, _) = verify_sort_order( - reader, - |bam| extract_coordinate_key_inline(bam, nref), - |key, prev| key < prev, - )?; - - assert!(total > 0); - assert!(violations > 0, "unsorted file should fail coordinate verify"); - Ok(()) - } - - /// Verifies that template-coordinate sort groups reads by CB when pairs share - /// the same outer genomic coordinates. CB is used as a hash-based tiebreaker, - /// so reads with the same CB value must appear contiguously in the output. - #[test] - fn test_template_coordinate_sorts_by_cell_barcode() -> Result<()> { - use crate::commands::command::Command; - use crate::sam::builder::SamBuilder; - use bstr::ByteSlice; - - let dir = tempfile::tempdir()?; - let input = dir.path().join("input.bam"); - let output = dir.path().join("output.bam"); - - let mut builder = SamBuilder::new(); - // Three pairs at the same position: two with CB=A and one with CB=B interleaved. - // After sorting, the two CB=A pairs must be adjacent (not split by CB=B). - let _ = builder.add_pair().name("pair_a1").contig(0).start1(100).attr("CB", "A").build(); - let _ = builder.add_pair().name("pair_b").contig(0).start1(100).attr("CB", "B").build(); - let _ = builder.add_pair().name("pair_a2").contig(0).start1(100).attr("CB", "A").build(); - builder.write_bam(&input)?; - - let mut sort = make_sort(SortOrderArg::TemplateCoordinate); - sort.input = input; - sort.output = Some(output.clone()); - sort.execute("test")?; - - let mut reader = noodles::bam::io::reader::Builder.build_from_path(&output)?; - let header = reader.read_header()?; - let records: Vec<_> = reader.record_bufs(&header).collect::>>()?; - - assert_eq!(records.len(), 6, "should have 6 records (3 pairs × 2 reads)"); - - // Collect output record names and find positions of the two CB=A pairs. - // They must be contiguous — i.e. not interleaved with the CB=B pair. - let names: Vec = records - .iter() - .map(|r| { - r.name() - .map(|n| String::from_utf8_lossy(n.as_bytes()).into_owned()) - .unwrap_or_default() - }) - .collect(); - let a_positions: Vec = names - .iter() - .enumerate() - .filter(|(_, n)| n.starts_with("pair_a")) - .map(|(i, _)| i) - .collect(); - assert_eq!(a_positions.len(), 4, "expected 4 reads for the two CB=A pairs"); - let min = a_positions[0]; - let max = *a_positions.last().unwrap(); - assert_eq!( - max - min, - 3, - "CB=A reads must be grouped together; got positions {a_positions:?}" - ); - - Ok(()) - } -} diff --git a/src/lib/pipeline/chains/builder.rs b/src/lib/pipeline/chains/builder.rs index c712a4b57..7a390a419 100644 --- a/src/lib/pipeline/chains/builder.rs +++ b/src/lib/pipeline/chains/builder.rs @@ -2020,6 +2020,7 @@ impl<'a> ChainBuilder<'a> { input_path: input_path.clone(), output_path: output_path.clone(), num_sorter_threads, + effective_memory, output_compression: self.spec.compression.compression_level, command_line: self.spec.command_line.clone(), stats_slot: Arc::clone(&stats_slot), @@ -2133,10 +2134,11 @@ impl<'a> ChainBuilder<'a> { let affinity = if num_threads.max(2) >= 3 { Affinity::Worker(1) } else { Affinity::Reader }; - let and_spill = SortAndSpill::from_sorter(sorter, &self.header, 64) - .map_err(|e| anyhow!("SortAndSpill::from_sorter: {e}"))? - .with_affinity(affinity); - let decompress = SortSpillDecompress::new(64); + let and_spill = + SortAndSpill::from_sorter(sorter, &self.header, self.tuning.per_step_byte_limit) + .map_err(|e| anyhow!("SortAndSpill::from_sorter: {e}"))? + .with_affinity(affinity); + let decompress = SortSpillDecompress::new(self.tuning.per_step_byte_limit); let merge = SortMerge::new(sort_order, self.tuning.per_step_byte_limit).with_affinity(affinity); diff --git a/src/lib/pipeline/chains/commands/sort.rs b/src/lib/pipeline/chains/commands/sort.rs index 3ef9b47e0..3af507340 100644 --- a/src/lib/pipeline/chains/commands/sort.rs +++ b/src/lib/pipeline/chains/commands/sort.rs @@ -27,16 +27,20 @@ use std::sync::Arc; use anyhow::{Result, bail}; -use bytesize::ByteSize; -use fgumi_sort::{RawExternalSorter, SortOrder}; use log::info; -use crate::commands::common::{MemoryLimit, resolve_memory_budget}; -use crate::commands::sort::{SortOptions, TMP_DIRS_ENV, parse_cell_tag, resolve_tmp_dirs}; use crate::logging::OperationTimer; use crate::pipeline::chains::builder::ChainBuilder; use crate::pipeline::chains::{BuiltPipeline, ChainSpec, FinalizeHook, SinkSpec, SourceSpec}; -use crate::pipeline::steps::sort::SortBamFile; + +// The sort step factory, its captures bundle, and the startup-banner logger now +// live in the `fgumi-sort-cli` crate. Re-export them so the umbrella's +// `ChainBuilder::add_sort` keeps importing them from this module. Their +// signatures are independent of the umbrella's `FinalizeHook` trait, so they +// can be shared verbatim. The two finalize hooks below, by contrast, implement +// the umbrella's `FinalizeHook` trait (which `fgumi-sort-cli`'s plain-method +// versions do not), so they stay umbrella-local. +pub(crate) use fgumi_sort_cli::chains::{SortStepCaptures, build_sort_step, log_sort_start}; // ───────────────────────────────────────────────────────────────────────────── // SortFinalizeHook @@ -141,84 +145,11 @@ impl FinalizeHook for IndexBamFinalizeHook { // ───────────────────────────────────────────────────────────────────────────── // Step factory +// +// `SortStepCaptures`, `build_sort_step`, and `log_sort_start` are re-exported +// from `fgumi-sort-cli` at the top of this module. // ───────────────────────────────────────────────────────────────────────────── -/// Parameters for [`build_sort_step`]. Bundles all captures that the -/// factory needs so the call site in [`ChainBuilder::add_sort`] stays -/// readable. -/// -/// `pub(crate)` — consumed only by [`ChainBuilder::add_sort`]. -pub(crate) struct SortStepCaptures { - pub(crate) sort: SortOptions, - pub(crate) input_path: std::path::PathBuf, - pub(crate) output_path: std::path::PathBuf, - pub(crate) num_sorter_threads: usize, - pub(crate) output_compression: u32, - pub(crate) command_line: String, - /// Shared slot for post-run stats retrieval. The caller (finalize hook) - /// holds the other end; the step fills it when `RawExternalSorter::sort` - /// returns. - pub(crate) stats_slot: Arc>>, -} - -/// Build the `SortBamFile` step. -/// -/// Constructs and fully configures the [`RawExternalSorter`] then wraps -/// it in a [`SortBamFile`]. The step's `Outputs = ()` means it has no -/// pipeline output branches; `PipelineBuilder::build()` will not flag it -/// as unwired. The step also has `Input = ()`, so the caller uses -/// `PipelineBuilder::append_source` to register it. -/// -/// Returns `(SortBamFile, effective_memory)` — the effective memory value -/// is not used by the step itself but is logged by `add_sort` before this -/// function is called (the logging is done in `add_sort` to keep it all in -/// one place). -/// -/// `pub(crate)` — consumed only by [`ChainBuilder::add_sort`]. -/// -/// # Errors -/// -/// Returns an error if the `MemoryLimit::Auto` initial-capacity calculation -/// overflows. -pub(crate) fn build_sort_step(cap: SortStepCaptures) -> Result { - let sort_order: SortOrder = cap.sort.order.into(); - let cell_tag = parse_cell_tag(cap.sort.order)?; - - let effective_memory = resolve_memory_budget( - cap.sort.max_memory, - cap.sort.memory_reserve, - cap.num_sorter_threads, - cap.sort.memory_per_thread, - )?; - - let mut sorter = RawExternalSorter::new(sort_order) - .memory_limit(effective_memory) - .threads(cap.num_sorter_threads) - .output_compression(cap.output_compression) - .temp_compression(cap.sort.temp_compression) - .pg_info(crate::version::VERSION.to_string(), cap.command_line); - - if matches!(cap.sort.max_memory, MemoryLimit::Auto) { - let init = 768_usize - .checked_mul(1024 * 1024) - .and_then(|b| b.checked_mul(cap.num_sorter_threads)) - .ok_or_else(|| anyhow::anyhow!("initial auto buffer size overflowed"))?; - sorter = sorter.initial_capacity(effective_memory.min(init)); - } - - if let Some(ct) = cell_tag { - sorter = sorter.cell_tag(ct); - } - - let env_value = std::env::var(TMP_DIRS_ENV).ok(); - let resolved_tmp_dirs = resolve_tmp_dirs(&cap.sort.tmp_dirs, env_value.as_deref()); - if !resolved_tmp_dirs.is_empty() { - sorter = sorter.temp_dirs(resolved_tmp_dirs); - } - - Ok(SortBamFile::new(sorter, cap.input_path, cap.output_path, cap.stats_slot)) -} - // ───────────────────────────────────────────────────────────────────────────── // build_sort_chain — 10-line delegate // ───────────────────────────────────────────────────────────────────────────── @@ -245,61 +176,6 @@ pub fn build_sort_chain(spec: ChainSpec) -> Result { chain.build() } -// ───────────────────────────────────────────────────────────────────────────── -// Logging helpers re-used by add_sort -// ───────────────────────────────────────────────────────────────────────────── - -/// Log the sort startup banner. Called from [`ChainBuilder::add_sort`]. -/// -/// `pub(crate)` — not part of the public API. -pub(crate) fn log_sort_start( - sort: &SortOptions, - input_path: &std::path::Path, - output_path: &std::path::Path, - num_sorter_threads: usize, - effective_memory: usize, -) { - let cell_tag = parse_cell_tag(sort.order).unwrap_or(None); - info!("Starting Sort"); - info!("Input: {}", input_path.display()); - info!("Output: {}", output_path.display()); - info!("Sort order: {:?}", sort.order); - if let Some(ct) = cell_tag { - let ct_bytes = *ct; - info!("Cell tag: {}{}", ct_bytes[0] as char, ct_bytes[1] as char); - } - if let MemoryLimit::Fixed(per_thread) = sort.max_memory { - if sort.memory_per_thread { - info!( - "Max memory: {} ({}/thread x {} threads)", - ByteSize(effective_memory as u64), - ByteSize(per_thread as u64), - num_sorter_threads - ); - } else { - info!("Max memory: {} (fixed)", ByteSize(effective_memory as u64)); - } - } - info!("Threads: {num_sorter_threads}"); - info!("Temp compression level: {}", sort.temp_compression); - // The "Write index: enabled" banner that the pre-Phase-4 path - // emitted here is no longer surfaced from sort's startup block — - // `IndexBamFinalizeHook` (registered by `add_sort` when - // `SinkSpec::BamWithIndex` is set) emits its own "Indexing BAM: …" - // line when the hook actually runs, which carries the same - // information at the more accurate moment. - let env_value = std::env::var(TMP_DIRS_ENV).ok(); - let resolved_tmp_dirs = resolve_tmp_dirs(&sort.tmp_dirs, env_value.as_deref()); - if !resolved_tmp_dirs.is_empty() { - let joined = resolved_tmp_dirs - .iter() - .map(|p| p.display().to_string()) - .collect::>() - .join(", "); - info!("Temp directories: {joined}"); - } -} - // ───────────────────────────────────────────────────────────────────────────── // Source/sink path helpers re-used by add_sort // ───────────────────────────────────────────────────────────────────────────── diff --git a/src/lib/pipeline/steps/sink/write_bgzf.rs b/src/lib/pipeline/steps/sink/write_bgzf.rs index 25c8a3612..eed799ae6 100644 --- a/src/lib/pipeline/steps/sink/write_bgzf.rs +++ b/src/lib/pipeline/steps/sink/write_bgzf.rs @@ -1,515 +1,3 @@ -//! `WriteBgzfFile` sink step. `Serial` + `Affinity::Writer`. Receives -//! pre-compressed `BgzfBlock`s from `BgzfCompress` and writes them -//! directly to disk. -//! -//! Worker `N - 1` (the last worker) is the only worker that ever -//! attempts the sink's mutex; other workers `Skip` this step in -//! dispatch. Mirrors legacy's "T(N-1) owns Write" pattern in -//! `pipeline/scheduler/mod.rs:282-286`. Worker N-1 still -//! round-robins through every other step; the priority loop in the -//! driver brings it back to the sink as upstream produces work. -//! -//! ### Header -//! -//! Two construction modes, eager and lazy. -//! -//! **Eager** ([`WriteBgzfFile::new`]) writes the SAM header bytes at -//! construction time: the bytes are framed -//! (`BAM_MAGIC` + `l_text` + text + `n_ref` + per-ref records) by -//! [`fgumi_bam_io::write_bam_header`], BGZF-compressed via an -//! `InlineBgzfCompressor`, and emitted to the output file before the -//! first `try_run`. Header bytes flush as their own BGZF block(s), -//! ensuring they're a clean prefix to the data blocks that follow. -//! -//! **Lazy** ([`WriteBgzfFile::new_with_handle`]) opens the file but -//! parks the header write until an upstream step resolves the -//! [`HeaderHandle`]. On every `try_run`, the writer first probes the -//! handle via [`HeaderHandle::try_get`]: if `None`, it returns -//! `StepOutcome::NoProgress`; if `Some(Ok(header))`, it writes the -//! header bytes the same way the eager path does and then proceeds to -//! drain `ctx.input`; if `Some(Err(e))`, it propagates the error. -//! Used by `AlignAndMergeStep` to inject the aligner's `@PG` (plus any -//! runtime `@RG`/`@CO`) into the merged output header. -//! -//! ### Trailer -//! -//! Once `ctx.input` is drained, `try_run` emits the 28-byte BGZF EOF marker, -//! flushes/closes the file, and returns `StepOutcome::Finished` (the writer -//! state is taken so a re-dispatch is a no-op). [`Drop`] remains the -//! best-effort net for shutdown paths that never reach a clean drain (a -//! sibling-step error or a panic): see the `Drop` impl. - -use std::fs::File; -use std::io::{self, BufWriter, Write}; -use std::path::Path; - -use fgumi_bgzf::{BGZF_EOF, InlineBgzfCompressor}; -use noodles::sam::Header; -use parking_lot::Mutex; - -use crate::pipeline::core::header::HeaderHandle; -use crate::pipeline::core::step::{Affinity, Step, StepCtx, StepKind, StepOutcome, StepProfile}; -use crate::pipeline::steps::types::BgzfBlock; - -/// `Exclusive + sticky` BAM sink that consumes pre-compressed `BgzfBlock`s. -/// -/// State is held behind a `Mutex>` because the `Step` trait's -/// uniform `Clone` bound applies. The runtime never `Clone`s `Exclusive` -/// steps; clones panic. The owning worker accesses `&mut self`. -pub struct WriteBgzfFile { - state: Mutex>, - name: &'static str, -} - -struct WriterState { - /// Buffered file. In the eager-header path the header has already - /// been written before `try_run` ever runs and `pending_header` is - /// `None`; in the lazy path `pending_header` carries the - /// `HeaderHandle` and the compression level used to frame the - /// header BGZF block(s), and `try_run` consumes it as soon as an - /// upstream step resolves the handle. - out: BufWriter, - pending_header: Option, -} - -struct PendingHeader { - handle: HeaderHandle, - compression_level: u32, -} - -impl WriteBgzfFile { - /// Open `path`, BGZF-compress and write the BAM header bytes, return - /// the sink ready to receive `BgzfBlock`s. The header must be exactly - /// what should appear in the output (caller has already applied any - /// `@PG`-record updates). - /// - /// `compression_level` is used **only for the header bytes**; data - /// `BgzfBlock`s arrive already compressed at whatever level the - /// upstream `BgzfCompress` step was configured with. - /// - /// # Errors - /// - /// Returns I/O errors from path open or header write. - pub fn new>( - path: P, - header: &Header, - compression_level: u32, - ) -> io::Result { - let file = File::create(path.as_ref())?; - // Match the legacy sort writer's 256 KiB output BufWriter - // (`crates/fgumi-sort/src/external.rs::sort_*` paths). The - // default 8 KiB buffer overflows on every BGZF block write - // (~64 KiB compressed) — 32× too small to amortize syscalls. - let mut out = BufWriter::with_capacity(256 * 1024, file); - - // BAM header bytes (BAM_MAGIC + l_text + text + n_ref + refs). - let mut header_bytes = Vec::new(); - fgumi_bam_io::write_bam_header(&mut header_bytes, header) - .map_err(|e| io::Error::other(format!("write_bam_header: {e}")))?; - - // BGZF-compress the header bytes via a one-shot inline compressor - // and emit them to the output file. This lets the header occupy - // its own BGZF block(s); the data blocks that follow are appended - // verbatim from upstream. - let mut hc = InlineBgzfCompressor::new(compression_level); - hc.write_all(&header_bytes)?; - hc.flush()?; - hc.write_blocks_to(&mut out)?; - - Ok(Self { - state: Mutex::new(Some(WriterState { out, pending_header: None })), - name: "WriteBgzfFile", - }) - } - - /// Open `path` and return the sink with the BAM header write - /// **deferred** until an upstream step resolves `handle`. - /// - /// On every `try_run`, the writer probes `handle.try_get()` until - /// it observes a resolved value. While unresolved, `try_run` - /// returns `StepOutcome::NoProgress` without popping from - /// `ctx.input`, so block-buffered backpressure builds upstream - /// naturally. Once resolved, the BAM header bytes are framed and - /// BGZF-compressed (using `compression_level`) and emitted before - /// any data block — identical on-disk layout to the eager path. - /// - /// `compression_level` is used **only for the header bytes**; - /// data `BgzfBlock`s arrive already compressed at whatever level - /// the upstream `BgzfCompress` step was configured with. - /// - /// # Errors - /// Returns I/O errors from path open. Header-write errors are - /// surfaced from `try_run` once the handle resolves. - pub fn new_with_handle>( - path: P, - handle: HeaderHandle, - compression_level: u32, - ) -> io::Result { - let file = File::create(path.as_ref())?; - let out = BufWriter::with_capacity(256 * 1024, file); - Ok(Self { - state: Mutex::new(Some(WriterState { - out, - pending_header: Some(PendingHeader { handle, compression_level }), - })), - name: "WriteBgzfFile", - }) - } - - /// Consume the pending header handle if present and resolved. - /// - /// Returns: - /// - `Ok(true)` if the header has already been written or was - /// written in this call — caller may proceed to drain input. - /// - `Ok(false)` if the handle is still unresolved — caller - /// should yield without popping from input. - /// - `Err` if the handle is poisoned or the underlying I/O fails. - fn try_write_pending_header(state: &mut WriterState) -> io::Result { - let Some(pending) = state.pending_header.as_ref() else { - return Ok(true); - }; - // Header clone is a one-time cost on the resolve transition. - // Cloning lets us drop the borrow on `pending_header` before - // mutating `state.out`, avoiding a borrow-checker conflict. - let header_clone = match pending.handle.try_get() { - None => return Ok(false), - Some(Err(e)) => return Err(e), - Some(Ok(h)) => h.clone(), - }; - let level = pending.compression_level; - - let mut header_bytes = Vec::new(); - fgumi_bam_io::write_bam_header(&mut header_bytes, &header_clone) - .map_err(|e| io::Error::other(format!("write_bam_header: {e}")))?; - let mut hc = InlineBgzfCompressor::new(level); - hc.write_all(&header_bytes)?; - hc.flush()?; - hc.write_blocks_to(&mut state.out)?; - - state.pending_header = None; - Ok(true) - } -} - -impl Step for WriteBgzfFile { - type Input = BgzfBlock; - type Outputs = (); - - fn profile(&self) -> StepProfile { - StepProfile { - name: self.name, - kind: StepKind::Serial, - // Worker N-1 (the Affinity::Writer target) drives the sink - // sticky — drains the writer queue in tight bursts before - // yielding to round-robin. Without this the writer is - // visited only once per priority loop pass, which lets the - // upstream BgzfCompress queue backpressure unevenly and - // produces large run-to-run wall-time variance on the - // 2.4M-record benchmark. - sticky: true, - output_queues: vec![], - branch_ordering: vec![], - } - } - - fn affinity(&self) -> Affinity { - // Worker N-1 is the dedicated writer; other workers Skip this step. - // Matches legacy `scheduler/mod.rs:282-286` where T(N-1) owns Write. - Affinity::Writer - } - - fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { - let mut guard = self.state.lock(); - // State already taken (trailer written on a prior pass) — the writer - // is closed; a re-dispatch is an idempotent no-op. - let Some(state) = guard.as_mut() else { - return Ok(StepOutcome::Finished); - }; - - // Lazy-header path: resolve + write the header before popping input, - // otherwise we'd consume a block we can't yet flush in the correct - // on-disk order. `header_ready` is false while the handle is unresolved. - let header_ready = Self::try_write_pending_header(state)?; - if header_ready { - if let Some(block) = ctx.input.pop() { - state.out.write_all(&block.bytes)?; - return Ok(StepOutcome::Progress); - } - } - - // No block written this call. If upstream is drained, emit the trailer - // and finish. - if ctx.input.is_drained() { - // The header must be resolved by now (drain happens after every - // upstream is drained). If it isn't, the step that owns it never - // set it (typically a crashed sibling) — surface explicitly rather - // than emit a truncated header-less BAM. - if !header_ready { - return Err(io::Error::other( - "WriteBgzfFile: input drained before HeaderHandle was resolved", - )); - } - state.out.write_all(&BGZF_EOF)?; - state.out.flush()?; - // Close the writer so `Drop` (and any re-dispatch) is a no-op. - let _ = guard.take(); - return Ok(StepOutcome::Finished); - } - Ok(StepOutcome::NoProgress) - } -} - -impl Drop for WriteBgzfFile { - /// Best-effort BGZF EOF emission for shutdown paths that bypass the clean - /// `try_run` completion (where the trailer is normally written) — e.g., a - /// pipeline error in another step that triggers cancellation before this - /// sink observes its input drained, or a panic anywhere in the worker pool. - /// Without this, a partial run leaves a header-only BAM that downstream - /// tools (samtools, noodles) flag as truncated. - /// - /// **Lazy-header path:** if the writer was constructed with - /// `new_with_handle` and the handle is still unresolved at Drop - /// time, the BAM header bytes were never written. Emitting just - /// the 28-byte BGZF EOF marker would produce an unrecognisable - /// artifact (no `BAM_MAGIC`, no `@SQ` table). We deliberately - /// skip EOF in that case: the file is left at 0 bytes — clearly - /// "nothing was written" rather than a corrupt fragment. - fn drop(&mut self) { - let mut guard = self.state.lock(); - if let Some(mut state) = guard.take() { - if state.pending_header.is_some() { - return; - } - // Ignore errors: we're already on a shutdown path and there's - // no useful place to surface them. - let _ = state.out.write_all(&BGZF_EOF); - let _ = state.out.flush(); - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - fn empty_header() -> Header { - Header::default() - } - - #[test] - fn profile_advertises_serial_writer_sink() { - let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); - let header = empty_header(); - let step = WriteBgzfFile::new(&path, &header, 1).unwrap(); - let profile = step.profile(); - assert_eq!(profile.name, "WriteBgzfFile"); - assert_eq!(profile.kind, StepKind::Serial); - assert!(profile.sticky); - assert_eq!(step.affinity(), Affinity::Writer); - assert_eq!(profile.output_queues.len(), 0); - assert_eq!(profile.branch_ordering.len(), 0); - } - - #[test] - fn header_only_round_trip() { - let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); - let header = empty_header(); - let step = WriteBgzfFile::new(&path, &header, 1).unwrap(); - // Write the trailer (BGZF EOF) directly, as `try_run`'s drained - // completion would. Manual cleanup since constructing a real StepCtx in - // unit tests is awkward; framework integration is tested in tests.rs. - let mut guard = step.state.lock(); - let mut state = guard.take().expect("state present"); - state.out.write_all(&BGZF_EOF).unwrap(); - state.out.flush().unwrap(); - drop(guard); - - let bytes = std::fs::read(&path).unwrap(); - assert!(bytes.len() >= 28, "BGZF EOF + header should be at least 28 bytes"); - assert_eq!(&bytes[0..2], &[0x1f, 0x8b], "BGZF/gzip magic at start"); - let tail = &bytes[bytes.len() - 28..]; - assert_eq!(tail, &BGZF_EOF, "file ends with BGZF EOF marker"); - } - - #[test] - fn new_with_handle_defers_header_until_resolved() { - let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); - let handle = HeaderHandle::new(); - let step = WriteBgzfFile::new_with_handle(&path, handle.clone(), 1).unwrap(); - - // Before the handle resolves the file should be empty. - let bytes_before = std::fs::read(&path).unwrap(); - assert_eq!(bytes_before.len(), 0, "no bytes written until header resolves"); - - // First try_write_pending_header probe must observe pending. - { - let mut guard = step.state.lock(); - let state = guard.as_mut().expect("state present"); - assert!(state.pending_header.is_some(), "handle still pending"); - let wrote = WriteBgzfFile::try_write_pending_header(state).unwrap(); - assert!(!wrote, "unresolved handle should yield without writing"); - assert!(state.pending_header.is_some(), "still pending after no-op probe"); - } - let bytes_mid = std::fs::read(&path).unwrap(); - assert_eq!(bytes_mid.len(), 0, "still nothing on disk after no-op probe"); - - // Resolve the handle and probe again — header bytes flush. - handle.set(empty_header()).expect("first set"); - { - let mut guard = step.state.lock(); - let state = guard.as_mut().expect("state present"); - let wrote = WriteBgzfFile::try_write_pending_header(state).unwrap(); - assert!(wrote, "resolved handle should write"); - assert!(state.pending_header.is_none(), "pending slot cleared"); - // BufWriter buffers up to 256 KiB; force a flush so the - // on-disk inspection below sees the bytes. - state.out.flush().unwrap(); - } - let bytes_after = std::fs::read(&path).unwrap(); - assert!(bytes_after.len() >= 2, "header BGZF block emitted"); - assert_eq!(&bytes_after[0..2], &[0x1f, 0x8b], "BGZF/gzip magic at start"); - - // Subsequent probe is a no-op. - { - let mut guard = step.state.lock(); - let state = guard.as_mut().expect("state present"); - let wrote = WriteBgzfFile::try_write_pending_header(state).unwrap(); - assert!(wrote, "already-resolved fast path returns true"); - } - } - - #[test] - fn new_with_handle_propagates_poison() { - let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); - let handle = HeaderHandle::new(); - let step = WriteBgzfFile::new_with_handle(&path, handle.clone(), 1).unwrap(); - - handle.poison(io::Error::new(io::ErrorKind::BrokenPipe, "aligner died")).unwrap(); - let mut guard = step.state.lock(); - let state = guard.as_mut().expect("state present"); - let err = WriteBgzfFile::try_write_pending_header(state).expect_err("poison"); - assert_eq!(err.kind(), io::ErrorKind::BrokenPipe); - assert_eq!(err.to_string(), "aligner died"); - } - - #[test] - fn lazy_path_byte_identical_to_eager_via_delayed_set() { - // Genuine parity test for the deferred path: probe returns - // `None`, then we `set`, then we probe again, then EOF. This - // exercises every branch of `try_write_pending_header` — - // unlike `new_with_handle_static_header_round_trip` which - // fast-paths through a pre-set handle. - let header = empty_header(); - - let lazy_path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); - let handle = HeaderHandle::new(); - let lazy_step = WriteBgzfFile::new_with_handle(&lazy_path, handle.clone(), 1).unwrap(); - { - let mut guard = lazy_step.state.lock(); - let state = guard.as_mut().expect("state present"); - let wrote = WriteBgzfFile::try_write_pending_header(state).unwrap(); - assert!(!wrote, "first probe with unresolved handle"); - } - handle.set(header.clone()).expect("first set"); - { - let mut guard = lazy_step.state.lock(); - let state = guard.as_mut().expect("state present"); - let wrote = WriteBgzfFile::try_write_pending_header(state).unwrap(); - assert!(wrote, "second probe writes header"); - state.out.write_all(&BGZF_EOF).unwrap(); - state.out.flush().unwrap(); - let _ = guard.take(); - } - let lazy_bytes = std::fs::read(&lazy_path).unwrap(); - - let eager_path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); - let eager_step = WriteBgzfFile::new(&eager_path, &header, 1).unwrap(); - { - let mut guard = eager_step.state.lock(); - let state = guard.as_mut().expect("state present"); - state.out.write_all(&BGZF_EOF).unwrap(); - state.out.flush().unwrap(); - let _ = guard.take(); - } - let eager_bytes = std::fs::read(&eager_path).unwrap(); - - assert_eq!( - lazy_bytes, eager_bytes, - "lazy + delayed-set must match eager bytes after a None→resolve→Ok probe sequence" - ); - } - - #[test] - fn unresolved_handle_yields_false_from_helper() { - // The drained-before-resolve path in `try_run` is: - // `if !try_write_pending_header(state)? { return Err(...) }`. - // We can't easily construct a `StepCtx` in a unit test, so - // we verify the load-bearing precondition directly: the - // helper returns Ok(false) (no panic, no error) for an - // unresolved handle, which is what that `!` branch keys - // off of. Integration coverage of the full error surface - // lives in `pipeline/steps/tests.rs`. - let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); - let handle = HeaderHandle::new(); - let step = WriteBgzfFile::new_with_handle(&path, handle, 1).unwrap(); - - let mut guard = step.state.lock(); - let state = guard.as_mut().expect("state present"); - let wrote = WriteBgzfFile::try_write_pending_header(state).unwrap(); - assert!(!wrote, "unresolved handle yields false (drain converts this to Err)"); - assert!(state.pending_header.is_some(), "slot still pending"); - } - - #[test] - fn drop_with_unresolved_handle_leaves_empty_file() { - let path = tempfile::NamedTempFile::new().unwrap(); - let path_buf = path.path().to_path_buf(); - let handle = HeaderHandle::new(); - let step = WriteBgzfFile::new_with_handle(&path_buf, handle, 1).unwrap(); - drop(step); - - let bytes = std::fs::read(&path_buf).unwrap(); - assert_eq!(bytes.len(), 0, "Drop with unresolved handle must skip EOF — see Drop doc"); - } - - #[test] - fn new_with_handle_static_header_round_trip() { - // Lazy path with a handle that was pre-set via `from_header` - // must produce a byte-identical file to the eager path. Note: - // this is a unit test of the internal helper; it does not - // drive `Step::try_run`. The genuine - // None→resolve→Ok lazy-probe path is exercised by - // `lazy_path_byte_identical_to_eager_via_delayed_set`. - let header = empty_header(); - - let lazy_path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); - let lazy_step = WriteBgzfFile::new_with_handle( - &lazy_path, - HeaderHandle::from_header(header.clone()), - 1, - ) - .unwrap(); - { - let mut guard = lazy_step.state.lock(); - let state = guard.as_mut().expect("state present"); - WriteBgzfFile::try_write_pending_header(state).unwrap(); - state.out.write_all(&BGZF_EOF).unwrap(); - state.out.flush().unwrap(); - let _ = guard.take(); - } - let lazy_bytes = std::fs::read(&lazy_path).unwrap(); - - let eager_path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); - let eager_step = WriteBgzfFile::new(&eager_path, &header, 1).unwrap(); - { - let mut guard = eager_step.state.lock(); - let state = guard.as_mut().expect("state present"); - state.out.write_all(&BGZF_EOF).unwrap(); - state.out.flush().unwrap(); - let _ = guard.take(); - } - let eager_bytes = std::fs::read(&eager_path).unwrap(); - - assert_eq!(lazy_bytes, eager_bytes, "lazy + pre-set handle must match eager bytes"); - } -} +//! Re-export shim. The `WriteBgzfFile` sink step now lives in the +//! `fgumi-pipeline-io` crate. +pub use fgumi_pipeline_io::sink::write_bgzf::WriteBgzfFile; diff --git a/src/lib/pipeline/steps/sort/and_spill/tests.rs b/src/lib/pipeline/steps/sort/and_spill/tests.rs deleted file mode 100644 index b4f91c058..000000000 --- a/src/lib/pipeline/steps/sort/and_spill/tests.rs +++ /dev/null @@ -1,61 +0,0 @@ -//! Profile-pin tests for `SortAndSpill`. -//! -//! The end-to-end behavior (record-byte parity vs `RawExternalSorter::sort()`) -//! is covered by the 3-step chain integration tests in -//! `src/lib/pipeline/steps/sort/tests.rs`. These tests pin the -//! step's `StepProfile` and basic affinity-override behavior — failures -//! here flag a contract change without having to scroll through the -//! integration-test output. - -use fgumi_sort::{RawExternalSorter, SortOrder}; -use noodles::sam::Header; - -use super::*; - -#[test] -fn profile_advertises_serial_count_bounded() { - let header = Header::default(); - let sorter = RawExternalSorter::new(SortOrder::Coordinate).memory_limit(1024 * 1024); - let s = SortAndSpill::from_sorter(sorter, &header, 32).expect("from_sorter"); - let p = s.profile(); - assert_eq!(p.name, "SortAndSpill"); - assert_eq!(p.kind, StepKind::Serial); - assert!(!p.sticky); - assert_eq!(p.output_queues.len(), 1); - match p.output_queues[0] { - QueueSpec::CountBounded { capacity } => assert_eq!(capacity, 32), - other => panic!("expected CountBounded capacity, got {other:?}"), - } - assert_eq!(p.branch_ordering, vec![BranchOrdering::None]); -} - -#[test] -fn affinity_default_is_none_and_overridable() { - let header = Header::default(); - let sorter = RawExternalSorter::new(SortOrder::Coordinate).memory_limit(1024 * 1024); - let s = SortAndSpill::from_sorter(sorter, &header, 32).expect("from_sorter"); - assert_eq!(s.affinity(), Affinity::None); - - let sorter = RawExternalSorter::new(SortOrder::Coordinate).memory_limit(1024 * 1024); - let s = SortAndSpill::from_sorter(sorter, &header, 32) - .expect("from_sorter") - .with_affinity(Affinity::Worker(2)); - assert_eq!(s.affinity(), Affinity::Worker(2)); -} - -#[test] -fn from_sorter_builds_for_all_sort_orders() { - use fgumi_sort::QuerynameComparator; - - let header = Header::default(); - for order in [ - SortOrder::Coordinate, - SortOrder::Queryname(QuerynameComparator::Lexicographic), - SortOrder::Queryname(QuerynameComparator::Natural), - SortOrder::TemplateCoordinate, - ] { - let sorter = RawExternalSorter::new(order).memory_limit(1024 * 1024); - let s = SortAndSpill::from_sorter(sorter, &header, 32); - assert!(s.is_ok(), "from_sorter must succeed for {order:?}"); - } -} diff --git a/src/lib/pipeline/steps/sort/merge/tests.rs b/src/lib/pipeline/steps/sort/merge/tests.rs deleted file mode 100644 index 1b79e5bfa..000000000 --- a/src/lib/pipeline/steps/sort/merge/tests.rs +++ /dev/null @@ -1,58 +0,0 @@ -//! Profile-pin tests for `SortMerge`. -//! -//! End-to-end record-byte parity vs `RawExternalSorter::sort()` is covered -//! by the 3-step chain integration tests in -//! `src/lib/pipeline/steps/sort/tests.rs`. These tests pin the -//! step's `StepProfile` and basic affinity / construction behavior. - -use fgumi_sort::{QuerynameComparator, SortOrder}; - -use super::*; - -#[test] -fn profile_advertises_serial_byordinal_byte_bounded() { - let s = SortMerge::new(SortOrder::Coordinate, 64 * 1024); - let p = s.profile(); - assert_eq!(p.name, "SortMerge"); - assert_eq!(p.kind, StepKind::Serial); - assert!(!p.sticky); - assert_eq!(p.output_queues.len(), 1); - match p.output_queues[0] { - QueueSpec::ByteBounded { limit_bytes } => assert_eq!(limit_bytes, 64 * 1024), - other => panic!("expected ByteBounded limit, got {other:?}"), - } - assert_eq!(p.branch_ordering, vec![BranchOrdering::ByItemOrdinal]); -} - -#[test] -fn affinity_default_is_none_and_overridable() { - let s = SortMerge::new(SortOrder::Coordinate, 64 * 1024); - assert_eq!(s.affinity(), Affinity::None); - - let s = SortMerge::new(SortOrder::Coordinate, 64 * 1024).with_affinity(Affinity::Worker(2)); - assert_eq!(s.affinity(), Affinity::Worker(2)); -} - -#[test] -fn builds_for_all_sort_orders() { - for order in [ - SortOrder::Coordinate, - SortOrder::Queryname(QuerynameComparator::Lexicographic), - SortOrder::Queryname(QuerynameComparator::Natural), - SortOrder::TemplateCoordinate, - ] { - // Construction itself doesn't fail for any order; the typed - // dispatch happens inside `build_driver` on transition to - // Merging (which is exercised in the integration tests). - let _ = SortMerge::new(order, 64 * 1024); - let _ = SortMerge::with_target_batch_count(order, 64 * 1024, 256); - } -} - -#[test] -fn target_batch_count_floor_is_one() { - // Defensive: a zero target_batch_count would cause the inner loop to - // never break by-count, so we clamp to 1 in the constructor. - let s = SortMerge::with_target_batch_count(SortOrder::Coordinate, 64 * 1024, 0); - assert_eq!(s.target_batch_count, 1); -} diff --git a/src/lib/pipeline/steps/sort/mod.rs b/src/lib/pipeline/steps/sort/mod.rs index 2ba126939..6d6bdab57 100644 --- a/src/lib/pipeline/steps/sort/mod.rs +++ b/src/lib/pipeline/steps/sort/mod.rs @@ -1,131 +1,15 @@ -//! Sort typed-steps for the unified pipeline. -//! -//! Two flavors: -//! -//! 1. **[`SortBamFile`]** — `Exclusive` single-step that wraps the -//! `RawExternalSorter::sort(input, output)` end-to-end. Used by -//! standalone `fgumi sort`. The chain is literally `[SortBamFile]` -//! — the framework provides one thread to run `sort()` on. -//! Concurrency profile is one driver thread + the sort engine's -//! own `SortWorkerPool` and rayon pools. -//! -//! 2. **Three-step chain** (`SortAndSpill` → `SortSpillDecompress` → -//! `SortMerge`) — middle-of-chain composition for `runall --sort` -//! fusion. The three steps split the legacy monolithic `Sort` body -//! so the spill decompress phase runs on the framework's -//! work-stealing pool (eliminating the +N `SortWorkerPool::Phase2` -//! OS-thread oversubscription during merge). Records stream in via -//! `SortAndSpill`, sorted records stream out the back of `SortMerge`, -//! and downstream consumers (Decode → Group → ...) run concurrently -//! with the merge tail. -//! -//! See [`and_spill`], [`spill_decompress`], and [`merge`] for the -//! per-step state machines, and `docs/design/sort-step-split.md` for -//! the locked design. - -pub mod and_spill; -pub mod merge; -pub mod protocol; -pub mod spill_decompress; - -pub use and_spill::SortAndSpill; -pub use merge::SortMerge; -pub use spill_decompress::SortSpillDecompress; - -use std::io; -use std::path::PathBuf; -use std::sync::Arc; - -use fgumi_sort::RawExternalSorter; -use parking_lot::Mutex; - -use crate::pipeline::core::step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}; - -// ───────────────────────────────────────────────────────────────────────────── -// SortBamFile — Exclusive single-step wrapping legacy sort end-to-end. -// ───────────────────────────────────────────────────────────────────────────── - -/// `Exclusive` step that drives a complete legacy -/// `RawExternalSorter::sort(input, output)` call to completion in a -/// single `try_run`. The framework gives this step a dedicated owner -/// worker, on which the entire sort runs. Concurrency profile is -/// identical to invoking the legacy sort directly: one driver thread -/// + the sort engine's `SortWorkerPool` and rayon pools. -/// -/// The chain at standalone use is just `[SortBamFile]` — there are no -/// other typed-steps. The pipeline framework adds *zero* concurrency -/// vs legacy. The wrapping exists purely so the operation is -/// composable. -/// -/// `--write-index` is *not* handled by this step: BAI generation is -/// decoupled into the chain-builder's `IndexBamFinalizeHook`, which runs -/// as a post-pipeline pass over the finished coordinate-sorted BAM. This -/// step only produces the sorted BGZF stream. -pub struct SortBamFile { - /// `Some` until the first `try_run` consumes it; `None` after. - /// Guards against a second `try_run` doing the sort again — once the step - /// reports `Finished` it is dropped from the worklist and not re-dispatched, - /// but the option makes the no-op path explicit. - sorter: Option, - input: PathBuf, - output: PathBuf, - /// Out-parameter slot for the `SortStats` produced by - /// `RawExternalSorter::sort`. Filled in by `try_run` after the - /// sort completes; the caller (e.g. standalone `Sort::execute`) - /// holds an `Arc` clone and reads the totals once the pipeline - /// returns so we can log the Records-processed / Records-written - /// / Temporary-chunks Summary block. - stats_out: Arc>>, +//! Re-export shim. The sort typed-steps now live in the `fgumi-pipeline-io` +//! crate (`fgumi_pipeline_io::sort`). +pub use fgumi_pipeline_io::sort::{ + SortAndSpill, SortBamFile, SortMerge, SortSpillDecompress, protocol, +}; + +pub mod and_spill { + pub use fgumi_pipeline_io::sort::and_spill::*; } - -impl SortBamFile { - /// Build a `SortBamFile` step. The `sorter` should be fully - /// configured identically to the standalone `Sort::execute_sort` - /// path (memory limit, threads, temp/output compression, spill - /// codec, temp dirs, max temp files, cell tag, key types, `@PG` - /// info, initial capacity, async reader) — `SortBamFile` does not - /// touch any of the sorter's tuning knobs. - /// - /// `stats_out` is the slot the step writes its `SortStats` into - /// after `sorter.sort()` returns. Pass `Arc::clone(&slot)` here and - /// take the inner value from the same slot after `Pipeline::run` - /// returns to recover the per-run counters. - #[must_use] - pub fn new( - sorter: RawExternalSorter, - input: PathBuf, - output: PathBuf, - stats_out: Arc>>, - ) -> Self { - Self { sorter: Some(sorter), input, output, stats_out } - } +pub mod merge { + pub use fgumi_pipeline_io::sort::merge::*; } - -impl Step for SortBamFile { - type Input = (); - type Outputs = (); - - fn profile(&self) -> StepProfile { - StepProfile { - name: "SortBamFile", - kind: StepKind::Exclusive, - sticky: false, - output_queues: vec![], - branch_ordering: vec![], - } - } - - fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { - let Some(sorter) = self.sorter.take() else { - return Ok(StepOutcome::Finished); - }; - let stats = sorter - .sort(&self.input, &self.output) - .map_err(|e| io::Error::other(format!("SortBamFile: sort failed: {e:#}")))?; - *self.stats_out.lock() = Some(stats); - Ok(StepOutcome::Finished) - } +pub mod spill_decompress { + pub use fgumi_pipeline_io::sort::spill_decompress::*; } - -#[cfg(test)] -mod tests; diff --git a/src/lib/pipeline/steps/sort/spill_decompress.rs b/src/lib/pipeline/steps/sort/spill_decompress.rs deleted file mode 100644 index f32ada194..000000000 --- a/src/lib/pipeline/steps/sort/spill_decompress.rs +++ /dev/null @@ -1,355 +0,0 @@ -//! `SortSpillDecompress` — Parallel typed step that reads spill chunk -//! files, decompresses them inline (codec-aware: BGZF blocks or zstd -//! frames per `slot.codec`, detected from the file magic at slot-open), -//! and pushes the decompressed bytes into per-slot bounded queues on -//! `SortMergeSlot`. Replaces the legacy +N `SortWorkerPool::Phase2` -//! OS threads with framework workers from the unified pipeline's -//! work-stealing pool. -//! -//! See `docs/design/sort-step-split-parity-fix.md` for the v4 design. -//! -//! ## Step body (`try_run`) -//! -//! 1. **Drain held output** if any (legacy `HeldSlot` pattern — pushes -//! that failed last call due to a full output queue retry first). -//! 2. **Pop one input event** if available: -//! * `SpillReady` → register the slot in the shared registry, AND -//! forward verbatim to `SortMerge` for slot-table install. -//! * `MemoryChunk` → forward verbatim to `SortMerge`. -//! * `AllAnnounced` → forward verbatim. -//! 3. **Greedy slot fill** (`try_fill_some_slot`): walk a snapshot of -//! the registry; for the first slot that has room in its -//! `decompressed` queue (i.e., `decompressed.len() < CAP`) AND a -//! free reader (`reader.try_lock()` succeeds), read up to -//! `min(CAP - len, MAX_BATCH)` raw BGZF blocks from disk, -//! decompress them inline on this worker, and push the batch into -//! `decompressed`. If the read returns fewer than requested, set -//! `queue_eof = true` (last batch). -//! 4. If no slot accepted reader/fill work: return `Contention` if -//! ANY registered slot is still not `queue_eof` (consumer may -//! drain space later); otherwise, once the input edge is drained -//! too, report `Finished`. -//! -//! Producer is **non-blocking** — never waits on consumer activity. -//! The consumer (`MergeDriver::try_step` via `slot_try_load_block`) is -//! also non-blocking: when a slot's queue is empty and not yet -//! `queue_eof` it reports `WouldBlock` and the cooperative `SortMerge` -//! step yields, so the framework re-dispatches it after this producer -//! has pushed more blocks or set `queue_eof`. No condvar / notify. -//! -//! ## Atomic ordering for `queue_eof` / `decomp_error` -//! -//! Producers' BOTH success and error paths hold `slot.decompressed` -//! while storing the atomics. Consumer reads them under the same -//! mutex. The mutex release-acquire chain establishes happens-before -//! for both atomics simultaneously; consumer's load order does not -//! matter for correctness. - -use std::io; -use std::sync::Arc; - -use fgumi_sort::{PHASE2_DECOMP_CAP, SortMergeSlot, SpillBlockDecompressor}; -use parking_lot::Mutex; - -use crate::pipeline::core::Unpushed; -use crate::pipeline::core::held::HeldSlot; -use crate::pipeline::core::outputs::Single; -use crate::pipeline::core::queues::QueueSpec; -use crate::pipeline::core::reorder::BranchOrdering; -use crate::pipeline::core::step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}; -use crate::pipeline::steps::sort::protocol::{SortPhase1Event, SortPhase2Event}; - -/// Max raw blocks read+decompressed by a single `try_fill_some_slot` -/// call. Bounds the per-call work so the framework's round-robin -/// dispatch stays responsive (other steps don't starve while this -/// worker hogs a slot). At ~ms-per-block decompress and N=4, this is -/// ~4ms of work per call. -const MAX_BATCH_PER_CALL: usize = 4; - -/// A slot announced via `SpillReady` and now eligible for greedy fill. -/// Stored in the step's shared registry; all Parallel worker copies -/// read it via the same `Arc`. -struct RegisteredSpill { - slot: Arc, -} - -/// Parallel step that reads + decompresses spill chunk files and -/// pushes results into per-slot queues. Forwards `SortPhase1Event`s -/// verbatim to `SortMerge`. -/// -/// All Parallel worker copies share: -/// -/// * `registry: Arc>>` — files announced -/// via `SpillReady`. Append-only; per-worker `try_run` walks this -/// list when looking for fill work. -/// -/// Per-worker: -/// -/// * `block_dec: SpillBlockDecompressor` — fresh per worker copy; the -/// codec-aware (BGZF block / zstd frame) decompressor that reads + -/// inflates a batch of blocks and returns owned `Vec`s for the -/// per-slot queue (mimalloc size-class buffer reuse inside). -/// * `held: HeldSlot>` — backpressure retry -/// slot for the forwarded events (`SpillReady`/`MemoryChunk`/ -/// `AllAnnounced`). The per-slot decompressed queue handles its -/// own backpressure separately (producer SKIPs when at cap). -pub struct SortSpillDecompress { - registry: Arc>>, - /// Per-worker codec-aware decompressor (BGZF blocks or zstd frames). - block_dec: SpillBlockDecompressor, - held: HeldSlot>, - /// Max in-flight forwarded events. `SortSpillDecompress` emits at - /// most one event per `try_run` (forward from input pop). 64 - /// covers any realistic spill count. - output_capacity: usize, -} - -impl SortSpillDecompress { - /// Construct a fresh step with an empty registry. - #[must_use] - pub fn new(output_capacity: usize) -> Self { - Self { - registry: Arc::new(Mutex::new(Vec::new())), - block_dec: SpillBlockDecompressor::new(), - held: HeldSlot::new(), - output_capacity, - } - } - - /// Attempt to deliver any held output event. Returns `true` if - /// the held slot is now empty. - fn flush_held(&mut self, ctx: &mut StepCtx<'_, Self>) -> bool { - let Some(unpushed) = self.held.take() else { - return true; - }; - match ctx.outputs.retry(unpushed) { - Ok(()) => true, - Err(again) => { - self.held.put(again); - false - } - } - } - - /// Push an output event, stashing it on backpressure. Returns - /// `true` on success. - fn push_or_hold(&mut self, ctx: &mut StepCtx<'_, Self>, event: SortPhase2Event) -> bool { - match ctx.outputs.push(event) { - Ok(()) => true, - Err(unpushed) => { - self.held.put(unpushed); - false - } - } - } - - /// Snapshot the registry's slots so the caller can walk them - /// without holding the registry mutex across disk I/O or - /// decompression. - fn snapshot_registry(&self) -> Vec> { - let registry = self.registry.lock(); - registry.iter().map(|e| Arc::clone(&e.slot)).collect() - } - - /// Greedy slot-fill step. Walk registered slots; for the first - /// one with `decompressed.len() < CAP` AND a free reader, read + - /// decompress + push a batch of up to `MAX_BATCH_PER_CALL` blocks. - /// Returns `Ok(true)` if any work happened. - /// - /// **Atomic ordering**: the success path sets `queue_eof = true` - /// (when reader hits EOF) WHILE HOLDING the `decompressed` mutex, - /// so the consumer's next lock-acquire synchronizes-with this - /// release and sees both the pushed batch AND the eof flag. - /// The error path same discipline (sets `decomp_error` + - /// `queue_eof` while holding `decompressed`). - fn try_fill_some_slot(&mut self) -> io::Result { - for slot in self.snapshot_registry() { - // Skip slots already done producing. - if slot.queue_eof.load(std::sync::atomic::Ordering::Acquire) { - continue; - } - - // Try to acquire the slot's reader. If contended, move - // to the next slot — another worker is feeding this one. - let Ok(mut reader_guard) = slot.reader.try_lock() else { - continue; - }; - - // Compute available queue space M = CAP - decompressed.len(). - // Brief decompressed-lock for the len read; skip slot if - // queue is at cap (consumer must drain first). - let m = { - let dec = slot.decompressed.lock().expect("decompressed mutex poisoned"); - PHASE2_DECOMP_CAP.saturating_sub(dec.len()) - }; - if m == 0 { - // Queue at cap. Drop reader and try another slot. - continue; - } - let m = m.min(MAX_BATCH_PER_CALL); - - // Read + decompress up to m blocks from disk, codec-aware: BGZF - // self-framed blocks or zstd `[len][frame]` records, per `slot.codec` - // (detected from the file magic at slot-open). `read_blocks` returns - // fewer than `m` (incl. 0) at EOF. - // - // **We hold `reader_guard` through the whole read+decompress + push** - // for this batch — NOT just through the read. This is the - // "I'm the producer for this slot" lock: while held, no sibling - // worker can claim this slot, read another batch, and prematurely set - // `queue_eof` while our batch is still in flight. Releasing the reader - // before decompression would let a sibling worker see an empty read on - // a slot we'd already exhausted and set `queue_eof` before our batch - // lands — silently dropping our decompressed records. - // - // Cost: less per-slot parallelism (one worker decompresses per slot - // at a time). Cross-slot parallelism preserved — different workers - // handle different slots. - let decompressed_batch = - match self.block_dec.read_blocks(&mut reader_guard.inner, slot.codec, m) { - Ok(b) => b, - Err(e) => { - // Mark the slot as errored + EOF so the consumer exits - // cleanly with Err. Hold `reader_guard` through the flag - // block (lock order: reader → decompressed, same as the - // success paths below) and drop it AFTER. Dropping the - // reader first would let a sibling worker claim the slot, - // read a clean EOF, and set `queue_eof` WITHOUT - // `decomp_error` in the window — which the consumer (checks - // `decomp_error` then `queue_eof`) would see as a clean - // drain, silently dropping this error and the records after - // it. - { - let _g = slot.decompressed.lock().expect("decompressed mutex poisoned"); - slot.decomp_error.store(true, std::sync::atomic::Ordering::Release); - slot.queue_eof.store(true, std::sync::atomic::Ordering::Release); - } - drop(reader_guard); - return Err(e); - } - }; - let got = decompressed_batch.len(); - let hit_eof = got < m; - - if got == 0 { - // Empty read at EOF. Set queue_eof under the decompressed lock - // (still holding reader so no sibling can race). - { - let _g = slot.decompressed.lock().expect("decompressed mutex poisoned"); - slot.queue_eof.store(true, std::sync::atomic::Ordering::Release); - } - drop(reader_guard); - return Ok(true); - } - - // Push all decompressed blocks under one decompressed-lock - // acquisition. Set queue_eof under the same lock if this - // batch is the last (hit_eof). Drop reader AFTER releasing - // decompressed (lock order: reader → decompressed). - { - let mut dec = slot.decompressed.lock().expect("decompressed mutex poisoned"); - for b in decompressed_batch { - dec.push_back(b); - } - if hit_eof { - slot.queue_eof.store(true, std::sync::atomic::Ordering::Release); - } - } - drop(reader_guard); - return Ok(true); - } - Ok(false) - } -} - -impl Clone for SortSpillDecompress { - fn clone(&self) -> Self { - Self { - // Per Parallel-step semantics: each worker copy shares - // the registry so SpillReady popped by any worker is - // immediately visible to all others. - registry: Arc::clone(&self.registry), - block_dec: SpillBlockDecompressor::new(), - held: HeldSlot::new(), - output_capacity: self.output_capacity, - } - } -} - -impl Step for SortSpillDecompress { - type Input = SortPhase1Event; - type Outputs = Single; - - fn profile(&self) -> StepProfile { - StepProfile { - name: "SortSpillDecompress", - kind: StepKind::Parallel, - sticky: false, - output_queues: vec![QueueSpec::CountBounded { capacity: self.output_capacity }], - branch_ordering: vec![BranchOrdering::None], - } - } - - fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { - // 1. Drain held output first. - if !self.flush_held(ctx) { - return Ok(StepOutcome::Contention); - } - - // 2. Pop one input event, register if SpillReady, forward all. - if let Some(event) = ctx.input.pop() { - let forwarded = match event { - SortPhase1Event::SpillReady { slot, path, records_ingested_so_far } => { - self.registry.lock().push(RegisteredSpill { slot: Arc::clone(&slot) }); - SortPhase2Event::SpillReady { slot, path, records_ingested_so_far } - } - SortPhase1Event::MemoryChunk { chunk, records_ingested_so_far } => { - SortPhase2Event::MemoryChunk { chunk, records_ingested_so_far } - } - SortPhase1Event::AllAnnounced { slot_count, memory_chunk_count, total_records } => { - SortPhase2Event::AllAnnounced { slot_count, memory_chunk_count, total_records } - } - }; - let _ = self.push_or_hold(ctx, forwarded); - return Ok(StepOutcome::Progress); - } - - // 3. Greedy slot-fill step. - if self.try_fill_some_slot()? { - return Ok(StepOutcome::Progress); - } - - // 4. No fill work. Stay alive (return Contention) while ANY - // slot is still producing (i.e., !queue_eof). Otherwise, if the - // input edge is drained too, this step is done — report Finished. - // - // See `docs/design/sort-step-split-parity-fix.md` Change - // 2 — staying alive (Contention) while slots are producing is - // what prevents workers from Skipping prematurely during the - // cap-bound transient window that deadlocked v3.1. - let any_alive = self - .snapshot_registry() - .iter() - .any(|slot| !slot.queue_eof.load(std::sync::atomic::Ordering::Acquire)); - if any_alive { - return Ok(StepOutcome::Contention); - } - - // No fill work, no input this call, no slot still producing. If the - // input edge is drained too, this step will never push again — report - // Finished (only the last Parallel clone closes the shared output, - // gated by the StepDrainCounter in the driver). Otherwise wait for - // more SpillReady events. - if ctx.input.is_drained() { - return Ok(StepOutcome::Finished); - } - Ok(StepOutcome::NoProgress) - } - - fn new_worker_copy(&self) -> Self { - self.clone() - } -} - -#[cfg(test)] -mod tests; diff --git a/src/lib/pipeline/steps/sort/spill_decompress/tests.rs b/src/lib/pipeline/steps/sort/spill_decompress/tests.rs deleted file mode 100644 index 0494b5a74..000000000 --- a/src/lib/pipeline/steps/sort/spill_decompress/tests.rs +++ /dev/null @@ -1,29 +0,0 @@ -//! Tests for `SortSpillDecompress` (v4 — per-slot bounded queue design). -//! -//! TODO(P6): rewrite the test suite per spec v4. Pre-v4 tests covered -//! gap-filler admission, `try_claim_one_task` reservation, and -//! `decomp_in_flight` accounting — all removed in v4. The two -//! retained tests below pin construction + registry sharing across -//! Parallel-step worker copies. - -use super::*; - -#[test] -fn profile_advertises_parallel_no_ordering() { - let step = SortSpillDecompress::new(16); - let profile = step.profile(); - assert_eq!(profile.name, "SortSpillDecompress"); - assert_eq!(profile.kind, StepKind::Parallel); - assert!(!profile.sticky); - assert_eq!(profile.branch_ordering.len(), 1); - assert_eq!(profile.branch_ordering[0], BranchOrdering::None); -} - -#[test] -fn clone_shares_registry_across_worker_copies() { - use std::sync::Arc as StdArc; - let step = SortSpillDecompress::new(16); - let cloned = step.new_worker_copy(); - // Both clones should point at the same registry Arc. - assert!(StdArc::ptr_eq(&step.registry, &cloned.registry)); -} diff --git a/src/lib/pipeline/steps/source/read_bam.rs b/src/lib/pipeline/steps/source/read_bam.rs index 040ee3331..e503a7fdd 100644 --- a/src/lib/pipeline/steps/source/read_bam.rs +++ b/src/lib/pipeline/steps/source/read_bam.rs @@ -1,357 +1,6 @@ -//! `ReadBgzfBlocks` source step + `read_bam(path)` convenience helper. -//! -//! Reads raw BGZF blocks from a file (no decompression) and emits them as -//! `BgzfBlock` items with monotonically increasing `batch_serial`. The -//! header bytes are NOT skipped here — they pass through as part of the -//! first block(s); `FindBamBoundaries` strips them downstream. Same model -//! as the legacy `pipeline/bam.rs` block-reading thread. -//! -//! `Serial` + `Affinity::Reader`. Worker 0 is the only worker that ever -//! attempts the source's mutex; other workers `Skip` this step in -//! dispatch — no `try_lock` thrash. Mirrors legacy's "sticky read on -//! T0" pattern (`pipeline/bam.rs:3541` + `base.rs:4360-4392`). -//! Worker 0 still round-robins through every other step in the chain; -//! the priority-restart logic in the driver keeps it returning to the -//! source whenever the source has more data and downstream isn't -//! backpressured. - -use std::collections::VecDeque; -use std::fs::File; -use std::io; -use std::path::Path; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; - -use fgumi_bam_io::PipelineReaderOpts; -use fgumi_bgzf::reader::read_raw_blocks; -use noodles::sam::Header; -use parking_lot::Mutex; - -use crate::pipeline::core::Unpushed; -use crate::pipeline::core::held::HeldSlot; -use crate::pipeline::core::outputs::OrderedBytesSingle; -use crate::pipeline::core::queues::QueueSpec; -use crate::pipeline::core::reorder::BranchOrdering; -use crate::pipeline::core::step::{Affinity, Step, StepCtx, StepKind, StepOutcome, StepProfile}; -use crate::pipeline::steps::types::BgzfBlock; - -/// Legacy default blocks-per-batch (kept for compatibility with the -/// original `read_bam(..., DEFAULT_BLOCKS_PER_BATCH, ...)` callers). -/// **New code should use [`crate::pipeline::steps::tuning::BamPipelineTuning::auto_tuned`]** -/// which scales the value with thread count (16/32/48/64) to match the -/// legacy pipeline's `auto_tuned` defaults. -pub const DEFAULT_BLOCKS_PER_BATCH: usize = 16; - -/// `Exclusive + sticky` source step that reads raw BGZF blocks from a -/// file. Each `try_run` call reads up to `blocks_per_batch` blocks and -/// emits them one at a time as `BgzfBlock` items. -pub struct ReadBgzfBlocks { - /// Held inside `Mutex>` because the `Step` trait requires - /// `Clone` (which the runtime never invokes for `Exclusive` steps but - /// the bound is uniform). Clone panics; the worker that owns the - /// step holds the `&mut self` and never contends. - reader: Arc>>>, - blocks_per_batch: usize, - next_serial: u64, - /// Pending block batch when emission is mid-flight (held across - /// retries until each block is pushed). FIFO so that consecutive - /// `pop_front` calls preserve the read-order serials assigned in - /// `try_run`. - pending: VecDeque, - held: HeldSlot>, - output_byte_limit: u64, - finished: Arc, -} - -impl ReadBgzfBlocks { - #[must_use] - pub fn new( - reader: Box, - blocks_per_batch: usize, - output_byte_limit: u64, - ) -> Self { - Self { - reader: Arc::new(Mutex::new(Some(reader))), - blocks_per_batch: blocks_per_batch.max(1), - next_serial: 0, - pending: VecDeque::new(), - held: HeldSlot::new(), - output_byte_limit, - finished: Arc::new(AtomicBool::new(false)), - } - } -} - -impl Step for ReadBgzfBlocks { - type Input = (); - type Outputs = OrderedBytesSingle; - - fn profile(&self) -> StepProfile { - StepProfile { - name: "ReadBgzfBlocks", - kind: StepKind::Serial, - // Worker 0 (the Affinity::Reader target) drives the source - // sticky — fills the downstream queue in tight bursts before - // yielding to round-robin. Mirrors legacy `bam.rs:3541` + - // `base.rs:4360-4392` sticky read. - sticky: true, - output_queues: vec![QueueSpec::ByteBounded { limit_bytes: self.output_byte_limit }], - branch_ordering: vec![BranchOrdering::ByItemOrdinal], - } - } - - fn affinity(&self) -> Affinity { - // Worker 0 is the dedicated reader; other workers Skip this step. - // Matches legacy `bam.rs:3541` (`is_reader: thread_id == 0`). - Affinity::Reader - } - - fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { - // 1. Drain the held slot first. - if let Some(unpushed) = self.held.take() { - match ctx.outputs.retry(unpushed) { - Ok(()) => {} - Err(again) => { - self.held.put(again); - return Ok(StepOutcome::Contention); - } - } - } - - // 2. Drain pending blocks (one per call iteration). - if let Some(block) = self.pending.pop_front() { - match ctx.outputs.push(block) { - Ok(()) => return Ok(StepOutcome::Progress), - Err(unpushed) => { - self.held.put(unpushed); - return Ok(StepOutcome::Progress); - } - } - } - - if self.finished.load(Ordering::Acquire) { - return Ok(StepOutcome::Finished); - } - - // 3. Read up to `blocks_per_batch` raw BGZF blocks. - let raw_blocks = { - let mut guard = self.reader.lock(); - let reader = - guard.as_mut().expect("ReadBgzfBlocks: reader missing — was clone() called?"); - read_raw_blocks(reader.as_mut(), self.blocks_per_batch)? - }; - - if raw_blocks.is_empty() { - self.finished.store(true, Ordering::Release); - return Ok(StepOutcome::Finished); - } - - // Assign serials in read order; `pop_front` then yields them in - // the same order without any reversal trick. - for raw in raw_blocks { - let serial = self.next_serial; - self.next_serial += 1; - self.pending.push_back(BgzfBlock { - batch_serial: serial, - uncompressed_size: u32::try_from(raw.uncompressed_size()).unwrap_or(u32::MAX), - bytes: raw.data, - }); - } - - // Emit one block from the freshly-filled queue this iteration so - // we don't return `NoProgress` while having work in hand. - if let Some(block) = self.pending.pop_front() { - match ctx.outputs.push(block) { - Ok(()) => Ok(StepOutcome::Progress), - Err(unpushed) => { - self.held.put(unpushed); - Ok(StepOutcome::Progress) - } - } - } else { - Ok(StepOutcome::NoProgress) - } - } -} - -/// Build a [`ReadBgzfBlocks`] step from an already-prepared reader + header. -/// -/// The reader MUST be positioned at byte 0 of the BAM stream (i.e. include -/// the BAM header bytes). `FindBamBoundaries::new()` strips them downstream. -/// Used by both [`read_bam`] (file path) and [`read_bam_stdin`] (stdin via -/// `create_bam_reader_for_pipeline_with_opts`'s tee-replay reader). -#[must_use] -pub fn read_bam_from_reader( - reader: Box, - header: Header, - blocks_per_batch: usize, - output_byte_limit: u64, -) -> (ReadBgzfBlocks, Header) { - (ReadBgzfBlocks::new(reader, blocks_per_batch, output_byte_limit), header) -} - -/// Convenience helper: open a BAM file, parse its header, return the -/// `(step, header)` pair. Header parsing reads the first BGZF block(s) -/// once, then the file is reopened from byte 0 for the source's raw-block -/// reads. The header bytes pass through the chain a second time (the -/// `FindBamBoundaries` step strips them); the double-read of header -/// bytes is small overhead since headers are typically a few KB. -/// -/// **Limitation (Phase 3):** this helper requires a seekable file — -/// stdin/pipe inputs use [`read_bam_stdin`] instead, which buffers the -/// header bytes via a `TeeReader`/`ChainedReader` pair so the source's -/// raw-block reads start at byte 0 of the replayed stream. -/// A Phase 4 follow-up will share the post-header reader between header -/// parsing and raw-block reads (per the design doc's read-once contract). -/// -/// # Errors -/// -/// Returns I/O errors from file open or BAM-header parse. -pub fn read_bam>( - path: P, - opts: PipelineReaderOpts, - blocks_per_batch: usize, - output_byte_limit: u64, -) -> io::Result<(ReadBgzfBlocks, Header)> { - let path = path.as_ref(); - // Parse header via the high-level reader (consumes the first BGZF - // block(s) but the resulting reader is dropped right after). - let (_, header) = fgumi_bam_io::create_raw_bam_reader_with_opts(path, 1, opts) - .map_err(|e| io::Error::other(format!("create_raw_bam_reader_with_opts: {e}")))?; - - // Reopen from byte 0 for the source's raw-block reads. Wrap in a - // 2 MiB `BufReader` to amortize disk-I/O syscalls — every BGZF - // block read involves several short `read(2)` calls (header + extra - // fields + body); a raw `File` would syscall on each. Matches - // the legacy spill-merger's `BufReader::with_capacity(2 * 1024 * - // 1024, file)` (`crates/fgumi-sort/src/external.rs:552`). - let file = File::open(path)?; - let reader: Box = - Box::new(io::BufReader::with_capacity(2 * 1024 * 1024, file)); - Ok(read_bam_from_reader(reader, header, blocks_per_batch, output_byte_limit)) -} - -/// Stdin counterpart to [`read_bam`]. Uses -/// [`fgumi_bam_io::create_bam_reader_for_pipeline_with_opts`] to parse the -/// BAM header off stdin via a `TeeReader`, then replays the buffered header -/// bytes ahead of the remaining stdin stream so the source emits BGZF -/// blocks starting at byte 0 (same shape as the file path). -/// -/// # Errors -/// -/// Returns I/O errors from stdin read or BAM-header parse. -pub fn read_bam_stdin( - opts: PipelineReaderOpts, - blocks_per_batch: usize, - output_byte_limit: u64, -) -> io::Result<(ReadBgzfBlocks, Header)> { - let (reader, header) = - fgumi_bam_io::create_bam_reader_for_pipeline_with_opts(Path::new("-"), opts).map_err( - |e| io::Error::other(format!("create_bam_reader_for_pipeline_with_opts: {e}")), - )?; - Ok(read_bam_from_reader(reader, header, blocks_per_batch, output_byte_limit)) -} - -/// Path-aware dispatcher: routes to [`read_bam_stdin`] when `path` is a -/// stdin sentinel (`-` or `/dev/stdin`) and to [`read_bam`] otherwise. -/// -/// Used by every command's typed-step pipeline as the single entry point -/// for opening a BAM source. Replaces the legacy "reject stdin on the new -/// path" branches with a uniform dispatch. -/// -/// # Errors -/// -/// Returns I/O errors from file open, stdin read, or BAM-header parse. -pub fn read_bam_auto>( - path: P, - opts: PipelineReaderOpts, - blocks_per_batch: usize, - output_byte_limit: u64, -) -> io::Result<(ReadBgzfBlocks, Header)> { - if fgumi_bam_io::is_stdin_path(path.as_ref()) { - read_bam_stdin(opts, blocks_per_batch, output_byte_limit) - } else { - read_bam(path, opts, blocks_per_batch, output_byte_limit) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn profile_advertises_serial_reader_byordinal() { - let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); - let header = noodles::sam::Header::default(); - let writer = fgumi_bam_io::create_raw_bam_writer(&path, &header, 1, 1).unwrap(); - writer.finish().unwrap(); - - let (step, _hdr) = - read_bam(&path, PipelineReaderOpts::default(), DEFAULT_BLOCKS_PER_BATCH, 1024 * 1024) - .unwrap(); - let profile = step.profile(); - assert_eq!(profile.name, "ReadBgzfBlocks"); - assert_eq!(profile.kind, StepKind::Serial); - assert!(profile.sticky); - assert_eq!(step.affinity(), Affinity::Reader); - assert_eq!(profile.branch_ordering, vec![BranchOrdering::ByItemOrdinal]); - assert!(matches!(profile.output_queues[0], QueueSpec::ByteBounded { .. })); - } - - /// Round-trip a known BAM through [`read_bam_from_reader`] — the same - /// reader-from-`Box` shape that [`read_bam_stdin`] - /// produces. Drives the step to completion against the in-memory - /// reader and asserts that the emitted BGZF blocks reconstruct the - /// original BAM byte-for-byte (sans the BGZF EOF marker, which - /// `read_raw_blocks` filters out). - #[test] - fn read_bam_from_reader_round_trips_bytes() { - const BGZF_EOF_LEN: usize = 28; - - // Write a non-empty BAM file. `create_raw_bam_writer` emits the - // BAM header + BGZF EOF marker. - let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); - let header = noodles::sam::Header::default(); - let writer = fgumi_bam_io::create_raw_bam_writer(&path, &header, 1, 1).unwrap(); - writer.finish().unwrap(); - - let on_disk = std::fs::read(&path).unwrap(); - assert!(!on_disk.is_empty(), "BAM file should contain header + EOF block"); - - // Build the source from a `Cursor>` reader (stand-in for - // the chained-stdin reader shape) and drain the held queue. - let cursor = std::io::Cursor::new(on_disk.clone()); - let reader: Box = Box::new(cursor); - let (mut step, _hdr) = - read_bam_from_reader(reader, header, DEFAULT_BLOCKS_PER_BATCH, 1024 * 1024); - - // Drain the source via the same `read_raw_blocks` call its - // `try_run` uses. The concatenated block bytes equal everything - // in the file except the trailing 28-byte EOF block. - let mut collected = Vec::with_capacity(on_disk.len()); - let mut last_serial: Option = None; - loop { - let raw_blocks = { - let mut guard = step.reader.lock(); - let reader = guard.as_mut().expect("reader present"); - fgumi_bgzf::reader::read_raw_blocks(reader.as_mut(), DEFAULT_BLOCKS_PER_BATCH) - .unwrap() - }; - if raw_blocks.is_empty() { - break; - } - for raw in raw_blocks { - let serial = step.next_serial; - step.next_serial += 1; - if let Some(prev) = last_serial { - assert_eq!(serial, prev + 1, "serials must be monotonic"); - } - last_serial = Some(serial); - collected.extend_from_slice(&raw.data); - } - } - let expected = &on_disk[..on_disk.len() - BGZF_EOF_LEN]; - assert_eq!(collected, expected, "concatenated blocks must equal source bytes minus EOF"); - assert!(last_serial.is_some(), "should have read at least one block"); - } -} +//! Re-export shim. The `ReadBgzfBlocks` source step and `read_bam*` helpers +//! now live in the `fgumi-pipeline-io` crate. +pub use fgumi_pipeline_io::source::read_bam::{ + DEFAULT_BLOCKS_PER_BATCH, ReadBgzfBlocks, read_bam, read_bam_auto, read_bam_from_reader, + read_bam_stdin, +}; diff --git a/src/lib/pipeline/steps/types.rs b/src/lib/pipeline/steps/types.rs index 02ada8779..e61dc77a9 100644 --- a/src/lib/pipeline/steps/types.rs +++ b/src/lib/pipeline/steps/types.rs @@ -17,244 +17,20 @@ //! `Vec