Hand CPU work from async tasks to a thread pool

You'll build a small prime-counting service where async request handlers hand their work to a pool of plain OS threads over one mpmc channel. The handlers use the async end of the channel, the worker threads use the sync end and a separate sync thread submits a backlog in one batch. Along the way the workers free their memory when idle with recv_timeout and you'll swap the bounded queue for an unbounded one to see what changes.

fibre = "0"

Set up the project#

Start a new binary crate:

cargo new prime-service
cd prime-service

Add fibre and Tokio to Cargo.toml. The service runs Tokio's single-threaded runtime and uses its timers. It doesn't need the macros feature because main builds the runtime by hand.

[dependencies]
fibre = "0"
tokio = { version = "1", features = ["rt", "time"] }

The problem this solves is a common one. An async runtime runs many tasks on a few threads and each task is expected to yield often. A task that spends 100ms counting primes holds its runtime thread for all of that time and every other task on that thread waits. The fix is to move the heavy work onto threads you own and give the async side a way to hand jobs over and wait for answers without blocking. An mpmc channel does the handing over: many handlers send, many worker threads receive and each job goes to exactly one worker.

Send a job from an async task to a thread#

Start with one worker thread and three requests. Replace src/main.rs with this:

use fibre::{mpmc, oneshot};
use std::thread;
use std::time::{Duration, Instant};

struct Job {
    below: usize,
    reply: oneshot::Sender<usize>,
}

fn ms(start: Instant) -> u128 {
    (start.elapsed().as_millis() + 50) / 100 * 100
}

fn count_primes(below: usize, sieve: &mut Vec<bool>) -> usize {
    sieve.clear();
    sieve.resize(below, true);
    let mut count = 0;
    for i in 2..below {
        if sieve[i] {
            count += 1;
            let mut multiple = i * i;
            while multiple < below {
                sieve[multiple] = false;
                multiple += i;
            }
        }
    }
    count
}

fn worker(job_rx: mpmc::Receiver<Job>) -> usize {
    let mut sieve = Vec::new();
    let mut jobs = 0;
    while let Ok(job) = job_rx.recv() {
        thread::sleep(Duration::from_millis(100));
        let count = count_primes(job.below, &mut sieve);
        let _ = job.reply.send(count);
        jobs += 1;
    }
    jobs
}

async fn handle(id: usize, below: usize, job_tx: mpmc::AsyncSender<Job>, start: Instant) -> String {
    let (reply_tx, reply_rx) = oneshot::oneshot();
    if job_tx.send(Job { below, reply: reply_tx }).await.is_err() {
        return format!("request {id}: no workers left");
    }
    match reply_rx.recv().await {
        Ok(count) => format!("{:>4}ms request {id}: {count} primes below {below}", ms(start)),
        Err(_) => format!("request {id}: worker dropped the job"),
    }
}

async fn serve(job_tx: mpmc::AsyncSender<Job>, start: Instant) {
    let mut handlers = Vec::new();
    for (id, below) in [100, 1_000, 10_000].into_iter().enumerate() {
        handlers.push(tokio::spawn(handle(id, below, job_tx.clone(), start)));
    }
    drop(job_tx);

    for handler in handlers {
        println!("{}", handler.await.unwrap());
    }
}

fn main() {
    let (job_tx, job_rx) = mpmc::bounded::<Job>(4);
    let start = Instant::now();

    let worker = thread::spawn(move || worker(job_rx));

    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_time()
        .build()
        .unwrap();
    runtime.block_on(serve(job_tx.to_async(), start));

    let jobs = worker.join().unwrap();
    println!("worker: queue closed after {jobs} jobs");
}

mpmc::bounded(4) returns a sync Sender and a sync Receiver for a queue that holds at most 4 jobs. The worker thread keeps the sync receiver: recv blocks the thread until a job arrives, which is fine because the thread has nothing else to do. The handlers need the async side, so main passes job_tx.to_async() into the runtime. to_async turns the sync sender into an AsyncSender on the same channel at no cost. Its send returns a future, so a handler waiting for queue space yields to the runtime instead of blocking it.

Each job carries a oneshot::Sender for its answer. The handler keeps the matching receiver and awaits it, so it gets its own result back even when several handlers share one queue.

count_primes is a sieve that reuses the worker's Vec<bool> from job to job. At these sizes it finishes in well under a millisecond, so the worker sleeps 100ms first to stand in for real CPU work and keep the timings easy to read. ms rounds the elapsed time to the nearest 100ms.

main is a plain function. It owns the worker threads, runs the async part with block_on and joins the threads once the runtime is done. Run it:

$ cargo run
 100ms request 0: 25 primes below 100
 200ms request 1: 168 primes below 1000
 300ms request 2: 1229 primes below 10000
worker: queue closed after 3 jobs

One worker takes the jobs one at a time in the order they were sent. The worker exits on its own: each handler owns a clone of the sender and drops it when it returns, serve drops the original and once the last sender is gone recv returns Err(RecvError::Disconnected).

