cs_utils/utils/futures/

with_retries.rs

use std::{
    error::Error as StdError,
    fmt::Debug,
};

use futures::Future;
use thiserror::Error as ThisError;

use crate::futures::wait;

#[derive(ThisError, Debug)]
pub enum BaseError<S: AsRef<str> + ToString + Send = String> {
    #[error("{0}")]
    Generic(#[from] Box<dyn StdError + Send + Sync>),

    #[error("{0}")]
    Msg(S),
}

impl<T, S: AsRef<str> + ToString + Send> From<BaseError<S>> for Result<T, BaseError<S>> {
    fn from(error: BaseError<S>) -> Self {
        return Err(error);
    }
}

pub struct WithRetriesOptions<
    TError: Into<Box<dyn StdError + Send + Sync>> = BaseError,
> {
    pub(crate) should_stop: Box<dyn Fn(&TError, usize) -> bool + Send>,
    pub(crate) retries: usize,
    pub(crate) retry_delay: u64,
}

impl<TError: Into<Box<dyn StdError + Send + Sync>>> WithRetriesOptions<TError> {
    pub fn retries(
        self,
        retries: usize,
    ) -> WithRetriesOptions<TError> {
        return WithRetriesOptions {
            retries,
            ..self
        };
    }

    pub fn decrement_retries(
        &mut self,
    ) {
        // make sure we don't violate `usize` type by substracting from zero
        assert!(
            self.retries > 0,
            "Retries must be greater than 0.",
        );

        self.retries -= 1;
    }

    pub fn retry_delay(
        self,
        retry_delay_ms: u64,
    ) -> WithRetriesOptions<TError> {
        return WithRetriesOptions {
            retry_delay: retry_delay_ms,
            ..self
        };
    }

    pub fn should_stop_retries(
        self,
        should_stop_retries: impl Fn(&TError, usize) -> bool + Send + 'static,
    ) -> WithRetriesOptions<TError> {
        return WithRetriesOptions {
            should_stop: Box::new(should_stop_retries),
            ..self
        };
    }
}

impl<TError: Into<Box<dyn StdError + Send + Sync>>> Default for WithRetriesOptions<TError> {
    fn default() -> WithRetriesOptions<TError> {
        return WithRetriesOptions {
            retries: 2,
            retry_delay: 0,
            should_stop: Box::new(|_, _| {
                return false;
            }),
        }
    }
}

pub async fn with_retries_inner<
    T: Send,
    TError: Into<Box<dyn StdError + Send + Sync>>,
    TFuture: Future<Output = Result<T, TError>> + Send,
>(
    job: TFuture,
    retries: usize,
    // options: &WithRetriesOptions,
) -> (Result<T, TError>, bool) {
    // let retries = options.retries;
    // let should_stop = &options.should_stop;

    // if success, return the future result
    let error = match job.await {
        Ok(result) => return (Ok(result), true),
        Err(err) => err,
    };

    // if no retries left, return the error result
    if retries == 0 {
        return (Err(error), true);
    }

    // // pass `error` reference and `retry` number to the closure,
    // // if `true` is returned, - return the error result and stop
    // if should_stop(&error, retries) {
    //     return (Err(error), true);
    // }

    // return error result but not stop
    return (Err(error), false);
}

/// !Experimental! Might be unstable in regard of runtime stability and API surface.
/// 
/// Executes a future with retries if it fails.
pub async fn with_retries<
    T: Send,
    TError: Into<Box<dyn StdError + Send + Sync>>,
    TFuture: Future<Output = Result<T, TError>> + Send,
>(
    mut future_factory: impl FnMut(usize) -> TFuture,
    mut options: WithRetriesOptions<TError>,
) -> Result<T, TError> {
    loop {
        let (result, is_done) = with_retries_inner(
            future_factory(options.retries),
            options.retries,
        ).await;

        // if done, return the result
        if is_done {
            return result;
        }
        
        // get error value and assert that the result is not `Ok()`
        let error = match result {
            Err(err) => err,
            Ok(_) => {
                panic!("Result cannot be `Ok` here.");
            },
        };

        // pass `error` reference and `retry` number to the closure,
        // if `true` is returned, - return the error result and stop
        if (&options.should_stop)(&error, options.retries) {
            return Err(error);
        }

        // wait for `retry_delay` milliseconds before retrying
        wait(options.retry_delay).await;

        // remove one retry
        options.decrement_retries();
    }
}

#[cfg(test)]
mod tests {
    use std::error::Error as StdError;

    use cs_utils::{
        random_number,
        futures::{wait_random, with_retries, WithRetriesOptions},
    };
    use anyhow::anyhow;
    use thiserror::Error as ThisError;

    use crate::futures::with_retries::BaseError;

    mod future_factory {
        use super::*;

