iota_data_ingestion_core/worker_pool.rs
1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{
6 collections::{BTreeSet, HashMap, VecDeque},
7 fmt::Debug,
8 sync::Arc,
9 time::Instant,
10};
11
12use backoff::{ExponentialBackoff, backoff::Backoff};
13use futures::StreamExt;
14use iota_metrics::spawn_monitored_task;
15use iota_types::{
16 full_checkpoint_content::CheckpointData, messages_checkpoint::CheckpointSequenceNumber,
17};
18use tokio::{sync::mpsc, task::JoinHandle};
19use tokio_stream::wrappers::ReceiverStream;
20use tokio_util::sync::CancellationToken;
21use tracing::{debug, error, info, warn};
22
23use crate::{
24 IngestionError, IngestionResult, Reducer, Worker, executor::MAX_CHECKPOINTS_IN_PROGRESS,
25 reducer::reduce, util::reset_backoff,
26};
27
28type TaskName = String;
29type WorkerID = usize;
30
31/// Represents the possible message types a [`WorkerPool`] can communicate with
32/// external components.
33#[derive(Debug, Clone)]
34pub enum WorkerPoolStatus {
35 /// Message with information (e.g. `(<task-name>,
36 /// checkpoint_sequence_number)`) about the ingestion progress.
37 Running((TaskName, CheckpointSequenceNumber)),
38 /// Message with information (e.g. `<task-name>`) about shutdown status.
39 Shutdown(String),
40}
41
42/// Represents the possible message types a [`Worker`] can communicate with
43/// external components
44#[derive(Debug, Clone, Copy)]
45enum WorkerStatus<M> {
46 /// Message with information (e.g. `(<worker-id>`,
47 /// `checkpoint_sequence_number`, Option<[`Worker::Message`]>) about the
48 /// ingestion progress.
49 ///
50 /// The `Option<M>` is used to indicate that the worker skipped
51 /// processing the checkpoint. Useful for filtered checkpoints where non
52 /// matching checkpoints should not be forwarded to worker. In this case the
53 /// `checkpoint_sequence_number` is needed to track the progress.
54 Running((WorkerID, CheckpointSequenceNumber, Option<M>)),
55 /// Message with information (e.g. `<worker-id>`) about shutdown status.
56 Shutdown(WorkerID),
57}
58
59/// A pool of [`Worker`]'s that process checkpoints concurrently.
60///
61/// This struct manages a collection of workers that process checkpoints in
62/// parallel. It handles checkpoint distribution, progress tracking, and
63/// graceful shutdown. It can optionally use a [`Reducer`] to aggregate and
64/// process worker [`Messages`](Worker::Message).
65///
66/// # Examples
67/// ## Direct Processing (Without Batching)
68/// ```rust,no_run
69/// use std::sync::Arc;
70///
71/// use async_trait::async_trait;
72/// use iota_data_ingestion_core::{Worker, WorkerPool};
73/// use iota_types::full_checkpoint_content::{CheckpointData, CheckpointTransaction};
74/// #
75/// # struct DatabaseClient;
76/// #
77/// # impl DatabaseClient {
78/// # pub fn new() -> Self {
79/// # Self
80/// # }
81/// #
82/// # pub async fn store_transaction(&self,
83/// # _transactions: &CheckpointTransaction,
84/// # ) -> Result<(), DatabaseError> {
85/// # Ok(())
86/// # }
87/// # }
88/// #
89/// # #[derive(Debug, Clone)]
90/// # struct DatabaseError;
91/// #
92/// # impl std::fmt::Display for DatabaseError {
93/// # fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
94/// # write!(f, "database error")
95/// # }
96/// # }
97/// #
98/// # fn extract_transaction(checkpoint: &CheckpointData) -> CheckpointTransaction {
99/// # checkpoint.transactions.first().unwrap().clone()
100/// # }
101///
102/// struct DirectProcessor {
103/// // generic Database client.
104/// client: Arc<DatabaseClient>,
105/// }
106///
107/// #[async_trait]
108/// impl Worker for DirectProcessor {
109/// type Message = ();
110/// type Error = DatabaseError;
111///
112/// async fn process_checkpoint(
113/// &self,
114/// checkpoint: Arc<CheckpointData>,
115/// ) -> Result<Self::Message, Self::Error> {
116/// // extract a particulat transaction we care about.
117/// let tx: CheckpointTransaction = extract_transaction(checkpoint.as_ref());
118/// // store the transaction in our database of choice.
119/// self.client.store_transaction(&tx).await?;
120/// Ok(())
121/// }
122/// }
123///
124/// // configure worker pool for direct processing.
125/// let processor = DirectProcessor {
126/// client: Arc::new(DatabaseClient::new()),
127/// };
128/// let pool = WorkerPool::new(processor, "direct_processing".into(), 5, Default::default());
129/// ```
130///
131/// ## Batch Processing (With Reducer)
132/// ```rust,no_run
133/// use std::sync::Arc;
134///
135/// use async_trait::async_trait;
136/// use iota_data_ingestion_core::{Reducer, Worker, WorkerPool};
137/// use iota_types::full_checkpoint_content::{CheckpointData, CheckpointTransaction};
138/// # struct DatabaseClient;
139/// #
140/// # impl DatabaseClient {
141/// # pub fn new() -> Self {
142/// # Self
143/// # }
144/// #
145/// # pub async fn store_transactions_batch(&self,
146/// # _transactions: &Vec<CheckpointTransaction>,
147/// # ) -> Result<(), DatabaseError> {
148/// # Ok(())
149/// # }
150/// # }
151/// #
152/// # #[derive(Debug, Clone)]
153/// # struct DatabaseError;
154/// #
155/// # impl std::fmt::Display for DatabaseError {
156/// # fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
157/// # write!(f, "database error")
158/// # }
159/// # }
160///
161/// // worker that accumulates transactions for batch processing.
162/// struct BatchProcessor;
163///
164/// #[async_trait]
165/// impl Worker for BatchProcessor {
166/// type Message = Vec<CheckpointTransaction>;
167/// type Error = DatabaseError;
168///
169/// async fn process_checkpoint(
170/// &self,
171/// checkpoint: Arc<CheckpointData>,
172/// ) -> Result<Self::Message, Self::Error> {
173/// // collect all checkpoint transactions for batch processing.
174/// Ok(checkpoint.transactions.clone())
175/// }
176/// }
177///
178/// // batch reducer for efficient storage.
179/// struct TransactionBatchReducer {
180/// batch_size: usize,
181/// // generic Database client.
182/// client: Arc<DatabaseClient>,
183/// }
184///
185/// #[async_trait]
186/// impl Reducer<BatchProcessor> for TransactionBatchReducer {
187/// async fn commit(&self, batch: &[Vec<CheckpointTransaction>]) -> Result<(), DatabaseError> {
188/// let flattened: Vec<CheckpointTransaction> = batch.iter().flatten().cloned().collect();
189/// // store the transaction batch in the database of choice.
190/// self.client.store_transactions_batch(&flattened).await?;
191/// Ok(())
192/// }
193///
194/// fn should_close_batch(
195/// &self,
196/// batch: &[Vec<CheckpointTransaction>],
197/// _: Option<&Vec<CheckpointTransaction>>,
198/// ) -> bool {
199/// batch.iter().map(|b| b.len()).sum::<usize>() >= self.batch_size
200/// }
201/// }
202///
203/// // configure worker pool with batch processing.
204/// let processor = BatchProcessor;
205/// let reducer = TransactionBatchReducer {
206/// batch_size: 1000,
207/// client: Arc::new(DatabaseClient::new()),
208/// };
209/// let pool = WorkerPool::new_with_reducer(
210/// processor,
211/// "batch_processing".into(),
212/// 5,
213/// Default::default(),
214/// reducer,
215/// );
216/// ```
217pub struct WorkerPool<W: Worker> {
218 /// An unique name of the WorkerPool task.
219 pub task_name: String,
220 /// How many instances of the current [`Worker`] to create, more workers are
221 /// created more checkpoints they can process in parallel.
222 concurrency: usize,
223 /// The actual [`Worker`] instance itself.
224 worker: Arc<W>,
225 /// The reducer instance, responsible for batch processing.
226 reducer: Option<Box<dyn Reducer<W>>>,
227 backoff: Arc<ExponentialBackoff>,
228}
229
230impl<W: Worker + 'static> WorkerPool<W> {
231 /// Creates a new `WorkerPool` without a reducer.
232 pub fn new(
233 worker: W,
234 task_name: String,
235 concurrency: usize,
236 backoff: ExponentialBackoff,
237 ) -> Self {
238 Self {
239 task_name,
240 concurrency,
241 worker: Arc::new(worker),
242 reducer: None,
243 backoff: Arc::new(backoff),
244 }
245 }
246
247 /// Creates a new `WorkerPool` with a reducer.
248 pub fn new_with_reducer<R>(
249 worker: W,
250 task_name: String,
251 concurrency: usize,
252 backoff: ExponentialBackoff,
253 reducer: R,
254 ) -> Self
255 where
256 R: Reducer<W> + 'static,
257 {
258 Self {
259 task_name,
260 concurrency,
261 worker: Arc::new(worker),
262 reducer: Some(Box::new(reducer)),
263 backoff: Arc::new(backoff),
264 }
265 }
266
267 /// Runs the worker pool main logic.
268 pub async fn run(
269 mut self,
270 watermark: CheckpointSequenceNumber,
271 mut checkpoint_receiver: mpsc::Receiver<Arc<CheckpointData>>,
272 pool_status_sender: mpsc::Sender<WorkerPoolStatus>,
273 token: CancellationToken,
274 ) {
275 info!(
276 "Starting indexing pipeline {} with concurrency {}. Current watermark is {watermark}.",
277 self.task_name, self.concurrency
278 );
279 // This channel will be used to send progress data from Workers to WorkerPool
280 // main loop.
281 let (progress_sender, mut progress_receiver) = mpsc::channel(MAX_CHECKPOINTS_IN_PROGRESS);
282 // This channel will be used to send Workers progress data from WorkerPool to
283 // watermark tracking task.
284 let (watermark_sender, watermark_receiver) = mpsc::channel(MAX_CHECKPOINTS_IN_PROGRESS);
285 let mut idle: BTreeSet<_> = (0..self.concurrency).collect();
286 let mut checkpoints = VecDeque::new();
287 let mut workers_shutdown_signals = Vec::with_capacity(self.concurrency);
288 let (workers, workers_join_handles) = self.spawn_workers(progress_sender, token.clone());
289 // Spawn a task that tracks checkpoint processing progress. The task:
290 // - Receives (checkpoint_number, message) pairs from workers.
291 // - Maintains checkpoint sequence order.
292 // - Reports progress either:
293 // * After processing each chunk (simple tracking).
294 // * After committing batches (with reducer).
295 let watermark_handle = self.spawn_watermark_tracking(
296 watermark,
297 watermark_receiver,
298 pool_status_sender.clone(),
299 token.clone(),
300 );
301 // main worker pool loop.
302 loop {
303 tokio::select! {
304 Some(worker_progress_msg) = progress_receiver.recv() => {
305 match worker_progress_msg {
306 WorkerStatus::Running((worker_id, checkpoint_number, message)) => {
307 idle.insert(worker_id);
308 // Try to send progress to reducer. If it fails (reducer has exited),
309 // we just continue - we still need to wait for all workers to shutdown.
310 let _ = watermark_sender.send((checkpoint_number, message)).await;
311
312 // By checking if token was not cancelled we ensure that no
313 // further checkpoints will be sent to the workers.
314 while !token.is_cancelled() && !checkpoints.is_empty() && !idle.is_empty() {
315 let checkpoint = checkpoints.pop_front().unwrap();
316 let worker_id = idle.pop_first().unwrap();
317 if workers[worker_id].send(checkpoint).await.is_err() {
318 // The worker channel closing is a sign we need to exit this inner loop.
319 break;
320 }
321 }
322 }
323 WorkerStatus::Shutdown(worker_id) => {
324 // Track workers that have initiated shutdown.
325 workers_shutdown_signals.push(worker_id);
326 }
327 }
328 }
329 // Adding an if guard to this branch ensure that no checkpoints
330 // will be sent to workers once the token has been cancelled.
331 Some(checkpoint) = checkpoint_receiver.recv(), if !token.is_cancelled() => {
332 let sequence_number = checkpoint.checkpoint_summary.sequence_number;
333 if sequence_number < watermark {
334 continue;
335 }
336
337 if !Self::should_skip_filtered_checkpoint(&checkpoint) {
338 self.worker
339 .preprocess_hook(checkpoint.clone())
340 .map_err(|err| IngestionError::CheckpointHookProcessing(err.to_string()))
341 .expect("failed to preprocess task");
342 }
343
344 if idle.is_empty() {
345 checkpoints.push_back(checkpoint);
346 } else {
347 let worker_id = idle.pop_first().unwrap();
348 // If worker channel is closed, put the checkpoint back in queue
349 // and continue - we still need to wait for all worker shutdown signals.
350 if let Err(send_error) = workers[worker_id].send(checkpoint).await {
351 checkpoints.push_front(send_error.0);
352 };
353 }
354 }
355 }
356 // Once all workers have signaled completion, start the graceful shutdown
357 // process.
358 if workers_shutdown_signals.len() == self.concurrency {
359 break self
360 .workers_graceful_shutdown(
361 workers_join_handles,
362 watermark_handle,
363 pool_status_sender,
364 watermark_sender,
365 )
366 .await;
367 }
368 }
369 }
370
371 /// Spawn workers based on `self.concurrency` to process checkpoints
372 /// in parallel.
373 fn spawn_workers(
374 &self,
375 progress_sender: mpsc::Sender<WorkerStatus<W::Message>>,
376 token: CancellationToken,
377 ) -> (Vec<mpsc::Sender<Arc<CheckpointData>>>, Vec<JoinHandle<()>>) {
378 let mut worker_senders = Vec::with_capacity(self.concurrency);
379 let mut workers_join_handles = Vec::with_capacity(self.concurrency);
380
381 for worker_id in 0..self.concurrency {
382 let (worker_sender, mut worker_recv) =
383 mpsc::channel::<Arc<CheckpointData>>(MAX_CHECKPOINTS_IN_PROGRESS);
384 let cloned_progress_sender = progress_sender.clone();
385 let task_name = self.task_name.clone();
386 worker_senders.push(worker_sender);
387
388 let token = token.clone();
389
390 let worker = self.worker.clone();
391 let backoff = self.backoff.clone();
392 let join_handle = spawn_monitored_task!(async move {
393 loop {
394 tokio::select! {
395 // Once token is cancelled, notify worker's shutdown to the main loop
396 _ = token.cancelled() => {
397 _ = cloned_progress_sender.send(WorkerStatus::Shutdown(worker_id)).await;
398 break
399 },
400 Some(checkpoint) = worker_recv.recv() => {
401 let sequence_number = checkpoint.checkpoint_summary.sequence_number;
402 debug!("received checkpoint for processing {sequence_number} for workflow {task_name}", );
403 let start_time = Instant::now();
404 let status = Self::process_checkpoint_with_retry(worker_id, &worker, checkpoint, reset_backoff(&backoff), &token).await;
405 if matches!(status, WorkerStatus::Running((_,_, None))) {
406 debug!("checkpoint {sequence_number} for workflow {task_name} filtered out");
407 }
408 let trigger_shutdown = matches!(status, WorkerStatus::Shutdown(_));
409 if cloned_progress_sender.send(status).await.is_err() || trigger_shutdown {
410 break;
411 }
412 debug!(
413 "finished checkpoint processing {sequence_number} for workflow {task_name} in {:?}",
414 start_time.elapsed()
415 );
416 }
417 }
418 }
419 });
420 // Keep all join handles to ensure all workers are terminated before exiting
421 workers_join_handles.push(join_handle);
422 }
423 (worker_senders, workers_join_handles)
424 }
425
426 /// Returns `true` if the checkpoint was entirely stripped of its
427 /// transactions by a server-side filter, indicating a filtered-out
428 /// checkpoint.
429 ///
430 /// The fullnode's gRPC `stream_checkpoints` endpoint applies configured
431 /// `TransactionFilter`s to the expanded `transactions` payload but
432 /// leaves `checkpoint_contents` (the list of all transaction digests in
433 /// the original checkpoint) completely untouched.
434 fn should_skip_filtered_checkpoint(checkpoint: &CheckpointData) -> bool {
435 !checkpoint.checkpoint_contents.transactions().is_empty()
436 && checkpoint.transactions.is_empty()
437 }
438
439 /// Attempts to process a checkpoint with exponential backoff retries on
440 /// failure.
441 ///
442 /// This function repeatedly calls the
443 /// [`process_checkpoint`](Worker::process_checkpoint) method of the
444 /// provided [`Worker`] until either:
445 /// - The checkpoint processing succeeds, returning `WorkerStatus::Running` with the processed
446 /// message.
447 /// - A cancellation signal is received via the [`CancellationToken`], returning
448 /// `WorkerStatus::Shutdown(<worker-id>)`.
449 /// - All retry attempts are exhausted within backoff's maximum elapsed time: the pool is
450 /// cancelled and `WorkerStatus::Shutdown(<worker-id>)` is returned.
451 ///
452 /// # Retry Mechanism:
453 /// - Uses [`ExponentialBackoff`](backoff::ExponentialBackoff) to introduce increasing delays
454 /// between retry attempts.
455 /// - Checks for cancellation both before and after each processing attempt.
456 /// - If a cancellation signal is received during a backoff delay, the function exits
457 /// immediately with `WorkerStatus::Shutdown(<worker-id>)`.
458 async fn process_checkpoint_with_retry(
459 worker_id: WorkerID,
460 worker: &W,
461 checkpoint: Arc<CheckpointData>,
462 mut backoff: ExponentialBackoff,
463 token: &CancellationToken,
464 ) -> WorkerStatus<W::Message> {
465 let sequence_number = checkpoint.checkpoint_summary.sequence_number;
466
467 if Self::should_skip_filtered_checkpoint(&checkpoint) {
468 return if token.is_cancelled() {
469 WorkerStatus::Shutdown(worker_id)
470 } else {
471 WorkerStatus::Running((worker_id, sequence_number, None))
472 };
473 }
474
475 loop {
476 // check for cancellation before attempting processing.
477 if token.is_cancelled() {
478 return WorkerStatus::Shutdown(worker_id);
479 }
480
481 // attempt to process checkpoint.
482 match worker.process_checkpoint(checkpoint.clone()).await {
483 Ok(message) => {
484 return WorkerStatus::Running((worker_id, sequence_number, Some(message)));
485 }
486 Err(err) => {
487 let err = IngestionError::CheckpointProcessing(err.to_string());
488 warn!(
489 "transient worker execution error {err:?} for checkpoint {sequence_number}"
490 );
491 // check for cancellation after failed processing.
492 if token.is_cancelled() {
493 return WorkerStatus::Shutdown(worker_id);
494 }
495 }
496 }
497 // Retrying the same checkpoint any further cannot help, and a panic
498 // would abort the whole process in release builds, so stop the pool.
499 let Some(duration) = backoff.next_backoff() else {
500 error!(
501 "checkpoint {sequence_number} processing failed and the retries are exhausted, shutting down"
502 );
503 token.cancel();
504 return WorkerStatus::Shutdown(worker_id);
505 };
506 // if cancellation occurs during backoff wait, exit early with Shutdown.
507 // Otherwise (if timeout expires), continue with the next retry attempt.
508 if tokio::time::timeout(duration, token.cancelled())
509 .await
510 .is_ok()
511 {
512 return WorkerStatus::Shutdown(worker_id);
513 }
514 }
515 }
516
517 /// Spawns a task that tracks the progress of checkpoint processing,
518 /// optionally with message reduction.
519 ///
520 /// This function spawns one of two types of tracking tasks:
521 ///
522 /// 1. Simple Watermark Tracking (when reducer = None):
523 /// - Reports watermark after processing each chunk.
524 ///
525 /// 2. Batch Processing (when reducer = Some):
526 /// - Reports progress only after successful batch commits.
527 /// - A batch is committed based on [`should_close_batch`](Reducer::should_close_batch)
528 /// policy.
529 fn spawn_watermark_tracking(
530 &mut self,
531 watermark: CheckpointSequenceNumber,
532 watermark_receiver: mpsc::Receiver<(CheckpointSequenceNumber, Option<W::Message>)>,
533 executor_progress_sender: mpsc::Sender<WorkerPoolStatus>,
534 token: CancellationToken,
535 ) -> JoinHandle<Result<(), IngestionError>> {
536 let task_name = self.task_name.clone();
537 let backoff = self.backoff.clone();
538 if let Some(reducer) = self.reducer.take() {
539 return spawn_monitored_task!(reduce::<W>(
540 task_name,
541 watermark,
542 watermark_receiver,
543 executor_progress_sender,
544 reducer,
545 backoff,
546 token
547 ));
548 };
549 spawn_monitored_task!(simple_watermark_tracking::<W>(
550 task_name,
551 watermark,
552 watermark_receiver,
553 executor_progress_sender
554 ))
555 }
556
557 /// Start the workers graceful shutdown.
558 ///
559 /// - Awaits all worker handles.
560 /// - Awaits the reducer handle.
561 /// - Send `WorkerPoolStatus::Shutdown(<task-name>)` message notifying external components that
562 /// Worker Pool has been shutdown.
563 async fn workers_graceful_shutdown(
564 &self,
565 workers_join_handles: Vec<JoinHandle<()>>,
566 watermark_handle: JoinHandle<Result<(), IngestionError>>,
567 executor_progress_sender: mpsc::Sender<WorkerPoolStatus>,
568 watermark_sender: mpsc::Sender<(CheckpointSequenceNumber, Option<<W as Worker>::Message>)>,
569 ) {
570 for worker in workers_join_handles {
571 _ = worker
572 .await
573 .inspect_err(|err| tracing::error!("worker task panicked: {err}"));
574 }
575 // by dropping the sender we make sure that the stream will be closed and the
576 // watermark tracker task will exit its loop.
577 drop(watermark_sender);
578 _ = watermark_handle
579 .await
580 .inspect_err(|err| tracing::error!("watermark task panicked: {err}"));
581 _ = executor_progress_sender
582 .send(WorkerPoolStatus::Shutdown(self.task_name.clone()))
583 .await;
584 tracing::info!("Worker pool `{}` terminated gracefully", self.task_name);
585 }
586}
587
588/// Tracks checkpoint progress without reduction logic.
589///
590/// This function maintains a watermark of processed checkpoints by worker:
591/// 1. Receiving batches of progress status from workers.
592/// 2. Processing them in sequence order.
593/// 3. Reporting progress to the executor after each chunk from the stream.
594async fn simple_watermark_tracking<W: Worker>(
595 task_name: String,
596 mut current_checkpoint_number: CheckpointSequenceNumber,
597 watermark_receiver: mpsc::Receiver<(CheckpointSequenceNumber, Option<W::Message>)>,
598 executor_progress_sender: mpsc::Sender<WorkerPoolStatus>,
599) -> IngestionResult<()> {
600 // convert to a stream of MAX_CHECKPOINTS_IN_PROGRESS size. This way, each
601 // iteration of the loop will process all ready messages.
602 let mut stream =
603 ReceiverStream::new(watermark_receiver).ready_chunks(MAX_CHECKPOINTS_IN_PROGRESS);
604 // store unprocessed progress messages from workers.
605 let mut unprocessed = HashMap::new();
606 // track the next unprocessed checkpoint number for reporting progress
607 // after each chunk of messages is received from the stream.
608 let mut progress_update = None;
609
610 while let Some(update_batch) = stream.next().await {
611 unprocessed.extend(update_batch);
612 // Process messages sequentially based on checkpoint sequence number.
613 // This ensures in-order processing and maintains progress integrity.
614 while unprocessed.remove(¤t_checkpoint_number).is_some() {
615 current_checkpoint_number += 1;
616 progress_update = Some(current_checkpoint_number);
617 }
618 // report progress update to executor.
619 if let Some(watermark) = progress_update.take() {
620 executor_progress_sender
621 .send(WorkerPoolStatus::Running((task_name.clone(), watermark)))
622 .await
623 .map_err(|_| IngestionError::Channel("unable to send worker pool progress updates to executor, receiver half closed".into()))?;
624 }
625 }
626 Ok(())
627}