Rayon ThreadPool Worker State?

[tqwewe]

Hi! I'm working on a work-stealing thread pool using Crossbeam and have run into a challenge.

In my design, worker threads not only steal work but also need to receive messages from a channel to update their internal state periodically. To achieve this, I'm using:

• Crossbeam's channel for state updates.
• Crossbeam's Injector for work-stealing task processing.

The Challenge

How can worker threads avoid consuming CPU while waiting for work?

Ideally, I'd implement some kind of wake-up mechanism when jobs are submitted so that idle threads don’t spin-wait unnecessarily. However, designing a proper notification system seems quite complex.

Alternative Approach

I'm considering using Rayon's ThreadPool instead, but I have a few concerns:

1. Worker State: I couldn't find a straightforward way for Rayon workers to maintain per-thread state. I could use thread_local, but that wouldn’t help if the state needs external updates.
2. State Updates: If the state isn’t thread_local, how would I update it in a predicable way between work being processed?

Use Case

The goal of this thread pool is to perform atomic file reads using pread concurrently.

Below is my current Crossbeam implementation:

Crossbeam Work-Stealing Threadpool

use std::{
    collections::HashMap,
    io,
    panic::{catch_unwind, AssertUnwindSafe},
    sync::Arc,
    thread,
};

use crossbeam::{
    channel,
    deque::{Injector, Steal},
};
use tracing::error;

use super::{BucketSegmentReader, Record};

type ReadFn = Box<dyn FnOnce(io::Result<Option<Record>>) + Send>;

#[derive(Clone, Debug)]
pub struct ReaderThreadPool {
    txs: Arc<[channel::Sender<(u16, u32, BucketSegmentReader)>]>,
    injector: Arc<Injector<ReadRequest>>,
}

impl ReaderThreadPool {
    /// Spawns threads to process read requests in a thread pool.
    pub fn spawn(num_workers: u16) -> Self {
        let mut txs = Vec::with_capacity(num_workers as usize);
        let injector = Arc::new(Injector::<ReadRequest>::new());
        for _ in 0..num_workers {
            let (tx, rx) = channel::unbounded();
            let worker = ReaderWorker::new(Arc::clone(&injector), rx);
            thread::spawn(move || worker.run());
            txs.push(tx);
        }

        ReaderThreadPool {
            txs: Arc::from(txs),
            injector,
        }
    }

    /// Adds a bucket segment reader to all workers in the thread pool.
    pub fn add_bucket_segment(
        &self,
        bucket_id: u16,
        segment: u32,
        reader: &BucketSegmentReader,
    ) -> io::Result<()> {
        for tx in self.txs.as_ref() {
            let _ = tx.send((bucket_id, segment, reader.try_clone()?));
        }

        Ok(())
    }

    /// Reads a record, utilizing a thread from the thread pool.
    pub fn read_record(
        &self,
        bucket_id: u16,
        segment: u32,
        offset: u64,
        handler: impl FnOnce(io::Result<Option<Record>>) + Send + 'static,
    ) {
        self.injector.push(ReadRequest {
            bucket_id,
            segment,
            offset,
            handler: Box::new(handler),
        });
    }
}

struct ReaderWorker {
    injector: Arc<Injector<ReadRequest>>,
    readers: HashMap<(u16, u32), BucketSegmentReader>,
    rx: channel::Receiver<(u16, u32, BucketSegmentReader)>,
}

impl ReaderWorker {
    fn new(
        injector: Arc<Injector<ReadRequest>>,
        rx: channel::Receiver<(u16, u32, BucketSegmentReader)>,
    ) -> Self {
        Self {
            injector,
            readers: HashMap::new(),
            rx,
        }
    }

