foundry_common/
retry.rs

Help
1//! Retry utilities.
2
3use eyre::{Error, Report, Result};
4use std::{future::Future, time::Duration};
5
6/// Error type for Retry.
7#[derive(Debug, thiserror::Error)]
8pub enum RetryError<E = Report> {
9    /// Continues operation without decrementing retries.
10    Continue(E),
11    /// Keeps retrying operation.
12    Retry(E),
13    /// Stops retrying operation immediately.
14    Break(E),
15}
16
17/// A type that keeps track of attempts.
18#[derive(Clone, Debug)]
19pub struct Retry {
20    retries: u32,
21    delay: Duration,
22}
23
24impl Retry {
25    /// Creates a new `Retry` instance.
26    pub fn new(retries: u32, delay: Duration) -> Self {
27        Self { retries, delay }
28    }
29
30    /// Creates a new `Retry` instance with no delay between retries.
31    pub fn new_no_delay(retries: u32) -> Self {
32        Self::new(retries, Duration::ZERO)
33    }
34
35    /// Runs the given closure in a loop, retrying if it fails up to the specified number of times.
36    pub fn run<F: FnMut() -> Result<T>, T>(mut self, mut callback: F) -> Result<T> {
37        loop {
38            match callback() {
39                Err(e) if self.retries > 0 => {
40                    self.handle_err(e);
41                    if !self.delay.is_zero() {
42                        std::thread::sleep(self.delay);
43                    }
44                }
45                res => return res,
46            }
47        }
48    }
49
50    /// Runs the given async closure in a loop, retrying if it fails up to the specified number of
51    /// times.
52    pub async fn run_async<F, Fut, T>(mut self, mut callback: F) -> Result<T>
53    where
54        F: FnMut() -> Fut,
55        Fut: Future<Output = Result<T>>,
56    {
57        loop {
58            match callback().await {
59                Err(e) if self.retries > 0 => {
60                    self.handle_err(e);
61                    if !self.delay.is_zero() {
62                        tokio::time::sleep(self.delay).await;
63                    }
64                }
65                res => return res,
66            };
67        }
68    }
69
70    /// Runs the given async closure in a loop, retrying if it fails up to the specified number of
71    /// times or immediately returning an error if the closure returned [`RetryError::Break`].
72    pub async fn run_async_until_break<F, Fut, T>(mut self, mut callback: F) -> Result<T>
73    where
74        F: FnMut() -> Fut,
75        Fut: Future<Output = Result<T, RetryError>>,
76    {
77        loop {
78            match callback().await {
79                Err(RetryError::Continue(e)) => {
80                    self.log(e, false);
81                    if !self.delay.is_zero() {
82                        tokio::time::sleep(self.delay).await;
83                    }
84                }
85                Err(RetryError::Retry(e)) if self.retries > 0 => {
86                    self.handle_err(e);
87                    if !self.delay.is_zero() {
88                        tokio::time::sleep(self.delay).await;
89                    }
90                }
91                Err(RetryError::Retry(e) | RetryError::Break(e)) => return Err(e),
92                Ok(t) => return Ok(t),
93            };
94        }
95    }
96
97    fn handle_err(&mut self, err: Error) {
98        debug_assert!(self.retries > 0);
99        self.retries -= 1;
100        self.log(err, true);
101    }
102
103    fn log(&self, err: Error, warn: bool) {
104        let msg = format!(
105            "{err}{delay} ({retries} tries remaining)",
106            delay = if self.delay.is_zero() {
107                String::new()
108            } else {
109                format!("; waiting {} seconds before trying again", self.delay.as_secs())
110            },
111            retries = self.retries,
112        );
113        if warn {
114            let _ = sh_warn!("{msg}");
115        } else {
116            tracing::info!("{msg}");
117        }
118    }
119}
foundry_common/retry.rs

foundry_common/
retry.rs
