Merge pull request #42 from flatcar/dongsu/client-timeout-retry

download_sysext: HTTP client timeout and retry

src/bin/download_sysext.rs

@@ -7,6 +7,7 @@ use std::io;
 use std::io::BufReader;
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
+use std::time::Duration;
 
 #[macro_use]
 extern crate log;
@@ -318,6 +319,9 @@ impl Args {
     }
 }
 
+const HTTP_CONN_TIMEOUT: u64 = 20;
+const DOWNLOAD_TIMEOUT: u64 = 3600;
+
 fn main() -> Result<(), Box<dyn Error>> {
     env_logger::init();
 
@@ -337,7 +341,7 @@ fn main() -> Result<(), Box<dyn Error>> {
     fs::create_dir_all(&temp_dir)?;
 
     // The default policy of reqwest Client supports max 10 attempts on HTTP redirect.
-    let client = Client::builder().redirect(Policy::default()).build()?;
+    let client = Client::builder().connect_timeout(Duration::from_secs(HTTP_CONN_TIMEOUT)).timeout(Duration::from_secs(DOWNLOAD_TIMEOUT)).redirect(Policy::default()).build()?;
 
     // If input_xml exists, simply read it.
     // If not, try to read from payload_url.

src/download.rs

@@ -10,6 +10,8 @@ use reqwest::blocking::Client;
 
 use sha2::{Sha256, Digest};
 
+const MAX_DOWNLOAD_RETRY: u32 = 20;
+
 pub struct DownloadResult {
     pub hash: omaha::Hash<omaha::Sha256>,
     pub data: File,
@@ -57,17 +59,17 @@ pub fn hash_on_disk_sha256(path: &Path, maxlen: Option<usize>) -> Result<omaha::
     Ok(omaha::Hash::from_bytes(hasher.finalize().into()))
 }
 
-pub fn download_and_hash<U>(client: &Client, url: U, path: &Path, print_progress: bool) -> Result<DownloadResult>
+fn do_download_and_hash<U>(client: &Client, url: U, path: &Path, print_progress: bool) -> Result<DownloadResult>
 where
     U: reqwest::IntoUrl + Clone,
     Url: From<U>,
 {
     let client_url = url.clone();
 
     #[rustfmt::skip]
-    let mut res = client.get(url)
+    let mut res = client.get(url.clone())
         .send()
-        .context(format!("client get and send({:?}) failed", client_url.as_str()))?;
+        .context(format!("client get & send{:?} failed ", client_url.as_str()))?;
 
     // Redirect was already handled at this point, so there is no need to touch
     // response or url again. Simply print info and continue.
@@ -98,3 +100,14 @@ where
         data: file,
     })
 }
+
+pub fn download_and_hash<U>(client: &Client, url: U, path: &Path, print_progress: bool) -> Result<DownloadResult>
+where
+    U: reqwest::IntoUrl + Clone,
+    Url: From<U>,
+{
+    crate::retry_loop(
+        || do_download_and_hash(client, url.clone(), path, print_progress),
+        MAX_DOWNLOAD_RETRY,
+    )
+}

src/lib.rs

@@ -3,4 +3,7 @@ pub use download::DownloadResult;
 pub use download::download_and_hash;
 pub use download::hash_on_disk_sha256;
 
+mod util;
+pub use util::retry_loop;
+
 pub mod request;

src/util.rs

@@ -0,0 +1,25 @@
+use core::time::Duration;
+use std::thread::sleep;
+
+const RETRY_INTERVAL_MSEC: u64 = 1000;
+
+pub fn retry_loop<F, T, E>(mut func: F, max_tries: u32) -> Result<T, E>
+where
+    F: FnMut() -> Result<T, E>,
+{
+    let mut tries = 0;
+
+    loop {
+        match func() {
+            ok @ Ok(_) => return ok,
+            err @ Err(_) => {
+                tries += 1;
+
+                if tries >= max_tries {
+                    return err;
+                }
+                sleep(Duration::from_millis(RETRY_INTERVAL_MSEC));
+            }
+        }
+    }
+}