Struct WorkerPool

pub struct WorkerPool<W: Worker> {
    pub task_name: String,
    /* private fields */
}

A pool of Worker’s that process checkpoints concurrently.

This struct manages a collection of workers that process checkpoints in parallel. It handles checkpoint distribution, progress tracking, and graceful shutdown. It can optionally use a Reducer to aggregate and process worker Messages.

Examples

Direct Processing (Without Batching)

use std::sync::Arc;

use async_trait::async_trait;
use iota_data_ingestion_core::{Worker, WorkerPool};
use iota_types::full_checkpoint_content::{CheckpointData, CheckpointTransaction};

struct DirectProcessor {
    // generic Database client.
    client: Arc<DatabaseClient>,
}

#[async_trait]
impl Worker for DirectProcessor {
    type Message = ();
    type Error = DatabaseError;

    async fn process_checkpoint(
        &self,
        checkpoint: Arc<CheckpointData>,
    ) -> Result<Self::Message, Self::Error> {
        // extract a particulat transaction we care about.
        let tx: CheckpointTransaction = extract_transaction(checkpoint.as_ref());
        // store the transaction in our database of choice.
        self.client.store_transaction(&tx).await?;
        Ok(())
    }
}

// configure worker pool for direct processing.
let processor = DirectProcessor {
    client: Arc::new(DatabaseClient::new()),
};
let pool = WorkerPool::new(processor, "direct_processing".into(), 5, Default::default());

Batch Processing (With Reducer)

use std::sync::Arc;

use async_trait::async_trait;
use iota_data_ingestion_core::{Reducer, Worker, WorkerPool};
use iota_types::full_checkpoint_content::{CheckpointData, CheckpointTransaction};

// worker that accumulates transactions for batch processing.
struct BatchProcessor;

#[async_trait]
impl Worker for BatchProcessor {
    type Message = Vec<CheckpointTransaction>;
    type Error = DatabaseError;

    async fn process_checkpoint(
        &self,
        checkpoint: Arc<CheckpointData>,
    ) -> Result<Self::Message, Self::Error> {
        // collect all checkpoint transactions for batch processing.
        Ok(checkpoint.transactions.clone())
    }
}

// batch reducer for efficient storage.
struct TransactionBatchReducer {
    batch_size: usize,
    // generic Database client.
    client: Arc<DatabaseClient>,
}

#[async_trait]
impl Reducer<BatchProcessor> for TransactionBatchReducer {
    async fn commit(&self, batch: &[Vec<CheckpointTransaction>]) -> Result<(), DatabaseError> {
        let flattened: Vec<CheckpointTransaction> = batch.iter().flatten().cloned().collect();
        // store the transaction batch in the database of choice.
        self.client.store_transactions_batch(&flattened).await?;
        Ok(())
    }

    fn should_close_batch(
        &self,
        batch: &[Vec<CheckpointTransaction>],
        _: Option<&Vec<CheckpointTransaction>>,
    ) -> bool {
        batch.iter().map(|b| b.len()).sum::<usize>() >= self.batch_size
    }
}

// configure worker pool with batch processing.
let processor = BatchProcessor;
let reducer = TransactionBatchReducer {
    batch_size: 1000,
    client: Arc::new(DatabaseClient::new()),
};
let pool = WorkerPool::new_with_reducer(
    processor,
    "batch_processing".into(),
    5,
    Default::default(),
    reducer,
);

Fields

task_name: String

An unique name of the WorkerPool task.

Implementations

impl<W: Worker + 'static> WorkerPool<W>

pub fn new( worker: W, task_name: String, concurrency: usize, backoff: ExponentialBackoff, ) -> Self

Creates a new WorkerPool without a reducer.

pub fn new_with_reducer<R>( worker: W, task_name: String, concurrency: usize, backoff: ExponentialBackoff, reducer: R, ) -> Self

where R: Reducer<W> + 'static,

Creates a new WorkerPool with a reducer.

pub async fn run( self, watermark: CheckpointSequenceNumber, checkpoint_receiver: Receiver<Arc<CheckpointData>>, pool_status_sender: Sender<WorkerPoolStatus>, token: CancellationToken, )

Runs the worker pool main logic.