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::Http(e) => classify_transport(
91 e.is_timeout(),
92 e.status().map(|s| s.as_u16()),
93 e.is_connect(),
94 ),
95 _ => None,
96 }
97}
98
99fn classify_transport(
111 is_timeout: bool,
112 status: Option<u16>,
113 _is_connect: bool,
114) -> Option<RetryClass> {
115 if is_timeout {
116 Some(RetryClass::Timeout)
117 } else if let Some(status) = status {
118 if status >= 500 {
119 Some(RetryClass::Http5xx)
120 } else if status == 429 {
121 Some(RetryClass::RateLimited)
122 } else {
123 None
124 }
125 } else {
126 Some(RetryClass::Connection)
129 }
130}
131
132#[cfg(test)]
133mod tests {
134 use super::*;
135 use std::time::Duration;
136
137 #[test]
138 fn classifies_5xx_429_and_rate_limited() {
139 assert_eq!(
140 classify(&FaucetError::HttpStatus {
141 status: 503,
142 url: "u".into(),
143 body: "".into()
144 }),
145 Some(RetryClass::Http5xx)
146 );
147 assert_eq!(
148 classify(&FaucetError::HttpStatus {
149 status: 429,
150 url: "u".into(),
151 body: "".into()
152 }),
153 Some(RetryClass::RateLimited)
154 );
155 assert_eq!(
156 classify(&FaucetError::RateLimited(Duration::from_secs(1))),
157 Some(RetryClass::RateLimited)
158 );
159 }
160
161 #[test]
162 fn transport_classification_covers_every_arm() {
163 assert_eq!(
165 classify_transport(true, None, false),
166 Some(RetryClass::Timeout)
167 );
168 assert_eq!(
169 classify_transport(true, Some(500), true),
170 Some(RetryClass::Timeout)
171 );
172 assert_eq!(
174 classify_transport(false, Some(503), false),
175 Some(RetryClass::Http5xx)
176 );
177 assert_eq!(
178 classify_transport(false, Some(429), false),
179 Some(RetryClass::RateLimited)
180 );
181 assert_eq!(classify_transport(false, Some(404), false), None);
183 assert_eq!(
185 classify_transport(false, None, true),
186 Some(RetryClass::Connection)
187 );
188 assert_eq!(
189 classify_transport(false, None, false),
190 Some(RetryClass::Connection)
191 );
192 }
193
194 #[test]
195 fn non_retriable_errors_classify_none() {
196 assert_eq!(classify(&FaucetError::Auth("x".into())), None);
197 assert_eq!(classify(&FaucetError::Config("x".into())), None);
198 assert_eq!(classify(&FaucetError::Sink("x".into())), None);
199 }
200
201 #[test]
202 fn retry_class_serde_round_trips() {
203 use serde_json::json;
204 assert_eq!(
207 serde_json::from_value::<RetryClass>(json!("http_5xx")).unwrap(),
208 RetryClass::Http5xx
209 );
210 assert_eq!(
211 serde_json::to_value(RetryClass::Http5xx).unwrap(),
212 json!("http_5xx")
213 );
214 for (s, c) in [
216 ("rate_limited", RetryClass::RateLimited),
217 ("connection", RetryClass::Connection),
218 ("timeout", RetryClass::Timeout),
219 ] {
220 assert_eq!(serde_json::from_value::<RetryClass>(json!(s)).unwrap(), c);
221 assert_eq!(serde_json::to_value(c).unwrap(), json!(s));
222 }
223 }
224
225 #[test]
226 fn set_contains_and_default_covers_all_four() {
227 let set = RetryClassSet::default();
228 for c in [
229 RetryClass::Http5xx,
230 RetryClass::RateLimited,
231 RetryClass::Connection,
232 RetryClass::Timeout,
233 ] {
234 assert!(set.contains(c), "default set should contain {c:?}");
235 }
236 let only5xx = RetryClassSet::from_iter([RetryClass::Http5xx]);
237 assert!(only5xx.contains(RetryClass::Http5xx));
238 assert!(!only5xx.contains(RetryClass::RateLimited));
239 }
240}