feldera-adapterlib 0.325.0

//! A job queue that dispatches work to a pool of tokio tasks while preserving
//! output ordering.
//!
//! Connectors use this to parse record batches in parallel: several worker
//! tasks parse concurrently, but their results are consumed in the order the
//! jobs were enqueued, so records reach the circuit in their original order.

use std::future::Future;
use std::pin::Pin;

use futures::channel::oneshot;
use tokio::{spawn, task::JoinHandle};

/// A job queue that dispatches work to a pool of tokio tasks.
///
/// While the jobs can complete out-of-order, their outputs are consumed in the same order
/// they were enqueued. This is useful for implementing parallel parsing, where parsed
/// records must be fed into the circuit in the order they were received.
/*
                          sync
   ┌─────────────────────────────────────────────────────────┐
   │                           ┌───────┐                     │
   │                        ┌─►│worker1├────────┐            │
   │              jobs      │  └───────┘        │            │
   ▼          ┌─┬─┬─┬─┬─┬─┐ │  ┌───────┐        │       ┌────┴───┐
producer  ─┬─►│ │ │ │ │ │ ├─┼─►│worker2├────┐   │       │consumer│
           │  └─┴─┴─┴─┴─┴─┘ │  └───────┘    │   │       └────────┘
           │                │  ┌───────┐    │   │            ▲
           │                └─►│worker3├──┐ │   │            │
           │                   └───────┘  ▼ ▼   ▼            │
           │                             ┌─┬─┬─┬─┬─┬─┐       │
           └────────────────────────────►│ │ │ │ │ │ ├───────┘
                                         └─┴─┴─┴─┴─┴─┘
                                          completions

* Producer adds jobs to the job queue. A job consists of an input value and
  a completion one-shot channel where the worker will send the output of the job.

* The receiving side of the channel is pushed to the completions queue.

* Worker tasks dequeue jobs from the jobs queue and send the result to the
  one-shot channel associated with each job. The consumer receives the next
  item from the completion queue and waits for the output of the job.

* When the producer needs to wait for all jobs in the queue to complete, it
  sends a special Sync message to the completions queue. Upon receiving
  this message, the consumer sends an acknowledgement to the sync channel.
*/
pub struct JobQueue<I, O> {
    /// The producer side of the job queue.
    job_sender: async_channel::Sender<(I, oneshot::Sender<O>)>,

    /// The producer side of the completions queue.
    completion_sender: async_channel::Sender<Completion<O>>,

    /// The receiving side of the sync channel.
    sync_receiver: async_channel::Receiver<()>,

    /// Worker tasks.
    workers: Vec<JoinHandle<()>>,

    /// Consumer task.
    consumer: JoinHandle<()>,
}

// Every message in the completions channel contains the receiving side of the
// oneshot channel that will contain the result of the completed job, or a special
// Sync message.
enum Completion<O> {
    Completion(oneshot::Receiver<O>),
    Sync,
}

impl<I, O> JobQueue<I, O>
where
    I: Send + 'static,
    O: Send + 'static,
{
    /// Create a job queue.
    ///
    /// # Arguments
    ///
    /// * `num_workers` - the number of threads in the worker pool. Must be >0.
    /// * `worker_builder` - a closure that returns a closure that each worker will execute for each job.
    ///   The outer closure is needed to allocate any resources needed by the worker.
    /// * `consumer_func` - closure that the consumer will execute for each completed job.
    pub fn new(
        num_workers: usize,
        worker_builder: impl Fn() -> Box<dyn FnMut(I) -> Pin<Box<dyn Future<Output = O> + Send>> + Send>,
        mut consumer_func: impl FnMut(O) + Send + 'static,
    ) -> Self {
        assert_ne!(num_workers, 0);

        // The jobs queue length is equal to the number of workers. This way, workers don't get
        // starved, but we also don't queue more data than necessary to keep the workers busy.
        // TODO: does it make sense to make queue length separately configurable?
        let (job_sender, job_receiver) =
            async_channel::bounded::<(I, oneshot::Sender<O>)>(num_workers);

        // The completion queue can contain at most one message per worker + one message per
        // job in the jobs queue plus the Sync message.
        let (completion_sender, completion_receiver) = async_channel::bounded(2 * num_workers + 1);
        let (sync_sender, sync_receiver) = async_channel::bounded(1);

        let workers = (0..num_workers)
            .map(move |_| {
                let mut worker_fn = worker_builder();

                let job_receiver = job_receiver.clone();
                spawn(async move {
                    loop {
                        let Ok((input, completion_sender)) = job_receiver.recv().await else {
                            return;
                        };
                        let result = worker_fn(input).await;
                        if completion_sender.send(result).is_err() {
                            return;
                        };
                    }
                })
            })
            .collect();

        let consumer = spawn(async move {
            loop {
                match completion_receiver.recv().await {
                    Err(_) => {
                        return;
                    }
                    Ok(Completion::Completion(receiver)) => {
                        let Ok(v) = receiver.await else {
                            continue;
                        };
                        consumer_func(v);
                    }
                    Ok(Completion::Sync) => {
                        if sync_sender.send(()).await.is_err() {
                            return;
                        }
                    }
                }
            }
        });

        Self {
            job_sender,
            completion_sender,
            sync_receiver,
            workers,
            consumer,
        }
    }

    /// Push a new job to the queue. Blocks until there is space in the queue.
    pub async fn push_job(&self, job: I) {
        let (completion_sender, completion_receiver) = oneshot::channel();
        let _ = self.job_sender.send((job, completion_sender)).await;
        let _ = self
            .completion_sender
            .send(Completion::Completion(completion_receiver))
            .await;
    }

    /// Wait for all previously queued jobs to complete.
    pub async fn flush(&self) {
        let _ = self.completion_sender.send(Completion::Sync).await;
        let _ = self.sync_receiver.recv().await;
    }
}

impl<I, O> Drop for JobQueue<I, O> {
    fn drop(&mut self) {
        self.consumer.abort();
        for worker in self.workers.drain(..) {
            worker.abort();
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::{Arc, Mutex};

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn test_job_queue() {
        let result = Arc::new(Mutex::new(Vec::new()));
        let result_clone = result.clone();

        let job_queue = super::JobQueue::new(
            6,
            || Box::new(|i: u32| Box::pin(async move { i })),
            move |i| result_clone.lock().unwrap().push(i),
        );

        for i in 0..1000000 {
            job_queue.push_job(i).await;
        }

        job_queue.flush().await;

        let expected: Vec<u32> = (0..1000000).collect();

        assert_eq!(&*result.lock().unwrap(), &*expected);
    }
}