diff --git a/crates/codecs/src/codecs/tiff/bands.rs b/crates/codecs/src/codecs/tiff/bands.rs --- a/crates/codecs/src/codecs/tiff/bands.rs +++ b/crates/codecs/src/codecs/tiff/bands.rs @@ -110,16 +110,50 @@ impl std::ops::Add for Tally { } } +/// Bytes an open `tiff` decoder keeps per chunk: its offset and its byte count, as `u64`s. +pub(super) const DECODER_BYTES_PER_CHUNK: u64 = 16; + +/// Threads the band decode may run on. +pub(super) fn worker_threads() -> usize { + #[cfg(not(target_arch = "wasm32"))] + return rayon::current_num_threads().max(1); + #[cfg(target_arch = "wasm32")] + return 1; +} + +/// How many band workers, each holding `per_worker` bytes, fit in `max_alloc` beside the `held` +/// bytes already allocated (at most `threads`, at least one), or a limit error when none does. +pub(super) fn budget_workers(held: u64, per_worker: u64, max_alloc: u64, threads: usize) -> Result { + let threads = threads.max(1); + let room = max_alloc.saturating_sub(held); + let fit = room.checked_div(per_worker).map_or(threads, |n| usize::try_from(n).unwrap_or(usize::MAX)); + if fit == 0 { + return Err(CodecError::LimitExceeded(format!( + "decoding needs {per_worker} bytes of strip or tile scratch beside {held} bytes of image, exceeding max_alloc {max_alloc}" + ))); + } + Ok(fit.min(threads)) +} + +/// Bands each parallel task decodes, one after the other: one band per task when every thread +/// may hold a worker, else the bands split into `workers` runs so that no more than `workers` +/// decoders and scratch buffers are ever alive at once. +fn bands_per_task(bands: usize, workers: usize, threads: usize) -> usize { + if workers >= threads { 1 } else { bands.div_ceil(workers.max(1)).max(1) } +} + /// Decodes every band into `data`, in parallel on native targets (one decoder and scratch -/// buffer per worker). `extra` receives samples that are not part of the colour image, one -/// band of rows at a time, when it is non-empty. +/// buffer per worker, at most `workers` of them at once). `extra` receives samples that are not +/// part of the colour image, one band of rows at a time, when it is non-empty. #[cfg(not(target_arch = "wasm32"))] +#[allow(clippy::too_many_arguments)] pub(super) fn decode_bands<'a, O>( lay: &Layout, data: &mut [u8], band_bytes: usize, extra: &mut [u8], extra_band: usize, + workers: usize, _first: &mut Decoder>, open: &O, ) -> Result @@ -127,40 +161,40 @@ where O: Fn() -> Result>, CodecError> + Sync, { use rayon::prelude::*; + let band_bytes = band_bytes.max(1); + let per_task = bands_per_task(data.len().div_ceil(band_bytes), workers, worker_threads()); + let task_bytes = band_bytes.saturating_mul(per_task); + let init = || open().map(|d| (d, Vec::new())).map_err(|e| e.to_string()); if extra.is_empty() { return data - .par_chunks_mut(band_bytes.max(1)) + .par_chunks_mut(task_bytes) .enumerate() - .map_init( - || open().map(|d| (d, Vec::new())).map_err(|e| e.to_string()), - |worker, (b, band)| { - let (dec, scratch) = worker.as_mut().map_err(|e| err(e.clone()))?; - decode_band(lay, dec, scratch, b, band, None) - }, - ) + .map_init(init, |worker, (t, task)| { + let (dec, scratch) = worker.as_mut().map_err(|e| err(e.clone()))?; + decode_task(lay, dec, scratch, t.saturating_mul(per_task), task, band_bytes, None, extra_band) + }) .try_reduce(Tally::default, |a, b| Ok(a + b)); } - data.par_chunks_mut(band_bytes.max(1)) - .zip(extra.par_chunks_mut(extra_band.max(1))) + data.par_chunks_mut(task_bytes) + .zip(extra.par_chunks_mut(extra_band.max(1).saturating_mul(per_task))) .enumerate() - .map_init( - || open().map(|d| (d, Vec::new())).map_err(|e| e.to_string()), - |worker, (b, (band, extra_band))| { - let (dec, scratch) = worker.as_mut().map_err(|e| err(e.clone()))?; - decode_band(lay, dec, scratch, b, band, Some(extra_band)) - }, - ) + .map_init(init, |worker, (t, (task, extra_task))| { + let (dec, scratch) = worker.as_mut().map_err(|e| err(e.clone()))?; + decode_task(lay, dec, scratch, t.saturating_mul(per_task), task, band_bytes, Some(extra_task), extra_band) + }) .try_reduce(Tally::default, |a, b| Ok(a + b)) } /// Decodes every band into `data`, one after the other. #[cfg(target_arch = "wasm32")] +#[allow(clippy::too_many_arguments)] pub(super) fn decode_bands<'a, O>( lay: &Layout, data: &mut [u8], band_bytes: usize, extra: &mut [u8], extra_band: usize, + _workers: usize, first: &mut Decoder>, _open: &O, ) -> Result @@ -168,15 +202,30 @@ where O: Fn() -> Result>, CodecError> + Sync, { let mut scratch = Vec::new(); + let extra = if extra.is_empty() { None } else { Some(extra) }; + decode_task(lay, first, &mut scratch, 0, data, band_bytes.max(1), extra, extra_band) +} + +/// Decodes the consecutive bands in `task` (the first is band `b0`), one after the other. +#[allow(clippy::too_many_arguments)] +fn decode_task( + lay: &Layout, + dec: &mut Decoder>, + scratch: &mut Vec, + b0: usize, + task: &mut [u8], + band_bytes: usize, + extra: Option<&mut [u8]>, + extra_band: usize, +) -> Result { let mut tally = Tally::default(); - if extra.is_empty() { - for (b, band) in data.chunks_mut(band_bytes.max(1)).enumerate() { - tally = tally + decode_band(lay, first, &mut scratch, b, band, None)?; - } - return Ok(tally); - } - for (b, (band, extra_band)) in data.chunks_mut(band_bytes.max(1)).zip(extra.chunks_mut(extra_band.max(1))).enumerate() { - tally = tally + decode_band(lay, first, &mut scratch, b, band, Some(extra_band))?; + let mut extra = extra.map(|e| e.chunks_mut(extra_band.max(1))); + for (i, band) in task.chunks_mut(band_bytes.max(1)).enumerate() { + let extra_band = match extra.as_mut() { + Some(bands) => Some(bands.next().ok_or_else(|| err("extra-sample band out of range"))?), + None => None, + }; + tally = tally + decode_band(lay, dec, scratch, b0.saturating_add(i), band, extra_band)?; } Ok(tally) } @@ -485,3 +534,36 @@ fn finish_band(lay: &Layout, band: &mut [u8]) { } } } + +#[cfg(test)] +mod scratch_budget { + use super::{CodecError, bands_per_task, budget_workers}; + + #[test] + fn workers_are_capped_by_the_room_beside_the_image() { + // NV-9: four threads, a 128 MiB strip scratch each, 8 MiB of image, 160 MiB of budget: + // one worker fits, not four. + let mib = 1u64 << 20; + assert_eq!(budget_workers(8 * mib, 128 * mib, 160 * mib, 4).unwrap(), 1); + assert_eq!(budget_workers(8 * mib, 32 * mib, 160 * mib, 4).unwrap(), 4); + assert_eq!(budget_workers(8 * mib, 32 * mib, 104 * mib, 16).unwrap(), 3); + assert!(matches!(budget_workers(8 * mib, 160 * mib, 160 * mib, 4), Err(CodecError::LimitExceeded(_)))); + assert!(matches!(budget_workers(200 * mib, 1, 160 * mib, 4), Err(CodecError::LimitExceeded(_)))); + // No scratch: every thread; no threads reported: still one worker. + assert_eq!(budget_workers(8 * mib, 0, 160 * mib, 4).unwrap(), 4); + assert_eq!(budget_workers(0, 1, u64::MAX, 0).unwrap(), 1); + } + + #[test] + fn a_short_budget_groups_bands_into_as_many_tasks_as_workers() { + assert_eq!(bands_per_task(16, 4, 4), 1); + assert_eq!(bands_per_task(16, 8, 4), 1); + assert_eq!(bands_per_task(16, 1, 4), 16); + assert_eq!(bands_per_task(17, 3, 4), 6); + assert_eq!(bands_per_task(0, 0, 4), 1); + // Never more tasks than workers. + for (bands, workers) in [(1usize, 1usize), (7, 2), (100, 3), (5, 5)] { + assert!(bands.div_ceil(bands_per_task(bands, workers, 8)) <= workers); + } + } +} diff --git a/crates/codecs/src/codecs/tiff/decode.rs b/crates/codecs/src/codecs/tiff/decode.rs --- a/crates/codecs/src/codecs/tiff/decode.rs +++ b/crates/codecs/src/codecs/tiff/decode.rs @@ -1,4 +1,4 @@ -use super::bands::{convert_ab, decode_bands, extra_sample_indices, split_extra_samples, zeroed}; +use super::bands::{DECODER_BYTES_PER_CHUNK, budget_workers, convert_ab, decode_bands, extra_sample_indices, split_extra_samples, worker_threads, zeroed}; use super::tags::{F, TAG_IMAGE_SOURCE_DATA, TAG_PHOTOSHOP, TAG_XMP, TEXT_TAGS, photoshop_tags}; use super::*; @@ -443,7 +443,8 @@ fn decode_ifd(f: &File<'_>, dir: &Ifd, limits: &Limits) -> Result 4, 32773 => 130, _ => 4096, @@ -491,29 +492,57 @@ fn decode_ifd(f: &File<'_>, dir: &Ifd, limits: &Limits) -> Result 1 { + let plane1 = u32::try_from(chunks_per_plane).map_err(|_| err("too many tiles"))?; + largest = largest.max(dec.image_chunk_buffer_layout(plane1).map_err(map_err)?.len); + } + let in_place = lay.direct() && first.len == band_bytes && first.row_stride.map(|s| s.get()) == Some(out_row); + let scratch = if in_place { 0 } else { largest }; + if scratch > lay.scratch_cap { + return Err(CodecError::LimitExceeded(format!("a {scratch} byte strip or tile"))); + } + lay.scratch_cap = scratch; + let tables = u64::from(total).saturating_mul(DECODER_BYTES_PER_CHUNK); + let jpeg_copy = if matches!(compression, 6 | 7) { largest as u64 } else { 0 }; + let per_worker = (scratch as u64).saturating_add(jpeg_copy).saturating_add(tables); + let len = out_row.checked_mul(hu).ok_or_else(|| CodecError::LimitExceeded("image size".into()))?; + // Extra samples share the allocation budget with the image: their interleaved buffer and the - // planes split from it are held next to the image, so both count against `max_alloc`. When - // they don't fit, the image still opens (it passed its own check above) without them. + // planes split from it are held next to the image, and the interleaved buffer next to the + // decode's working memory, so all of it counts against `max_alloc`. When they don't fit, the + // image still opens (if it fits on its own) without them. let mut extras_dropped = 0; if !lay.extra_at.is_empty() { let per_pixel = out_layout.channels().saturating_add(lay.extra_at.len().saturating_mul(2)).saturating_mul(out.bytes()); - if limits.check_bytes(w, h, per_pixel as u64).is_err() { + let interleaved = (lay.extra_row() as u64).saturating_mul(hu as u64); + let decoding = (len as u64).saturating_add(interleaved).saturating_add(tables).saturating_add(per_worker); + if limits.check_bytes(w, h, per_pixel as u64).is_err() || decoding > limits.max_alloc { extras_dropped = u32::try_from(lay.extra_at.len()).unwrap_or(u32::MAX); lay.extra_at.clear(); } } + let extra_row = lay.extra_row(); + let extra_len = extra_row.checked_mul(hu).ok_or_else(|| CodecError::LimitExceeded("extra samples".into()))?; + let extra_band = extra_row.saturating_mul(ch.min(hu)); + // The output, the extra samples and the first decoder (kept open next to the workers). + let held = (len as u64).saturating_add(extra_len as u64).saturating_add(tables); + let workers = budget_workers(held, per_worker, limits.max_alloc, worker_threads())?; // The output, allocated once and fallibly. - let out_row = lay.out_row(); - let len = out_row.checked_mul(hu).ok_or_else(|| CodecError::LimitExceeded("image size".into()))?; let mut data = zeroed(len)?; - let band_bytes = out_row.checked_mul(ch.min(hu)).ok_or_else(|| CodecError::LimitExceeded("band size".into()))?; - - let extra_row = lay.extra_row(); - let extra_len = extra_row.checked_mul(hu).ok_or_else(|| CodecError::LimitExceeded("extra samples".into()))?; let mut extra_data = if lay.extra_at.is_empty() { Vec::new() } else { zeroed(extra_len)? }; - let extra_band = extra_row.saturating_mul(ch.min(hu)); - let tally = decode_bands(&lay, &mut data, band_bytes, extra_data.as_mut_slice(), extra_band, &mut dec, &open)?; + let tally = decode_bands(&lay, &mut data, band_bytes, extra_data.as_mut_slice(), extra_band, workers, &mut dec, &open)?; if tally.decoded == 0 && tally.missing > 0 { return Err(err("the image data is missing (cut off or past the end of the file)")); } diff --git a/crates/codecs/src/options.rs b/crates/codecs/src/options.rs --- a/crates/codecs/src/options.rs +++ b/crates/codecs/src/options.rs @@ -16,7 +16,10 @@ pub struct Limits { /// PNG also applies this to the aggregate decoded text/XMP UTF-8 bytes, /// including keywords, separately from pixels. Animated WebP charges /// every retained frame's full pixel copy plus its own copy of the - /// file's ICC profile against this budget. This is not a total RSS cap. + /// file's ICC profile against this budget. TIFF charges its decode working memory (each + /// band worker's strip or tile scratch buffer and decoder state) next to the output, runs + /// only as many workers at once as fit, and refuses the image when not even one does. + /// This is not a total RSS cap. pub max_alloc: u64, } diff --git a/crates/codecs/tests/it/tiff_pages.rs b/crates/codecs/tests/it/tiff_pages.rs --- a/crates/codecs/tests/it/tiff_pages.rs +++ b/crates/codecs/tests/it/tiff_pages.rs @@ -659,3 +659,41 @@ fn zoo_long_random_mutations() { poke(&g); } } + +/// A 64-bit gray page whose strips decode through a scratch buffer four times the size of their +/// output rows: every sample of row `y` is `y << 48` (it reads back as the 16-bit value `y`). +fn wide_gray(w: u32, h: u32, rows: u32) -> Vec { + let px: Vec = (0..h).flat_map(|y| (0..w).flat_map(move |_| (u64::from(y) << 48).to_le_bytes())).collect(); + build(true, false, &[strips(w, h, 64, 1, 1, rows, &px)], &[0]).bytes +} + +fn with_max_alloc(max_alloc: u64) -> DecodeOptions { + DecodeOptions { limits: Limits { max_alloc, ..Limits::default() }, ..Default::default() } +} + +#[test] +fn strip_scratch_counts_against_max_alloc() { + // 64×64: an 8 KiB 16-bit output, but its one strip is 32 KiB of 64-bit samples, decoded + // into a scratch buffer held next to the output. The two together are over 36 KiB. + let file = wide_gray(64, 64, 64); + let r = decode_with(&file, &with_max_alloc(36 << 10)); + assert!(matches!(r, Err(CodecError::LimitExceeded(_))), "output plus strip scratch over max_alloc: {:?}", r.map(|i| i.data().len())); + // With room for both it decodes. + let img = decode_with(&file, &with_max_alloc(48 << 10)).unwrap(); + assert_eq!(img.to_u16_samples().unwrap()[64 * 63], 63); +} + +#[test] +fn a_tight_budget_decodes_its_strips_one_after_another() { + // 64×512 in eight 32 KiB strips: a 64 KiB output with room for only one strip's scratch + // beside it, so the bands are decoded by one worker, in order, and still land in place. + let file = wide_gray(64, 512, 64); + let img = decode_with(&file, &with_max_alloc((64 + 40) << 10)).unwrap(); + let v = img.to_u16_samples().unwrap(); + for y in 0..512 { + assert_eq!(v[y * 64], y as u16, "row {y}"); + assert_eq!(v[y * 64 + 63], y as u16, "row {y}"); + } + // Unlimited, the same file decodes the same in parallel. + assert_eq!(decode(&file).unwrap().to_u16_samples().unwrap(), v); +} diff --git a/crates/codecs/tests/it/tiff_photoshop.rs b/crates/codecs/tests/it/tiff_photoshop.rs --- a/crates/codecs/tests/it/tiff_photoshop.rs +++ b/crates/codecs/tests/it/tiff_photoshop.rs @@ -218,12 +218,13 @@ fn a_wrong_sized_plane_is_refused_before_a_sample_type_conversion() { } /// Extra samples count against `max_alloc` with the image (their interleaved buffer and the -/// planes split from it): a budget that holds the RGB image but not sixteen extra planes opens -/// the image without them, with a warning, instead of allocating past it. +/// planes split from it): a budget that holds the RGB image and one strip's decode scratch but not +/// sixteen extra planes opens the image without them, with a warning, instead of allocating past +/// it. #[test] fn extra_samples_count_against_the_allocation_limit() { - // Big enough for several strips (about 1 MB each): one strip of all 19 samples fits the budget - // of the image alone, as in real files. + // Big enough for several strips (about 1 MB each). Each strip of all 19 samples is decoded + // through a scratch buffer that counts against the budget too: 1 MiB more than the image. let (w, h) = (1024u32, 1024u32); let n = (w * h) as usize; let mut img = synth(w, h, ChannelLayout::Rgb, SampleType::U8, 1, 0.0); @@ -232,18 +233,21 @@ fn extra_samples_count_against_the_allocation_limit() { let with = |max_alloc: u64| decode_with(&bytes, &DecodeOptions { limits: Limits { max_alloc, ..Limits::default() }, ..DecodeOptions::default() }); // RGB is 3 bytes a pixel; the extra samples add 2 x 16 more. let image_only = 3 * n as u64; + let strip_room = 1u64 << 20; let with_extras = (3 + 2 * MAX_EXTRA_SAMPLES as u64) * n as u64; let back = with(with_extras).expect("everything fits"); assert_eq!(back.meta.extra_samples.len(), MAX_EXTRA_SAMPLES); assert!(back.warnings.is_empty()); // The image fits, its extra samples don't: it opens as before they were read, with a warning. - for max_alloc in [image_only, with_extras - 1] { + for max_alloc in [image_only + strip_room, with_extras - 1] { let only = with(max_alloc).expect("the image alone fits"); assert!(only.meta.extra_samples.is_empty()); assert_eq!(only.warnings, vec![DecodeWarning::ExtraSamplesDropped { format: Format::Tiff, count: MAX_EXTRA_SAMPLES as u32 }]); assert_eq!(only.data(), back.data()); } assert!(matches!(with(image_only - 1), Err(CodecError::LimitExceeded(_))), "the image itself must still fit"); + // The image fits but a strip's scratch beside it doesn't: refused, not decoded past the budget. + assert!(matches!(with(image_only), Err(CodecError::LimitExceeded(_))), "the strip scratch counts too"); } /// The value field of a little-endian classic TIFF's ExtraSamples entry (tag 338, SHORT) holding