Add retries for Python downloads (#9274)

## Summary This uses the same approach as in the rest of uv, but with another dedicated method for retries. Closes #8525.
astral-sh · Nov 20, 2024 · 1b13036 · 1b13036
1 parent 289771e
commit 1b13036
Show file tree

Hide file tree

Showing 16 changed files with 107 additions and 54 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/uv-client/src/base_client.rs b/crates/uv-client/src/base_client.rs
@@ -408,11 +408,6 @@ impl BaseClient {
         self.connectivity
     }
 
-    /// The number of retries to attempt on transient errors.
-    pub fn retries(&self) -> u32 {
-        self.retries
-    }
-
     /// The [`RetryPolicy`] for the client.
     pub fn retry_policy(&self) -> ExponentialBackoff {
         ExponentialBackoff::builder().build_with_max_retries(self.retries)
@@ -460,7 +455,7 @@ impl RetryableStrategy for UvRetryableStrategy {
 /// Check for additional transient error kinds not supported by the default retry strategy in `reqwest_retry`.
 ///
 /// These cases should be safe to retry with [`Retryable::Transient`].
-pub(crate) fn is_extended_transient_error(err: &dyn Error) -> bool {
+pub fn is_extended_transient_error(err: &dyn Error) -> bool {
     trace!("Attempting to retry error: {err:?}");
 
     if let Some(err) = find_source::<WrappedReqwestError>(&err) {

diff --git a/crates/uv-client/src/lib.rs b/crates/uv-client/src/lib.rs
@@ -1,5 +1,6 @@
 pub use base_client::{
-    AuthIntegration, BaseClient, BaseClientBuilder, UvRetryableStrategy, DEFAULT_RETRIES,
+    is_extended_transient_error, AuthIntegration, BaseClient, BaseClientBuilder,
+    UvRetryableStrategy, DEFAULT_RETRIES,
 };
 pub use cached_client::{CacheControl, CachedClient, CachedClientError, DataWithCachePolicy};
 pub use error::{Error, ErrorKind, WrappedReqwestError};

diff --git a/crates/uv-python/Cargo.toml b/crates/uv-python/Cargo.toml
@@ -45,6 +45,7 @@ owo-colors = { workspace = true }
 regex = { workspace = true }
 reqwest = { workspace = true }
 reqwest-middleware = { workspace = true }
+reqwest-retry = { workspace = true }
 rmp-serde = { workspace = true }
 same-file = { workspace = true }
 schemars = { workspace = true, optional = true }

diff --git a/crates/uv-python/src/downloads.rs b/crates/uv-python/src/downloads.rs
@@ -1,18 +1,22 @@
-use futures::TryStreamExt;
-use owo_colors::OwoColorize;
 use std::fmt::Display;
 use std::io;
 use std::path::{Path, PathBuf};
 use std::pin::Pin;
 use std::str::FromStr;
 use std::task::{Context, Poll};
+use std::time::{Duration, SystemTime};
+
+use futures::TryStreamExt;
+use owo_colors::OwoColorize;
+use reqwest_retry::RetryPolicy;
 use thiserror::Error;
 use tokio::io::{AsyncRead, ReadBuf};
 use tokio_util::compat::FuturesAsyncReadCompatExt;
 use tokio_util::either::Either;
 use tracing::{debug, instrument};
 use url::Url;
-use uv_client::WrappedReqwestError;
+
+use uv_client::{is_extended_transient_error, WrappedReqwestError};
 use uv_distribution_filename::{ExtensionError, SourceDistExtension};
 use uv_extract::hash::Hasher;
 use uv_fs::{rename_with_retry, Simplified};
@@ -417,6 +421,7 @@ impl FromStr for PythonDownloadRequest {
 
 include!("downloads.inc");
 
+#[derive(Debug, Clone)]
 pub enum DownloadResult {
     AlreadyAvailable(PathBuf),
     Fetched(PathBuf),
@@ -458,16 +463,66 @@ impl ManagedPythonDownload {
         self.sha256
     }
 
-    /// Download and extract
+    /// Download and extract a Python distribution, retrying on failure.
+    #[instrument(skip(client, installation_dir, cache_dir, reporter), fields(download = % self.key()))]
+    pub async fn fetch_with_retry(
+        &self,
+        client: &uv_client::BaseClient,
+        installation_dir: &Path,
+        cache_dir: &Path,
+        reinstall: bool,
+        python_install_mirror: Option<&str>,
+        pypy_install_mirror: Option<&str>,
+        reporter: Option<&dyn Reporter>,
+    ) -> Result<DownloadResult, Error> {
+        let mut n_past_retries = 0;
+        let start_time = SystemTime::now();
+        let retry_policy = client.retry_policy();
+        loop {
+            let result = self
+                .fetch(
+                    client,
+                    installation_dir,
+                    cache_dir,
+                    reinstall,
+                    python_install_mirror,
+                    pypy_install_mirror,
+                    reporter,
+                )
+                .await;
+            if result
+                .as_ref()
+                .err()
+                .is_some_and(|err| is_extended_transient_error(err))
+            {
+                let retry_decision = retry_policy.should_retry(start_time, n_past_retries);
+                if let reqwest_retry::RetryDecision::Retry { execute_after } = retry_decision {
+                    debug!(
+                        "Transient failure while handling response for {}; retrying...",
+                        self.key()
+                    );
+                    let duration = execute_after
+                        .duration_since(SystemTime::now())
+                        .unwrap_or_else(|_| Duration::default());
+                    tokio::time::sleep(duration).await;
+                    n_past_retries += 1;
+                    continue;
+                }
+            }
+            return result;
+        }
+    }
+
+    /// Download and extract a Python distribution.
     #[instrument(skip(client, installation_dir, cache_dir, reporter), fields(download = % self.key()))]
     pub async fn fetch(
         &self,
         client: &uv_client::BaseClient,
         installation_dir: &Path,
         cache_dir: &Path,
         reinstall: bool,
-        python_install_mirror: Option<String>,
-        pypy_install_mirror: Option<String>,
+        python_install_mirror: Option<&str>,
+        pypy_install_mirror: Option<&str>,
         reporter: Option<&dyn Reporter>,
     ) -> Result<DownloadResult, Error> {
         let url = self.download_url(python_install_mirror, pypy_install_mirror)?;
@@ -492,7 +547,7 @@ impl ManagedPythonDownload {
 
         debug!(
             "Downloading {url} to temporary location: {}",
-            temp_dir.path().simplified().display()
+            temp_dir.path().simplified_display()
         );
 
         let mut hashers = self
@@ -589,8 +644,8 @@ impl ManagedPythonDownload {
     /// appropriate environment variable, use it instead.
     fn download_url(
         &self,
-        python_install_mirror: Option<String>,
-        pypy_install_mirror: Option<String>,
+        python_install_mirror: Option<&str>,
+        pypy_install_mirror: Option<&str>,
     ) -> Result<Url, Error> {
         match self.key.implementation {
             LenientImplementationName::Known(ImplementationName::CPython) => {

diff --git a/crates/uv-python/src/installation.rs b/crates/uv-python/src/installation.rs
@@ -86,8 +86,8 @@ impl PythonInstallation {
         client_builder: &BaseClientBuilder<'a>,
         cache: &Cache,
         reporter: Option<&dyn Reporter>,
-        python_install_mirror: Option<String>,
-        pypy_install_mirror: Option<String>,
+        python_install_mirror: Option<&str>,
+        pypy_install_mirror: Option<&str>,
     ) -> Result<Self, Error> {
         let request = request.unwrap_or_else(|| &PythonRequest::Default);
 
@@ -132,8 +132,8 @@ impl PythonInstallation {
         client_builder: &BaseClientBuilder<'a>,
         cache: &Cache,
         reporter: Option<&dyn Reporter>,
-        python_install_mirror: Option<String>,
-        pypy_install_mirror: Option<String>,
+        python_install_mirror: Option<&str>,
+        pypy_install_mirror: Option<&str>,
     ) -> Result<Self, Error> {
         let installations = ManagedPythonInstallations::from_settings()?.init()?;
         let installations_dir = installations.root();
@@ -145,7 +145,7 @@ impl PythonInstallation {
 
         info!("Fetching requested Python...");
         let result = download
-            .fetch(
+            .fetch_with_retry(
                 &client,
                 installations_dir,
                 &cache_dir,

diff --git a/crates/uv/src/commands/build_frontend.rs b/crates/uv/src/commands/build_frontend.rs
@@ -430,8 +430,8 @@ async fn build_package(
         client_builder,
         cache,
         Some(&PythonDownloadReporter::single(printer)),
-        install_mirrors.python_install_mirror,
-        install_mirrors.pypy_install_mirror,
+        install_mirrors.python_install_mirror.as_deref(),
+        install_mirrors.pypy_install_mirror.as_deref(),
     )
     .await?
     .into_interpreter();

diff --git a/crates/uv/src/commands/project/add.rs b/crates/uv/src/commands/project/add.rs
@@ -176,8 +176,8 @@ pub(crate) async fn add(
             &client_builder,
             cache,
             Some(&reporter),
-            install_mirrors.python_install_mirror,
-            install_mirrors.pypy_install_mirror,
+            install_mirrors.python_install_mirror.as_deref(),
+            install_mirrors.pypy_install_mirror.as_deref(),
         )
         .await?
         .into_interpreter();

diff --git a/crates/uv/src/commands/project/init.rs b/crates/uv/src/commands/project/init.rs
@@ -423,8 +423,8 @@ async fn init_project(
                         &client_builder,
                         cache,
                         Some(&reporter),
-                        install_mirrors.python_install_mirror,
-                        install_mirrors.pypy_install_mirror,
+                        install_mirrors.python_install_mirror.as_deref(),
+                        install_mirrors.pypy_install_mirror.as_deref(),
                     )
                     .await?
                     .into_interpreter();
@@ -447,8 +447,8 @@ async fn init_project(
                     &client_builder,
                     cache,
                     Some(&reporter),
-                    install_mirrors.python_install_mirror,
-                    install_mirrors.pypy_install_mirror,
+                    install_mirrors.python_install_mirror.as_deref(),
+                    install_mirrors.pypy_install_mirror.as_deref(),
                 )
                 .await?
                 .into_interpreter();
@@ -509,8 +509,8 @@ async fn init_project(
                 &client_builder,
                 cache,
                 Some(&reporter),
-                install_mirrors.python_install_mirror,
-                install_mirrors.pypy_install_mirror,
+                install_mirrors.python_install_mirror.as_deref(),
+                install_mirrors.pypy_install_mirror.as_deref(),
             )
             .await?
             .into_interpreter();
@@ -533,8 +533,8 @@ async fn init_project(
             &client_builder,
             cache,
             Some(&reporter),
-            install_mirrors.python_install_mirror,
-            install_mirrors.pypy_install_mirror,
+            install_mirrors.python_install_mirror.as_deref(),
+            install_mirrors.pypy_install_mirror.as_deref(),
         )
         .await?
         .into_interpreter();

diff --git a/crates/uv/src/commands/project/mod.rs b/crates/uv/src/commands/project/mod.rs
@@ -650,8 +650,8 @@ impl ProjectInterpreter {
             &client_builder,
             cache,
             Some(&reporter),
-            install_mirrors.python_install_mirror,
-            install_mirrors.pypy_install_mirror,
+            install_mirrors.python_install_mirror.as_deref(),
+            install_mirrors.pypy_install_mirror.as_deref(),
         )
         .await?;
 
@@ -1551,8 +1551,8 @@ pub(crate) async fn init_script_python_requirement(
         client_builder,
         cache,
         Some(reporter),
-        install_mirrors.python_install_mirror,
-        install_mirrors.pypy_install_mirror,
+        install_mirrors.python_install_mirror.as_deref(),
+        install_mirrors.pypy_install_mirror.as_deref(),
     )
     .await?
     .into_interpreter();

diff --git a/crates/uv/src/commands/project/run.rs b/crates/uv/src/commands/project/run.rs
@@ -203,8 +203,8 @@ pub(crate) async fn run(
             &client_builder,
             cache,
             Some(&download_reporter),
-            install_mirrors.python_install_mirror.clone(),
-            install_mirrors.pypy_install_mirror.clone(),
+            install_mirrors.python_install_mirror.as_deref(),
+            install_mirrors.pypy_install_mirror.as_deref(),
         )
         .await?
         .into_interpreter();
@@ -555,8 +555,8 @@ pub(crate) async fn run(
                     &client_builder,
                     cache,
                     Some(&download_reporter),
-                    install_mirrors.python_install_mirror,
-                    install_mirrors.pypy_install_mirror,
+                    install_mirrors.python_install_mirror.as_deref(),
+                    install_mirrors.pypy_install_mirror.as_deref(),
                 )
                 .await?
                 .into_interpreter();
@@ -784,8 +784,8 @@ pub(crate) async fn run(
                     &client_builder,
                     cache,
                     Some(&download_reporter),
-                    install_mirrors.python_install_mirror,
-                    install_mirrors.pypy_install_mirror,
+                    install_mirrors.python_install_mirror.as_deref(),
+                    install_mirrors.pypy_install_mirror.as_deref(),
                 )
                 .await?;
 

diff --git a/crates/uv/src/commands/python/install.rs b/crates/uv/src/commands/python/install.rs
@@ -231,13 +231,13 @@ pub(crate) async fn install(
             (
                 download.key(),
                 download
-                    .fetch(
+                    .fetch_with_retry(
                         &client,
                         installations_dir,
                         &cache_dir,
                         reinstall,
-                        python_install_mirror.clone(),
-                        pypy_install_mirror.clone(),
+                        python_install_mirror.as_deref(),
+                        pypy_install_mirror.as_deref(),
                         Some(&reporter),
                     )
                     .await,

diff --git a/crates/uv/src/commands/tool/install.rs b/crates/uv/src/commands/tool/install.rs
@@ -75,8 +75,8 @@ pub(crate) async fn install(
         &client_builder,
         &cache,
         Some(&reporter),
-        install_mirrors.python_install_mirror,
-        install_mirrors.pypy_install_mirror,
+        install_mirrors.python_install_mirror.as_deref(),
+        install_mirrors.pypy_install_mirror.as_deref(),
     )
     .await?
     .into_interpreter();

diff --git a/crates/uv/src/commands/tool/run.rs b/crates/uv/src/commands/tool/run.rs
@@ -459,8 +459,8 @@ async fn get_or_create_environment(
         &client_builder,
         cache,
         Some(&reporter),
-        install_mirrors.python_install_mirror,
-        install_mirrors.pypy_install_mirror,
+        install_mirrors.python_install_mirror.as_deref(),
+        install_mirrors.pypy_install_mirror.as_deref(),
     )
     .await?
     .into_interpreter();

diff --git a/crates/uv/src/commands/tool/upgrade.rs b/crates/uv/src/commands/tool/upgrade.rs
@@ -85,8 +85,8 @@ pub(crate) async fn upgrade(
                 &client_builder,
                 cache,
                 Some(&reporter),
-                install_mirrors.python_install_mirror,
-                install_mirrors.pypy_install_mirror,
+                install_mirrors.python_install_mirror.as_deref(),
+                install_mirrors.pypy_install_mirror.as_deref(),
             )
             .await?
             .into_interpreter(),

diff --git a/crates/uv/src/commands/venv.rs b/crates/uv/src/commands/venv.rs
@@ -209,8 +209,8 @@ async fn venv_impl(
         &client_builder,
         cache,
         Some(&reporter),
-        install_mirrors.python_install_mirror,
-        install_mirrors.pypy_install_mirror,
+        install_mirrors.python_install_mirror.as_deref(),
+        install_mirrors.pypy_install_mirror.as_deref(),
     )
     .await
     .into_diagnostic()?;