reuse batch compressors

.changeset/reuse-batch-compressors.md

@@ -0,0 +1,5 @@
+---
+"@medicomind/rolldown-compression": patch
+---
+
+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.

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<mbrotli::Compressor>,
+    zstd: Option<inner_zstd::ZstdContext>,
+}
+
+// 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<u32>,
+        section_size: Option<u32>,
+        input: &[u8],
+    ) -> Result<Vec<u8>, 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<u32>,
     input: &[u8],
 ) -> Result<Vec<u8>, 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)]

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<Compressor>,
     level: u32,
     window_bits: Option<u32>,
     section_size: Option<u32>,
@@ -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<Option<Compressor>> = 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<Vec<u8>, 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<Compressor>,
+    config: EncoderConfig,
+    input: &[u8],
+) -> Result<Vec<u8>, 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<u8> = 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),

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<Vec<u8>, 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<Vec<u8>, 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)
 }

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<ZstdContext> = Default::default();
-}

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::{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<String>,
     }
 
-    struct ParsedAlgorithm {
-        algorithm: Algorithm,
-        level: u32,
-        window_bits: Option<u32>,
-        section_size: Option<u32>,
-    }
-
     pub struct CompressWorker {
         files: Vec<CompressFile>,
-        algorithms: Vec<ParsedAlgorithm>,
+        algorithms: Vec<BatchAlgorithm>,
         skip_if_larger_or_equal: bool,
     }
 
@@ -115,45 +108,29 @@ 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 (metadata, items): (Vec<_>, Vec<BatchItem>) =
-                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| {
-                                (
-                                    (file.file_name.clone(), config.algorithm, original_size),
-                                    BatchItem {
-                                        algorithm: config.algorithm,
-                                        level: config.level,
-                                        window_bits: config.window_bits,
-                                        section_size: config.section_size,
-                                        input: Arc::clone(&input),
-                                    },
-                                )
-                            })
-                        })
-                        .unzip()
-                });
-
-            let outcomes = run_batch(items, skip_if_larger_or_equal);
+            let (metadata, inputs): (Vec<_>, Vec<_>) = files
+                .into_iter()
+                .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);
 
-            Ok(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())
         }
 
@@ -268,7 +245,7 @@ mod binding {
                     validate_section_size(section_size)?;
                 }
             }
-            parsed.push(ParsedAlgorithm {
+            parsed.push(BatchAlgorithm {
                 algorithm,
                 level,
                 window_bits: config.window_bits,