mecha10-core 0.8.0

//! `BatchBuilder`: size/timeout-based message batching.

use crate::context::Receiver;
use crate::messages::Message;
use std::time::Duration;
use tokio::sync::mpsc;

pub struct BatchBuilder<T: Message> {
    receiver: Receiver<T>,
    batch_size: usize,
    batch_timeout: Duration,
}

impl<T: Message + Send + Clone + 'static> BatchBuilder<T> {
    crate::new_with_defaults!(
        /// Create a new batch builder
        pub fn new(receiver: Receiver<T>, batch_size: usize, batch_timeout: Duration) -> Self
    );

    /// Build the batched receiver
    ///
    /// Spawns a task that collects messages into batches and emits them
    /// when either the batch size is reached or the timeout expires.
    pub fn build(self) -> Receiver<Vec<T>>
    where
        T: Clone,
    {
        let (tx, rx) = mpsc::unbounded_channel();

        tokio::spawn(async move {
            self.run_batch_pipeline(tx).await;
        });

        Receiver::from(rx)
    }

    /// Internal method to run the batch collection pipeline
    async fn run_batch_pipeline(mut self, tx: mpsc::UnboundedSender<Vec<T>>) {
        let mut batch = Vec::with_capacity(self.batch_size);
        let mut deadline = tokio::time::Instant::now() + self.batch_timeout;

        loop {
            tokio::select! {
                msg_opt = self.receiver.recv() => {
                    match msg_opt {
                        Some(msg) => {
                            batch.push(msg);

                            if batch.len() >= self.batch_size {
                                if tx.send(batch.clone()).is_err() {
                                    break;
                                }
                                batch.clear();
                                deadline = tokio::time::Instant::now() + self.batch_timeout;
                            }
                        }
                        None => {
                            if !batch.is_empty() {
                                let _ = tx.send(batch);
                            }
                            break;
                        }
                    }
                }

                _ = tokio::time::sleep_until(deadline) => {
                    if !batch.is_empty() {
                        if tx.send(batch.clone()).is_err() {
                            break;
                        }
                        batch.clear();
                    }
                    deadline = tokio::time::Instant::now() + self.batch_timeout;
                }
            }
        }
    }
}