From 1b13036674c1da241491e418c54878ad67582d51 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 20 Nov 2024 09:42:42 -0500 Subject: [PATCH] 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 https://github.com/astral-sh/uv/issues/8525. --- Cargo.lock | 1 + crates/uv-client/src/base_client.rs | 7 +-- crates/uv-client/src/lib.rs | 3 +- crates/uv-python/Cargo.toml | 1 + crates/uv-python/src/downloads.rs | 73 +++++++++++++++++++++--- crates/uv-python/src/installation.rs | 10 ++-- crates/uv/src/commands/build_frontend.rs | 4 +- crates/uv/src/commands/project/add.rs | 4 +- crates/uv/src/commands/project/init.rs | 16 +++--- crates/uv/src/commands/project/mod.rs | 8 +-- crates/uv/src/commands/project/run.rs | 12 ++-- crates/uv/src/commands/python/install.rs | 6 +- crates/uv/src/commands/tool/install.rs | 4 +- crates/uv/src/commands/tool/run.rs | 4 +- crates/uv/src/commands/tool/upgrade.rs | 4 +- crates/uv/src/commands/venv.rs | 4 +- 16 files changed, 107 insertions(+), 54 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d1a15f1a0063..60ef3945991f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5297,6 +5297,7 @@ dependencies = [ "regex", "reqwest", "reqwest-middleware", + "reqwest-retry", "rmp-serde", "same-file", "schemars", diff --git a/crates/uv-client/src/base_client.rs b/crates/uv-client/src/base_client.rs index f4ade9787976..d9acc8964163 100644 --- 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::(&err) { diff --git a/crates/uv-client/src/lib.rs b/crates/uv-client/src/lib.rs index bd496aac1213..54b7e9fa7785 100644 --- 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 index a1cb0b509e09..3b8384f6baff 100644 --- 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 index d19018f8f374..80202f526893 100644 --- 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,7 +463,57 @@ 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 { + 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, @@ -466,8 +521,8 @@ impl ManagedPythonDownload { installation_dir: &Path, cache_dir: &Path, reinstall: bool, - python_install_mirror: Option, - pypy_install_mirror: Option, + python_install_mirror: Option<&str>, + pypy_install_mirror: Option<&str>, reporter: Option<&dyn Reporter>, ) -> Result { 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, - pypy_install_mirror: Option, + python_install_mirror: Option<&str>, + pypy_install_mirror: Option<&str>, ) -> Result { match self.key.implementation { LenientImplementationName::Known(ImplementationName::CPython) => { diff --git a/crates/uv-python/src/installation.rs b/crates/uv-python/src/installation.rs index d507992be874..b1da7c252242 100644 --- 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, - pypy_install_mirror: Option, + python_install_mirror: Option<&str>, + pypy_install_mirror: Option<&str>, ) -> Result { 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, - pypy_install_mirror: Option, + python_install_mirror: Option<&str>, + pypy_install_mirror: Option<&str>, ) -> Result { 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 index 3e8c89bfb0d6..8067c23f6671 100644 --- 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 index 40fd201431ab..4886f34ea778 100644 --- 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 index a9427e569e74..57fe0832faa0 100644 --- 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 index 516fe62f5b45..505963645557 100644 --- 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 index 1c3553a0073a..ad5271f43921 100644 --- 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 index 2059b28096df..52cc529768e6 100644 --- 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 index e751b9682360..f4d0a87363b3 100644 --- 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 index 297fea8b5465..b6b76bc8ab95 100644 --- 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 index 47b997aae731..6641195e8e33 100644 --- 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 index cc68cbb0266d..93cb6af415ae 100644 --- 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()?;