Skip to main content

rust_zero_core/
fx.rs

1use std::{future::Future, time::Duration};
2
3/// Exponential-backoff settings for a retryable operation.
4#[derive(Debug, Clone, Copy)]
5pub struct RetryPolicy {
6    pub max_attempts: usize,
7    pub initial_delay: Duration,
8    pub max_delay: Duration,
9}
10
11impl RetryPolicy {
12    pub fn new(max_attempts: usize, initial_delay: Duration) -> Self {
13        assert!(
14            max_attempts > 0,
15            "maximum attempts must be greater than zero"
16        );
17        assert!(
18            !initial_delay.is_zero(),
19            "initial delay must be greater than zero"
20        );
21        Self {
22            max_attempts,
23            initial_delay,
24            max_delay: Duration::from_secs(30),
25        }
26    }
27
28    pub fn with_max_delay(mut self, max_delay: Duration) -> Self {
29        assert!(
30            !max_delay.is_zero(),
31            "maximum delay must be greater than zero"
32        );
33        self.max_delay = max_delay;
34        self
35    }
36}
37
38/// Retries an async operation with bounded exponential backoff.
39pub async fn retry<T, E, F, Fut>(policy: RetryPolicy, mut operation: F) -> Result<T, E>
40where
41    F: FnMut() -> Fut,
42    Fut: Future<Output = Result<T, E>>,
43{
44    let mut delay = policy.initial_delay;
45
46    for attempt in 1..=policy.max_attempts {
47        match operation().await {
48            Ok(value) => return Ok(value),
49            Err(error) if attempt == policy.max_attempts => return Err(error),
50            Err(_) => {
51                tokio::time::sleep(delay).await;
52                delay = delay.saturating_mul(2).min(policy.max_delay);
53            }
54        }
55    }
56
57    unreachable!("a retry policy always has at least one attempt")
58}
59
60/// Runs an async operation with a hard deadline.
61pub async fn timeout<T, Fut>(
62    duration: Duration,
63    operation: Fut,
64) -> Result<T, tokio::time::error::Elapsed>
65where
66    Fut: Future<Output = T>,
67{
68    tokio::time::timeout(duration, operation).await
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use std::sync::atomic::{AtomicUsize, Ordering};
75
76    #[tokio::test]
77    async fn retries_until_the_operation_succeeds() {
78        let attempts = AtomicUsize::new(0);
79        let result = retry(RetryPolicy::new(3, Duration::from_millis(1)), || async {
80            let attempt = attempts.fetch_add(1, Ordering::SeqCst);
81            if attempt < 2 {
82                Err("unavailable")
83            } else {
84                Ok("connected")
85            }
86        })
87        .await;
88
89        assert_eq!(result, Ok("connected"));
90        assert_eq!(attempts.load(Ordering::SeqCst), 3);
91    }
92
93    #[tokio::test]
94    async fn deadline_cancels_slow_work() {
95        let result = timeout(Duration::from_millis(1), async {
96            tokio::time::sleep(Duration::from_millis(10)).await;
97        })
98        .await;
99
100        assert!(result.is_err());
101    }
102}