orx_parallel/runner/

par_runner.rs

use crate::parameters::Params;
use crate::pools::{ThreadPool, max_num_threads_for_computation};
#[cfg(feature = "std")]
use crate::runner::runner_variants::WithDiagnostics;
use orx_concurrent_iter::ConcurrentIter;

/// Controls how parallel computations are divided into chunks and run on a thread pool.
pub trait ParRunner: Sized + Sync {
    /// Underlying thread pool.
    type Pool: ThreadPool;

    /// Parallel computation state that is shared among thread computations.
    type State: Send + Sync;

    /// State local to a chunk, created when the chunk begins and consumed when it completes.
    type ChunkState;

    // required

    /// Reference to the underlying thread pool.
    fn pool(&self) -> &Self::Pool;

    /// Mutable reference to the underlying thread pool.
    fn pool_mut(&mut self) -> &mut Self::Pool;

    // required - orchestration

    /// Returns Some of the index of the new thread if it is beneficial to spawn a new thread
    /// for the remaining computation; it returns None otherwise.
    ///
    /// Since this method is called sequentially, returned value will always be `Some(spawned)` if
    /// it returns some.
    fn do_spawn_new(spawned: usize, state: &Self::State) -> Option<usize>;

    /// Creates an initial state for a new parallel computation.
    fn new_state(
        &mut self,
        params: Params,
        max_num_threads: usize,
        size_hint: (usize, Option<usize>),
    ) -> Self::State;

    /// Adjusts runner state for an input source that is consumed serially.
    fn configure_for_serialized_input(state: &mut Self::State, size_hint: (usize, Option<usize>));

    /// Called when a worker thread begins processing the computation.
    fn begin_thread(state: &Self::State, th_idx: usize);

    /// Returns the next chunk size to be pulled from the input with remaining length
    /// provided by the `size_hint` for the current `state`.
    fn next_chunk_size(state: &Self::State, size_hint: (usize, Option<usize>)) -> usize;

    /// Creates state for a chunk before its items are processed.
    fn begin_chunk(th_idx: usize, chunk_size: usize) -> Self::ChunkState;

    /// Called when a chunk completes, with the state returned by [`Self::begin_chunk`].
    fn complete_chunk(state: &Self::State, chunk_state: Self::ChunkState);

    /// Called when a worker thread finishes processing the computation.
    fn complete_thread(state: &Self::State, th_idx: usize);

    /// Called after all worker threads finish, consuming the shared computation state.
    fn complete_computation(state: Self::State);

    // provided

    /// Wraps this runner to print diagnostics about thread usage and chunk sizes (`std` only).
    #[cfg(feature = "std")]
    fn with_diagnostics(self) -> WithDiagnostics<Self> {
        WithDiagnostics::new(self)
    }

    // provided - helpers

    /// Creates computation state and returns the effective thread limit with that state.
    ///
    /// Applies the pool limit and any per-computation limit, then configures the state for
    /// serialized input when needed.
    fn nt_state(
        &mut self,
        params: Params,
        is_source_serialized: bool,
        size_hint: (usize, Option<usize>),
        computation_max_nt: Option<usize>,
    ) -> (usize, Self::State) {
        let max_nt = max_num_threads_for_computation(self.pool(), params, size_hint);

        let max_nt = match computation_max_nt {
            Some(0) => 1,
            Some(comp_nt) if comp_nt < max_nt => comp_nt,
            _ => max_nt,
        };

        let mut state = self.new_state(params, max_nt, size_hint);
        if is_source_serialized {
            Self::configure_for_serialized_input(&mut state, size_hint);
        }
        (max_nt, state)
    }

    /// Exhausts the iterator and completes the active chunk when processing stops early.
    fn broadcast_stop<I: ConcurrentIter>(
        iter: &I,
        state: &Self::State,
        chunk_state: Self::ChunkState,
    ) {
        iter.skip_to_end();
        Self::complete_chunk(state, chunk_state);
    }
}

impl<P: ParRunner> ParRunner for &mut P {
    type Pool = P::Pool;

    type State = P::State;

    type ChunkState = P::ChunkState;

    fn pool(&self) -> &Self::Pool {
        <P as ParRunner>::pool(self)
    }

    fn pool_mut(&mut self) -> &mut Self::Pool {
        <P as ParRunner>::pool_mut(self)
    }

    fn do_spawn_new(spawned: usize, state: &Self::State) -> Option<usize> {
        <P as ParRunner>::do_spawn_new(spawned, state)
    }

    fn new_state(
        &mut self,
        params: Params,
        max_num_threads: usize,
        size_hint: (usize, Option<usize>),
    ) -> Self::State {
        <P as ParRunner>::new_state(self, params, max_num_threads, size_hint)
    }

    fn begin_thread(state: &Self::State, th_idx: usize) {
        <P as ParRunner>::begin_thread(state, th_idx);
    }

    fn next_chunk_size(state: &Self::State, size_hint: (usize, Option<usize>)) -> usize {
        <P as ParRunner>::next_chunk_size(state, size_hint)
    }

    fn begin_chunk(th_idx: usize, chunk_size: usize) -> Self::ChunkState {
        <P as ParRunner>::begin_chunk(th_idx, chunk_size)
    }

    fn complete_chunk(state: &Self::State, chunk_state: Self::ChunkState) {
        <P as ParRunner>::complete_chunk(state, chunk_state);
    }

    fn complete_thread(state: &Self::State, th_idx: usize) {
        <P as ParRunner>::complete_thread(state, th_idx);
    }

    fn complete_computation(state: Self::State) {
        <P as ParRunner>::complete_computation(state);
    }

    fn configure_for_serialized_input(state: &mut Self::State, size_hint: (usize, Option<usize>)) {
        <P as ParRunner>::configure_for_serialized_input(state, size_hint);
    }
}