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, 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`
446 /// with the processed message.
447 /// - A cancellation signal is received via the [`CancellationToken`],
448 /// returning `WorkerStatus::Shutdown(<worker-id>)`.
449 /// - All retry attempts are exhausted within backoff's maximum elapsed
450 /// time, causing a panic.
451 ///
452 /// # Retry Mechanism:
453 /// - Uses [`ExponentialBackoff`](backoff::ExponentialBackoff) to introduce
454 /// increasing delays 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
457 /// function exits immediately with `WorkerStatus::Shutdown(<worker-id>)`.
458 ///
459 /// # Panics:
460 /// - If all retry attempts are exhausted within the backoff's maximum
461 /// elapsed time, indicating a persistent failure.
462 async fn process_checkpoint_with_retry(
463 worker_id: WorkerID,
464 worker: &W,
465 checkpoint: Arc<CheckpointData>,
466 mut backoff: ExponentialBackoff,
467 token: &CancellationToken,
468 ) -> WorkerStatus<W::Message> {
469 let sequence_number = checkpoint.checkpoint_summary.sequence_number;
470
471 if Self::should_skip_filtered_checkpoint(&checkpoint) {
472 return if token.is_cancelled() {
473 WorkerStatus::Shutdown(worker_id)
474 } else {
475 WorkerStatus::Running((worker_id, sequence_number, None))
476 };
477 }
478
479 loop {
480 // check for cancellation before attempting processing.
481 if token.is_cancelled() {
482 return WorkerStatus::Shutdown(worker_id);
483 }
484
485 // attempt to process checkpoint.
486 match worker.process_checkpoint(checkpoint.clone()).await {
487 Ok(message) => {
488 return WorkerStatus::Running((worker_id, sequence_number, Some(message)));
489 }
490 Err(err) => {
491 let err = IngestionError::CheckpointProcessing(err.to_string());
492 warn!(
493 "transient worker execution error {err:?} for checkpoint {sequence_number}"
494 );
495 // check for cancellation after failed processing.
496 if token.is_cancelled() {
497 return WorkerStatus::Shutdown(worker_id);
498 }
499 }
500 }
501 // get next backoff duration or panic if max retries exceeded.
502 let duration = backoff
503 .next_backoff()
504 .expect("max elapsed time exceeded: checkpoint processing failed for checkpoint {sequence_number}");
505 // if cancellation occurs during backoff wait, exit early with Shutdown.
506 // Otherwise (if timeout expires), continue with the next retry attempt.
507 if tokio::time::timeout(duration, token.cancelled())
508 .await
509 .is_ok()
510 {
511 return WorkerStatus::Shutdown(worker_id);
512 }
513 }
514 }
515
516 /// Spawns a task that tracks the progress of checkpoint processing,
517 /// optionally with message reduction.
518 ///
519 /// This function spawns one of two types of tracking tasks:
520 ///
521 /// 1. Simple Watermark Tracking (when reducer = None):
522 /// - Reports watermark after processing each chunk.
523 ///
524 /// 2. Batch Processing (when reducer = Some):
525 /// - Reports progress only after successful batch commits.
526 /// - A batch is committed based on
527 /// [`should_close_batch`](Reducer::should_close_batch) policy.
528 fn spawn_watermark_tracking(
529 &mut self,
530 watermark: CheckpointSequenceNumber,
531 watermark_receiver: mpsc::Receiver<(CheckpointSequenceNumber, Option<W::Message>)>,
532 executor_progress_sender: mpsc::Sender<WorkerPoolStatus>,
533 token: CancellationToken,
534 ) -> JoinHandle<Result<(), IngestionError>> {
535 let task_name = self.task_name.clone();
536 let backoff = self.backoff.clone();
537 if let Some(reducer) = self.reducer.take() {
538 return spawn_monitored_task!(reduce::<W>(
539 task_name,
540 watermark,
541 watermark_receiver,
542 executor_progress_sender,
543 reducer,
544 backoff,
545 token
546 ));
547 };
548 spawn_monitored_task!(simple_watermark_tracking::<W>(
549 task_name,
550 watermark,
551 watermark_receiver,
552 executor_progress_sender
553 ))
554 }
555
556 /// Start the workers graceful shutdown.
557 ///
558 /// - Awaits all worker handles.
559 /// - Awaits the reducer handle.
560 /// - Send `WorkerPoolStatus::Shutdown(<task-name>)` message notifying
561 /// external components that Worker Pool has been shutdown.
562 async fn workers_graceful_shutdown(
563 &self,
564 workers_join_handles: Vec<JoinHandle<()>>,
565 watermark_handle: JoinHandle<Result<(), IngestionError>>,
566 executor_progress_sender: mpsc::Sender<WorkerPoolStatus>,
567 watermark_sender: mpsc::Sender<(CheckpointSequenceNumber, Option<<W as Worker>::Message>)>,
568 ) {
569 for worker in workers_join_handles {
570 _ = worker
571 .await
572 .inspect_err(|err| tracing::error!("worker task panicked: {err}"));
573 }
574 // by dropping the sender we make sure that the stream will be closed and the
575 // watermark tracker task will exit its loop.
576 drop(watermark_sender);
577 _ = watermark_handle
578 .await
579 .inspect_err(|err| tracing::error!("watermark task panicked: {err}"));
580 _ = executor_progress_sender
581 .send(WorkerPoolStatus::Shutdown(self.task_name.clone()))
582 .await;
583 tracing::info!("Worker pool `{}` terminated gracefully", self.task_name);
584 }
585}
586
587/// Tracks checkpoint progress without reduction logic.
588///
589/// This function maintains a watermark of processed checkpoints by worker:
590/// 1. Receiving batches of progress status from workers.
591/// 2. Processing them in sequence order.
592/// 3. Reporting progress to the executor after each chunk from the stream.
593async fn simple_watermark_tracking<W: Worker>(
594 task_name: String,
595 mut current_checkpoint_number: CheckpointSequenceNumber,
596 watermark_receiver: mpsc::Receiver<(CheckpointSequenceNumber, Option<W::Message>)>,
597 executor_progress_sender: mpsc::Sender<WorkerPoolStatus>,
598) -> IngestionResult<()> {
599 // convert to a stream of MAX_CHECKPOINTS_IN_PROGRESS size. This way, each
600 // iteration of the loop will process all ready messages.
601 let mut stream =
602 ReceiverStream::new(watermark_receiver).ready_chunks(MAX_CHECKPOINTS_IN_PROGRESS);
603 // store unprocessed progress messages from workers.
604 let mut unprocessed = HashMap::new();
605 // track the next unprocessed checkpoint number for reporting progress
606 // after each chunk of messages is received from the stream.
607 let mut progress_update = None;
608
609 while let Some(update_batch) = stream.next().await {
610 unprocessed.extend(update_batch);
611 // Process messages sequentially based on checkpoint sequence number.
612 // This ensures in-order processing and maintains progress integrity.
613 while unprocessed.remove(¤t_checkpoint_number).is_some() {
614 current_checkpoint_number += 1;
615 progress_update = Some(current_checkpoint_number);
616 }
617 // report progress update to executor.
618 if let Some(watermark) = progress_update.take() {
619 executor_progress_sender
620 .send(WorkerPoolStatus::Running((task_name.clone(), watermark)))
621 .await
622 .map_err(|_| IngestionError::Channel("unable to send worker pool progress updates to executor, receiver half closed".into()))?;
623 }
624 }
625 Ok(())
626}