From dd8beae0b7b059c592b6849d03d57008e390de5e Mon Sep 17 00:00:00 2001
From: nick42d <133559267+nick42d@users.noreply.github.com>
Date: Mon, 5 Aug 2024 00:45:54 +0800
Subject: [PATCH] Make max retries configurable (#33)

* Implement customisable number of retries

* Better name for constant

* Fix rustdoc link

* Better name for option

* Fix rustdoc field name
---
 src/constants.rs               |  3 +++
 src/info.rs                    |  8 +++++---
 src/stream/streams/live.rs     |  4 ++--
 src/stream/streams/non_live.rs |  4 ++--
 src/structs.rs                 | 15 +++++++++++++++
 5 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/src/constants.rs b/src/constants.rs
index ec3d6e2..0762c51 100644
--- a/src/constants.rs
+++ b/src/constants.rs
@@ -1234,3 +1234,6 @@ pub static FORMATS: Lazy<HashMap<&str, StaticFormat>> = Lazy::new(|| {
         ),
     ])
 });
+
+/// Default max number of retries for a web reqwest.
+pub const DEFAULT_MAX_RETRIES: u32 = 3;
diff --git a/src/info.rs b/src/info.rs
index 6a2781d..b6b0120 100644
--- a/src/info.rs
+++ b/src/info.rs
@@ -15,7 +15,7 @@ use crate::stream::{LiveStream, LiveStreamOptions};
 use crate::structs::FFmpegArgs;
 
 use crate::{
-    constants::BASE_URL,
+    constants::{BASE_URL, DEFAULT_MAX_RETRIES},
     info_extras::{get_media, get_related_videos},
     stream::{NonLiveStream, NonLiveStreamOptions, Stream},
     structs::{
@@ -52,7 +52,7 @@ impl Video {
 
         let retry_policy = ExponentialBackoff::builder()
             .retry_bounds(Duration::from_millis(1000), Duration::from_millis(30000))
-            .build_with_max_retries(3);
+            .build_with_max_retries(DEFAULT_MAX_RETRIES);
         let client = ClientBuilder::new(client)
             .with(RetryTransientMiddleware::new_with_policy_and_strategy(
                 retry_policy,
@@ -102,9 +102,11 @@ impl Video {
             }
         };
 
+        let max_retries = options.request_options.max_retries.unwrap_or(DEFAULT_MAX_RETRIES);
+
         let retry_policy = ExponentialBackoff::builder()
             .retry_bounds(Duration::from_millis(1000), Duration::from_millis(30000))
-            .build_with_max_retries(3);
+            .build_with_max_retries(max_retries);
         let client = ClientBuilder::new(client)
             .with(RetryTransientMiddleware::new_with_policy_and_strategy(
                 retry_policy,
diff --git a/src/stream/streams/live.rs b/src/stream/streams/live.rs
index cf97c93..e7d0c86 100644
--- a/src/stream/streams/live.rs
+++ b/src/stream/streams/live.rs
@@ -1,4 +1,4 @@
-use crate::constants::DEFAULT_HEADERS;
+use crate::constants::{DEFAULT_HEADERS, DEFAULT_MAX_RETRIES};
 use crate::stream::{
     encryption::Encryption, media_format::MediaFormat, remote_data::RemoteData, segment::Segment,
     streams::Stream,
@@ -41,7 +41,7 @@ impl LiveStream {
                     std::time::Duration::from_millis(1000),
                     std::time::Duration::from_millis(30000),
                 )
-                .build_with_max_retries(3);
+                .build_with_max_retries(DEFAULT_MAX_RETRIES);
             reqwest_middleware::ClientBuilder::new(client)
                 .with(
                     reqwest_retry::RetryTransientMiddleware::new_with_policy_and_strategy(
diff --git a/src/stream/streams/non_live.rs b/src/stream/streams/non_live.rs
index d59a0d6..99eab8d 100644
--- a/src/stream/streams/non_live.rs
+++ b/src/stream/streams/non_live.rs
@@ -8,7 +8,7 @@ use std::sync::Arc;
 use tokio::sync::Mutex;
 use tokio::sync::RwLock;
 
-use crate::constants::DEFAULT_HEADERS;
+use crate::constants::{DEFAULT_HEADERS, DEFAULT_MAX_RETRIES};
 use crate::stream::streams::Stream;
 use crate::structs::{CustomRetryableStrategy, VideoError};
 
@@ -62,7 +62,7 @@ impl NonLiveStream {
                     std::time::Duration::from_millis(1000),
                     std::time::Duration::from_millis(30000),
                 )
-                .build_with_max_retries(3);
+                .build_with_max_retries(DEFAULT_MAX_RETRIES);
             reqwest_middleware::ClientBuilder::new(client)
                 .with(
                     reqwest_retry::RetryTransientMiddleware::new_with_policy_and_strategy(
diff --git a/src/structs.rs b/src/structs.rs
index f0b97d8..00d04f0 100644
--- a/src/structs.rs
+++ b/src/structs.rs
@@ -218,6 +218,21 @@ pub struct RequestOptions {
     ///     };
     /// ```
     pub ipv6_block: Option<String>,
+    /// Override the default number of retries to allow per web request (ie, per chunk downloaded)
+    /// Default is [`crate::constants::DEFAULT_MAX_RETRIES`].
+    ///
+    /// # Example
+    /// ```ignore
+    ///     // Allow 5 retries per chunk.
+    ///     let video_options = VideoOptions {
+    ///          request_options: RequestOptions {
+    ///               max_retries: Some(5),
+    ///                ..Default::default()
+    ///          },
+    ///          ..Default::default()
+    ///     };
+    /// ```
+    pub max_retries: Option<u32>,
 }
 
 #[derive(thiserror::Error, Debug)]