post_archiver_utils/
request.rs

1use futures::StreamExt;
2use governor::{
3    Jitter, Quota, RateLimiter,
4    clock::{QuantaClock, QuantaInstant},
5    middleware::NoOpMiddleware,
6    state::{InMemoryState, NotKeyed},
7};
8use http::Method;
9use reqwest::{Client, IntoUrl, Request, Response};
10use reqwest_middleware::{ClientWithMiddleware, Middleware, Next, RequestBuilder};
11use reqwest_retry::{RetryTransientMiddleware, policies::ExponentialBackoff};
12use serde::de::DeserializeOwned;
13use std::{
14    fs::File,
15    io::{BufWriter, Write},
16    num::NonZeroU32,
17    ops::{Deref, DerefMut},
18    sync::Arc,
19    time::Duration,
20};
21use tokio::sync::Semaphore;
22
23use crate::{Error, Result};
24
25const RETRY_LIMIT: u32 = 3;
26
27#[derive(Debug, Clone)]
28pub struct ArchiveClient(ClientWithMiddleware);
29
30impl ArchiveClient {
31    pub fn new(client: Client, limit: usize) -> Self {
32        let retry_policy = ExponentialBackoff::builder().build_with_max_retries(RETRY_LIMIT);
33        let client = reqwest_middleware::ClientBuilder::new(client)
34            .with(SemaphoreMiddleware::new(limit))
35            .with(RetryTransientMiddleware::new_with_policy(retry_policy))
36            .build();
37
38        Self(client)
39    }
40
41    pub async fn fetch_with_method<T: DeserializeOwned>(
42        &self,
43        method: Method,
44        url: impl IntoUrl,
45    ) -> Result<T> {
46        let request = self.0.request(method, url);
47        let response = request.send().await?;
48        let response = response.bytes().await?;
49        serde_json::from_slice(&response).map_err(|e| {
50            Error::UnexpectedResponse(e, String::from_utf8(response.to_vec()).unwrap())
51        })
52    }
53
54    pub async fn fetch<T: DeserializeOwned>(&self, url: impl IntoUrl) -> Result<T> {
55        self.fetch_with_method(Method::GET, url).await
56    }
57
58    pub async fn download_with_method(
59        &self,
60        method: Method,
61        url: impl IntoUrl + Clone,
62        file: &mut File,
63    ) -> Result<()> {
64        async fn handle(request: RequestBuilder, file: &mut File) -> Result<()> {
65            file.set_len(0)?;
66
67            let response = request.send().await?;
68            let mut stream = response.bytes_stream();
69
70            let mut buffer = BufWriter::new(file);
71            while let Some(bytes) = stream.next().await {
72                let bytes = bytes?;
73                buffer.write_all(&bytes)?;
74            }
75            buffer.flush()?;
76            Ok(())
77        }
78
79        let mut err = Ok(());
80        for _ in 0..=RETRY_LIMIT {
81            let request = self.0.request(method.clone(), url.clone());
82            match handle(request, file).await {
83                Ok(_) => return Ok(()),
84                Err(e) => err = Err(e),
85            }
86        }
87        err
88    }
89
90    pub async fn download(&self, url: impl IntoUrl + Clone, file: &mut File) -> Result<()> {
91        self.download_with_method(Method::GET, url, file).await
92    }
93}
94
95impl Deref for ArchiveClient {
96    type Target = ClientWithMiddleware;
97
98    fn deref(&self) -> &Self::Target {
99        &self.0
100    }
101}
102
103impl DerefMut for ArchiveClient {
104    fn deref_mut(&mut self) -> &mut Self::Target {
105        &mut self.0
106    }
107}
108
109type ArchiveRateLimiter =
110    RateLimiter<NotKeyed, InMemoryState, QuantaClock, NoOpMiddleware<QuantaInstant>>;
111#[derive(Debug, Clone)]
112pub struct SemaphoreMiddleware(Arc<(Semaphore, ArchiveRateLimiter)>);
113
114impl SemaphoreMiddleware {
115    pub fn new(limit: usize) -> Self {
116        let semaphore = Semaphore::new(5);
117        let rate_limiter =
118            RateLimiter::direct(Quota::per_minute(NonZeroU32::new(limit as u32).unwrap()));
119        Self(Arc::new((semaphore, rate_limiter)))
120    }
121}
122
123#[async_trait::async_trait]
124impl Middleware for SemaphoreMiddleware {
125    async fn handle(
126        &self,
127        req: Request,
128        extensions: &mut http::Extensions,
129        next: Next<'_>,
130    ) -> reqwest_middleware::Result<Response> {
131        let (semaphore, rate_limiter) = self.0.as_ref();
132        let _ = semaphore.acquire().await.unwrap();
133        rate_limiter
134            .until_ready_with_jitter(Jitter::up_to(Duration::from_millis(800)))
135            .await;
136        next.run(req, extensions).await
137    }
138}
post_archiver_utils/request.rs

post_archiver_utils/
request.rs
