diff --git a/src/exponential.rs b/src/exponential.rs index abaa28d..8dfe819 100644 --- a/src/exponential.rs +++ b/src/exponential.rs @@ -12,6 +12,8 @@ pub struct ExponentialBackoff { pub current_interval: Duration, /// The initial retry interval. pub initial_interval: Duration, + /// The current attempt count. + pub attempt_count: usize, /// The randomization factor to use for creating a range around the retry interval. /// /// A randomization factor of 0.5 results in a random period ranging between 50% below and 50% @@ -28,6 +30,9 @@ pub struct ExponentialBackoff { /// The maximum elapsed time after instantiating [`ExponentialBackfff`](struct.ExponentialBackoff.html) or calling /// [`reset`](trait.Backoff.html#method.reset) after which [`next_backoff`](../trait.Backoff.html#method.reset) returns `None`. pub max_elapsed_time: Option, + /// The maximum number of attempts after instantiating [`ExponentialBackfff`](struct.ExponentialBackoff.html) or calling + /// [`reset`](trait.Backoff.html#method.reset) after which [`next_backoff`](../trait.Backoff.html#method.reset) returns `None`. + pub max_attempts: Option, /// The clock used to get the current time. pub clock: C, } @@ -44,8 +49,10 @@ where multiplier: default::MULTIPLIER, max_interval: Duration::from_millis(default::MAX_INTERVAL_MILLIS), max_elapsed_time: Some(Duration::from_millis(default::MAX_ELAPSED_TIME_MILLIS)), + max_attempts: None, clock: C::default(), start_time: Instant::now(), + attempt_count: 0, }; eb.reset(); eb @@ -106,6 +113,7 @@ where fn reset(&mut self) { self.current_interval = self.initial_interval; self.start_time = self.clock.now(); + self.attempt_count = 0 } fn next_backoff(&mut self) -> Option { @@ -114,6 +122,12 @@ where match self.max_elapsed_time { Some(v) if elapsed_time > v => None, _ => { + if self.max_attempts == Some(self.attempt_count) { + return None; + } + + self.attempt_count += 1; + let random = rand::random::(); let randomized_interval = Self::get_random_value_from_interval( self.randomization_factor, @@ -156,6 +170,7 @@ pub struct ExponentialBackoffBuilder { multiplier: f64, max_interval: Duration, max_elapsed_time: Option, + max_attempt_count: Option, _clock: PhantomData, } @@ -168,6 +183,7 @@ impl Default for ExponentialBackoffBuilder { max_interval: Duration::from_millis(default::MAX_INTERVAL_MILLIS), max_elapsed_time: Some(Duration::from_millis(default::MAX_ELAPSED_TIME_MILLIS)), _clock: PhantomData, + max_attempt_count: None, } } } @@ -208,6 +224,13 @@ where self } + /// The maximum number of attempts after instantiating [`ExponentialBackfff`](struct.ExponentialBackoff.html) or calling + /// [`reset`](trait.Backoff.html#method.reset) after which [`next_backoff`](../trait.Backoff.html#method.reset) returns `None`. + pub fn with_max_attempts(&mut self, max_attempt_count: usize) -> &mut Self { + self.max_attempt_count = Some(max_attempt_count); + self + } + /// The maximum elapsed time after instantiating [`ExponentialBackfff`](struct.ExponentialBackoff.html) or calling /// [`reset`](trait.Backoff.html#method.reset) after which [`next_backoff`](../trait.Backoff.html#method.reset) returns `None`. pub fn with_max_elapsed_time(&mut self, max_elapsed_time: Option) -> &mut Self { @@ -225,6 +248,8 @@ where max_elapsed_time: self.max_elapsed_time, clock: C::default(), start_time: Instant::now(), + attempt_count: 0, + max_attempts: self.max_attempt_count, } } } @@ -246,6 +271,25 @@ fn get_randomized_interval() { assert_eq!(Duration::new(0, 3), f(0.5, 0.99, Duration::new(0, 2))); } +#[test] +fn max_attempt_count() { + let mut backoff: ExponentialBackoff = ExponentialBackoffBuilder::new() + .with_max_attempts(2) + .build(); + + assert!(backoff.next_backoff().is_some()); + assert!(backoff.next_backoff().is_some()); + assert!(backoff.next_backoff().is_none()); + assert!(backoff.next_backoff().is_none()); + + backoff.reset(); + + assert!(backoff.next_backoff().is_some()); + assert!(backoff.next_backoff().is_some()); + assert!(backoff.next_backoff().is_none()); + assert!(backoff.next_backoff().is_none()); +} + #[test] fn exponential_backoff_builder() { let initial_interval = Duration::from_secs(1);