forest/utils/stream.rs
1// Copyright 2019-2026 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3use futures::{Stream, StreamExt};
4use tokio_util::task::AbortOnDropHandle;
5
6/// Decouple stream generation and stream consumption into separate threads,
7/// keeping not-yet-consumed elements in a bounded queue. This is similar to
8/// [`stream::buffered`](https://docs.rs/futures/latest/futures/stream/trait.StreamExt.html#method.buffered)
9/// and
10/// [`sink::buffer`](https://docs.rs/futures/latest/futures/sink/trait.SinkExt.html#method.buffer).
11/// The key difference is that [`par_buffer`] is parallel rather than concurrent
12/// and will make use of multiple cores when both the stream and the stream
13/// consumer are CPU-bound. Because a new thread is spawned, the stream has to
14/// be [`Sync`], [`Send`] and `'static`.
15pub fn par_buffer<V: Send + Sync + 'static>(
16 cap: usize,
17 stream: impl Stream<Item = V> + Send + Sync + 'static,
18) -> (
19 impl Stream<Item = V>,
20 AbortOnDropHandle<Result<(), flume::SendError<V>>>,
21) {
22 let (send, recv) = flume::bounded(cap);
23 let handle =
24 AbortOnDropHandle::new(tokio::task::spawn(stream.map(Ok).forward(send.into_sink())));
25 (recv.into_stream(), handle)
26}