diff --git a/CHANGELOG.md b/CHANGELOG.md index ef92f16c6..23d41b9f3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ All notable changes to this project will be documented in this file. ### Bug Fixes - Cap `fgumi group` metric memory by merging per-position-group counts into per-thread accumulators instead of a `SegQueue` that grew one `AHashMap` per position group ([#285](https://github.com/fulcrumgenomics/fgumi/issues/285)). +- Replace unbounded `SegQueue` buffering in `filter`, `clip`, `correct`, `simplex`, `duplex`, and `codec` with a shared `PerThreadAccumulator` helper, capping retained metric memory at `O(threads × distinct keys)` across all pipeline commands. Rejects buffering in the consensus commands is unchanged and tracked separately. ### Breaking Changes diff --git a/src/lib/commands/clip.rs b/src/lib/commands/clip.rs index 5a248a59d..d2953597a 100644 --- a/src/lib/commands/clip.rs +++ b/src/lib/commands/clip.rs @@ -14,6 +14,7 @@ use crate::grouper::TemplateGrouper; use crate::logging::OperationTimer; use crate::metrics::clip::{ClipCounts, ClippingMetricsCollection}; use crate::metrics::writer::write_metrics as write_metrics_tsv; +use crate::per_thread_accumulator::PerThreadAccumulator; use crate::progress::ProgressTracker; use crate::reference::ReferenceReader; use crate::template::{TemplateBatch, TemplateIterator}; @@ -23,7 +24,6 @@ use crate::unified_pipeline::{ use crate::validation::validate_file_exists; use anyhow::Result; use clap::Parser; -use crossbeam_queue::SegQueue; use log::info; use noodles::sam::Header; use noodles::sam::alignment::io::Write as AlignmentWrite; @@ -641,8 +641,8 @@ impl Clip { num_threads, )?; - // Lock-free metrics collection - let collected_metrics: Arc> = Arc::new(SegQueue::new()); + // Per-thread metrics accumulator: bounded memory, no unbounded queue. + let collected_metrics = PerThreadAccumulator::::new(num_threads); let collected_for_serialize = Arc::clone(&collected_metrics); // Configuration for closures @@ -789,11 +789,11 @@ impl Clip { header: &Header, output: &mut Vec| -> io::Result { - // Push metrics to lock-free queue - collected_for_serialize.push(CollectedClipMetrics { - total_templates: processed.templates_count, - overlap_clipped: processed.overlap_clipped_count, - extend_clipped: processed.extend_clipped_count, + // Merge per-batch counts into this worker's accumulator slot + collected_for_serialize.with_slot(|m| { + m.total_templates += processed.templates_count; + m.overlap_clipped += processed.overlap_clipped_count; + m.extend_clipped += processed.extend_clipped_count; }); // Serialize clipped records into the provided buffer @@ -817,10 +817,11 @@ impl Clip { let mut total_overlap_clipped = 0u64; let mut total_extend_clipped = 0u64; - while let Some(metrics) = collected_metrics.pop() { - total_templates += metrics.total_templates; - total_overlap_clipped += metrics.overlap_clipped; - total_extend_clipped += metrics.extend_clipped; + for slot in collected_metrics.slots() { + let m = slot.lock(); + total_templates += m.total_templates; + total_overlap_clipped += m.overlap_clipped; + total_extend_clipped += m.extend_clipped; } info!("Total templates processed: {total_templates}"); diff --git a/src/lib/commands/codec.rs b/src/lib/commands/codec.rs index c04e56941..11cd46b2c 100644 --- a/src/lib/commands/codec.rs +++ b/src/lib/commands/codec.rs @@ -12,9 +12,9 @@ use crate::bam_io::{ }; use crate::commands::command::Command; use crate::commands::consensus_runner::{ConsensusStatsOps, create_unmapped_consensus_header}; +use crate::per_thread_accumulator::PerThreadAccumulator; use anyhow::{Context, Result, bail}; use clap::Parser; -use crossbeam_queue::SegQueue; use fgoxide::io::DelimFile; use super::common::{ @@ -502,8 +502,9 @@ impl Codec { num_threads, )?; - // Lock-free metrics collection - let collected_metrics: Arc> = Arc::new(SegQueue::new()); + // Per-thread metrics accumulator: bounded metric memory, no unbounded + // queue. Rejects buffering semantics are preserved (see follow-up). + let collected_metrics = PerThreadAccumulator::::new(num_threads); let collected_metrics_for_serialize = Arc::clone(&collected_metrics); // Parse cell tag @@ -604,11 +605,14 @@ impl Codec { _header: &Header, output: &mut Vec| -> io::Result { - // Collect metrics (lock-free) - collected_metrics_for_serialize.push(CollectedCodecMetrics { - stats: processed.stats, - groups_processed: processed.groups_count, - rejects: processed.rejects, + // Merge per-batch metrics + rejects into this worker's slot + let batch_stats = processed.stats; + let groups_count = processed.groups_count; + let mut batch_rejects = processed.rejects; + collected_metrics_for_serialize.with_slot(|m| { + m.stats.merge(&batch_stats); + m.groups_processed += groups_count; + m.rejects.append(&mut batch_rejects); }); // Serialize consensus reads @@ -624,11 +628,12 @@ impl Codec { let mut merged_stats = CodecConsensusStats::default(); let mut all_rejects: Vec> = Vec::new(); - while let Some(metrics) = collected_metrics.pop() { - total_groups += metrics.groups_processed; - merged_stats.merge(&metrics.stats); + for slot in collected_metrics.slots() { + let mut m = slot.lock(); + total_groups += m.groups_processed; + merged_stats.merge(&m.stats); if track_rejects { - all_rejects.extend(metrics.rejects); + all_rejects.append(&mut m.rejects); } } diff --git a/src/lib/commands/correct.rs b/src/lib/commands/correct.rs index 7b2278674..a83334cd8 100644 --- a/src/lib/commands/correct.rs +++ b/src/lib/commands/correct.rs @@ -48,6 +48,7 @@ use crate::dna::reverse_complement_str; use crate::grouper::TemplateGrouper; use crate::logging::OperationTimer; use crate::metrics::correct::UmiCorrectionMetrics; +use crate::per_thread_accumulator::PerThreadAccumulator; use crate::progress::ProgressTracker; use crate::sam::SamTag; use crate::sort::bam_fields; @@ -57,7 +58,6 @@ use crate::validation::validate_file_exists; use ahash::AHashMap; use anyhow::{Result, bail}; use clap::Parser; -use crossbeam_queue::SegQueue; use fgumi_raw_bam::RawRecord; use log::{info, warn}; use lru::LruCache; @@ -788,8 +788,8 @@ impl CorrectUmis { pipeline_config.group_key_config = Some(GroupKeyConfig::new_raw_no_cell(library_index)); } - // Lock-free metrics collection - let collected_metrics: Arc> = Arc::new(SegQueue::new()); + // Per-thread metrics accumulator: bounded memory, no unbounded queue. + let collected_metrics = PerThreadAccumulator::::new(num_threads); let collected_for_serialize = Arc::clone(&collected_metrics); // Configuration for closures @@ -959,18 +959,23 @@ impl CorrectUmis { }; // Serialize function: convert records to bytes and collect metrics - let serialize_fn = move |processed: CorrectProcessedBatch, + let serialize_fn = move |mut processed: CorrectProcessedBatch, _header: &Header, output: &mut Vec| -> io::Result { - // Push metrics to lock-free queue - collected_for_serialize.push(CollectedCorrectMetrics { - templates_processed: processed.templates_count, - missing_umis: processed.missing_umis, - wrong_length: processed.wrong_length, - mismatched: processed.mismatched, - umi_matches: processed.umi_matches, - raw_rejects: processed.rejected_raw_records, + // Merge per-batch counts + UMI matches + rejects into this + // worker's accumulator slot. + let umi_matches = std::mem::take(&mut processed.umi_matches); + let mut rejects = std::mem::take(&mut processed.rejected_raw_records); + collected_for_serialize.with_slot(|m| { + m.templates_processed += processed.templates_count; + m.missing_umis += processed.missing_umis; + m.wrong_length += processed.wrong_length; + m.mismatched += processed.mismatched; + for (umi, counts) in umi_matches { + merge_umi_counts(&mut m.umi_matches, umi, &counts); + } + m.raw_rejects.append(&mut rejects); }); // Serialize kept records @@ -1008,25 +1013,19 @@ impl CorrectUmis { let mut merged_umi_matches: AHashMap = AHashMap::new(); let mut all_raw_rejects: Vec> = Vec::new(); - while let Some(metrics) = collected_metrics.pop() { - total_templates += metrics.templates_processed; - total_missing += metrics.missing_umis; - total_wrong_length += metrics.wrong_length; - total_mismatched += metrics.mismatched; - - for (umi, counts) in metrics.umi_matches { - let entry = merged_umi_matches - .entry(umi.clone()) - .or_insert_with(|| UmiCorrectionMetrics::new(umi)); - entry.total_matches += counts.total_matches; - entry.perfect_matches += counts.perfect_matches; - entry.one_mismatch_matches += counts.one_mismatch_matches; - entry.two_mismatch_matches += counts.two_mismatch_matches; - entry.other_matches += counts.other_matches; + for slot in collected_metrics.slots() { + let mut m = slot.lock(); + total_templates += m.templates_processed; + total_missing += m.missing_umis; + total_wrong_length += m.wrong_length; + total_mismatched += m.mismatched; + + for (umi, counts) in m.umi_matches.drain() { + merge_umi_counts(&mut merged_umi_matches, umi, &counts); } if track_rejects { - all_raw_rejects.extend(metrics.raw_rejects); + all_raw_rejects.append(&mut m.raw_rejects); } } @@ -1362,6 +1361,22 @@ impl CorrectUmis { /// assert_eq!(count_mismatches_with_max(b"AAAAAA", b"AAAAAT", 10), 1); /// assert_eq!(count_mismatches_with_max(b"AAAAAA", b"CCCCCC", 2), 3); /// ``` +/// Merges `counts` into `dst[umi]`, creating a zero-initialized entry if the +/// key is absent. Uses `or_insert_with_key` so `umi` is only cloned on insert, +/// not on hit. +fn merge_umi_counts( + dst: &mut AHashMap, + umi: String, + counts: &UmiCorrectionMetrics, +) { + let entry = dst.entry(umi).or_insert_with_key(|k| UmiCorrectionMetrics::new(k.clone())); + entry.total_matches += counts.total_matches; + entry.perfect_matches += counts.perfect_matches; + entry.one_mismatch_matches += counts.one_mismatch_matches; + entry.two_mismatch_matches += counts.two_mismatch_matches; + entry.other_matches += counts.other_matches; +} + #[must_use] pub fn count_mismatches_with_max(a: &[u8], b: &[u8], max_mismatches: usize) -> usize { let mut mismatches = 0; diff --git a/src/lib/commands/duplex.rs b/src/lib/commands/duplex.rs index ebcc51233..8e5c75ee2 100644 --- a/src/lib/commands/duplex.rs +++ b/src/lib/commands/duplex.rs @@ -29,6 +29,7 @@ use crate::overlapping_consensus::{ AgreementStrategy, CorrectionStats, DisagreementStrategy, OverlappingBasesConsensusCaller, apply_overlapping_consensus_raw, }; +use crate::per_thread_accumulator::PerThreadAccumulator; use crate::progress::ProgressTracker; use crate::read_info::LibraryIndex; use crate::sam::SamTag; @@ -38,7 +39,6 @@ use crate::unified_pipeline::{ GroupKeyConfig, Grouper, MemoryEstimate, run_bam_pipeline_from_reader, }; use crate::validation::validate_file_exists; -use crossbeam_queue::SegQueue; use fgumi_raw_bam::{RawRecord, RawRecordView}; use log::info; use noodles::sam::Header; @@ -574,8 +574,9 @@ impl Duplex { num_threads, )?; - // Lock-free metrics collection - let collected_metrics: Arc> = Arc::new(SegQueue::new()); + // Per-thread metrics accumulator: bounded metric memory, no unbounded + // queue. Rejects buffering semantics are preserved (see follow-up). + let collected_metrics = PerThreadAccumulator::::new(num_threads); let collected_metrics_for_serialize = Arc::clone(&collected_metrics); // Capture configuration for closures @@ -731,12 +732,18 @@ impl Duplex { _header: &Header, output: &mut Vec| -> io::Result { - // Collect metrics (lock-free) - collected_metrics_for_serialize.push(CollectedDuplexMetrics { - stats: processed.stats, - overlapping_stats: processed.overlapping_stats, - groups_processed: processed.groups_count, - rejects: processed.rejects, + // Merge per-batch metrics + rejects into this worker's slot + let batch_stats = processed.stats; + let batch_overlapping = processed.overlapping_stats; + let groups_count = processed.groups_count; + let mut batch_rejects = processed.rejects; + collected_metrics_for_serialize.with_slot(|m| { + m.stats.merge(&batch_stats); + if let Some(o) = batch_overlapping { + m.overlapping_stats.get_or_insert_with(CorrectionStats::new).merge(&o); + } + m.groups_processed += groups_count; + m.rejects.append(&mut batch_rejects); }); let count = processed.consensus_output.count as u64; @@ -752,14 +759,15 @@ impl Duplex { let mut merged_overlapping_stats = CorrectionStats::new(); let mut all_rejects: Vec> = Vec::new(); - while let Some(metrics) = collected_metrics.pop() { - total_groups += metrics.groups_processed; - merged_stats.merge(&metrics.stats); - if let Some(ref ocs) = metrics.overlapping_stats { + for slot in collected_metrics.slots() { + let mut m = slot.lock(); + total_groups += m.groups_processed; + merged_stats.merge(&m.stats); + if let Some(ref ocs) = m.overlapping_stats { merged_overlapping_stats.merge(ocs); } if track_rejects { - all_rejects.extend(metrics.rejects); + all_rejects.append(&mut m.rejects); } } diff --git a/src/lib/commands/filter.rs b/src/lib/commands/filter.rs index 9f83991e6..4160d7383 100644 --- a/src/lib/commands/filter.rs +++ b/src/lib/commands/filter.rs @@ -21,6 +21,7 @@ use crate::consensus_filter::{ }; use crate::grouper::{SingleRawRecordGrouper, TemplateGrouper}; use crate::logging::OperationTimer; +use crate::per_thread_accumulator::PerThreadAccumulator; use crate::read_info::LibraryIndex; use crate::reference::ReferenceReader; use crate::sort::bam_fields; @@ -34,7 +35,6 @@ use crate::validation::validate_file_exists; use ahash::AHashMap; use anyhow::{Result, bail}; use clap::Parser; -use crossbeam_queue::SegQueue; use fgumi_raw_bam::RawRecordView; use log::info; use noodles::sam::Header; @@ -260,7 +260,7 @@ fn serialize_raw_records(records: &[Vec], output: &mut Vec) -> io::Resul Ok(records.len() as u64) } -/// Metrics collected from filter processing, aggregated post-pipeline. +/// Per-thread accumulator merged into final counts after pipeline completion. #[derive(Default)] struct CollectedFilterMetrics { /// Total records processed. @@ -278,7 +278,7 @@ struct FilterPipelineSetup { pipeline_config: BamPipelineConfig, config: Arc, reference: Option>, - collected_metrics: Arc>, + collected_metrics: Arc>, progress_counter: Arc, } @@ -406,7 +406,7 @@ impl Filter { None => None, }; - let collected_metrics: Arc> = Arc::new(SegQueue::new()); + let collected_metrics = PerThreadAccumulator::::new(num_threads); let progress_counter = Arc::new(AtomicU64::new(0)); Ok(FilterPipelineSetup { @@ -475,11 +475,11 @@ impl Filter { _header: &Header, output: &mut Vec| -> io::Result { - collected_for_serialize.push(CollectedFilterMetrics { - total_records: processed.records_count, - passed_records: processed.passed_count, - failed_records: processed.records_count - processed.passed_count, - total_bases_masked: processed.bases_masked, + collected_for_serialize.with_slot(|m| { + m.total_records += processed.records_count; + m.passed_records += processed.passed_count; + m.failed_records += processed.records_count - processed.passed_count; + m.total_bases_masked += processed.bases_masked; }); serialize_raw_records(&processed.kept_records, output) @@ -523,11 +523,12 @@ impl Filter { let mut failed_reads = 0u64; let mut total_bases_masked = 0u64; - while let Some(metrics) = setup.collected_metrics.pop() { - total_reads += metrics.total_records; - passed_reads += metrics.passed_records; - failed_reads += metrics.failed_records; - total_bases_masked += metrics.total_bases_masked; + for slot in setup.collected_metrics.slots() { + let m = slot.lock(); + total_reads += m.total_records; + passed_reads += m.passed_records; + failed_reads += m.failed_records; + total_bases_masked += m.total_bases_masked; } if let Some(stats_path) = &self.stats { diff --git a/src/lib/commands/group.rs b/src/lib/commands/group.rs index d795d4ca3..a71474a31 100644 --- a/src/lib/commands/group.rs +++ b/src/lib/commands/group.rs @@ -13,6 +13,7 @@ use crate::grouper::{ }; use crate::logging::{OperationTimer, log_umi_grouping_summary}; use crate::metrics::group::{FamilySizeMetrics, PositionGroupSizeMetrics, UmiGroupingMetrics}; +use crate::per_thread_accumulator::PerThreadAccumulator; use crate::progress::ProgressTracker; use crate::read_info::{LibraryIndex, compute_group_key}; use crate::sam::{is_sorted, is_template_coordinate_sorted, unclipped_five_prime_position}; @@ -29,7 +30,6 @@ use anyhow::{Context, Result, bail}; use bstr::BString; use clap::Parser; use fgoxide::io::DelimFile; -use parking_lot::Mutex; // MemoryEstimate is gated because it's only used in memory-debug blocks below use crate::sam::SamTag; #[cfg(feature = "memory-debug")] @@ -88,26 +88,6 @@ impl GroupMetricsAccumulator { } } -// Thread-slot assignment for [`GroupMetricsAccumulator`] lookup. Each worker -// thread lazily claims a slot index from a process-wide counter and caches it -// in TLS; callers wrap via `% num_slots`, so reused threads and counter growth -// across test runs stay correct (collisions just share a slot). -static GROUP_SLOT_COUNTER: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); -std::thread_local! { - static GROUP_SLOT: std::cell::Cell> = const { std::cell::Cell::new(None) }; -} - -#[inline] -fn group_thread_slot(num_slots: usize) -> usize { - GROUP_SLOT.with(|c| { - c.get().unwrap_or_else(|| { - let s = GROUP_SLOT_COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - c.set(Some(s)); - s - }) - }) % num_slots.max(1) -} - /// Configuration for template filtering during group processing. #[derive(Clone)] struct GroupFilterConfig { @@ -1071,11 +1051,7 @@ impl Command for GroupReadsByUmi { // so retained memory is O(threads × distinct sizes) rather than growing // one hashmap per position group (see issue #285). let num_threads = self.threading.num_threads(); - let accumulators: Arc>> = Arc::new( - (0..num_threads.max(1)) - .map(|_| Mutex::new(GroupMetricsAccumulator::default())) - .collect(), - ); + let accumulators = PerThreadAccumulator::::new(num_threads); // Clone values needed by closures let strategy = effective_strategy; @@ -1414,10 +1390,9 @@ impl Command for GroupReadsByUmi { // Merge per-group metrics into this worker's accumulator slot. // Memory stays O(threads × distinct sizes) instead of growing // one hashmap per position group. - let slot = group_thread_slot(accumulators_clone.len()); - accumulators_clone[slot] - .lock() - .record_group(processed.family_sizes, &processed.filter_metrics); + accumulators_clone.with_slot(|acc| { + acc.record_group(processed.family_sizes, &processed.filter_metrics); + }); // Save input record count for progress tracking let input_record_count = processed.input_record_count; @@ -1555,7 +1530,7 @@ impl Command for GroupReadsByUmi { let mut position_group_size_counter: AHashMap = AHashMap::with_capacity(50); let mut total_filter_metrics = FilterMetrics::new(); - for slot in accumulators.iter() { + for slot in accumulators.slots() { let acc = slot.lock(); for (&size, &count) in &acc.family_sizes { *family_size_counter.entry(size).or_insert(0) += count; diff --git a/src/lib/commands/simplex.rs b/src/lib/commands/simplex.rs index 08725f454..a91c43eac 100644 --- a/src/lib/commands/simplex.rs +++ b/src/lib/commands/simplex.rs @@ -27,9 +27,9 @@ use clap::Parser; use fgoxide::io::DelimFile; use fgumi_raw_bam::RawRecord; // RejectionTracker now used via ConsensusStatsOps trait in consensus_runner +use crate::per_thread_accumulator::PerThreadAccumulator; use crate::sam::SamTag; use crate::vanilla_consensus_caller::{VanillaUmiConsensusCaller, VanillaUmiConsensusOptions}; -use crossbeam_queue::SegQueue; use log::info; use noodles::sam::Header; @@ -79,10 +79,10 @@ impl MemoryEstimate for SimplexProcessedBatch { } } -/// Metrics collected from each batch during parallel processing. +/// Per-thread accumulator for simplex consensus metrics and rejects. /// -/// Used with a lock-free queue (`SegQueue`) to collect metrics from -/// parallel workers for post-pipeline aggregation. +/// Merged into final aggregates after the pipeline completes; one instance +/// per worker slot (see [`PerThreadAccumulator`]). #[derive(Default)] struct CollectedSimplexMetrics { /// Consensus calling statistics @@ -91,7 +91,9 @@ struct CollectedSimplexMetrics { overlapping_stats: Option, /// Number of MI groups processed groups_processed: u64, - /// Rejected reads as raw BAM bytes for deferred writing + /// Rejected reads as raw BAM bytes for deferred writing. + /// NB: rejects buffering is left unchanged here; streaming rejects to + /// disk during the pipeline is a follow-up (tracked separately). rejects: Vec>, } @@ -495,8 +497,9 @@ impl Simplex { num_threads, )?; - // Lock-free metrics collection - let collected_metrics: Arc> = Arc::new(SegQueue::new()); + // Per-thread metrics accumulator: bounded metric memory, no unbounded + // queue. Rejects buffering semantics are preserved (see follow-up). + let collected_metrics = PerThreadAccumulator::::new(num_threads); let collected_metrics_for_serialize = Arc::clone(&collected_metrics); // Capture configuration for closures @@ -643,12 +646,18 @@ impl Simplex { _header: &Header, output: &mut Vec| -> io::Result { - // Collect metrics (lock-free) - collected_metrics_for_serialize.push(CollectedSimplexMetrics { - stats: processed.stats, - overlapping_stats: processed.overlapping_stats, - groups_processed: processed.groups_count, - rejects: processed.rejects, + // Merge per-batch metrics + rejects into this worker's slot + let batch_stats = processed.stats; + let batch_overlapping = processed.overlapping_stats; + let groups_count = processed.groups_count; + let mut batch_rejects = processed.rejects; + collected_metrics_for_serialize.with_slot(|m| { + m.stats.merge(&batch_stats); + if let Some(o) = batch_overlapping { + m.overlapping_stats.get_or_insert_with(CorrectionStats::new).merge(&o); + } + m.groups_processed += groups_count; + m.rejects.append(&mut batch_rejects); }); // Serialize consensus reads @@ -665,14 +674,15 @@ impl Simplex { let mut merged_overlapping_stats = CorrectionStats::new(); let mut all_rejects: Vec> = Vec::new(); - while let Some(metrics) = collected_metrics.pop() { - total_groups += metrics.groups_processed; - merged_stats.merge(&metrics.stats); - if let Some(ref ocs) = metrics.overlapping_stats { + for slot in collected_metrics.slots() { + let mut m = slot.lock(); + total_groups += m.groups_processed; + merged_stats.merge(&m.stats); + if let Some(ref ocs) = m.overlapping_stats { merged_overlapping_stats.merge(ocs); } if track_rejects { - all_rejects.extend(metrics.rejects); + all_rejects.append(&mut m.rejects); } } diff --git a/src/lib/mod.rs b/src/lib/mod.rs index 2267b0844..e8cb3808e 100644 --- a/src/lib/mod.rs +++ b/src/lib/mod.rs @@ -146,6 +146,7 @@ pub mod metrics; pub mod mi_group; pub use fgumi_consensus::phred; pub mod os_hints; +pub mod per_thread_accumulator; pub mod prefetch_reader; pub mod progress; pub mod read_info; diff --git a/src/lib/per_thread_accumulator.rs b/src/lib/per_thread_accumulator.rs new file mode 100644 index 000000000..3d06c3697 --- /dev/null +++ b/src/lib/per_thread_accumulator.rs @@ -0,0 +1,167 @@ +#![deny(unsafe_code)] + +//! Sharded per-thread accumulator used to cap metric memory in parallel +//! pipeline stages. +//! +//! Pipeline commands historically collected per-group metric structs into an +//! unbounded `SegQueue`, retaining one entry per position/MI group for the +//! whole run. At real-data scale (hundreds of millions of records, tens of +//! millions of groups) the backing `AHashMap`s and `SegQueue` nodes grew into +//! tens of gigabytes even though downstream reduction only needed a handful of +//! counters (see issue #285). +//! +//! [`PerThreadAccumulator`] replaces that pattern with a fixed number of +//! [`Mutex`] slots — one per worker thread. Each worker claims a slot on +//! first use and merges per-group results into it immediately, so retained +//! memory is `O(threads × distinct keys)` instead of `O(groups)`. After the +//! pipeline completes, callers fold the slots into the final metric output. +//! +//! The slot index is assigned lazily from a process-wide atomic counter +//! ([`SLOT_COUNTER`]) and cached in thread-local storage ([`THREAD_SLOT`]). +//! Indices are taken modulo the instance's slot count, so reused threads +//! across test runs, or long-lived threads that touch multiple accumulators, +//! remain correct — at worst two threads share a slot and merge under the same +//! mutex. Slot distribution is only uniform when `num_slots` is at least the +//! count of distinct global indices ever assigned; with a smaller `num_slots` +//! the modulo mapping deterministically clusters multiple threads onto the +//! same slot and adds lock contention, but does not affect correctness. + +use parking_lot::Mutex; +use std::cell::Cell; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +static SLOT_COUNTER: AtomicUsize = AtomicUsize::new(0); +std::thread_local! { + static THREAD_SLOT: Cell> = const { Cell::new(None) }; +} + +/// Per-thread accumulator with a fixed number of mutex-protected slots. +/// +/// See [module docs][self] for rationale and the memory-scaling contract. +#[derive(Debug)] +pub struct PerThreadAccumulator { + slots: Vec>, +} + +impl PerThreadAccumulator { + /// Allocates `num_slots` default-initialized slots. `num_slots` is clamped + /// to at least 1; callers typically pass the pipeline's worker count. + #[must_use] + pub fn new(num_slots: usize) -> Arc { + let slots = (0..num_slots.max(1)).map(|_| Mutex::new(A::default())).collect(); + Arc::new(Self { slots }) + } + + /// Runs `f` against the calling thread's slot. + /// + /// The mutex is held only for the duration of `f`; do not perform blocking + /// I/O inside the closure. Slot assignment is lazy and persisted in TLS, + /// so repeated calls from the same thread contend on the same slot. + #[inline] + pub fn with_slot(&self, f: F) -> R + where + F: FnOnce(&mut A) -> R, + { + let slot = global_thread_index() % self.slots.len(); + let mut guard = self.slots[slot].lock(); + f(&mut guard) + } + + /// Borrows each slot in order. Callers typically reduce into a final + /// aggregate after the pipeline has returned. + /// + /// Safe to call while other `Arc` holders exist, unlike [`Self::into_slots`]. + #[must_use] + pub fn slots(&self) -> &[Mutex] { + &self.slots + } + + /// Consumes the accumulator and yields each slot's inner value. + /// + /// Requires unique `Arc` ownership — i.e. the pipeline closures that held + /// clones have been dropped. Debug builds assert this invariant. If other + /// `Arc` holders remain in release builds, falls back to `std::mem::take` + /// under `A: Default`, which is lossy under concurrent use: any + /// `with_slot` call racing from a surviving holder after the take will land + /// in the freshly defaulted slot and be dropped when that holder releases + /// its `Arc`. Callers must quiesce all writers before invoking. + pub fn into_slots(self: Arc) -> Vec { + debug_assert_eq!( + Arc::strong_count(&self), + 1, + "into_slots called with outstanding Arc holders; fallback is lossy under concurrent writes", + ); + match Arc::try_unwrap(self) { + Ok(inner) => inner.slots.into_iter().map(Mutex::into_inner).collect(), + Err(arc) => arc.slots.iter().map(|m| std::mem::take(&mut *m.lock())).collect(), + } + } +} + +/// Returns a process-wide unique index for the calling thread, lazily +/// assigned from [`SLOT_COUNTER`] and cached in [`THREAD_SLOT`]. Callers map +/// the index into their own slot range with modulo. +#[inline] +fn global_thread_index() -> usize { + THREAD_SLOT.with(|c| { + c.get().unwrap_or_else(|| { + let s = SLOT_COUNTER.fetch_add(1, Ordering::Relaxed); + c.set(Some(s)); + s + }) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::thread; + + #[derive(Default, Debug)] + struct Counter(u64); + + #[test] + fn single_thread_accumulates_in_one_slot() { + let acc: Arc> = PerThreadAccumulator::new(4); + for _ in 0..100 { + acc.with_slot(|c| c.0 += 1); + } + let total: u64 = acc.slots().iter().map(|s| s.lock().0).sum(); + assert_eq!(total, 100); + } + + #[test] + fn parallel_threads_share_total_count() { + let acc: Arc> = PerThreadAccumulator::new(8); + thread::scope(|s| { + for _ in 0..8 { + let acc = Arc::clone(&acc); + s.spawn(move || { + for _ in 0..1_000 { + acc.with_slot(|c| c.0 += 1); + } + }); + } + }); + let total: u64 = acc.slots().iter().map(|s| s.lock().0).sum(); + assert_eq!(total, 8_000); + } + + #[test] + fn num_slots_zero_clamped_to_one() { + let acc: Arc> = PerThreadAccumulator::new(0); + assert_eq!(acc.slots().len(), 1); + acc.with_slot(|c| c.0 = 42); + assert_eq!(acc.slots()[0].lock().0, 42); + } + + #[test] + fn into_slots_drains_under_unique_ownership() { + let acc: Arc> = PerThreadAccumulator::new(4); + acc.with_slot(|c| c.0 = 7); + let slots = acc.into_slots(); + assert_eq!(slots.len(), 4); + assert_eq!(slots.iter().map(|c| c.0).sum::(), 7); + } +}