From d14fbb0bd8e2ee3eebf6456a75343284ee83b22d Mon Sep 17 00:00:00 2001 From: konstin Date: Mon, 23 Sep 2024 10:28:48 +0200 Subject: [PATCH] Add retries to `uv publish` --- Cargo.lock | 1 + crates/uv-client/src/base_client.rs | 6 ++- crates/uv-client/src/lib.rs | 4 +- crates/uv-publish/Cargo.toml | 1 + crates/uv-publish/src/lib.rs | 65 ++++++++++++++++++----------- crates/uv/src/commands/publish.rs | 7 ++-- 6 files changed, 54 insertions(+), 30 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f15a04f23c44..754faa03db44 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 4dc6ecde7415..9a5085fcafb3 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; + /// Selectively skip parts or the entire auth middleware. #[derive(Debug, Clone, Copy, Default)] pub enum AuthIntegration { @@ -65,7 +67,7 @@ impl BaseClientBuilder<'_> { allow_insecure_host: vec![], native_tls: false, connectivity: Connectivity::Online, - retries: 3, + retries: DEFAULT_RETRIES, client: None, markers: None, platform: None, @@ -374,7 +376,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 fadf86aed78f..bd496aac1213 100644 --- a/crates/uv-client/src/lib.rs +++ b/crates/uv-client/src/lib.rs @@ -1,4 +1,6 @@ -pub use base_client::{AuthIntegration, BaseClient, BaseClientBuilder}; +pub use base_client::{ + AuthIntegration, 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 377dce0b4ba2..ba8b6a8ee2ea 100644 --- a/crates/uv-publish/Cargo.toml +++ b/crates/uv-publish/Cargo.toml @@ -27,6 +27,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 5b8ede6d1c96..a2402e1319c5 100644 --- a/crates/uv-publish/src/lib.rs +++ b/crates/uv-publish/src/lib.rs @@ -13,6 +13,7 @@ use reqwest::header::AUTHORIZATION; use reqwest::multipart::Part; use reqwest::{Body, Response, StatusCode}; use reqwest_middleware::{ClientWithMiddleware, RequestBuilder}; +use reqwest_retry::{Retryable, RetryableStrategy}; use rustc_hash::FxHashSet; use serde::Deserialize; use sha2::{Digest, Sha256}; @@ -23,8 +24,9 @@ use std::{env, 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::UvRetryableStrategy; use uv_configuration::{KeyringProviderType, TrustedPublishing}; use uv_fs::{ProgressReader, Simplified}; use uv_metadata::read_metadata_async_seek; @@ -273,11 +275,14 @@ pub async fn check_trusted_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: &ClientWithMiddleware, + retries: u32, username: Option<&str>, password: Option<&str>, reporter: Arc, @@ -285,26 +290,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. @@ -465,12 +482,12 @@ async fn build_request( client: &ClientWithMiddleware, 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?; @@ -687,7 +704,7 @@ mod tests { &BaseClientBuilder::new().build().client(), Some("ferris"), Some("F3RR!S"), - form_metadata, + &form_metadata, Arc::new(DummyReporter), ) .await @@ -830,7 +847,7 @@ mod tests { &BaseClientBuilder::new().build().client(), 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 ca52af6f8cc7..50f7f3c40e1e 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::{AuthIntegration, BaseClientBuilder, Connectivity}; +use uv_client::{AuthIntegration, BaseClientBuilder, Connectivity, DEFAULT_RETRIES}; use uv_configuration::{KeyringProviderType, TrustedHost, TrustedPublishing}; use uv_publish::{check_trusted_publishing, files_for_publishing, upload}; @@ -34,9 +34,9 @@ pub(crate) async fn publish( n => writeln!(printer.stderr(), "Publishing {n} files {publish_url}")?, } - // * For the uploads themselves, we can't use retries due to + // * For the uploads themselves, we roll our own retries due to // https://github.com/seanmonstar/reqwest/issues/2416, but for trusted publishing, we want - // retires. + // the default retries. // * We want to allow configuring TLS for the registry, while for trusted publishing we know the // defaults are correct. // * For the uploads themselves, we know we need an authorization header and we can't nor @@ -86,6 +86,7 @@ pub(crate) async fn publish( &filename, &publish_url, &upload_client.client(), + DEFAULT_RETRIES, username.as_deref(), password.as_deref(), // Needs to be an `Arc` because the reqwest `Body` static lifetime requirement