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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
// This file is adapted from reqwest-retry, which was a bit too opinionated for our use.
// https://github.com/TrueLayer/reqwest-middleware
use async_trait::async_trait;
use http::Extensions;
use rand::{rng, RngExt};
use reqwest::{Request, Response, StatusCode};
use reqwest_middleware::{Middleware, Next, Result};
use std::time::Duration;
/// Middleware for retrying requests.
pub struct CustomRetryMiddleware {
max_retries: u32,
max_delay_ms: u64,
initial_delay_ms: u64,
}
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl Middleware for CustomRetryMiddleware {
async fn handle(
&self,
req: Request,
extensions: &mut Extensions,
next: Next<'_>,
) -> Result<Response> {
self.execute_with_retry(req, next, extensions).await
}
}
impl CustomRetryMiddleware {
/// Create a new retry middleware instance.
pub fn new(max_retries: u32, max_delay_ms: u64, initial_delay_ms: u64) -> Self {
Self {
max_retries: max_retries.min(10),
max_delay_ms,
initial_delay_ms,
}
}
async fn execute_with_retry<'a>(
&'a self,
req: Request,
next: Next<'a>,
ext: &'a mut Extensions,
) -> Result<Response> {
let mut n_past_retries = 0;
let mut last_req_401 = false;
loop {
let duplicate_request = match req.try_clone() {
Some(x) => x,
None => return next.run(req, ext).await,
};
let result = next.clone().run(duplicate_request, ext).await;
// Check if the error can be retried.
break match Retryable::from_reqwest_response(&result) {
Some(retryable)
if (retryable == Retryable::Transient
|| retryable == Retryable::Unauthorized && !last_req_401)
&& n_past_retries < self.max_retries =>
{
last_req_401 = retryable == Retryable::Unauthorized;
// If the response failed and the error type was transient
// we can safely try to retry the request.
let mut retry_delay = self.initial_delay_ms * 2u64.pow(n_past_retries);
if retry_delay > self.max_delay_ms {
retry_delay = self.max_delay_ms;
}
// Jitter so we land between initial * 2 ** attempt * 3/4 and initial * 2 ** attempt * 5/4
retry_delay = retry_delay / 4 * 3 + rng().random_range(0..=(retry_delay / 2));
futures_timer::Delay::new(Duration::from_millis(retry_delay)).await;
n_past_retries += 1;
continue;
}
Some(_) | None => result,
};
}
}
}
#[derive(PartialEq, Eq)]
pub(crate) enum Retryable {
/// The failure was due to something that might resolve in the future.
Transient,
/// Unresolvable error.
Fatal,
/// Unauthorized. This is _maybe_ resolvable, if the last request wasn't also a 401.
Unauthorized,
}
impl Retryable {
/// Try to map a `reqwest` response into `Retryable`.
///
/// Returns `None` if the response object does not contain any errors.
///
/// # Arguments
///
/// * `res` - Request response.
pub fn from_reqwest_response(
res: &reqwest_middleware::Result<reqwest::Response>,
) -> Option<Self> {
match res {
Ok(success) => {
let status = success.status();
if status.is_success() {
None
} else if status == StatusCode::UNAUTHORIZED {
Some(Retryable::Unauthorized)
} else if status.is_server_error()
|| status == StatusCode::REQUEST_TIMEOUT
|| status == StatusCode::TOO_MANY_REQUESTS
|| success
.headers()
.get("cdf-is-auto-retryable")
.and_then(|v| v.to_str().ok())
.is_some_and(|v| v == "true")
{
Some(Retryable::Transient)
} else {
Some(Retryable::Fatal)
}
}
Err(error) => match error {
reqwest_middleware::Error::Middleware(_) => Some(Retryable::Fatal),
reqwest_middleware::Error::Reqwest(error) => {
#[cfg(not(target_arch = "wasm32"))]
let is_connect = error.is_connect();
#[cfg(target_arch = "wasm32")]
let is_connect = false;
if error.is_timeout() || is_connect {
Some(Retryable::Transient)
} else if error.is_body()
|| error.is_decode()
|| error.is_builder()
|| error.is_redirect()
|| error.is_request()
{
Some(Retryable::Fatal)
} else {
// We omit checking if error.is_status() since we check that already.
// However, if Response::error_for_status is used the status will still
// remain in the response object.
None
}
}
},
}
}
}