pdf_engine/
batch.rs

1//! Batch helpers for processing many PDFs with a bounded worker pool.
2
3// Copyright (c) 2026 Innovation Trigger B.V.
4//
5// PDFluent is available under two licences, at your option: the GNU AGPLv3, or
6// the PDFluent Commercial Licence. See the LICENSE file in this repository --
7// that file travels with the copy you received, which a URL does not.
8
9use crate::{PdfDocument, Result};
10use std::fs;
11use std::path::{Path, PathBuf};
12use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
13use std::sync::{mpsc, Arc};
14use std::thread;
15use std::time::{Duration, Instant};
16
17/// How batch processing should classify per-file failures.
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
19pub enum ErrorStrategy {
20    /// Collect failures in [`BatchResult::failures`] and continue (default).
21    #[default]
22    Collect,
23    /// Record failures in [`BatchResult::skipped`] and continue.
24    Skip,
25    /// Stop scheduling new work after the first failure.
26    FailFast,
27}
28
29/// Configuration for [`process_batch`] and [`PdfBatch`] helpers.
30pub struct BatchConfig {
31    /// Number of worker threads to use.
32    pub workers: usize,
33    /// Maximum time allowed per input PDF.
34    pub timeout: Duration,
35    /// How per-file failures should be recorded.
36    pub on_error: ErrorStrategy,
37    /// Optional progress callback invoked as `(done, total)`.
38    pub on_progress: Option<Box<dyn Fn(usize, usize) + Send + Sync + 'static>>,
39}
40
41impl Default for BatchConfig {
42    fn default() -> Self {
43        Self {
44            workers: thread::available_parallelism()
45                .map(|n| n.get())
46                .unwrap_or(1),
47            timeout: Duration::from_secs(30),
48            on_error: ErrorStrategy::Collect,
49            on_progress: None,
50        }
51    }
52}
53
54/// Aggregated outcome of a batch run.
55#[derive(Debug)]
56pub struct BatchResult<T> {
57    /// Successfully processed PDFs and their outputs.
58    pub successes: Vec<(PathBuf, T)>,
59    /// PDFs that failed to process.
60    pub failures: Vec<(PathBuf, String)>,
61    /// PDFs skipped due to [`ErrorStrategy::Skip`].
62    pub skipped: Vec<(PathBuf, String)>,
63    /// End-to-end runtime for the batch.
64    pub duration: Duration,
65}
66
67impl<T> Default for BatchResult<T> {
68    fn default() -> Self {
69        Self {
70            successes: Vec::new(),
71            failures: Vec::new(),
72            skipped: Vec::new(),
73            duration: Duration::default(),
74        }
75    }
76}
77
78impl<T> BatchResult<T> {
79    /// Total number of completed inputs.
80    pub fn total(&self) -> usize {
81        self.successes.len() + self.failures.len() + self.skipped.len()
82    }
83}
84
85/// A progress callback shared with the worker threads.
86type SharedProgress = Arc<Box<dyn Fn(usize, usize) + Send + Sync + 'static>>;
87
88/// Convenience helpers built on top of [`process_batch`].
89#[derive(Debug, Default, Clone, Copy)]
90pub struct PdfBatch;
91
92impl PdfBatch {
93    /// Recursively collect `.pdf` files from a directory tree.
94    pub fn collect_pdfs(root: impl AsRef<Path>) -> std::io::Result<Vec<PathBuf>> {
95        let mut paths = Vec::new();
96        collect_pdfs_recursive(root.as_ref(), &mut paths)?;
97        paths.sort();
98        Ok(paths)
99    }
100
101    /// Extract combined text from every PDF under `root`.
102    pub fn extract_text(root: impl AsRef<Path>, config: BatchConfig) -> BatchResult<String> {
103        let root = root.as_ref().to_path_buf();
104        match Self::collect_pdfs(&root) {
105            Ok(paths) => process_batch(&paths, config, |doc| Ok(doc.extract_all_text())),
106            Err(err) => {
107                let mut result = BatchResult::default();
108                result.failures.push((root, err.to_string()));
109                result
110            }
111        }
112    }
113
114    /// Count pages for every PDF under `root`.
115    pub fn page_counts(root: impl AsRef<Path>, config: BatchConfig) -> BatchResult<usize> {
116        let root = root.as_ref().to_path_buf();
117        match Self::collect_pdfs(&root) {
118            Ok(paths) => process_batch(&paths, config, |doc| Ok(doc.page_count())),
119            Err(err) => {
120                let mut result = BatchResult::default();
121                result.failures.push((root, err.to_string()));
122                result
123            }
124        }
125    }
126}
127
128/// Process a list of PDF paths with a fixed-size worker pool.
129pub fn process_batch<P, F, T>(paths: &[P], config: BatchConfig, processor: F) -> BatchResult<T>
130where
131    P: AsRef<Path>,
132    F: Fn(&PdfDocument) -> Result<T> + Send + Sync + 'static,
133    T: Send + 'static,
134{
135    run_batch(paths, config, None, processor)
136}
137
138/// [`process_batch`] with a per-page progress callback, invoked as
139/// `(done, total)` with the current file's page counts: a three-page file
140/// reports `(1, 3)` through `(3, 3)` once it has been read. It fires only for
141/// files that open and process successfully, from the worker thread; with
142/// several workers the files interleave. The per-file callback in
143/// [`BatchConfig::on_progress`] still fires as before.
144///
145/// A separate function rather than a [`BatchConfig`] field: `BatchConfig` is
146/// built with struct literals, so a new public field would break every such
147/// literal written against 1.0.
148pub fn process_batch_with_page_progress<P, F, T>(
149    paths: &[P],
150    config: BatchConfig,
151    on_page_progress: Box<dyn Fn(usize, usize) + Send + Sync + 'static>,
152    processor: F,
153) -> BatchResult<T>
154where
155    P: AsRef<Path>,
156    F: Fn(&PdfDocument) -> Result<T> + Send + Sync + 'static,
157    T: Send + 'static,
158{
159    run_batch(paths, config, Some(on_page_progress), processor)
160}
161
162fn run_batch<P, F, T>(
163    paths: &[P],
164    config: BatchConfig,
165    on_page_progress: Option<Box<dyn Fn(usize, usize) + Send + Sync + 'static>>,
166    processor: F,
167) -> BatchResult<T>
168where
169    P: AsRef<Path>,
170    F: Fn(&PdfDocument) -> Result<T> + Send + Sync + 'static,
171    T: Send + 'static,
172{
173    let started = Instant::now();
174    let input_paths = paths
175        .iter()
176        .map(|path| path.as_ref().to_path_buf())
177        .collect::<Vec<_>>();
178    let total = input_paths.len();
179
180    if total == 0 {
181        return BatchResult {
182            duration: started.elapsed(),
183            ..Default::default()
184        };
185    }
186
187    let workers = config.workers.max(1).min(total);
188    let next_index = Arc::new(AtomicUsize::new(0));
189    let stop = Arc::new(AtomicBool::new(false));
190    let input_paths = Arc::new(input_paths);
191    let processor = Arc::new(processor);
192    let on_page = on_page_progress.map(Arc::new);
193    let (tx, rx) = mpsc::channel();
194    let mut handles = Vec::with_capacity(workers);
195
196    for _ in 0..workers {
197        let tx = tx.clone();
198        let next_index = Arc::clone(&next_index);
199        let stop = Arc::clone(&stop);
200        let input_paths = Arc::clone(&input_paths);
201        let processor = Arc::clone(&processor);
202        let on_page = on_page.clone();
203        let timeout = config.timeout;
204
205        handles.push(thread::spawn(move || loop {
206            if stop.load(Ordering::Relaxed) {
207                break;
208            }
209
210            let index = next_index.fetch_add(1, Ordering::Relaxed);
211            if index >= input_paths.len() {
212                break;
213            }
214
215            let path = input_paths[index].clone();
216            let outcome = process_one(path, timeout, Arc::clone(&processor), on_page.clone());
217            if tx.send(outcome).is_err() {
218                break;
219            }
220        }));
221    }
222    drop(tx);
223
224    let mut result = BatchResult::default();
225    let mut done = 0usize;
226
227    while done < total {
228        let Ok(outcome) = rx.recv() else {
229            break;
230        };
231        done += 1;
232
233        match outcome {
234            BatchOutcome::Success(path, value) => result.successes.push((path, value)),
235            BatchOutcome::Failure(path, err) => match config.on_error {
236                ErrorStrategy::Collect => result.failures.push((path, err)),
237                ErrorStrategy::Skip => result.skipped.push((path, err)),
238                ErrorStrategy::FailFast => {
239                    result.failures.push((path, err));
240                    stop.store(true, Ordering::Relaxed);
241                }
242            },
243        }
244
245        if let Some(callback) = config.on_progress.as_ref() {
246            callback(done, total);
247        }
248
249        if config.on_error == ErrorStrategy::FailFast && !result.failures.is_empty() {
250            break;
251        }
252    }
253
254    stop.store(true, Ordering::Relaxed);
255    for handle in handles {
256        let _ = handle.join();
257    }
258
259    result.duration = started.elapsed();
260    result
261}
262
263#[derive(Debug)]
264enum BatchOutcome<T> {
265    Success(PathBuf, T),
266    Failure(PathBuf, String),
267}
268
269fn process_one<F, T>(
270    path: PathBuf,
271    timeout: Duration,
272    processor: Arc<F>,
273    on_page: Option<SharedProgress>,
274) -> BatchOutcome<T>
275where
276    F: Fn(&PdfDocument) -> Result<T> + Send + Sync + 'static,
277    T: Send + 'static,
278{
279    let bytes = match fs::read(&path) {
280        Ok(bytes) => bytes,
281        Err(err) => return BatchOutcome::Failure(path, err.to_string()),
282    };
283
284    let (tx, rx) = mpsc::sync_channel(1);
285    thread::spawn(move || {
286        let outcome = (|| {
287            let doc = PdfDocument::open(bytes).map_err(|err| err.to_string())?;
288            let value = processor(&doc).map_err(|err| err.to_string())?;
289            if let Some(callback) = on_page.as_ref() {
290                let total = doc.page_count();
291                for done in 1..=total {
292                    callback(done, total);
293                }
294            }
295            Ok(value)
296        })();
297        let _ = tx.send(outcome);
298    });
299
300    match rx.recv_timeout(timeout) {
301        Ok(Ok(value)) => BatchOutcome::Success(path, value),
302        Ok(Err(err)) => BatchOutcome::Failure(path, err),
303        Err(mpsc::RecvTimeoutError::Timeout) => {
304            BatchOutcome::Failure(path, format!("timed out after {}s", timeout.as_secs_f64()))
305        }
306        Err(mpsc::RecvTimeoutError::Disconnected) => {
307            BatchOutcome::Failure(path, "worker disconnected".into())
308        }
309    }
310}
311
312fn collect_pdfs_recursive(root: &Path, out: &mut Vec<PathBuf>) -> std::io::Result<()> {
313    for entry in fs::read_dir(root)? {
314        let entry = entry?;
315        let path = entry.path();
316        if path.is_dir() {
317            collect_pdfs_recursive(&path, out)?;
318            continue;
319        }
320
321        if path
322            .extension()
323            .and_then(|ext| ext.to_str())
324            .is_some_and(|ext| ext.eq_ignore_ascii_case("pdf"))
325        {
326            out.push(path);
327        }
328    }
329
330    Ok(())
331}
332
333#[cfg(test)]
334mod tests {
335    use super::{process_batch, BatchConfig, ErrorStrategy, PdfBatch};
336    use std::fs;
337    use std::path::PathBuf;
338    use std::time::{Duration, SystemTime, UNIX_EPOCH};
339
340    #[test]
341    fn batch_process_counts_pages() {
342        let dir = make_temp_dir("counts");
343        let first = dir.join("first.pdf");
344        let second = dir.join("second.pdf");
345        fs::write(&first, minimal_pdf_bytes()).expect("write fixture");
346        fs::write(&second, minimal_pdf_bytes()).expect("write fixture");
347
348        let result = process_batch(
349            &[first.clone(), second.clone()],
350            BatchConfig {
351                workers: 2,
352                timeout: Duration::from_secs(5),
353                on_error: ErrorStrategy::Collect,
354                on_progress: None,
355            },
356            |doc| Ok(doc.page_count()),
357        );
358
359        assert_eq!(result.successes.len(), 2);
360        assert!(result.failures.is_empty());
361        assert!(result.skipped.is_empty());
362        assert!(result.successes.iter().all(|(_, count)| *count == 1));
363
364        fs::remove_dir_all(dir).ok();
365    }
366
367    #[test]
368    fn batch_process_skip_strategy_records_skips() {
369        let dir = make_temp_dir("skip");
370        let good = dir.join("good.pdf");
371        let bad = dir.join("bad.pdf");
372        fs::write(&good, minimal_pdf_bytes()).expect("write fixture");
373        fs::write(&bad, b"not a pdf").expect("write invalid fixture");
374
375        let result = process_batch(
376            &[good, bad],
377            BatchConfig {
378                workers: 1,
379                timeout: Duration::from_secs(5),
380                on_error: ErrorStrategy::Skip,
381                on_progress: None,
382            },
383            |doc| Ok(doc.page_count()),
384        );
385
386        assert_eq!(result.successes.len(), 1);
387        assert!(result.failures.is_empty());
388        assert_eq!(result.skipped.len(), 1);
389
390        fs::remove_dir_all(dir).ok();
391    }
392
393    #[test]
394    fn pdf_batch_collects_nested_pdfs() {
395        let dir = make_temp_dir("walk");
396        let nested = dir.join("nested");
397        fs::create_dir_all(&nested).expect("create nested dir");
398        fs::write(dir.join("root.pdf"), minimal_pdf_bytes()).expect("write root pdf");
399        fs::write(nested.join("child.pdf"), minimal_pdf_bytes()).expect("write child pdf");
400        fs::write(nested.join("ignore.txt"), b"hello").expect("write non-pdf");
401
402        let pdfs = PdfBatch::collect_pdfs(&dir).expect("collect pdfs");
403        assert_eq!(pdfs.len(), 2);
404
405        fs::remove_dir_all(dir).ok();
406    }
407
408    fn make_temp_dir(label: &str) -> PathBuf {
409        let nanos = SystemTime::now()
410            .duration_since(UNIX_EPOCH)
411            .expect("time")
412            .as_nanos();
413        let dir = std::env::temp_dir().join(format!(
414            "pdf-engine-batch-{label}-{}-{nanos}",
415            std::process::id()
416        ));
417        fs::create_dir_all(&dir).expect("create temp dir");
418        dir
419    }
420
421    fn minimal_pdf_bytes() -> Vec<u8> {
422        br#"%PDF-1.4
4231 0 obj
424<< /Type /Catalog /Pages 2 0 R >>
425endobj
4262 0 obj
427<< /Type /Pages /Kids [3 0 R] /Count 1 >>
428endobj
4293 0 obj
430<< /Type /Page /Parent 2 0 R /MediaBox [0 0 200 200] >>
431endobj
432xref
4330 4
4340000000000 65535 f 
4350000000009 00000 n 
4360000000058 00000 n 
4370000000115 00000 n 
438trailer
439<< /Size 4 /Root 1 0 R >>
440startxref
441186
442%%EOF
443"#
444        .to_vec()
445    }
446}
pdf_engine/batch.rs

pdf_engine/
batch.rs
