Skip to main content

rskit_stream/operators/
concurrent.rs

1use futures::Stream;
2use rskit_errors::{AppError, AppResult};
3use std::future::Future;
4
5/// Process up to `concurrency` items concurrently.
6///
7/// Output order is not preserved (items complete as they finish).
8/// Uses `futures::stream::FuturesUnordered` for zero-copy task tracking.
9pub fn parallel<S, T, O, F, Fut>(
10    stream: S,
11    concurrency: usize,
12    f: F,
13) -> impl Stream<Item = AppResult<O>> + Send + 'static
14where
15    S: Stream<Item = T> + Send + 'static,
16    T: Send + 'static,
17    O: Send + 'static,
18    F: Fn(T) -> Fut + Clone + Send + Sync + 'static,
19    Fut: Future<Output = AppResult<O>> + Send + 'static,
20{
21    use futures::StreamExt;
22    if concurrency == 0 {
23        futures::stream::once(async {
24            Err(AppError::invalid_input(
25                "concurrency",
26                "must be greater than zero",
27            ))
28        })
29        .left_stream()
30    } else {
31        stream.map(f).buffer_unordered(concurrency).right_stream()
32    }
33}
34
35/// Apply multiple functions to the same item concurrently, collecting results.
36///
37/// All functions are applied to a clone of each item;
38/// all results are collected into a `Vec` in the order of `fns`.
39/// Branch futures are polled concurrently by the stream task
40/// and short-circuit on the first branch error.
41pub fn fan_out<S, T, O, F, Fut>(
42    stream: S,
43    max_branches: usize,
44    fns: Vec<F>,
45) -> impl Stream<Item = AppResult<Vec<O>>> + Send + 'static
46where
47    S: Stream<Item = T> + Send + 'static,
48    T: Clone + Send + 'static,
49    O: Send + 'static,
50    F: Fn(T) -> Fut + Clone + Send + Sync + 'static,
51    Fut: Future<Output = AppResult<O>> + Send + 'static,
52{
53    use futures::StreamExt;
54    if max_branches == 0 || fns.len() > max_branches {
55        let error = AppError::invalid_input(
56            "max_branches",
57            format!("branch count {} exceeds limit {max_branches}", fns.len()),
58        );
59        futures::stream::once(async { Err(error) }).left_stream()
60    } else {
61        let fns = std::sync::Arc::new(fns);
62        stream
63            .then(move |item| {
64                let fns = fns.clone();
65                async move { futures::future::try_join_all(fns.iter().map(|f| f(item.clone()))).await }
66            })
67            .right_stream()
68    }
69}