foundry_common/
retry.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
//! Retry utilities.

use eyre::{Error, Report, Result};
use std::{future::Future, time::Duration};

/// Error type for Retry.
#[derive(Debug, thiserror::Error)]
pub enum RetryError<E = Report> {
    /// Keeps retrying operation.
    Retry(E),
    /// Stops retrying operation immediately.
    Break(E),
}

/// A type that keeps track of attempts.
#[derive(Clone, Debug)]
pub struct Retry {
    retries: u32,
    delay: Option<Duration>,
}

impl Retry {
    /// Creates a new `Retry` instance.
    pub fn new(retries: u32, delay: Option<Duration>) -> Self {
        Self { retries, delay }
    }

    /// Runs the given closure in a loop, retrying if it fails up to the specified number of times.
    pub fn run<F: FnMut() -> Result<T>, T>(mut self, mut callback: F) -> Result<T> {
        loop {
            match callback() {
                Err(e) if self.retries > 0 => {
                    self.handle_err(e);
                    if let Some(delay) = self.delay {
                        std::thread::sleep(delay);
                    }
                }
                res => return res,
            }
        }
    }

    /// Runs the given async closure in a loop, retrying if it fails up to the specified number of
    /// times.
    pub async fn run_async<F, Fut, T>(mut self, mut callback: F) -> Result<T>
    where
        F: FnMut() -> Fut,
        Fut: Future<Output = Result<T>>,
    {
        loop {
            match callback().await {
                Err(e) if self.retries > 0 => {
                    self.handle_err(e);
                    if let Some(delay) = self.delay {
                        tokio::time::sleep(delay).await;
                    }
                }
                res => return res,
            };
        }
    }

    /// Runs the given async closure in a loop, retrying if it fails up to the specified number of
    /// times or immediately returning an error if the closure returned [`RetryError::Break`].
    pub async fn run_async_until_break<F, Fut, T>(mut self, mut callback: F) -> Result<T>
    where
        F: FnMut() -> Fut,
        Fut: Future<Output = Result<T, RetryError>>,
    {
        loop {
            match callback().await {
                Err(RetryError::Retry(e)) if self.retries > 0 => {
                    self.handle_err(e);
                    if let Some(delay) = self.delay {
                        tokio::time::sleep(delay).await;
                    }
                }
                Err(RetryError::Retry(e) | RetryError::Break(e)) => return Err(e),
                Ok(t) => return Ok(t),
            };
        }
    }

    fn handle_err(&mut self, err: Error) {
        self.retries -= 1;
        let _ = sh_warn!("{} ({} tries remaining)", err.root_cause(), self.retries);
    }
}