Skip to content

Commit

Permalink
Add retries to uv publish
Browse files Browse the repository at this point in the history
  • Loading branch information
konstin committed Sep 24, 2024
1 parent c053dc8 commit 9054af4
Show file tree
Hide file tree
Showing 6 changed files with 50 additions and 30 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

6 changes: 4 additions & 2 deletions crates/uv-client/src/base_client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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<Response, reqwest_middleware::Error>) -> Option<Retryable> {
Expand Down
2 changes: 1 addition & 1 deletion crates/uv-client/src/lib.rs
Original file line number Diff line number Diff line change
@@ -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};
Expand Down
1 change: 1 addition & 0 deletions crates/uv-publish/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Expand Down
65 changes: 40 additions & 25 deletions crates/uv-publish/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand All @@ -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;

Expand Down Expand Up @@ -214,38 +215,52 @@ 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<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,)
.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.
Expand Down Expand Up @@ -406,12 +421,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<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?;
Expand Down Expand Up @@ -629,7 +644,7 @@ mod tests {
&BaseClientBuilder::new().build(),
Some("ferris"),
Some("F3RR!S"),
form_metadata,
&form_metadata,
Arc::new(DummyReporter),
)
.await
Expand Down Expand Up @@ -772,7 +787,7 @@ mod tests {
&BaseClientBuilder::new().build(),
Some("ferris"),
Some("F3RR!S"),
form_metadata,
&form_metadata,
Arc::new(DummyReporter),
)
.await
Expand Down
5 changes: 3 additions & 2 deletions crates/uv/src/commands/publish.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};

Expand All @@ -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)
Expand All @@ -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
Expand Down

0 comments on commit 9054af4

Please sign in to comment.