    fn run(mut self) {
        let mut request: Option<ReadRequest> = None;
        loop {
            // Check if a new bucket segment reader has been registered
            if let Ok((bucket_id, segment, reader)) = self.rx.try_recv() {
                self.add_bucket_segment(bucket_id, segment, reader);
                continue;
            }

            // Check if we need to retry a previous request
            if let Some(req) = request.take() {
                if let Some(handler) =
                    self.process(req.bucket_id, req.segment, req.offset, req.handler)
                {
                    handler(Err(io::Error::new(
                        io::ErrorKind::Other,
                        "bucket segment reader not found",
                    )));
                }
                continue;
            }

            // Steal any work pending to be done
            match self.injector.steal() {
                Steal::Success(req) => {
                    // Try to process the work, otherwise if the reader wasn't found then we'll read from the channel and retry
                    request = self
                        .process(req.bucket_id, req.segment, req.offset, req.handler)
                        .map(|handler| ReadRequest {
                            bucket_id: req.bucket_id,
                            segment: req.segment,
                            offset: req.offset,
                            handler,
                        });
                }
                Steal::Retry => continue,
                Steal::Empty => {
                    thread::yield_now();
                }
            }
        }
    }

    fn add_bucket_segment(&mut self, bucket_id: u16, segment: u32, reader: BucketSegmentReader) {
        self.readers.insert((bucket_id, segment), reader);
    }

    fn process(
        &mut self,
        bucket_id: u16,
        segment: u32,
        offset: u64,
        handler: ReadFn,
    ) -> Option<ReadFn> {
        match self.readers.get_mut(&(bucket_id, segment)) {
            Some(reader) => {
                let res = reader.read_record(offset, false);
                let panic_res = catch_unwind(AssertUnwindSafe(move || handler(res)));
                if let Err(err) = panic_res {
                    error!("reader panicked: {err:?}");
                }
                None
            }
            None => {
                error!("tried to read from unknown bucket segment in threadpool");
                Some(handler)
            }
        }
    }
}

struct ReadRequest {
    bucket_id: u16,
    segment: u32,
    offset: u64,
    handler: ReadFn,
}

Any advice or suggestions would be greatly appreciated!

[adamreichold]

I could use thread_local, but that wouldn’t help if the state needs external updates.

If you have a way of dispatching arbitrary code to worker threads, there is an old trick to do this. The worker keeps everything in thread-local storage and you do not use a separate channel for the state updates, but the same injector which also dispatches other work items. The work item then runs on the worker thread accessing the thread-local state and updating the state.

Of course, this implies being able to dispatch something to specific or at at least all workers, c.f. e.g. rayon::broadcast. Additionally, you might want to use work item priority to ensure the same latency properties as with the separate channel.

[tqwewe]

So I actually gave this a try, pretty much the exact same:

thread_local! {
    static READERS: RefCell<HashMap<(u16, u32), BucketSegmentReader>> = RefCell::new(HashMap::new());
}

/// Adds a bucket segment reader to all workers in the thread pool.
pub fn add_bucket_segment(
    &self,
    bucket_id: u16,
    segment: u32,
    reader: &BucketSegmentReader,
) -> io::Result<()> {
    self.pool.broadcast(|_| {
        READERS.with_borrow_mut(|readers| {
            readers.insert((bucket_id, segment), reader.try_clone().unwrap())
        });
    });

    Ok(())
}

/// Reads a record, utilizing a thread from the thread pool.
pub fn read_record(
    &self,
    bucket_id: u16,
    segment: u32,
    offset: u64,
    handler: impl FnOnce(io::Result<Option<Record>>) + Send + 'static,
) {
    self.pool.install(move || {
        READERS.with_borrow_mut(|readers| match readers.get_mut(&(bucket_id, segment)) {
            // perform read...
        });
    });
}

However, doing a benchmark comparing reads with 8 workers from my old crossbeam implementation (with thread::yield_now()) vs rayon ThreadPool (the code above), I got the following results:

Benchmark	Time	vs Single-threaded
Single-threaded Reads	2.78 µs	1x
Multi-threaded Reads Crossbeam Pool	1.61 µs	1.7x (faster)
Multi-threaded Reads Rayon Pool	15.88 µs	5.6x (slower)

As you can see in the results, it crossbeam gives a decent performance increase vs single-threaded reads. But using rayon gave a huge slow down for some reason.

[tqwewe]

Ah! I've realised this is because I'm using install on rayon's ThreadPool, which blocks until the operation completes. Whereas on my crossbeam implementation it wasn't blocking.

A simple switch from pool.install(...) to pool.spawn(...) solved this and gave me the same performance as the crossbeam implementation – very happy!