Crate rayon_iter_concurrent_limit 

Limit the concurrency of a rayon parallel iterator.

Example

The concurrent_limit method of the ConcurrentLimit extension trait limits the concurrency of everything chained after it.

use rayon::iter::{IntoParallelIterator, ParallelIterator};
use rayon_iter_concurrent_limit::ConcurrentLimit;
const N: usize = 1000;
let output = (0..100)
    .into_par_iter()
    .concurrent_limit(2) // limits everything chained after it
    .map(|i| {
        let alloc = vec![i; N]; // max of 2 concurrent allocations
        alloc.into_par_iter().sum::<usize>() // runs on all threads
    })
    .map(|alloc_sum| {
        alloc_sum / N // max of 2 concurrent executions
    })
    .collect::<Vec<usize>>();
assert_eq!(output, (0..100).collect::<Vec<usize>>());

Motivation

Consider this example:

use rayon::iter::{IntoParallelIterator, ParallelIterator};
let op = |_: usize| {
    // operation involving a large allocation
};
(0..100).into_par_iter().for_each(op);

In this case, it may be necessary to limit the number of concurrent executions of op due to memory constraints. The number of threads could be limited with rayon::ThreadPool::install like so:

let thread_pool = rayon::ThreadPoolBuilder::new().num_threads(1).build()?;
thread_pool.install(|| {
    (0..100).into_par_iter().for_each(op);
});

However, this constrains more than intended and has a footgun. Any parallel operations within op use the same thread-limited pool, and the iterator must be consumed inside the install scope or it will not use that pool at all. Calling install internally with a different pool avoids the first problem but introduces a worse one: op can then yield, so multiple instances of op may run concurrently on a single thread, as detailed here in the install documentation.

How it works

This crate provides ConcurrentLimit, an extension trait implemented for every rayon::iter::IndexedParallelIterator. Its single method, concurrent_limit, limits the concurrency of every subsequent method in the chain, while parallel operations within the supplied function continue to use the whole thread pool.

Concurrency is limited by reducing the number of work items available to rayon, so that a chained operation runs sequentially within a work item but in parallel across them. concurrent_limit splits the iterator into exactly concurrent_limit pieces of near-equal size. Nothing is allocated and items are consumed lazily.

Reaching the limit exactly requires this crate to drive the iterator itself, via rayon::iter::plumbing. rayon’s own driver always splits a producer at its midpoint and decides whether to split with a boolean, which rounds the number of work items to a power of two; splitting proportionally to a target piece count instead hits any piece count exactly.

ConcurrencyLimited is an IndexedParallelIterator, so indexed methods such as zip, enumerate, and collect_into_vec remain available. Whether the limit stays exact depends on what is chained after it:

• Adaptors that merely wrap the consumer — map, filter, filter_map, update, inspect, cloned, copied — leave the split to ConcurrencyLimited, so the limit stays exact. So do the methods that consume the iterator, such as for_each, collect, collect_into_vec, and reduce.
• Adaptors that need a Producer of their own — zip, enumerate, rev, skip, take, step_by, chunks — take over the splitting. The limit then degrades from exact to an upper bound, applied as an IndexedParallelIterator::with_min_len hint: rayon never splits below the minimum length, but because it splits by halving, the concurrency can undershoot the limit (as low as approximately half of it).

Interaction with with_min/max_len or a second concurrent_limit

A minimum length set further up the chain is a hard floor on the size of a work item, so it caps the number of work items at len / min_len. concurrent_limit honours that floor, which means the tightest constraint in the chain wins and the limit is always an upper bound, never an override:

// 64 items, at least 32 per work item, so at most 2 work items — not 16.
(0..64)
    .into_par_iter()
    .with_min_len(32)
    .concurrent_limit(16)
    .for_each(|_i| {
        // at most 2 concurrent executions
    });

The same rule governs two concurrent_limit calls in one chain. The first one degrades to its with_min_len fallback (an adaptor downstream of it took the producer), so the operations between the two calls keep the tighter of the two limits:

let _output = (0..64)
    .into_par_iter()
    .concurrent_limit(2)
    .map(|i| {
        i // at most 2 concurrent executions, not 8
    })
    .concurrent_limit(8)
    .map(|i| {
        i // also at most 2: the tighter limit upstream still applies
    })
    .collect::<Vec<usize>>();

Relaxing a limit part-way through a chain is therefore not possible; split the chain into two separate iterators instead. Note that with_max_len is not honoured, since asking for smaller work items is the opposite of what concurrent_limit is for.

Alternatives

Chunking the iterator is the closest equivalent without this crate, and is what earlier versions of this crate did:

let sum_iter = (0..100)
    .into_par_iter()
    .chunks(100_usize.div_ceil(2))
    .flat_map_iter(|chunk| chunk)
    .map(op);

That form allocates a Vec per chunk, produces an unindexed iterator, and can fall short of the requested concurrency because the chunk size is rounded up.

IndexedParallelIterator::by_uniform_blocks (rayon 1.9.0) can also bound the number of concurrent executions of an operation. iterator.by_uniform_blocks(limit) processes blocks of limit items sequentially, with parallelism within each block. However, every block ends with a synchronisation point, so a single slow item stalls the entire pipeline at each block boundary. The approach of this crate lets each of the limit concurrent streams proceed independently, which suits expensive operations with variable cost.

Structs

ConcurrencyLimited

A parallel iterator which yields the items of its base iterator, split into exactly concurrent_limit work items.

Traits

ConcurrentLimit

An extension trait which limits the concurrency of an iterator chain.