Skip to main content

snix_castore/combinators/
race.rs

1//! Contains helper function to run operations across multiple services as the same time.
2
3use futures::{Stream, StreamExt, TryStreamExt, stream::FuturesUnordered};
4
5/// Runs a unary operation on all services.
6///
7/// The passed async function describes the operation to run on each service.
8/// The control struct can be used to decide whether to ignore a response
9/// (continuing with other backends) or return it.
10///
11/// Errors short-circuit.
12///
13/// FUTUREWORK: allow configuring errors to log only (and skip)
14pub async fn race_unary<'a: 'f, 'f, SVC: 'a, SVCS, T, E, F, Fut>(
15    services: SVCS,
16    mut f: F,
17) -> Result<T, Error<E>>
18where
19    F: FnMut(SVC) -> Fut + Copy + 'a,
20    Fut: Future<Output = Option<Result<T, E>>> + 'a,
21    SVCS: IntoIterator<Item = SVC> + 'a,
22    T: Default,
23{
24    let mut requests = services
25        .into_iter()
26        .enumerate()
27        .map(|(backend_idx, svc)| async move {
28            f(svc)
29                .await
30                .map(|resp| resp.map_err(|err| Error::Backend(backend_idx, err)))
31        })
32        .collect::<FuturesUnordered<_>>();
33
34    while let Some(resp) = requests.next().await {
35        if let Some(resp) = resp {
36            return resp;
37        }
38    }
39
40    Ok(T::default())
41}
42
43/// Runs an operation returning a stream on all services.
44///
45/// The passed async function describes the operation to run on each service.
46///
47/// The returned `Option<_>` controls whether to ignore a response, such as an error,
48/// or an empty stream (and continue with other backends) or return it.
49pub fn race_stream<'a, SVC, SVCS, F, Fut, S, T, E>(
50    services: SVCS,
51    mut f: F,
52) -> impl Stream<Item = Result<T, Error<E>>> + 'a
53where
54    F: FnMut(SVC) -> Fut + Copy + 'a,
55    Fut: Future<Output = Option<S>> + 'a,
56    SVCS: IntoIterator<Item = SVC> + 'a,
57    SVC: 'a,
58    S: Stream<Item = Result<T, E>>,
59    E: Send + 'a,
60    T: Send + 'a,
61{
62    let mut requests = services
63        .into_iter()
64        .enumerate()
65        .map(|(backend_idx, svc)| async move {
66            f(svc)
67                .await
68                .map(|stream| stream.map_err(move |err| Error::Backend(backend_idx, err)))
69        })
70        .collect::<FuturesUnordered<_>>();
71
72    async_stream::stream! {
73        while let Some(maybe_stream) = requests.next().await {
74            if let Some(stream) = maybe_stream {
75                // yield from the stream
76                let mut stream = std::pin::pin!(stream);
77                while let Some(elem) = stream.next().await {
78                    yield elem
79                }
80            }
81        }
82    }
83}
84
85#[derive(thiserror::Error, Debug)]
86pub enum Error<E> {
87    #[error("error from service at index {0}")]
88    Backend(usize, #[source] E),
89}