Use a consistent buffer size for downloads (#5569)

astral-sh · Jul 29, 2024 · 05b1f51 · 05b1f51
1 parent 48162de
commit 05b1f51
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/crates/uv-extract/src/stream.rs b/crates/uv-extract/src/stream.rs
@@ -8,6 +8,8 @@ use tracing::warn;
 
 use crate::Error;
 
+const DEFAULT_BUF_SIZE: usize = 128 * 1024;
+
 /// Unzip a `.zip` archive into the target directory, without requiring `Seek`.
 ///
 /// This is useful for unzipping files as they're being downloaded. If the archive
@@ -18,7 +20,7 @@ pub async fn unzip<R: tokio::io::AsyncRead + Unpin>(
     target: impl AsRef<Path>,
 ) -> Result<(), Error> {
     let target = target.as_ref();
-    let mut reader = futures::io::BufReader::with_capacity(128 * 1024, reader.compat());
+    let mut reader = futures::io::BufReader::with_capacity(DEFAULT_BUF_SIZE, reader.compat());
     let mut zip = async_zip::base::read::stream::ZipFileReader::new(&mut reader);
 
     let mut directories = FxHashSet::default();
@@ -155,7 +157,7 @@ pub async fn untar_gz<R: tokio::io::AsyncRead + Unpin>(
     reader: R,
     target: impl AsRef<Path>,
 ) -> Result<(), Error> {
-    let reader = tokio::io::BufReader::new(reader);
+    let reader = tokio::io::BufReader::with_capacity(DEFAULT_BUF_SIZE, reader);
     let decompressed_bytes = async_compression::tokio::bufread::GzipDecoder::new(reader);
 
     let mut archive = tokio_tar::ArchiveBuilder::new(decompressed_bytes)
@@ -172,7 +174,7 @@ pub async fn untar_bz2<R: tokio::io::AsyncRead + Unpin>(
     reader: R,
     target: impl AsRef<Path>,
 ) -> Result<(), Error> {
-    let reader = tokio::io::BufReader::new(reader);
+    let reader = tokio::io::BufReader::with_capacity(DEFAULT_BUF_SIZE, reader);
     let decompressed_bytes = async_compression::tokio::bufread::BzDecoder::new(reader);
 
     let mut archive = tokio_tar::ArchiveBuilder::new(decompressed_bytes)
@@ -189,7 +191,7 @@ pub async fn untar_zst<R: tokio::io::AsyncRead + Unpin>(
     reader: R,
     target: impl AsRef<Path>,
 ) -> Result<(), Error> {
-    let reader = tokio::io::BufReader::new(reader);
+    let reader = tokio::io::BufReader::with_capacity(DEFAULT_BUF_SIZE, reader);
     let decompressed_bytes = async_compression::tokio::bufread::ZstdDecoder::new(reader);
 
     let mut archive = tokio_tar::ArchiveBuilder::new(decompressed_bytes)
@@ -205,7 +207,7 @@ pub async fn untar_xz<R: tokio::io::AsyncRead + Unpin>(
     reader: R,
     target: impl AsRef<Path>,
 ) -> Result<(), Error> {
-    let reader = tokio::io::BufReader::new(reader);
+    let reader = tokio::io::BufReader::with_capacity(DEFAULT_BUF_SIZE, reader);
     let decompressed_bytes = async_compression::tokio::bufread::XzDecoder::new(reader);
 
     let mut archive = tokio_tar::ArchiveBuilder::new(decompressed_bytes)