From fc3e9d662a4243216d549fe52237012b64031fcd Mon Sep 17 00:00:00 2001 From: Mithronn Date: Sat, 3 Aug 2024 16:58:22 +0300 Subject: [PATCH] 03.08.2024 * Custom request retry strategy implemented (4XX-5XX) --- src/info.rs | 18 +++++++++++++----- src/stream/streams/live.rs | 24 +++++++++++++----------- src/stream/streams/non_live.rs | 15 +++++++++------ src/structs.rs | 33 +++++++++++++++++++++++++++++++++ 4 files changed, 68 insertions(+), 22 deletions(-) diff --git a/src/info.rs b/src/info.rs index 7dac2c0..6a2781d 100644 --- 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 index 0bd7fb0..cf97c93 100644 --- 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 index 49df89b..d59a0d6 100644 --- 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 index 9d65a28..f0b97d8 100644 --- a/src/structs.rs +++ b/src/structs.rs @@ -993,3 +993,36 @@ pub struct YTConfig { #[serde(rename = "STS")] pub sts: Option, } + +pub struct CustomRetryableStrategy; + +impl reqwest_retry::RetryableStrategy for CustomRetryableStrategy { + fn handle( + &self, + res: &reqwest_middleware::Result, + ) -> Option { + 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 { + 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) + } +}