diff --git a/Cargo.lock b/Cargo.lock index 8b2df3bfa7280..65a58ba875539 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5020,6 +5020,7 @@ dependencies = [ "pypi-types", "reqwest", "reqwest-middleware", + "reqwest-retry", "rustc-hash", "serde", "serde_json", diff --git a/crates/uv-client/src/base_client.rs b/crates/uv-client/src/base_client.rs index c9a2b5d596a59..e2416e633d665 100644 --- a/crates/uv-client/src/base_client.rs +++ b/crates/uv-client/src/base_client.rs @@ -25,6 +25,8 @@ use crate::middleware::OfflineMiddleware; use crate::tls::read_identity; use crate::Connectivity; +pub const DEFAULT_RETRIES: u32 = 3; + /// A builder for an [`BaseClient`]. #[derive(Debug, Clone)] pub struct BaseClientBuilder<'a> { @@ -52,7 +54,7 @@ impl BaseClientBuilder<'_> { allow_insecure_host: vec![], native_tls: false, connectivity: Connectivity::Online, - retries: 3, + retries: DEFAULT_RETRIES, client: None, markers: None, platform: None, @@ -322,7 +324,7 @@ impl BaseClient { } /// Extends [`DefaultRetryableStrategy`], to log transient request failures and additional retry cases. -struct UvRetryableStrategy; +pub struct UvRetryableStrategy; impl RetryableStrategy for UvRetryableStrategy { fn handle(&self, res: &Result) -> Option { diff --git a/crates/uv-client/src/lib.rs b/crates/uv-client/src/lib.rs index dce88b0acccf0..0099583564718 100644 --- a/crates/uv-client/src/lib.rs +++ b/crates/uv-client/src/lib.rs @@ -1,4 +1,4 @@ -pub use base_client::{BaseClient, BaseClientBuilder}; +pub use base_client::{BaseClient, BaseClientBuilder, UvRetryableStrategy, DEFAULT_RETRIES}; pub use cached_client::{CacheControl, CachedClient, CachedClientError, DataWithCachePolicy}; pub use error::{Error, ErrorKind, WrappedReqwestError}; pub use flat_index::{FlatIndexClient, FlatIndexEntries, FlatIndexError}; diff --git a/crates/uv-publish/Cargo.toml b/crates/uv-publish/Cargo.toml index 20224a4a2bc6d..606b3d22ed202 100644 --- a/crates/uv-publish/Cargo.toml +++ b/crates/uv-publish/Cargo.toml @@ -25,6 +25,7 @@ itertools = { workspace = true } krata-tokio-tar = { workspace = true } reqwest = { workspace = true } reqwest-middleware = { workspace = true } +reqwest-retry = { workspace = true } rustc-hash = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } diff --git a/crates/uv-publish/src/lib.rs b/crates/uv-publish/src/lib.rs index cbe0e9c34e2de..c9270f2a041d7 100644 --- a/crates/uv-publish/src/lib.rs +++ b/crates/uv-publish/src/lib.rs @@ -10,6 +10,7 @@ use reqwest::header::AUTHORIZATION; use reqwest::multipart::Part; use reqwest::{Body, Response, StatusCode}; use reqwest_middleware::RequestBuilder; +use reqwest_retry::{Retryable, RetryableStrategy}; use rustc_hash::FxHashSet; use serde::Deserialize; use sha2::{Digest, Sha256}; @@ -20,9 +21,9 @@ use std::{fmt, io}; use thiserror::Error; use tokio::io::AsyncReadExt; use tokio_util::io::ReaderStream; -use tracing::{debug, enabled, trace, Level}; +use tracing::{debug, enabled, trace, warn, Level}; use url::Url; -use uv_client::BaseClient; +use uv_client::{BaseClient, UvRetryableStrategy}; use uv_fs::{ProgressReader, Simplified}; use uv_metadata::read_metadata_async_seek; @@ -214,11 +215,14 @@ pub fn files_for_publishing( /// Upload a file to a registry. /// /// Returns `true` if the file was newly uploaded and `false` if it already existed. +/// +/// Implements a custom retry flow since the request isn't cloneable. pub async fn upload( file: &Path, filename: &DistFilename, registry: &Url, client: &BaseClient, + retries: u32, username: Option<&str>, password: Option<&str>, reporter: Arc, @@ -226,26 +230,38 @@ pub async fn upload( let form_metadata = form_metadata(file, filename) .await .map_err(|err| PublishError::PublishPrepare(file.to_path_buf(), Box::new(err)))?; - let request = build_request( - file, - filename, - registry, - client, - username, - password, - form_metadata, - reporter, - ) - .await - .map_err(|err| PublishError::PublishPrepare(file.to_path_buf(), Box::new(err)))?; - - let response = request.send().await.map_err(|err| { - PublishError::PublishSend(file.to_path_buf(), registry.clone(), err.into()) - })?; - - handle_response(registry, response) + + // Retry loop + let mut attempt = 0; + loop { + attempt += 1; + let request = build_request( + file, + filename, + registry, + client, + username, + password, + &form_metadata, + reporter.clone(), + ) .await - .map_err(|err| PublishError::PublishSend(file.to_path_buf(), registry.clone(), err)) + .map_err(|err| PublishError::PublishPrepare(file.to_path_buf(), Box::new(err)))?; + + let result = request.send().await; + if attempt <= retries && UvRetryableStrategy.handle(&result) == Some(Retryable::Transient) { + warn!("Transient request failure for {}, retrying", registry); + continue; + } + + let response = result.map_err(|err| { + PublishError::PublishSend(file.to_path_buf(), registry.clone(), err.into()) + })?; + + return handle_response(registry, response) + .await + .map_err(|err| PublishError::PublishSend(file.to_path_buf(), registry.clone(), err)); + } } /// Calculate the SHA256 of a file. @@ -406,12 +422,12 @@ async fn build_request( client: &BaseClient, username: Option<&str>, password: Option<&str>, - form_metadata: Vec<(&'static str, String)>, + form_metadata: &[(&'static str, String)], reporter: Arc, ) -> Result { let mut form = reqwest::multipart::Form::new(); for (key, value) in form_metadata { - form = form.text(key, value); + form = form.text(*key, value.clone()); } let file = fs_err::tokio::File::open(file).await?; @@ -629,7 +645,7 @@ mod tests { &BaseClientBuilder::new().build(), Some("ferris"), Some("F3RR!S"), - form_metadata, + &form_metadata, Arc::new(DummyReporter), ) .await @@ -772,7 +788,7 @@ mod tests { &BaseClientBuilder::new().build(), Some("ferris"), Some("F3RR!S"), - form_metadata, + &form_metadata, Arc::new(DummyReporter), ) .await diff --git a/crates/uv/src/commands/publish.rs b/crates/uv/src/commands/publish.rs index b675721ce61b0..d55c3f5179d66 100644 --- a/crates/uv/src/commands/publish.rs +++ b/crates/uv/src/commands/publish.rs @@ -7,7 +7,7 @@ use std::fmt::Write; use std::sync::Arc; use tracing::info; use url::Url; -use uv_client::{BaseClientBuilder, Connectivity}; +use uv_client::{BaseClientBuilder, Connectivity, DEFAULT_RETRIES}; use uv_configuration::{KeyringProviderType, TrustedHost}; use uv_publish::{files_for_publishing, upload}; @@ -34,7 +34,7 @@ pub(crate) async fn publish( } let client = BaseClientBuilder::new() - // Don't try cloning the request for retries. + // Don't try cloning the request for retries, we roll our own retry loop // https://github.com/seanmonstar/reqwest/issues/2416 .retries(0) .keyring(keyring_provider) @@ -59,6 +59,7 @@ pub(crate) async fn publish( &filename, &publish_url, &client, + DEFAULT_RETRIES, username.as_deref(), password.as_deref(), // Needs to be an `Arc` because the reqwest `Body` static lifetime requirement