j2k_jpeg/
batch_session.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! Persistent JPEG batch decode session.
4
5use alloc::vec::Vec;
6use core::num::NonZeroUsize;
7use j2k_core::{BatchInfrastructureError, PixelFormat, TileBatchOptions};
8use std::sync::Mutex;
9
10use crate::decoder::{
11    planned_jpeg_tile_decode_live_bytes, DecodeOutcome, DecodedTile, PreparedJpegTileJob,
12    TileBatchError, TileDecodeJob, TileRegionScaledDecodeJob, TileScaledDecodeJob,
13};
14use crate::error::JpegError;
15use crate::info::DecodeOptions;
16
17mod allocation;
18mod collection;
19mod planning;
20mod runtime;
21mod scheduler;
22mod worker;
23
24use allocation::vec_capacity_bytes;
25use collection::{
26    collect_per_tile_results, collect_results, decode_outcome_retained_bytes,
27    decoded_tile_retained_bytes, retained_live_bytes,
28};
29use planning::{min_output_len, plan_per_tile_jobs, plan_regular_jobs, planned_job_chunk};
30use scheduler::{run_chunks_rayon, run_chunks_scoped};
31use worker::WorkerSlot;
32
33const SMALL_OUTPUT_DEFAULT_WORKER_CAP: usize = 4;
34const SMALL_OUTPUT_BYTES: usize = 32 * 1024;
35
36type BatchResultSlot<T> = j2k_core::BatchResultSlot<T, JpegError>;
37
38/// Reusable JPEG tile-batch runtime for WSI viewport loops.
39///
40/// The session keeps one decoder context and scratch pool per active worker
41/// during a batch and retains worker slots across calls. Before planning the
42/// next batch it may preserve one bounded decoder context, while evicting
43/// stale scratch and other contexts so the planning decoder retains the full
44/// codec memory allowance. Callers continue to own compressed inputs and
45/// decoded output buffers.
46#[derive(Debug)]
47pub struct JpegBatchSession {
48    options: TileBatchOptions,
49    workers: Vec<Mutex<WorkerSlot>>,
50    active_workers: usize,
51    cap_small_output_default_workers: bool,
52}
53
54impl Default for JpegBatchSession {
55    fn default() -> Self {
56        Self::new(TileBatchOptions::default())
57    }
58}
59
60impl JpegBatchSession {
61    /// Create a session using the given CPU batch worker options.
62    #[must_use]
63    pub fn new(options: TileBatchOptions) -> Self {
64        Self {
65            options,
66            workers: Vec::new(),
67            active_workers: 0,
68            cap_small_output_default_workers: false,
69        }
70    }
71
72    pub(crate) fn new_one_shot(options: TileBatchOptions) -> Self {
73        Self {
74            cap_small_output_default_workers: true,
75            ..Self::new(options)
76        }
77    }
78
79    /// Return current worker options.
80    #[must_use]
81    pub fn options(&self) -> TileBatchOptions {
82        self.options
83    }
84
85    /// Replace worker options for future decode calls.
86    pub fn set_options(&mut self, options: TileBatchOptions) {
87        self.options = options;
88    }
89
90    /// Number of active workers used by the most recent non-empty decode call.
91    #[must_use]
92    #[doc(hidden)]
93    pub fn worker_count(&self) -> usize {
94        self.active_workers
95    }
96
97    /// Number of worker slots retained by the session.
98    #[must_use]
99    #[doc(hidden)]
100    pub fn retained_worker_slots(&self) -> usize {
101        self.workers.len()
102    }
103
104    /// Clear worker-local decode contexts and scratch pools while retaining slots.
105    pub fn reset(&mut self) {
106        for slot in &self.workers {
107            match slot.lock() {
108                Ok(mut worker) => worker.reset(),
109                Err(poisoned) => {
110                    poisoned.into_inner().reset();
111                    slot.clear_poison();
112                }
113            }
114        }
115        self.active_workers = 0;
116    }
117
118    /// Decode full JPEG tiles into caller-owned output buffers.
119    ///
120    /// # Errors
121    /// Returns [`TileBatchError`] with the first codec failure in input order,
122    /// or a typed infrastructure failure when no tile index applies.
123    pub fn decode_tiles_into(
124        &mut self,
125        jobs: &mut [TileDecodeJob<'_, '_>],
126        fmt: PixelFormat,
127    ) -> Result<Vec<DecodeOutcome>, TileBatchError> {
128        self.decode_tiles_into_with_options(jobs, fmt, DecodeOptions::default())
129    }
130
131    /// Decode full JPEG tiles with explicit JPEG decode options.
132    ///
133    /// # Errors
134    /// Returns [`TileBatchError`] with the first codec failure in input order,
135    /// or a typed infrastructure failure when no tile index applies.
136    pub fn decode_tiles_into_with_options(
137        &mut self,
138        jobs: &mut [TileDecodeJob<'_, '_>],
139        fmt: PixelFormat,
140        decode_options: DecodeOptions,
141    ) -> Result<Vec<DecodeOutcome>, TileBatchError> {
142        if jobs.is_empty() {
143            return Ok(Vec::new());
144        }
145        let job_count = jobs.len();
146        let planning_metadata = self.prepare_job_planning(job_count)?;
147        let planning_context = self.planning_context()?;
148        let plans = plan_regular_jobs(jobs, planning_metadata, planning_context, |job, ctx| {
149            planned_jpeg_tile_decode_live_bytes(
150                job.input,
151                ctx,
152                fmt,
153                None,
154                j2k_core::Downscale::None,
155                decode_options,
156            )
157        })?;
158        let batch = self.prepare_batch::<DecodeOutcome, DecodeOutcome>(
159            &plans,
160            vec_capacity_bytes(&plans)?,
161            min_output_len(jobs),
162        )?;
163        let planned_jobs = &plans;
164        let decode_chunk =
165            |worker: &mut WorkerSlot,
166             start_index: usize,
167             chunk: &mut [_],
168             results: &mut [BatchResultSlot<DecodeOutcome>]| {
169                let chunk_plans = planned_job_chunk(planned_jobs, start_index, chunk.len())?;
170                worker.decode_tile_job_chunk(chunk, results, chunk_plans, fmt, decode_options)
171            };
172        let results = if self.options.workers.is_some() {
173            run_chunks_scoped(&self.workers, jobs, batch, decode_chunk)?
174        } else {
175            run_chunks_rayon(&self.workers, jobs, batch, decode_chunk)?
176        };
177        let retained = retained_live_bytes(
178            &self.workers,
179            vec_capacity_bytes(&self.workers)?,
180            vec_capacity_bytes(&plans)?,
181            &results,
182            decode_outcome_retained_bytes,
183        )?;
184        collect_results(job_count, results, retained)
185    }
186
187    /// Decode prepared TIFF/WSI JPEG tiles into caller-owned RGB8 buffers.
188    ///
189    /// Results preserve the caller's input order and retain each tile's error
190    /// independently instead of collapsing the batch to the first failure.
191    ///
192    /// # Errors
193    ///
194    /// Returns a typed batch infrastructure error for planning, allocation,
195    /// scheduling, or ordered-collection failures.
196    pub fn decode_prepared_jpeg_tiles_rgb8(
197        &mut self,
198        jobs: &mut [PreparedJpegTileJob<'_, '_>],
199    ) -> Result<Vec<Result<DecodedTile, JpegError>>, BatchInfrastructureError> {
200        if jobs.is_empty() {
201            return Ok(Vec::new());
202        }
203        let job_count = jobs.len();
204        let planning_metadata = self.prepare_job_planning(job_count)?;
205        let planning_context = self.planning_context()?;
206        let plans = plan_per_tile_jobs(jobs, planning_metadata, planning_context, |job, ctx| {
207            planned_jpeg_tile_decode_live_bytes(
208                job.input.as_bytes(),
209                ctx,
210                PixelFormat::Rgb8,
211                None,
212                j2k_core::Downscale::None,
213                job.options,
214            )
215        })?;
216        let batch = self.prepare_batch::<DecodedTile, Result<DecodedTile, JpegError>>(
217            &plans,
218            vec_capacity_bytes(&plans)?,
219            min_output_len(jobs),
220        )?;
221        let planned_jobs = &plans;
222        let decode_chunk =
223            |worker: &mut WorkerSlot,
224             start_index: usize,
225             chunk: &mut [_],
226             results: &mut [BatchResultSlot<DecodedTile>]| {
227                let chunk_plans = planned_job_chunk(planned_jobs, start_index, chunk.len())?;
228                worker.decode_prepared_tile_job_chunk(chunk, results, chunk_plans)
229            };
230        let results = if self.options.workers.is_some() {
231            run_chunks_scoped(&self.workers, jobs, batch, decode_chunk)?
232        } else {
233            run_chunks_rayon(&self.workers, jobs, batch, decode_chunk)?
234        };
235        let retained = retained_live_bytes(
236            &self.workers,
237            vec_capacity_bytes(&self.workers)?,
238            vec_capacity_bytes(&plans)?,
239            &results,
240            decoded_tile_retained_bytes,
241        )?;
242        collect_per_tile_results(job_count, results, retained)
243    }
244
245    /// Decode scaled JPEG tiles into caller-owned output buffers.
246    ///
247    /// # Errors
248    /// Returns [`TileBatchError`] with the first codec failure in input order,
249    /// or a typed infrastructure failure when no tile index applies.
250    pub fn decode_tiles_scaled_into(
251        &mut self,
252        jobs: &mut [TileScaledDecodeJob<'_, '_>],
253        fmt: PixelFormat,
254    ) -> Result<Vec<DecodeOutcome>, TileBatchError> {
255        self.decode_tiles_scaled_into_with_options(jobs, fmt, DecodeOptions::default())
256    }
257
258    /// Decode scaled JPEG tiles with explicit JPEG decode options.
259    ///
260    /// # Errors
261    /// Returns [`TileBatchError`] with the first codec failure in input order,
262    /// or a typed infrastructure failure when no tile index applies.
263    pub fn decode_tiles_scaled_into_with_options(
264        &mut self,
265        jobs: &mut [TileScaledDecodeJob<'_, '_>],
266        fmt: PixelFormat,
267        decode_options: DecodeOptions,
268    ) -> Result<Vec<DecodeOutcome>, TileBatchError> {
269        if jobs.is_empty() {
270            return Ok(Vec::new());
271        }
272        let job_count = jobs.len();
273        let planning_metadata = self.prepare_job_planning(job_count)?;
274        let planning_context = self.planning_context()?;
275        let plans = plan_regular_jobs(jobs, planning_metadata, planning_context, |job, ctx| {
276            planned_jpeg_tile_decode_live_bytes(
277                job.input,
278                ctx,
279                fmt,
280                None,
281                job.scale,
282                decode_options,
283            )
284        })?;
285        let batch = self.prepare_batch::<DecodeOutcome, DecodeOutcome>(
286            &plans,
287            vec_capacity_bytes(&plans)?,
288            min_output_len(jobs),
289        )?;
290        let planned_jobs = &plans;
291        let decode_chunk =
292            |worker: &mut WorkerSlot,
293             start_index: usize,
294             chunk: &mut [_],
295             results: &mut [BatchResultSlot<DecodeOutcome>]| {
296                let chunk_plans = planned_job_chunk(planned_jobs, start_index, chunk.len())?;
297                worker.decode_tile_scaled_job_chunk(
298                    chunk,
299                    results,
300                    chunk_plans,
301                    fmt,
302                    decode_options,
303                )
304            };
305        let results = if self.options.workers.is_some() {
306            run_chunks_scoped(&self.workers, jobs, batch, decode_chunk)?
307        } else {
308            run_chunks_rayon(&self.workers, jobs, batch, decode_chunk)?
309        };
310        let retained = retained_live_bytes(
311            &self.workers,
312            vec_capacity_bytes(&self.workers)?,
313            vec_capacity_bytes(&plans)?,
314            &results,
315            decode_outcome_retained_bytes,
316        )?;
317        collect_results(job_count, results, retained)
318    }
319
320    /// Decode scaled JPEG tile regions into caller-owned output buffers.
321    ///
322    /// # Errors
323    /// Returns [`TileBatchError`] with the first codec failure in input order,
324    /// or a typed infrastructure failure when no tile index applies.
325    pub fn decode_tiles_region_scaled_into(
326        &mut self,
327        jobs: &mut [TileRegionScaledDecodeJob<'_, '_>],
328        fmt: PixelFormat,
329    ) -> Result<Vec<DecodeOutcome>, TileBatchError> {
330        self.decode_tiles_region_scaled_into_with_options(jobs, fmt, DecodeOptions::default())
331    }
332
333    /// Decode scaled JPEG tile regions with explicit JPEG decode options.
334    ///
335    /// # Errors
336    /// Returns [`TileBatchError`] with the first codec failure in input order,
337    /// or a typed infrastructure failure when no tile index applies.
338    pub fn decode_tiles_region_scaled_into_with_options(
339        &mut self,
340        jobs: &mut [TileRegionScaledDecodeJob<'_, '_>],
341        fmt: PixelFormat,
342        decode_options: DecodeOptions,
343    ) -> Result<Vec<DecodeOutcome>, TileBatchError> {
344        if jobs.is_empty() {
345            return Ok(Vec::new());
346        }
347        let job_count = jobs.len();
348        let planning_metadata = self.prepare_job_planning(job_count)?;
349        let planning_context = self.planning_context()?;
350        let plans = plan_regular_jobs(jobs, planning_metadata, planning_context, |job, ctx| {
351            planned_jpeg_tile_decode_live_bytes(
352                job.input,
353                ctx,
354                fmt,
355                Some(job.roi.into()),
356                job.scale,
357                decode_options,
358            )
359        })?;
360        let batch = self.prepare_batch::<DecodeOutcome, DecodeOutcome>(
361            &plans,
362            vec_capacity_bytes(&plans)?,
363            min_output_len(jobs),
364        )?;
365        let planned_jobs = &plans;
366        let decode_chunk =
367            |worker: &mut WorkerSlot,
368             start_index: usize,
369             chunk: &mut [_],
370             results: &mut [BatchResultSlot<DecodeOutcome>]| {
371                let chunk_plans = planned_job_chunk(planned_jobs, start_index, chunk.len())?;
372                worker.decode_tile_region_scaled_job_chunk(
373                    chunk,
374                    results,
375                    chunk_plans,
376                    fmt,
377                    decode_options,
378                )
379            };
380        let results = if self.options.workers.is_some() {
381            run_chunks_scoped(&self.workers, jobs, batch, decode_chunk)?
382        } else {
383            run_chunks_rayon(&self.workers, jobs, batch, decode_chunk)?
384        };
385        let retained = retained_live_bytes(
386            &self.workers,
387            vec_capacity_bytes(&self.workers)?,
388            vec_capacity_bytes(&plans)?,
389            &results,
390            decode_outcome_retained_bytes,
391        )?;
392        collect_results(job_count, results, retained)
393    }
394}
395
396fn available_tile_batch_workers() -> usize {
397    std::thread::available_parallelism().map_or(1, NonZeroUsize::get)
398}
399
400#[cfg(test)]
401mod tests;
j2k_jpeg/batch_session.rs

j2k_jpeg/
batch_session.rs
