Rayon and Tokio together

Tokio runs async code on a few threads, and it can switch from one task to another only at an .await. rayon runs work on its own pool of threads: par_iter() hands the pieces of a collection to those threads and returns when they are all done. Put the two together and one question matters: which thread is waiting?

RememberWork with no .await in it, a parallel loop included, should not run for long on a runtime thread. A blocking call goes to spawn_blocking. Parallel CPU work goes to rayon::spawn, and the async side waits on a oneshot channel.

Project layout

rayon-tokio/
├── src/
│   └── main.rs                 the program on this page
├── Cargo.lock                  created by Cargo
└── Cargo.toml                  dependencies go here

RustCargo.toml

[package]
name = "rayon-tokio"
version = "0.1.0"
edition = "2024"

[dependencies]
rayon = "1.12.0"
tokio = { version = "1.54.1", features = ["full"] }

Rustsrc/main.rs

use rayon::prelude::*;
use tokio::sync::oneshot;

#[derive(Debug)]
struct Reading {
    sensor: String,
    value: i64,
}

fn parse_line(line: &str) -> Result<Reading, String> {
    let (sensor, value) = line.split_once(',').ok_or(format!("{line:?}: no comma"))?;
    let value = value
        .trim()
        .parse::<i64>()
        .map_err(|e| format!("{line:?}: {e}"))?;
    Ok(Reading {
        sensor: sensor.trim().to_string(),
        value,
    })
}

// CPU-bound work goes to rayon's pool; the async fn only waits on a oneshot channel.
async fn parse_all(lines: Vec<String>) -> Result<Vec<Reading>, String> {
    let (tx, rx) = oneshot::channel();
    rayon::spawn(move || {
        let result: Result<Vec<Reading>, String> =
            lines.par_iter().map(|l| parse_line(l)).collect();
        let _ = tx.send(result);
    });
    rx.await.expect("the rayon job dropped its sender")
}

#[tokio::main]
async fn main() {
    let requests: Vec<Vec<String>> = vec![
        vec!["a,1".into(), "b,2".into(), "c,3".into()],
        vec!["d,10".into(), "e,20".into()],
        vec!["f,7".into(), "g,oops".into()],
    ];

    // three "requests" handled at once, as three tasks on the Tokio runtime
    let handles: Vec<_> = requests
        .into_iter()
        .map(|lines| tokio::spawn(parse_all(lines)))
        .collect();
    for (i, handle) in handles.into_iter().enumerate() {
        match handle.await.unwrap() {
            Ok(readings) => {
                let total: i64 = readings.iter().map(|r| r.value).sum();
                let names: Vec<&str> = readings.iter().map(|r| r.sensor.as_str()).collect();
                println!("request {i}: {} total {total}", names.join(","));
            }
            Err(e) => println!("request {i}: error: {e}"),
        }
    }

    // a blocking call (a file read) goes to Tokio's blocking pool instead
    let lines = tokio::task::spawn_blocking(|| {
        std::fs::read_to_string("Cargo.toml").map(|s| s.lines().count())
    })
    .await
    .unwrap()
    .unwrap();
    println!("Cargo.toml has {lines} lines");
}

$ cargo run

request 0: a,b,c total 6
request 1: d,e total 30
request 2: error: "g,oops": invalid digit found in string
Cargo.toml has 8 lines

Gomain.go

package main

import (
  "fmt"
  "os"
  "strconv"
  "strings"
  "sync"
)

type reading struct {
  sensor string
  value  int
}

func parseLine(line string) (reading, error) {
  sensor, value, ok := strings.Cut(line, ",")
  if !ok {
    return reading{}, fmt.Errorf("%q: no comma", line)
  }
  n, err := strconv.Atoi(strings.TrimSpace(value))
  if err != nil {
    return reading{}, fmt.Errorf("%q: %w", line, err)
  }
  return reading{strings.TrimSpace(sensor), n}, nil
}

// One goroutine per line. Each writes to its own slot, so the order is kept.
func parseAll(lines []string) ([]reading, error) {
  out := make([]reading, len(lines))
  errs := make([]error, len(lines))
  var wg sync.WaitGroup
  for i, line := range lines {
    wg.Add(1)
    go func() {
      defer wg.Done()
      out[i], errs[i] = parseLine(line)
    }()
  }
  wg.Wait()
  for _, err := range errs {
    if err != nil {
      return nil, err
    }
  }
  return out, nil
}

