reqwest_partial_retry/
lib.rs

1//! Wrapper around reqwest to allow for easy partial retries
2//!
3//! # Example
4//!
5//! ```
6//! use futures_util::StreamExt;
7//! use reqwest_partial_retry::ClientExt;
8//!
9//! # #[tokio::main]
10//! # async fn main() -> Result<(), Box<dyn std::error::Error>> {
11//! let client = reqwest::Client::new().resumable();
12//! let request = client.get("http://httpbin.org/ip").build().unwrap();
13//! let mut stream = client
14//!     .execute_resumable(request)
15//!     .await?
16//!     .bytes_stream_resumable();
17//!
18//! while let Some(item) = stream.next().await {
19//!     println!("Bytes: {:?}", item?);
20//! }
21//! # Ok(())
22//! # }
23//! ```
24
25#![warn(
26    missing_docs,
27    clippy::missing_errors_doc,
28    clippy::missing_panics_doc,
29    clippy::missing_const_for_fn,
30    clippy::future_not_send,
31    clippy::large_futures
32)]
33
34use std::fmt::Display;
35use std::future::Future;
36use std::ops::Deref;
37use std::pin::Pin;
38use std::sync::Arc;
39use std::task::{Context, Poll};
40use std::time::Duration;
41
42use bytes::Bytes;
43use chrono::Utc;
44use futures_core::{ready, Stream};
45use reqwest::Request;
46use reqwest_retry::{DefaultRetryableStrategy, RetryPolicy, Retryable, RetryableStrategy};
47use retry_policies::RetryDecision;
48
49/// Config for the resumable [`Client`]
50pub struct Config {
51    inner: Arc<ConfigInner>,
52}
53
54struct ConfigInner {
55    stream_timeout: Option<Duration>,
56    retry_policy: Box<dyn RetryPolicy + Send + Sync>,
57    retryable_strategy: Box<dyn RetryableStrategy + Send + Sync>,
58}
59
60impl Clone for Config {
61    #[inline]
62    fn clone(&self) -> Self {
63        Self {
64            inner: self.inner.clone(),
65        }
66    }
67}
68
69impl Default for Config {
70    /// The default config has no stream timeout,
71    /// uses an exponential backoff with 10 retries,
72    /// and uses the [`DefaultRetryableStrategy`]
73    #[inline]
74    fn default() -> Self {
75        Self {
76            inner: Arc::new(ConfigInner {
77                stream_timeout: None,
78                retry_policy: Box::new(
79                    retry_policies::policies::ExponentialBackoffBuilder::default().build_with_max_retries(10),
80                ),
81                retryable_strategy: Box::new(DefaultRetryableStrategy),
82            }),
83        }
84    }
85}
86
87impl Config {
88    /// Create a `ConfigBuilder`
89    #[inline]
90    pub fn builder() -> ConfigBuilder {
91        ConfigBuilder::default()
92    }
93
94    /// Get the stream timeout
95    #[inline]
96    pub fn stream_timeout(&self) -> &Option<Duration> {
97        &self.inner.stream_timeout
98    }
99
100    /// Get the retry policy
101    #[inline]
102    pub fn retry_policy(&self) -> &(dyn RetryPolicy + Send + Sync) {
103        self.inner.retry_policy.as_ref()
104    }
105
106    /// Get the retryable strategy
107    #[inline]
108    pub fn retryable_strategy(&self) -> &(dyn RetryableStrategy + Send + Sync) {
109        self.inner.retryable_strategy.as_ref()
110    }
111}
112
113/// ConfigBuilder for the [`Config`] of the resumable [`Client`]
114pub struct ConfigBuilder {
115    stream_timeout: Option<Duration>,
116    retry_policy: Box<dyn RetryPolicy + Send + Sync>,
117    retryable_strategy: Box<dyn RetryableStrategy + Send + Sync>,
118}
119
120impl Default for ConfigBuilder {
121    /// The default config builder has no stream timeout,
122    /// uses an exponential backoff with 10 retries,
123    /// and uses the [`DefaultRetryableStrategy`]
124    #[inline]
125    fn default() -> Self {
126        Self {
127            stream_timeout: None,
128            retry_policy: Box::new(
129                retry_policies::policies::ExponentialBackoffBuilder::default().build_with_max_retries(10),
130            ),
131            retryable_strategy: Box::new(DefaultRetryableStrategy),
132        }
133    }
134}
135
136impl ConfigBuilder {
137    /// Set the timeout for the
138    /// [`bytes_stream_resumable`](ResumableResponse::bytes_stream_resumable).
139    ///
140    /// If no new data has been received for the specified duration, it times out,
141    /// and depending on the `RetryPolicy` and `RetryableStrategy` a retry is tried.
142    ///
143    /// If the value is `None`, it will never time out.
144    #[inline]
145    pub const fn stream_timeout(mut self, stream_timeout: Option<Duration>) -> Self {
146        self.stream_timeout = stream_timeout;
147        self
148    }
149
150    /// Set the retry policy.
151    ///
152    /// Retries are counted in [`execute_resumable`](Client::execute_resumable)
153    /// and [`bytes_stream_resumable`](ResumableResponse::bytes_stream_resumable),
154    /// and will be reset whenever new data has been received.
155    #[inline]
156    pub fn retry_policy<P: RetryPolicy + Send + Sync + 'static>(mut self, retry_policy: P) -> Self {
157        self.retry_policy = Box::new(retry_policy);
158        self
159    }
160
161    /// Set the retryable strategy.
162    #[inline]
163    pub fn retryable_strategy<S: RetryableStrategy + Send + Sync + 'static>(mut self, retryable_strategy: S) -> Self {
164        self.retryable_strategy = Box::new(retryable_strategy);
165        self
166    }
167
168    /// Build a `Config`
169    #[inline]
170    pub fn build(self) -> Config {
171        let Self {
172            stream_timeout,
173            retry_policy,
174            retryable_strategy,
175        } = self;
176
177        Config {
178            inner: Arc::new(ConfigInner {
179                stream_timeout,
180                retry_policy,
181                retryable_strategy,
182            }),
183        }
184    }
185}
186
187/// Extension to [`reqwest::Client`] that provides methods to convert it into a resumable [`Client`]
188pub trait ClientExt {
189    /// Convert a [`reqwest::Client`] into a
190    /// [`reqwest_partial_retry::Client`](Client)
191    fn resumable(self) -> Client;
192
193    /// Convert a [`reqwest::Client`] into a
194    /// [`reqwest_partial_retry::Client`](Client) with a config
195    fn resumable_with_config(self, config: Config) -> Client;
196}
197
198impl ClientExt for reqwest::Client {
199    #[inline]
200    fn resumable(self) -> Client {
201        Client {
202            client: self,
203            config: Config::default(),
204        }
205    }
206
207    #[inline]
208    fn resumable_with_config(self, config: Config) -> Client {
209        Client { client: self, config }
210    }
211}
212
213/// A wrapper for [`reqwest::Client`] to allow for resumable byte streams
214pub struct Client {
215    client: reqwest::Client,
216    config: Config,
217}
218
219impl Clone for Client {
220    #[inline]
221    fn clone(&self) -> Self {
222        Self {
223            client: self.client.clone(),
224            config: self.config.clone(),
225        }
226    }
227}
228
229impl Deref for Client {
230    type Target = reqwest::Client;
231
232    #[inline]
233    fn deref(&self) -> &Self::Target {
234        &self.client
235    }
236}
237
238impl Client {
239    /// Executes a `Request`.
240    ///
241    /// A `Request` can be built manually with `Request::new()` or obtained
242    /// from a RequestBuilder with `RequestBuilder::build()`.
243    ///
244    /// # Errors
245    ///
246    /// This method fails if the `Request` is not cloneable (i.e. if the body is
247    /// a stream).
248    ///
249    /// It also fails if the [`RetryableStrategy`] decides that the error is fatal,
250    /// or it is transient and the retry limit has been reached.
251    pub async fn execute_resumable(&self, request: Request) -> Result<ResumableResponse, Error> {
252        let mut current_retries = 0;
253
254        'outer_loop: loop {
255            let Some(cloned_request) = request.try_clone() else {
256                return Err(Error::RequestNotCloneable);
257            };
258
259            let response = self
260                .client
261                .execute(cloned_request)
262                .await
263                .map_err(reqwest_middleware::Error::Reqwest);
264
265            match (self.config.inner.retryable_strategy.handle(&response), response) {
266                (Some(Retryable::Fatal), response) => {
267                    let reqwest_error = if let Err(reqwest_middleware::Error::Reqwest(err)) = response {
268                        Some(err)
269                    } else {
270                        None
271                    };
272
273                    return Err(Error::UnresolvableError(reqwest_error));
274                }
275                (None, Ok(response)) => {
276                    let headers = hyperx::Headers::from(response.headers());
277                    let accept_byte_ranges = if let Some(hyperx::header::AcceptRanges(ranges)) = headers.get() {
278                        ranges.iter().any(|u| *u == hyperx::header::RangeUnit::Bytes)
279                    } else {
280                        false
281                    };
282
283                    return Ok(ResumableResponse {
284                        client: self.clone(),
285                        request,
286                        response,
287                        accept_byte_ranges,
288                    });
289                }
290                (Some(Retryable::Transient), response) | (None, response @ Err(_)) => {
291                    let retry_decision = self.config.inner.retry_policy.should_retry(current_retries);
292                    current_retries += 1;
293
294                    match retry_decision {
295                        RetryDecision::Retry { execute_after } => {
296                            if let Ok(duration) = (execute_after - Utc::now()).to_std() {
297                                tokio::time::sleep(duration).await;
298                            }
299
300                            continue 'outer_loop;
301                        }
302                        RetryDecision::DoNotRetry => {
303                            let reqwest_error = if let Err(reqwest_middleware::Error::Reqwest(err)) = response {
304                                Some(RetryError::RequestError(err))
305                            } else {
306                                None
307                            };
308
309                            return Err(Error::MaxRetries(reqwest_error));
310                        }
311                    }
312                }
313            }
314        }
315    }
316}
317
318/// A wrapper for [`reqwest::Response`] to allow for resumable byte streams
319pub struct ResumableResponse {
320    client: Client,
321    request: reqwest::Request,
322    response: reqwest::Response,
323    accept_byte_ranges: bool,
324}
325
326impl Deref for ResumableResponse {
327    type Target = reqwest::Response;
328
329    #[inline]
330    fn deref(&self) -> &Self::Target {
331        &self.response
332    }
333}
334
335impl ResumableResponse {
336    /// Get the underlying [`reqwest::Response`]
337    #[inline]
338    pub fn response(self) -> reqwest::Response {
339        self.response
340    }
341
342    /// Convert the response into a `Stream` of `Bytes` from the body.
343    ///
344    /// # Example
345    ///
346    /// ```
347    /// use futures_util::StreamExt;
348    /// use reqwest_partial_retry::ClientExt;
349    ///
350    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
351    /// let client = reqwest::Client::new().resumable();
352    /// let request = client.get("http://httpbin.org/ip").build().unwrap();
353    /// let mut stream = client
354    ///     .execute_resumable(request)
355    ///     .await?
356    ///     .bytes_stream_resumable();
357    ///
358    /// while let Some(item) = stream.next().await {
359    ///     println!("Bytes: {:?}", item?);
360    /// }
361    /// # Ok(())
362    /// # }
363    /// ```
364    #[inline]
365    pub fn bytes_stream_resumable(self) -> impl Stream<Item = Result<Bytes, Error>> + Send + Unpin {
366        ResumableBytesStream {
367            client: self.client,
368            request: self.request,
369            accept_byte_ranges: self.accept_byte_ranges,
370            pos: 0,
371            read_bytes: 0,
372            current_retries: 0,
373            stream: Box::pin(self.response.bytes_stream()),
374            next_request: None,
375            sleep: Box::pin(tokio::time::sleep(Duration::ZERO)),
376            run_sleep: false,
377            stream_sleep: Box::pin(tokio::time::sleep(Duration::ZERO)),
378            tried_stream_poll: false,
379        }
380    }
381}
382
383struct ResumableBytesStream {
384    client: Client,
385    request: reqwest::Request,
386    accept_byte_ranges: bool,
387    pos: u64,
388    read_bytes: u64,
389    current_retries: u32,
390    stream: Pin<Box<dyn Stream<Item = reqwest::Result<Bytes>> + Send>>,
391    next_request: Option<Pin<Box<dyn Future<Output = Result<reqwest::Response, reqwest::Error>> + Send>>>,
392    sleep: Pin<Box<tokio::time::Sleep>>,
393    run_sleep: bool,
394    stream_sleep: Pin<Box<tokio::time::Sleep>>,
395    tried_stream_poll: bool,
396}
397
398/// The errors that may occur
399#[derive(thiserror::Error, Debug)]
400pub enum Error {
401    /// The error was unresolvable
402    UnresolvableError(Option<reqwest::Error>),
403    /// The maximum retry limit has been reached
404    MaxRetries(Option<RetryError>),
405    /// The request is not cloneable
406    RequestNotCloneable,
407}
408
409impl Display for Error {
410    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
411        match self {
412            Error::UnresolvableError(err) => {
413                if let Some(err) = err {
414                    write!(f, "encountered fatal error: {err}")
415                } else {
416                    write!(f, "encountered fatal error")
417                }
418            }
419            Error::MaxRetries(err) => {
420                if let Some(err) = err {
421                    write!(f, "maximum amount of retries reached: {err}")
422                } else {
423                    write!(f, "maximum amount of retries reached")
424                }
425            }
426            Error::RequestNotCloneable => write!(f, "request is not cloneable"),
427        }
428    }
429}
430
431/// The error that possibly caused the retry
432#[derive(thiserror::Error, Debug)]
433pub enum RetryError {
434    /// A request error
435    #[error("request error: {0}")]
436    RequestError(reqwest::Error),
437    /// Requested range is not satisfiable
438    #[error("requested range is not satisfiable")]
439    RangeNotSatisfiable,
440    /// A stream error
441    #[error("stream error: {0}")]
442    StreamError(reqwest::Error),
443    /// The stream timed out
444    #[error("stream timed out")]
445    StreamTimeout,
446}
447
448impl Stream for ResumableBytesStream {
449    type Item = Result<Bytes, Error>;
450
451    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
452        macro_rules! prepare_retry {
453            ($retry_after:expr $(,)?) => {
454                if let Ok(duration) = ($retry_after - Utc::now()).to_std() {
455                    self.sleep.as_mut().reset(tokio::time::Instant::now() + duration);
456                    self.run_sleep = true;
457                }
458
459                let Some(mut cloned_request) = self.request.try_clone() else {
460                    return Poll::Ready(Some(Err(Error::RequestNotCloneable)));
461                };
462
463                if self.accept_byte_ranges {
464                    let range_value = reqwest::header::HeaderValue::from_str(
465                        &hyperx::header::Range::Bytes(vec![hyperx::header::ByteRangeSpec::AllFrom(self.pos)])
466                            .to_string(),
467                    )
468                    .unwrap();
469                    let _ = cloned_request
470                        .headers_mut()
471                        .insert(reqwest::header::RANGE, range_value);
472                }
473
474                self.next_request = Some(Box::pin(self.client.client.execute(cloned_request)));
475            };
476        }
477
478        'outer_loop: loop {
479            if self.run_sleep {
480                ready!(self.sleep.as_mut().poll(cx));
481                self.run_sleep = false;
482            }
483
484            if let Some(next_request) = &mut self.next_request {
485                let response = ready!(next_request.as_mut().poll(cx)).map_err(reqwest_middleware::Error::Reqwest);
486                self.next_request = None;
487
488                if let Ok(response) = &response {
489                    if matches!(
490                        response.status(),
491                        reqwest::StatusCode::OK | reqwest::StatusCode::RANGE_NOT_SATISFIABLE
492                    ) {
493                        self.accept_byte_ranges = false;
494                    }
495                }
496
497                match (self.client.config.inner.retryable_strategy.handle(&response), response) {
498                    (Some(Retryable::Fatal), response) => {
499                        let reqwest_error = if let Err(reqwest_middleware::Error::Reqwest(err)) = response {
500                            Some(err)
501                        } else {
502                            None
503                        };
504
505                        return Poll::Ready(Some(Err(Error::UnresolvableError(reqwest_error))));
506                    }
507                    (None, Ok(response)) => {
508                        if !self.accept_byte_ranges {
509                            self.pos = 0;
510                        }
511
512                        if response.status() == reqwest::StatusCode::RANGE_NOT_SATISFIABLE {
513                            let retry_decision =
514                                self.client.config.inner.retry_policy.should_retry(self.current_retries);
515                            self.current_retries += 1;
516
517                            match retry_decision {
518                                RetryDecision::Retry { execute_after } => {
519                                    prepare_retry!(execute_after);
520                                    continue 'outer_loop;
521                                }
522                                RetryDecision::DoNotRetry => {
523                                    return Poll::Ready(Some(Err(Error::MaxRetries(Some(
524                                        RetryError::RangeNotSatisfiable,
525                                    )))));
526                                }
527                            }
528                        }
529
530                        self.stream = Box::pin(response.bytes_stream());
531                        self.current_retries = 0;
532
533                        if let Some(timeout_duration) = self.client.config.inner.stream_timeout {
534                            let stream_sleep_deadline = tokio::time::Instant::now() + timeout_duration;
535                            self.stream_sleep.as_mut().reset(stream_sleep_deadline);
536                        }
537                    }
538                    (Some(Retryable::Transient), response) | (None, response @ Err(_)) => {
539                        let retry_decision = self.client.config.inner.retry_policy.should_retry(self.current_retries);
540                        self.current_retries += 1;
541
542                        match retry_decision {
543                            RetryDecision::Retry { execute_after } => {
544                                prepare_retry!(execute_after);
545                                continue 'outer_loop;
546                            }
547                            RetryDecision::DoNotRetry => {
548                                let reqwest_error = if let Err(reqwest_middleware::Error::Reqwest(err)) = response {
549                                    Some(RetryError::RequestError(err))
550                                } else {
551                                    None
552                                };
553
554                                return Poll::Ready(Some(Err(Error::MaxRetries(reqwest_error))));
555                            }
556                        }
557                    }
558                }
559            }
560
561            if !self.tried_stream_poll {
562                self.tried_stream_poll = true;
563
564                if let Some(timeout_duration) = self.client.config.inner.stream_timeout {
565                    let stream_sleep_deadline = tokio::time::Instant::now() + timeout_duration;
566                    self.stream_sleep.as_mut().reset(stream_sleep_deadline);
567                }
568            }
569
570            'stream_loop: loop {
571                match self.stream.as_mut().poll_next(cx) {
572                    Poll::Ready(Some(Err(err))) => {
573                        let retry_decision = self.client.config.inner.retry_policy.should_retry(self.current_retries);
574                        self.current_retries += 1;
575
576                        match retry_decision {
577                            RetryDecision::Retry { execute_after } => {
578                                prepare_retry!(execute_after);
579                                continue 'outer_loop;
580                            }
581                            RetryDecision::DoNotRetry => {
582                                return Poll::Ready(Some(Err(Error::MaxRetries(Some(RetryError::StreamError(err))))));
583                            }
584                        }
585                    }
586                    Poll::Ready(Some(Ok(bytes))) => {
587                        if !bytes.is_empty() {
588                            self.current_retries = 0;
589
590                            if let Some(timeout_duration) = self.client.config.inner.stream_timeout {
591                                let stream_sleep_deadline = tokio::time::Instant::now() + timeout_duration;
592                                self.stream_sleep.as_mut().reset(stream_sleep_deadline);
593                            }
594
595                            let start_pos = self.pos;
596                            self.pos += bytes.len() as u64;
597
598                            if self.pos > self.read_bytes {
599                                let diff_read = (self.read_bytes - start_pos) as usize;
600                                self.read_bytes = self.pos;
601                                return Poll::Ready(Some(Ok(bytes.slice(diff_read..))));
602                            }
603                        }
604
605                        continue 'stream_loop;
606                    }
607                    Poll::Ready(None) => return Poll::Ready(None),
608                    Poll::Pending => {
609                        if self.client.config.inner.stream_timeout.is_some() {
610                            ready!(self.stream_sleep.as_mut().poll(cx));
611
612                            let retry_decision =
613                                self.client.config.inner.retry_policy.should_retry(self.current_retries);
614                            self.current_retries += 1;
615
616                            match retry_decision {
617                                RetryDecision::Retry { execute_after } => {
618                                    prepare_retry!(execute_after);
619                                    continue 'outer_loop;
620                                }
621                                RetryDecision::DoNotRetry => {
622                                    return Poll::Ready(Some(Err(Error::MaxRetries(Some(RetryError::StreamTimeout)))));
623                                }
624                            }
625                        }
626
627                        return Poll::Pending;
628                    }
629                }
630            }
631        }
632    }
633}
reqwest_partial_retry/lib.rs

reqwest_partial_retry/
lib.rs
