Resume server-interrupted downloads

Cargo.toml

@@ -92,7 +92,6 @@ tempfile = "3.8"
 thiserror = "2"
 threadpool = "1"
 tokio = { version = "1.26.0", default-features = false, features = ["macros", "rt-multi-thread", "sync"] }
-tokio-retry = "0.3.0"
 tokio-stream = "0.1.14"
 toml = "1.0"
 tracing = "0.1"

doc/dev-guide/src/tips-and-tricks.md

@@ -52,7 +52,9 @@ should one arise with the backtrack limits.
 When downloading a file, rustup will retry the download a number of times. The
 default is 3 times, but if this variable is set to a valid usize then it is the
 max retry count. A value of `0` means no retries, thus the default of `3` will
-mean a download is tried a total of four times before failing out.
+mean a download is tried a total of four times before failing out. Attempts
+that download at least 1 MiB of the file before failing are resumed by the
+next attempt and don't count towards this limit.
 
 ## `RUSTUP_BACKTRACE`
 

src/dist/download.rs

@@ -64,6 +64,8 @@ impl<'a> DownloadCfg<'a> {
     ) -> anyhow::Result<File> {
         utils::ensure_dir_exists("Download Directory", self.download_dir)?;
         let target_file = self.download_dir.join(Path::new(hash));
+        let partial_file_path = self.partial_file_path(hash);
+        let partial_file_existed = partial_file_path.exists();
 
         if target_file.exists() {
             let cached_result = file_hash(&target_file)?;
@@ -77,17 +79,6 @@ impl<'a> DownloadCfg<'a> {
             }
         }
 
-        let partial_file_path = target_file.with_file_name(
-            target_file
-                .file_name()
-                .map(|s| s.to_str().unwrap_or("_"))
-                .unwrap_or("_")
-                .to_owned()
-                + ".partial",
-        );
-
-        let partial_file_existed = partial_file_path.exists();
-
         let mut hasher = Sha256::new();
         let mut download = DownloadOptions::try_from(self.process)?
             .start(url, &partial_file_path)
@@ -132,6 +123,18 @@ impl<'a> DownloadCfg<'a> {
         }
     }
 