func main() {
  requests := [][]string{
    {"a,1", "b,2", "c,3"},
    {"d,10", "e,20"},
    {"f,7", "g,oops"},
  }

  // three "requests" handled at once
  type result struct {
    readings []reading
    err      error
  }
  results := make([]result, len(requests))
  var wg sync.WaitGroup
  for i, lines := range requests {
    wg.Add(1)
    go func() {
      defer wg.Done()
      results[i].readings, results[i].err = parseAll(lines)
    }()
  }
  wg.Wait()
  for i, r := range results {
    if r.err != nil {
      fmt.Printf("request %d: error: %v\n", i, r.err)
      continue
    }
    total := 0
    names := []string{}
    for _, rd := range r.readings {
      total += rd.value
      names = append(names, rd.sensor)
    }
    fmt.Printf("request %d: %s total %d\n", i, strings.Join(names, ","), total)
  }

  // a blocking call is just a call: the goroutine waits, the others keep running
  data, err := os.ReadFile("go.mod")
  if err != nil {
    panic(err)
  }
  fmt.Printf("go.mod has %d lines\n", strings.Count(string(data), "\n"))
}

$ go run .

request 0: a,b,c total 6
request 1: d,e total 30
request 2: error: "g,oops": strconv.Atoi: parsing "oops": invalid syntax
go.mod has 3 lines

rayon::spawn(move || ...) puts the closure in rayon’s job queue and returns. The closure must be Send + 'static (docs), the bounds thread::spawn has (Send and Sync), so lines is moved in.
oneshot carries one value back. Its docs say “Since the send method is not async, it can be used anywhere”, so a rayon thread can call send, and rx.await makes the task wait without holding a thread. The Tokio docs give this as the hint for rayon.
spawn_blocking runs a closure on a thread from Tokio’s separate blocking pool. The docs say it is for “non-async operations that eventually finish on their own”, and the default limit on those threads is 512 (max_blocking_threads). Here it reads a file.
The handles are awaited in order, so the output is in request order. The Go version gets the same effect with one result slot per request. In an axum handler, which is an async fn too, you would call parse_all(...).await the same way (axum).
Is it worth it? Alice Ryhl’s article on blocking, about a sum, says that “unless the array is very very large” computing it “is probably cheap enough that you can just do it directly in Tokio”. The three-line inputs here are only to keep the output short. Neither source gives a size to switch at, so measure your own workload.

Two pools, two sizes

Each pool sizes itself. Tokio’s default is “the number of cores available to the system”, from available_parallelism (I read that call in the tokio 1.54.1 source; the Tokio docs only say “cores”), and rayon’s is the number of logical CPUs, unless TOKIO_WORKER_THREADS or RAYON_NUM_THREADS is set (num_threads). The run asserts that both defaults matched in this image, then sets 2 and 3 by hand and shows par_iter staying on the rayon threads. Neither pool sizes itself from the other, so on a busy N-core machine there can be N runtime workers and N rayon threads wanting to run. I did not measure what that costs.

Rustsrc/main.rs

use rayon::prelude::*;
use std::thread;

fn thread_name() -> String {
    thread::current().name().unwrap_or("?").to_string()
}

fn main() {
    // Defaults: each pool sizes itself from the core count, and neither asks the other.
    let cores = thread::available_parallelism().unwrap().get();
    let default_rt = tokio::runtime::Runtime::new().unwrap();
    assert_eq!(default_rt.metrics().num_workers(), cores);
    assert_eq!(rayon::current_num_threads(), cores);
    println!("by default, Tokio workers and rayon threads each match available_parallelism");
    drop(default_rt);

    // Set both by hand: two numbers you chose, two sets of threads.
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(2)
        .thread_name("io")
        .build()
        .unwrap();
    let pool = rayon::ThreadPoolBuilder::new()
        .num_threads(3)
        .thread_name(|i| format!("cpu-{i}"))
        .build()
        .unwrap();
    println!("Tokio workers: {}", rt.metrics().num_workers());
    println!("rayon threads: {}", pool.current_num_threads());

    // Where does code run? An async task runs on a Tokio worker; par_iter inside
    // install() runs on the rayon pool's threads.
    let task = rt.block_on(async { tokio::spawn(async { thread_name() }).await.unwrap() });
    let mut rayon_names: Vec<String> =
        pool.install(|| (0..100).into_par_iter().map(|_| thread_name()).collect());
    rayon_names.sort();
    rayon_names.dedup();
    println!("async task ran on a thread named {task:?}");
    println!(
        "par_iter ran only on cpu-* threads: {}",
        rayon_names.iter().all(|n| n.starts_with("cpu-"))
    );
    assert_eq!(task, "io");
    assert!(rayon_names.iter().all(|n| n.starts_with("cpu-")) && rayon_names.len() <= 3);
}

$ cargo run

by default, Tokio workers and rayon threads each match available_parallelism
Tokio workers: 2
rayon threads: 3
async task ran on a thread named "io"
par_iter ran only on cpu-* threads: true

