use crate::error::ForgeError;
use crate::recovery::backoff::{Backoff, ExponentialBackoff, FixedBackoff, LinearBackoff};
use std::marker::PhantomData;
use std::thread;
use std::time::Duration;
pub type RetryPredicate<E> = Box<dyn Fn(&E) -> bool + Send + Sync + 'static>;
pub enum BackoffStrategy {
Exponential(ExponentialBackoff),
Linear(LinearBackoff),
Fixed(FixedBackoff),
}
impl BackoffStrategy {
fn next_delay(&self, attempt: usize) -> Duration {
match self {
BackoffStrategy::Exponential(b) => b.next_delay(attempt),
BackoffStrategy::Linear(b) => b.next_delay(attempt),
BackoffStrategy::Fixed(b) => b.next_delay(attempt),
}
}
}
pub struct RetryExecutor<E> {
max_retries: usize,
backoff: BackoffStrategy,
retry_if: Option<RetryPredicate<E>>,
_marker: PhantomData<E>,
}
impl<E> RetryExecutor<E>
where
E: std::error::Error + 'static,
{
pub fn new_exponential() -> Self {
Self {
max_retries: 3,
backoff: BackoffStrategy::Exponential(ExponentialBackoff::default()),
retry_if: None,
_marker: PhantomData,
}
}
pub fn new_linear() -> Self {
Self {
max_retries: 3,
backoff: BackoffStrategy::Linear(LinearBackoff::default()),
retry_if: None,
_marker: PhantomData,
}
}
pub fn new_fixed(delay_ms: u64) -> Self {
Self {
max_retries: 3,
backoff: BackoffStrategy::Fixed(FixedBackoff::new(delay_ms)),
retry_if: None,
_marker: PhantomData,
}
}
pub fn with_max_retries(mut self, max_retries: usize) -> Self {
self.max_retries = max_retries;
self
}
pub fn with_retry_if<F>(mut self, predicate: F) -> Self
where
F: Fn(&E) -> bool + Send + Sync + 'static,
{
self.retry_if = Some(Box::new(predicate));
self
}
pub fn retry<F, T>(&self, mut operation: F) -> Result<T, E>
where
F: FnMut() -> Result<T, E>,
{
let mut attempt = 0;
loop {
match operation() {
Ok(value) => return Ok(value),
Err(err) => {
if attempt >= self.max_retries {
return Err(err);
}
let should_retry = match &self.retry_if {
Some(predicate) => predicate(&err),
None => true,
};
if !should_retry {
return Err(err);
}
let delay = self.backoff.next_delay(attempt);
thread::sleep(delay);
attempt += 1;
}
}
}
}
pub fn retry_with_handler<F, H, T>(&self, mut operation: F, mut on_error: H) -> Result<T, E>
where
F: FnMut() -> Result<T, E>,
H: FnMut(&E, usize, Duration),
{
let mut attempt = 0;
loop {
match operation() {
Ok(value) => return Ok(value),
Err(err) => {
if attempt >= self.max_retries {
return Err(err);
}
let should_retry = match &self.retry_if {
Some(predicate) => predicate(&err),
None => true,
};
if !should_retry {
return Err(err);
}
let delay = self.backoff.next_delay(attempt);
on_error(&err, attempt, delay);
thread::sleep(delay);
attempt += 1;
}
}
}
}
}
pub struct RetryPolicy {
max_retries: usize,
backoff_type: BackoffType,
}
pub enum BackoffType {
Exponential,
Linear,
Fixed(u64),
}
impl RetryPolicy {
pub fn new_exponential() -> Self {
Self {
max_retries: 3,
backoff_type: BackoffType::Exponential,
}
}
pub fn new_linear() -> Self {
Self {
max_retries: 3,
backoff_type: BackoffType::Linear,
}
}
pub fn new_fixed(delay_ms: u64) -> Self {
Self {
max_retries: 3,
backoff_type: BackoffType::Fixed(delay_ms),
}
}
pub fn with_max_retries(mut self, max_retries: usize) -> Self {
self.max_retries = max_retries;
self
}
pub fn executor<E>(&self) -> RetryExecutor<E>
where
E: std::error::Error + 'static,
{
let executor = match self.backoff_type {
BackoffType::Exponential => RetryExecutor::new_exponential(),
BackoffType::Linear => RetryExecutor::new_linear(),
BackoffType::Fixed(delay_ms) => RetryExecutor::new_fixed(delay_ms),
};
executor.with_max_retries(self.max_retries)
}
pub fn forge_executor<E>(&self) -> RetryExecutor<E>
where
E: ForgeError,
{
self.executor::<E>().with_retry_if(|err| err.is_retryable())
}
pub fn retry<F, T, E>(&self, operation: F) -> Result<T, E>
where
F: FnMut() -> Result<T, E>,
E: std::error::Error + 'static,
{
self.executor::<E>().retry(operation)
}
}
impl Default for RetryPolicy {
fn default() -> Self {
Self::new_exponential()
}
}