+    /// The number of bytes already downloaded for the file with the given hash
+    pub(crate) fn partial_len(&self, hash: &str) -> u64 {
+        match fs::metadata(self.partial_file_path(hash)) {
+            Ok(metadata) => metadata.len(),
+            Err(_) => 0,
+        }
+    }
+
+    fn partial_file_path(&self, hash: &str) -> PathBuf {
+        self.download_dir.join(format!("{hash}.partial"))
+    }
+
     pub(crate) fn clean(&self, hashes: &[impl AsRef<Path>]) -> anyhow::Result<()> {
         for hash in hashes.iter() {
             let used_file = self.download_dir.join(hash);
@@ -366,9 +369,10 @@ pub(crate) struct DownloadStatus {
 }
 
 impl DownloadStatus {
-    pub(crate) fn received_length(&self, len: u64) {
+    pub(crate) fn received_length(&self, len: u64, resumed: u64) {
         self.progress.reset();
         self.progress.set_length(len);
+        self.progress.set_position(resumed);
     }
 
     pub(crate) fn received_data(&self, len: usize) {

src/dist/manifestation.rs

@@ -770,32 +770,45 @@ impl<'a> ComponentBinary<'a> {
         }))
     }
 
-    async fn download(self, max_retries: usize) -> anyhow::Result<(ComponentInstall, &'a str)> {
-        use tokio_retry::{RetryIf, strategy::FixedInterval};
-
+    async fn download(self, mut retries: usize) -> anyhow::Result<(ComponentInstall, &'a str)> {
         let url = self.download_cfg.url(&self.binary.url)?;
-        let installer = RetryIf::start(
-            FixedInterval::from_millis(0).take(max_retries),
-            || {
-                self.download_cfg
-                    .download(&url, &self.binary.hash, &self.status)
-            },
-            |e: &anyhow::Error| {
-                // retry only known retriable cases
-                match e.downcast_ref::<RustupError>() {
-                    Some(RustupError::BrokenPartialFile)
-                    | Some(RustupError::DownloadingFile { .. }) => {
-                        self.status.retrying();
-                        true
-                    }
-                    _ => false,
+        let installer = loop {
+            let downloaded = self.download_cfg.partial_len(&self.binary.hash);
+            let future = self
+                .download_cfg
+                .download(&url, &self.binary.hash, &self.status);
+            let err = match future.await {
+                Ok(file) => break file,
+                Err(err) => err,
+            };
+
+            // retry only known retriable cases
+            match err.downcast_ref::<RustupError>() {
+                Some(RustupError::BrokenPartialFile)
+                | Some(RustupError::IncompletePartialFile)
+                | Some(RustupError::DownloadingFile { .. }) => {}
+                _ => {
+                    return Err(err).with_context(|| {
+                        RustupError::ComponentDownloadFailed(self.manifest.name(&self.component))
+                    });
                 }
-            },
-        )
-        .await
-        .with_context(|| {
-            RustupError::ComponentDownloadFailed(self.manifest.name(&self.component))
-        })?;
+            }
+
+            let progress = self
+                .download_cfg
+                .partial_len(&self.binary.hash)
+                .saturating_sub(downloaded);
+            if progress < Self::MIN_PROGRESS {
+                if retries == 0 {
+                    return Err(err).with_context(|| {
+                        RustupError::ComponentDownloadFailed(self.manifest.name(&self.component))
+                    });
+                }
+                retries -= 1;
+            }
+
+            self.status.retrying();
+        };
 
         let install = ComponentInstall {
             status: self.status,
@@ -812,6 +825,8 @@ impl<'a> ComponentBinary<'a> {
 
         Ok((install, &self.binary.hash))
     }
+
+    const MIN_PROGRESS: u64 = 1024 * 1024;
 }
 
 struct ComponentInstall {

src/download/mod.rs

@@ -310,13 +310,13 @@ impl<'a> Download<'a> {
         if let Some(len) = res.content_length() {
             let len = len + resume_from;
             if let Some(status) = self.status {
-                status.received_length(len);
+                status.received_length(len, resume_from);
             }
         }
 
         let mut stream = res.bytes_stream();
         while let Some(item) = stream.next().await {
-            let bytes = item.map_err(DownloadError::Reqwest)?;
+            let bytes = item.map_err(DownloadError::Interrupted)?;
             file.write_all(&bytes)
                 .context("unable to write download to disk")?;
             self.data_received(&bytes);
@@ -338,6 +338,8 @@ pub(crate) fn is_network_failure(err: &anyhow::Error) -> bool {
     match err.downcast_ref::<DownloadError>() {
         #[cfg(any(feature = "reqwest-rustls-tls", feature = "reqwest-native-tls"))]
         Some(DownloadError::Reqwest(e)) => e.is_timeout() || e.is_connect(),
+        #[cfg(any(feature = "reqwest-rustls-tls", feature = "reqwest-native-tls"))]
+        Some(DownloadError::Interrupted(_)) => true,
         _ => false,
     }
 }
@@ -447,6 +449,9 @@ enum DownloadError {
     #[cfg(any(feature = "reqwest-rustls-tls", feature = "reqwest-native-tls"))]
     #[error(transparent)]
     Reqwest(#[from] reqwest::Error),
+    #[cfg(any(feature = "reqwest-rustls-tls", feature = "reqwest-native-tls"))]
+    #[error(transparent)]
+    Interrupted(reqwest::Error),
     #[cfg(feature = "reqwest-rustls-tls")]
     #[error("failed to initialize platform verifier")]
     PlatformVerifierInit(#[source] rustls::Error),

src/download/tests.rs

@@ -2,7 +2,7 @@ use std::{
     convert::Infallible,
     env::remove_var,
     fs, io,
-    net::SocketAddr,
+    net::{SocketAddr, TcpListener},
     path::Path,
     sync::{
         LazyLock,
@@ -30,7 +30,7 @@ mod reqwest {
     use reqwest::{Client, Proxy};
     use url::Url;
 
-    use super::{scrub_env, serve_file, tmp_dir, write_file};
+    use super::{scrub_env, serve_file, serve_truncated, tmp_dir, write_file};
     use crate::download::{DownloadOptions, Tls};
 
     const OPTIONS: DownloadOptions = DownloadOptions {
@@ -170,6 +170,35 @@ mod reqwest {
         assert!(target_path.exists(), "partial file should not be deleted");
         assert_eq!(std::fs::read_to_string(&target_path).unwrap(), "123");
     }
+
+    #[tokio::test]
+    async fn interrupted_download_resumes_from_partial_file() {
+        let _guard = scrub_env().await;
+        let tmpdir = tmp_dir();
+        let target_path = tmpdir.path().join("downloaded.partial");
+
+        let addr = serve_truncated(b"123".to_vec(), 5);
+        let from_url = format!("http://{addr}").parse().unwrap();
+        OPTIONS
+            .start(&from_url, &target_path)
+            .with_resume()
+            .download()
+            .await
+            .expect_err("download should fail if the body is cut short");
+
+        assert_eq!(std::fs::read_to_string(&target_path).unwrap(), "123");
+
+        let addr = serve_file(b"12345".to_vec(), true);
+        let from_url = format!("http://{addr}").parse().unwrap();
+        OPTIONS
+            .start(&from_url, &target_path)
+            .with_resume()
+            .download()
+            .await
+            .expect("resumed download should succeed");
+
+        assert_eq!(std::fs::read_to_string(&target_path).unwrap(), "12345");
+    }
 }
 
 pub fn tmp_dir() -> TempDir {
@@ -284,6 +313,31 @@ fn serve_contents(
     res
 }
 
+/// Serve a single response that announces `len` bytes, then closes the connection after `contents`
+pub fn serve_truncated(contents: Vec<u8>, len: usize) -> SocketAddr {
+    let listener = TcpListener::bind("127.0.0.1:0").expect("can not bind");
+    let addr = listener.local_addr().unwrap();
+
+    thread::spawn(move || {
+        let (mut stream, _) = listener.accept().expect("could not accept connection");
+        let mut request = Vec::new();
+        let mut buf = [0; 1024];
+        loop {
+            let n = io::Read::read(&mut stream, &mut buf).expect("could not read request");
+            request.extend_from_slice(&buf[..n]);
+            if n == 0 || request.ends_with(b"\r\n\r\n") {
+                break;
+            }
+        }
+
+        let head = format!("HTTP/1.1 200 OK\r\nContent-Length: {len}\r\n\r\n");
+        io::Write::write_all(&mut stream, head.as_bytes()).expect("could not write response");
+        io::Write::write_all(&mut stream, &contents).expect("could not write response");
+    });
+
+    addr
+}
+
 /// Clear proxy-related environment variables
 ///
 /// Every test using a proxy-sensitive URL should call this and hold the returned guard,