From 641f04f66f3d83e006ed3852ce5cf9a6b5ceabd9 Mon Sep 17 00:00:00 2001 From: Tim Date: Thu, 22 Oct 2026 17:41:54 -0700 Subject: [PATCH] perf(decode): reuse component planes across decodes on the same thread A decode to RGB sized one plane per component (~3.1 MB for a 1080p 4:2:0 frame) and dropped them as soon as the image was assembled, so a caller that builds a fresh Decoder per image -- thumbnailers, MJPEG frame loops, the ArtCraft apps in remade_ffmpeg_rs#48 -- paid for those pages to be faulted in on every decode. Under rusty_alloc/mimalloc that was nearly free; under the Windows system heap it was +0.43 ms per 1080p frame and left 0.4.4 1.13-1.17x behind zune-jpeg on one core (reported by ArtCraft, reproduced 16/16). `compute_image` now borrows the planes and hands them to a small thread-local pool (`decode::plane_pool`); the next decode on the thread takes the best-fitting one before allocating. Bounded to 4 planes and 64 MB (an oversized image's planes are not kept), disabled with RUSTY_JPEG_PLANE_POOL=0, a no-op without std. Planes are refilled, never trusted: every block is written before a plane is read, which the cross-image check confirms. Byte-identical: 30-file fixture set per file, plus all 30 decoded forwards then backwards in ONE process (planes passing between images of different sizes and subsamplings), single- and multi-threaded, under both allocators: 60/60 identical to 0.4.3. Measured vs 0.4.4, default decode() pinned to 1 core, ABBA 20x20, wall: system allocator new/0.4.4 (paired) new/zune-jpeg 0.5.15 frame1080 0.863, 19/20 (z -4.02) 1.016 (zune 16/20) photo2000s 0.883, 20/20 (z -4.47) 0.978, 17/20 photo2000 0.945, 19/20 (z -4.02) 0.946, 19/20 rusty_alloc new/0.4.4 (paired) new/zune-jpeg 0.5.15 frame1080 0.964, 14/20 (z -1.79) 0.956, 17/20 photo2000s 0.986, 14/20 (z -1.79) 0.892, 20/20 photo2000 1.001, 10/20 (neutral) 0.932, 17/20 Co-Authored-By: Claude Opus 5.5 --- src/decode/decoder.rs | 22 ++++++- src/decode/mod.rs | 1 + src/decode/plane_pool.rs | 125 +++++++++++++++++++++++++++++++++++++ src/decode/worker/mod.rs | 6 +- src/decode/worker/rayon.rs | 4 +- 5 files changed, 150 insertions(+), 8 deletions(-) create mode 100644 src/decode/plane_pool.rs diff --git a/src/decode/decoder.rs b/src/decode/decoder.rs index af5044b..73c8b5b 100644 --- a/src/decode/decoder.rs +++ b/src/decode/decoder.rs @@ -947,7 +947,7 @@ impl Decoder { // Get the worker prepared let row_data = RowData { index: i, - recycled: self.plane_pool.pop(), + recycled: Self::take_plane(&mut self.plane_pool, component), component: component.clone(), quantization_table, }; @@ -1001,6 +1001,19 @@ impl Decoder { ) } + /// A plane to refill for `component`: the caller's (`recycle_planes`) + /// first, else one from this thread's pool (see `plane_pool`). + fn take_plane(pool: &mut Vec>, component: &Component) -> Option> { + if let Some(plane) = pool.pop() { + return Some(plane); + } + let needed = component.block_size.width as usize + * component.block_size.height as usize + * component.dct_scale + * component.dct_scale; + crate::decode::plane_pool::take(needed) + } + fn determine_color_transform(&self) -> ColorTransform { if let Some(color_transform) = self.color_transform { return color_transform; @@ -1157,7 +1170,7 @@ impl Decoder { if finished[i] { let row_data = RowData { index: i, - recycled: self.plane_pool.pop(), + recycled: Self::take_plane(&mut self.plane_pool, component), component: component.clone(), quantization_table: self.quantization_tables [component.quantization_table_index] @@ -1780,7 +1793,10 @@ fn compute_image( decoded.resize(size, 0); Ok(decoded) } else { - compute_image_parallel(components, data, output_size, color_transform, parallel) + let image = + compute_image_parallel(components, &data, output_size, color_transform, parallel)?; + crate::decode::plane_pool::give_back(data); + Ok(image) } } diff --git a/src/decode/mod.rs b/src/decode/mod.rs index a02e40f..ea0c9bd 100644 --- a/src/decode/mod.rs +++ b/src/decode/mod.rs @@ -20,6 +20,7 @@ mod huffman; pub(crate) mod idct; mod marker; mod parser; +mod plane_pool; mod upsampler; mod worker; diff --git a/src/decode/plane_pool.rs b/src/decode/plane_pool.rs new file mode 100644 index 0000000..6c7d918 --- /dev/null +++ b/src/decode/plane_pool.rs @@ -0,0 +1,125 @@ +//! Component planes kept for reuse by the next decode on the same thread. +//! +//! A `decode()` to RGB sizes one plane per component -- ~3.1 MB for a 1080p +//! 4:2:0 frame -- and drops them as soon as the image is assembled. Callers +//! that construct a fresh `Decoder` per image (the common shape: a thumbnailer, +//! an MJPEG frame loop) therefore paid for those pages to be faulted in again +//! on every decode. Under an allocator that recycles freed memory (rusty_alloc, +//! mimalloc) that was nearly free; under the Windows system heap a 1080p frame +//! paid +0.43 ms for it, enough to put the whole decode 1.17x behind zune-jpeg +//! on one core. Keeping the planes here makes the cost independent of the +//! allocator. +//! +//! The pool is bounded (see `MAX_PLANES` / `MAX_BYTES`): after an unusually +//! large image it keeps nothing rather than pinning that memory, and +//! `RUSTY_JPEG_PLANE_POOL=0` turns it off. Planes are refilled, never trusted: +//! the decoder overwrites every block before a plane is read (see +//! `ImmediateWorker::start_immediate`), so reuse cannot leak one image into the +//! next. Without `std` there is no thread-local storage and this is a no-op. + +use alloc::vec::Vec; + +/// At most this many planes are kept: one image's worth (Y, Cb, Cr, and a +/// fourth for CMYK/YCCK). +#[cfg(feature = "std")] +const MAX_PLANES: usize = 4; +/// ...and at most this many bytes of them (a 4K 4:4:4 frame is ~25 MB). +#[cfg(feature = "std")] +const MAX_BYTES: usize = 64 << 20; + +#[cfg(feature = "std")] +std::thread_local! { + static POOL: core::cell::RefCell>> = const { core::cell::RefCell::new(Vec::new()) }; +} + +#[cfg(feature = "std")] +fn enabled() -> bool { + crate::cached_knob!(bool, { + crate::knob("RUSTY_JPEG_PLANE_POOL").as_deref() != Some("0") + }) +} + +/// The pooled plane whose capacity fits `len` most tightly, if any. +pub(crate) fn take(len: usize) -> Option> { + #[cfg(feature = "std")] + { + if !enabled() { + return None; + } + POOL.try_with(|pool| { + let mut pool = pool.borrow_mut(); + let best = pool + .iter() + .enumerate() + .filter(|(_, plane)| plane.capacity() >= len) + .min_by_key(|(_, plane)| plane.capacity()) + .map(|(i, _)| i)?; + Some(pool.swap_remove(best)) + }) + .ok() + .flatten() + } + #[cfg(not(feature = "std"))] + { + let _ = len; + None + } +} + +/// Hand planes back after a decode. Keeps them only within the bounds. +pub(crate) fn give_back(planes: Vec>) { + #[cfg(feature = "std")] + { + if !enabled() { + return; + } + let _ = POOL.try_with(|pool| { + let mut pool = pool.borrow_mut(); + for plane in planes { + if plane.capacity() == 0 { + continue; + } + let held: usize = pool.iter().map(Vec::capacity).sum(); + if pool.len() >= MAX_PLANES || held + plane.capacity() > MAX_BYTES { + // Full, or this image was large: keep the newest set + // rather than growing. + pool.clear(); + if plane.capacity() > MAX_BYTES { + continue; + } + } + pool.push(plane); + } + }); + } + #[cfg(not(feature = "std"))] + { + let _ = planes; + } +} + +#[cfg(all(test, feature = "std"))] +mod tests { + use super::*; + + #[test] + fn best_fit_and_bounds() { + give_back(vec![ + Vec::with_capacity(100), + Vec::with_capacity(10), + Vec::with_capacity(50), + ]); + let a = take(40).unwrap(); + assert!( + a.capacity() >= 40 && a.capacity() < 100, + "best fit, got {}", + a.capacity() + ); + assert!(take(1000).is_none()); + // An oversized plane is never kept. + give_back(vec![Vec::with_capacity(MAX_BYTES + 1)]); + assert!(take(MAX_BYTES).is_none()); + // Drain. + while take(0).is_some() {} + } +} diff --git a/src/decode/worker/mod.rs b/src/decode/worker/mod.rs index 7ad475c..1c5fe60 100644 --- a/src/decode/worker/mod.rs +++ b/src/decode/worker/mod.rs @@ -249,7 +249,7 @@ const MAX_OUTPUT_THREADS: usize = 4; pub fn compute_image_parallel( components: &[Component], - data: Vec>, + data: &[Vec], output_size: Dimensions, color_transform: ColorTransform, parallel: bool, @@ -286,7 +286,7 @@ pub fn compute_image_parallel( bands, |rows, lines| { upsampler.upsample_and_interleave_rows_uninit( - &data, + data, rows, width, lines, @@ -314,7 +314,7 @@ pub fn compute_image_parallel( let mut image = vec![0u8; len]; in_bands(&mut image, line_size, bands, |rows, lines| { upsampler.upsample_and_interleave_rows( - &data, + data, rows, width, lines, diff --git a/src/decode/worker/rayon.rs b/src/decode/worker/rayon.rs index 8b06105..8a6d561 100644 --- a/src/decode/worker/rayon.rs +++ b/src/decode/worker/rayon.rs @@ -192,7 +192,7 @@ impl ComponentMetadata { pub fn compute_image_parallel( components: &[Component], - data: Vec>, + data: &[Vec], output_size: Dimensions, color_transform: ColorTransform, ) -> Result> { @@ -211,7 +211,7 @@ pub fn compute_image_parallel( .for_each(|(band, lines)| { let first = band * BAND; upsampler.upsample_and_interleave_rows( - &data, + data, first..first + lines.len() / line_size, output_size.width as usize, lines,