1use crate::error::FaucetError;
5
6#[derive(
9 Debug,
10 Clone,
11 Copy,
12 PartialEq,
13 Eq,
14 Hash,
15 serde::Serialize,
16 serde::Deserialize,
17 schemars::JsonSchema,
18)]
19#[serde(rename_all = "snake_case")]
20pub enum RetryClass {
21 #[serde(rename = "http_5xx")]
23 Http5xx,
24 RateLimited,
26 Connection,
28 Timeout,
30}
31
32impl RetryClass {
33 pub fn as_str(self) -> &'static str {
35 match self {
36 RetryClass::Http5xx => "http_5xx",
37 RetryClass::RateLimited => "rate_limited",
38 RetryClass::Connection => "connection",
39 RetryClass::Timeout => "timeout",
40 }
41 }
42}
43
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub struct RetryClassSet([bool; 4]);
48
49impl RetryClassSet {
50 fn idx(c: RetryClass) -> usize {
51 match c {
52 RetryClass::Http5xx => 0,
53 RetryClass::RateLimited => 1,
54 RetryClass::Connection => 2,
55 RetryClass::Timeout => 3,
56 }
57 }
58
59 pub fn contains(&self, c: RetryClass) -> bool {
61 self.0[Self::idx(c)]
62 }
63}
64
65impl Default for RetryClassSet {
66 fn default() -> Self {
68 RetryClassSet([true; 4])
69 }
70}
71
72impl FromIterator<RetryClass> for RetryClassSet {
73 fn from_iter<I: IntoIterator<Item = RetryClass>>(iter: I) -> Self {
74 let mut set = RetryClassSet([false; 4]);
75 for c in iter {
76 set.0[Self::idx(c)] = true;
77 }
78 set
79 }
80}
81
82pub fn classify(err: &FaucetError) -> Option<RetryClass> {
86 match err {
87 FaucetError::HttpStatus { status, .. } if *status == 429 => Some(RetryClass::RateLimited),
88 FaucetError::HttpStatus { status, .. } if *status >= 500 => Some(RetryClass::Http5xx),
89 FaucetError::RateLimited(_) => Some(RetryClass::RateLimited),
90 FaucetError::Timeout { .. } => Some(RetryClass::Timeout),
91 FaucetError::Http(e) => classify_transport(
92 e.is_timeout(),
93 e.status().map(|s| s.as_u16()),
94 e.is_connect(),
95 ),
96 _ => None,
97 }
98}
99
100fn classify_transport(
112 is_timeout: bool,
113 status: Option<u16>,
114 _is_connect: bool,
115) -> Option<RetryClass> {
116 if is_timeout {
117 Some(RetryClass::Timeout)
118 } else if let Some(status) = status {
119 if status >= 500 {
120 Some(RetryClass::Http5xx)
121 } else if status == 429 {
122 Some(RetryClass::RateLimited)
123 } else {
124 None
125 }
126 } else {
127 Some(RetryClass::Connection)
130 }
131}
132
133#[cfg(test)]
134mod tests {
135 use super::*;
136 use std::time::Duration;
137
138 #[test]
139 fn classifies_5xx_429_and_rate_limited() {
140 assert_eq!(
141 classify(&FaucetError::HttpStatus {
142 status: 503,
143 url: "u".into(),
144 body: "".into()
145 }),
146 Some(RetryClass::Http5xx)
147 );
148 assert_eq!(
149 classify(&FaucetError::HttpStatus {
150 status: 429,
151 url: "u".into(),
152 body: "".into()
153 }),
154 Some(RetryClass::RateLimited)
155 );
156 assert_eq!(
157 classify(&FaucetError::RateLimited(Duration::from_secs(1))),
158 Some(RetryClass::RateLimited)
159 );
160 }
161
162 #[test]
163 fn a_typed_timeout_is_the_timeout_class() {
164 assert_eq!(
165 classify(&FaucetError::Timeout {
166 operation: "rabbitmq publish".into(),
167 after: Duration::from_secs(60),
168 }),
169 Some(RetryClass::Timeout)
170 );
171 }
172
173 #[test]
174 fn transport_classification_covers_every_arm() {
175 assert_eq!(
177 classify_transport(true, None, false),
178 Some(RetryClass::Timeout)
179 );
180 assert_eq!(
181 classify_transport(true, Some(500), true),
182 Some(RetryClass::Timeout)
183 );
184 assert_eq!(
186 classify_transport(false, Some(503), false),
187 Some(RetryClass::Http5xx)
188 );
189 assert_eq!(
190 classify_transport(false, Some(429), false),
191 Some(RetryClass::RateLimited)
192 );
193 assert_eq!(classify_transport(false, Some(404), false), None);
195 assert_eq!(
197 classify_transport(false, None, true),
198 Some(RetryClass::Connection)
199 );
200 assert_eq!(
201 classify_transport(false, None, false),
202 Some(RetryClass::Connection)
203 );
204 }
205
206 #[test]
207 fn non_retriable_errors_classify_none() {
208 assert_eq!(classify(&FaucetError::Auth("x".into())), None);
209 assert_eq!(classify(&FaucetError::Config("x".into())), None);
210 assert_eq!(classify(&FaucetError::Sink("x".into())), None);
211 }
212
213 #[test]
214 fn retry_class_serde_round_trips() {
215 use serde_json::json;
216 assert_eq!(
219 serde_json::from_value::<RetryClass>(json!("http_5xx")).unwrap(),
220 RetryClass::Http5xx
221 );
222 assert_eq!(
223 serde_json::to_value(RetryClass::Http5xx).unwrap(),
224 json!("http_5xx")
225 );
226 for (s, c) in [
228 ("rate_limited", RetryClass::RateLimited),
229 ("connection", RetryClass::Connection),
230 ("timeout", RetryClass::Timeout),
231 ] {
232 assert_eq!(serde_json::from_value::<RetryClass>(json!(s)).unwrap(), c);
233 assert_eq!(serde_json::to_value(c).unwrap(), json!(s));
234 }
235 }
236
237 #[test]
238 fn set_contains_and_default_covers_all_four() {
239 let set = RetryClassSet::default();
240 for c in [
241 RetryClass::Http5xx,
242 RetryClass::RateLimited,
243 RetryClass::Connection,
244 RetryClass::Timeout,
245 ] {
246 assert!(set.contains(c), "default set should contain {c:?}");
247 }
248 let only5xx = RetryClassSet::from_iter([RetryClass::Http5xx]);
249 assert!(only5xx.contains(RetryClass::Http5xx));
250 assert!(!only5xx.contains(RetryClass::RateLimited));
251 }
252}