use std::num::NonZeroU32;
use std::time::Duration;
#[cfg(feature = "config")]
use qubit_config::ConfigReader;
#[cfg(feature = "config")]
use super::retry_config_values::RetryConfigValues;
use super::{
AttemptTimeoutOption,
EffectiveAttemptTimeout,
};
use crate::constants::{
DEFAULT_RETRY_MAX_ATTEMPTS,
DEFAULT_RETRY_MAX_OPERATION_ELAPSED,
DEFAULT_RETRY_MAX_TOTAL_ELAPSED,
DEFAULT_RETRY_WORKER_CANCEL_GRACE_MILLIS,
KEY_ATTEMPT_TIMEOUT_MILLIS,
KEY_DELAY,
KEY_JITTER_FACTOR,
KEY_MAX_ATTEMPTS,
};
use crate::{
AttemptFailureDecision,
AttemptTimeoutSource,
RetryConfigError,
RetryDelay,
RetryErrorReason,
RetryJitter,
};
#[derive(Debug, Clone, PartialEq)]
pub struct RetryOptions {
pub(crate) max_attempts: NonZeroU32,
pub(crate) max_operation_elapsed: Option<Duration>,
pub(crate) max_total_elapsed: Option<Duration>,
pub(crate) delay: RetryDelay,
pub(crate) jitter: RetryJitter,
pub(crate) attempt_timeout: Option<AttemptTimeoutOption>,
pub(crate) worker_cancel_grace: Duration,
}
impl RetryOptions {
#[inline]
pub fn max_attempts(&self) -> u32 {
self.max_attempts.get()
}
#[inline]
pub fn max_operation_elapsed(&self) -> Option<Duration> {
self.max_operation_elapsed
}
#[inline]
pub fn max_total_elapsed(&self) -> Option<Duration> {
self.max_total_elapsed
}
#[inline]
pub fn delay(&self) -> &RetryDelay {
&self.delay
}
#[inline]
pub fn jitter(&self) -> RetryJitter {
self.jitter
}
#[inline]
pub fn attempt_timeout(&self) -> Option<AttemptTimeoutOption> {
self.attempt_timeout
}
#[inline]
pub fn worker_cancel_grace(&self) -> Duration {
self.worker_cancel_grace
}
pub fn new(
max_attempts: u32,
max_operation_elapsed: Option<Duration>,
max_total_elapsed: Option<Duration>,
delay: RetryDelay,
jitter: RetryJitter,
) -> Result<Self, RetryConfigError> {
Self::new_with_attempt_timeout(
max_attempts,
max_operation_elapsed,
max_total_elapsed,
delay,
jitter,
None,
)
}
pub fn new_with_attempt_timeout(
max_attempts: u32,
max_operation_elapsed: Option<Duration>,
max_total_elapsed: Option<Duration>,
delay: RetryDelay,
jitter: RetryJitter,
attempt_timeout: Option<AttemptTimeoutOption>,
) -> Result<Self, RetryConfigError> {
let max_attempts = NonZeroU32::new(max_attempts).ok_or_else(|| {
RetryConfigError::invalid_value(
KEY_MAX_ATTEMPTS,
"max_attempts must be greater than zero",
)
})?;
let options = Self {
max_attempts,
max_operation_elapsed,
max_total_elapsed,
delay,
jitter,
attempt_timeout,
worker_cancel_grace: Duration::from_millis(DEFAULT_RETRY_WORKER_CANCEL_GRACE_MILLIS),
};
options.validate()?;
Ok(options)
}
#[cfg(feature = "config")]
pub fn from_config<R>(config: &R) -> Result<Self, RetryConfigError>
where
R: ConfigReader + ?Sized,
{
let default = Self::default();
let values = RetryConfigValues::new(config).map_err(RetryConfigError::from)?;
values.to_options(&default)
}
pub fn validate(&self) -> Result<(), RetryConfigError> {
self.delay
.validate()
.map_err(|message| RetryConfigError::invalid_value(KEY_DELAY, message))?;
self.jitter
.validate()
.map_err(|message| RetryConfigError::invalid_value(KEY_JITTER_FACTOR, message))?;
if let Some(attempt_timeout) = self.attempt_timeout {
attempt_timeout.validate().map_err(|message| {
RetryConfigError::invalid_value(KEY_ATTEMPT_TIMEOUT_MILLIS, message)
})?;
}
Ok(())
}
pub fn base_delay_for_attempt(&self, attempt: u32) -> Duration {
self.delay.base_delay(attempt)
}
pub fn delay_for_attempt(&self, attempt: u32) -> Duration {
self.jitter.delay_for_attempt(&self.delay, attempt)
}
pub fn next_base_delay_from_current(&self, current: Duration) -> Duration {
match &self.delay {
RetryDelay::None => Duration::ZERO,
RetryDelay::Fixed(delay) => *delay,
RetryDelay::Random { .. } => self.delay.base_delay(1),
RetryDelay::Exponential {
initial,
max,
multiplier,
} => {
if current.is_zero() {
return *initial;
}
let bounded_current = current.min(*max);
let next = bounded_current.mul_f64(*multiplier);
if next > *max { *max } else { next }
}
}
}
pub fn jittered_delay(&self, base_delay: Duration) -> Duration {
self.jitter.apply(base_delay)
}
pub fn next_delay_from_current(&self, current: Duration) -> Duration {
self.jittered_delay(self.next_base_delay_from_current(current))
}
#[inline]
pub(crate) fn attempt_timeout_duration(&self) -> Option<Duration> {
self.attempt_timeout
.map(|attempt_timeout| attempt_timeout.timeout())
}
pub(crate) fn effective_attempt_timeout(
&self,
operation_elapsed: Duration,
total_elapsed: Duration,
) -> EffectiveAttemptTimeout {
let candidates = [
self.attempt_timeout_duration()
.map(|duration| (duration, AttemptTimeoutSource::Configured)),
self.remaining_operation_elapsed(operation_elapsed)
.map(|duration| (duration, AttemptTimeoutSource::MaxOperationElapsed)),
self.remaining_total_elapsed(total_elapsed)
.map(|duration| (duration, AttemptTimeoutSource::MaxTotalElapsed)),
];
let selected = candidates
.into_iter()
.flatten()
.min_by(|left, right| left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1)));
match selected {
Some((duration, source)) => EffectiveAttemptTimeout::new(Some(duration), Some(source)),
None => EffectiveAttemptTimeout::none(),
}
}
#[inline]
pub(crate) fn elapsed_error_reason(
&self,
operation_elapsed: Duration,
total_elapsed: Duration,
) -> Option<RetryErrorReason> {
if self
.max_operation_elapsed
.is_some_and(|max_operation_elapsed| operation_elapsed >= max_operation_elapsed)
{
Some(RetryErrorReason::MaxOperationElapsedExceeded)
} else if self
.max_total_elapsed
.is_some_and(|max_total_elapsed| total_elapsed >= max_total_elapsed)
{
Some(RetryErrorReason::MaxTotalElapsedExceeded)
} else {
None
}
}
#[inline]
pub(crate) fn retry_sleep_exhausts_total_elapsed(
&self,
total_elapsed: Duration,
delay: Duration,
) -> bool {
if delay.is_zero() {
return false;
}
let Some(max_total_elapsed) = self.max_total_elapsed else {
return false;
};
total_elapsed.saturating_add(delay) >= max_total_elapsed
}
pub(crate) fn retry_delay(
&self,
decision: AttemptFailureDecision,
attempts: u32,
hint: Option<Duration>,
) -> Duration {
match decision {
AttemptFailureDecision::RetryAfter(delay) => delay,
AttemptFailureDecision::UseDefault => {
hint.unwrap_or_else(|| self.jitter.delay_for_attempt(&self.delay, attempts))
}
AttemptFailureDecision::Retry | AttemptFailureDecision::Abort => {
self.jitter.delay_for_attempt(&self.delay, attempts)
}
}
}
#[inline]
fn remaining_operation_elapsed(&self, operation_elapsed: Duration) -> Option<Duration> {
self.max_operation_elapsed
.map(|max_operation_elapsed| max_operation_elapsed.saturating_sub(operation_elapsed))
}
#[inline]
fn remaining_total_elapsed(&self, total_elapsed: Duration) -> Option<Duration> {
self.max_total_elapsed
.map(|max_total_elapsed| max_total_elapsed.saturating_sub(total_elapsed))
}
}
impl Default for RetryOptions {
#[inline]
fn default() -> Self {
Self {
max_attempts: NonZeroU32::new(DEFAULT_RETRY_MAX_ATTEMPTS)
.expect("default retry attempts must be non-zero"),
max_operation_elapsed: DEFAULT_RETRY_MAX_OPERATION_ELAPSED,
max_total_elapsed: DEFAULT_RETRY_MAX_TOTAL_ELAPSED,
delay: RetryDelay::default(),
jitter: RetryJitter::default(),
attempt_timeout: None,
worker_cancel_grace: Duration::from_millis(DEFAULT_RETRY_WORKER_CANCEL_GRACE_MILLIS),
}
}
}