        use tokio::sync::mpsc;

        /// The `expected_result` variable is external for the future factory here.
        #[tokio::test]
        async fn factory_depends_on_external_variable() {
            let expected_result = random_number(0..=u64::MAX);

            let result = with_retries(|_| {
                    return async move {
                        wait_random(1..=5).await;

                        // help rust to infer the `Result` type
                        if false {
                            return BaseError::Msg("Something bad happened.").into();
                        }

                        return Ok(expected_result);
                    };
                },
                Default::default(),
            ).await.expect("Must succeed.");

            assert_eq!(
                result,
                expected_result,
                "Must return the expected result.",
            );
        }

        #[tokio::test]
        async fn factory_depends_on_retries_left() {
            let result = with_retries(|retries_left| async move {
                    wait_random(1..=5).await;
                    
                    // help rust to infer the `Result` type
                    if false {
                        return Err(BaseError::Msg("Something bad happened."));
                    }

                    return Ok(retries_left);
                },
                Default::default(),
            ).await.expect("Must succeed.");

            assert_eq!(
                result,
                2,
                "Must return the default `retries` value result.",
            );
        }

        #[tokio::test]
        async fn factory_without_move() {
            let result = with_retries(|_| async {
                    wait_random(1..=5).await;

                    // help rust to infer the `Result` type
                    if false {
                        return Err(BaseError::Msg("Something bad happened."));
                    }

                    return Ok(10);
                },
                Default::default(),
            ).await.expect("Must succeed.");

            assert_eq!(
                result,
                10,
                "Must return the expected result.",
            );
        }

        #[tokio::test]
        async fn factory_takes_number_of_remaining_retries() {
            let (sender, mut receiver) = mpsc::channel(1024);

            tokio::try_join!(
                tokio::spawn(async move {
                    return with_retries(move |retries_left| {
                            let tx = sender.clone();
            
                            return async move {
                                wait_random(1..=5).await;
            
                                tx
                                    .send(retries_left).await
                                    .expect("Cannot send retries left number.");
            
                                if retries_left > 0 {
                                    let err: Box<dyn StdError + Send + Sync> = anyhow!("Some error.").into();
                                    return Err(err);
                                }
            
                                return Ok(());
                            };
                        },
                        Default::default(),
                    ).await.expect("Must succeed.");
                }),
                tokio::spawn(async move {
                    let mut expected_retries_left = 2;

                    while let Some(retries_left) = receiver.recv().await {
                        assert_eq!(
                            retries_left,
                            expected_retries_left,
                            "Must return the expected retries left.",
                        );

                        if expected_retries_left == 0 {
                            break;
                        }

                        expected_retries_left -= 1;
                    }
                }),
            ).unwrap();
        }
    }
    
    mod should_stop_retries {
        use crate::{traits::Random, test::random_vec, random_str};

        use super::*;

        #[rstest::rstest]
        #[case(2, vec![2], true)]
        #[case(1, vec![2, 1], true)]
        #[case(0, vec![2, 1, 0], false)]
        #[tokio::test]
        async fn defines_if_should_stop_retries(
            #[case] stop_at: usize,
            #[case] expected_retries_left: Vec<usize>,
            #[case] is_error: bool,
        ) {
            let options = WithRetriesOptions::default()
                .should_stop_retries(move |_error, retries_left| {
                    return retries_left == stop_at;
                });

            let mut items = vec![];
            let result = with_retries(|retries_left| {
                    items.push(retries_left);
                    return async move {
                        wait_random(1..=5).await;

                        if retries_left > 0 {
                            return Err(anyhow!("Some error."));
                        }
    
                        return Ok(());
                    };
                },
                options,
            ).await;

            assert_eq!(
                items,
                expected_retries_left,
            );

            assert_eq!(
                result.is_err(),
                is_error,
                "Must return correct result.",
            );
        }

        #[tokio::test]
        async fn receives_original_error() {
            let options = WithRetriesOptions::default()
                .should_stop_retries(move |error, retry_number| {
                    assert_eq!(
                        format!("{}", error),
                        format!("Oh no! #{}", retry_number),
                        "Must accept original error.",
                    );

                    return false;
                });

            with_retries(|retries_left| {
                    return async move {
                        wait_random(1..=5).await;

                        if retries_left > 0 {
                            return Err(BaseError::Msg(format!("Oh no! #{}", retries_left)));
                        }
    
                        return Ok(());
                    };
                },
                options,
            ).await.expect("Must succeed.");
        }

        #[derive(ThisError, Debug)]
        pub enum TestError<S: AsRef<str> + ToString = String> {
            // #[error("data store disconnected")]
            // Disconnect(#[from] io::Error),

            // #[error("the data for key `{0}` is not available")]
            // Redaction(String),