Run a pool of worker threads#

The sync Receiver is Clone. Every clone pulls from the same queue and each job goes to whichever worker asks first, so a pool is a loop of clones. Double the requests in serve:

    for (id, below) in [100, 1_000, 10_000, 50, 500, 5_000].into_iter().enumerate() {

Then replace main with a version that starts three workers:

fn main() {
    let (job_tx, job_rx) = mpmc::bounded::<Job>(4);
    let start = Instant::now();

    let mut workers = Vec::new();
    for _ in 0..3 {
        let job_rx = job_rx.clone();
        workers.push(thread::spawn(move || worker(job_rx)));
    }
    drop(job_rx);

    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_time()
        .build()
        .unwrap();
    runtime.block_on(serve(job_tx.to_async(), start));

    for (n, worker) in workers.into_iter().enumerate() {
        let jobs = worker.join().unwrap();
        println!("worker {n}: queue closed after {jobs} jobs");
    }
}

 100ms request 0: 25 primes below 100
 100ms request 1: 168 primes below 1000
 100ms request 2: 1229 primes below 10000
 200ms request 3: 15 primes below 50
 200ms request 4: 95 primes below 500
 200ms request 5: 669 primes below 5000
worker 0: queue closed after 2 jobs
worker 1: queue closed after 2 jobs
worker 2: queue closed after 2 jobs

Three workers finish six jobs in 200ms. Which worker gets which job changes from run to run, which is why the handlers print the results and not the workers. Each worker still takes exactly two because a job keeps a worker busy for 100ms and the queue always has another job ready when it comes back.

Six requests don't fit in a queue of 4. They all get in here because the workers take the first three off the queue almost at once. With more requests the extra handlers would wait in send until a worker freed a slot. They wait as async tasks, so the runtime keeps serving everything else in the meantime.

The drop(job_rx) after the loop matters. The original receiver would otherwise stay alive in main and keep the channel open. Senders would never see it as closed even after every worker had gone.

Free idle memory with recv_timeout#

Each worker keeps its sieve between jobs so it doesn't allocate every time. After a big job that buffer can be large. A worker that sits idle for minutes keeps it all that time. recv_timeout waits for a job for a limited time and returns RecvErrorTimeout::Timeout if none comes, which gives the worker a moment to tidy up. Add the import at the top:

use fibre::error::RecvErrorTimeout;

Then replace worker with a version that frees the sieve after 300ms without work and reports what it did:

#[derive(Default)]
struct Stats {
    jobs: usize,
    released: usize,
}

fn worker(job_rx: mpmc::Receiver<Job>) -> Stats {
    let mut sieve = Vec::new();
    let mut stats = Stats::default();
    loop {
        match job_rx.recv_timeout(Duration::from_millis(300)) {
            Ok(job) => {
                thread::sleep(Duration::from_millis(100));
                let count = count_primes(job.below, &mut sieve);
                let _ = job.reply.send(count);
                stats.jobs += 1;
            }
            Err(RecvErrorTimeout::Timeout) => {
                if sieve.capacity() > 0 {
                    sieve = Vec::new();
                    stats.released += 1;
                }
            }
            Err(RecvErrorTimeout::Disconnected) => return stats,
        }
    }
}

To give the workers a quiet spell, let the second three requests arrive 500ms after the first. Give handle a delay:

async fn handle(id: usize, delay: u64, below: usize, job_tx: mpmc::AsyncSender<Job>, start: Instant) -> String {
    tokio::time::sleep(Duration::from_millis(delay)).await;

In serve, replace the for line with a list of (delay, below) pairs:

    let requests = [(0, 100), (0, 1_000), (0, 10_000), (500, 50), (500, 500), (500, 5_000)];
    for (id, (delay, below)) in requests.into_iter().enumerate() {
        handlers.push(tokio::spawn(handle(id, delay, below, job_tx.clone(), start)));

Finally print the stats in main:

        let stats = worker.join().unwrap();
        println!("worker {n}: jobs={} releases={}", stats.jobs, stats.released);

 100ms request 0: 25 primes below 100
 100ms request 1: 168 primes below 1000
 100ms request 2: 1229 primes below 10000
 600ms request 3: 15 primes below 50
 600ms request 4: 95 primes below 500
 600ms request 5: 669 primes below 5000
worker 0: jobs=2 releases=1
worker 1: jobs=2 releases=1
worker 2: jobs=2 releases=1

The workers finish the first batch at 100ms and wait. At 400ms their 300ms timeout runs out and each one drops its sieve. The second batch arrives at 500ms, cuts the next wait short and finishes at 600ms. Then the last handler returns, the channel disconnects and every waiting worker gets Disconnected right away. They don't sit out the rest of their timeout.

Submit a backlog from a sync thread#

Not every sender is async. Say the service also has a backfill job that runs on its own thread at startup and needs nine counts. It can use the sync Sender directly and hand all nine jobs over in one call with send_batch, which puts the whole batch through under one lock and wakes the workers together. Add this function:

fn backfill(job_tx: mpmc::Sender<Job>, start: Instant) {
    let mut replies = Vec::new();
    let mut jobs = Vec::new();
    for k in 1..=9 {
        let (reply_tx, reply_rx) = oneshot::oneshot();
        jobs.push(Job { below: k * 1_000, reply: reply_tx });
        replies.push(reply_rx);
    }

    if job_tx.send_batch(jobs).is_err() {
        println!("backfill: no workers left");
        return;
    }
    println!("{:>4}ms backfill: queued 9 jobs", ms(start));
    drop(job_tx);

    let mut total = 0;
    for reply_rx in replies {
        total += reply_rx.recv_blocking().unwrap_or(0);
    }
    println!("{:>4}ms backfill: done, {total} primes counted", ms(start));
}

send_batch on a sync sender blocks until every job is in the queue. It returns Ok(9) when they all got in. If the workers all go away partway through, it returns a SendBatchError with the count that was sent and the jobs that weren't. The thread then waits for each answer with the oneshot's recv_blocking, the sync counterpart of the recv().await the handlers use.

In main, clone a sync sender for the backfill thread before the original is turned into an async one, then join the thread after the runtime finishes:

    let backfill_tx = job_tx.clone();
    let backfiller = thread::spawn(move || backfill(backfill_tx, start));

    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_time()
        .build()
        .unwrap();
    runtime.block_on(serve(job_tx.to_async(), start));
    backfiller.join().unwrap();

Move the requests later so they arrive after the backlog is done and the workers have been idle for a while:

    let requests = [(800, 100), (800, 1_000), (800, 10_000)];

 100ms backfill: queued 9 jobs
 300ms backfill: done, 5927 primes counted
 900ms request 0: 25 primes below 100
 900ms request 1: 168 primes below 1000
 900ms request 2: 1229 primes below 10000
worker 0: jobs=4 releases=1
worker 1: jobs=4 releases=1
worker 2: jobs=4 releases=1

Nine jobs don't fit in a queue of 4. The workers take three and four fill the queue, so send_batch blocks the backfill thread with two jobs left. At 100ms the workers take three more jobs and the last two get in. That's backpressure on a sync sender: the thread blocks, where an async handler would have yielded. The workers finish the backlog at 300ms, free their sieves at 600ms and serve the three requests at 800ms.

The channel now has three kinds of handle on it at once: an async sender in each handler, a sync sender on the backfill thread and a sync receiver on each worker. They are all the same queue.

Try the unbounded channel#

mpmc::unbounded has no capacity limit. A send never waits, so the backfill thread would never block and neither would a handler. The unbounded handles take &mut self for sending and for waiting receives, because each one keeps some per-handle state to avoid locking. Every thread and task in this program already owns its own clone, so the switch is a new constructor and a mut on each handle parameter:

fn worker(mut job_rx: mpmc::UnboundedSyncReceiver<Job>) -> Stats {

fn backfill(mut job_tx: mpmc::UnboundedSyncSender<Job>, start: Instant) {

async fn handle(id: usize, delay: u64, below: usize, mut job_tx: mpmc::UnboundedAsyncSender<Job>, start: Instant) -> String {

async fn serve(job_tx: mpmc::UnboundedAsyncSender<Job>, start: Instant) {

    let (job_tx, job_rx) = mpmc::unbounded::<Job>();

Nothing else changes. to_async, clone, send_batch, recv_timeout and the batch error all work the same way on the unbounded handles.

   0ms backfill: queued 9 jobs
 300ms backfill: done, 5927 primes counted
 900ms request 0: 25 primes below 100
 900ms request 1: 168 primes below 1000
 900ms request 2: 1229 primes below 10000
worker 0: jobs=4 releases=1
worker 1: jobs=4 releases=1
worker 2: jobs=4 releases=1

The only difference is the first line. send_batch returns at once and six jobs sit in the queue until workers come for them. The work still takes as long as the workers need, so the backfill is done at 300ms either way.

That makes the choice about where waiting is allowed to happen. With a bounded queue a burst of requests makes the senders wait, which is how the service tells its callers to slow down. With an unbounded queue a burst turns into memory: every job that can't be served yet sits in the queue. Nothing stops the queue growing if requests keep coming faster than the workers can clear them. Use bounded for a queue fed by outside requests. Unbounded fits a sender that must never block, such as a callback from a library that doesn't expect to wait, where you know the total amount of work is limited.

Where to go next#

The channels guide at https://github.com/excsn/fibre/blob/main/channels/README.GUIDE.md covers the rest of mpmc, including try_send, try_recv, the in-place batch calls send_batch_mut and recv_batch_mut and the zero-capacity mpmc::rendezvous channel, along with the other channel types and the full disconnect rules.

The full API, including every handle type and error, is on docs.rs at https://docs.rs/fibre.