From 8a30e4f7e12a7edae0890950ad00b3cec1b119e6 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Wed, 20 Nov 2024 09:17:04 -0500
Subject: [PATCH] Use exponential backoff for publish retries

---
 crates/uv-publish/src/lib.rs      | 31 +++++++++++++++++++++----------
 crates/uv/src/commands/publish.rs |  5 +----
 2 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/crates/uv-publish/src/lib.rs b/crates/uv-publish/src/lib.rs
index 4ecc17f89a863..db68af0bdc187 100644
--- a/crates/uv-publish/src/lib.rs
+++ b/crates/uv-publish/src/lib.rs
@@ -11,18 +11,22 @@ use reqwest::header::AUTHORIZATION;
 use reqwest::multipart::Part;
 use reqwest::{Body, Response, StatusCode};
 use reqwest_middleware::RequestBuilder;
-use reqwest_retry::{Retryable, RetryableStrategy};
+use reqwest_retry::{RetryPolicy, Retryable, RetryableStrategy};
 use rustc_hash::FxHashSet;
 use serde::Deserialize;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
+use std::time::{Duration, SystemTime};
 use std::{env, fmt, io};
 use thiserror::Error;
 use tokio::io::{AsyncReadExt, BufReader};
 use tokio_util::io::ReaderStream;
 use tracing::{debug, enabled, trace, Level};
 use url::Url;
-use uv_client::{BaseClient, OwnedArchive, RegistryClientBuilder, UvRetryableStrategy};
+use uv_client::{
+    BaseClient, OwnedArchive, RegistryClientBuilder,
+    UvRetryableStrategy,
+};
 use uv_configuration::{KeyringProviderType, TrustedPublishing};
 use uv_distribution_filename::{DistFilename, SourceDistExtension, SourceDistFilename};
 use uv_fs::{ProgressReader, Simplified};
@@ -365,7 +369,6 @@ pub async fn upload(
     filename: &DistFilename,
     registry: &Url,
     client: &BaseClient,
-    retries: u32,
     username: Option<&str>,
     password: Option<&str>,
     check_url_client: Option<&CheckUrlClient<'_>>,
@@ -375,8 +378,9 @@ pub async fn upload(
         .await
         .map_err(|err| PublishError::PublishPrepare(file.to_path_buf(), Box::new(err)))?;
 
-    // Retry loop
-    let mut attempt = 0;
+    let mut n_past_retries = 0;
+    let start_time = SystemTime::now();
+    let retry_policy = client.retry_policy();
     loop {
         let (request, idx) = build_request(
             file,
@@ -393,11 +397,18 @@ pub async fn upload(
         .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) {
-            reporter.on_download_complete(idx);
-            warn_user!("Transient request failure for {}, retrying", registry);
-            attempt += 1;
-            continue;
+        if UvRetryableStrategy.handle(&result) == Some(Retryable::Transient) {
+            let retry_decision = retry_policy.should_retry(start_time, n_past_retries);
+            if let reqwest_retry::RetryDecision::Retry { execute_after } = retry_decision {
+                warn_user!("Transient failure while handling response for {registry}; retrying...",);
+                reporter.on_download_complete(idx);
+                let duration = execute_after
+                    .duration_since(SystemTime::now())
+                    .unwrap_or_else(|_| Duration::default());
+                tokio::time::sleep(duration).await;
+                n_past_retries += 1;
+                continue;
+            }
         }
 
         let response = result.map_err(|err| {
diff --git a/crates/uv/src/commands/publish.rs b/crates/uv/src/commands/publish.rs
index 3d37a4fe52104..71bc2ef9af1bb 100644
--- a/crates/uv/src/commands/publish.rs
+++ b/crates/uv/src/commands/publish.rs
@@ -11,9 +11,7 @@ use std::time::Duration;
 use tracing::info;
 use url::Url;
 use uv_cache::Cache;
-use uv_client::{
-    AuthIntegration, BaseClientBuilder, Connectivity, RegistryClientBuilder, DEFAULT_RETRIES,
-};
+use uv_client::{AuthIntegration, BaseClientBuilder, Connectivity, RegistryClientBuilder};
 use uv_configuration::{KeyringProviderType, TrustedHost, TrustedPublishing};
 use uv_distribution_types::{Index, IndexCapabilities, IndexLocations, IndexUrl};
 use uv_publish::{
@@ -176,7 +174,6 @@ pub(crate) async fn publish(
             &filename,
             &publish_url,
             &upload_client,
-            DEFAULT_RETRIES,
             username.as_deref(),
             password.as_deref(),
             check_url_client.as_ref(),