            // #[error("invalid header (expected {expected:?}, found {found:?})")]
            // InvalidHeader {
            //     expected: String,
            //     found: String,
            // },

            #[error("{0}")]
            Msg(S),

            #[error("{0}")]
            Anyhow(anyhow::Error),

            #[error("Unknown error. [code: {0}]")]
            Unknown(u128),
        }

        impl Clone for TestError {
            fn clone(&self) -> TestError {
                match self {
                    TestError::Msg(msg) => {
                        return TestError::Msg(msg.clone());
                    },
                    TestError::Anyhow(msg) => {
                        return TestError::Anyhow(anyhow!(msg.to_string()));
                    },
                    TestError::Unknown(id) => return TestError::Unknown(id.clone()),
                };
            }
        }

        impl PartialEq for TestError {
            fn eq(&self, other: &TestError) -> bool {
                match (self, other) {
                    (TestError::Msg(left_msg), TestError::Msg(right_msg)) => {
                        return left_msg == right_msg;
                    },
                    (TestError::Anyhow(left_err), TestError::Anyhow(right_err)) => {
                        return format!("{}", left_err) == format!("{}", right_err);
                    },
                    (TestError::Unknown(left_id), TestError::Unknown(right_id)) => {
                        return left_id == right_id;
                    },
                    _ => return false,
                };
            }
        }

        impl Random for TestError {
            fn random() -> TestError {
                match random_number(0..=2) {
                    0 => {
                        return TestError::Msg(
                            format!("Something bad happened. [code: {}]", random_str(16)),
                        );
                    },
                    1 => {
                        return TestError::Anyhow(anyhow!(
                            format!("Oh no! [code: {}]", random_str(16)),
                        ));
                    },
                    2 => {
                        return TestError::Unknown(random_number(0..=u128::MAX))
                    },
                    _ => unreachable!(),
                };
            }
        }

        #[rstest::rstest]
        #[case(random_number(0..=8))]
        #[case(random_number(0..=8))]
        #[case(random_number(0..=8))]
        #[case(random_number(0..=8))]
        #[case(random_number(0..=8))]
        #[case(random_number(0..=8))]
        #[case(random_number(0..=8))]
        #[case(random_number(0..=8))]
        #[tokio::test]
        async fn receives_original_error_of_type(
            #[case] retry_count: u32,
        ) {
            let errors: Vec<TestError> = random_vec(retry_count + 1);

            let expected_errors = errors.clone();
            let options = WithRetriesOptions::default()
                .should_stop_retries(move |error, retry_number| {
                    assert_eq!(
                        error,
                        &expected_errors[retry_number],
                        "Must receive original errors.",
                    );

                    return false;
                })
                .retries(retry_count as usize);

            let result = with_retries(|retries_left|
                {
                    let sent_errors = errors.clone();
                    return async move {
                        wait_random(1..=5).await;

                        if retries_left > 0 {
                            return Err(sent_errors[retries_left].clone());
                        }
    
                        return Ok(retries_left);
                    };
                },
                options,
            ).await.expect("Must succeed.");

            assert_eq!(
                result,
                0,
                "Must exhaust all retry attempts.",
            );
        }

        #[rstest::rstest]
        #[case(random_number(0..=2), TestError::random())]
        #[case(random_number(0..=2), TestError::random())]
        #[case(random_number(0..=2), TestError::random())]
        #[case(random_number(0..=2), TestError::random())]
        #[case(random_number(0..=2), TestError::random())]
        #[case(random_number(0..=2), TestError::random())]
        #[case(random_number(0..=2), TestError::random())]
        #[case(random_number(0..=2), TestError::random())]
        #[tokio::test]
        async fn can_stop_by_error(
            #[case] stop_at: usize,
            #[case] test_error: TestError,
        ) {
            let expected_error = test_error.clone();

            let options = WithRetriesOptions::default()
                .should_stop_retries(move |error, retry_number| {
                    assert!(
                        retry_number >= stop_at,
                        "Must not run after `true` is returned",
                    );

                    if error == &expected_error {
                        assert_eq!(
                            retry_number,
                            stop_at,
                            "The expected error must be thrown at the expected iteration.",
                        );

                        return true;
                    }
                    
                    return false;
                });

            let result = with_retries(|retries_left|
                {
                    let sent_error = test_error.clone();
                    return async move {
                        wait_random(1..=5).await;

                        if retries_left == stop_at {
                            return Err(sent_error);
                        }

                        if retries_left > 0 {
                            return Err(TestError::random());
                        }
    
                        return Ok(());
                    };
                },
                options,
            ).await.expect_err("Must fail.");

            assert_eq!(
                result,
                test_error,
                "Must complete with the expected error.",
            );
        }
    }

    mod retry_delay {
        #[tokio::test]
        async fn waits_before_retry() {
            todo!();
        }
    }
}