use std::time::Duration;
#[derive(Debug, Clone)]
pub struct BackoffConfig {
pub base_interval_ms: u64,
pub max_backoff_ms: u64,
pub max_exponent: u32,
}
impl Default for BackoffConfig {
fn default() -> Self {
Self {
base_interval_ms: 5000, max_backoff_ms: 60000, max_exponent: 4, }
}
}
impl BackoffConfig {
pub fn new(base_interval_ms: u64, max_backoff_ms: u64, max_exponent: u32) -> Self {
Self {
base_interval_ms,
max_backoff_ms,
max_exponent,
}
}
pub fn fast() -> Self {
Self {
base_interval_ms: 1000, max_backoff_ms: 10000, max_exponent: 3, }
}
pub fn aggressive() -> Self {
Self {
base_interval_ms: 10000, max_backoff_ms: 300000, max_exponent: 5, }
}
}
pub fn calculate_backoff(config: &BackoffConfig, consecutive_failures: u32) -> Duration {
if consecutive_failures == 0 {
return Duration::from_millis(config.base_interval_ms);
}
let exponent = (consecutive_failures.saturating_sub(1)).min(config.max_exponent);
let multiplier = 2u64.saturating_pow(exponent);
let delay_ms = config
.base_interval_ms
.saturating_mul(multiplier)
.min(config.max_backoff_ms);
Duration::from_millis(delay_ms)
}
pub fn calculate_backoff_ms(
base_interval_ms: u64,
max_backoff_ms: u64,
consecutive_failures: u32,
) -> u64 {
let config = BackoffConfig {
base_interval_ms,
max_backoff_ms,
max_exponent: 4, };
calculate_backoff(&config, consecutive_failures).as_millis() as u64
}
#[derive(Debug, Clone)]
pub enum PollResult<T> {
Continue,
Terminal(T),
Error(String),
}
#[derive(Debug, Clone)]
pub struct PollConfig {
pub backoff: BackoffConfig,
pub max_consecutive_errors: u32,
pub timeout_secs: u64,
pub poll_interval_ms: u64,
}
impl Default for PollConfig {
fn default() -> Self {
Self {
backoff: BackoffConfig::default(),
max_consecutive_errors: 5,
timeout_secs: 0, poll_interval_ms: 5000, }
}
}
#[derive(Debug)]
pub struct PollState {
pub consecutive_errors: u32,
pub total_attempts: u32,
start_time: std::time::Instant,
config: PollConfig,
}
impl PollState {
pub fn new(config: PollConfig) -> Self {
Self {
consecutive_errors: 0,
total_attempts: 0,
start_time: std::time::Instant::now(),
config,
}
}
pub fn record_success(&mut self) {
self.consecutive_errors = 0;
self.total_attempts += 1;
}
pub fn record_error(&mut self) {
self.consecutive_errors += 1;
self.total_attempts += 1;
}
pub fn should_give_up(&self) -> bool {
self.consecutive_errors >= self.config.max_consecutive_errors
}
pub fn is_timed_out(&self) -> bool {
if self.config.timeout_secs == 0 {
return false;
}
self.start_time.elapsed().as_secs() >= self.config.timeout_secs
}
pub fn next_delay(&self) -> Duration {
if self.consecutive_errors > 0 {
calculate_backoff(&self.config.backoff, self.consecutive_errors)
} else {
Duration::from_millis(self.config.poll_interval_ms)
}
}
pub fn elapsed(&self) -> Duration {
self.start_time.elapsed()
}
}
#[derive(Debug, Clone)]
pub struct RetryConfig {
pub max_attempts: u32,
pub initial_delay_ms: u64,
pub backoff_multiplier: f64,
pub max_delay_ms: u64,
}
impl Default for RetryConfig {
fn default() -> Self {
Self {
max_attempts: 3,
initial_delay_ms: 1000,
backoff_multiplier: 2.0,
max_delay_ms: 30000,
}
}
}
impl RetryConfig {
pub fn delay_for_attempt(&self, attempt: u32) -> Duration {
if attempt == 0 {
return Duration::from_millis(self.initial_delay_ms);
}
let multiplier = self.backoff_multiplier.powi(attempt as i32);
let delay_ms = (self.initial_delay_ms as f64 * multiplier) as u64;
Duration::from_millis(delay_ms.min(self.max_delay_ms))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_calculate_backoff_default() {
let config = BackoffConfig::default();
assert_eq!(calculate_backoff(&config, 1).as_millis(), 5000);
assert_eq!(calculate_backoff(&config, 2).as_millis(), 10000);
assert_eq!(calculate_backoff(&config, 3).as_millis(), 20000);
assert_eq!(calculate_backoff(&config, 4).as_millis(), 40000);
assert_eq!(calculate_backoff(&config, 5).as_millis(), 60000);
assert_eq!(calculate_backoff(&config, 6).as_millis(), 60000);
assert_eq!(calculate_backoff(&config, 10).as_millis(), 60000);
}
#[test]
fn test_calculate_backoff_zero_failures() {
let config = BackoffConfig::default();
assert_eq!(calculate_backoff(&config, 0).as_millis(), 5000);
}
#[test]
fn test_calculate_backoff_fast() {
let config = BackoffConfig::fast();
assert_eq!(calculate_backoff(&config, 1).as_millis(), 1000); assert_eq!(calculate_backoff(&config, 2).as_millis(), 2000); assert_eq!(calculate_backoff(&config, 3).as_millis(), 4000); assert_eq!(calculate_backoff(&config, 4).as_millis(), 8000); assert_eq!(calculate_backoff(&config, 5).as_millis(), 8000); assert_eq!(calculate_backoff(&config, 10).as_millis(), 8000); }
#[test]
fn test_poll_state() {
let config = PollConfig {
max_consecutive_errors: 3,
timeout_secs: 0,
..Default::default()
};
let mut state = PollState::new(config);
assert!(!state.should_give_up());
assert_eq!(state.consecutive_errors, 0);
state.record_error();
assert!(!state.should_give_up());
state.record_error();
assert!(!state.should_give_up());
state.record_error();
assert!(state.should_give_up());
state.record_success();
assert!(!state.should_give_up());
assert_eq!(state.consecutive_errors, 0);
}
#[test]
fn test_retry_delay() {
let config = RetryConfig {
max_attempts: 5,
initial_delay_ms: 1000,
backoff_multiplier: 2.0,
max_delay_ms: 10000,
};
assert_eq!(config.delay_for_attempt(0).as_millis(), 1000);
assert_eq!(config.delay_for_attempt(1).as_millis(), 2000);
assert_eq!(config.delay_for_attempt(2).as_millis(), 4000);
assert_eq!(config.delay_for_attempt(3).as_millis(), 8000);
assert_eq!(config.delay_for_attempt(4).as_millis(), 10000);
}
}