work_dispatcher 0.1.1

//! # WorkDispatcher
//!
//! A simple yet powerful concurrent data processing dispatch framework based on the Actor model.
//!
//! This crate provides a generic `WorkDispatcher` to set up a dispatcher that reads items
//! from a `Producer` and processes them concurrently using a pool of `Processor` actors.
//!
//! # Example
//!
//! ```no_run
//! // 1. Define your data item
//! struct MyDataItem(i32);
//!
//! // 2. Implement the Producer
//! use work_dispatcher::{Producer, Processor};
//! use flume::Sender;
//! use anyhow::Result;
//! use async_trait::async_trait;
//!
//! struct MyProducer;
//! #[async_trait]
//! impl Producer for MyProducer {
//!     type Item = MyDataItem;
//!     async fn run(self, sender: Sender<Self::Item>) {
//!         for i in 0..100 {
//!             sender.send_async(MyDataItem(i))?;
//!         }
//!         Ok(())
//!     }
//! }
//!
//! // 3. Implement the Processor
//! #[derive(Clone)]
//! struct MyProcessor;
//! #[async_trait]
//! impl Processor for MyProcessor {
//!     type Item = MyDataItem;
//!     type Context = String; // Example context
//!
//!     async fn process(&self, item: Self::Item, context: &Self::Context) {
//!         println!("Processing item #{} with context '{}'", item.0, context);
//!         tokio::time::sleep(std::time::Duration::from_millis(10)).await;
//!         Ok(())
//!     }
//! }
//!
//! // 4. Run the dispatcher
//! #[tokio::main]
//! async fn main() -> Result<()> {
//!     let producer = MyProducer;
//!     let processor = MyProcessor;
//!     let context = "MyAppContext".to_string();
//!
//!     work_dispatcher::WorkDispatcher::new(producer, processor, context)
//!         .workers(4)
//!         .run()
//!         .await?;
//!
//!     Ok(())
//! }
//! ```

use anyhow::Result;
use async_trait::async_trait;
use flume::Sender;
use tokio::task::JoinHandle;
use tracing::{debug, error};

/// `Producer` defines how to generate items for the dispatcher.
/// It runs in a dedicated blocking thread.
#[async_trait]
pub trait Producer {
    type Item: Send;

    /// The main loop for the producer. It should send items into the provided channel.
    async fn run(self, sender: Sender<Self::Item>);
}

/// `Processor` defines the logic for processing a single item.
/// It runs concurrently in multiple async tasks (actors).
#[async_trait]
pub trait Processor {
    type Item: Send;
    type Context: Send + Sync + Clone;

    /// The core logic for processing one item.
    async fn process(&self, item: Self::Item, context: &Self::Context);
}

/// The `WorkDispatcher` orchestrates the entire concurrent dispatcher.
pub struct WorkDispatcher<P: Producer, C: Processor<Item = P::Item>> {
    producer: P,
    processor: C,
    context: C::Context,
    num_workers: usize,
    channel_buffer: usize,
}

impl<P, C> WorkDispatcher<P, C>
where
    P: Producer + Send + 'static,
    P::Item: Send + 'static,
    C: Processor<Item = P::Item> + Clone + Send + Sync + 'static,
{
    /// Creates a new `WorkDispatcher` with a producer, a processor, and a shared context.
    pub fn new(producer: P, processor: C, context: C::Context) -> Self {
        let default_workers = num_cpus::get().max(1); // Ensure at least 1 worker
        Self {
            producer,
            processor,
            context,
            num_workers: default_workers,
            channel_buffer: default_workers * 100,
        }
    }

    /// Sets the number of concurrent worker tasks.
    /// Defaults to the number of logical CPU cores.
    pub fn workers(mut self, count: usize) -> Self {
        self.num_workers = count.max(1); // Must have at least one worker
        self
    }

    /// Sets the size of the bounded channel connecting the producer and processors.
    /// Defaults to `num_workers * 100`.
    pub fn buffer(mut self, size: usize) -> Self {
        self.channel_buffer = size;
        self
    }

    /// Runs the entire dispatcher and waits for it to complete.
    pub async fn run(self) -> Result<()> {
        debug!(
            "Starting WorkDispatcher with {} workers and a channel buffer of {}...",
            self.num_workers, self.channel_buffer
        );

        let (sender, receiver) = flume::bounded::<P::Item>(self.channel_buffer);

        //  Start the Producer Task
        let producer_handle = self.producer.run(sender);

        //  Start the Processor (Worker) Tasks
        let mut worker_handles = Vec::with_capacity(self.num_workers);
        for i in 0..self.num_workers {
            let receiver_clone = receiver.clone();
            let processor_clone = self.processor.clone();
            let context_clone = self.context.clone();

            let handle: JoinHandle<()> = tokio::spawn(async move {
                while let Ok(item) = receiver_clone.recv_async().await {
                    processor_clone.process(item, &context_clone).await;
                }
                debug!("[Worker {}] Channel closed, shutting down.", i);
            });
            worker_handles.push(handle);
        }

        // Wait for graceful shutdown ---
        producer_handle.await;
        debug!("Producer has finished successfully.");

        // Wait for all workers to finish processing the remaining items in the channel.
        for (i, handle) in worker_handles.into_iter().enumerate() {
            if let Err(e) = handle.await {
                error!("[Worker {}] Panicked: {}", i, e);
            }
        }
        debug!("All workers have finished.");

        Ok(())
    }
}