1use std::{future::Future, time::Duration};
2
3#[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
38pub 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
60pub 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}