03.08.2024

* Custom request retry strategy implemented (4XX-5XX)
Mithronn · Aug 3, 2024 · fc3e9d6 · fc3e9d6
1 parent 7606f78
commit fc3e9d6
Show file tree

Hide file tree

Showing 4 changed files with 68 additions and 22 deletions.
diff --git a/src/info.rs b/src/info.rs
@@ -18,7 +18,9 @@ use crate::{
     constants::BASE_URL,
     info_extras::{get_media, get_related_videos},
     stream::{NonLiveStream, NonLiveStreamOptions, Stream},
-    structs::{PlayerResponse, VideoError, VideoInfo, VideoOptions, YTConfig},
+    structs::{
+        CustomRetryableStrategy, PlayerResponse, VideoError, VideoInfo, VideoOptions, YTConfig,
+    },
     utils::{
         between, choose_format, clean_video_details, get_functions, get_html, get_html5player,
         get_random_v6_ip, get_video_id, get_ytconfig, is_age_restricted_from_html,
@@ -49,10 +51,13 @@ impl Video {
         let client = Client::builder().build().map_err(VideoError::Reqwest)?;
 
         let retry_policy = ExponentialBackoff::builder()
-            .retry_bounds(Duration::from_millis(500), Duration::from_millis(10000))
+            .retry_bounds(Duration::from_millis(1000), Duration::from_millis(30000))
             .build_with_max_retries(3);
         let client = ClientBuilder::new(client)
-            .with(RetryTransientMiddleware::new_with_policy(retry_policy))
+            .with(RetryTransientMiddleware::new_with_policy_and_strategy(
+                retry_policy,
+                CustomRetryableStrategy,
+            ))
             .build();
 
         Ok(Self {
@@ -98,10 +103,13 @@ impl Video {
         };
 
         let retry_policy = ExponentialBackoff::builder()
-            .retry_bounds(Duration::from_millis(500), Duration::from_millis(10000))
+            .retry_bounds(Duration::from_millis(1000), Duration::from_millis(30000))
             .build_with_max_retries(3);
         let client = ClientBuilder::new(client)
-            .with(RetryTransientMiddleware::new_with_policy(retry_policy))
+            .with(RetryTransientMiddleware::new_with_policy_and_strategy(
+                retry_policy,
+                CustomRetryableStrategy,
+            ))
             .build();
 
         Ok(Self {

diff --git a/src/stream/streams/live.rs b/src/stream/streams/live.rs
@@ -1,10 +1,9 @@
 use crate::constants::DEFAULT_HEADERS;
-use crate::stream::encryption::Encryption;
-use crate::stream::media_format::MediaFormat;
-use crate::stream::remote_data::RemoteData;
-use crate::stream::segment::Segment;
-use crate::stream::streams::Stream;
-use crate::structs::VideoError;
+use crate::stream::{
+    encryption::Encryption, media_format::MediaFormat, remote_data::RemoteData, segment::Segment,
+    streams::Stream,
+};
+use crate::structs::{CustomRetryableStrategy, VideoError};
 use crate::utils::{get_html, make_absolute_url};
 
 use async_trait::async_trait;
@@ -39,14 +38,17 @@ impl LiveStream {
 
             let retry_policy = reqwest_retry::policies::ExponentialBackoff::builder()
                 .retry_bounds(
-                    std::time::Duration::from_millis(500),
-                    std::time::Duration::from_millis(10000),
+                    std::time::Duration::from_millis(1000),
+                    std::time::Duration::from_millis(30000),
                 )
                 .build_with_max_retries(3);
             reqwest_middleware::ClientBuilder::new(client)
-                .with(reqwest_retry::RetryTransientMiddleware::new_with_policy(
-                    retry_policy,
-                ))
+                .with(
+                    reqwest_retry::RetryTransientMiddleware::new_with_policy_and_strategy(
+                        retry_policy,
+                        CustomRetryableStrategy,
+                    ),
+                )
                 .build()
         };
 

diff --git a/src/stream/streams/non_live.rs b/src/stream/streams/non_live.rs
@@ -10,7 +10,7 @@ use tokio::sync::RwLock;
 
 use crate::constants::DEFAULT_HEADERS;
 use crate::stream::streams::Stream;
-use crate::structs::VideoError;
+use crate::structs::{CustomRetryableStrategy, VideoError};
 
 #[cfg(feature = "ffmpeg")]
 use crate::structs::FFmpegArgs;
@@ -59,14 +59,17 @@ impl NonLiveStream {
 
             let retry_policy = reqwest_retry::policies::ExponentialBackoff::builder()
                 .retry_bounds(
-                    std::time::Duration::from_millis(500),
-                    std::time::Duration::from_millis(10000),
+                    std::time::Duration::from_millis(1000),
+                    std::time::Duration::from_millis(30000),
                 )
                 .build_with_max_retries(3);
             reqwest_middleware::ClientBuilder::new(client)
-                .with(reqwest_retry::RetryTransientMiddleware::new_with_policy(
-                    retry_policy,
-                ))
+                .with(
+                    reqwest_retry::RetryTransientMiddleware::new_with_policy_and_strategy(
+                        retry_policy,
+                        CustomRetryableStrategy,
+                    ),
+                )
                 .build()
         };
 

diff --git a/src/structs.rs b/src/structs.rs
@@ -993,3 +993,36 @@ pub struct YTConfig {
     #[serde(rename = "STS")]
     pub sts: Option<u64>,
 }
+
+pub struct CustomRetryableStrategy;
+
+impl reqwest_retry::RetryableStrategy for CustomRetryableStrategy {
+    fn handle(
+        &self,
+        res: &reqwest_middleware::Result<reqwest::Response>,
+    ) -> Option<reqwest_retry::Retryable> {
+        match res {
+            // retry if 201
+            Ok(success) => custom_on_request_success(success),
+            Err(error) => reqwest_retry::default_on_request_failure(error),
+        }
+    }
+}
+
+/// Custom request success retry strategy.
+///
+/// Will only retry if:
+/// * The status was 5XX (server error)
+/// * The status was 4XX (client error)
+///
+/// Note that success here means that the request finished without interruption, not that it was logically OK.
+fn custom_on_request_success(success: &reqwest::Response) -> Option<reqwest_retry::Retryable> {
+    let status = success.status();
+    if status.is_server_error() || status.is_client_error() {
+        Some(reqwest_retry::Retryable::Transient)
+    } else if status.is_success() {
+        None
+    } else {
+        Some(reqwest_retry::Retryable::Fatal)
+    }
+}