From 75904cfe43279cd44014d464227d36020582d060 Mon Sep 17 00:00:00 2001
From: nicolaiunrein <n.unrein@gmail.com>
Date: Thu, 20 Aug 2020 21:14:52 +0200
Subject: [PATCH] Fix RestartStrategy::timeout < 1s (#265)

---
 src/bastion/src/supervisor.rs         | 47 +++++++++++++---------
 src/bastion/tests/restart_strategy.rs | 58 +++++++++++++++++++++++++++
 2 files changed, 86 insertions(+), 19 deletions(-)

diff --git a/src/bastion/src/supervisor.rs b/src/bastion/src/supervisor.rs
index 575f5fd2..d12a086d 100644
--- a/src/bastion/src/supervisor.rs
+++ b/src/bastion/src/supervisor.rs
@@ -189,13 +189,13 @@ pub enum RestartPolicy {
 ///
 /// The default strategy used is `ActorRestartStrategy::Immediate`
 /// with the `RestartPolicy::Always` restart policy.
-#[derive(Debug, Clone, Eq, PartialEq)]
+#[derive(Debug, Clone, PartialEq)]
 pub struct RestartStrategy {
     restart_policy: RestartPolicy,
     strategy: ActorRestartStrategy,
 }
 
-#[derive(Debug, Clone, Eq, PartialEq)]
+#[derive(Debug, Clone, PartialEq)]
 /// The strategy for restating an actor as far as it
 /// returned an failure.
 ///
@@ -219,10 +219,31 @@ pub enum ActorRestartStrategy {
         /// An initial delay before the restarting an actor.
         timeout: Duration,
         /// Defines a multiplier how fast the timeout will be increasing.
-        multiplier: u64,
+        multiplier: f64,
     },
 }
 
+impl ActorRestartStrategy {
+    /// Calculate the expected restart delay for given strategy after n restarts.
+    pub fn calculate(&self, restarts_count: usize) -> Option<Duration> {
+        match *self {
+            ActorRestartStrategy::LinearBackOff { timeout } => {
+                let delay = timeout.mul_f64(restarts_count as f64);
+                Some(timeout + delay)
+            }
+            ActorRestartStrategy::ExponentialBackOff {
+                timeout,
+                multiplier,
+            } => {
+                let factor = multiplier * restarts_count as f64;
+                let delay = timeout.mul_f64(factor);
+                Some(timeout + delay)
+            }
+            _ => None,
+        }
+    }
+}
+
 impl Supervisor {
     pub(crate) fn new(bcast: Broadcast) -> Self {
         debug!("Supervisor({}): Initializing.", bcast.id());
@@ -733,7 +754,7 @@ impl Supervisor {
     ///         .with_actor_restart_strategy(           
     ///             ActorRestartStrategy::ExponentialBackOff {
     ///                 timeout: Duration::from_millis(5000),
-    ///                 multiplier: 3,
+    ///                 multiplier: 3.0,
     ///             }
     ///         )
     /// )
@@ -1892,21 +1913,9 @@ impl RestartStrategy {
     }
 
     pub(crate) async fn apply_strategy(&self, restarts_count: usize) {
-        match self.strategy {
-            ActorRestartStrategy::LinearBackOff { timeout } => {
-                let start_in = timeout.as_secs() + (timeout.as_secs() * restarts_count as u64);
-                Delay::new(Duration::from_secs(start_in)).await;
-            }
-            ActorRestartStrategy::ExponentialBackOff {
-                timeout,
-                multiplier,
-            } => {
-                let start_in =
-                    timeout.as_secs() + (timeout.as_secs() * multiplier * restarts_count as u64);
-                Delay::new(Duration::from_secs(start_in)).await;
-            }
-            _ => {}
-        };
+        if let Some(dur) = self.strategy.calculate(restarts_count) {
+            Delay::new(dur).await;
+        }
     }
 }
 
diff --git a/src/bastion/tests/restart_strategy.rs b/src/bastion/tests/restart_strategy.rs
index a875044b..28516c59 100644
--- a/src/bastion/tests/restart_strategy.rs
+++ b/src/bastion/tests/restart_strategy.rs
@@ -43,3 +43,61 @@ fn override_restart_strategy_and_policy() {
     assert_eq!(restart_strategy.restart_policy(), policy);
     assert_eq!(restart_strategy.strategy(), strategy);
 }
+
+#[test]
+fn calculate_immediate_strategy() {
+    let strategy = ActorRestartStrategy::Immediate;
+
+    assert_eq!(strategy.calculate(0), None);
+    assert_eq!(strategy.calculate(1), None);
+    assert_eq!(strategy.calculate(100), None);
+}
+
+#[test]
+fn calculate_linear_strategy() {
+    let strategy = ActorRestartStrategy::LinearBackOff {
+        timeout: Duration::from_millis(100),
+    };
+
+    assert_eq!(strategy.calculate(0), Some(Duration::from_millis(100)));
+
+    assert_eq!(
+        strategy.calculate(1),
+        Some(Duration::from_millis(100 + 1 * 100))
+    );
+    assert_eq!(
+        strategy.calculate(99),
+        Some(Duration::from_millis(100 + 99 * 100))
+    );
+}
+
+#[test]
+fn calculate_exp_strategy_with_multiplier_zero() {
+    let strategy = ActorRestartStrategy::ExponentialBackOff {
+        timeout: Duration::from_millis(100),
+        multiplier: 0.0,
+    };
+
+    assert_eq!(strategy.calculate(0), Some(Duration::from_millis(100)));
+    assert_eq!(strategy.calculate(1), Some(Duration::from_millis(100)));
+    assert_eq!(strategy.calculate(100), Some(Duration::from_millis(100)));
+}
+
+#[test]
+fn calculate_exp_strategy_with_multiplier_non_zero() {
+    let strategy = ActorRestartStrategy::ExponentialBackOff {
+        timeout: Duration::from_millis(100),
+        multiplier: 5.0,
+    };
+
+    assert_eq!(strategy.calculate(0), Some(Duration::from_millis(100)));
+
+    assert_eq!(
+        strategy.calculate(1),
+        Some(Duration::from_millis(100 + 1 * 5 * 100))
+    );
+    assert_eq!(
+        strategy.calculate(99),
+        Some(Duration::from_millis(100 + 99 * 5 * 100))
+    );
+}