Add retries to uv publish (#7635)

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<Response, reqwest_middleware::Error>) -> Option<Retryable> {

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};

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 }

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,38 +275,53 @@ 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<impl Reporter>,
 ) -> Result<bool, PublishError> {
     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<impl Reporter>,
 ) -> Result<RequestBuilder, PublishPrepareError> {
     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

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