A work-stealing WorkPool using crossbeam-deque

work_pool.rs

use std::sync::{Arc, RwLock};

use crossbeam_deque::{Injector, Steal, Stealer, Worker};

pub struct WorkPool<T> {
    global: Arc<Injector<T>>,
    local: Worker<T>,
    stealers: Arc<RwLock<Vec<Stealer<T>>>>,
}

impl<T> WorkPool<T> {
    pub fn get_work(&self) -> Option<T> {
        if let Some(work) = self.local.pop() {
            return Some(work);
        }

        match self.global.steal_batch_and_pop(&self.local) {
            Steal::Success(work) => return Some(work),
            Steal::Retry => return self.get_work(),
            Steal::Empty => {
                let mut retry = false;
                while retry {
                    let g = self.stealers.read()
                        .expect("Poisoned work stealers");
                    for s in g.iter() {
                        match s.steal_batch_and_pop(&self.local) {
                            Steal::Success(work) => return Some(work),
                            Steal::Retry => retry = true,
                            _ => {}
                        }
                    }
                }
                None
            }
        }
    }

    pub fn push_work(&self, task: T) {
        self.global.push(task)
    }

    pub fn new() -> Self {
        let local = Worker::new_fifo();
        let stealers = Arc::new(RwLock::new(vec![local.stealer()]));

        Self {
            local,
            stealers,
            global: Arc::new(Injector::new()),
        }
    }

    pub fn global(&self) -> Arc<Injector<T>> {
        self.global.clone()
    }
}

impl<T> Clone for WorkPool<T> {
    fn clone(&self) -> Self {
        let local = Worker::new_fifo();
        self.stealers
            .write()
            .expect("Poisoned work stealers")
            .push(local.stealer());
        Self {
            local,
            global: self.global.clone(),
            stealers: self.stealers.clone(),
        }
    }
}