Add retry utility in prep for s3_store

config/backends.rs

@@ -72,3 +72,47 @@ pub struct EvictionPolicy {
     #[serde(default)]
     pub max_count: u64,
 }
+
+/// Retry configuration. This configuration is exponential and each iteration
+/// a jitter as a percentage is applied of the calculated delay. For example:
+/// ```
+/// Retry{
+///   max_retries: 7
+///   delay: .1,
+///   jitter: .5
+/// }
+/// ```
+/// will result in:
+/// Attempt - Delay
+/// 1         0ms
+/// 2         75ms - 125ms
+/// 3         150ms - 250ms
+/// 4         300ms - 500ms
+/// 5         600ms - 1s
+/// 6         1.2s - 2s
+/// 7         2.4s - 4s
+/// 8         4.8s - 8s
+/// Remember that to get total results is additive, meaning the above results
+/// would mean a single request would have a total delay of 9.525s - 15.875s.
+#[derive(Deserialize, Clone, Debug, Default)]
+pub struct Retry {
+    /// Maximum number of retries until retrying stops.
+    /// Setting this to zero will always attempt 1 time, but not retry.
+    #[serde(default)]
+    pub max_retries: usize,
+
+    /// Delay in seconds for exponential back off.
+    #[serde(default)]
+    pub delay: f32,
+
+    /// Amount of jitter to add as a percentage in decimal form. This will
+    /// change the formula like:
+    /// ```
+    /// random(
+    ///    2 ^ {attempt_number} * {delay}) * (1 - (jitter / 2)),
+    ///    2 ^ {attempt_number} * {delay}) * (1 + (jitter / 2)),
+    /// )
+    /// ```
+    #[serde(default)]
+    pub jitter: f32,
+}

util/BUILD

@@ -28,6 +28,17 @@ rust_library(
     visibility = ["//visibility:public"],
 )
 
+rust_library(
+    name = "retry",
+    srcs = ["retry.rs"],
+    deps = [
+        "//third_party:futures",
+        "//third_party:tokio",
+        ":error",
+    ],
+    visibility = ["//visibility:public"],
+)
+
 rust_library(
     name = "evicting_map",
     srcs = ["evicting_map.rs"],
@@ -86,6 +97,18 @@ rust_test(
     ],
 )
 
+rust_test(
+    name = "retry_test",
+    srcs = ["tests/retry_test.rs"],
+    deps = [
+        "//third_party:futures",
+        "//third_party:pretty_assertions",
+        "//third_party:tokio",
+        ":error",
+        ":retry",
+    ],
+)
+
 rust_test(
     name = "async_read_taker_test",
     srcs = ["tests/async_read_taker_test.rs"],

util/retry.rs

@@ -0,0 +1,72 @@
+// Copyright 2021 Nathan (Blaise) Bruer.  All rights reserved.
+
+use futures::future::Future;
+use futures::stream::StreamExt;
+use std::pin::Pin;
+use std::time::Duration;
+
+use error::{make_err, Code, Error};
+
+pub struct ExponentialBackoff {
+    current: Duration,
+}
+
+impl ExponentialBackoff {
+    pub fn new(base: Duration) -> Self {
+        ExponentialBackoff { current: base }
+    }
+}
+
+impl Iterator for ExponentialBackoff {
+    type Item = Duration;
+
+    fn next(&mut self) -> Option<Duration> {
+        self.current = self.current * 2;
+        Some(self.current)
+    }
+}
+
+type SleepFn = Box<dyn Fn(Duration) -> Pin<Box<dyn Future<Output = ()> + Send>> + Sync + Send>;
+
+#[derive(PartialEq, Eq, Debug)]
+pub enum RetryResult<T> {
+    Ok(T),
+    Retry(Error),
+    Err(Error),
+}
+
+/// Class used to retry a job with a sleep function in between each retry.
+pub struct Retrier {
+    sleep_fn: SleepFn,
+}
+
+impl Retrier {
+    pub fn new(sleep_fn: SleepFn) -> Self {
+        Retrier { sleep_fn }
+    }
+
+    pub fn retry<'a, T, Fut, I>(
+        &'a self,
+        duration_iter: I,
+        operation: Fut,
+    ) -> Pin<Box<dyn Future<Output = Result<T, Error>> + 'a + Send>>
+    where
+        Fut: futures::stream::Stream<Item = RetryResult<T>> + Send + 'a,
+        I: IntoIterator<Item = Duration> + Send + 'a,
+        <I as IntoIterator>::IntoIter: Send,
+        T: Send,
+    {
+        Box::pin(async move {
+            let mut iter = duration_iter.into_iter();
+            let mut operation = Box::pin(operation);
+            loop {
+                match operation.next().await {
+                    None => return Err(make_err!(Code::Internal, "Retry stream ended abruptly",)),
+                    Some(RetryResult::Ok(value)) => return Ok(value),
+                    Some(RetryResult::Err(e)) => return Err(e),
+                    Some(RetryResult::Retry(e)) => (self.sleep_fn)(iter.next().ok_or_else(|| e)?).await,
+                }
+            }
+        })
+    }
+}