For CPU-bound work the Tokio docs advise “a separate thread pool dedicated to CPU bound tasks”, rayon for example, if you want to limit the threads it uses. The rayon docs warn the default may change and say to set num_threads if you rely on a fixed number. Neither gives a ratio between the two pools.

Dropping the receiver does not stop the job

If the request is cancelled, the task awaiting the oneshot is dropped, and with it the receiver. The rayon job is not told. The run drops the receiver after a job’s first of five chunks. The first job ignores that and does all five. The second asks tx.is_closed() before each chunk, which the docs say is true once the receiver is dropped, and stops after one. The halfway, go and finished channels in the code (a std::sync::mpsc channel and two oneshots) only force that order. Compare dropping a future. The Tokio docs say a started spawn_blocking task cannot be aborted either.

Rustsrc/main.rs

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, mpsc};
use tokio::sync::oneshot;

// A rayon job that works through 5 chunks. After chunk 1 it tells us, then waits for a
// go-ahead (this only forces the order: the receiver is dropped while the job is mid-way).
// `polite` jobs look at `tx.is_closed()` before each chunk; the others do not.
async fn run(polite: bool) -> usize {
    let done = Arc::new(AtomicUsize::new(0));
    let (tx, rx) = oneshot::channel::<usize>();
    let (halfway_tx, halfway) = oneshot::channel::<()>();
    let (finished_tx, finished) = oneshot::channel::<()>();
    let (go_tx, go) = mpsc::channel::<()>(); // the job blocks on this one, which is fine on a rayon thread

    let counter = Arc::clone(&done);
    let mut halfway_tx = Some(halfway_tx);
    rayon::spawn(move || {
        for _ in 0..5 {
            if polite && tx.is_closed() {
                break;
            }
            counter.fetch_add(1, Ordering::SeqCst); // "process" one chunk
            if let Some(tx) = halfway_tx.take() {
                tx.send(()).unwrap();
                go.recv().unwrap();
            }
        }
        let sent = tx.send(5);
        println!(
            "  job over: sending the result {}",
            if sent.is_ok() {
                "worked"
            } else {
                "failed, nobody is listening"
            }
        );
        finished_tx.send(()).unwrap();
    });

    halfway.await.unwrap(); // the job has finished chunk 1
    drop(rx); // the async side gives up, like a request that was cancelled
    go_tx.send(()).unwrap();
    finished.await.unwrap(); // wait for the job to end, so the count is final
    done.load(Ordering::SeqCst)
}

#[tokio::main]
async fn main() {
    println!("job that ignores cancellation:");
    let a = run(false).await;
    println!("  chunks processed: {a}");
    println!("job that checks tx.is_closed():");
    let b = run(true).await;
    println!("  chunks processed: {b}");
    assert_eq!((a, b), (5, 1));
}

$ cargo run

job that ignores cancellation:
  job over: sending the result failed, nobody is listening
  chunks processed: 5
job that checks tx.is_closed():
  job over: sending the result failed, nobody is listening
  chunks processed: 1

The same job in your language

Each row except Go moved the same sum (1,000,000 numbers, printed 2999997) off the calling thread, in the image behind this page. The Go row is the parse program above, which prints different lines.

Language	Usual choice	What I ran
Go	goroutines; no separate pool to choose	the Go program above, and the gotcha’s preempted loop
Python	`asyncio` with an executor	`asyncio.to_thread`, and `run_in_executor` with a `ProcessPoolExecutor`
TypeScript	`worker_threads`	a `Worker` started from the same file
Java	`CompletableFuture`, parallel streams	`supplyAsync` and a parallel `LongStream`
Kotlin	coroutines, `Dispatchers.Default`	`withContext(Dispatchers.Default)`, using the kotlinx.coroutines jar in the Kotlin compiler’s `lib` directory (not a project dependency). The same file also runs the gotcha’s loop called directly from a coroutine on `runBlocking`: the log has the Rust order (a Kotlin tab is not in the gotcha above, because its fixed command line has no way to include that jar)
C#	`Task.Run`, PLINQ	`await Task.Run(...)` around `AsParallel().Sum`
C	no async runtime; threads	one `pthread`
C++	`std::async`	`std::async(std::launch::async, ...)` and a future

the same sum, off the calling thread

