xmtp_common/
retry.rs

1//! A retry strategy that works with rusts native [`std::error::Error`] type.
2//!
3//! TODO: Could make the impl of `RetryableError` trait into a proc-macro to auto-derive Retryable
4//! on annotated enum variants.
5//! ```ignore
6//! #[derive(Debug, Error)]
7//! enum ErrorFoo {
8//!     #[error("I am retryable")]
9//!     #[retryable]
10//!     Retryable,
11//!     #[error("Nested errors are retryable")]
12//!     #[retryable(inherit)]
13//!     NestedRetryable(AnotherErrorWithRetryableVariants),
14//!     #[error("Always fail")]
15//!     NotRetryable
16//! }
17//! ```
18
19use crate::time::Duration;
20use crate::{MaybeSend, MaybeSync};
21use rand::RngExt;
22use std::error::Error;
23use std::sync::Arc;
24
25// Rust 1.86 added Trait upcasting, so we can add these infallible conversions
26// which is useful when getting error messages
27impl From<Box<dyn RetryableError>> for Box<dyn Error> {
28    fn from(retryable: Box<dyn RetryableError>) -> Box<dyn Error> {
29        retryable
30    }
31}
32
33// NOTE: From<> implementation is not possible here b/c of rust orphan rules (relaxed for Box
34// types)
35/// Convert an `Arc<[RetryableError]>` to a Standard Library `Arc<Error>`
36pub fn arc_retryable_to_error(retryable: Arc<dyn RetryableError>) -> Arc<dyn Error> {
37    retryable
38}
39
40pub type BoxedRetry = Retry<Box<dyn Strategy>>;
41
42pub struct NotSpecialized;
43
44/// Specifies which errors are retryable.
45/// All Errors are not retryable by-default.
46pub trait RetryableError<SP = NotSpecialized>: std::error::Error + MaybeSend + MaybeSync {
47    fn is_retryable(&self) -> bool;
48}
49
50impl<T> RetryableError for &'_ T
51where
52    T: RetryableError,
53{
54    fn is_retryable(&self) -> bool {
55        (**self).is_retryable()
56    }
57}
58
59impl<E: RetryableError> RetryableError for Box<E> {
60    fn is_retryable(&self) -> bool {
61        (**self).is_retryable()
62    }
63}
64
65impl RetryableError for core::convert::Infallible {
66    fn is_retryable(&self) -> bool {
67        unreachable!()
68    }
69}
70
71/// Options to specify how to retry a function
72#[derive(Debug, Clone)]
73pub struct Retry<S = ExponentialBackoff> {
74    retries: usize,
75    strategy: S,
76}
77
78impl Default for Retry {
79    fn default() -> Retry {
80        Retry {
81            retries: 5,
82            strategy: ExponentialBackoff::default(),
83        }
84    }
85}
86
87impl<S: Strategy> Retry<S> {
88    /// Get the number of retries this is configured with.
89    pub fn retries(&self) -> usize {
90        self.retries
91    }
92
93    pub fn backoff(&self, attempts: usize, time_spent: crate::time::Instant) -> Option<Duration> {
94        self.strategy.backoff(attempts, time_spent)
95    }
96}
97
98impl<S: Strategy + 'static> Retry<S> {
99    pub fn boxed(self) -> Retry<Box<dyn Strategy>> {
100        Retry {
101            strategy: Box::new(self.strategy),
102            retries: self.retries,
103        }
104    }
105}
106
107/// The strategy interface
108pub trait Strategy: MaybeSend + MaybeSync {
109    /// A time that this retry should backoff
110    /// Returns None when we should no longer backoff,
111    /// despite possibly being below attempts
112    fn backoff(&self, attempts: usize, time_spent: crate::time::Instant) -> Option<Duration>;
113}
114
115impl Strategy for () {
116    fn backoff(&self, _attempts: usize, _time_spent: crate::time::Instant) -> Option<Duration> {
117        Some(Duration::ZERO)
118    }
119}
120
121impl<S: ?Sized + Strategy> Strategy for Box<S> {
122    fn backoff(&self, attempts: usize, time_spent: crate::time::Instant) -> Option<Duration> {
123        (**self).backoff(attempts, time_spent)
124    }
125}
126
127#[derive(Clone, Debug)]
128pub struct ExponentialBackoff {
129    /// The amount to multiply the duration on each subsequent attempt
130    multiplier: u32,
131    /// Duration to be multiplied
132    duration: Duration,
133    /// jitter to add randomness
134    max_jitter: Duration,
135    /// upper limit on time to wait for all retries
136    total_wait_max: Duration,
137    /// upper limit on time to wait between retries
138    individual_wait_max: Duration,
139}
140
141impl ExponentialBackoff {
142    pub fn builder() -> ExponentialBackoffBuilder {
143        ExponentialBackoffBuilder::default()
144    }
145}
146
147impl Default for ExponentialBackoff {
148    fn default() -> Self {
149        Self {
150            // total wait time == two minutes
151            multiplier: 3,
152            duration: Duration::from_millis(50),
153            total_wait_max: Duration::from_secs(120),
154            individual_wait_max: Duration::from_secs(30),
155            max_jitter: Duration::from_millis(25),
156        }
157    }
158}
159
160#[derive(Default)]
161pub struct ExponentialBackoffBuilder {
162    duration: Option<Duration>,
163    max_jitter: Option<Duration>,
164    multiplier: Option<u32>,
165    total_wait_max: Option<Duration>,
166}
167
168impl ExponentialBackoffBuilder {
169    pub fn duration(mut self, duration: Duration) -> Self {
170        self.duration = Some(duration);
171        self
172    }
173
174    pub fn max_jitter(mut self, max_jitter: Duration) -> Self {
175        self.max_jitter = Some(max_jitter);
176        self
177    }
178
179    pub fn multiplier(mut self, multiplier: u32) -> Self {
180        self.multiplier = Some(multiplier);
181        self
182    }
183
184    pub fn total_wait_max(mut self, total_wait_max: Duration) -> Self {
185        self.total_wait_max = Some(total_wait_max);
186        self
187    }
188
189    pub fn build(self) -> ExponentialBackoff {
190        ExponentialBackoff {
191            duration: self.duration.unwrap_or(Duration::from_millis(25)),
192            max_jitter: self.max_jitter.unwrap_or(Duration::from_millis(25)),
193            multiplier: self.multiplier.unwrap_or(3),
194            total_wait_max: self.total_wait_max.unwrap_or(Duration::from_secs(120)),
195            individual_wait_max: Duration::from_secs(30),
196        }
197    }
198}
199
200impl Strategy for ExponentialBackoff {
201    fn backoff(&self, attempts: usize, time_spent: crate::time::Instant) -> Option<Duration> {
202        if time_spent.elapsed() > self.total_wait_max {
203            return None;
204        }
205        let mut duration = self.duration;
206        for _ in 0..(attempts.saturating_sub(1)) {
207            duration *= self.multiplier;
208            if duration > self.individual_wait_max {
209                duration = self.individual_wait_max;
210            }
211        }
212        let distr = rand::distr::Uniform::new_inclusive(Duration::ZERO, self.max_jitter).unwrap();
213        let jitter = rand::rng().sample(distr);
214        let wait = duration + jitter;
215        Some(wait)
216    }
217}
218
219/// Builder for [`Retry`]
220#[derive(Default, Debug, Copy, Clone)]
221pub struct RetryBuilder<S> {
222    retries: Option<usize>,
223    strategy: S,
224}
225
226impl RetryBuilder<ExponentialBackoff> {
227    pub fn new() -> Self {
228        Self {
229            retries: Some(5),
230            strategy: ExponentialBackoff::default(),
231        }
232    }
233}
234
235/// Builder for [`Retry`].
236///
237/// # Example
238/// ```ignore
239/// use xmtp_common::retry::RetryBuilder;
240///
241/// RetryBuilder::default()
242///     .retries(5)
243///     .with_strategy(xmtp_common::ExponentialBackoff::default())
244///     .build();
245/// ```
246impl<S: Strategy> RetryBuilder<S> {
247    pub fn build(self) -> Retry<S> {
248        let mut retry = Retry {
249            retries: 5usize,
250            strategy: self.strategy,
251        };
252
253        if let Some(retries) = self.retries {
254            retry.retries = retries;
255        }
256
257        retry
258    }
259
260    /// Specify the  of retries to allow
261    pub fn retries(mut self, retries: usize) -> Self {
262        self.retries = Some(retries);
263        self
264    }
265
266    pub fn with_strategy<St: Strategy>(self, strategy: St) -> RetryBuilder<St> {
267        RetryBuilder {
268            retries: self.retries,
269            strategy,
270        }
271    }
272}
273
274impl Retry {
275    /// Get the builder for [`Retry`]
276    pub fn builder() -> RetryBuilder<ExponentialBackoff> {
277        RetryBuilder::new()
278    }
279}
280
281/// Retry but for an async context
282/// ```
283/// use xmtp_common::{retry_async, retry::{RetryableError, Retry}};
284/// use thiserror::Error;
285/// use tokio::sync::mpsc;
286///
287/// #[derive(Debug, Error)]
288/// enum MyError {
289///     #[error("A retryable error")]
290///     Retryable,
291///     #[error("An error we don't want to retry")]
292///     NotRetryable
293/// }
294///
295/// impl RetryableError for MyError {
296///     fn is_retryable(&self) -> bool {
297///         match self {
298///             Self::Retryable => true,
299///             _=> false,
300///         }
301///     }
302/// }
303///
304/// async fn fallable_fn(rx: &mut mpsc::Receiver<usize>) -> Result<(), MyError> {
305///     if rx.recv().await.unwrap() == 2 {
306///         return Ok(());
307///     }
308///     Err(MyError::Retryable)
309/// }
310///
311/// #[tokio::main(flavor = "current_thread")]
312/// async fn main() -> Result<(), MyError> {
313///
314///     let (tx, mut rx) = mpsc::channel(3);
315///
316///     for i in 0..3 {
317///         tx.send(i).await.unwrap();
318///     }
319///     retry_async!(Retry::default(), (async {
320///         fallable_fn(&mut rx).await
321///     }))
322/// }
323/// ```
324#[macro_export]
325macro_rules! retry_async {
326    ($retry: expr, $code: tt) => {{
327        use tracing::Instrument as _;
328        #[allow(unused)]
329        use $crate::retry::RetryableError;
330        let mut attempts = 0;
331        let time_spent = $crate::time::Instant::now();
332        let span = tracing::trace_span!("retry");
333        loop {
334            let span = span.clone();
335            #[allow(clippy::redundant_closure_call)]
336            let res = $code.instrument(span).await;
337            match res {
338                Ok(v) => break Ok(v),
339                Err(e) => {
340                    if (&e).is_retryable() && attempts < $retry.retries() {
341                        // A single retry is normal transient behavior; keep it at debug so a
342                        // flapping endpoint doesn't flood prod warn logs (only exhaustion warns).
343                        tracing::debug!(
344                            attempt = attempts,
345                            "retrying function that failed with error={}",
346                            e.to_string()
347                        );
348                        if let Some(d) = $retry.backoff(attempts, time_spent) {
349                            attempts += 1;
350                            $crate::time::sleep(d).await;
351                        } else {
352                            tracing::warn!(
353                                attempts,
354                                elapsed_ms = time_spent.elapsed().as_millis(),
355                                "retry strategy exceeded max wait time, giving up: {}",
356                                e.to_string()
357                            );
358                            break Err(e);
359                        }
360                    } else {
361                        tracing::trace!("error is not retryable. {}", e);
362                        break Err(e);
363                    }
364                }
365            }
366        }
367    }};
368}
369
370#[macro_export]
371macro_rules! retryable {
372    ($error: ident) => {{
373        #[allow(unused)]
374        use $crate::retry::RetryableError;
375        $error.is_retryable()
376    }};
377    ($error: expr) => {{
378        use $crate::retry::RetryableError;
379        $error.is_retryable()
380    }};
381}
382
383#[cfg(test)]
384pub(crate) mod tests {
385    use super::*;
386
387    use thiserror::Error;
388    use tokio::sync::mpsc;
389
390    #[derive(Debug, Error)]
391    enum SomeError {
392        #[error("this is a retryable error")]
393        ARetryableError,
394        #[error("Dont retry")]
395        DontRetryThis,
396    }
397
398    impl RetryableError for SomeError {
399        fn is_retryable(&self) -> bool {
400            matches!(self, Self::ARetryableError)
401        }
402    }
403
404    fn retry_error_fn() -> Result<(), SomeError> {
405        Err(SomeError::ARetryableError)
406    }
407
408    fn retryable_with_args(foo: usize, name: String, list: &Vec<String>) -> Result<(), SomeError> {
409        println!("I am {foo} of {name} with items {list:?}");
410        Err(SomeError::ARetryableError)
411    }
412
413    #[xmtp_macro::test]
414    async fn it_retries_twice_and_succeeds() {
415        let mut i = 0;
416        let mut test_fn = || -> Result<(), SomeError> {
417            if i == 2 {
418                return Ok(());
419            }
420            i += 1;
421            retry_error_fn()?;
422            Ok(())
423        };
424
425        retry_async!(Retry::default(), (async { test_fn() })).unwrap();
426    }
427
428    #[xmtp_macro::test]
429    async fn it_works_with_random_args() {
430        let mut i = 0;
431        let list = vec!["String".into(), "Foo".into()];
432        let mut test_fn = || -> Result<(), SomeError> {
433            if i == 2 {
434                return Ok(());
435            }
436            i += 1;
437            retryable_with_args(i, "Hello".to_string(), &list)
438        };
439
440        retry_async!(Retry::default(), (async { test_fn() })).unwrap();
441    }
442
443    #[xmtp_macro::test]
444    async fn it_fails_on_three_retries() {
445        let closure = || -> Result<(), SomeError> {
446            retry_error_fn()?;
447            Ok(())
448        };
449        let result: Result<(), SomeError> = retry_async!(Retry::default(), (async { closure() }));
450
451        assert!(result.is_err())
452    }
453
454    #[xmtp_macro::test]
455    async fn it_only_runs_non_retryable_once() {
456        let mut attempts = 0;
457        let mut test_fn = || -> Result<(), SomeError> {
458            attempts += 1;
459            Err(SomeError::DontRetryThis)
460        };
461
462        let _r = retry_async!(Retry::default(), (async { test_fn() }));
463
464        assert_eq!(attempts, 1);
465    }
466
467    #[xmtp_macro::test]
468    async fn it_works_async() {
469        async fn retryable_async_fn(rx: &mut mpsc::Receiver<usize>) -> Result<(), SomeError> {
470            let val = rx.recv().await.unwrap();
471            if val == 2 {
472                return Ok(());
473            }
474            // do some work
475            crate::time::sleep(core::time::Duration::from_nanos(100)).await;
476            Err(SomeError::ARetryableError)
477        }
478
479        let (tx, mut rx) = mpsc::channel(3);
480
481        for i in 0..3 {
482            tx.send(i).await.unwrap();
483        }
484        retry_async!(
485            Retry::default(),
486            (async { retryable_async_fn(&mut rx).await })
487        )
488        .unwrap();
489        assert!(rx.is_empty());
490    }
491
492    #[xmtp_macro::test]
493    async fn it_works_async_mut() {
494        async fn retryable_async_fn(data: &mut usize) -> Result<(), SomeError> {
495            if *data == 2 {
496                return Ok(());
497            }
498            *data += 1;
499            // do some work
500            crate::time::sleep(core::time::Duration::from_nanos(100)).await;
501            Err(SomeError::ARetryableError)
502        }
503
504        let mut data: usize = 0;
505        retry_async!(
506            Retry::default(),
507            (async { retryable_async_fn(&mut data).await })
508        )
509        .unwrap();
510    }
511
512    #[xmtp_macro::test]
513    fn backoff_retry() {
514        let backoff_retry = Retry::default();
515        let time_spent = crate::time::Instant::now();
516        assert!(backoff_retry.backoff(1, time_spent).unwrap().as_millis() - 50 <= 25);
517        assert!(backoff_retry.backoff(2, time_spent).unwrap().as_millis() - 150 <= 25);
518        assert!(backoff_retry.backoff(3, time_spent).unwrap().as_millis() - 450 <= 25);
519    }
520}
xmtp_common/retry.rs

xmtp_common/
retry.rs
