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/// Convenience helpers built on top of [`process_batch`].
86#[derive(Debug, Default, Clone, Copy)]
87pub struct PdfBatch;
88
89impl PdfBatch {
90    /// Recursively collect `.pdf` files from a directory tree.
91    pub fn collect_pdfs(root: impl AsRef<Path>) -> std::io::Result<Vec<PathBuf>> {
92        let mut paths = Vec::new();
93        collect_pdfs_recursive(root.as_ref(), &mut paths)?;
94        paths.sort();
95        Ok(paths)
96    }
97
98    /// Extract combined text from every PDF under `root`.
99    pub fn extract_text(root: impl AsRef<Path>, config: BatchConfig) -> BatchResult<String> {
100        let root = root.as_ref().to_path_buf();
101        match Self::collect_pdfs(&root) {
102            Ok(paths) => process_batch(&paths, config, |doc| Ok(doc.extract_all_text())),
103            Err(err) => {
104                let mut result = BatchResult::default();
105                result.failures.push((root, err.to_string()));
106                result
107            }
108        }
109    }
110
111    /// Count pages for every PDF under `root`.
112    pub fn page_counts(root: impl AsRef<Path>, config: BatchConfig) -> BatchResult<usize> {
113        let root = root.as_ref().to_path_buf();
114        match Self::collect_pdfs(&root) {
115            Ok(paths) => process_batch(&paths, config, |doc| Ok(doc.page_count())),
116            Err(err) => {
117                let mut result = BatchResult::default();
118                result.failures.push((root, err.to_string()));
119                result
120            }
121        }
122    }
123}
124
125/// Process a list of PDF paths with a fixed-size worker pool.
126pub fn process_batch<P, F, T>(paths: &[P], config: BatchConfig, processor: F) -> BatchResult<T>
127where
128    P: AsRef<Path>,
129    F: Fn(&PdfDocument) -> Result<T> + Send + Sync + 'static,
130    T: Send + 'static,
131{
132    let started = Instant::now();
133    let input_paths = paths
134        .iter()
135        .map(|path| path.as_ref().to_path_buf())
136        .collect::<Vec<_>>();
137    let total = input_paths.len();
138
139    if total == 0 {
140        return BatchResult {
141            duration: started.elapsed(),
142            ..Default::default()
143        };
144    }
145
146    let workers = config.workers.max(1).min(total);
147    let next_index = Arc::new(AtomicUsize::new(0));
148    let stop = Arc::new(AtomicBool::new(false));
149    let input_paths = Arc::new(input_paths);
150    let processor = Arc::new(processor);
151    let (tx, rx) = mpsc::channel();
152    let mut handles = Vec::with_capacity(workers);
153
154    for _ in 0..workers {
155        let tx = tx.clone();
156        let next_index = Arc::clone(&next_index);
157        let stop = Arc::clone(&stop);
158        let input_paths = Arc::clone(&input_paths);
159        let processor = Arc::clone(&processor);
160        let timeout = config.timeout;
161
162        handles.push(thread::spawn(move || loop {
163            if stop.load(Ordering::Relaxed) {
164                break;
165            }
166
167            let index = next_index.fetch_add(1, Ordering::Relaxed);
168            if index >= input_paths.len() {
169                break;
170            }
171
172            let path = input_paths[index].clone();
173            let outcome = process_one(path, timeout, Arc::clone(&processor));
174            if tx.send(outcome).is_err() {
175                break;
176            }
177        }));
178    }
179    drop(tx);
180
181    let mut result = BatchResult::default();
182    let mut done = 0usize;
183
184    while done < total {
185        let Ok(outcome) = rx.recv() else {
186            break;
187        };
188        done += 1;
189
190        match outcome {
191            BatchOutcome::Success(path, value) => result.successes.push((path, value)),
192            BatchOutcome::Failure(path, err) => match config.on_error {
193                ErrorStrategy::Collect => result.failures.push((path, err)),
194                ErrorStrategy::Skip => result.skipped.push((path, err)),
195                ErrorStrategy::FailFast => {
196                    result.failures.push((path, err));
197                    stop.store(true, Ordering::Relaxed);
198                }
199            },
200        }
201
202        if let Some(callback) = config.on_progress.as_ref() {
203            callback(done, total);
204        }
205
206        if config.on_error == ErrorStrategy::FailFast && !result.failures.is_empty() {
207            break;
208        }
209    }
210
211    stop.store(true, Ordering::Relaxed);
212    for handle in handles {
213        let _ = handle.join();
214    }
215
216    result.duration = started.elapsed();
217    result
218}
219
220#[derive(Debug)]
221enum BatchOutcome<T> {
222    Success(PathBuf, T),
223    Failure(PathBuf, String),
224}
225
226fn process_one<F, T>(path: PathBuf, timeout: Duration, processor: Arc<F>) -> BatchOutcome<T>
227where
228    F: Fn(&PdfDocument) -> Result<T> + Send + Sync + 'static,
229    T: Send + 'static,
230{
231    let bytes = match fs::read(&path) {
232        Ok(bytes) => bytes,
233        Err(err) => return BatchOutcome::Failure(path, err.to_string()),
234    };
235
236    let (tx, rx) = mpsc::sync_channel(1);
237    thread::spawn(move || {
238        let outcome = (|| {
239            let doc = PdfDocument::open(bytes).map_err(|err| err.to_string())?;
240            processor(&doc).map_err(|err| err.to_string())
241        })();
242        let _ = tx.send(outcome);
243    });
244
245    match rx.recv_timeout(timeout) {
246        Ok(Ok(value)) => BatchOutcome::Success(path, value),
247        Ok(Err(err)) => BatchOutcome::Failure(path, err),
248        Err(mpsc::RecvTimeoutError::Timeout) => {
249            BatchOutcome::Failure(path, format!("timed out after {}s", timeout.as_secs_f64()))
250        }
251        Err(mpsc::RecvTimeoutError::Disconnected) => {
252            BatchOutcome::Failure(path, "worker disconnected".into())
253        }
254    }
255}
256
257fn collect_pdfs_recursive(root: &Path, out: &mut Vec<PathBuf>) -> std::io::Result<()> {
258    for entry in fs::read_dir(root)? {
259        let entry = entry?;
260        let path = entry.path();
261        if path.is_dir() {
262            collect_pdfs_recursive(&path, out)?;
263            continue;
264        }
265
266        if path
267            .extension()
268            .and_then(|ext| ext.to_str())
269            .is_some_and(|ext| ext.eq_ignore_ascii_case("pdf"))
270        {
271            out.push(path);
272        }
273    }
274
275    Ok(())
276}
277
278#[cfg(test)]
279mod tests {
280    use super::{process_batch, BatchConfig, ErrorStrategy, PdfBatch};
281    use std::fs;
282    use std::path::PathBuf;
283    use std::time::{Duration, SystemTime, UNIX_EPOCH};
284
285    #[test]
286    fn batch_process_counts_pages() {
287        let dir = make_temp_dir("counts");
288        let first = dir.join("first.pdf");
289        let second = dir.join("second.pdf");
290        fs::write(&first, minimal_pdf_bytes()).expect("write fixture");
291        fs::write(&second, minimal_pdf_bytes()).expect("write fixture");
292
293        let result = process_batch(
294            &[first.clone(), second.clone()],
295            BatchConfig {
296                workers: 2,
297                timeout: Duration::from_secs(5),
298                on_error: ErrorStrategy::Collect,
299                on_progress: None,
300            },
301            |doc| Ok(doc.page_count()),
302        );
303
304        assert_eq!(result.successes.len(), 2);
305        assert!(result.failures.is_empty());
306        assert!(result.skipped.is_empty());
307        assert!(result.successes.iter().all(|(_, count)| *count == 1));
308
309        fs::remove_dir_all(dir).ok();
310    }
311
312    #[test]
313    fn batch_process_skip_strategy_records_skips() {
314        let dir = make_temp_dir("skip");
315        let good = dir.join("good.pdf");
316        let bad = dir.join("bad.pdf");
317        fs::write(&good, minimal_pdf_bytes()).expect("write fixture");
318        fs::write(&bad, b"not a pdf").expect("write invalid fixture");
319
320        let result = process_batch(
321            &[good, bad],
322            BatchConfig {
323                workers: 1,
324                timeout: Duration::from_secs(5),
325                on_error: ErrorStrategy::Skip,
326                on_progress: None,
327            },
328            |doc| Ok(doc.page_count()),
329        );
330
331        assert_eq!(result.successes.len(), 1);
332        assert!(result.failures.is_empty());
333        assert_eq!(result.skipped.len(), 1);
334
335        fs::remove_dir_all(dir).ok();
336    }
337
338    #[test]
339    fn pdf_batch_collects_nested_pdfs() {
340        let dir = make_temp_dir("walk");
341        let nested = dir.join("nested");
342        fs::create_dir_all(&nested).expect("create nested dir");
343        fs::write(dir.join("root.pdf"), minimal_pdf_bytes()).expect("write root pdf");
344        fs::write(nested.join("child.pdf"), minimal_pdf_bytes()).expect("write child pdf");
345        fs::write(nested.join("ignore.txt"), b"hello").expect("write non-pdf");
346
347        let pdfs = PdfBatch::collect_pdfs(&dir).expect("collect pdfs");
348        assert_eq!(pdfs.len(), 2);
349
350        fs::remove_dir_all(dir).ok();
351    }
352
353    fn make_temp_dir(label: &str) -> PathBuf {
354        let nanos = SystemTime::now()
355            .duration_since(UNIX_EPOCH)
356            .expect("time")
357            .as_nanos();
358        let dir = std::env::temp_dir().join(format!(
359            "pdf-engine-batch-{label}-{}-{nanos}",
360            std::process::id()
361        ));
362        fs::create_dir_all(&dir).expect("create temp dir");
363        dir
364    }
365
366    fn minimal_pdf_bytes() -> Vec<u8> {
367        br#"%PDF-1.4
3681 0 obj
369<< /Type /Catalog /Pages 2 0 R >>
370endobj
3712 0 obj
372<< /Type /Pages /Kids [3 0 R] /Count 1 >>
373endobj
3743 0 obj
375<< /Type /Page /Parent 2 0 R /MediaBox [0 0 200 200] >>
376endobj
377xref
3780 4
3790000000000 65535 f 
3800000000009 00000 n 
3810000000058 00000 n 
3820000000115 00000 n 
383trailer
384<< /Size 4 /Root 1 0 R >>
385startxref
386186
387%%EOF
388"#
389        .to_vec()
390    }
391}
pdf_engine/batch.rs

pdf_engine/
batch.rs
