Skip to main content

tuwunel_core/utils/stream/
try_parallel.rs

1//! Parallelism stream combinator extensions to futures::Stream
2
3use futures::{TryFutureExt, stream::TryStream};
4use tokio::{runtime, task::JoinError};
5
6use super::TryBroadbandExt;
7use crate::{Error, Result, utils::sys::available_parallelism};
8
9/// Adds unordered parallel transformations for fallible streams.
10///
11/// Each closure runs on Tokio's blocking pool, making these combinators
12/// suitable for CPU-bound work rather than asynchronous I/O. Concurrency
13/// defaults to the host's available parallelism.
14pub trait TryParallelExt<T, E>
15where
16	Self: TryStream<Ok = T, Error = E, Item = Result<T, E>> + Send + Sized,
17	E: From<JoinError> + From<Error> + Send + 'static,
18	T: Send + 'static,
19{
20	/// Runs a synchronous fallible transform on the blocking pool.
21	///
22	/// `n` controls concurrent jobs and defaults to available parallelism; an
23	/// explicit zero cannot make progress. Job results are completion ordered,
24	/// while source errors bypass `f` and may overtake queued jobs. Spawned
25	/// jobs can continue after the adapter is dropped.
26	///
27	/// # Panics
28	///
29	/// Panics when no handle is supplied outside a Tokio runtime context.
30	fn paralleln_and_then<U, F, N, H>(
31		self,
32		h: H,
33		n: N,
34		f: F,
35	) -> impl TryStream<Ok = U, Error = E, Item = Result<U, E>> + Send
36	where
37		N: Into<Option<usize>>,
38		H: Into<Option<runtime::Handle>>,
39		F: Fn(Self::Ok) -> Result<U, E> + Clone + Send + 'static,
40		U: Send + 'static;
41
42	/// Runs a synchronous fallible transform at the default parallelism.
43	///
44	/// The optional handle defaults to the current Tokio runtime. Job results
45	/// are completion ordered, while source errors may overtake them; join
46	/// failures convert into `E`. Spawned jobs can continue after the adapter
47	/// is dropped.
48	///
49	/// # Panics
50	///
51	/// Panics when no handle is supplied outside a Tokio runtime context.
52	fn parallel_and_then<U, F, H>(
53		self,
54		h: H,
55		f: F,
56	) -> impl TryStream<Ok = U, Error = E, Item = Result<U, E>> + Send
57	where
58		H: Into<Option<runtime::Handle>>,
59		F: Fn(Self::Ok) -> Result<U, E> + Clone + Send + 'static,
60		U: Send + 'static,
61	{
62		self.paralleln_and_then(h, None, f)
63	}
64}
65
66impl<T, E, S> TryParallelExt<T, E> for S
67where
68	S: TryStream<Ok = T, Error = E, Item = Result<T, E>> + Send + Sized,
69	E: From<JoinError> + From<Error> + Send + 'static,
70	T: Send + 'static,
71{
72	fn paralleln_and_then<U, F, N, H>(
73		self,
74		h: H,
75		n: N,
76		f: F,
77	) -> impl TryStream<Ok = U, Error = E, Item = Result<U, E>> + Send
78	where
79		N: Into<Option<usize>>,
80		H: Into<Option<runtime::Handle>>,
81		F: Fn(Self::Ok) -> Result<U, E> + Clone + Send + 'static,
82		U: Send + 'static,
83	{
84		let n = n.into().unwrap_or_else(available_parallelism);
85		let h = h.into().unwrap_or_else(runtime::Handle::current);
86		self.broadn_and_then(n, move |val| {
87			let (h, f) = (h.clone(), f.clone());
88			async move {
89				h.spawn_blocking(move || f(val))
90					.map_err(E::from)
91					.await?
92			}
93		})
94	}
95}