Skip to content

Commit

Permalink
Use a consistent buffer size when writing out zip files (#5570)
Browse files Browse the repository at this point in the history
  • Loading branch information
charliermarsh committed Jul 29, 2024
1 parent 05b1f51 commit f70501a
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 8 deletions.
12 changes: 6 additions & 6 deletions crates/uv-extract/src/stream.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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?;
}
Expand Down
11 changes: 9 additions & 2 deletions crates/uv-extract/src/sync.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down Expand Up @@ -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.
Expand Down

0 comments on commit f70501a

Please sign in to comment.