From f70501a22e4468ec87b073ffa0b5dc96c6de4edb Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Mon, 29 Jul 2024 14:45:33 -0400
Subject: [PATCH] Use a consistent buffer size when writing out zip files
 (#5570)

---
 crates/uv-extract/src/stream.rs | 12 ++++++------
 crates/uv-extract/src/sync.rs   | 11 +++++++++--
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/crates/uv-extract/src/stream.rs b/crates/uv-extract/src/stream.rs
index 31ca74134..f89767e93 100644
--- a/crates/uv-extract/src/stream.rs
+++ b/crates/uv-extract/src/stream.rs
@@ -45,12 +45,12 @@ pub async fn unzip<R: tokio::io::AsyncRead + Unpin>(
 
             // We don't know the file permissions here, because we haven't seen the central directory yet.
             let file = fs_err::tokio::File::create(&path).await?;
-            let mut writer =
-                if let Ok(size) = usize::try_from(entry.reader().entry().uncompressed_size()) {
-                    tokio::io::BufWriter::with_capacity(size, file)
-                } else {
-                    tokio::io::BufWriter::new(file)
-                };
+            let size = entry.reader().entry().uncompressed_size();
+            let mut writer = if let Ok(size) = usize::try_from(size) {
+                tokio::io::BufWriter::with_capacity(std::cmp::min(size, 1024 * 1024), file)
+            } else {
+                tokio::io::BufWriter::new(file)
+            };
             let mut reader = entry.reader_mut().compat();
             tokio::io::copy(&mut reader, &mut writer).await?;
         }
diff --git a/crates/uv-extract/src/sync.rs b/crates/uv-extract/src/sync.rs
index 07dfd1a63..bcf13a5a0 100644
--- a/crates/uv-extract/src/sync.rs
+++ b/crates/uv-extract/src/sync.rs
@@ -14,6 +14,7 @@ pub fn unzip<R: Send + std::io::Read + std::io::Seek + HasLength>(
     target: &Path,
 ) -> Result<(), Error> {
     // Unzip in parallel.
+    let reader = std::io::BufReader::new(reader);
     let archive = ZipArchive::new(CloneableSeekableReader::new(reader))?;
     let directories = Mutex::new(FxHashSet::default());
     (0..archive.len())
@@ -45,8 +46,14 @@ pub fn unzip<R: Send + std::io::Read + std::io::Seek + HasLength>(
             }
 
             // Copy the file contents.
-            let mut outfile = fs_err::File::create(&path)?;
-            std::io::copy(&mut file, &mut outfile)?;
+            let outfile = fs_err::File::create(&path)?;
+            let size = file.size();
+            let mut writer = if let Ok(size) = usize::try_from(size) {
+                std::io::BufWriter::with_capacity(std::cmp::min(size, 1024 * 1024), outfile)
+            } else {
+                std::io::BufWriter::new(outfile)
+            };
+            std::io::copy(&mut file, &mut writer)?;
 
             // See `uv_extract::stream::unzip`. For simplicity, this is identical with the code there except for being
             // sync.