use crate::error::FaucetError;
#[derive(
Debug,
Clone,
Copy,
PartialEq,
Eq,
Hash,
serde::Serialize,
serde::Deserialize,
schemars::JsonSchema,
)]
#[serde(rename_all = "snake_case")]
pub enum RetryClass {
#[serde(rename = "http_5xx")]
Http5xx,
RateLimited,
Connection,
Timeout,
}
impl RetryClass {
pub fn as_str(self) -> &'static str {
match self {
RetryClass::Http5xx => "http_5xx",
RetryClass::RateLimited => "rate_limited",
RetryClass::Connection => "connection",
RetryClass::Timeout => "timeout",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RetryClassSet([bool; 4]);
impl RetryClassSet {
fn idx(c: RetryClass) -> usize {
match c {
RetryClass::Http5xx => 0,
RetryClass::RateLimited => 1,
RetryClass::Connection => 2,
RetryClass::Timeout => 3,
}
}
pub fn contains(&self, c: RetryClass) -> bool {
self.0[Self::idx(c)]
}
}
impl Default for RetryClassSet {
fn default() -> Self {
RetryClassSet([true; 4])
}
}
impl FromIterator<RetryClass> for RetryClassSet {
fn from_iter<I: IntoIterator<Item = RetryClass>>(iter: I) -> Self {
let mut set = RetryClassSet([false; 4]);
for c in iter {
set.0[Self::idx(c)] = true;
}
set
}
}
pub fn classify(err: &FaucetError) -> Option<RetryClass> {
match err {
FaucetError::HttpStatus { status, .. } if *status == 429 => Some(RetryClass::RateLimited),
FaucetError::HttpStatus { status, .. } if *status >= 500 => Some(RetryClass::Http5xx),
FaucetError::RateLimited(_) => Some(RetryClass::RateLimited),
FaucetError::Http(e) => classify_transport(
e.is_timeout(),
e.status().map(|s| s.as_u16()),
e.is_connect(),
),
_ => None,
}
}
fn classify_transport(
is_timeout: bool,
status: Option<u16>,
_is_connect: bool,
) -> Option<RetryClass> {
if is_timeout {
Some(RetryClass::Timeout)
} else if let Some(status) = status {
if status >= 500 {
Some(RetryClass::Http5xx)
} else if status == 429 {
Some(RetryClass::RateLimited)
} else {
None
}
} else {
Some(RetryClass::Connection)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
#[test]
fn classifies_5xx_429_and_rate_limited() {
assert_eq!(
classify(&FaucetError::HttpStatus {
status: 503,
url: "u".into(),
body: "".into()
}),
Some(RetryClass::Http5xx)
);
assert_eq!(
classify(&FaucetError::HttpStatus {
status: 429,
url: "u".into(),
body: "".into()
}),
Some(RetryClass::RateLimited)
);
assert_eq!(
classify(&FaucetError::RateLimited(Duration::from_secs(1))),
Some(RetryClass::RateLimited)
);
}
#[test]
fn transport_classification_covers_every_arm() {
assert_eq!(
classify_transport(true, None, false),
Some(RetryClass::Timeout)
);
assert_eq!(
classify_transport(true, Some(500), true),
Some(RetryClass::Timeout)
);
assert_eq!(
classify_transport(false, Some(503), false),
Some(RetryClass::Http5xx)
);
assert_eq!(
classify_transport(false, Some(429), false),
Some(RetryClass::RateLimited)
);
assert_eq!(classify_transport(false, Some(404), false), None);
assert_eq!(
classify_transport(false, None, true),
Some(RetryClass::Connection)
);
assert_eq!(
classify_transport(false, None, false),
Some(RetryClass::Connection)
);
}
#[test]
fn non_retriable_errors_classify_none() {
assert_eq!(classify(&FaucetError::Auth("x".into())), None);
assert_eq!(classify(&FaucetError::Config("x".into())), None);
assert_eq!(classify(&FaucetError::Sink("x".into())), None);
}
#[test]
fn retry_class_serde_round_trips() {
use serde_json::json;
assert_eq!(
serde_json::from_value::<RetryClass>(json!("http_5xx")).unwrap(),
RetryClass::Http5xx
);
assert_eq!(
serde_json::to_value(RetryClass::Http5xx).unwrap(),
json!("http_5xx")
);
for (s, c) in [
("rate_limited", RetryClass::RateLimited),
("connection", RetryClass::Connection),
("timeout", RetryClass::Timeout),
] {
assert_eq!(serde_json::from_value::<RetryClass>(json!(s)).unwrap(), c);
assert_eq!(serde_json::to_value(c).unwrap(), json!(s));
}
}
#[test]
fn set_contains_and_default_covers_all_four() {
let set = RetryClassSet::default();
for c in [
RetryClass::Http5xx,
RetryClass::RateLimited,
RetryClass::Connection,
RetryClass::Timeout,
] {
assert!(set.contains(c), "default set should contain {c:?}");
}
let only5xx = RetryClassSet::from_iter([RetryClass::Http5xx]);
assert!(only5xx.contains(RetryClass::Http5xx));
assert!(!only5xx.contains(RetryClass::RateLimited));
}
}