timvw_rust_runtime/
retry.rs

1//! A retry policy for HTTP APIs, for tower's `RetryLayer`.
2//!
3//! `Client::default()` already retries with `RetryPolicy::default()`. For another policy, start from the bare service:
4//!
5//! ```no_run
6//! # #[cfg(feature = "reqwest")]
7//! # fn demo() {
8//! use timvw_rust_runtime::{retry::RetryPolicy, Client, ReqwestService};
9//!
10//! let client = Client::from_service(ReqwestService::default())
11//!     .layer(tower::retry::RetryLayer::new(RetryPolicy::new(5)));
12//! # }
13//! ```
14
15use std::future::Future;
16use std::pin::Pin;
17use std::time::Duration;
18
19use http::header::RETRY_AFTER;
20use http::StatusCode;
21use tower::retry::Policy;
22
23use crate::{Body, BoxError};
24
25/// Retries `429 Too Many Requests`, `408 Request Timeout` and `5xx` responses, and failures to connect.
26///
27/// It waits as long as a `Retry-After` header (in seconds) asks, or else backs off exponentially from
28/// [`RetryPolicy::base_delay`] up to [`RetryPolicy::max_delay`]. A request whose body is a stream (a streamed file
29/// upload) cannot be sent again and is not retried.
30#[derive(Clone, Debug)]
31pub struct RetryPolicy {
32    max_retries: usize,
33    attempts: usize,
34    base_delay: Duration,
35    max_delay: Duration,
36}
37
38impl RetryPolicy {
39    /// A policy that retries at most `max_retries` times (after the first attempt).
40    pub fn new(max_retries: usize) -> Self {
41        Self {
42            max_retries,
43            attempts: 0,
44            base_delay: Duration::from_millis(200),
45            max_delay: Duration::from_secs(8),
46        }
47    }
48
49    /// The first backoff delay, doubled on each retry (default 200 ms).
50    pub fn base_delay(mut self, delay: Duration) -> Self {
51        self.base_delay = delay;
52        self
53    }
54
55    /// The longest backoff delay, also the cap on `Retry-After` (default 8 s).
56    pub fn max_delay(mut self, delay: Duration) -> Self {
57        self.max_delay = delay;
58        self
59    }
60
61    /// The number of retries allowed.
62    pub fn max_retries(&self) -> usize {
63        self.max_retries
64    }
65
66    fn delay(&self, retry_after: Option<Duration>) -> Duration {
67        let backoff = self
68            .base_delay
69            .saturating_mul(1u32 << self.attempts.min(16));
70        retry_after.unwrap_or(backoff).min(self.max_delay)
71    }
72}
73
74/// Three retries.
75impl Default for RetryPolicy {
76    fn default() -> Self {
77        Self::new(3)
78    }
79}
80
81/// Whether a response or error is worth retrying, and how long the server asked to wait.
82fn retryable<B>(result: &Result<http::Response<B>, BoxError>) -> Option<Option<Duration>> {
83    match result {
84        Ok(response) => {
85            let status = response.status();
86            let transient = status == StatusCode::TOO_MANY_REQUESTS
87                || status == StatusCode::REQUEST_TIMEOUT
88                || status.is_server_error();
89            transient.then(|| {
90                response
91                    .headers()
92                    .get(RETRY_AFTER)
93                    .and_then(|value| value.to_str().ok())
94                    .and_then(|value| value.trim().parse::<u64>().ok())
95                    .map(Duration::from_secs)
96            })
97        }
98        Err(error) => is_connect(error.as_ref()).then_some(None),
99    }
100}
101
102#[allow(unused_variables)]
103fn is_connect(error: &(dyn std::error::Error + 'static)) -> bool {
104    #[cfg(feature = "reqwest")]
105    {
106        let mut source = Some(error);
107        while let Some(error) = source {
108            if error
109                .downcast_ref::<reqwest::Error>()
110                .is_some_and(reqwest::Error::is_connect)
111            {
112                return true;
113            }
114            source = error.source();
115        }
116    }
117    false
118}
119
120impl<B> Policy<http::Request<Body>, http::Response<B>, BoxError> for RetryPolicy {
121    type Future = Pin<Box<dyn Future<Output = ()> + Send>>;
122
123    fn retry(
124        &mut self,
125        _request: &mut http::Request<Body>,
126        result: &mut Result<http::Response<B>, BoxError>,
127    ) -> Option<Self::Future> {
128        if self.attempts >= self.max_retries {
129            return None;
130        }
131        let retry_after = retryable(result)?;
132        let delay = self.delay(retry_after);
133        self.attempts += 1;
134        #[cfg(feature = "tracing")]
135        tracing::debug!(
136            target: "timvw_rust_runtime::trace",
137            attempt = self.attempts,
138            max_retries = self.max_retries,
139            delay_ms = delay.as_millis() as u64,
140            status = result.as_ref().ok().map(|response| response.status().as_u16()),
141            "retrying"
142        );
143        Some(Box::pin(tokio::time::sleep(delay)))
144    }
145
146    fn clone_request(&mut self, request: &http::Request<Body>) -> Option<http::Request<Body>> {
147        let body = request.body().try_clone()?;
148        let mut clone = http::Request::new(body);
149        *clone.method_mut() = request.method().clone();
150        *clone.uri_mut() = request.uri().clone();
151        *clone.version_mut() = request.version();
152        *clone.headers_mut() = request.headers().clone();
153        Some(clone)
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160
161    fn response(status: u16, retry_after: Option<&str>) -> Result<http::Response<()>, BoxError> {
162        let mut builder = http::Response::builder().status(status);
163        if let Some(value) = retry_after {
164            builder = builder.header(RETRY_AFTER, value);
165        }
166        Ok(builder.body(()).unwrap())
167    }
168
169    #[test]
170    fn retries_transient_statuses_only() {
171        assert_eq!(retryable(&response(429, None)), Some(None));
172        assert_eq!(
173            retryable(&response(503, Some("2"))),
174            Some(Some(Duration::from_secs(2)))
175        );
176        assert_eq!(retryable(&response(408, None)), Some(None));
177        assert_eq!(retryable(&response(400, None)), None);
178        assert_eq!(retryable(&response(200, None)), None);
179        assert_eq!(retryable::<()>(&Err("boom".into())), None);
180    }
181
182    #[test]
183    fn backs_off_and_caps() {
184        let mut policy = RetryPolicy::new(10)
185            .base_delay(Duration::from_millis(100))
186            .max_delay(Duration::from_secs(1));
187        assert_eq!(policy.delay(None), Duration::from_millis(100));
188        policy.attempts = 2;
189        assert_eq!(policy.delay(None), Duration::from_millis(400));
190        policy.attempts = 8;
191        assert_eq!(policy.delay(None), Duration::from_secs(1));
192        assert_eq!(
193            policy.delay(Some(Duration::from_secs(30))),
194            Duration::from_secs(1)
195        );
196    }
197
198    #[test]
199    fn only_bodies_in_memory_are_cloned() {
200        let mut policy = RetryPolicy::default();
201        let mut clone_request = |request: &http::Request<Body>| {
202            Policy::<_, http::Response<()>, BoxError>::clone_request(&mut policy, request)
203        };
204        let request = http::Request::post("https://example.com")
205            .header("x", "1")
206            .body(Body::from("a"))
207            .unwrap();
208        let clone = clone_request(&request).unwrap();
209        assert_eq!(clone.uri(), request.uri());
210        assert_eq!(clone.headers()["x"], "1");
211        let chunks =
212            futures_util::stream::iter(vec![Ok::<_, std::io::Error>(bytes::Bytes::from("a"))]);
213        let streamed = http::Request::post("https://example.com")
214            .body(Body::wrap_stream(chunks))
215            .unwrap();
216        assert!(clone_request(&streamed).is_none());
217    }
218}
timvw_rust_runtime/retry.rs

timvw_rust_runtime/
retry.rs
