Skip to main content

pubhubs/misc/
task.rs

1//! Tools for dealing with tokio tasks.
2
3use std::future::Future;
4use tokio::time::Duration;
5
6/// Options for [retry].
7pub struct RetryOptions {
8    /// Wait this amount of time after the first `Ok(None)` is returned.
9    pub initial_wait_time: Duration,
10
11    /// Increase the wait time by this factor each time `Ok(None)` is returned.
12    pub backoff_factor: f32,
13
14    /// Don't wait longer than this
15    pub max_wait_time: Duration,
16
17    /// Try `max_retries+1` number of times to get a non-`Ok(None)` answer.
18    pub max_retries: Option<usize>,
19}
20
21impl Default for RetryOptions {
22    fn default() -> Self {
23        Self {
24            initial_wait_time: Duration::from_millis(10),
25            backoff_factor: 2f32,
26            max_wait_time: Duration::from_millis(5_000),
27            max_retries: None, // Some(10), TODO: some?
28        }
29    }
30}
31
32impl RetryOptions {
33    pub async fn retry<T, E, Fut: Future<Output = Result<Option<T>, E>>>(
34        &self,
35        f: impl Fn() -> Fut,
36    ) -> Result<Option<T>, E> {
37        let mut retries_left: usize = self.max_retries.unwrap_or(usize::MAX);
38        let mut wait_time: Duration = self.initial_wait_time;
39        let backoff_factor = self.backoff_factor;
40
41        loop {
42            let res = f().await;
43
44            // return when res is Ok(Some(v)) or Err(ec)
45            match res.as_ref() {
46                Ok(None) => {}
47                _ => return res,
48            };
49
50            if retries_left == 0 {
51                return Ok(None);
52            }
53
54            retries_left -= 1;
55
56            log::trace!("retrying after {wait_time:.1?}");
57
58            tokio::time::sleep(wait_time).await;
59
60            wait_time = std::cmp::min(wait_time.mul_f32(backoff_factor), self.max_wait_time)
61        }
62    }
63}
64
65/// Calls the given function `f` until it no longer returns `Ok(None)`, and returns the last result
66/// which is thus either an `Ok(Some(value))` or an `Err(err)`.
67///
68/// When a maximal number of retries is reached (see [RetryOptions::max_retries]), `Ok(None)` is returned.
69///
70/// See [RetryOptions::retry] for more options.
71pub async fn retry<T, E, Fut: Future<Output = Result<Option<T>, E>>>(
72    f: impl Fn() -> Fut,
73) -> Result<Option<T>, E> {
74    RetryOptions::default().retry(f).await
75}