qubit_batch/process/impls/chunked_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 cmp,
12 num::NonZeroUsize,
13 sync::Arc,
14 time::Duration,
15};
16
17use qubit_progress::{
18 Progress,
19 reporter::ProgressReporter,
20};
21
22use crate::process::{
23 BatchProcessResult,
24 BatchProcessState,
25 BatchProcessor,
26 ChunkedBatchProcessError,
27 PROCESS_PROGRESS_METRIC_ID,
28 PROCESS_PROGRESS_METRIC_NAME,
29};
30
31use super::ChunkedBatchProcessorBuilder;
32
33/// Processes input items by submitting fixed-size chunks to a delegate.
34///
35/// `ChunkedBatchProcessor` is useful when the caller has a large logical batch
36/// but the real target must receive smaller batches, such as SQL batch insert
37/// operations with a maximum row count per statement.
38///
39/// The delegate must return a result whose `item_count` and `completed_count`
40/// match the submitted chunk length whenever it returns `Ok`. The delegate may
41/// still report a lower `processed_count`, such as when a database reports
42/// fewer affected rows than submitted rows. If the delegate cannot reach a
43/// terminal outcome for every item in the chunk, it should return `Err`;
44/// inconsistent `Ok` results are returned as
45/// [`ChunkedBatchProcessError::InvalidChunkResult`].
46///
47/// # Type Parameters
48///
49/// * `P` - Delegate processor receiving each chunk.
50///
51/// ```rust
52/// use std::{
53/// num::NonZeroUsize,
54/// time::Duration,
55/// };
56///
57/// use qubit_batch::{
58/// BatchProcessResult,
59/// BatchProcessResultBuilder,
60/// BatchProcessor,
61/// ChunkedBatchProcessor,
62/// };
63///
64/// struct InsertChunk;
65///
66/// impl BatchProcessor<i32> for InsertChunk {
67/// type Error = &'static str;
68///
69/// fn process_with_count<I>(
70/// &mut self,
71/// rows: I,
72/// count: usize,
73/// ) -> Result<BatchProcessResult, Self::Error>
74/// where
75/// I: IntoIterator<Item = i32>,
76/// {
77/// let processed = rows.into_iter().count();
78/// BatchProcessResultBuilder::builder(count)
79/// .completed_count(processed)
80/// .processed_count(processed)
81/// .chunk_count(1)
82/// .elapsed(Duration::ZERO)
83/// .build()
84/// .map_err(|_| "invalid process result")
85/// }
86/// }
87///
88/// let mut processor = ChunkedBatchProcessor::new(
89/// InsertChunk,
90/// NonZeroUsize::new(2).expect("chunk size should be non-zero"),
91/// );
92///
93/// let result = processor
94/// .process([1, 2, 3, 4, 5])
95/// .expect("array length should be exact");
96///
97/// assert_eq!(result.completed_count(), 5);
98/// assert_eq!(result.chunk_count(), 3);
99/// ```
100///
101pub struct ChunkedBatchProcessor<P> {
102 /// Delegate processor receiving each chunk.
103 pub(crate) delegate: P,
104 /// Maximum number of items submitted to the delegate at once.
105 pub(crate) chunk_size: NonZeroUsize,
106 /// Minimum interval between progress callbacks.
107 pub(crate) report_interval: Duration,
108 /// Reporter receiving batch lifecycle callbacks.
109 pub(crate) reporter: Arc<dyn ProgressReporter>,
110}
111
112impl<P> ChunkedBatchProcessor<P> {
113 /// Default interval between progress callbacks.
114 pub const DEFAULT_REPORT_INTERVAL: Duration = Duration::from_secs(5);
115
116 /// Creates a chunked batch processor.
117 ///
118 /// # Parameters
119 ///
120 /// * `delegate` - Processor receiving each chunk.
121 /// * `chunk_size` - Maximum number of items submitted in one chunk.
122 ///
123 /// # Returns
124 ///
125 /// A chunked processor using no-op progress reporting.
126 ///
127 /// # Type Constraints
128 ///
129 /// This constructor only stores `delegate`; it intentionally does not
130 /// require `P: BatchProcessor<Item>` because the item type is not part of
131 /// construction. That bound is enforced when this wrapper is used as a
132 /// [`BatchProcessor<Item>`], such as when calling [`BatchProcessor::process`].
133 /// Therefore, a value can be constructed with any delegate type, but it can
134 /// only process items for item types that the delegate actually supports.
135 #[inline]
136 pub fn new(delegate: P, chunk_size: NonZeroUsize) -> Self {
137 Self::builder(delegate, chunk_size).build()
138 }
139
140 /// Creates a builder for configuring a chunked batch processor.
141 ///
142 /// # Parameters
143 ///
144 /// * `delegate` - Processor receiving each chunk.
145 /// * `chunk_size` - Maximum number of items submitted in one chunk.
146 ///
147 /// # Returns
148 ///
149 /// A builder initialized with default settings.
150 #[inline]
151 pub fn builder(delegate: P, chunk_size: NonZeroUsize) -> ChunkedBatchProcessorBuilder<P> {
152 ChunkedBatchProcessorBuilder::new(delegate, chunk_size)
153 }
154
155 /// Returns the configured chunk size.
156 ///
157 /// # Returns
158 ///
159 /// The maximum number of items submitted to the delegate at once.
160 #[inline]
161 pub const fn chunk_size(&self) -> NonZeroUsize {
162 self.chunk_size
163 }
164
165 /// Returns the configured progress-report interval.
166 ///
167 /// # Returns
168 ///
169 /// The minimum time between due-based running progress callbacks.
170 #[inline]
171 pub const fn report_interval(&self) -> Duration {
172 self.report_interval
173 }
174
175 /// Returns the configured progress reporter.
176 ///
177 /// # Returns
178 ///
179 /// A shared reference to the configured progress reporter.
180 #[inline]
181 pub fn reporter(&self) -> &Arc<dyn ProgressReporter> {
182 &self.reporter
183 }
184
185 /// Returns a shared reference to the delegate processor.
186 ///
187 /// # Returns
188 ///
189 /// The wrapped delegate processor.
190 #[inline]
191 pub const fn delegate(&self) -> &P {
192 &self.delegate
193 }
194
195 /// Returns a mutable reference to the delegate processor.
196 ///
197 /// # Returns
198 ///
199 /// The wrapped delegate processor.
200 #[inline]
201 pub fn delegate_mut(&mut self) -> &mut P {
202 &mut self.delegate
203 }
204
205 /// Consumes this wrapper and returns the delegate processor.
206 ///
207 /// # Returns
208 ///
209 /// The wrapped delegate processor.
210 #[inline]
211 pub fn into_delegate(self) -> P {
212 self.delegate
213 }
214}
215
216impl<Item, P> BatchProcessor<Item> for ChunkedBatchProcessor<P>
217where
218 P: BatchProcessor<Item>,
219{
220 type Error = ChunkedBatchProcessError<P::Error>;
221
222 /// Processes items by delegating fixed-size chunks.
223 ///
224 /// # Parameters
225 ///
226 /// * `items` - Item source for the logical batch.
227 /// * `count` - Declared number of items expected from `items`.
228 ///
229 /// # Returns
230 ///
231 /// A result aggregating all successfully delegated chunks.
232 ///
233 /// # Errors
234 ///
235 /// Returns [`ChunkedBatchProcessError`] when the source count does not
236 /// match `count`, when the delegate fails for one chunk, or when a delegate
237 /// `Ok` result does not describe the submitted chunk.
238 fn process_with_count<I>(&mut self, items: I, count: usize) -> Result<BatchProcessResult, Self::Error>
239 where
240 I: IntoIterator<Item = Item>,
241 {
242 let reporter = Arc::clone(&self.reporter);
243 let mut progress = Progress::single_metric(
244 reporter.as_ref(),
245 self.report_interval,
246 PROCESS_PROGRESS_METRIC_ID,
247 PROCESS_PROGRESS_METRIC_NAME,
248 );
249 let state = BatchProcessState::new(count);
250 progress.report_started(|event| event.counters(state.progress_counters()));
251 let capacity = cmp::min(self.chunk_size.get(), count.max(1));
252 let mut chunk = Vec::with_capacity(capacity);
253
254 for item in items {
255 let observed_count = state.record_item_observed();
256 if observed_count > count {
257 if !chunk.is_empty() {
258 self.process_chunk(&mut chunk, &state, &mut progress)?;
259 }
260 let failed = progress.report_failed(|event| event.counters(state.progress_counters()));
261 let result = state.to_chunked_result(failed.elapsed());
262 return Err(ChunkedBatchProcessError::CountExceeded {
263 expected: count,
264 observed_at_least: observed_count,
265 result,
266 });
267 }
268 chunk.push(item);
269 if chunk.len() == self.chunk_size.get() {
270 self.process_chunk(&mut chunk, &state, &mut progress)?;
271 }
272 }
273
274 if !chunk.is_empty() {
275 self.process_chunk(&mut chunk, &state, &mut progress)?;
276 }
277
278 if state.observed_count() < count {
279 let failed = progress.report_failed(|event| event.counters(state.progress_counters()));
280 let result = state.to_chunked_result(failed.elapsed());
281 Err(ChunkedBatchProcessError::CountShortfall {
282 expected: count,
283 actual: state.observed_count(),
284 result,
285 })
286 } else {
287 let finished = progress.report_finished(|event| event.counters(state.progress_counters()));
288 let result = state.to_chunked_result(finished.elapsed());
289 Ok(result)
290 }
291 }
292}
293
294impl<P> ChunkedBatchProcessor<P> {
295 /// Submits one collected chunk to the delegate and updates aggregate state.
296 ///
297 /// # Parameters
298 ///
299 /// * `chunk` - Buffered items waiting to be submitted.
300 /// * `state` - Aggregate counters updated after successful delegation.
301 /// * `progress` - Progress run used for lifecycle and periodic callbacks.
302 ///
303 /// # Returns
304 ///
305 /// Returns `Ok(())` after the delegate accepts the chunk.
306 ///
307 /// # Errors
308 ///
309 /// Returns [`ChunkedBatchProcessError::ChunkFailed`] when the delegate
310 /// returns an error.
311 fn process_chunk<Item>(
312 &mut self,
313 chunk: &mut Vec<Item>,
314 state: &BatchProcessState,
315 progress: &mut Progress<'_>,
316 ) -> Result<(), ChunkedBatchProcessError<P::Error>>
317 where
318 P: BatchProcessor<Item>,
319 {
320 let chunk_len = chunk.len();
321 let start_index = state.completed_count();
322 let chunk_index = state.chunk_count();
323 let current_chunk = std::mem::take(chunk);
324 match self.delegate.process_with_count(current_chunk, chunk_len) {
325 Ok(chunk_result) => {
326 if chunk_result.item_count() != chunk_len || chunk_result.completed_count() != chunk_len {
327 let failed = progress.report_failed(|event| event.counters(state.progress_counters()));
328 let result = state.to_chunked_result(failed.elapsed());
329 return Err(ChunkedBatchProcessError::InvalidChunkResult {
330 chunk_index,
331 start_index,
332 chunk_len,
333 item_count: chunk_result.item_count(),
334 completed_count: chunk_result.completed_count(),
335 result,
336 });
337 }
338 state.record_chunk_processed(chunk_len, chunk_result.processed_count());
339 let _ = progress.report_running_if_due(|event| event.counters(state.running_chunk_progress_counters()));
340 Ok(())
341 }
342 Err(source) => {
343 let failed = progress.report_failed(|event| event.counters(state.progress_counters()));
344 let result = state.to_chunked_result(failed.elapsed());
345 Err(ChunkedBatchProcessError::ChunkFailed {
346 chunk_index,
347 start_index,
348 chunk_len,
349 source,
350 result,
351 })
352 }
353 }
354 }
355}