use std::time::Duration;
#[derive(Debug, Clone)]
pub struct StageRetryConfig {
pub max_attempts: usize,
pub initial_delay: Duration,
pub max_delay: Duration,
pub multiplier: f64,
}
impl Default for StageRetryConfig {
fn default() -> Self {
Self {
max_attempts: 3,
initial_delay: Duration::from_millis(100),
max_delay: Duration::from_secs(10),
multiplier: 2.0,
}
}
}
impl StageRetryConfig {
pub fn new() -> Self {
Self::default()
}
pub fn with_max_attempts(mut self, n: usize) -> Self {
self.max_attempts = n.max(1);
self
}
pub fn with_initial_delay(mut self, delay: Duration) -> Self {
self.initial_delay = delay;
self
}
pub fn with_max_delay(mut self, delay: Duration) -> Self {
self.max_delay = delay;
self
}
pub fn with_multiplier(mut self, multiplier: f64) -> Self {
self.multiplier = multiplier;
self
}
pub fn delay_for_attempt(&self, attempt: usize) -> Duration {
let delay_ms =
(self.initial_delay.as_millis() as f64) * self.multiplier.powi(attempt as i32);
let capped_ms = delay_ms.min(self.max_delay.as_millis() as f64);
Duration::from_millis(capped_ms as u64)
}
}
#[derive(Debug, Clone)]
pub enum FailurePolicy {
Fail,
Skip,
Retry(StageRetryConfig),
}
impl Default for FailurePolicy {
fn default() -> Self {
Self::Fail
}
}
impl FailurePolicy {
pub fn fail() -> Self {
Self::Fail
}
pub fn skip() -> Self {
Self::Skip
}
pub fn retry() -> Self {
Self::Retry(StageRetryConfig::default())
}
pub fn retry_with(config: StageRetryConfig) -> Self {
Self::Retry(config)
}
pub fn allows_continuation(&self) -> bool {
matches!(self, Self::Skip)
}
pub fn has_retry(&self) -> bool {
matches!(self, Self::Retry(_))
}
pub fn retry_config(&self) -> Option<&StageRetryConfig> {
match self {
Self::Retry(config) => Some(config),
_ => None,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_retry_config() {
let config = StageRetryConfig::default();
assert_eq!(config.max_attempts, 3);
assert_eq!(config.initial_delay, Duration::from_millis(100));
assert_eq!(config.max_delay, Duration::from_secs(10));
}
#[test]
fn test_retry_config_builder() {
let config = StageRetryConfig::new()
.with_max_attempts(5)
.with_initial_delay(Duration::from_millis(200))
.with_max_delay(Duration::from_secs(30));
assert_eq!(config.max_attempts, 5);
assert_eq!(config.initial_delay, Duration::from_millis(200));
assert_eq!(config.max_delay, Duration::from_secs(30));
}
#[test]
fn test_delay_for_attempt() {
let config = StageRetryConfig::new()
.with_initial_delay(Duration::from_millis(100))
.with_multiplier(2.0);
assert_eq!(config.delay_for_attempt(0), Duration::from_millis(100));
assert_eq!(config.delay_for_attempt(1), Duration::from_millis(200));
assert_eq!(config.delay_for_attempt(2), Duration::from_millis(400));
}
#[test]
fn test_delay_respects_max() {
let config = StageRetryConfig::new()
.with_initial_delay(Duration::from_secs(1))
.with_max_delay(Duration::from_secs(5))
.with_multiplier(10.0);
assert_eq!(config.delay_for_attempt(0), Duration::from_secs(1));
assert_eq!(config.delay_for_attempt(1), Duration::from_secs(5)); assert_eq!(config.delay_for_attempt(2), Duration::from_secs(5)); }
#[test]
fn test_failure_policy_constructors() {
assert!(matches!(FailurePolicy::fail(), FailurePolicy::Fail));
assert!(matches!(FailurePolicy::skip(), FailurePolicy::Skip));
assert!(matches!(FailurePolicy::retry(), FailurePolicy::Retry(_)));
}
#[test]
fn test_allows_continuation() {
assert!(!FailurePolicy::fail().allows_continuation());
assert!(FailurePolicy::skip().allows_continuation());
assert!(!FailurePolicy::retry().allows_continuation());
}
}