use crate::error::{Result, TorshError};
use std::fmt;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RecoveryStrategy {
FailFast,
Retry {
max_attempts: usize,
delay_ms: Option<u64>,
},
Fallback,
Skip,
UseDefault,
Custom(String),
}
#[derive(Debug, Clone)]
pub struct RecoveryContext {
pub error: TorshError,
pub operation: String,
pub attempt_count: usize,
pub max_attempts: usize,
pub context: std::collections::HashMap<String, String>,
}
#[derive(Debug)]
pub enum RecoveryResult<T> {
Success(T),
Degraded(T, String),
Retry,
Failed(TorshError),
}
pub trait ErrorRecovery<T> {
fn recover(&self, context: &RecoveryContext) -> RecoveryResult<T>;
fn can_recover(&self, error: &TorshError) -> bool;
fn strategy(&self) -> &RecoveryStrategy;
}
#[derive(Debug)]
pub struct BasicRecovery<T> {
strategy: RecoveryStrategy,
fallback_value: Option<T>,
}
impl<T: Clone> BasicRecovery<T> {
pub fn new(strategy: RecoveryStrategy) -> Self {
Self {
strategy,
fallback_value: None,
}
}
pub fn with_default(default_value: T) -> Self {
Self {
strategy: RecoveryStrategy::UseDefault,
fallback_value: Some(default_value),
}
}
pub fn retry(max_attempts: usize, delay_ms: Option<u64>) -> Self {
Self {
strategy: RecoveryStrategy::Retry {
max_attempts,
delay_ms,
},
fallback_value: None,
}
}
}
impl<T: Clone + fmt::Debug> ErrorRecovery<T> for BasicRecovery<T> {
fn recover(&self, context: &RecoveryContext) -> RecoveryResult<T> {
match &self.strategy {
RecoveryStrategy::FailFast => RecoveryResult::Failed(context.error.clone()),
RecoveryStrategy::Retry {
max_attempts,
delay_ms,
} => {
if context.attempt_count < *max_attempts {
if let Some(delay) = delay_ms {
std::thread::sleep(std::time::Duration::from_millis(*delay));
}
RecoveryResult::Retry
} else {
RecoveryResult::Failed(TorshError::RuntimeError(format!(
"Failed after {max_attempts} retry attempts"
)))
}
}
RecoveryStrategy::UseDefault => {
if let Some(ref default_val) = self.fallback_value {
RecoveryResult::Degraded(
default_val.clone(),
"Using default value due to error".to_string(),
)
} else {
RecoveryResult::Failed(TorshError::RuntimeError(
"No default value available for recovery".to_string(),
))
}
}
RecoveryStrategy::Skip => RecoveryResult::Failed(TorshError::RuntimeError(
"Skip recovery not supported for this type".to_string(),
)),
RecoveryStrategy::Fallback => RecoveryResult::Failed(TorshError::RuntimeError(
"Fallback recovery not implemented".to_string(),
)),
RecoveryStrategy::Custom(name) => RecoveryResult::Failed(TorshError::RuntimeError(
format!("Custom recovery '{name}' not implemented"),
)),
}
}
fn can_recover(&self, _error: &TorshError) -> bool {
match self.strategy {
RecoveryStrategy::FailFast => false,
RecoveryStrategy::UseDefault => self.fallback_value.is_some(),
_ => true,
}
}
fn strategy(&self) -> &RecoveryStrategy {
&self.strategy
}
}
pub struct RecoveryManager<T> {
strategies: Vec<Box<dyn ErrorRecovery<T>>>,
default_strategy: RecoveryStrategy,
}
impl<T> fmt::Debug for RecoveryManager<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("RecoveryManager")
.field("strategy_count", &self.strategies.len())
.field("default_strategy", &self.default_strategy)
.finish()
}
}
impl<T: Clone + fmt::Debug + 'static> RecoveryManager<T> {
pub fn new() -> Self {
Self {
strategies: Vec::new(),
default_strategy: RecoveryStrategy::FailFast,
}
}
pub fn add_strategy(mut self, strategy: Box<dyn ErrorRecovery<T>>) -> Self {
self.strategies.push(strategy);
self
}
pub fn with_default_strategy(mut self, strategy: RecoveryStrategy) -> Self {
self.default_strategy = strategy;
self
}
pub fn attempt_recovery(&self, mut context: RecoveryContext) -> RecoveryResult<T> {
for strategy in &self.strategies {
if strategy.can_recover(&context.error) {
match strategy.recover(&context) {
RecoveryResult::Success(result) => return RecoveryResult::Success(result),
RecoveryResult::Degraded(result, msg) => {
return RecoveryResult::Degraded(result, msg)
}
RecoveryResult::Retry => {
context.attempt_count += 1;
if context.attempt_count < context.max_attempts {
return RecoveryResult::Retry;
}
}
RecoveryResult::Failed(_) => continue,
}
}
}
let default_recovery = BasicRecovery::new(self.default_strategy.clone());
default_recovery.recover(&context)
}
}
impl<T: Clone + fmt::Debug + 'static> Default for RecoveryManager<T> {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug)]
pub struct RecoverableOperation<T> {
recovery_manager: RecoveryManager<T>,
operation_name: String,
max_attempts: usize,
}
impl<T: Clone + fmt::Debug + 'static> RecoverableOperation<T> {
pub fn new(operation_name: impl Into<String>) -> Self {
Self {
recovery_manager: RecoveryManager::new(),
operation_name: operation_name.into(),
max_attempts: 3,
}
}
pub fn with_recovery_manager(mut self, manager: RecoveryManager<T>) -> Self {
self.recovery_manager = manager;
self
}
pub fn with_max_attempts(mut self, attempts: usize) -> Self {
self.max_attempts = attempts;
self
}
pub fn execute<F>(&self, mut operation: F) -> Result<T>
where
F: FnMut() -> Result<T>,
{
let mut attempt_count = 0;
loop {
match operation() {
Ok(result) => return Ok(result),
Err(error) => {
let context = RecoveryContext {
error: error.clone(),
operation: self.operation_name.clone(),
attempt_count,
max_attempts: self.max_attempts,
context: std::collections::HashMap::new(),
};
match self.recovery_manager.attempt_recovery(context) {
RecoveryResult::Success(result) => return Ok(result),
RecoveryResult::Degraded(result, _msg) => {
return Ok(result);
}
RecoveryResult::Retry => {
attempt_count += 1;
if attempt_count >= self.max_attempts {
return Err(TorshError::RuntimeError(format!(
"Operation '{}' failed after {} attempts",
self.operation_name, self.max_attempts
)));
}
continue;
}
RecoveryResult::Failed(recovery_error) => {
return Err(TorshError::wrap_with_location(
recovery_error,
format!("Recovery failed for operation '{}'", self.operation_name),
));
}
}
}
}
}
}
pub fn execute_with_context<F>(
&self,
mut operation: F,
context: std::collections::HashMap<String, String>,
) -> Result<T>
where
F: FnMut() -> Result<T>,
{
let mut attempt_count = 0;
loop {
match operation() {
Ok(result) => return Ok(result),
Err(error) => {
let recovery_context = RecoveryContext {
error: error.clone(),
operation: self.operation_name.clone(),
attempt_count,
max_attempts: self.max_attempts,
context: context.clone(),
};
match self.recovery_manager.attempt_recovery(recovery_context) {
RecoveryResult::Success(result) => return Ok(result),
RecoveryResult::Degraded(result, _msg) => return Ok(result),
RecoveryResult::Retry => {
attempt_count += 1;
if attempt_count >= self.max_attempts {
return Err(error);
}
continue;
}
RecoveryResult::Failed(recovery_error) => return Err(recovery_error),
}
}
}
}
}
}
pub fn with_retry<T, F>(operation: F, max_attempts: usize) -> Result<T>
where
F: FnMut() -> Result<T>,
T: Clone + fmt::Debug + 'static,
{
let recovery_manager = RecoveryManager::new()
.add_strategy(Box::new(BasicRecovery::retry(max_attempts, Some(100))));
let recoverable = RecoverableOperation::new("retry_operation")
.with_recovery_manager(recovery_manager)
.with_max_attempts(max_attempts);
recoverable.execute(operation)
}
pub fn with_fallback<T, F>(operation: F, fallback: T) -> Result<T>
where
F: FnMut() -> Result<T>,
T: Clone + fmt::Debug + 'static,
{
let recovery_manager =
RecoveryManager::new().add_strategy(Box::new(BasicRecovery::with_default(fallback)));
let recoverable =
RecoverableOperation::new("fallback_operation").with_recovery_manager(recovery_manager);
recoverable.execute(operation)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_basic_recovery_retry() {
let recovery = BasicRecovery::<i32>::retry(3, Some(1));
let context = RecoveryContext {
error: TorshError::RuntimeError("Test error".to_string()),
operation: "test_op".to_string(),
attempt_count: 1,
max_attempts: 3,
context: std::collections::HashMap::new(),
};
match recovery.recover(&context) {
RecoveryResult::Retry => {} _ => panic!("Expected retry result"),
}
}
#[test]
fn test_basic_recovery_default() {
let recovery = BasicRecovery::with_default(42);
let context = RecoveryContext {
error: TorshError::RuntimeError("Test error".to_string()),
operation: "test_op".to_string(),
attempt_count: 1,
max_attempts: 3,
context: std::collections::HashMap::new(),
};
match recovery.recover(&context) {
RecoveryResult::Degraded(value, _) => assert_eq!(value, 42),
_ => panic!("Expected degraded result with default value"),
}
}
#[test]
fn test_recoverable_operation_success() {
let recoverable = RecoverableOperation::new("test_operation");
let result = recoverable.execute(|| Ok(42));
assert_eq!(result.expect("execute should succeed"), 42);
}
#[test]
fn test_with_retry_convenience() {
let mut counter = 0;
let result = with_retry(
|| {
counter += 1;
if counter < 3 {
Err(TorshError::RuntimeError("Not ready".to_string()))
} else {
Ok(42)
}
},
5,
);
assert_eq!(result.expect("with_retry should succeed"), 42);
assert_eq!(counter, 3);
}
#[test]
fn test_with_fallback_convenience() {
let result = with_fallback(
|| Err(TorshError::RuntimeError("Operation failed".to_string())),
100,
);
assert_eq!(result.expect("with_fallback should succeed"), 100);
}
#[test]
fn test_recovery_manager() {
let manager = RecoveryManager::new()
.add_strategy(Box::new(BasicRecovery::retry(2, Some(1))))
.add_strategy(Box::new(BasicRecovery::with_default(999)));
let context = RecoveryContext {
error: TorshError::RuntimeError("Test error".to_string()),
operation: "test_op".to_string(),
attempt_count: 3, max_attempts: 3,
context: std::collections::HashMap::new(),
};
match manager.attempt_recovery(context) {
RecoveryResult::Degraded(value, _) => assert_eq!(value, 999),
_ => panic!("Expected degraded result with default value"),
}
}
}