qubit_batch/process/impls/parallel_batch_processor.rs
1/*******************************************************************************
2 *
3 * Copyright (c) 2025 - 2026 Haixing Hu.
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * Licensed under the Apache License, Version 2.0.
8 *
9 ******************************************************************************/
10use std::{
11 num::NonZeroUsize,
12 sync::Arc,
13 thread,
14 time::Duration,
15};
16
17use qubit_function::{
18 ArcConsumer,
19 Consumer,
20};
21use qubit_progress::{
22 Progress,
23 reporter::ProgressReporter,
24};
25
26use crate::process::{
27 BatchProcessError,
28 BatchProcessResult,
29 BatchProcessState,
30 BatchProcessor,
31 PROCESS_PROGRESS_METRIC_ID,
32 PROCESS_PROGRESS_METRIC_NAME,
33};
34use crate::utils::run_scoped_parallel;
35
36use super::parallel_batch_processor_builder::ParallelBatchProcessorBuilder;
37
38/// Processes batch items with sequential fallback and scoped standard threads.
39///
40/// The processor stores the supplied consumer as an [`ArcConsumer`] so every
41/// worker can share it safely. By default, small batches run sequentially to
42/// avoid thread setup overhead. Larger batches use scoped worker threads for
43/// each [`BatchProcessor::process`] call, therefore input items may borrow data
44/// from the caller as long as they are [`Send`]. Running progress is reported
45/// between items on the sequential path and from a scoped reporter thread on
46/// the parallel path.
47///
48/// # Type Parameters
49///
50/// * `Item` - Item type consumed by the stored consumer.
51///
52/// ```rust
53/// use std::{
54/// sync::{
55/// Arc,
56/// atomic::{
57/// AtomicUsize,
58/// Ordering,
59/// },
60/// },
61/// };
62///
63/// use qubit_batch::{
64/// BatchProcessor,
65/// ParallelBatchProcessor,
66/// };
67///
68/// let total = Arc::new(AtomicUsize::new(0));
69/// let total_for_consumer = Arc::clone(&total);
70/// let mut processor = ParallelBatchProcessor::builder(move |item: &usize| {
71/// total_for_consumer.fetch_add(*item, Ordering::Relaxed);
72/// })
73/// .thread_count(2)
74/// .sequential_threshold(0)
75/// .build()
76/// .expect("parallel processor configuration should be valid");
77///
78/// let result = processor
79/// .process([1, 2, 3])
80/// .expect("array length should be exact");
81///
82/// assert!(result.is_success());
83/// assert_eq!(total.load(Ordering::Relaxed), 6);
84/// ```
85pub struct ParallelBatchProcessor<Item> {
86 /// Consumer shared by all scoped workers.
87 pub(crate) consumer: ArcConsumer<Item>,
88 /// Fixed worker-thread count used by each processing call.
89 pub(crate) thread_count: NonZeroUsize,
90 /// Maximum batch size that still uses sequential processing.
91 pub(crate) sequential_threshold: usize,
92 /// Minimum interval between progress callbacks.
93 pub(crate) report_interval: Duration,
94 /// Reporter receiving batch lifecycle callbacks.
95 pub(crate) reporter: Arc<dyn ProgressReporter>,
96}
97
98impl<Item> ParallelBatchProcessor<Item> {
99 /// Default interval between progress callbacks.
100 pub const DEFAULT_REPORT_INTERVAL: Duration = Duration::from_secs(5);
101
102 /// Default maximum batch size that still uses sequential processing.
103 pub const DEFAULT_SEQUENTIAL_THRESHOLD: usize = 100;
104
105 /// Creates a parallel consumer-backed batch processor.
106 ///
107 /// # Parameters
108 ///
109 /// * `consumer` - Thread-safe consumer invoked once for each accepted item.
110 ///
111 /// # Returns
112 ///
113 /// A processor storing `consumer` as an [`ArcConsumer`] and using
114 /// [`Self::default_thread_count`] workers.
115 #[inline]
116 pub fn new<C>(consumer: C) -> Self
117 where
118 C: Consumer<Item> + Send + Sync + 'static,
119 {
120 Self::builder(consumer)
121 .build()
122 .expect("default parallel batch processor should build")
123 }
124
125 /// Creates a builder for configuring a parallel consumer-backed processor.
126 ///
127 /// # Parameters
128 ///
129 /// * `consumer` - Thread-safe consumer invoked once for each accepted item.
130 ///
131 /// # Returns
132 ///
133 /// A builder initialized with default settings.
134 #[inline]
135 pub fn builder<C>(consumer: C) -> ParallelBatchProcessorBuilder<Item>
136 where
137 C: Consumer<Item> + Send + Sync + 'static,
138 {
139 ParallelBatchProcessorBuilder::new(consumer)
140 }
141
142 /// Returns the default worker-thread count.
143 ///
144 /// # Returns
145 ///
146 /// The available CPU parallelism, or `1` if it cannot be detected.
147 #[inline]
148 pub fn default_thread_count() -> usize {
149 thread::available_parallelism().map(usize::from).unwrap_or(1)
150 }
151
152 /// Returns the configured worker-thread count.
153 ///
154 /// # Returns
155 ///
156 /// The maximum number of scoped worker threads used for one batch.
157 #[inline]
158 pub const fn thread_count(&self) -> usize {
159 self.thread_count.get()
160 }
161
162 /// Returns the configured sequential fallback threshold.
163 ///
164 /// # Returns
165 ///
166 /// The maximum item count that still runs sequentially.
167 #[inline]
168 pub const fn sequential_threshold(&self) -> usize {
169 self.sequential_threshold
170 }
171
172 /// Returns the configured progress-report interval.
173 ///
174 /// # Returns
175 ///
176 /// The minimum time between due-based running progress callbacks.
177 #[inline]
178 pub const fn report_interval(&self) -> Duration {
179 self.report_interval
180 }
181
182 /// Returns the configured progress reporter.
183 ///
184 /// # Returns
185 ///
186 /// A shared reference to the configured progress reporter.
187 #[inline]
188 pub fn reporter(&self) -> &Arc<dyn ProgressReporter> {
189 &self.reporter
190 }
191
192 /// Returns the stored consumer.
193 ///
194 /// # Returns
195 ///
196 /// A shared reference to the arc-backed consumer.
197 #[inline]
198 pub const fn consumer(&self) -> &ArcConsumer<Item> {
199 &self.consumer
200 }
201
202 /// Consumes this processor and returns the stored consumer.
203 ///
204 /// # Returns
205 ///
206 /// The arc-backed consumer used by this processor.
207 #[inline]
208 pub fn into_consumer(self) -> ArcConsumer<Item> {
209 self.consumer
210 }
211}
212
213impl<Item> BatchProcessor<Item> for ParallelBatchProcessor<Item>
214where
215 Item: Send,
216{
217 type Error = BatchProcessError;
218
219 /// Processes items sequentially for small batches or on scoped workers.
220 ///
221 /// # Parameters
222 ///
223 /// * `items` - Item source for the batch.
224 /// * `count` - Declared number of items expected from `items`.
225 ///
226 /// # Returns
227 ///
228 /// A result with completed and processed counts equal to the number of
229 /// consumer calls when the input source yields exactly `count` items.
230 ///
231 /// # Errors
232 ///
233 /// Returns [`BatchProcessError::CountShortfall`] when the source ends before
234 /// `count`, or [`BatchProcessError::CountExceeded`] when the source yields an
235 /// extra item. Extra items are observed but not passed to the consumer.
236 ///
237 /// # Panics
238 ///
239 /// Propagates any panic raised by the stored consumer from the caller thread
240 /// or a worker thread, or by the configured progress reporter.
241 fn process_with_count<I>(&mut self, items: I, count: usize) -> Result<BatchProcessResult, Self::Error>
242 where
243 I: IntoIterator<Item = Item>,
244 {
245 let state = Arc::new(BatchProcessState::new(count));
246 let mut progress = Progress::single_metric(
247 self.reporter.as_ref(),
248 self.report_interval,
249 PROCESS_PROGRESS_METRIC_ID,
250 PROCESS_PROGRESS_METRIC_NAME,
251 );
252 progress.report_started(|event| event.counters(state.progress_counters()));
253
254 if count > 0 {
255 if count <= self.sequential_threshold {
256 self.process_sequential(items, count, state.as_ref(), &mut progress);
257 } else {
258 self.process_parallel_non_empty(items, count, Arc::clone(&state), &progress);
259 }
260 } else if items.into_iter().next().is_some() {
261 state.record_item_observed();
262 }
263
264 if state.observed_count() < count {
265 let failed = progress.report_failed(|event| event.counters(state.progress_counters()));
266 let result = state.to_direct_result(failed.elapsed());
267 Err(BatchProcessError::CountShortfall {
268 expected: count,
269 actual: state.observed_count(),
270 result,
271 })
272 } else if state.observed_count() > count {
273 let failed = progress.report_failed(|event| event.counters(state.progress_counters()));
274 let result = state.to_direct_result(failed.elapsed());
275 Err(BatchProcessError::CountExceeded {
276 expected: count,
277 observed_at_least: state.observed_count(),
278 result,
279 })
280 } else {
281 let finished = progress.report_finished(|event| event.counters(state.progress_counters()));
282 let result = state.to_direct_result(finished.elapsed());
283 Ok(result)
284 }
285 }
286}
287
288impl<Item> ParallelBatchProcessor<Item>
289where
290 Item: Send,
291{
292 /// Processes a declared batch on the caller thread.
293 ///
294 /// # Parameters
295 ///
296 /// * `items` - Item source for the batch.
297 /// * `count` - Declared item count.
298 /// * `state` - Processing state updated by this method.
299 /// * `progress` - Progress run used for between-item running callbacks.
300 ///
301 /// # Panics
302 ///
303 /// Propagates any panic raised while invoking the stored consumer.
304 fn process_sequential<I>(&self, items: I, count: usize, state: &BatchProcessState, progress: &mut Progress<'_>)
305 where
306 I: IntoIterator<Item = Item>,
307 {
308 for item in items {
309 let observed_count = state.record_item_observed();
310 if observed_count > count {
311 break;
312 }
313 state.record_item_started();
314 self.consumer.accept(&item);
315 state.record_item_processed();
316 let _ = progress.report_running_if_due(|event| event.counters(state.progress_counters()));
317 }
318 }
319
320 /// Processes a non-empty declared batch through scoped workers.
321 ///
322 /// # Parameters
323 ///
324 /// * `items` - Item source for the batch.
325 /// * `count` - Declared item count.
326 /// * `state` - Shared processing state updated by producer and workers.
327 /// * `progress` - Progress run used to spawn the running reporter.
328 ///
329 /// # Panics
330 ///
331 /// Propagates any worker panic raised while invoking the stored consumer.
332 fn process_parallel_non_empty<I>(
333 &self,
334 items: I,
335 count: usize,
336 state: Arc<BatchProcessState>,
337 progress: &Progress<'_>,
338 ) where
339 I: IntoIterator<Item = Item>,
340 {
341 thread::scope(|scope| {
342 let reporter_state = Arc::clone(&state);
343 let running_progress = progress.spawn_running_reporter(scope, move || reporter_state.progress_counters());
344 let running_point_handle = running_progress.point_handle();
345
346 let worker_count = self.thread_count.get().min(count);
347 let observer_state = Arc::clone(&state);
348 let worker_state = Arc::clone(&state);
349 let consumer = self.consumer.clone();
350 run_scoped_parallel(
351 items,
352 count,
353 worker_count,
354 move || observer_state.record_item_observed(),
355 move |_index, item| {
356 worker_state.record_item_started();
357 consumer.accept(&item);
358 worker_state.record_item_processed();
359 running_point_handle.report();
360 },
361 );
362 running_progress.stop_and_join();
363 });
364 }
365}