1use crate::Error;
2use serde::{Deserialize, Serialize};
3use std::time::Duration;
4
5const RETRY_DELAY: Duration = Duration::from_millis(250);
6const MAX_RETRY_AFTER: Duration = Duration::from_secs(5);
7
8pub fn retry_delay(status: reqwest::StatusCode, retry_after: Option<&str>) -> Option<Duration> {
9 match status {
10 reqwest::StatusCode::TOO_MANY_REQUESTS => Some(
11 retry_after
12 .and_then(|value| value.parse::<u64>().ok())
13 .map(Duration::from_secs)
14 .unwrap_or(RETRY_DELAY)
15 .min(MAX_RETRY_AFTER),
16 ),
17 reqwest::StatusCode::BAD_GATEWAY
18 | reqwest::StatusCode::SERVICE_UNAVAILABLE
19 | reqwest::StatusCode::GATEWAY_TIMEOUT => Some(RETRY_DELAY),
20 _ => None,
21 }
22}
23
24pub async fn send_with_retry(
25 request: reqwest::RequestBuilder,
26) -> reqwest::Result<reqwest::Response> {
27 let retry = request.try_clone();
28 match request.send().await {
29 Ok(response) => {
30 let delay = retry_delay(
31 response.status(),
32 response
33 .headers()
34 .get(reqwest::header::RETRY_AFTER)
35 .and_then(|value| value.to_str().ok()),
36 );
37 if let (Some(request), Some(delay)) = (retry, delay) {
38 tokio::time::sleep(delay).await;
39 request.send().await
40 } else {
41 Ok(response)
42 }
43 }
44 Err(error) if error.is_connect() || error.is_timeout() => {
45 if let Some(request) = retry {
46 tokio::time::sleep(RETRY_DELAY).await;
47 request.send().await
48 } else {
49 Err(error)
50 }
51 }
52 Err(error) => Err(error),
53 }
54}
55
56#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
57pub struct ProviderOutcome {
58 pub provider: String,
59 pub status: ProviderStatus,
60 #[serde(skip_serializing_if = "Option::is_none")]
61 pub message: Option<String>,
62}
63
64#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
65#[serde(rename_all = "snake_case")]
66pub enum ProviderStatus {
67 Success,
68 Unavailable,
69 RateLimited,
70 Failed,
71}
72
73impl ProviderOutcome {
74 pub fn success(provider: impl Into<String>) -> Self {
75 Self {
76 provider: provider.into(),
77 status: ProviderStatus::Success,
78 message: None,
79 }
80 }
81
82 pub fn from_error(provider: impl Into<String>, error: &Error) -> Self {
83 let status = match error {
84 Error::Reqwest(error)
85 if error.status() == Some(reqwest::StatusCode::TOO_MANY_REQUESTS) =>
86 {
87 ProviderStatus::RateLimited
88 }
89 Error::Reqwest(error) if error.is_timeout() || error.is_connect() => {
90 ProviderStatus::Unavailable
91 }
92 _ => ProviderStatus::Failed,
93 };
94 Self {
95 provider: provider.into(),
96 status,
97 message: Some(error.to_string()),
98 }
99 }
100
101 pub fn succeeded(&self) -> bool {
102 self.status == ProviderStatus::Success
103 }
104
105 pub fn aggregate(provider: impl Into<String>, outcomes: &[Self]) -> Self {
106 let status = outcomes
107 .iter()
108 .map(|outcome| outcome.status)
109 .max_by_key(|status| match status {
110 ProviderStatus::Success => 0,
111 ProviderStatus::Unavailable => 1,
112 ProviderStatus::RateLimited => 2,
113 ProviderStatus::Failed => 3,
114 })
115 .unwrap_or(ProviderStatus::Success);
116 let messages: Vec<_> = outcomes
117 .iter()
118 .filter_map(|outcome| outcome.message.as_deref())
119 .collect();
120 Self {
121 provider: provider.into(),
122 status,
123 message: (!messages.is_empty()).then(|| messages.join("; ")),
124 }
125 }
126}
127
128pub fn default_success() -> ProviderOutcome {
129 ProviderOutcome::success("legacy")
130}
131
132impl std::fmt::Display for ProviderStatus {
133 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
134 f.write_str(match self {
135 Self::Success => "success",
136 Self::Unavailable => "unavailable",
137 Self::RateLimited => "rate_limited",
138 Self::Failed => "failed",
139 })
140 }
141}