$ python3 offload.py
process pool: 2999997
to_thread: 2999997
$ tsc --strict --target es2022 --module nodenext --types node offload.ts && node offload.js
worker thread: 2999997
$ javac Offload.java && java Offload
common pool: 2999997
$ java -cp offload.jar:kotlinx-coroutines-core-jvm.jar OffloadKt   # Kotlin, kotlinx.coroutines
called directly: tick 1, sum done, total 2999997, tick 2, tick 3
Dispatchers.Default: 2999997
$ dotnet run
Task.Run + PLINQ: 2999997
$ gcc -Wall -Wextra offload.c -o offload_c -pthread && ./offload_c
pthread: 2999997
$ g++ -std=c++20 -Wall -Wextra offload.cpp -o offload_cpp -pthread && ./offload_cpp
std::async: 2999997

Gotcha

`par_iter` inside an `async fn`

Looks fine, and it compiles in Rust too. But the two give different results.

Same result in Go. The same trap.Not run in Kotlin on this page.Same result in Python. The same trap.Same result in TypeScript. The same trap.Not run in Java on this page.Not run in C# on this page.Not run in C on this page.Not run in C++ on this page.

Rustsrc/main.rs✓ compiles · different result

use rayon::prelude::*;
use std::sync::{Arc, Mutex};

// One runtime thread, so "the runtime is busy" has a single meaning.
#[tokio::main(flavor = "current_thread")]
async fn main() {
    let log = Arc::new(Mutex::new(Vec::new()));

    // A heartbeat task: it records a tick, then lets the runtime run something else.
    let beat = Arc::clone(&log);
    let heartbeat = tokio::spawn(async move {
        for i in 1..=3 {
            beat.lock().unwrap().push(format!("tick {i}"));
            tokio::task::yield_now().await;
        }
    });
    tokio::task::yield_now().await; // let the heartbeat start

    // CPU-bound work, called straight from an async fn.
    let total: u64 = (0..1_000_000u64).into_par_iter().map(|n| n % 7).sum();
    log.lock()
        .unwrap()
        .push(format!("par_iter done, total {total}"));

    heartbeat.await.unwrap();
    let log = log.lock().unwrap();
    for entry in log.iter() {
        println!("{entry}");
    }
    assert_eq!(
        *log,
        ["tick 1", "par_iter done, total 2999997", "tick 2", "tick 3"]
    );
}

package main

import (
  "fmt"
  "runtime"
  "sync/atomic"
)

func main() {
  runtime.GOMAXPROCS(1) // one thread runs every goroutine, like the single-thread runtime

  var ticks atomic.Int32
  go func() { // the heartbeat
    for i := 0; i < 3; i++ {
      ticks.Add(1)
      runtime.Gosched()
    }
  }()

  // CPU-bound work in the main goroutine, with no function call that could yield.
  // It spins until the heartbeat has ticked (the cap is only a safety net).
  n := 0
  for ticks.Load() == 0 && n < 2_000_000_000 {
    n++
  }
  fmt.Println("the heartbeat ticked while the loop was still running:", ticks.Load() > 0)
}

$ go run main.go✓ runs

the heartbeat ticked while the loop was still running: true

There is no Kotlin version of this one.

import asyncio

log = []


async def heartbeat():
    for i in range(1, 4):
        log.append(f"tick {i}")
        await asyncio.sleep(0)  # let the event loop run something else


async def main():
    hb = asyncio.create_task(heartbeat())
    await asyncio.sleep(0)  # let the heartbeat start

    # CPU-bound work, called straight from a coroutine
    total = sum(n % 7 for n in range(1_000_000))
    log.append(f"sum done, total {total}")

    await hb
    print("\n".join(log))


asyncio.run(main())

$ python3 main.py✓ runs

tick 1
sum done, total 2999997
tick 2
tick 3

const log: string[] = [];
const yieldToLoop = () => new Promise<void>((resolve) => setImmediate(resolve));

async function heartbeat() {
  for (let i = 1; i <= 3; i++) {
    log.push(`tick ${i}`);
    await yieldToLoop(); // let the event loop run something else
  }
}

async function main() {
  const hb = heartbeat(); // a promise runs when it is created: tick 1 is logged here

  // CPU-bound work, called straight from an async function
  let total = 0;
  for (let n = 0; n < 1_000_000; n++) total += n % 7;
  log.push(`sum done, total ${total}`);

  await hb;
  console.log(log.join("\n"));
}

main();

$ tsc --strict --module nodenext --types node main.ts && node main.js✓ runs

tick 1
sum done, total 2999997
tick 2
tick 3

There is no Java version of this one.

There is no C# version of this one.

There is no C version of this one.

There is no C++ version of this one.

$ cargo run

tick 1
par_iter done, total 2999997
tick 2
tick 3

What happens in Go

The Go scheduler can preempt a goroutine that is spinning. With GOMAXPROCS(1) the heartbeat ticked while the loop was still running. (A Rust runtime cannot interrupt a task between awaits.)

Why it looks fine