olympipe_rs/

pipeline.rs

use crossbeam_channel::{bounded, unbounded, Receiver, RecvTimeoutError, Sender};
use std::sync::{Arc, Mutex};
use std::thread::{self, JoinHandle};
use std::time::Duration;

const DEFAULT_BUFFER: usize = 4;

/// A parallel data pipeline backed by bounded crossbeam channels and worker threads.
///
/// Each stage method consumes the pipeline and returns a new one whose type
/// reflects the output of that stage.  All worker threads are tracked in a
/// shared `Arc<Mutex<Vec<JoinHandle>>>` so that calling one of the terminal
/// methods (`collect`, `for_each`, `wait_for_completion`) joins every thread
/// spawned by the whole graph.
pub struct Pipeline<T: Send + 'static> {
    pub(crate) receiver: Receiver<T>,
    pub(crate) handles: Arc<Mutex<Vec<JoinHandle<()>>>>,
    pub(crate) buffer_size: usize,
}

// ─── constructor ──────────────────────────────────────────────────────────────

impl<T: Send + 'static> Pipeline<T> {
    /// Create a pipeline from any iterator.  A feeder thread is spawned
    /// immediately to push items into the first bounded channel.
    pub fn new(iter: impl IntoIterator<Item = T> + Send + 'static) -> Self {
        let (tx, rx) = bounded(DEFAULT_BUFFER);
        let handles: Arc<Mutex<Vec<JoinHandle<()>>>> = Arc::new(Mutex::new(Vec::new()));
        let h = thread::spawn(move || {
            for item in iter {
                if tx.send(item).is_err() {
                    break;
                }
            }
        });
        handles.lock().unwrap().push(h);
        Pipeline {
            receiver: rx,
            handles,
            buffer_size: DEFAULT_BUFFER,
        }
    }

    /// Override the inter-stage channel capacity (default: 4, matching olympipe).
    pub fn with_buffer(mut self, size: usize) -> Self {
        self.buffer_size = size;
        self
    }
}

// ─── internal helpers ─────────────────────────────────────────────────────────

impl<T: Send + 'static> Pipeline<T> {
    /// Spawn `count` worker threads that each read from `self.receiver` and
    /// write to a fresh bounded channel.  `worker` must be `Clone` so it can
    /// be duplicated per thread.  Use `spawn_single` for stages that run
    /// exactly one thread and whose closure cannot be cloned.
    fn spawn_stage<U, W>(self, count: usize, worker: W) -> Pipeline<U>
    where
        U: Send + 'static,
        W: Fn(Receiver<T>, Sender<U>) + Send + Clone + 'static,
    {
        let (tx, rx) = bounded::<U>(self.buffer_size);
        let handles = Arc::clone(&self.handles);

        for _ in 0..count {
            let w = worker.clone();
            let in_rx = self.receiver.clone();
            let out_tx = tx.clone();
            let h = thread::spawn(move || {
                let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                    w(in_rx, out_tx);
                }));
                if let Err(payload) = result {
                    let msg = payload
                        .downcast_ref::<&str>()
                        .copied()
                        .or_else(|| payload.downcast_ref::<String>().map(String::as_str))
                        .unwrap_or("(unknown panic)");
                    eprintln!("[olympipe-rs] worker panic: {msg}");
                }
            });
            handles.lock().unwrap().push(h);
        }

        Pipeline {
            receiver: rx,
            handles,
            buffer_size: self.buffer_size,
        }
    }

    /// Like `spawn_stage` but for a single worker thread whose closure does
    /// not need to implement `Clone`.
    fn spawn_single<U, W>(self, worker: W) -> Pipeline<U>
    where
        U: Send + 'static,
        W: FnOnce(Receiver<T>, Sender<U>) + Send + 'static,
    {
        let (tx, rx) = bounded::<U>(self.buffer_size);
        let handles = Arc::clone(&self.handles);
        let in_rx = self.receiver;

        let h = thread::spawn(move || {
            let result =
                std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| worker(in_rx, tx)));
            if let Err(payload) = result {
                let msg = payload
                    .downcast_ref::<&str>()
                    .copied()
                    .or_else(|| payload.downcast_ref::<String>().map(String::as_str))
                    .unwrap_or("(unknown panic)");
                eprintln!("[olympipe-rs] worker panic: {msg}");
            }
        });
        handles.lock().unwrap().push(h);

        Pipeline {
            receiver: rx,
            handles,
            buffer_size: self.buffer_size,
        }
    }
}

// ─── stage methods ────────────────────────────────────────────────────────────

