recover_flatten.rs

use super::Error;
use async_stream::try_stream;
use futures::{prelude::*, stream};

// Flattens a stream-of-streams until an unrecoverable error is encountered.
pub fn flatten<S, T, R>(
    mut stream_of_streams: S,
    recover: R,
) -> impl TryStream<Ok = T::Ok, Error = Error> + 'static
where
    S: TryStream<Ok = T> + Unpin + 'static,
    S::Error: Into<Error>,
    T: TryStream + Unpin + 'static,
    T::Ok: Unpin,
    T::Error: Into<Error>,
    R: Recover + 'static,
    R::Backoff: Unpin,
{
    try_stream! {
        let mut backoff: Option<R::Backoff> = None;

        loop {
            if let Some(mut b) = backoff.take() {
                if b.next().await.is_some() {
                    backoff = Some(b);
                }
            }

            match stream_of_streams.try_next().await {
                Ok(Some(mut stream)) => {
                    backoff = None;
                    loop {
                        match stream.try_next().await {
                            Ok(Some(item)) => yield item,
                            Ok(None) => break,
                            Err(e) => {
                                let new_backoff = recover.recover(e.into())?;
                                backoff = Some(new_backoff);
                            }
                        }
                    }
                }
                Ok(None) =>  return,
                Err(e) => {
                    let new_backoff = recover.recover(e.into())?;
                    backoff = Some(backoff.unwrap_or(new_backoff));
                }
            }
        }
    }
}