tuwunel_core/utils/stream/
try_parallel.rs1use futures::{TryFutureExt, stream::TryStream};
4use tokio::{runtime, task::JoinError};
5
6use super::TryBroadbandExt;
7use crate::{Error, Result, utils::sys::available_parallelism};
8
9pub 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 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 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}