From a6c001703c651d8d36ca38bc2b2198cbfaf35a78 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 | 2 +- crates/uv-publish/Cargo.toml | 1 + crates/uv-publish/src/lib.rs | 64 ++++++++++++++++++----------- crates/uv/src/commands/publish.rs | 5 ++- 6 files changed, 50 insertions(+), 29 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c707253d7e21a..370850cb55db5 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 6e186b978dd33..a6d5f51481e9d 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 76b00b942a4b2..d15ab5014ff5d 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}; @@ -18,9 +19,9 @@ use std::path::{Path, PathBuf}; use std::{fmt, io}; use thiserror::Error; use tokio::io::AsyncReadExt; -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::Simplified; use uv_metadata::read_metadata_async_seek; @@ -205,36 +206,51 @@ 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>, ) -> Result { 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, - ) - .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, + ) .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. @@ -395,11 +411,11 @@ async fn build_request( client: &BaseClient, username: Option<&str>, password: Option<&str>, - form_metadata: Vec<(&'static str, String)>, + form_metadata: &[(&'static str, String)], ) -> 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: tokio::fs::File = fs_err::tokio::File::open(file).await?.into(); @@ -601,7 +617,7 @@ mod tests { &BaseClientBuilder::new().build(), Some("ferris"), Some("F3RR!S"), - form_metadata, + &form_metadata, ) .await .unwrap(); @@ -743,7 +759,7 @@ mod tests { &BaseClientBuilder::new().build(), Some("ferris"), Some("F3RR!S"), - form_metadata, + &form_metadata, ) .await .unwrap(); diff --git a/crates/uv/src/commands/publish.rs b/crates/uv/src/commands/publish.rs index 84b5be894446f..5f01d6f12007e 100644 --- a/crates/uv/src/commands/publish.rs +++ b/crates/uv/src/commands/publish.rs @@ -5,7 +5,7 @@ use owo_colors::OwoColorize; use std::fmt::Write; 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}; @@ -32,7 +32,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) @@ -56,6 +56,7 @@ pub(crate) async fn publish( &filename, &publish_url, &client, + DEFAULT_RETRIES, username.as_deref(), password.as_deref(), )