gwseq_io/

parallel.rs

//! Thread pools.
//!
//! One [`Executor`] over a `rayon::ThreadPool`.
//!
//! rayon's scoped parallelism covers what this library needs: a `par_iter`
//! over batches writes into indexed slots, which *is* submission order. There
//! is no completion-ordered result stream and no handle-checkout semaphore,
//! because nothing here wants either.
//!
//! One pool per reader, owned for its lifetime, so that a request reading a
//! few blocks does not pay to start and join threads. `close()` drops the
//! executor and that is what gives the threads back — but dropping a rayon pool
//! only *signals* its workers, so [`Executor`] keeps every `JoinHandle` and
//! joins them itself. See its `Drop`.

use std::sync::Arc;

use parking_lot::Mutex;

use crate::error::{Error, Result};

/// One thread per core, capped here.
pub const RECOMMENDED_MAX_THREADS: usize = 12;

/// Turn the `parallel` argument of the public API — where zero or less spells
/// "you decide" — into a thread count.
pub fn resolve_parallel(parallel: i64) -> usize {
    if parallel > 0 {
        return parallel as usize;
    }
    std::thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(1)
        .min(RECOMMENDED_MAX_THREADS)
}

pub struct Executor {
    /// `None` only during [`Drop`], which takes it to trigger the shutdown.
    pool: Option<rayon::ThreadPool>,
    /// The worker threads, so [`Drop`] can join them.
    ///
    /// Dropping a `rayon::ThreadPool` *signals* its workers to stop and returns
    /// without waiting for them, so a caller counting process threads
    /// immediately after `close()` still sees them. The pool is therefore
    /// built with a `spawn_handler` that keeps every `JoinHandle`, and
    /// dropping the pool is followed by joining them here.
    handles: Arc<Mutex<Vec<std::thread::JoinHandle<()>>>>,
    parallel: usize,
}

impl Drop for Executor {
    fn drop(&mut self) {
        // Order matters: dropping the pool is what tells the workers to stop,
        // and joining them is what makes `close()` mean they are gone. Joining
        // first would wait forever.
        drop(self.pool.take());
        for handle in self.handles.lock().drain(..) {
            let _ = handle.join();
        }
        // Nothing after the joins. On Darwin `pthread_join` returns before the
        // exiting thread has finished `__bsdthread_terminate`, so the process's
        // own thread count can briefly still include a worker this has already
        // joined — but the join is what "the threads are gone" means, and a
        // sleep in a close path to make an observer of `proc_pidinfo` see the
        // number it expects is measuring the OS, not this library.
    }
}

/// Hand-written: `rayon::ThreadPool` is not `Debug`, and a reader that holds an
/// executor still has to be printable.
impl std::fmt::Debug for Executor {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Executor")
            .field("parallel", &self.parallel)
            .finish()
    }
}

impl Executor {
    pub fn new(parallel: i64) -> Result<Self> {
        let parallel = resolve_parallel(parallel);
        let handles: Arc<Mutex<Vec<std::thread::JoinHandle<()>>>> =
            Arc::new(Mutex::new(Vec::with_capacity(parallel)));
        let sink = handles.clone();
        let pool = rayon::ThreadPoolBuilder::new()
            .num_threads(parallel)
            .spawn_handler(move |thread| {
                let name = thread
                    .name()
                    .map(str::to_string)
                    .unwrap_or_else(|| format!("gwseq-io-{}", thread.index()));
                let handle = std::thread::Builder::new()
                    .name(name)
                    .spawn(move || thread.run())?;
                sink.lock().push(handle);
                Ok(())
            })
            .thread_name(|i| format!("gwseq-io-{i}"))
            // rayon's default is to abort the process when a `spawn` job
            // panics, which would make the recovery below it unreachable: the
            // writer's `Promise::wait` returns `None` for a worker that dropped
            // its half without filling it, and turns that into "a block failed
            // to compress" on the call that is running. Swallowed here because
            // rayon has already printed the panic; what this restores is the
            // caller's chance to see an error instead of a signal.
            .panic_handler(|_| {})
            .build()
            // `Io`, not `InvalidArgument`: a thread that will not spawn is the
            // operating system refusing, not the caller asking for something
            // impossible — `parallel` was clamped to something sane long before
            // here. Python surfaces this as `SourceError`/`OSError`, which is
            // what it is.
            .map_err(|e| {
                Error::io(
                    format!("could not start {parallel} threads"),
                    std::io::Error::other(e.to_string()),
                )
            })?;
        Ok(Self {
            pool: Some(pool),
            handles,
            parallel,
        })
    }

