Refactor retry (#458)

* Refactor retry

* Forgot to run cargo fmt

Author: rylev

sdk/core/src/policies/retry_policies/exponential_retry.rs

@@ -1,8 +1,4 @@
-use crate::policies::{Policy, PolicyResult, Request, Response};
-use crate::sleep::sleep;
-use crate::PipelineContext;
 use chrono::{DateTime, Local};
-use std::sync::Arc;
 use std::time::Duration;
 
 /// Retry policy with exponential back-off.
@@ -27,56 +23,24 @@ impl ExponentialRetryPolicy {
             max_delay,
         }
     }
+}
 
-    fn is_expired(
-        &self,
-        first_retry_time: &mut Option<DateTime<Local>>,
-        current_retries: &u32,
-    ) -> bool {
-        if *current_retries > self.max_retries {
+impl super::RetryPolicy for ExponentialRetryPolicy {
+    fn is_expired(&self, first_retry_time: &mut Option<DateTime<Local>>, retry_count: u32) -> bool {
+        if retry_count > self.max_retries {
             return true;
         }
 
-        if first_retry_time.is_none() {
-            *first_retry_time = Some(Local::now());
-        }
+        let first_retry_time = first_retry_time.get_or_insert_with(|| Local::now());
+        let max_delay = chrono::Duration::from_std(self.max_delay)
+            .unwrap_or_else(|_| chrono::Duration::max_value());
 
-        Local::now()
-            > first_retry_time.unwrap() + chrono::Duration::from_std(self.max_delay).unwrap()
+        Local::now() > *first_retry_time + max_delay
     }
-}
 
-#[async_trait::async_trait]
-impl<C> Policy<C> for ExponentialRetryPolicy
-where
-    C: Send + Sync,
-{
-    async fn send(
-        &self,
-        ctx: &mut PipelineContext<C>,
-        request: &mut Request,
-        next: &[Arc<dyn Policy<C>>],
-    ) -> PolicyResult<Response> {
-        let mut first_retry_time = None;
-        let mut current_retries = 0;
-
-        loop {
-            match next[0].send(ctx, request, &next[1..]).await {
-                Ok(response) => return Ok(response),
-                Err(error) => {
-                    log::error!("Error occurred when making request: {}", error);
-                    if self.is_expired(&mut first_retry_time, &current_retries) {
-                        return Err(error);
-                    } else {
-                        current_retries += 1;
-
-                        let sleep_ms = self.delay.as_millis() as u64
-                            * u64::pow(2u64, current_retries - 1)
-                            + rand::random::<u8>() as u64;
-                        sleep(Duration::from_millis(sleep_ms)).await;
-                    }
-                }
-            }
-        }
+    fn sleep_duration(&self, retry_count: u32) -> Duration {
+        let sleep_ms = self.delay.as_millis() as u64 * u64::pow(2u64, retry_count - 1)
+            + rand::random::<u8>() as u64;
+        Duration::from_millis(sleep_ms)
     }
 }

sdk/core/src/policies/retry_policies/fixed_retry.rs

@@ -1,8 +1,4 @@
-use crate::policies::{Policy, PolicyResult, Request, Response};
-use crate::sleep::sleep;
-use crate::PipelineContext;
 use chrono::{DateTime, Local};
-use std::sync::Arc;
 use std::time::Duration;
 
 /// Retry policy with fixed back-off.
@@ -26,54 +22,23 @@ impl FixedRetryPolicy {
             max_delay,
         }
     }
+}
 
-    fn is_expired(
-        &self,
-        first_retry_time: &mut Option<DateTime<Local>>,
-        current_retries: &u32,
-    ) -> bool {
-        if *current_retries > self.max_retries {
+impl super::RetryPolicy for FixedRetryPolicy {
+    fn is_expired(&self, first_retry_time: &mut Option<DateTime<Local>>, retry_count: u32) -> bool {
+        if retry_count > self.max_retries {
             return true;
         }
 
-        if first_retry_time.is_none() {
-            *first_retry_time = Some(Local::now());
-        }
+        let first_retry_time = first_retry_time.get_or_insert_with(Local::now);
+        let max_delay = chrono::Duration::from_std(self.max_delay)
+            .unwrap_or_else(|_| chrono::Duration::max_value());
 
-        Local::now()
-            > first_retry_time.unwrap() + chrono::Duration::from_std(self.max_delay).unwrap()
+        Local::now() > *first_retry_time + max_delay
     }
-}
 
-#[async_trait::async_trait]
-impl<C> Policy<C> for FixedRetryPolicy
-where
-    C: Send + Sync,
-{
-    async fn send(
-        &self,
-        ctx: &mut PipelineContext<C>,
-        request: &mut Request,
-        next: &[Arc<dyn Policy<C>>],
-    ) -> PolicyResult<Response> {
-        let mut first_retry_time = None;
-        let mut current_retries = 0;
-
-        loop {
-            match next[0].send(ctx, request, &next[1..]).await {
-                Ok(response) => return Ok(response),
-                Err(error) => {
-                    log::error!("Error occurred when making request: {}", error);
-                    if self.is_expired(&mut first_retry_time, &current_retries) {
-                        return Err(error);
-                    } else {
-                        current_retries += 1;
-
-                        let sleep_ms = self.delay.as_millis() as u64 + rand::random::<u8>() as u64;
-                        sleep(Duration::from_millis(sleep_ms)).await;
-                    }
-                }
-            }
-        }
+    fn sleep_duration(&self, _retry_count: u32) -> Duration {
+        let sleep_ms = self.delay.as_millis() as u64 + rand::random::<u8>() as u64;
+        Duration::from_millis(sleep_ms)
     }
 }

sdk/core/src/policies/retry_policies/mod.rs

@@ -1,7 +1,9 @@
 mod exponential_retry;
 mod fixed_retry;
 mod no_retry;
+mod retry_policy;
 
 pub use exponential_retry::*;
 pub use fixed_retry::*;
 pub use no_retry::*;
+use retry_policy::RetryPolicy;

sdk/core/src/policies/retry_policies/retry_policy.rs

@@ -0,0 +1,44 @@
+use crate::policies::{Policy, PolicyResult, Request, Response};
+use crate::sleep::sleep;
+use crate::PipelineContext;
+use chrono::{DateTime, Local};
+use std::sync::Arc;
+use std::time::Duration;
+
+pub trait RetryPolicy {
+    fn is_expired(&self, first_retry_time: &mut Option<DateTime<Local>>, retry_count: u32) -> bool;
+    fn sleep_duration(&self, retry_count: u32) -> Duration;
+}
+
+#[async_trait::async_trait]
+impl<T, C> Policy<C> for T
+where
+    T: RetryPolicy + std::fmt::Debug + Send + Sync,
+    C: Send + Sync,
+{
+    async fn send(
+        &self,
+        ctx: &mut PipelineContext<C>,
+        request: &mut Request,
+        next: &[Arc<dyn Policy<C>>],
+    ) -> PolicyResult<Response> {
+        let mut first_retry_time = None;
+        let mut retry_count = 0;
+
+        loop {
+            match next[0].send(ctx, request, &next[1..]).await {
+                Ok(response) => return Ok(response),
+                Err(error) => {
+                    log::error!("Error occurred when making request: {}", error);
+                    if self.is_expired(&mut first_retry_time, retry_count) {
+                        return Err(error);
+                    } else {
+                        retry_count += 1;
+
+                        sleep(self.sleep_duration(retry_count)).await;
+                    }
+                }
+            }
+        }
+    }
+}