Login

u/gilfyole

r/rust r/learnrust
tokio for I/O, rayon for CPU: how we bridge them in a Rust search engine
▲ 131 r/learnrust+1 crossposts

tokio for I/O, rayon for CPU: how we bridge them in a Rust search engine

Disclosure: I work on infino, an Apache-2.0 embedded retrieval engine in Rust. This is an internals post about how we split work between tokio and rayon, and a couple of spots where we got it wrong before benchmarks caught it.

Quick context: infino stores data in superfiles, standard Parquet files with a BM25 index and a vector index embedded just before the footer, so they stay fully readable by any normal Parquet reader (more details in the storage-format post). A supertable is a manifest referencing many append-only, snapshot-isolated superfiles.

A query against a supertable opens up some of those superfiles from object storage. That work splits into two different jobs. Opening the files, sending GET ranges to S3, Azure, or disk, and prefetching tombstone sidecars is I/O: awaitable, and none of it CPU heavy. Decoding the Parquet pages, scoring BM25 postings, computing vector distances, and reranking is CPU work: synchronous, and wanting a core to itself for a few milliseconds at a time.

Put both the above tasks on one tokio runtime and you have a problem: tokio's scheduler only yields at .await points, so a CPU-bound task blocks everything else on that worker thread until it finishes. spawn_blocking gets you out of that, but it hands out one OS thread per task, not the fixed-size, work-stealing pool you actually want for chunked parallel compute. To resolve this problem, our approach uses both tokio and rayon: tokio owns I/O, rayon owns CPU.

Query fan-out is one tokio::spawn per superfile on a shared multi-thread runtime, joined with try_join_all (supertable/query/dispatch.rs):

let handles = units.into_iter().map(|(entry, params)| {
    let store = Arc::clone(&store);
    let handle = tokio::spawn(async move {
        let r = open_reader(&store, disk_cache.as_ref(), storage.as_ref(), &entry).await?;
        body(r, entry, tombstone_cache, now, params).await
    });
    async move { handle.await.map_err(|e| QueryError::Store(format!("fan-out task join: {e}")))? }
});
try_join_all(handles).await

CPU work runs on a pair of rayon pools sized to the machine, one for reads and one for writes, built once and shared by every open Supertable in the process rather than per-handle or falling back to rayon's global pool (supertable/options.rs):

fn default_reader_thread_count() -> usize {
    num_cpus::get().max(1)
}

static SHARED_READER_POOL: OnceLock<Arc<ThreadPool>> = OnceLock::new();

fn shared_reader_pool() -> Arc<ThreadPool> {
    Arc::clone(SHARED_READER_POOL.get_or_init(|| {
        Arc::new(
            ThreadPoolBuilder::new()
                .num_threads(default_reader_thread_count())
                .thread_name(|i| format!("supertable-reader-{i}"))
                .build()
                .expect("invariant: rayon pool build only fails on thread-spawn failure"),
        )
    }))
}

pool.install(|| ... .par_iter() ...) runs shard builds and warm-read Parquet decode on it (supertbale/build.rs, query/exec/common.rs).

Bridging the two with a oneshot channel

The tricky part is the seam between them. An async task that needs CPU work can't call par_iter() inline, because that blocks the tokio worker until it's done and stalls every other task on it. So instead we hand the closure to rayon and await a tokio::sync::oneshot for the result. Here's the coarse-to-fine vector scoring path doing exactly that (superfile/vector/reader.rs):

let (tx, rx) = oneshot::channel();
rayon::spawn(move || {
    let acc = meta_owned
        .par_chunks(chunk)
        .zip(blocks_owned.par_chunks(chunk))
        .map(|(meta_chunk, block_chunk)| score_cluster_codes_into_heap(meta_chunk, block_chunk))
        .reduce(/* ... */);
    let _ = tx.send(acc);
});
let acc = rx.await?;

The tokio worker is free to poll other tasks while rayon does the scoring. The rule: no single thread should ever be both an async-runtime driver and a rayon worker at the same time. Other search engines like Meilisearch have hit exactly this bug in production; their writeup is worth a read (linked below).

Sync code calling back into async

That bridge only goes one way. Infino also needs the reverse: the public API is sync (Supertable::append, bm25_search, vector_search), but the storage layer underneath is async (object_store is an async trait). Calls from rayon threads, the CLI, and the Python bindings all need to drop into async code to do I/O, then hand back a plain value.

runtime_bridge.rs handles this. If there's already a multi_thread runtime running, it uses block_in_place + Handle::block_on. For the rayon-thread case, where there's no ambient runtime, it builds a current_thread runtime and drives the future on that instead.

We got this wrong at first. An earlier version of the code routed every sync caller through one shared multi_thread().worker_threads(1) runtime instead of a current_thread one, on the theory that reusing a runtime beats building one per call. That part's true, but a multi_thread runtime's block_on adds per-poll coordination with its worker thread that current_thread doesn't pay, since it just polls inline with no handoff. The commit that fixed it recorded the cost: +6-17% on multi-term FTS search at 10M docs, worse the more async fan-out a query did; single-term queries were unaffected. Switching back to current_thread fixed it.

Which runtime flavor you pick for this bridge matters as much as which pool you pick for the other one, and it wasn't obvious which one would win until we measured it.

A second tokio runtime instead of rayon

While researching this I ran into Andrew Lamb's (InfluxDB IOx / DataFusion) case for the opposite bridge: a second, dedicated tokio runtime for CPU work, instead of rayon. The point here is that rayon has no idea it's part of an async system, since it has no cancellation and no yielding, while a second tokio runtime is at least built the same way as the first one. This work is still pushing this upstream (tokio#8085, a proposed spawn_compute API).

We haven't hit the problems that argument is solving for. Our compute bursts are short and we don't need to cancel them mid-flight, unlike DataFusion's long-running operators. However, this might be an area of exploration for us in the future.

Where to read the code

Two things that would be interesting to get this sub's take on as we continue exploring:

  1. Rayon+oneshot versus a second tokio runtime for CPU work. For short bursts like ours, does the second runtime's cancellation support actually matter, or is it solving a problem we don't have?
  2. We used to build a reader pool, a writer pool, and a query runtime per Supertable, and it kept one table's queries from starving another's, until you open more than a couple tables and now it's N pools elbowing each other for the same cores. So we merged all three into process-wide singletons. Fixes the oversubscription, but a busy table can now crowd out a quiet one, since there's no isolation between them anymore. Anyone actually solved fair-share on a shared rayon pool? Rayon has no notion of priority itself, so my guess is whatever works has to live above the pool, not inside it. Curious if that's right.

References:

(Disclosure again: infino is Apache-2.0 OSS, and there's a commercial hosted version coming. This post is about the engineering, specifically about the use of tokio and rayon in the core engine.)

u/gilfyole — 4 days ago