impl<T: Send + 'static> Pipeline<T> {
    /// Map each item through `f`.  `count` worker threads run concurrently,
    /// sharing the input channel (MPMC) and writing to a shared output channel.
    pub fn task<U, F>(self, f: F, count: usize) -> Pipeline<U>
    where
        U: Send + 'static,
        F: Fn(T) -> U + Send + Clone + 'static,
    {
        let count = count.max(1);
        self.spawn_stage(count, move |rx, tx| {
            for item in rx {
                let out = f(item);
                if tx.send(out).is_err() {
                    break;
                }
            }
        })
    }

    /// Map each item through `f` which may fail.  On error, `on_error` is
    /// called with a clone of the original item and the error; if it returns
    /// `Some(v)`, `v` is forwarded downstream; `None` skips the item.
    pub fn task_or<U, E, F, H>(self, f: F, on_error: H) -> Pipeline<U>
    where
        T: Clone,
        U: Send + 'static,
        E: 'static,
        F: Fn(T) -> Result<U, E> + Send + 'static,
        H: Fn(T, E) -> Option<U> + Send + 'static,
    {
        self.spawn_single(move |rx, tx| {
            for item in rx {
                let cloned = item.clone();
                match f(item) {
                    Ok(out) => {
                        if tx.send(out).is_err() {
                            break;
                        }
                    }
                    Err(e) => {
                        if on_error(cloned, e)
                            .map(|fallback| tx.send(fallback).is_err())
                            .unwrap_or(false)
                        {
                            break;
                        }
                    }
                }
            }
        })
    }

    /// Keep only items for which `f` returns `true`.
    pub fn filter<F>(self, f: F) -> Pipeline<T>
    where
        F: Fn(&T) -> bool + Send + 'static,
    {
        self.spawn_single(move |rx, tx| {
            for item in rx {
                if f(&item) && tx.send(item).is_err() {
                    break;
                }
            }
        })
    }

    /// Group items into `Vec<T>` of at most `size` elements.
    /// The last (potentially incomplete) batch is always emitted.
    pub fn batch(self, size: usize) -> Pipeline<Vec<T>> {
        assert!(size >= 1, "batch size must be >= 1");
        self.spawn_single(move |rx, tx| {
            let mut buf: Vec<T> = Vec::with_capacity(size);
            for item in rx {
                buf.push(item);
                if buf.len() >= size {
                    let full = std::mem::replace(&mut buf, Vec::with_capacity(size));
                    if tx.send(full).is_err() {
                        return;
                    }
                }
            }
            if !buf.is_empty() {
                let _ = tx.send(buf);
            }
        })
    }

    /// Collect items arriving within `window` of each other into a `Vec<T>`.
    /// A new batch begins whenever the inter-item gap exceeds `window`, or when
    /// the upstream channel disconnects.
    pub fn temporal_batch(self, window: Duration) -> Pipeline<Vec<T>> {
        self.spawn_single(move |rx, tx| {
            loop {
                // Wait for the first item of a new batch (no timeout — block).
                let first = match rx.recv() {
                    Ok(item) => item,
                    Err(_) => break, // upstream done, no pending batch
                };
                let mut batch = vec![first];

                // Drain additional items within `window`.
                loop {
                    match rx.recv_timeout(window) {
                        Ok(item) => batch.push(item),
                        Err(RecvTimeoutError::Timeout) => break,
                        Err(RecvTimeoutError::Disconnected) => {
                            let _ = tx.send(batch);
                            return;
                        }
                    }
                }

                if tx.send(batch).is_err() {
                    return;
                }
            }
        })
    }

    /// Apply `f` to each item and forward every element yielded by the
    /// returned iterator (flatMap).
    pub fn explode<U, I, F>(self, f: F) -> Pipeline<U>
    where
        U: Send + 'static,
        I: IntoIterator<Item = U>,
        F: Fn(T) -> I + Send + 'static,
    {
        self.spawn_single(move |rx, tx| {
            for item in rx {
                for out in f(item) {
                    if tx.send(out).is_err() {
                        return;
                    }
                }
            }
        })
    }

    /// Route each item to one or both of two output pipelines according to `f`.
    ///
    /// Both returned pipelines share the same handle registry, so a single
    /// terminal call on either one will join all threads.
    ///
    /// The two output channels are **unbounded** so that the router thread can
    /// fully drain its input even when the two branches are consumed
    /// sequentially (calling `collect()` on one branch at a time). If you add
    /// further bounded stages after `split`, back-pressure is restored there.
    pub fn split<A, B, F>(self, f: F) -> (Pipeline<A>, Pipeline<B>)
    where
        A: Send + 'static,
        B: Send + 'static,
        F: Fn(T) -> (Option<A>, Option<B>) + Send + 'static,
    {
        let buf = self.buffer_size;
        let (tx_a, rx_a) = unbounded::<A>();
        let (tx_b, rx_b) = unbounded::<B>();
        let handles = Arc::clone(&self.handles);

        let h = thread::spawn(move || {
            for item in self.receiver {
                let (a, b) = f(item);
                if let Some(v) = a {
                    // If the A-branch is closed, keep feeding B.
                    let _ = tx_a.send(v);
                }
                if let Some(v) = b {
                    let _ = tx_b.send(v);
                }
            }
        });
        handles.lock().unwrap().push(h);

        (
            Pipeline {
                receiver: rx_a,
                handles: Arc::clone(&handles),
                buffer_size: buf,
            },
            Pipeline {
                receiver: rx_b,
                handles,
                buffer_size: buf,
            },
        )
    }

    /// Merge `self` and all pipelines in `others` into a single output stream.
    /// One forwarder thread is spawned per source; all handles are merged into
    /// the returned pipeline's handle registry.
    pub fn gather(self, others: Vec<Pipeline<T>>) -> Pipeline<T> {
        let buf = self.buffer_size;
        let (tx, rx) = bounded::<T>(buf);
        let handles = Arc::clone(&self.handles);

        // Forward self.
        let tx0 = tx.clone();
        let self_rx = self.receiver;
        let h = thread::spawn(move || {
            for item in self_rx {
                if tx0.send(item).is_err() {
                    break;
                }
            }
        });
        handles.lock().unwrap().push(h);

        // Forward each other pipeline, merging their handles first.
        for other in others {
            let mut other_handles = other.handles.lock().unwrap();
            let drained: Vec<_> = other_handles.drain(..).collect();
            drop(other_handles);
            handles.lock().unwrap().extend(drained);

            let tx_n = tx.clone();
            let other_rx = other.receiver;
            let h = thread::spawn(move || {
                for item in other_rx {
                    if tx_n.send(item).is_err() {
                        break;
                    }
                }
            });
            handles.lock().unwrap().push(h);
        }

        Pipeline {
            receiver: rx,
            handles,
            buffer_size: buf,
        }
    }

    /// Fold all items into a single accumulated value and emit it as the sole
    /// output item.
    pub fn reduce<U, F>(self, init: U, f: F) -> Pipeline<U>
    where
        U: Send + 'static,
        F: Fn(U, T) -> U + Send + 'static,
    {
        self.spawn_single(move |rx, tx| {
            let mut acc = init;
            for item in rx {
                acc = f(acc, item);
            }
            let _ = tx.send(acc);
        })
    }

    /// Pass at most `n` items downstream, then stop.
    pub fn limit(self, n: usize) -> Pipeline<T> {
        self.spawn_single(move |rx, tx| {
            for (i, item) in rx.iter().enumerate() {
                if tx.send(item).is_err() {
                    break;
                }
                if i + 1 >= n {
                    break;
                }
            }
        })
    }

    /// Error (drop the sender, propagating disconnection downstream) if no
    /// item arrives within `duration` between two consecutive items.
    pub fn timeout(self, duration: Duration) -> Pipeline<T> {
        self.spawn_single(move |rx, tx| {
            // First item: block indefinitely.
            match rx.recv() {
                Err(_) => return,
                Ok(first) => {
                    if tx.send(first).is_err() {
                        return;
                    }
                }
            }
            loop {
                match rx.recv_timeout(duration) {
                    Ok(item) => {
                        if tx.send(item).is_err() {
                            break;
                        }
                    }
                    Err(RecvTimeoutError::Timeout) => {
                        eprintln!(
                            "[olympipe-rs] timeout: no item received within {:?}",
                            duration
                        );
                        break;
                    }
                    Err(RecvTimeoutError::Disconnected) => break,
                }
            }
        })
    }

    /// Print each item to stdout and forward it unchanged.
    pub fn debug(self) -> Pipeline<T>
    where
        T: std::fmt::Debug,
    {
        self.spawn_single(move |rx, tx| {
            for item in rx {
                println!("[olympipe-rs] {:?}", item);
                if tx.send(item).is_err() {
                    break;
                }
            }
        })
    }
}

// ─── terminal methods ─────────────────────────────────────────────────────────

impl<T: Send + 'static> Pipeline<T> {
    /// Collect all output items into a `Vec`, then join every worker thread.
    pub fn collect(self) -> Vec<T> {
        let items: Vec<T> = self.receiver.into_iter().collect();
        let mut handles = self.handles.lock().unwrap();
        for h in handles.drain(..) {
            let _ = h.join();
        }
        items
    }

    /// Apply `f` to each output item, then join every worker thread.
    pub fn for_each<F>(self, mut f: F)
    where
        F: FnMut(T),
    {
        for item in self.receiver {
            f(item);
        }
        let mut handles = self.handles.lock().unwrap();
        for h in handles.drain(..) {
            let _ = h.join();
        }
    }

    /// Drain and discard all output, then join every worker thread.
    pub fn wait_for_completion(self) {
        for _ in self.receiver {}
        let mut handles = self.handles.lock().unwrap();
        for h in handles.drain(..) {
            let _ = h.join();
        }
    }
}