wayback_rs/util/
retries.rs

1//! Tools for attaching retry logic to error types.
2use core::pin::Pin;
3use futures::{
4    task::{Context, Poll},
5    Future,
6};
7use log::{log, Level};
8use std::fmt::Debug;
9use std::marker::PhantomData;
10use std::time::Duration;
11use tryhard::{
12    backoff_strategies::BackoffStrategy, OnRetry, RetryFuture, RetryFutureConfig, RetryPolicy,
13};
14
15/// Execute a future with retries where the error type is `Retryable`.
16pub fn retry_future<F, Fut, T, E>(f: F) -> RetryFuture<F, Fut, ErrorBackoff<E>, LogOnRetry>
17where
18    F: FnMut() -> Fut,
19    Fut: Future<Output = Result<T, E>>,
20    E: Retryable,
21{
22    tryhard::retry_fn(f).with_config(E::retry_config())
23}
24
25pub struct LogFuture {
26    level: Option<Level>,
27    message: Option<String>,
28}
29
30impl Future for LogFuture {
31    type Output = ();
32    fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
33        if let Some(level) = self.level {
34            log!(
35                level,
36                "{}",
37                self.message
38                    .take()
39                    .expect("LogFuture polled after completion")
40            );
41        }
42
43        Poll::Ready(())
44    }
45}
46
47pub struct LogOnRetry {
48    level: Option<Level>,
49}
50
51impl<E: Debug> OnRetry<E> for LogOnRetry {
52    type Future = LogFuture;
53
54    fn on_retry(
55        &mut self,
56        attempts: u32,
57        next_delay: Option<Duration>,
58        previous_error: &E,
59    ) -> Self::Future {
60        match next_delay {
61            Some(delay) => {
62                let message = if self.level.is_none() {
63                    None
64                } else {
65                    Some(format!(
66                        "Retry {}; waiting {:?} after error: {:?}",
67                        attempts, delay, previous_error
68                    ))
69                };
70                LogFuture {
71                    level: self.level,
72                    message,
73                }
74            }
75            None => LogFuture {
76                level: None,
77                message: None,
78            },
79        }
80    }
81}
82
83pub struct ErrorBackoff<E>
84where
85    E: ?Sized,
86{
87    delay: Duration,
88    _error: PhantomData<E>,
89}
90
91impl<'a, E: Retryable> BackoffStrategy<'a, E> for ErrorBackoff<E> {
92    type Output = RetryPolicy;
93
94    fn delay(&mut self, _attempt: u32, error: &'a E) -> RetryPolicy {
95        error.custom_retry_policy().unwrap_or_else(|| {
96            let prev_delay = self.delay;
97            self.delay *= 2;
98            RetryPolicy::Delay(prev_delay)
99        })
100    }
101}
102
103/// The `Retryable` trait allows an error type to define retry logic for
104/// specific errors.
105pub trait Retryable {
106    /// Return the maximum number of retries.
107    fn max_retries() -> u32;
108
109    /// Return the default initial delay.
110    fn default_initial_delay() -> Duration;
111
112    /// Return the log level for this error type (an empty value indicates that
113    /// no logging will be done).
114    fn log_level() -> Option<Level>;
115
116    /// Return a retry policy for the given error value.
117    ///
118    /// An empty value represents the default.
119    fn custom_retry_policy(&self) -> Option<RetryPolicy>;
120
121    /// Generate a new backoff strategy instance.
122    fn new_backoff() -> ErrorBackoff<Self> {
123        ErrorBackoff {
124            delay: Self::default_initial_delay(),
125            _error: PhantomData,
126        }
127    }
128
129    /// Generate a new retry configuration instance.
130    fn retry_config() -> RetryFutureConfig<ErrorBackoff<Self>, LogOnRetry> {
131        RetryFutureConfig::new(Self::max_retries())
132            .on_retry(LogOnRetry {
133                level: Self::log_level(),
134            })
135            .custom_backoff(Self::new_backoff())
136    }
137}
wayback_rs/util/retries.rs

wayback_rs/util/
retries.rs