util/tests/retry_test.rs

@@ -0,0 +1,177 @@
+// Copyright 2021 Nathan (Blaise) Bruer.  All rights reserved.
+use std::pin::Pin;
+use std::sync::atomic::{AtomicI32, Ordering};
+use std::sync::Arc;
+
+use error::{make_err, Code, Error};
+
+use tokio::time::Duration;
+
+use futures::future::ready;
+use futures::stream::repeat_with;
+
+use retry::{Retrier, RetryResult};
+
+struct MockDurationIterator {
+    duration: Duration,
+}
+
+impl MockDurationIterator {
+    pub fn new(duration: Duration) -> Self {
+        MockDurationIterator { duration: duration }
+    }
+}
+
+impl Iterator for MockDurationIterator {
+    type Item = Duration;
+
+    fn next(&mut self) -> Option<Duration> {
+        Some(self.duration)
+    }
+}
+
+#[cfg(test)]
+mod retry_tests {
+    use super::*;
+    use pretty_assertions::assert_eq; // Must be declared in every module.
+
+    #[tokio::test]
+    async fn retry_simple_success() -> Result<(), Error> {
+        let retrier = Retrier::new(Box::new(|_duration| Box::pin(ready(()))));
+        let retry_config = MockDurationIterator::new(Duration::from_millis(1));
+        let run_count = Arc::new(AtomicI32::new(0));
+
+        let result = Pin::new(&retrier)
+            .retry(
+                retry_config,
+                repeat_with(|| {
+                    run_count.fetch_add(1, Ordering::Relaxed);
+                    RetryResult::Ok(true)
+                }),
+            )
+            .await?;
+        assert_eq!(
+            run_count.load(Ordering::Relaxed),
+            1,
+            "Expected function to be called once"
+        );
+        assert_eq!(result, true, "Expected result to succeed");
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn retry_fails_after_3_runs() -> Result<(), Error> {
+        let retrier = Retrier::new(Box::new(|_duration| Box::pin(ready(()))));
+        let retry_config = MockDurationIterator::new(Duration::from_millis(1)).take(2); // .take() will run X times + 1.
+        let run_count = Arc::new(AtomicI32::new(0));
+
+        let result = Pin::new(&retrier)
+            .retry(
+                retry_config,
+                repeat_with(|| {
+                    run_count.fetch_add(1, Ordering::Relaxed);
+                    RetryResult::<bool>::Retry(make_err!(Code::Unavailable, "Dummy failure",))
+                }),
+            )
+            .await;
+        assert_eq!(run_count.load(Ordering::Relaxed), 3, "Expected function to be called");
+        assert_eq!(result.is_err(), true, "Expected result to error");
+        assert_eq!(
+            result.unwrap_err().to_string(),
+            "Error { code: Unavailable, messages: [\"Dummy failure\"] }"
+        );
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn retry_success_after_2_runs() -> Result<(), Error> {
+        let retrier = Retrier::new(Box::new(|_duration| Box::pin(ready(()))));
+        let retry_config = MockDurationIterator::new(Duration::from_millis(1)).take(5); // .take() will run X times + 1.
+        let run_count = Arc::new(AtomicI32::new(0));
+
+        let result = Pin::new(&retrier)
+            .retry(
+                retry_config,
+                repeat_with(|| {
+                    run_count.fetch_add(1, Ordering::Relaxed);
+                    if run_count.load(Ordering::Relaxed) == 2 {
+                        return RetryResult::Ok(true);
+                    }
+                    RetryResult::<bool>::Retry(make_err!(Code::Unavailable, "Dummy failure",))
+                }),
+            )
+            .await?;
+        assert_eq!(run_count.load(Ordering::Relaxed), 2, "Expected function to be called");
+        assert_eq!(result, true, "Expected result to succeed");
+
+        Ok(())
+    }
+
+    // test-prefix/987cc59b2d364596f94bd82f250d02aebb716c3a100163ca784b54a50e0dfde2-142
+    #[tokio::test]
+    async fn retry_calls_sleep_fn() -> Result<(), Error> {
+        const EXPECTED_MS: u64 = 71;
+        let sleep_fn_run_count = Arc::new(AtomicI32::new(0));
+        let sleep_fn_run_count_copy = sleep_fn_run_count.clone();
+        let retrier = Retrier::new(Box::new(move |duration| {
+            // Note: Need to make another copy to make the compiler happy.
+            let sleep_fn_run_count_copy = sleep_fn_run_count_copy.clone();
+            Box::pin(async move {
+                // Remember: This function is called only on retries, not the first run.
+                sleep_fn_run_count_copy.fetch_add(1, Ordering::Relaxed);
+                assert_eq!(duration, Duration::from_millis(EXPECTED_MS));
+                ()
+            })
+        }));
+
+        // s3://blaisebruer-cas-store/test-prefix/987cc59b2d364596f94bd82f250d02aebb716c3a100163ca784b54a50e0dfde2-142
+        {
+            // Try with retry limit hit.
+            let retry_config = MockDurationIterator::new(Duration::from_millis(EXPECTED_MS)).take(5);
+            let result = Pin::new(&retrier)
+                .retry(
+                    retry_config,
+                    repeat_with(|| RetryResult::<bool>::Retry(make_err!(Code::Unavailable, "Dummy failure",))),
+                )
+                .await;
+
+            assert_eq!(result.is_err(), true, "Expected the retry to fail");
+            assert_eq!(
+                sleep_fn_run_count.load(Ordering::Relaxed),
+                5,
+                "Expected the sleep_fn to be called twice"
+            );
+        }
+        sleep_fn_run_count.store(0, Ordering::Relaxed); // Reset our counter.
+        {
+            // Try with 3 retries.
+            let retry_config = MockDurationIterator::new(Duration::from_millis(EXPECTED_MS)).take(5);
+            let run_count = Arc::new(AtomicI32::new(0));
+            let result = Pin::new(&retrier)
+                .retry(
+                    retry_config,
+                    repeat_with(|| {
+                        run_count.fetch_add(1, Ordering::Relaxed);
+                        // Remember: This function is only called every time, not just retries.
+                        // We run the first time, then retry 2 additional times meaning 3 runs.
+                        if run_count.load(Ordering::Relaxed) == 3 {
+                            return RetryResult::Ok(true);
+                        }
+                        RetryResult::<bool>::Retry(make_err!(Code::Unavailable, "Dummy failure",))
+                    }),
+                )
+                .await?;
+
+            assert_eq!(result, true, "Expected results to pass");
+            assert_eq!(
+                sleep_fn_run_count.load(Ordering::Relaxed),
+                2,
+                "Expected the sleep_fn to be called twice"
+            );
+        }
+
+        Ok(())
+    }
+}