From e2e9ad6cdd39f6fb7b367d5ec76f80018e81d0da Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Fri, 19 Jun 2026 20:59:08 -0700 Subject: [PATCH 1/4] refactor(cli-common): single-source shared CLI helpers via umbrella shims Make fgumi-cli-common the authoritative home for the helpers the sort extraction had duplicated (FgumiError/Result, validate_file_exists, parse_memory_size, OperationTimer, format_duration/rate/count, detect_total_memory/detect_cpu_count). The umbrella errors.rs, logging.rs, system.rs, and validation.rs become re-export shims so existing call-site paths resolve unchanged, leaving one FgumiError type and no duplicated helper bodies. Also harden detect_total_memory's 32-bit fallback and add resolve_memory_budget Auto-path test coverage (CodeRabbit). --- crates/fgumi-cli-common/src/lib.rs | 69 ++++++++++++-- src/lib/errors.rs | 120 +----------------------- src/lib/logging.rs | 143 ++--------------------------- src/lib/system.rs | 80 +--------------- src/lib/validation.rs | 134 +-------------------------- 5 files changed, 80 insertions(+), 466 deletions(-) diff --git a/crates/fgumi-cli-common/src/lib.rs b/crates/fgumi-cli-common/src/lib.rs index e4c35a0a6..cef0e4b53 100644 --- a/crates/fgumi-cli-common/src/lib.rs +++ b/crates/fgumi-cli-common/src/lib.rs @@ -5,7 +5,6 @@ // Command trait // ───────────────────────────────────────────────────────────────────────────── -use anyhow::Result; use enum_dispatch::enum_dispatch; /// Trait implemented by all fgumi CLI commands. @@ -15,7 +14,7 @@ use enum_dispatch::enum_dispatch; #[enum_dispatch] pub trait Command { #[allow(clippy::missing_errors_doc)] - fn execute(&self, command_line: &str) -> Result<()>; + fn execute(&self, command_line: &str) -> anyhow::Result<()>; } // ───────────────────────────────────────────────────────────────────────────── @@ -24,9 +23,16 @@ pub trait Command { use thiserror::Error; -/// Result type alias for fgumi operations +/// Result type alias for fgumi operations (preferred in standalone crates). pub type FgumiResult = std::result::Result; +/// Unqualified result alias used by the umbrella crate's error/validation modules. +/// +/// Both `FgumiResult` and `Result` are the same type; the two names exist +/// so callers that shadow `std::result::Result` with `use crate::errors::Result` +/// (umbrella convention) resolve to the same `FgumiError`-based alias. +pub type Result = std::result::Result; + /// Error type for fgumi operations #[derive(Error, Debug)] pub enum FgumiError { @@ -99,7 +105,11 @@ pub fn detect_total_memory() -> usize { 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) + // Saturate at usize::MAX / 2 rather than usize::MAX on 32-bit platforms so + // the downstream `budget > total` overflow check in `resolve_memory_budget` + // can fire correctly (no value can exceed usize::MAX, so using it as the + // fallback renders the check dead). + usize::try_from(bytes).unwrap_or(usize::MAX / 2) } /// Returns the number of logical CPUs available to this process. @@ -119,7 +129,8 @@ pub fn detect_cpu_count() -> usize { // ───────────────────────────────────────────────────────────────────────────── /// Format an integer with comma separators (e.g. 1234567 → "1,234,567"). -fn format_count(n: u64) -> String { +#[must_use] +pub 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; @@ -406,7 +417,7 @@ const AUTO_RESERVE_CAP: usize = 10 * 1024 * 1024 * 1024; /// # Errors /// /// Returns an error string if parsing fails. -pub fn parse_memory(s: &str) -> Result { +pub fn parse_memory(s: &str) -> std::result::Result { let s = s.trim(); if s.eq_ignore_ascii_case("auto") { return Ok(MemoryLimit::Auto); @@ -419,7 +430,7 @@ pub fn parse_memory(s: &str) -> Result { /// # Errors /// /// Returns an error string if parsing fails. -pub fn parse_memory_reserve(s: &str) -> Result { +pub fn parse_memory_reserve(s: &str) -> std::result::Result { let s = s.trim(); if s.eq_ignore_ascii_case("auto") { return Ok(MemoryReserve::Auto); @@ -517,7 +528,7 @@ fn resolve_memory_budget_with_total( } /// Parse a memory size string into `usize` bytes (private helper). -fn parse_memory_bytes(s: &str, label: &str) -> Result { +fn parse_memory_bytes(s: &str, label: &str) -> std::result::Result { let bytes = parse_memory_size(s).map_err(|e| e.to_string())?; usize::try_from(bytes).map_err(|_| format!("{label} too large: {bytes}")) } @@ -528,7 +539,7 @@ fn parse_memory_bytes(s: &str, label: &str) -> Result { /// # Errors /// /// Returns an error string if the input is not a recognized boolean. -pub fn parse_bool(s: &str) -> Result { +pub fn parse_bool(s: &str) -> std::result::Result { match s.to_ascii_lowercase().as_str() { "true" | "t" | "yes" | "y" => Ok(true), "false" | "f" | "no" | "n" => Ok(false), @@ -696,4 +707,44 @@ mod tests { assert!(CompressionHarness::try_parse_from(["prog", "--compression-level", "13"]).is_err()); assert!(CompressionHarness::try_parse_from(["prog", "--compression-level", "99"]).is_err()); } + + #[test] + fn test_resolve_memory_budget_auto_low_available() { + // When available memory per thread is below MIN_MEMORY_PER_THREAD, the + // budget is floored to MIN_MEMORY_PER_THREAD × threads, then capped at + // available (which is less), so the result equals available. + let total = 2 * MIN_MEMORY_PER_THREAD; // very tight: only 2 × floor per 4 threads + let reserve = 0; + let budget = resolve_memory_budget_with_total( + MemoryLimit::Auto, + MemoryReserve::Fixed(reserve), + 4, + true, + total, + ) + .unwrap(); + // available = total - 0 = 2×MIN; per-thread = 2×MIN/4 < MIN → floored to MIN; + // target = MIN × 4 = 4×MIN > available → capped at available. + assert_eq!(budget, total); + } + + #[test] + fn test_resolve_memory_budget_auto_margin_exceeds_total() { + // When the reserve margin >= total, saturating_sub → 0 available. + // Budget is then floored to MIN_MEMORY_PER_THREAD, capped at 0 + // (available), so result == 0 (the cap wins). + let total = 512 * 1024 * 1024_usize; // 512 MiB + let margin = total + 1; // margin exceeds total + let budget = resolve_memory_budget_with_total( + MemoryLimit::Auto, + MemoryReserve::Fixed(margin), + 1, + false, + total, + ) + .unwrap(); + // available = total.saturating_sub(margin) = 0; target = max(0, MIN) = MIN; + // budget = MIN.min(0) = 0. + assert_eq!(budget, 0); + } } diff --git a/src/lib/errors.rs b/src/lib/errors.rs index 190ef163e..1f4dcd0ea 100644 --- a/src/lib/errors.rs +++ b/src/lib/errors.rs @@ -1,116 +1,6 @@ //! Custom error types for fgumi operations. - -use thiserror::Error; - -/// Result type alias for fgumi operations -pub type Result = 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, - }, -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_invalid_parameter() { - let error = FgumiError::InvalidParameter { - parameter: "min-reads".to_string(), - reason: "must be >= 1".to_string(), - }; - let msg = format!("{error}"); - assert!(msg.contains("Invalid parameter 'min-reads'")); - assert!(msg.contains("must be >= 1")); - } - - #[test] - fn test_invalid_frequency() { - let error = FgumiError::InvalidFrequency { value: 1.5, min: 0.0, max: 1.0 }; - let msg = format!("{error}"); - assert!(msg.contains("1.5")); - assert!(msg.contains("between 0 and 1")); - } - - #[test] - fn test_invalid_file_format() { - let error = FgumiError::InvalidFileFormat { - file_type: "BAM".to_string(), - path: "/path/to/file.bam".to_string(), - reason: "truncated file".to_string(), - }; - let msg = format!("{error}"); - assert!(msg.contains("Invalid BAM file")); - assert!(msg.contains("truncated file")); - } - - #[test] - fn test_invalid_memory_size() { - let error = - FgumiError::InvalidMemorySize { reason: "Memory size cannot be empty".to_string() }; - let msg = format!("{error}"); - assert!(msg.contains("Invalid memory size")); - assert!(msg.contains("cannot be empty")); - } - - #[test] - fn test_reference_not_found() { - let error = FgumiError::ReferenceNotFound { ref_name: "chr1".to_string() }; - let msg = format!("{error}"); - assert!(msg.contains("Reference sequence 'chr1' not found")); - } -} +//! +//! The authoritative definitions live in `fgumi-cli-common`; this module +//! re-exports them so existing call-site paths (`crate::errors::FgumiError`, +//! `crate::errors::Result`) continue to resolve. +pub use fgumi_cli_common::{FgumiError, Result}; diff --git a/src/lib/logging.rs b/src/lib/logging.rs index f18b74fba..79e111c3e 100644 --- a/src/lib/logging.rs +++ b/src/lib/logging.rs @@ -2,23 +2,17 @@ //! //! This module provides consistent, user-friendly logging utilities for metrics, //! progress tracking, and operation summaries. +//! +//! `format_count`, `format_duration`, `format_rate`, and `OperationTimer` are +//! re-exported from `fgumi-cli-common` (single source of truth). `format_percent` +//! and the consensus/UMI summary helpers are umbrella-only and defined here. -use std::time::{Duration, Instant}; +pub use fgumi_cli_common::{OperationTimer, format_count, format_duration, format_rate}; use crate::metrics::{ConsensusMetrics, UmiGroupingMetrics}; -use crate::rejection::format_count; /// Formats a percentage with specified decimal places. /// -/// # Arguments -/// -/// * `value` - The fraction (0.0-1.0) to format as percentage -/// * `decimals` - Number of decimal places to include -/// -/// # Returns -/// -/// A string formatted as "XX.XX%" (e.g., "95.43%") -/// /// # Examples /// /// ``` @@ -33,88 +27,11 @@ pub fn format_percent(value: f64, decimals: usize) -> String { format!("{:.decimals$}%", value * 100.0, decimals = decimals) } -/// Formats a duration in human-readable form. -/// -/// # Arguments -/// -/// * `duration` - The duration to format -/// -/// # Returns -/// -/// A human-readable string (e.g., "2m 15s", "1h 30m", "45s") -/// -/// # Examples -/// -/// ``` -/// use fgumi_lib::logging::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: 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. -/// -/// # Arguments -/// -/// * `count` - Number of items processed -/// * `duration` - Time taken to process items -/// -/// # Returns -/// -/// A formatted rate string (e.g., "1,234 reads/s", "50 reads/min") -/// -/// # Examples -/// -/// ``` -/// use fgumi_lib::logging::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: 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") - } -} - /// Logs a formatted summary of consensus metrics. /// /// Outputs key metrics including input/output counts, rejection breakdown, /// and quality statistics. /// -/// # Arguments -/// -/// * `metrics` - The consensus metrics to summarize -/// /// # Examples /// /// ```no_run @@ -161,10 +78,6 @@ pub fn log_consensus_summary(metrics: &ConsensusMetrics) { /// /// Outputs grouping statistics including molecule counts and family sizes. /// -/// # Arguments -/// -/// * `metrics` - The UMI grouping metrics to summarize -/// /// # Examples /// /// ```no_run @@ -200,21 +113,17 @@ pub fn log_umi_grouping_summary(metrics: &UmiGroupingMetrics) { } // Log filter counts (matching fgbio's logging style) - // Non-PF only logged if > 0 if metrics.discarded_non_pf > 0 { log::info!("Filtered out {} non-PF records.", format_count(metrics.discarded_non_pf)); } - // Poor alignment always logged (like fgbio) log::info!( "Filtered out {} records due to mapping issues.", format_count(metrics.discarded_poor_alignment) ); - // Ns in UMI always logged (like fgbio) log::info!( "Filtered out {} records that contained one or more Ns in their UMIs.", format_count(metrics.discarded_ns_in_umi) ); - // UMI too short only logged if > 0 if metrics.discarded_umi_too_short > 0 { log::info!( "Filtered out {} records that contained UMIs that were too short.", @@ -223,51 +132,11 @@ pub fn log_umi_grouping_summary(metrics: &UmiGroupingMetrics) { } } -/// Operation timing and summary helper. -/// -/// Tracks operation timing and provides formatted summary output. -/// -/// # Examples -/// -/// ```no_run -/// use fgumi_lib::logging::OperationTimer; -/// -/// let timer = OperationTimer::new("Processing reads"); -/// -/// // ... do work ... -/// -/// timer.log_completion(10_000); // Log with item count -/// ``` -pub struct OperationTimer { - operation: String, - start_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: 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) - ); - } -} - #[cfg(test)] #[allow(clippy::cast_precision_loss)] mod tests { use super::*; + use std::time::Duration; #[test] fn test_format_percent() { diff --git a/src/lib/system.rs b/src/lib/system.rs index 45b821131..5f7423e88 100644 --- a/src/lib/system.rs +++ b/src/lib/system.rs @@ -1,78 +1,6 @@ //! Cgroup-aware system resource detection. //! -//! Provides [`detect_total_memory`] and [`detect_cpu_count`] which correctly -//! report resource limits when running inside Docker or Kubernetes containers. -//! Both functions take the minimum of the cgroup limit and the host value so -//! they behave correctly on bare-metal too. -//! -//! # Memory (`detect_total_memory`) -//! -//! Uses `sysinfo::System::cgroup_limits()` which reads: -//! - cgroup v2: `/sys/fs/cgroup/memory.max` -//! - cgroup v1: `/sys/fs/cgroup/memory/memory.limit_in_bytes` -//! -//! Falls back to `System::total_memory()` (physical RAM) on macOS and when -//! no cgroup limit is configured. -//! -//! # CPU (`detect_cpu_count`) -//! -//! Uses the `num_cpus` crate which reads the CFS quota: -//! - cgroup v2: `cpu.max` -//! - cgroup v1: `cpu.cfs_quota_us` / `cpu.cfs_period_us` -//! -//! Falls back to `/proc/cpuinfo` on Linux and `sysctl` on macOS. -//! Note: CFS quota is the *hard* limit set by `--cpus`; it is not the same as -//! `cpu.shares` which is a soft scheduling weight. - -/// 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. -/// -/// Respects cgroup CPU quotas (set by `--cpus` in Docker or resource limits in -/// Kubernetes), falling back to the physical core count. Returns at least 1. -#[must_use] -pub fn detect_cpu_count() -> usize { - num_cpus::get().max(1) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[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_total_memory_bounded_by_physical() { - let total = detect_total_memory(); - let mut system = sysinfo::System::new(); - system.refresh_memory(); - let physical = usize::try_from(system.total_memory()).unwrap_or(usize::MAX); - assert!(total <= physical, "cgroup-limited total {total} exceeded physical {physical}"); - } - - #[test] - fn test_detect_cpu_count_at_least_one() { - assert!(detect_cpu_count() >= 1); - } - - #[test] - fn test_detect_cpu_count_reasonable() { - // No machine has more than 65536 logical CPUs (yet). - assert!(detect_cpu_count() <= 65536); - } -} +//! The authoritative implementations live in `fgumi-cli-common`; this module +//! re-exports them so existing call-site paths (`crate::system::detect_total_memory`, +//! `crate::system::detect_cpu_count`) continue to resolve. +pub use fgumi_cli_common::{detect_cpu_count, detect_total_memory}; diff --git a/src/lib/validation.rs b/src/lib/validation.rs index 153fa05f1..2be3efd93 100644 --- a/src/lib/validation.rs +++ b/src/lib/validation.rs @@ -5,41 +5,16 @@ //! //! All validation functions now use structured error types from [`crate::errors`] to provide //! rich contextual information when validation fails. +//! +//! `validate_file_exists` and `parse_memory_size` are re-exported from `fgumi-cli-common` +//! (single source of truth); umbrella-only validators are defined here. + +pub use fgumi_cli_common::{parse_memory_size, validate_file_exists}; use crate::errors::{FgumiError, Result}; -use bytesize::ByteSize; use std::fmt::Display; use std::path::Path; -/// Validate that a file exists -/// -/// # Arguments -/// * `path` - Path to validate -/// * `description` - Human-readable description of the file (e.g., "Input file", "Reference") -/// -/// # Errors -/// Returns an error if the file does not exist -/// -/// # Example -/// ``` -/// use fgumi_lib::validation::validate_file_exists; -/// use std::path::Path; -/// -/// let result = validate_file_exists("/nonexistent/file.bam", "Input file"); -/// assert!(result.is_err()); -/// ``` -pub fn validate_file_exists>(path: P, description: &str) -> Result<()> { - 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(()) -} - /// Validate that multiple files exist /// /// # Arguments @@ -189,105 +164,6 @@ pub fn validate_positive(value: T, name: &str) -> Re 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, i.e. `n * 1024 * 1024`) -/// -/// # Examples -/// -/// ``` -/// # use fgumi_lib::validation::parse_memory_size; -/// assert_eq!(parse_memory_size("768").unwrap(), 768 * 1024 * 1024); -/// assert_eq!(parse_memory_size("2GB").unwrap(), 2 * 1000 * 1000 * 1000); -/// assert_eq!(parse_memory_size("1024MiB").unwrap(), 1024 * 1024 * 1024); -/// ``` -/// -/// # Errors -/// -/// Returns [`FgumiError::InvalidMemorySize`] if the string cannot be parsed as a valid size. -pub fn parse_memory_size(size_str: &str) -> Result { - let trimmed = size_str.trim(); - if trimmed.is_empty() { - return Err(FgumiError::InvalidMemorySize { - reason: "Memory size cannot be empty".to_string(), - }); - } - - // Handle negative values early - if trimmed.starts_with('-') { - return Err(FgumiError::InvalidMemorySize { - reason: format!("Memory size cannot be negative: '{trimmed}'"), - }); - } - - // First try parsing as a plain integer in MiB (backward compatibility) - // Only accept simple integers, not floats or scientific notation - 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 { - // Sanity guard: >1TB as a plain number likely means the user forgot a unit suffix. - 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"), - }); - } - - // Reject scientific notation (e.g. "1e3") but allow decimals in human-readable sizes (e.g. "1.5GB") - 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'." - ), - }); - } - - // Reject bare decimal numbers without a unit suffix (e.g. "1.5") since plain numbers are MiB - 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')." - ), - }); - } - - // Fall back to parsing as a human-readable size (like "2GB", "1024MiB") - 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'" - ), - }), - } -} - #[cfg(test)] mod tests { use super::*; From b4514b500ad786a31ae9fa78fc1a620619099539 Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Fri, 19 Jun 2026 20:59:10 -0700 Subject: [PATCH 2/4] fix(pipeline-core): address correctness and robustness review findings CodeRabbit findings: fix the ByteBoundedQueue byte-counter race (reserve-before-push, roll back on failure) that could underflow and wedge backpressure; bound the reorder overflow stash; defer worker-panic re-raise until monitor/rebalancer threads are joined; add topology wire bounds checks; cap the worker backoff; and correct the Step/Affinity docs. --- crates/fgumi-pipeline-core/src/builder.rs | 172 +++++++++++++++--- crates/fgumi-pipeline-core/src/erased.rs | 20 +- crates/fgumi-pipeline-core/src/handles.rs | 8 +- crates/fgumi-pipeline-core/src/queues.rs | 17 +- .../src/runtime/worker_core.rs | 2 +- crates/fgumi-pipeline-core/src/step.rs | 6 +- crates/fgumi-pipeline-core/src/topology.rs | 52 +++++- 7 files changed, 232 insertions(+), 45 deletions(-) diff --git a/crates/fgumi-pipeline-core/src/builder.rs b/crates/fgumi-pipeline-core/src/builder.rs index ab09e72b2..1c9bbc54c 100644 --- a/crates/fgumi-pipeline-core/src/builder.rs +++ b/crates/fgumi-pipeline-core/src/builder.rs @@ -903,6 +903,10 @@ impl Pipeline { (None, None) }; + // Holds the first worker panic payload; re-raised after helper threads + // are cleaned up so monitor/rebalancer shutdown always executes. + let mut worker_panic: Option> = None; + if n_threads == 1 { // Single-threaded fast path: run the worker loop directly on // the caller's thread instead of spawning + joining a fresh @@ -929,14 +933,26 @@ impl Pipeline { let sticky_owner = sticky_owners[0]; let mut worker = WorkerCore::new(0, exclusive_owner, sticky_owner); let mut entries_local = entries; - run_worker_loop( - &mut worker, - &mut entries_local, - &contexts, - &drain_counters, - &signal_arc, - stats_arc.as_ref(), - ); + // Defer a panic on the single-threaded fast path the same way the + // multi-worker join loop does: capture the payload, signal + // cancellation, and let the common monitor/rebalancer shutdown run + // before re-raising at step 7. Without this, a worker-loop panic + // unwinds straight through the caller and leaks the helper threads. + // `AssertUnwindSafe` is sound: after a panic we never touch `worker` + // or `entries_local` again — the run is shutting down. + if let Err(panic) = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + run_worker_loop( + &mut worker, + &mut entries_local, + &contexts, + &drain_counters, + &signal_arc, + stats_arc.as_ref(), + ); + })) { + signal_arc.cancel(); + worker_panic = Some(panic); + } } else { // 5. Spawn worker threads. let mut handles = Vec::with_capacity(n_threads); @@ -957,26 +973,43 @@ impl Pipeline { .spawn(move || { let mut worker = WorkerCore::new(worker_id, exclusive_owner, sticky_owner); let mut entries_local = entries; - run_worker_loop( - &mut worker, - &mut entries_local, - &contexts_clone, - &drain_counters_clone, - &signal_clone, - stats_clone.as_ref(), - ); + // Catch a worker-loop panic so we can signal cancellation + // *before* unwinding. A peer parked in its retry loop on a + // full/empty queue only exits when it observes + // `signal.is_done()`; without an early `cancel()` here, the + // join loop below could block forever on an earlier, + // now-wedged worker and never reach this thread's panic. + // We re-raise after signalling so the join still collects + // the payload (preserving the deferred re-raise at step 7). + // `AssertUnwindSafe` is sound: on panic the run is tearing + // down and neither local is used again. + if let Err(panic) = + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + run_worker_loop( + &mut worker, + &mut entries_local, + &contexts_clone, + &drain_counters_clone, + &signal_clone, + stats_clone.as_ref(), + ); + })) + { + signal_clone.cancel(); + std::panic::resume_unwind(panic); + } }) .expect("failed to spawn worker thread"); handles.push(handle); } - // 6. Join workers. If a worker panicked, re-raise its original - // payload via `resume_unwind` so the main thread aborts with the - // worker's actual panic message and location — not the opaque - // `Any { .. }` that `join().expect(...)` would print. + // 6. Join workers. Capture the first worker panic payload so cleanup + // can proceed; re-raise after monitor/rebalancer threads are stopped. for h in handles { if let Err(panic) = h.join() { - std::panic::resume_unwind(panic); + if worker_panic.is_none() { + worker_panic = Some(panic); + } } } } @@ -1014,7 +1047,13 @@ impl Pipeline { } } - // 7. Surface error or cancellation. PipelineError isn't Clone + // 7. Re-raise worker panics after helper threads are cleaned up so + // monitor/rebalancer shutdown code always executes. + if let Some(panic) = worker_panic { + std::panic::resume_unwind(panic); + } + + // 8. Surface error or cancellation. PipelineError isn't Clone // (io::Error isn't Clone); `to_result` reconstructs the recorded // outcome and, for an external cancel whose payload isn't yet visible // to this thread, synthesizes `Cancelled` from the terminal state. @@ -1804,6 +1843,95 @@ mod tests { assert_eq!(received.load(AtomicOrd::Relaxed), 50); } + /// Sink whose worker loop panics the moment it pops an item — used to drive + /// the worker-panic deferral paths in `Pipeline::run`. + #[derive(Clone)] + struct PanickingSink; + impl Step for PanickingSink { + type Input = u32; + type Outputs = (); + fn profile(&self) -> StepProfile { + StepProfile { + name: "PanickingSink", + kind: StepKind::Parallel, + sticky: false, + output_queues: vec![], + branch_ordering: vec![], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> std::io::Result { + match ctx.input.pop() { + Some(_) => panic!("intentional worker panic for test"), + None if ctx.input.is_drained() => Ok(StepOutcome::Finished), + None => Ok(StepOutcome::NoProgress), + } + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + + /// Run `f` with the panic hook silenced so an *expected* worker panic does + /// not spew a backtrace into the test log. Safe under `nextest`, which runs + /// each test in its own process. + fn with_silenced_panic_hook(f: impl FnOnce() -> R) -> R { + let prev = std::panic::take_hook(); + std::panic::set_hook(Box::new(|_| {})); + let result = f(); + std::panic::set_hook(prev); + result + } + + #[test] + fn pipeline_run_reraises_worker_panic_single_threaded() { + // A worker-loop panic on the single-threaded fast path must propagate + // out of `run` (after the common monitor/rebalancer shutdown), not be + // swallowed. The test completing at all proves the run did not hang. + let remaining = Arc::new(AtomicU32::new(10)); + let builder = PipelineBuilder::new(); + builder + .chain(SharedCountingSource { remaining: Arc::clone(&remaining) }) + .chain(PanickingSink) + .into_sink_marker(); + let pipeline = builder.build().unwrap(); + + let result = with_silenced_panic_hook(|| { + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + pipeline.run(PipelineConfig { threads: 1, ..Default::default() }) + })) + }); + assert!(result.is_err(), "single-threaded worker panic must propagate out of run()"); + } + + #[test] + fn pipeline_run_reraises_worker_panic_with_monitor_enabled() { + // With the deadlock monitor enabled (stats + non-zero timeout), a + // multi-worker panic must still re-raise — after the monitor is stopped + // and joined — rather than deadlocking the join loop or leaking the + // helper thread. The panicking worker signals cancellation so any wedged + // peer observes `is_done()` and exits, letting every join complete. + let remaining = Arc::new(AtomicU32::new(1_000)); + let builder = PipelineBuilder::new(); + builder + .chain(SharedCountingSource { remaining: Arc::clone(&remaining) }) + .chain(PanickingSink) + .into_sink_marker(); + let pipeline = builder.build().unwrap(); + let stats = pipeline.stats(); + + let result = with_silenced_panic_hook(|| { + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + pipeline.run(PipelineConfig { + threads: 4, + stats: Some(Arc::clone(&stats)), + deadlock_timeout_secs: 5, + ..Default::default() + }) + })) + }); + assert!(result.is_err(), "multi-worker panic must propagate out of run()"); + } + #[test] fn pipeline_stats_handle_matches_chain_size() { let builder = PipelineBuilder::new(); diff --git a/crates/fgumi-pipeline-core/src/erased.rs b/crates/fgumi-pipeline-core/src/erased.rs index 1f10e7b60..272d3ef2c 100644 --- a/crates/fgumi-pipeline-core/src/erased.rs +++ b/crates/fgumi-pipeline-core/src/erased.rs @@ -530,13 +530,19 @@ impl ErasedStep for TypedStep2 { p1_idx: usize, p1_branch: usize, ) -> Box { - assert_ne!( - p0_idx, - p1_idx, - "Step2 producers must be distinct steps (each upstream is its \ - own subchain). Got p0_idx == p1_idx == {p0_idx} for step '{}'.", - self.profile().name - ); + if p0_idx == p1_idx { + assert_ne!( + p0_branch, + p1_branch, + "Step2 inputs must consume distinct branches when they share \ + producer step {p0_idx} for step '{}'.", + self.profile().name + ); + let set = &mut producer_sets[p0_idx]; + let a: BranchInputHandle = set.take_typed_input::(p0_branch); + let b: BranchInputHandle = set.take_typed_input::(p1_branch); + return Box::new(TwoInputHandles::::new(a, b)); + } // Borrow two disjoint elements of `producer_sets` simultaneously. // `split_at_mut(lo+1)` puts producer_sets[lo] in the first half; // we index into the second half for the hi side. diff --git a/crates/fgumi-pipeline-core/src/handles.rs b/crates/fgumi-pipeline-core/src/handles.rs index 1bb4451f8..176c00034 100644 --- a/crates/fgumi-pipeline-core/src/handles.rs +++ b/crates/fgumi-pipeline-core/src/handles.rs @@ -300,7 +300,7 @@ pub(crate) fn build_branch( ordered_branch( transport, OrdinalSource::Allocated(Arc::new(AtomicU64::new(0))), - None, + Some(DEFAULT_REORDER_OVERFLOW_BYTES), None, ) } @@ -314,7 +314,7 @@ pub(crate) fn build_branch( ordered_branch( transport, OrdinalSource::Allocated(Arc::new(AtomicU64::new(0))), - None, + Some(DEFAULT_REORDER_OVERFLOW_BYTES), None, ) } @@ -356,7 +356,7 @@ pub(crate) fn build_branch_ordered( ordered_branch( transport, OrdinalSource::ItemSerial(|item: &T| item.ordinal()), - None, + Some(DEFAULT_REORDER_OVERFLOW_BYTES), None, ) } @@ -366,7 +366,7 @@ pub(crate) fn build_branch_ordered( ordered_branch( transport, OrdinalSource::ItemSerial(|item: &T| item.ordinal()), - None, + Some(DEFAULT_REORDER_OVERFLOW_BYTES), None, ) } diff --git a/crates/fgumi-pipeline-core/src/queues.rs b/crates/fgumi-pipeline-core/src/queues.rs index 2fb92548c..872b24ecb 100644 --- a/crates/fgumi-pipeline-core/src/queues.rs +++ b/crates/fgumi-pipeline-core/src/queues.rs @@ -285,16 +285,19 @@ impl ItemQueue for ByteBoundedQueue { return Err(item); } let size = item.heap_size() as u64; - // ArrayQueue::push returns Err((item, size)) on full; map back to - // the raw item for the producer to retry. (In practice this slot - // capacity should never be hit before the byte budget triggers a + // Reserve bytes before pushing so a concurrent consumer cannot pop and + // decrement the counter before we add our share, which would cause the + // counter to underflow and create permanent false backpressure. + self.current_bytes.fetch_add(size, Ordering::Relaxed); + // ArrayQueue::push returns Err((item, size)) on full; roll back the + // reservation and return the item to the caller for retry. (In practice + // the slot cap should never be hit before the byte budget triggers a // reject above, but defend against it anyway.) match self.inner.push((item, size)) { - Ok(()) => { - self.current_bytes.fetch_add(size, Ordering::Relaxed); - Ok(()) - } + Ok(()) => Ok(()), Err((item, _size)) => { + // Roll back the byte reservation — the item never entered the queue. + self.current_bytes.fetch_sub(size, Ordering::Relaxed); // The fixed 1024-slot backing was hit before the byte budget. // This degrades byte-backpressure into a hard count cap for // small items (heap_size ≲ limit/1024) — correctness is diff --git a/crates/fgumi-pipeline-core/src/runtime/worker_core.rs b/crates/fgumi-pipeline-core/src/runtime/worker_core.rs index 133392744..326786d9a 100644 --- a/crates/fgumi-pipeline-core/src/runtime/worker_core.rs +++ b/crates/fgumi-pipeline-core/src/runtime/worker_core.rs @@ -5,7 +5,7 @@ use std::time::Duration; use crate::topology::StepIdx; const BACKOFF_INITIAL_US: u64 = 1; -const BACKOFF_MAX_US: u64 = 1_000_000; // 1 second +const BACKOFF_MAX_US: u64 = 50_000; // 50 milliseconds pub struct WorkerCore { /// `0..n_workers` diff --git a/crates/fgumi-pipeline-core/src/step.rs b/crates/fgumi-pipeline-core/src/step.rs index 4cb32b288..38fe54c80 100644 --- a/crates/fgumi-pipeline-core/src/step.rs +++ b/crates/fgumi-pipeline-core/src/step.rs @@ -64,8 +64,8 @@ pub enum Affinity { /// benefits from thread locality. Writer, /// Restrict attempts to a specific worker index. Out-of-range values - /// (`>= n_threads`) are debug-asserted at run start and result in - /// a deadlock at runtime (no worker is eligible). + /// (`>= n_threads`) trigger an assertion failure (panic) at run start + /// via an always-on `assert!` in `build_worker_storage`. Worker(usize), } @@ -245,7 +245,7 @@ pub trait Step: Send + Sized + 'static { } /// Step body. Pop from `ctx.input`, push to `ctx.outputs`. Returns - /// `Progress` / `NoProgress` / `Contention`. Errors propagate via `Err`. + /// `Progress` / `NoProgress` / `Contention` / `Finished`. Errors propagate via `Err`. /// /// # Errors /// diff --git a/crates/fgumi-pipeline-core/src/topology.rs b/crates/fgumi-pipeline-core/src/topology.rs index 119ce6d65..8b407aeb0 100644 --- a/crates/fgumi-pipeline-core/src/topology.rs +++ b/crates/fgumi-pipeline-core/src/topology.rs @@ -82,7 +82,9 @@ impl ChainGraph { /// /// # Panics /// - /// Panics if the (producer, branch) is already wired (defensive). + /// Panics if the (producer, branch) is already wired (defensive), or if + /// `producer`, `branch`, `consumer`, or `consumer_input_slot` are out of + /// range. pub fn wire_to_slot( &mut self, producer: StepIdx, @@ -90,6 +92,19 @@ impl ChainGraph { consumer: StepIdx, consumer_input_slot: usize, ) { + assert!( + consumer.0 < self.input_arities.len(), + "consumer StepIdx({}) out of range (graph has {} steps)", + consumer.0, + self.input_arities.len() + ); + let consumer_arity = self.input_arities[consumer.0]; + assert!( + consumer_input_slot < consumer_arity, + "consumer_input_slot {consumer_input_slot} out of range for step '{}' \ + with input_arity {consumer_arity}", + self.step_names[consumer.0] + ); let slot = self.consumer_slot_index(producer, branch); assert!( self.consumers[slot].is_none(), @@ -164,6 +179,20 @@ impl ChainGraph { } fn consumer_slot_index(&self, producer: StepIdx, branch: BranchIdx) -> usize { + assert!( + producer.0 < self.branch_counts.len(), + "producer StepIdx({}) out of range (graph has {} steps)", + producer.0, + self.branch_counts.len() + ); + let branch_count = self.branch_counts[producer.0]; + assert!( + branch.0 < branch_count, + "branch BranchIdx({}) out of range for producer '{}' with {} branches", + branch.0, + self.step_names[producer.0], + branch_count + ); let mut offset = 0; for &count in &self.branch_counts[..producer.0] { offset += count; @@ -233,4 +262,25 @@ mod tests { g.wire(src, BranchIdx(0), sink1); g.wire(src, BranchIdx(0), sink2); // panics } + + #[test] + #[should_panic(expected = "out of range")] + fn wire_into_zero_arity_source_panics() { + // Sources register with `input_arity = 0` (their input is implicit), so + // even slot 0 must be rejected — there is no valid input branch to wire. + let mut g = ChainGraph::new(); + let producer = g.register_step("Producer", 1); + let source = g.register_step_with_input_arity("Source", 1, 0); + g.wire(producer, BranchIdx(0), source); // slot 0, but arity 0 → panics + } + + #[test] + #[should_panic(expected = "consumer StepIdx(5) out of range")] + fn wire_to_out_of_range_consumer_panics() { + // A consumer index past the registered steps must produce a + // deterministic range error rather than an opaque out-of-bounds panic. + let mut g = ChainGraph::new(); + let producer = g.register_step("Producer", 1); + g.wire(producer, BranchIdx(0), StepIdx(5)); // no step 5 registered + } } From 6b9addf888e7b998c76d34be3ae889484b3bc1b5 Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Fri, 19 Jun 2026 20:59:12 -0700 Subject: [PATCH 3/4] test(pipeline-core): cover sticky-owner assignment, affinity dispatch, and chain wiring Add unit coverage for assign_sticky_owners, the affinity-gated Serial dispatch eligibility, and build_chain_contexts wiring (CodeRabbit). --- Cargo.lock | 2 + crates/fgumi-pipeline-core/Cargo.toml | 2 + .../src/runtime/contexts.rs | 102 +++++++++++++-- .../fgumi-pipeline-core/src/runtime/pool.rs | 123 +++++++++++++++++- .../src/runtime/storage.rs | 58 ++++++++- 5 files changed, 276 insertions(+), 11 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 536cdc14a..022ef794d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -819,6 +819,8 @@ dependencies = [ "log", "noodles", "parking_lot", + "proptest", + "rstest", "trybuild", ] diff --git a/crates/fgumi-pipeline-core/Cargo.toml b/crates/fgumi-pipeline-core/Cargo.toml index 762671575..e66c815c4 100644 --- a/crates/fgumi-pipeline-core/Cargo.toml +++ b/crates/fgumi-pipeline-core/Cargo.toml @@ -15,6 +15,8 @@ log = "0" noodles = { version = "0.111.0", features = ["sam"] } [dev-dependencies] +proptest = "1.10" +rstest = "0" trybuild = "1.0" [lints.clippy] diff --git a/crates/fgumi-pipeline-core/src/runtime/contexts.rs b/crates/fgumi-pipeline-core/src/runtime/contexts.rs index a6b7d4032..14a9d2efb 100644 --- a/crates/fgumi-pipeline-core/src/runtime/contexts.rs +++ b/crates/fgumi-pipeline-core/src/runtime/contexts.rs @@ -289,6 +289,9 @@ mod tests { use super::*; use std::io; + use proptest::prelude::*; + use rstest::rstest; + use crate::erased::TypedStep; use crate::outputs::Single; use crate::step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}; @@ -331,18 +334,99 @@ mod tests { } } - #[test] - fn build_chain_contexts_for_two_step_chain() { + /// A `u32 → u32` pass-through step used to grow a linear chain to an + /// arbitrary length between the source and sink. + #[derive(Clone)] + struct MiddleStep; + impl Step for MiddleStep { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "Middle", + kind: StepKind::Serial, + sticky: false, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + /// Build an `n`-step linear chain `Source → Middle×(n - 2) → Sink`, + /// returning the erased step boxes alongside the fully wired graph. + /// + /// # Panics + /// + /// Panics if `n < 2` (a linear chain needs at least a source and a sink). + fn linear_chain(n: usize) -> (Vec>, ChainGraph) { + assert!(n >= 2, "linear chain needs at least a source and a sink"); + let mut graph = ChainGraph::new(); - let src = graph.register_step("Source", 1); - let sink = graph.register_step("Sink", 0); - graph.wire(src, BranchIdx(0), sink); + let mut steps: Vec> = Vec::with_capacity(n); + let mut indices: Vec = Vec::with_capacity(n); + + indices.push(graph.register_step("Source", 1)); + steps.push(Box::new(TypedStep::new(StubSource))); + for _ in 1..n - 1 { + indices.push(graph.register_step("Middle", 1)); + steps.push(Box::new(TypedStep::new(MiddleStep))); + } + indices.push(graph.register_step("Sink", 0)); + steps.push(Box::new(TypedStep::new(StubSink))); + + for pair in indices.windows(2) { + graph.wire(pair[0], BranchIdx(0), pair[1]); + } + + (steps, graph) + } + + /// Assert the invariants `build_chain_contexts` must uphold for an + /// `n`-step linear chain: one input/output slot per step, no byte-bounded + /// queues (the stubs only use `CountBounded` transport), the source's + /// input is a dummy pre-drained `BranchInputHandle<()>`, and every + /// downstream step receives a real `BranchInputHandle` wired from its + /// producer. + fn assert_linear_chain_invariants(ctx: &ChainContexts, n: usize) { + assert_eq!(ctx.inputs.len(), n); + assert_eq!(ctx.outputs.len(), n); + assert!(ctx.bounded_queues.is_empty()); - let steps: Vec> = - vec![Box::new(TypedStep::new(StubSource)), Box::new(TypedStep::new(StubSink))]; + assert!( + ctx.inputs[0].downcast_ref::>().is_some(), + "source must have a dummy unit input handle" + ); + for i in 1..n { + assert!( + ctx.inputs[i].downcast_ref::>().is_some(), + "downstream step {i} must have a BranchInputHandle wired from its producer" + ); + } + } + /// `build_chain_contexts` wires linear chains of varying length: the + /// three-step case also exercises `find_all_producers` for a middle step. + #[rstest] + #[case(2)] + #[case(3)] + #[case(4)] + fn build_chain_contexts_linear(#[case] n: usize) { + let (steps, graph) = linear_chain(n); let ctx = build_chain_contexts(&steps, &graph); - assert_eq!(ctx.inputs.len(), 2); - assert_eq!(ctx.outputs.len(), 2); + assert_linear_chain_invariants(&ctx, n); + } + + proptest! { + /// The typed-handle and bounded-queue invariants hold for linear + /// chains of any length, not just the hand-picked rstest cases. + #[test] + fn build_chain_contexts_linear_invariants(n in 2usize..=8) { + let (steps, graph) = linear_chain(n); + let ctx = build_chain_contexts(&steps, &graph); + assert_linear_chain_invariants(&ctx, n); + } } } diff --git a/crates/fgumi-pipeline-core/src/runtime/pool.rs b/crates/fgumi-pipeline-core/src/runtime/pool.rs index 27f13593d..e584bea60 100644 --- a/crates/fgumi-pipeline-core/src/runtime/pool.rs +++ b/crates/fgumi-pipeline-core/src/runtime/pool.rs @@ -100,11 +100,13 @@ mod tests { use super::*; use std::io; + use rstest::rstest; + use crate::erased::TypedStep; use crate::outputs::Single; use crate::queues::QueueSpec; use crate::reorder::BranchOrdering; - use crate::step::{Step, StepCtx, StepOutcome, StepProfile}; + use crate::step::{Affinity, Step, StepCtx, StepOutcome, StepProfile}; fn stub_step(kind: StepKind) -> Box { #[derive(Clone)] @@ -169,4 +171,123 @@ mod tests { Err(PipelineError::NotEnoughThreads { required: 3, available: 2 }) )); } + + fn sticky_exclusive_step(owner_idx: usize) -> Box { + #[derive(Clone)] + struct StickyExclusive; + impl Step for StickyExclusive { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "StickyExclusive", + kind: StepKind::Exclusive, + sticky: true, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + fn new_worker_copy(&self) -> Self { + self.clone() + } + } + let _ = owner_idx; // owner_idx used by caller, not embedded in step + Box::new(TypedStep::new(StickyExclusive)) + } + + fn sticky_serial_step(affinity: crate::step::Affinity) -> Box { + struct StickySerial(crate::step::Affinity); + impl Step for StickySerial { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + StepProfile { + name: "StickySerial", + kind: StepKind::Serial, + sticky: true, + output_queues: vec![QueueSpec::CountBounded { capacity: 4 }], + branch_ordering: vec![BranchOrdering::None], + } + } + fn affinity(&self) -> crate::step::Affinity { + self.0 + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + Box::new(TypedStep::new(StickySerial(affinity))) + } + + /// Assert that `sticky` holds `Some(StepIdx(0))` in exactly `expected_slot` + /// (when `Some`) and `None` everywhere else across `n_workers` slots. + fn assert_only_slot( + sticky: &[Option], + n_workers: usize, + expected_slot: Option, + ) { + for (slot, &got) in sticky.iter().enumerate().take(n_workers) { + let want = if Some(slot) == expected_slot { Some(StepIdx(0)) } else { None }; + assert_eq!(got, want, "slot {slot}; expected owner slot {expected_slot:?}"); + } + } + + #[rstest] + #[case::in_range(2, Some(2))] + #[case::out_of_range(10, None)] + fn sticky_exclusive_owner_maps_to_slot( + #[case] owner: usize, + #[case] expected_slot: Option, + ) { + // A sticky-exclusive step's slot is its (in-range) owner worker; an + // out-of-range owner is skipped, leaving every slot empty. + let steps = vec![sticky_exclusive_step(0), stub_step(StepKind::Parallel)]; + let exclusive_owners = vec![Some(owner), None]; + let sticky = assign_sticky_owners(&steps, &exclusive_owners, 4); + assert_only_slot(&sticky, 4, expected_slot); + } + + #[test] + fn sticky_exclusive_occupied_slot_not_overwritten() { + // Two sticky-exclusive steps competing for the same worker slot. + let steps = vec![sticky_exclusive_step(0), sticky_exclusive_step(0)]; + let exclusive_owners = vec![Some(0_usize), Some(0_usize)]; + let sticky = assign_sticky_owners(&steps, &exclusive_owners, 4); + // First step wins; second is skipped because slot[0] is already occupied. + assert_eq!(sticky[0], Some(StepIdx(0))); + } + + #[rstest] + #[case::reader(Affinity::Reader, Some(0))] + #[case::writer(Affinity::Writer, Some(3))] + #[case::worker_in_range(Affinity::Worker(2), Some(2))] + #[case::worker_out_of_range(Affinity::Worker(10), None)] + #[case::none(Affinity::None, None)] + fn sticky_serial_affinity_maps_to_slot( + #[case] affinity: Affinity, + #[case] expected_slot: Option, + ) { + // Serial affinity resolves to a single worker slot: Reader→0, + // Writer→last, Worker(i)→i; an out-of-range Worker index and None are + // skipped, leaving every slot empty. + let steps = vec![sticky_serial_step(affinity)]; + let exclusive_owners = vec![None]; + let sticky = assign_sticky_owners(&steps, &exclusive_owners, 4); + assert_only_slot(&sticky, 4, expected_slot); + } + + #[test] + fn sticky_exclusive_beats_sticky_serial_on_same_slot() { + // Exclusive pass runs first; serial pass only fills empty slots. + let exc = sticky_exclusive_step(0); + let ser = sticky_serial_step(crate::step::Affinity::Reader); // also targets slot 0 + let steps: Vec> = vec![exc, ser]; + let exclusive_owners = vec![Some(0_usize), None]; + let sticky = assign_sticky_owners(&steps, &exclusive_owners, 4); + // Exclusive (step 0) wins slot 0; Serial (step 1) is blocked. + assert_eq!(sticky[0], Some(StepIdx(0))); + } } diff --git a/crates/fgumi-pipeline-core/src/runtime/storage.rs b/crates/fgumi-pipeline-core/src/runtime/storage.rs index 673bed0d9..5b575c62f 100644 --- a/crates/fgumi-pipeline-core/src/runtime/storage.rs +++ b/crates/fgumi-pipeline-core/src/runtime/storage.rs @@ -167,11 +167,13 @@ mod tests { use super::*; use std::io; + use rstest::rstest; + use crate::erased::TypedStep; use crate::outputs::Single; use crate::queues::QueueSpec; use crate::reorder::BranchOrdering; - use crate::step::{Step, StepCtx, StepOutcome, StepProfile}; + use crate::step::{Affinity, Step, StepCtx, StepOutcome, StepProfile}; fn profile_for(name: &'static str, kind: StepKind, sticky: bool) -> StepProfile { StepProfile { @@ -274,6 +276,60 @@ mod tests { assert!(matches!(entries[3][0], WorkerStepEntry::Skip)); } + struct AffinitySerialStep(crate::step::Affinity); + impl Step for AffinitySerialStep { + type Input = u32; + type Outputs = Single; + fn profile(&self) -> StepProfile { + profile_for("AffinitySer", StepKind::Serial, false) + } + fn affinity(&self) -> crate::step::Affinity { + self.0 + } + fn try_run(&mut self, _ctx: &mut StepCtx<'_, Self>) -> io::Result { + Ok(StepOutcome::NoProgress) + } + } + + #[rstest] + #[case::reader(Affinity::Reader, 0)] + #[case::writer(Affinity::Writer, 2)] // last of 3 workers + #[case::worker_idx(Affinity::Worker(1), 1)] + fn serial_affinity_eligible_only_for_target_worker( + #[case] affinity: Affinity, + #[case] eligible: usize, + ) { + // A Serial step's affinity makes exactly one worker eligible (`Shared`); + // every other worker `Skip`s it. Reader→0, Writer→last, Worker(i)→i. + let steps: Vec> = + vec![Box::new(TypedStep::new(AffinitySerialStep(affinity)))]; + let owners = vec![None]; + let entries = build_worker_storage(steps, &owners, 3); + for (worker, entry) in entries.iter().enumerate() { + if worker == eligible { + assert!( + matches!(entry[0], WorkerStepEntry::Shared { .. }), + "worker {worker} should be eligible (Shared) for {affinity:?}" + ); + } else { + assert!( + matches!(entry[0], WorkerStepEntry::Skip), + "worker {worker} should Skip for {affinity:?}" + ); + } + } + } + + #[test] + #[should_panic(expected = "Serial step affinity")] + fn serial_out_of_range_worker_panics_in_storage() { + let steps: Vec> = + vec![Box::new(TypedStep::new(AffinitySerialStep(Affinity::Worker(99))))]; + let owners = vec![None]; + // build_worker_storage checks affinity in range via assert!; should panic. + let _ = build_worker_storage(steps, &owners, 3); + } + #[test] fn mixed_chain_assigns_correctly() { let steps: Vec> = vec![ From c8c116f63f5cfd7635a8bb274917ced44d6f7177 Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Fri, 19 Jun 2026 20:59:14 -0700 Subject: [PATCH 4/4] fix(pipeline-io): guard sort memory-chunk count against u32 overflow Use checked_add for memory_chunk_count so an overflow panics rather than silently wrapping and corrupting the AllAnnounced count (CodeRabbit). --- crates/fgumi-pipeline-io/src/sort/and_spill.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/crates/fgumi-pipeline-io/src/sort/and_spill.rs b/crates/fgumi-pipeline-io/src/sort/and_spill.rs index bf2cc38d5..1886f4cfe 100644 --- a/crates/fgumi-pipeline-io/src/sort/and_spill.rs +++ b/crates/fgumi-pipeline-io/src/sort/and_spill.rs @@ -167,7 +167,8 @@ impl SortAndSpill { /// # Panics /// - /// Panics if the number of spill slots exceeds `u32::MAX`. + /// Panics if the number of spill slots exceeds `u32::MAX`, or if the number + /// of in-memory chunks exceeds `u32::MAX`. fn finalize_into_pending( stream: SortStream, caller_label: &str, @@ -190,7 +191,8 @@ impl SortAndSpill { if chunk.is_empty() { continue; } - memory_chunk_count += 1; + memory_chunk_count = + memory_chunk_count.checked_add(1).expect("memory chunk count fits in u32"); pending_events.push_back(SortPhase1Event::MemoryChunk { chunk: Arc::new(chunk), records_ingested_so_far: total_records,