use std::time::Duration;
use crate::error::{KnafehError, RpcStatusCode};
#[derive(Debug, Clone)]
pub struct RetryPolicy {
pub max_retries: u32,
pub initial_backoff: Duration,
pub max_backoff: Duration,
pub backoff_multiplier: f64,
pub jitter: bool,
}
impl RetryPolicy {
pub fn exponential(max_retries: u32) -> Self {
Self {
max_retries,
initial_backoff: Duration::from_millis(100),
max_backoff: Duration::from_secs(10),
backoff_multiplier: 2.0,
jitter: true,
}
}
pub fn none() -> Self {
Self {
max_retries: 0,
initial_backoff: Duration::ZERO,
max_backoff: Duration::ZERO,
backoff_multiplier: 1.0,
jitter: false,
}
}
pub fn with_initial_backoff(mut self, duration: Duration) -> Self {
self.initial_backoff = duration;
self
}
pub fn with_max_backoff(mut self, duration: Duration) -> Self {
self.max_backoff = duration;
self
}
pub fn with_jitter(mut self, jitter: bool) -> Self {
self.jitter = jitter;
self
}
pub fn backoff_for(&self, attempt: u32) -> Duration {
let mut backoff =
self.initial_backoff.as_secs_f64() * self.backoff_multiplier.powi(attempt as i32);
if backoff > self.max_backoff.as_secs_f64() {
backoff = self.max_backoff.as_secs_f64();
}
if self.jitter {
let jitter_factor = 0.75 + (((attempt as f64 * 0.618).fract()) * 0.5);
backoff *= jitter_factor;
}
Duration::from_secs_f64(backoff)
}
pub fn is_retryable_status(code: RpcStatusCode) -> bool {
matches!(
code,
RpcStatusCode::Unavailable | RpcStatusCode::ResourceExhausted
)
}
pub fn is_retryable(err: &KnafehError) -> bool {
matches!(
err,
KnafehError::Transport(_)
| KnafehError::ConnectionClosed
| KnafehError::Timeout
| KnafehError::Service {
code: RpcStatusCode::Unavailable,
..
}
| KnafehError::Service {
code: RpcStatusCode::ResourceExhausted,
..
}
)
}
}
impl Default for RetryPolicy {
fn default() -> Self {
Self::none()
}
}