From 6ec822a6988ef9f9bacbf8d393e6bafa9a7b15e5 Mon Sep 17 00:00:00 2001 From: Mnwa Date: Wed, 9 Sep 2026 17:09:10 +0300 Subject: [PATCH 1/2] reuse batch compressors --- .changeset/reuse-batch-compressors.md | 5 ++ src/compress.rs | 53 ++++++++++++++++---- src/compress/inner_brotli.rs | 68 ++++++++++--------------- src/compress/inner_zstd.rs | 24 ++++----- src/compress/inner_zstd/context.rs | 10 +--- src/lib.rs | 39 +++++++++------ src/scheduler.rs | 72 ++++++++++++++++++++++----- 7 files changed, 170 insertions(+), 101 deletions(-) create mode 100644 .changeset/reuse-batch-compressors.md diff --git a/.changeset/reuse-batch-compressors.md b/.changeset/reuse-batch-compressors.md new file mode 100644 index 0000000..a1aea30 --- /dev/null +++ b/.changeset/reuse-batch-compressors.md @@ -0,0 +1,5 @@ +--- +"@medicomind/rolldown-compression": patch +--- + +Group compression tasks by algorithm settings and reuse Brotli and Zstd compressors within parallel file iterators. Release compressor scratch memory when processing completes instead of retaining it in thread-local storage, while preserving result order. diff --git a/src/compress.rs b/src/compress.rs index 9812cef..df411ef 100644 --- a/src/compress.rs +++ b/src/compress.rs @@ -96,6 +96,47 @@ impl Display for Algorithm { } } +/// Scratch encoders owned by one rayon iterator partition. +#[derive(Default)] +pub(crate) struct Compressors { + brotli: Option, + zstd: Option, +} + +// Rayon clones the state when splitting `map_with`. Encoded streams never +// depend on the scratch buffers, so each partition starts with empty caches. +impl Clone for Compressors { + fn clone(&self) -> Self { + Self::default() + } +} + +impl Compressors { + #[hotpath::measure] + pub(crate) fn compress( + &mut self, + algorithm: Algorithm, + level: u32, + window_bits: Option, + section_size: Option, + input: &[u8], + ) -> Result, Error> { + algorithm.validate_level(level)?; + let mut output = match algorithm { + Algorithm::Gzip => inner_gzip::compress(level, input), + Algorithm::Brotli => { + inner_brotli::compress(&mut self.brotli, level, window_bits, section_size, input) + } + Algorithm::Zstd => { + inner_zstd::compress(self.zstd.get_or_insert_with(Default::default), level, input) + } + }?; + // Results remain alive until JS drains the batch; release unused capacity. + output.shrink_to_fit(); + Ok(output) + } +} + /// Compress `input` with the given algorithm and level. /// /// `window_bits` and `section_size` are only used by brotli and ignored by @@ -111,17 +152,7 @@ pub fn compress( section_size: Option, input: &[u8], ) -> Result, Error> { - algorithm.validate_level(level)?; - let mut output = match algorithm { - Algorithm::Gzip => inner_gzip::compress(level, input), - Algorithm::Brotli => inner_brotli::compress(level, window_bits, section_size, input), - Algorithm::Zstd => inner_zstd::compress(level, input), - }?; - // Output buffers are sized for the worst case, so compressible input - // leaves most of the capacity unused; results are held until the JS side - // drains the batch, so hand back right-sized buffers. - output.shrink_to_fit(); - Ok(output) + Compressors::default().compress(algorithm, level, window_bits, section_size, input) } #[cfg(test)] diff --git a/src/compress/inner_brotli.rs b/src/compress/inner_brotli.rs index bf9a81e..84c460b 100644 --- a/src/compress/inner_brotli.rs +++ b/src/compress/inner_brotli.rs @@ -13,7 +13,6 @@ use mbrotli::compressor::parallel::{ BatchConfig, ParallelCompressor, ParallelConfig, SegmentSize, TaskCount, }; use mbrotli::{Compressor, EncoderConfig, Quality, Window}; -use std::cell::RefCell; use std::ops::RangeInclusive; /// Default brotli window size (log2), matching `BROTLI_DEFAULT_WINDOW`. @@ -59,6 +58,7 @@ pub fn validate_section_size(section_size: u32) -> Result<(), Error> { /// decision that leaves the bytes alone, so a given input and set of options /// compress the same everywhere, as gzip and zstd already did. pub fn compress( + compressor: &mut Option, level: u32, window_bits: Option, section_size: Option, @@ -104,7 +104,7 @@ pub fn compress( // The threshold follows the clamped segment, so the size that decides // whether to split is the same one that decides how. if input.len() < BROTLI_MIN_SECTIONS * segment.get() { - compress_single(config, input) + compress_single(compressor, config, input) } else { compress_parallel(config, segment, input) } @@ -130,21 +130,6 @@ fn segment_size(section_size: usize) -> SegmentSize { .unwrap_or(SegmentSize::DEFAULT) } -thread_local! { - /// The calling worker's serial encoder, reused across every file it takes. - /// - /// A `Compressor` exists to be reused: its second call at a given shape - /// allocates nothing the first already paid for, and a batch is overwhelmingly - /// files below [`parallel_threshold`] all landing here. Building one per file - /// instead threw that away and re-paid for brotli's hasher tables and ring - /// buffer every time. - /// - /// One per worker rather than one shared behind a lock: the encoder is - /// `&mut`-driven, so sharing would serialize the batch it is meant to - /// parallelize. - static COMPRESSOR: RefCell> = const { RefCell::new(None) }; -} - /// Compress as one uninterrupted stream on the calling thread. /// /// Splitting is not free under `mbrotli` the way it was under the previous @@ -159,17 +144,19 @@ thread_local! { /// share a shape; `reconfigure` is transactional — it drops every trace of the /// previous stream while keeping whatever buffers still apply — so the output /// is what a fresh `Compressor` would have produced. -fn compress_single(config: EncoderConfig, input: &[u8]) -> Result, Error> { - COMPRESSOR.with_borrow_mut(|slot| { - match slot { - Some(compressor) => compressor.reconfigure(config)?, - None => *slot = Some(Compressor::new(config)?), +fn compress_single( + slot: &mut Option, + config: EncoderConfig, + input: &[u8], +) -> Result, Error> { + let compressor = match slot { + Some(compressor) => { + compressor.reconfigure(config)?; + compressor } - let compressor = slot - .as_mut() - .expect("compressor is present after the match above"); - compressor.compress(input).map_err(Error::from) - }) + None => slot.insert(Compressor::new(config)?), + }; + compressor.compress(input).map_err(Error::from) } /// Compress by cutting `input` into `segment`-sized sections and spreading @@ -567,7 +554,7 @@ mod tests { #[test] fn reused_compressor_matches_a_fresh_one_across_shape_changes() { - // The thread-local encoder is reconfigured, not rebuilt, so a batch + // The partition-owned encoder is reconfigured, not rebuilt, so a batch // walks one `Compressor` through every quality and window it meets. // If any state survived a reconfigure the output would drift from // what a fresh encoder produces — silently, and only for whichever @@ -589,6 +576,7 @@ mod tests { (11, 24, 90_000), (5, 10, 0), ]; + let mut compressor = None; for _ in 0..2 { for (level, window_bits, len) in shapes { let input: Vec = b"export const value = 42; // padding padding\n" @@ -602,7 +590,8 @@ mod tests { .expect("compressor") .compress(input.as_ref()) .expect("compress"); - let reused = compress_single(cfg, input.as_ref()).expect("compress"); + let reused = + compress_single(&mut compressor, cfg, input.as_ref()).expect("compress"); assert_eq!( reused, fresh, "reused encoder drifted at quality {level}, window {window_bits}, len {len}" @@ -614,28 +603,21 @@ mod tests { #[test] fn every_worker_gets_its_own_compressor() { - // The cache is thread-local and the batch is a rayon fan-out, so the - // same encoder must not be reached from two workers at once, and a - // worker that steals a serial file while another compression is in - // flight must not find the RefCell already borrowed. + // Each rayon partition owns its encoder, including when a worker + // steals another partition while parallel compression is in flight. use rayon::prelude::*; let jobs = 32 * rayon::current_num_threads(); let outputs: Vec<_> = (0..jobs) .into_par_iter() - .map(|i| { + .map_with(crate::compress::Compressors::default(), |compressors, i| { // Vary the shape per job so workers keep reconfiguring. let level = (i % 12) as u32; let window_bits = 10 + (i % 15) as u32; let input = b"function chunk(a, b) { return a + b; }\n".repeat(100 + i % 500); - let compressed = compress_any( - Algorithm::Brotli, - level, - Some(window_bits), - None, - input.clone(), - ) - .expect("compress"); + let compressed = compressors + .compress(Algorithm::Brotli, level, Some(window_bits), None, &input) + .expect("compress"); assert_eq!(decompress(Algorithm::Brotli, &compressed), input); (level, window_bits, input, compressed) }) @@ -666,7 +648,7 @@ mod tests { let input = b"export const value = 42; // padding padding\n".repeat(20_000); assert!(input.len() < DEFAULT_SECTION_SIZE); assert_eq!( - compress_single(config(5, 22), input.as_ref()).expect("serial"), + compress_single(&mut None, config(5, 22), input.as_ref()).expect("serial"), compress_parallel( config(5, 22), segment_size(DEFAULT_SECTION_SIZE), diff --git a/src/compress/inner_zstd.rs b/src/compress/inner_zstd.rs index acecbe4..14df9c4 100644 --- a/src/compress/inner_zstd.rs +++ b/src/compress/inner_zstd.rs @@ -3,17 +3,17 @@ use crate::error::Error; mod context; +pub(super) use context::ZstdContext; + #[hotpath::measure(label = "compress_zstd")] -pub fn compress(level: u32, input: &[u8]) -> Result, Error> { - context::CONTEXT.with_borrow_mut(|context| { - let level = level as i32; - if context.level != level { - context - .compressor - .set_compression_level(level) - .map_err(Error::Zstd)?; - context.level = level; - } - context.compressor.compress(input).map_err(Error::Zstd) - }) +pub fn compress(context: &mut ZstdContext, level: u32, input: &[u8]) -> Result, Error> { + let level = level as i32; + if context.level != level { + context + .compressor + .set_compression_level(level) + .map_err(Error::Zstd)?; + context.level = level; + } + context.compressor.compress(input).map_err(Error::Zstd) } diff --git a/src/compress/inner_zstd/context.rs b/src/compress/inner_zstd/context.rs index 5c75a7b..6b86449 100644 --- a/src/compress/inner_zstd/context.rs +++ b/src/compress/inner_zstd/context.rs @@ -1,10 +1,8 @@ -use std::cell::RefCell; - -/// Reusable zstd compressor carried across the items a single rayon worker +/// Reusable zstd compressor carried across the items a rayon partition /// handles. /// /// A zstd context at the levels used here owns tens of megabytes of match -/// tables; keeping one per worker avoids reallocating them for every file. +/// tables; retaining them in the partition avoids rebuilding them for each file. /// `i32::MIN` marks a fresh context whose level is not yet configured /// (validated levels are all above it). pub struct ZstdContext { @@ -20,7 +18,3 @@ impl Default for ZstdContext { } } } - -thread_local! { - pub static CONTEXT: RefCell = Default::default(); -} diff --git a/src/lib.rs b/src/lib.rs index c3e7635..896fd4e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -28,7 +28,7 @@ mod binding { use crate::compress::{Algorithm, validate_section_size, validate_window_bits}; use crate::error::Error as CompressionError; - use crate::scheduler::{BatchItem, BatchOutcome, run_batch}; + use crate::scheduler::{BatchItem, BatchOutcome, run_batch, scatter_in_place}; /// One source file, passed once regardless of the number of algorithms. #[napi(object)] @@ -115,35 +115,38 @@ mod binding { let algorithms = &self.algorithms; let skip_if_larger_or_equal = self.skip_if_larger_or_equal; - // Expand file × algorithm on a worker thread. Only Arc handles - // are cloned: every task reads the same source allocation, which - // is released when the last algorithm for that file finishes. + let file_count = files.len(); + let algorithm_count = algorithms.len(); + // Share each source allocation across the algorithm groups. + let files: Vec<_> = files + .into_iter() + .map(|file| (file.file_name, Arc::new(file.data))) + .collect(); let (metadata, items): (Vec<_>, Vec) = hotpath::measure_block!("CompressWorker::split_tasks", { - files - .into_par_iter() - .flat_map_iter(|file| { - let original_size = file.data.len() as u32; - let input = Arc::new(file.data); - algorithms.iter().map(move |config| { + algorithms + .par_iter() + .flat_map_iter(|config| { + files.iter().map(move |(file_name, input)| { ( - (file.file_name.clone(), config.algorithm, original_size), + (file_name.clone(), config.algorithm, input.len() as u32), BatchItem { algorithm: config.algorithm, level: config.level, window_bits: config.window_bits, section_size: config.section_size, - input: Arc::clone(&input), + input: Arc::clone(input), }, ) }) }) .unzip() }); + drop(files); let outcomes = run_batch(items, skip_if_larger_or_equal); - Ok(metadata + let mut output: Vec<_> = metadata .into_par_iter() .zip(outcomes) .map( @@ -154,7 +157,15 @@ mod binding { outcome, }, ) - .collect()) + .collect(); + // Keep the public file-then-algorithm result order. + let mut order: Vec = (0..algorithm_count) + .flat_map(|algorithm| { + (0..file_count).map(move |file| (file * algorithm_count + algorithm) as u32) + }) + .collect(); + scatter_in_place(&mut output, &mut order); + Ok(output) } fn resolve(&mut self, _env: Env, output: Self::Output) -> Result { diff --git a/src/scheduler.rs b/src/scheduler.rs index a6cb6ed..7f7df08 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -9,7 +9,7 @@ use std::sync::Arc; use rayon::prelude::*; -use crate::compress::{Algorithm, InputBuffer, compress}; +use crate::compress::{Algorithm, Compressors, InputBuffer}; use crate::error::Error; /// A single unit of compression work. @@ -54,7 +54,7 @@ fn algorithm_rank(algorithm: Algorithm) -> u8 { /// Run every item of the batch in parallel and return outcomes in input order. /// /// Items are scheduled longest-job-first — brotli, then zstd, then gzip, and -/// the largest input first inside each algorithm — and the outcomes are put +/// the largest input first inside each configuration — and the outcomes are put /// back into input order before returning. /// /// Work runs on the caller's ambient rayon pool: the global one unless the @@ -75,17 +75,51 @@ pub fn run_batch(mut items: Vec, skip_if_larger_or_equal: bool) -> Ve let mut order: Vec = (0..items.len() as u32).collect(); order.sort_unstable_by_key(|&i| { let item = &items[i as usize]; - (algorithm_rank(item.algorithm), Reverse(item.input.len())) + ( + algorithm_rank(item.algorithm), + item.level, + item.window_bits, + item.section_size, + Reverse(item.input.len()), + ) }); gather_in_place(&mut items, &order); - let mut outcomes: Vec = Vec::with_capacity(items.len()); + // Each configuration owns a parallel file iterator and its reusable + // encoders. Let rayon keep multiple files in a partition so `map_with` + // can retain scratch allocations between them. + let mut groups: Vec> = Vec::new(); + for item in items { + if let Some(group) = groups.last_mut() + && let Some(last) = group.last() + && ( + last.algorithm, + last.level, + last.window_bits, + last.section_size, + ) == ( + item.algorithm, + item.level, + item.window_bits, + item.section_size, + ) + { + group.push(item); + } else { + groups.push(vec![item]); + } + } - items + let mut outcomes: Vec = groups .into_par_iter() - .with_max_len(1) - .map(|item| run_one(item, skip_if_larger_or_equal)) - .collect_into_vec(&mut outcomes); + .flat_map(|items| { + items + .into_par_iter() + .map_with(Compressors::default(), |compressors, item| { + run_one(compressors, item, skip_if_larger_or_equal) + }) + }) + .collect(); // Outcomes come back in scheduled order; `order` says where each belongs. scatter_in_place(&mut outcomes, &mut order); @@ -112,7 +146,7 @@ fn gather_in_place(data: &mut [T], order: &[u32]) { /// Move every `data[i]` to index `order[i]`, the inverse of /// [`gather_in_place`]. `order` is used as scratch and left as the identity /// permutation. -fn scatter_in_place(data: &mut [T], order: &mut [u32]) { +pub(crate) fn scatter_in_place(data: &mut [T], order: &mut [u32]) { for i in 0..data.len() { // Each swap parks at least one element at its final index, so the // inner loop runs at most `len` times in total. @@ -125,12 +159,16 @@ fn scatter_in_place(data: &mut [T], order: &mut [u32]) { } #[hotpath::measure] -fn run_one(item: BatchItem, skip_if_larger_or_equal: bool) -> BatchOutcome { +fn run_one( + compressors: &mut Compressors, + item: BatchItem, + skip_if_larger_or_equal: bool, +) -> BatchOutcome { let input_len = item.input.len(); let algorithm = item.algorithm; // The item owns one shared reference; finishing this task releases it. let result = catch_unwind(AssertUnwindSafe(|| { - compress( + compressors.compress( item.algorithm, item.level, item.window_bits, @@ -138,7 +176,11 @@ fn run_one(item: BatchItem, skip_if_larger_or_equal: bool) -> BatchOutcome { &item.input, ) })) - .unwrap_or(Err(Error::CompressionPanicked(algorithm))); + .unwrap_or_else(|_| { + // A panicking encoder may contain an unfinished stream. + *compressors = Compressors::default(); + Err(Error::CompressionPanicked(algorithm)) + }); match result { Ok(data) if skip_if_larger_or_equal && data.len() >= input_len => BatchOutcome { @@ -386,7 +428,11 @@ mod tests { section_size: None, input: Arc::clone(&input), }; - let failed = run_one(make_item(Algorithm::Brotli, 99), false); + let failed = run_one( + &mut Compressors::default(), + make_item(Algorithm::Brotli, 99), + false, + ); assert!(failed.error.is_some()); let items = vec![ make_item(Algorithm::Gzip, 6), From df422bcd03db83c71c6d3d6d72928840ef0dfe46 Mon Sep 17 00:00:00 2001 From: Mnwa Date: Wed, 9 Sep 2026 17:27:01 +0300 Subject: [PATCH 2/2] reduce batch ops --- .changeset/reuse-batch-compressors.md | 2 +- src/lib.rs | 78 ++-- src/scheduler.rs | 528 +++++++++++--------------- 3 files changed, 252 insertions(+), 356 deletions(-) diff --git a/.changeset/reuse-batch-compressors.md b/.changeset/reuse-batch-compressors.md index a1aea30..4c91af4 100644 --- a/.changeset/reuse-batch-compressors.md +++ b/.changeset/reuse-batch-compressors.md @@ -2,4 +2,4 @@ "@medicomind/rolldown-compression": patch --- -Group compression tasks by algorithm settings and reuse Brotli and Zstd compressors within parallel file iterators. Release compressor scratch memory when processing completes instead of retaining it in thread-local storage, while preserving result order. +Build compression groups directly from algorithm settings and files sorted by size, scheduling Brotli, Zstd, then gzip with larger files first. Reuse Brotli and Zstd compressors within parallel file iterators and release their scratch memory when processing completes instead of retaining it in thread-local storage, while preserving result order. diff --git a/src/lib.rs b/src/lib.rs index 896fd4e..3e250e3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -28,7 +28,7 @@ mod binding { use crate::compress::{Algorithm, validate_section_size, validate_window_bits}; use crate::error::Error as CompressionError; - use crate::scheduler::{BatchItem, BatchOutcome, run_batch, scatter_in_place}; + use crate::scheduler::{BatchAlgorithm, BatchOutcome, run_batch}; /// One source file, passed once regardless of the number of algorithms. #[napi(object)] @@ -85,16 +85,9 @@ mod binding { pub error: Option, } - struct ParsedAlgorithm { - algorithm: Algorithm, - level: u32, - window_bits: Option, - section_size: Option, - } - pub struct CompressWorker { files: Vec, - algorithms: Vec, + algorithms: Vec, skip_if_larger_or_equal: bool, } @@ -115,57 +108,30 @@ mod binding { let algorithms = &self.algorithms; let skip_if_larger_or_equal = self.skip_if_larger_or_equal; - let file_count = files.len(); - let algorithm_count = algorithms.len(); - // Share each source allocation across the algorithm groups. - let files: Vec<_> = files + let (metadata, inputs): (Vec<_>, Vec<_>) = files .into_iter() - .map(|file| (file.file_name, Arc::new(file.data))) - .collect(); - let (metadata, items): (Vec<_>, Vec) = - hotpath::measure_block!("CompressWorker::split_tasks", { - algorithms - .par_iter() - .flat_map_iter(|config| { - files.iter().map(move |(file_name, input)| { - ( - (file_name.clone(), config.algorithm, input.len() as u32), - BatchItem { - algorithm: config.algorithm, - level: config.level, - window_bits: config.window_bits, - section_size: config.section_size, - input: Arc::clone(input), - }, - ) - }) - }) - .unzip() - }); - drop(files); - - let outcomes = run_batch(items, skip_if_larger_or_equal); + .map(|file| { + ( + (file.file_name, file.data.len() as u32), + Arc::new(file.data), + ) + }) + .unzip(); + let outcomes = run_batch(inputs, algorithms, skip_if_larger_or_equal); - let mut output: Vec<_> = metadata + Ok(outcomes .into_par_iter() - .zip(outcomes) - .map( - |((file_name, algorithm, original_size), outcome)| WorkerOutcome { - file_name, - algorithm, - original_size, + .enumerate() + .map(|(index, outcome)| { + let (file_name, original_size) = &metadata[index / algorithms.len()]; + WorkerOutcome { + file_name: file_name.clone(), + algorithm: algorithms[index % algorithms.len()].algorithm, + original_size: *original_size, outcome, - }, - ) - .collect(); - // Keep the public file-then-algorithm result order. - let mut order: Vec = (0..algorithm_count) - .flat_map(|algorithm| { - (0..file_count).map(move |file| (file * algorithm_count + algorithm) as u32) + } }) - .collect(); - scatter_in_place(&mut output, &mut order); - Ok(output) + .collect()) } fn resolve(&mut self, _env: Env, output: Self::Output) -> Result { @@ -279,7 +245,7 @@ mod binding { validate_section_size(section_size)?; } } - parsed.push(ParsedAlgorithm { + parsed.push(BatchAlgorithm { algorithm, level, window_bits: config.window_bits, diff --git a/src/scheduler.rs b/src/scheduler.rs index 7f7df08..99ecf55 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -12,19 +12,26 @@ use rayon::prelude::*; use crate::compress::{Algorithm, Compressors, InputBuffer}; use crate::error::Error; -/// A single unit of compression work. -/// -/// Shares the source allocation with the other algorithms for its file. -/// The last completed item releases it without waiting for the whole batch. -pub struct BatchItem { +/// Algorithm settings shared by every file in a batch. +#[derive(Clone, Copy)] +pub struct BatchAlgorithm { pub algorithm: Algorithm, pub level: u32, pub window_bits: Option, pub section_size: Option, - pub input: Arc, } -/// The outcome of one [`BatchItem`]. +struct BatchFile { + result_index: u32, + input: Arc, +} + +struct BatchGroup<'a> { + config: &'a BatchAlgorithm, + files: Vec, +} + +/// The outcome of compressing one file with one algorithm configuration. /// /// Exactly one of the following holds: /// - `error` is `Some`: the task failed, `data` is empty and `skipped` is false; @@ -51,102 +58,76 @@ fn algorithm_rank(algorithm: Algorithm) -> u8 { } } -/// Run every item of the batch in parallel and return outcomes in input order. -/// -/// Items are scheduled longest-job-first — brotli, then zstd, then gzip, and -/// the largest input first inside each configuration — and the outcomes are put -/// back into input order before returning. -/// -/// Work runs on the caller's ambient rayon pool: the global one unless the -/// caller wraps this in [`rayon::ThreadPool::install`] to pin the batch to a -/// dedicated pool. +/// Sort algorithms and files once, then build the groups directly. +/// Each task owns a source handle, so the last task for a file releases it. +#[hotpath::measure] +fn prepare_groups<'a>( + inputs: Vec>, + algorithms: &'a [BatchAlgorithm], +) -> Vec> { + let mut files: Vec<_> = inputs.into_iter().enumerate().collect(); + files.sort_unstable_by_key(|(index, input)| (Reverse(input.len()), *index)); + let mut configs: Vec<_> = algorithms.iter().enumerate().collect(); + configs.sort_unstable_by_key(|(index, config)| (algorithm_rank(config.algorithm), *index)); + + configs + .into_iter() + .map(|(algorithm_index, config)| BatchGroup { + config, + files: files + .iter() + .map(|(file_index, input)| BatchFile { + result_index: (file_index * algorithms.len() + algorithm_index) as u32, + input: Arc::clone(input), + }) + .collect(), + }) + .collect() +} + +/// Compress every file with every algorithm on the caller's ambient rayon pool. /// -/// * `skip_if_larger_or_equal` — mark items whose compressed size would be -/// `>=` the input size as skipped instead of returning the bloated output. +/// Schedule brotli, then zstd, then gzip, with files largest first inside each +/// configuration. Return outcomes in file order, then configuration order. +/// The binding validates that the file × algorithm count fits u32. /// -/// A failure (or panic) of a single item never aborts the batch; it is -/// reported through [`BatchOutcome::error`]. +/// `skip_if_larger_or_equal` discards compressed output that would be at least +/// as large as its input. A failure or panic is reported per task and never +/// aborts the batch. #[hotpath::measure] -pub fn run_batch(mut items: Vec, skip_if_larger_or_equal: bool) -> Vec { - // `order[scheduled position] == input position`. Four bytes per item is - // the whole cost of the reordering: both permutations below run in place, - // so neither the items nor the outcomes are ever copied into a second - // buffer. The binding validates that the file × algorithm count fits u32. - let mut order: Vec = (0..items.len() as u32).collect(); - order.sort_unstable_by_key(|&i| { - let item = &items[i as usize]; - ( - algorithm_rank(item.algorithm), - item.level, - item.window_bits, - item.section_size, - Reverse(item.input.len()), - ) - }); - gather_in_place(&mut items, &order); - - // Each configuration owns a parallel file iterator and its reusable - // encoders. Let rayon keep multiple files in a partition so `map_with` - // can retain scratch allocations between them. - let mut groups: Vec> = Vec::new(); - for item in items { - if let Some(group) = groups.last_mut() - && let Some(last) = group.last() - && ( - last.algorithm, - last.level, - last.window_bits, - last.section_size, - ) == ( - item.algorithm, - item.level, - item.window_bits, - item.section_size, - ) - { - group.push(item); - } else { - groups.push(vec![item]); - } - } - - let mut outcomes: Vec = groups +pub fn run_batch( + inputs: Vec>, + algorithms: &[BatchAlgorithm], + skip_if_larger_or_equal: bool, +) -> Vec { + let groups = prepare_groups(inputs, algorithms); + let (mut order, mut outcomes): (Vec, Vec) = groups .into_par_iter() - .flat_map(|items| { - items - .into_par_iter() - .map_with(Compressors::default(), |compressors, item| { - run_one(compressors, item, skip_if_larger_or_equal) - }) + .flat_map(|group| { + group.files.into_par_iter().map_with( + Compressors::default(), + move |compressors, file| { + ( + file.result_index, + run_one( + compressors, + group.config, + file.input, + skip_if_larger_or_equal, + ), + ) + }, + ) }) - .collect(); + .unzip(); - // Outcomes come back in scheduled order; `order` says where each belongs. scatter_in_place(&mut outcomes, &mut order); - outcomes } -/// Rearrange `data` so that `data[i]` holds what used to be at `order[i]`. -/// -/// `order` must be a permutation of `0..data.len()`; it is left untouched, and -/// elements only ever move by swapping — nothing is cloned or buffered. -fn gather_in_place(data: &mut [T], order: &[u32]) { - for target in 0..data.len() { - // Slots below `target` are already final: whatever they held has been - // swapped further along, so follow the chain to where it sits now. - let mut source = order[target] as usize; - while source < target { - source = order[source] as usize; - } - data.swap(target, source); - } -} - -/// Move every `data[i]` to index `order[i]`, the inverse of -/// [`gather_in_place`]. `order` is used as scratch and left as the identity -/// permutation. -pub(crate) fn scatter_in_place(data: &mut [T], order: &mut [u32]) { +/// Move every `data[i]` to index `order[i]`. `order` is a permutation of +/// `0..data.len()`, used as scratch and left as the identity permutation. +fn scatter_in_place(data: &mut [T], order: &mut [u32]) { for i in 0..data.len() { // Each swap parks at least one element at its final index, so the // inner loop runs at most `len` times in total. @@ -161,19 +142,20 @@ pub(crate) fn scatter_in_place(data: &mut [T], order: &mut [u32]) { #[hotpath::measure] fn run_one( compressors: &mut Compressors, - item: BatchItem, + config: &BatchAlgorithm, + input: Arc, skip_if_larger_or_equal: bool, ) -> BatchOutcome { - let input_len = item.input.len(); - let algorithm = item.algorithm; - // The item owns one shared reference; finishing this task releases it. + let input_len = input.len(); + let algorithm = config.algorithm; + // Finishing this task releases its shared source handle. let result = catch_unwind(AssertUnwindSafe(|| { compressors.compress( - item.algorithm, - item.level, - item.window_bits, - item.section_size, - &item.input, + config.algorithm, + config.level, + config.window_bits, + config.section_size, + &input, ) })) .unwrap_or_else(|_| { @@ -205,16 +187,52 @@ fn run_one( mod tests { use super::*; + fn decompress(algorithm: Algorithm, input: &[u8]) -> Vec { + use std::io::Read; + match algorithm { + Algorithm::Gzip => { + let mut out = Vec::new(); + flate2::read::GzDecoder::new(input) + .read_to_end(&mut out) + .expect("gzip decode"); + out + } + Algorithm::Brotli => { + let mut out = Vec::new(); + simd_brotli::BrotliDecompress(&mut { input }, &mut out).expect("brotli decode"); + out + } + Algorithm::Zstd => zstd::stream::decode_all(input).expect("zstd decode"), + } + } + fn text_fixture(seed: usize) -> Vec { format!("export const value{seed} = {seed};\n") .repeat(200 + seed * 7) .into_bytes() } - /// Run a batch on a dedicated pool, the way the napi binding does. - /// `threads` of 0 means the rayon default (one per logical CPU). + fn config(algorithm: Algorithm, level: u32) -> BatchAlgorithm { + BatchAlgorithm { + algorithm, + level, + window_bits: None, + section_size: None, + } + } + + fn algorithms() -> [BatchAlgorithm; 4] { + [ + config(Algorithm::Gzip, 1), + config(Algorithm::Brotli, 4), + config(Algorithm::Zstd, 3), + config(Algorithm::Gzip, 9), + ] + } + fn run_batch_on_pool( - items: Vec, + inputs: Vec>, + algorithms: &[BatchAlgorithm], threads: usize, skip_if_larger_or_equal: bool, ) -> Vec { @@ -222,50 +240,45 @@ mod tests { .num_threads(threads) .build() .expect("build pool") - .install(|| run_batch(items, skip_if_larger_or_equal)) - } - - fn make_items(inputs: &[Vec]) -> Vec { - let algorithms = [Algorithm::Gzip, Algorithm::Brotli, Algorithm::Zstd]; - inputs - .iter() - .enumerate() - .map(|(i, input)| { - let algorithm = algorithms[i % algorithms.len()]; - BatchItem { - algorithm, - level: algorithm.default_level(), - window_bits: None, - section_size: None, - input: Arc::new(input.clone()), - } - }) - .collect() + .install(|| run_batch(inputs, algorithms, skip_if_larger_or_equal)) } #[test] - fn batch_preserves_input_order_and_succeeds() { - let inputs: Vec> = (0..24).map(text_fixture).collect(); - let outcomes = run_batch_on_pool(make_items(&inputs), 0, false); - assert_eq!(outcomes.len(), inputs.len()); - for outcome in &outcomes { - assert!(outcome.error.is_none()); - assert!(!outcome.skipped); - assert!(!outcome.data.is_empty()); + fn batch_preserves_file_then_configuration_order() { + let inputs: Vec<_> = [7, 0, 3, 3, 1] + .into_iter() + .map(|seed| Arc::new(text_fixture(seed))) + .chain([Arc::new(Vec::new())]) + .collect(); + let algorithms = algorithms(); + let outcomes = run_batch_on_pool(inputs.clone(), &algorithms, 4, false); + assert_eq!(outcomes.len(), inputs.len() * algorithms.len()); + for (file_outcomes, input) in outcomes.chunks(algorithms.len()).zip(&inputs) { + for (outcome, config) in file_outcomes.iter().zip(&algorithms) { + assert!(outcome.error.is_none()); + assert!(!outcome.skipped); + assert_eq!(decompress(config.algorithm, &outcome.data), **input); + // The two gzip levels have identical algorithm metadata. + let fresh = crate::compress::compress( + config.algorithm, + config.level, + config.window_bits, + config.section_size, + input, + ) + .expect("compress"); + assert_eq!(outcome.data, fresh); + } } } #[test] fn batch_is_deterministic_across_thread_counts() { - // Scheduling never affects output for inputs this size. Brotli inputs - // past `threads * sectionSize` are the one exception — they are cut - // into as many sections as the pool is wide — so keep the fixtures - // well under the sectioning threshold. - let inputs: Vec> = (0..24).map(text_fixture).collect(); - - let single = run_batch_on_pool(make_items(&inputs), 1, false); + let inputs: Vec<_> = (0..24).map(|i| Arc::new(text_fixture(i))).collect(); + let algorithms = algorithms(); + let single = run_batch_on_pool(inputs.clone(), &algorithms, 1, false); for threads in [2, 4, 8] { - let multi = run_batch_on_pool(make_items(&inputs), threads, false); + let multi = run_batch_on_pool(inputs.clone(), &algorithms, threads, false); assert_eq!(single.len(), multi.len()); for (a, b) in single.iter().zip(multi.iter()) { assert_eq!(a.data, b.data, "output differs with {threads} threads"); @@ -275,25 +288,19 @@ mod tests { #[test] fn skip_if_larger_or_equal_marks_incompressible_items() { - // 4 bytes of data always grow under any container format. - let input = vec![1u8, 2, 3, 4]; - let make_items = || { - vec![BatchItem { - algorithm: Algorithm::Gzip, - level: 6, - window_bits: None, - section_size: None, - input: Arc::new(input.clone()), - }] - }; - let outcomes = run_batch_on_pool(make_items(), 0, true); - assert!(outcomes[0].skipped); - assert!(outcomes[0].data.is_empty()); - assert!(outcomes[0].error.is_none()); - - let outcomes = run_batch_on_pool(make_items(), 0, false); - assert!(!outcomes[0].skipped); - assert!(outcomes[0].data.len() > input.len()); + let input = Arc::new(vec![1u8, 2, 3, 4]); + let algorithms = algorithms(); + let outcomes = run_batch_on_pool(vec![Arc::clone(&input)], &algorithms, 0, true); + for outcome in outcomes { + assert!(outcome.skipped); + assert!(outcome.data.is_empty()); + assert!(outcome.error.is_none()); + } + let outcomes = run_batch_on_pool(vec![Arc::clone(&input)], &algorithms, 0, false); + for outcome in outcomes { + assert!(!outcome.skipped); + assert!(outcome.data.len() > input.len()); + } } /// The `rank`-th permutation of `0..n` in Lehmer-code order. @@ -310,183 +317,106 @@ mod tests { } #[test] - fn gather_and_scatter_invert_each_other() { + fn scatter_restores_every_permutation() { for n in 1..=6usize { let factorial: usize = (1..=n).product(); for rank in 0..factorial { - let order = permutation(n, rank); - let source: Vec = (0..n as u32).map(|i| i * 10).collect(); - - let mut data = source.clone(); - gather_in_place(&mut data, &order); - for (target, &from) in order.iter().enumerate() { - assert_eq!( - data[target], source[from as usize], - "gather n={n} rank={rank} order={order:?}" - ); - } - - let mut scratch = order.clone(); - scatter_in_place(&mut data, &mut scratch); - assert_eq!(data, source, "scatter n={n} rank={rank} order={order:?}"); - assert_eq!(scratch, (0..n as u32).collect::>()); + let mut order = permutation(n, rank); + let mut data: Vec<_> = order.iter().map(|i| i * 10).collect(); + scatter_in_place(&mut data, &mut order); + assert_eq!(data, (0..n as u32).map(|i| i * 10).collect::>()); + assert_eq!(order, (0..n as u32).collect::>()); } } } #[test] fn schedules_brotli_then_zstd_then_gzip_largest_first() { - // Sizes are distinct per item so the schedule is fully determined. - let algorithms = [ - Algorithm::Gzip, - Algorithm::Brotli, - Algorithm::Zstd, - Algorithm::Gzip, - Algorithm::Brotli, - Algorithm::Zstd, - ]; - let mut items: Vec = algorithms - .iter() - .enumerate() - .map(|(i, &algorithm)| BatchItem { - algorithm, - level: algorithm.default_level(), - window_bits: None, - section_size: None, - input: Arc::new(vec![0u8; (i + 1) * 10]), - }) + let algorithms = algorithms(); + let inputs = [20, 50, 10] + .into_iter() + .map(|len| Arc::new(vec![0u8; len])) .collect(); - - let mut order: Vec = (0..items.len() as u32).collect(); - order.sort_unstable_by_key(|&i| { - let item = &items[i as usize]; - (algorithm_rank(item.algorithm), Reverse(item.input.len())) - }); - gather_in_place(&mut items, &order); - - let scheduled: Vec<(Algorithm, usize)> = items + let groups = prepare_groups(inputs, &algorithms); + let scheduled: Vec<_> = groups .iter() - .map(|item| (item.algorithm, item.input.len())) + .map(|group| { + ( + group.config.algorithm, + group.config.level, + group + .files + .iter() + .map(|file| file.input.len()) + .collect::>(), + group + .files + .iter() + .map(|file| file.result_index) + .collect::>(), + ) + }) .collect(); assert_eq!( scheduled, vec![ - (Algorithm::Brotli, 50), - (Algorithm::Brotli, 20), - (Algorithm::Zstd, 60), - (Algorithm::Zstd, 30), - (Algorithm::Gzip, 40), - (Algorithm::Gzip, 10), + (Algorithm::Brotli, 4, vec![50, 20, 10], vec![5, 1, 9]), + (Algorithm::Zstd, 3, vec![50, 20, 10], vec![6, 2, 10]), + (Algorithm::Gzip, 1, vec![50, 20, 10], vec![4, 0, 8]), + (Algorithm::Gzip, 9, vec![50, 20, 10], vec![7, 3, 11]), ] ); } - #[test] - fn outcomes_stay_paired_with_their_own_input() { - // Fixtures differ in size and algorithm, so the schedule reorders them - // heavily; every outcome must still decompress back to its own input. - let inputs: Vec> = (0..24).map(text_fixture).collect(); - let items = make_items(&inputs); - let algorithms: Vec = items.iter().map(|item| item.algorithm).collect(); - - let outcomes = run_batch_on_pool(items, 4, false); - assert_eq!(outcomes.len(), inputs.len()); - for ((outcome, input), algorithm) in outcomes.iter().zip(&inputs).zip(algorithms) { - assert!(outcome.error.is_none()); - assert_eq!(&decompress(algorithm, &outcome.data), input); - } - } - - fn decompress(algorithm: Algorithm, input: &[u8]) -> Vec { - use std::io::Read; - match algorithm { - Algorithm::Gzip => { - let mut out = Vec::new(); - flate2::read::GzDecoder::new(input) - .read_to_end(&mut out) - .expect("gzip decode"); - out - } - Algorithm::Brotli => { - let mut out = Vec::new(); - simd_brotli::BrotliDecompress(&mut { input }, &mut out).expect("brotli decode"); - out - } - Algorithm::Zstd => zstd::stream::decode_all(input).expect("zstd decode"), - } - } - #[test] fn shared_input_survives_a_failed_task_and_is_released_after_the_last_task() { let expected = text_fixture(1); let input = Arc::new(expected.clone()); let weak = Arc::downgrade(&input); - let make_item = |algorithm, level| BatchItem { - algorithm, - level, - window_bits: None, - section_size: None, - input: Arc::clone(&input), - }; let failed = run_one( &mut Compressors::default(), - make_item(Algorithm::Brotli, 99), + &config(Algorithm::Brotli, 99), + Arc::clone(&input), false, ); assert!(failed.error.is_some()); - let items = vec![ - make_item(Algorithm::Gzip, 6), - make_item(Algorithm::Zstd, 3), - make_item(Algorithm::Brotli, 4), - ]; - drop(input); assert!(weak.upgrade().is_some()); - let outcomes = run_batch_on_pool(items, 3, false); + let algorithms = algorithms(); + let outcomes = run_batch_on_pool(vec![input], &algorithms, 3, false); assert!(weak.upgrade().is_none()); - for (outcome, algorithm) in - outcomes - .iter() - .zip([Algorithm::Gzip, Algorithm::Zstd, Algorithm::Brotli]) - { + for (outcome, config) in outcomes.iter().zip(&algorithms) { assert!(outcome.error.is_none()); - assert_eq!(decompress(algorithm, &outcome.data), expected); + assert_eq!(decompress(config.algorithm, &outcome.data), expected); } } #[test] fn single_failure_does_not_abort_batch() { - let good = b"hello world hello world hello world".to_vec(); - let items = vec![ - BatchItem { - algorithm: Algorithm::Gzip, - level: 6, - window_bits: None, - section_size: None, - input: Arc::new(good.clone()), - }, - BatchItem { - // Invalid level sneaks past FFI validation only in theory, - // but the scheduler must still isolate the failure. - algorithm: Algorithm::Zstd, - level: 99, - window_bits: None, - section_size: None, - input: Arc::new(good.clone()), - }, - ]; - let outcomes = run_batch_on_pool(items, 0, false); - assert!(outcomes[0].error.is_none()); - assert!(!outcomes[0].data.is_empty()); - assert!(matches!( - outcomes[1].error, - Some(Error::InvalidLevel { - algorithm: Algorithm::Zstd, - level: 99, - .. - }) - )); - assert!(outcomes[1].data.is_empty()); - assert!(!outcomes[1].skipped); + let inputs = vec![Arc::new(text_fixture(0)), Arc::new(text_fixture(1))]; + let algorithms = [config(Algorithm::Gzip, 6), config(Algorithm::Zstd, 99)]; + let outcomes = run_batch_on_pool(inputs, &algorithms, 0, false); + assert_eq!(outcomes.len(), 4); + for file_outcomes in outcomes.chunks(2) { + assert!(file_outcomes[0].error.is_none()); + assert!(!file_outcomes[0].data.is_empty()); + assert!(matches!( + file_outcomes[1].error, + Some(Error::InvalidLevel { + algorithm: Algorithm::Zstd, + level: 99, + .. + }) + )); + assert!(file_outcomes[1].data.is_empty()); + assert!(!file_outcomes[1].skipped); + } + } + + #[test] + fn empty_files_or_algorithms_produce_no_results() { + assert!(run_batch(Vec::new(), &algorithms(), false).is_empty()); + assert!(run_batch(vec![Arc::new(text_fixture(0))], &[], false).is_empty()); + assert!(run_batch(Vec::new(), &[], false).is_empty()); } }