use crate::constants::{
DEFAULT_MAX_RETRIES, DEFAULT_RETRY_DELAY_SECS, DEPRECATED_INFINITE_RETRY_CAP,
MAX_RETRY_DELAY_SECS,
};
use crate::utils::config::get_env_or_none;
use pretty_simple_display::{DebugPretty, DisplaySimple};
use serde::{Deserialize, Serialize};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
#[derive(DebugPretty, DisplaySimple, Clone, Deserialize, Serialize)]
pub struct RetryConfig {
pub max_retry_count: Option<u32>,
pub retry_delay_secs: Option<u64>,
}
impl RetryConfig {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
#[deprecated(note = "unbounded retry is banned; use a finite RetryConfig")]
pub fn infinite() -> Self {
Self {
max_retry_count: Some(DEPRECATED_INFINITE_RETRY_CAP),
retry_delay_secs: get_env_or_none("RETRY_DELAY_SECS"),
}
}
#[must_use]
pub fn with_max_retries(max_retries: u32) -> Self {
Self {
max_retry_count: Some(max_retries),
retry_delay_secs: get_env_or_none("RETRY_DELAY_SECS"),
}
}
#[must_use]
#[deprecated(note = "unbounded retry is banned; use a finite RetryConfig")]
pub fn with_delay(delay_secs: u64) -> Self {
Self {
max_retry_count: Some(DEPRECATED_INFINITE_RETRY_CAP),
retry_delay_secs: Some(delay_secs),
}
}
#[must_use]
pub fn with_max_retries_and_delay(max_retries: u32, delay_secs: u64) -> Self {
Self {
max_retry_count: Some(max_retries),
retry_delay_secs: Some(delay_secs),
}
}
#[must_use]
pub fn max_retries(&self) -> u32 {
self.max_retry_count.unwrap_or(DEFAULT_MAX_RETRIES)
}
#[must_use]
pub fn delay_secs(&self) -> u64 {
self.retry_delay_secs.unwrap_or(DEFAULT_RETRY_DELAY_SECS)
}
#[must_use]
pub fn delay_for_attempt(&self, attempt: u32) -> Duration {
backoff_delay(Duration::from_secs(self.delay_secs()), attempt)
}
}
#[must_use]
pub(crate) fn backoff_delay(base: Duration, attempt: u32) -> Duration {
let cap = Duration::from_secs(MAX_RETRY_DELAY_SECS);
let base_ms = u64::try_from(base.as_millis()).unwrap_or(u64::MAX);
let cap_ms = u64::try_from(cap.as_millis()).unwrap_or(u64::MAX);
let factor = 1u64.checked_shl(attempt).unwrap_or(u64::MAX);
let scaled_ms = base_ms.saturating_mul(factor).min(cap_ms);
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, |d| u64::from(d.subsec_nanos()));
let max_jitter = scaled_ms / 4;
let jitter_ms = if max_jitter == 0 {
0
} else {
nanos % (max_jitter + 1)
};
Duration::from_millis(scaled_ms.saturating_add(jitter_ms))
}
impl Default for RetryConfig {
fn default() -> Self {
let max_retry_count: Option<u32> = get_env_or_none("MAX_RETRY_COUNT");
let retry_delay_secs: Option<u64> = get_env_or_none("RETRY_DELAY_SECS");
Self {
max_retry_count,
retry_delay_secs,
}
}
}
#[cfg(test)]
mod tests {
use super::RetryConfig;
use crate::constants::{
DEFAULT_MAX_RETRIES, DEFAULT_RETRY_DELAY_SECS, DEPRECATED_INFINITE_RETRY_CAP,
MAX_RETRY_DELAY_SECS,
};
use std::time::Duration;
#[test]
fn test_retry_config_default_is_finite() {
let config = RetryConfig {
max_retry_count: None,
retry_delay_secs: None,
};
assert_eq!(config.max_retries(), DEFAULT_MAX_RETRIES);
assert_eq!(config.delay_secs(), DEFAULT_RETRY_DELAY_SECS);
}
#[test]
fn test_retry_config_explicit_zero_is_not_infinite() {
let config = RetryConfig {
max_retry_count: Some(0),
retry_delay_secs: None,
};
assert_eq!(config.max_retries(), 0);
}
#[test]
fn test_delay_for_attempt_grows_and_caps() {
let config = RetryConfig {
max_retry_count: None,
retry_delay_secs: Some(1), };
for attempt in 0u32..4 {
let base_ms: u64 = 1000 * (1u64 << attempt);
let expected = base_ms.min(MAX_RETRY_DELAY_SECS * 1000);
let got =
u64::try_from(config.delay_for_attempt(attempt).as_millis()).unwrap_or(u64::MAX);
assert!(
got >= expected,
"attempt {attempt}: {got} < lower bound {expected}"
);
let upper = expected + expected / 4;
assert!(
got <= upper,
"attempt {attempt}: {got} > upper bound {upper}"
);
}
let d0 = config.delay_for_attempt(0);
let d1 = config.delay_for_attempt(1);
let d2 = config.delay_for_attempt(2);
assert!(d1 >= Duration::from_millis(2000) && d0 <= Duration::from_millis(1250));
assert!(d2 >= Duration::from_millis(4000) && d1 <= Duration::from_millis(2500));
}
#[test]
fn test_delay_for_attempt_saturates_at_cap() {
let config = RetryConfig {
max_retry_count: None,
retry_delay_secs: Some(10),
};
let got = config.delay_for_attempt(60);
let cap = Duration::from_secs(MAX_RETRY_DELAY_SECS);
assert!(got >= cap);
assert!(got <= cap + cap / 4);
}
#[test]
#[allow(deprecated)]
fn test_deprecated_infinite_is_finite() {
assert_eq!(
RetryConfig::infinite().max_retries(),
DEPRECATED_INFINITE_RETRY_CAP
);
assert_eq!(
RetryConfig::with_delay(5).max_retries(),
DEPRECATED_INFINITE_RETRY_CAP
);
assert_eq!(RetryConfig::with_delay(5).delay_secs(), 5);
}
}