    /// Worker handles still held, i.e. threads this executor will join when it
    /// is dropped. Test-only.
    #[cfg(test)]
    fn handle_count(&self) -> usize {
        self.handles.lock().len()
    }

    fn pool(&self) -> &rayon::ThreadPool {
        self.pool
            .as_ref()
            .expect("the pool is only taken while the executor is being dropped")
    }

    /// How many workers, which is also how many batches the readers split a
    /// request into.
    pub fn parallel(&self) -> usize {
        self.parallel
    }

    /// Run `f` on the pool. Rayon calls made inside it use these threads rather
    /// than the global pool, which is what keeps one reader's `parallel` a
    /// promise about this reader.
    pub fn install<R: Send>(&self, f: impl FnOnce() -> R + Send) -> R {
        self.pool().install(f)
    }

    /// Hand one job to the pool and return without waiting.
    ///
    /// What the writer's deflate pipeline is built on: a block is submitted,
    /// the caller goes on filling the next one, and the result is collected in
    /// submission order later. Everything else here is fork-join, which does
    /// not fit a producer that has to keep producing.
    pub fn spawn(&self, f: impl FnOnce() + Send + 'static) {
        self.pool().spawn(f);
    }

    /// Run one job per batch, collecting the first error.
    ///
    /// Batching is the caller's: `IndexedLocs::batches` splits loci into
    /// coverage-balanced groups, and that split decides which blocks each
    /// worker touches and so how much of the cache is shared. Letting rayon
    /// work-steal its own split would change both, and with them the order
    /// `f32` accumulators are summed in — which is visible in the last bits of
    /// every mean and standard deviation.
    pub fn for_each_batch<T: Send + Sync>(
        &self,
        batches: &[T],
        f: impl Fn(usize, &T) -> Result<()> + Send + Sync,
    ) -> Result<()> {
        use rayon::prelude::*;
        // `par_iter` over the batches lets rayon choose which thread runs which
        // batch, and nothing more: `f` is called once per batch, so the batch
        // boundaries — and with them the order values are accumulated inside
        // one — stay exactly where the caller put them.
        self.install(|| {
            batches
                .par_iter()
                .enumerate()
                .try_for_each(|(index, batch)| f(index, batch))
        })
    }

    /// The same, keeping each batch's result in submission order.
    ///
    /// What the extraction kernels use: a batch's loci own scattered slices of
    /// the output, not a contiguous one — the loci are sorted into file order
    /// while their output slices follow the request's order — so each worker
    /// fills a compact buffer of its own and the caller scatters afterwards.
    /// That keeps every bin accumulated by exactly one worker in one order,
    /// which is what makes the result bit-reproducible.
    pub fn map_batches<T: Send, B: Send + Sync>(
        &self,
        batches: &[B],
        f: impl Fn(usize, &B) -> Result<T> + Send + Sync,
    ) -> Result<Vec<T>> {
        use rayon::prelude::*;
        self.install(|| {
            batches
                .par_iter()
                .enumerate()
                .map(|(index, batch)| f(index, batch))
                .collect::<Result<Vec<T>>>()
        })
    }
}

/// A one-shot slot a worker fills and its submitter waits on.
///
/// What the writer's deflate pipeline hands out per block. `mpsc::Receiver` is
/// the obvious thing and is `Send` but not `Sync`, which a `#[pyclass]` holding
/// the writer needs; and a channel of one is more machinery than a slot and a
/// condvar anyway.
#[derive(Debug)]
pub struct Promise<T> {
    inner: Arc<(Mutex<Option<T>>, parking_lot::Condvar)>,
}

impl<T> Clone for Promise<T> {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
        }
    }
}

impl<T> Default for Promise<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> Promise<T> {
    pub fn new() -> Self {
        Self {
            inner: Arc::new((Mutex::new(None), parking_lot::Condvar::new())),
        }
    }

    /// Fill it, waking whoever is waiting. A second `set` is ignored.
    pub fn set(&self, value: T) {
        let mut slot = self.inner.0.lock();
        if slot.is_none() {
            *slot = Some(value);
            self.inner.1.notify_all();
        }
    }

