systemprompt_agent/services/shared/
resilience.rs

1use crate::services::shared::error::{AgentServiceError, Result};
2use std::time::Duration;
3use tokio::time::{sleep, timeout};
4
5#[derive(Debug, Clone, Copy)]
6pub struct RetryConfiguration {
7    pub max_attempts: u32,
8    pub initial_delay: Duration,
9    pub max_delay: Duration,
10    pub exponential_base: u32,
11}
12
13impl Default for RetryConfiguration {
14    fn default() -> Self {
15        Self {
16            max_attempts: 3,
17            initial_delay: Duration::from_millis(100),
18            max_delay: Duration::from_secs(10),
19            exponential_base: 2,
20        }
21    }
22}
23
24pub async fn retry_operation<F, Fut, T>(operation: F, config: RetryConfiguration) -> Result<T>
25where
26    F: Fn() -> Fut,
27    Fut: std::future::Future<Output = Result<T>>,
28{
29    let mut current_delay = config.initial_delay;
30
31    for attempt in 1..=config.max_attempts {
32        match operation().await {
33            Ok(result) => return Ok(result),
34            Err(error) if attempt == config.max_attempts => return Err(error),
35            Err(_) => {
36                sleep(current_delay).await;
37                current_delay = calculate_next_delay(current_delay, &config);
38            },
39        }
40    }
41
42    Err(AgentServiceError::Configuration(
43        "RetryConfiguration".to_string(),
44        "Retry configuration resulted in no attempts".to_string(),
45    ))
46}
47
48fn calculate_next_delay(current: Duration, config: &RetryConfiguration) -> Duration {
49    let next = current.saturating_mul(config.exponential_base);
50    if next > config.max_delay {
51        config.max_delay
52    } else {
53        next
54    }
55}
56
57pub async fn retry_operation_with_backoff<F, Fut, T>(
58    operation: F,
59    max_attempts: u32,
60    initial_delay: Duration,
61) -> Result<T>
62where
63    F: Fn() -> Fut,
64    Fut: std::future::Future<Output = Result<T>>,
65{
66    let config = RetryConfiguration {
67        max_attempts,
68        initial_delay,
69        ..Default::default()
70    };
71    retry_operation(operation, config).await
72}
73
74pub async fn execute_with_timeout<F, T>(duration: Duration, operation: F) -> Result<T>
75where
76    F: std::future::Future<Output = Result<T>>,
77{
78    timeout(duration, operation)
79        .await
80        .unwrap_or_else(|_| Err(AgentServiceError::Timeout(duration.as_millis() as u64)))
81}
82
83#[derive(Debug, Clone, Copy)]
84pub struct TimeoutConfiguration {
85    pub default_timeout: Duration,
86    pub connect_timeout: Duration,
87    pub read_timeout: Duration,
88    pub write_timeout: Duration,
89}
90
91impl Default for TimeoutConfiguration {
92    fn default() -> Self {
93        Self {
94            default_timeout: Duration::from_secs(30),
95            connect_timeout: Duration::from_secs(10),
96            read_timeout: Duration::from_secs(30),
97            write_timeout: Duration::from_secs(30),
98        }
99    }
100}
101
102pub async fn execute_with_custom_timeout<F, T>(
103    config: TimeoutConfiguration,
104    timeout_type: TimeoutType,
105    operation: F,
106) -> Result<T>
107where
108    F: std::future::Future<Output = Result<T>>,
109{
110    let duration = match timeout_type {
111        TimeoutType::Connect => config.connect_timeout,
112        TimeoutType::Read => config.read_timeout,
113        TimeoutType::Write => config.write_timeout,
114        TimeoutType::Default => config.default_timeout,
115    };
116
117    execute_with_timeout(duration, operation).await
118}
119
120#[derive(Debug, Clone, Copy)]
121pub enum TimeoutType {
122    Connect,
123    Read,
124    Write,
125    Default,
126}
systemprompt_agent/services/shared/resilience.rs

systemprompt_agent/services/shared/
resilience.rs
