From 38504dcaeec54cf497bb6cfd61f58770d1d0f143 Mon Sep 17 00:00:00 2001
From: messense <messense@icloud.com>
Date: Tue, 16 Jul 2024 21:00:46 +0800
Subject: [PATCH] Download wheel to disk when streaming unzip failed with HTTP
 streaming error (#5094)

## Summary

Workaround the `stream_wheel` not retry issue
[found](https://github.com/astral-sh/uv/issues/3514#issuecomment-2229820667)
in #3514, it's not a perfect solution but I think it's acceptable
because the error should not occur frequently.

## Test Plan

Manually using `iptables -A OUTPUT -p tcp -dport 3128 -j REJECT
--reject-with tcp-reset` to inject connection reset error to the HTTP
proxy that proxies PyPI requests.

```
error: Failed to prepare distributions
  Caused by: Failed to fetch wheel: piqp==0.4.1
  Caused by: Request failed after 3 retries
  Caused by: error sending request for url (https://files.pythonhosted.org/packages/94/4d/09ade94dfda5b57c1ca43564541871bd1a0d89dfd3c368ac505b6ca09831/piqp-0.4.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl)
  Caused by: client error (Connect)
  Caused by: tcp connect error: Connection refused (os error 111)
  Caused by: Connection refused (os error 111)
```
---
 Cargo.lock                                        |  1 +
 .../uv-distribution/src/distribution_database.rs  | 14 ++++++++++----
 crates/uv-extract/Cargo.toml                      |  1 +
 crates/uv-extract/src/error.rs                    | 15 +++++++++++++++
 4 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 5422a804b..4048247b1 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4804,6 +4804,7 @@ dependencies = [
  "md-5",
  "pypi-types",
  "rayon",
+ "reqwest",
  "rustc-hash 2.0.0",
  "sha2",
  "thiserror",
diff --git a/crates/uv-distribution/src/distribution_database.rs b/crates/uv-distribution/src/distribution_database.rs
index 0ba87e723..ac8a81ebb 100644
--- a/crates/uv-distribution/src/distribution_database.rs
+++ b/crates/uv-distribution/src/distribution_database.rs
@@ -208,10 +208,16 @@ impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> {
                         hashes: archive.hashes,
                         filename: wheel.filename.clone(),
                     }),
-                    Err(Error::Extract(err)) if err.is_http_streaming_unsupported() => {
-                        warn!(
-                            "Streaming unsupported for {dist}; downloading wheel to disk ({err})"
-                        );
+                    Err(Error::Extract(err)) => {
+                        if err.is_http_streaming_unsupported() {
+                            warn!(
+                                "Streaming unsupported for {dist}; downloading wheel to disk ({err})"
+                            );
+                        } else if err.is_http_streaming_failed() {
+                            warn!("Streaming failed for {dist}; downloading wheel to disk ({err})");
+                        } else {
+                            return Err(Error::Extract(err));
+                        }
 
                         // If the request failed because streaming is unsupported, download the
                         // wheel directly.
diff --git a/crates/uv-extract/Cargo.toml b/crates/uv-extract/Cargo.toml
index f5c3c5eb4..234eee8f6 100644
--- a/crates/uv-extract/Cargo.toml
+++ b/crates/uv-extract/Cargo.toml
@@ -21,6 +21,7 @@ fs-err = { workspace = true, features = ["tokio"] }
 futures = { workspace = true }
 md-5.workspace = true
 rayon = { workspace = true }
+reqwest = { workspace = true }
 rustc-hash = { workspace = true }
 sha2 = { workspace = true }
 thiserror = { workspace = true }
diff --git a/crates/uv-extract/src/error.rs b/crates/uv-extract/src/error.rs
index 275caed53..3cfaeb35a 100644
--- a/crates/uv-extract/src/error.rs
+++ b/crates/uv-extract/src/error.rs
@@ -28,4 +28,19 @@ impl Error {
             Self::AsyncZip(async_zip::error::ZipError::FeatureNotSupported(_))
         )
     }
+
+    /// Returns `true` if the error is due to HTTP streaming request failed.
+    pub fn is_http_streaming_failed(&self) -> bool {
+        match self {
+            Self::AsyncZip(async_zip::error::ZipError::UpstreamReadError(_)) => true,
+            Self::Io(err) => {
+                if let Some(inner) = err.get_ref() {
+                    inner.downcast_ref::<reqwest::Error>().is_some()
+                } else {
+                    false
+                }
+            }
+            _ => false,
+        }
+    }
 }