    /// Wait for it, or return `None` if every other holder was dropped without
    /// filling it — which is what a panicked worker leaves behind.
    pub fn wait(self) -> Option<T> {
        let mut slot = self.inner.0.lock();
        loop {
            if let Some(value) = slot.take() {
                return Some(value);
            }
            // One other holder is the worker; none means it is gone.
            if Arc::strong_count(&self.inner) <= 1 {
                return None;
            }
            self.inner
                .1
                .wait_for(&mut slot, std::time::Duration::from_millis(50));
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering};

    #[test]
    fn resolve_parallel_takes_a_positive_count_as_given() {
        assert_eq!(resolve_parallel(1), 1);
        assert_eq!(resolve_parallel(7), 7);
        assert_eq!(resolve_parallel(100), 100);
    }

    #[test]
    fn zero_or_less_means_one_per_core_capped() {
        for asked in [0, -1, -12] {
            let n = resolve_parallel(asked);
            assert!(
                (1..=RECOMMENDED_MAX_THREADS).contains(&n),
                "{asked} gave {n}"
            );
        }
        assert_eq!(resolve_parallel(0), resolve_parallel(-1));
    }

    #[test]
    fn every_batch_runs_exactly_once() {
        let executor = Executor::new(4).unwrap();
        assert_eq!(executor.parallel(), 4);
        let batches: Vec<usize> = (0..64).collect();
        let seen: Vec<AtomicUsize> = (0..64).map(|_| AtomicUsize::new(0)).collect();
        executor
            .for_each_batch(&batches, |index, batch| {
                assert_eq!(index, *batch);
                seen[index].fetch_add(1, Ordering::SeqCst);
                Ok(())
            })
            .unwrap();
        assert!(seen.iter().all(|c| c.load(Ordering::SeqCst) == 1));
    }

    #[test]
    fn a_failing_batch_surfaces_as_the_result() {
        let executor = Executor::new(4).unwrap();
        let batches: Vec<usize> = (0..32).collect();
        let err = executor
            .for_each_batch(&batches, |_, batch| {
                if *batch == 17 {
                    Err(Error::invalid("batch 17"))
                } else {
                    Ok(())
                }
            })
            .unwrap_err();
        assert!(err.to_string().contains("batch 17"));
    }

    #[test]
    fn every_worker_is_held_for_joining_and_the_join_happens_on_drop() {
        // Counting the process's threads after a `close()` would be the
        // end-to-end version, but that count is noise here: cargo runs the
        // other tests in this file at the same time. What is deterministic is
        // that a handle is kept per worker and `Drop` joins every one of them
        // before returning.
        let executor = Executor::new(4).unwrap();
        let batches: Vec<usize> = (0..64).collect();
        executor.for_each_batch(&batches, |_, _| Ok(())).unwrap();
        assert_eq!(
            executor.handle_count(),
            4,
            "rayon spawned workers this executor did not keep a handle for"
        );
        // `drop` returns only once every `JoinHandle::join` has, which is the
        // guarantee: rayon's own `Drop` merely signals.
        drop(executor);
    }

    #[test]
    fn map_batches_keeps_submission_order() {
        let executor = Executor::new(4).unwrap();
        let batches: Vec<usize> = (0..50).collect();
        let out = executor
            .map_batches(&batches, |index, batch| Ok(index * 10 + batch))
            .unwrap();
        assert_eq!(out, (0..50).map(|i| i * 11).collect::<Vec<_>>());
    }

    #[test]
    fn map_batches_surfaces_a_failure() {
        let executor = Executor::new(4).unwrap();
        let batches: Vec<usize> = (0..32).collect();
        let err = executor
            .map_batches(&batches, |_, batch| {
                if *batch == 5 {
                    Err(Error::invalid("batch 5"))
                } else {
                    Ok(*batch)
                }
            })
            .unwrap_err();
        assert!(err.to_string().contains("batch 5"));
    }

    #[test]
    fn an_empty_request_is_not_an_error() {
        let executor = Executor::new(2).unwrap();
        let batches: Vec<usize> = Vec::new();
        executor.for_each_batch(&batches, |_, _| Ok(())).unwrap();
    }
}