Skip to main content

paper_gap/
provider.rs

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}