Skip to main content

specado_core/
retry.rs

1use crate::error::Result;
2use std::future::Future;
3use std::time::Duration;
4
5#[derive(Debug, Clone)]
6pub struct RetryPolicy {
7    max_attempts: usize,
8    base_delay: Duration,
9    max_delay: Duration,
10}
11
12impl RetryPolicy {
13    pub fn new(max_attempts: usize, base_delay: Duration, max_delay: Duration) -> Self {
14        Self {
15            max_attempts: max_attempts.max(1),
16            base_delay,
17            max_delay,
18        }
19    }
20
21    pub fn max_attempts(&self) -> usize {
22        self.max_attempts
23    }
24
25    fn backoff_delay(&self, attempt: usize) -> Duration {
26        if attempt <= 1 {
27            return self.base_delay.min(self.max_delay);
28        }
29
30        let mut delay = self.base_delay;
31        for _ in 1..attempt {
32            delay = delay.checked_mul(2).unwrap_or(self.max_delay);
33            if delay >= self.max_delay {
34                return self.max_delay;
35            }
36        }
37
38        delay.min(self.max_delay)
39    }
40
41    pub async fn execute<F, Fut, T>(&self, mut operation: F) -> Result<T>
42    where
43        F: FnMut() -> Fut,
44        Fut: Future<Output = Result<T>> + Send,
45        T: Send,
46    {
47        let mut attempt = 0;
48        loop {
49            attempt += 1;
50            match operation().await {
51                Ok(value) => return Ok(value),
52                Err(err) => {
53                    if attempt >= self.max_attempts {
54                        return Err(err);
55                    }
56                    let delay = self.backoff_delay(attempt);
57                    if !delay.is_zero() {
58                        tokio::time::sleep(delay).await;
59                    } else {
60                        tokio::task::yield_now().await;
61                    }
62                }
63            }
64        }
65    }
66}
67
68impl Default for RetryPolicy {
69    fn default() -> Self {
70        Self {
71            max_attempts: 3,
72            base_delay: Duration::from_millis(100),
73            max_delay: Duration::from_secs(10),
74        }
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81    use crate::error::Error;
82    use std::sync::{Arc, Mutex};
83
84    #[tokio::test]
85    async fn retries_until_success() {
86        let policy = RetryPolicy::new(3, Duration::from_millis(0), Duration::from_millis(0));
87        let attempts = Arc::new(Mutex::new(0));
88        let result = policy
89            .execute({
90                let attempts = attempts.clone();
91                move || {
92                    let attempts = attempts.clone();
93                    Box::pin(async move {
94                        let mut guard = attempts.lock().unwrap();
95                        *guard += 1;
96                        if *guard < 2 {
97                            Err(Error::Transform("fail".into()))
98                        } else {
99                            Ok("ok")
100                        }
101                    })
102                }
103            })
104            .await;
105
106        assert_eq!(*attempts.lock().unwrap(), 2);
107        assert_eq!(result.unwrap(), "ok");
108    }
109
110    #[tokio::test]
111    async fn stops_after_max_attempts() {
112        let policy = RetryPolicy::new(2, Duration::from_millis(0), Duration::from_millis(0));
113        let attempts = Arc::new(Mutex::new(0));
114        let result: Result<()> = policy
115            .execute({
116                let attempts = attempts.clone();
117                move || {
118                    let attempts = attempts.clone();
119                    Box::pin(async move {
120                        let mut guard = attempts.lock().unwrap();
121                        *guard += 1;
122                        Err(Error::Transform("still failing".into()))
123                    })
124                }
125            })
126            .await;
127
128        assert!(result.is_err());
129        assert_eq!(*attempts.lock().unwrap(), 2);
130    }
131
132    #[test]
133    fn backoff_caps_at_max_delay() {
134        let policy = RetryPolicy::new(5, Duration::from_millis(100), Duration::from_millis(350));
135
136        assert_eq!(policy.backoff_delay(1), Duration::from_millis(100));
137        assert_eq!(policy.backoff_delay(2), Duration::from_millis(200));
138        assert_eq!(policy.backoff_delay(3), Duration::from_millis(350));
139        assert_eq!(policy.backoff_delay(4), Duration::from_millis(350));
140    }
141}