diff --git a/crates/puffin-distribution/src/distribution_database.rs b/crates/puffin-distribution/src/distribution_database.rs index 3baa5a9cc..e07d7dfea 100644 --- a/crates/puffin-distribution/src/distribution_database.rs +++ b/crates/puffin-distribution/src/distribution_database.rs @@ -4,7 +4,6 @@ use std::path::Path; use std::str::FromStr; use std::sync::Arc; -use bytesize::ByteSize; use fs_err::tokio as fs; use puffin_extract::unzip_no_seek; use thiserror::Error; @@ -25,9 +24,7 @@ use pypi_types::Metadata21; use crate::download::{BuiltWheel, UnzippedWheel}; use crate::locks::Locks; use crate::reporter::Facade; -use crate::{ - DiskWheel, InMemoryWheel, LocalWheel, Reporter, SourceDistCachedBuilder, SourceDistError, -}; +use crate::{DiskWheel, LocalWheel, Reporter, SourceDistCachedBuilder, SourceDistError}; #[derive(Debug, Error)] pub enum DistributionDatabaseError { @@ -124,11 +121,11 @@ impl<'a, Context: BuildContext + Send + Sync> DistributionDatabase<'a, Context> wheel_filename.stem(), ); - // Start the download - let reader = self.client.stream_external(&url).await?; - - // In all wheels we've seen so far, unzipping while downloading is the - // faster option. + // Download and unzip on the same tokio task + // + // In all wheels we've seen so far, unzipping while downloading is + // faster than downloading into a file and then unzipping on multiple + // threads. // // Writing to a file first may be faster if the wheel takes longer to // unzip than it takes to download. This may happen if the wheel is a @@ -139,89 +136,15 @@ impl<'a, Context: BuildContext + Send + Sync> DistributionDatabase<'a, Context> // for downloading and unzipping (with a buffer in between) and switch // to rayon if this buffer grows large by the time the file is fully // downloaded. - let unzip_while_downloading = true; - if unzip_while_downloading { - // Download and unzip to a temporary dir - let temp_dir = tempfile::tempdir_in(self.cache.root())?; - let temp_target = temp_dir.path().join(&wheel.file.filename); - unzip_no_seek(reader.compat(), &temp_target).await?; + let reader = self.client.stream_external(&url).await?; + let target = cache_entry.into_path_buf(); + unzip_no_seek(reader.compat(), &target).await?; - // Move the dir to the right place - fs::create_dir_all(&cache_entry.dir()).await?; - let target = cache_entry.into_path_buf(); - tokio::fs::rename(temp_target, &target).await?; - - return Ok(LocalWheel::Unzipped(UnzippedWheel { - dist: dist.clone(), - target, - filename: wheel_filename, - })); - } - - // If the file is greater than 5MB, write it to disk; otherwise, keep it in memory. - // - // TODO this is currently dead code. Consider deleting if there's no use for it. - let byte_size = wheel.file.size.map(ByteSize::b); - let local_wheel = if let Some(byte_size) = - byte_size.filter(|byte_size| *byte_size < ByteSize::mb(5)) - { - debug!("Fetching in-memory wheel from registry: {dist} ({byte_size})",); - - // Read into a buffer. - let mut buffer = Vec::with_capacity( - wheel - .file - .size - .unwrap_or(0) - .try_into() - .expect("5MB shouldn't be bigger usize::MAX"), - ); - let mut reader = tokio::io::BufReader::new(reader.compat()); - tokio::io::copy(&mut reader, &mut buffer).await?; - - LocalWheel::InMemory(InMemoryWheel { - dist: dist.clone(), - target: cache_entry.into_path_buf(), - buffer, - filename: wheel_filename, - }) - } else { - let size = - byte_size.map_or("unknown size".to_string(), |size| size.to_string()); - - debug!("Fetching disk-based wheel from registry: {dist} ({size})"); - - let filename = wheel_filename.to_string(); - - // Download the wheel to a temporary file. - let temp_dir = tempfile::tempdir_in(self.cache.root())?; - let temp_file = temp_dir.path().join(&filename); - let mut writer = - tokio::io::BufWriter::new(tokio::fs::File::create(&temp_file).await?); - tokio::io::copy(&mut reader.compat(), &mut writer).await?; - - // Move the temporary file to the cache. - let cache_entry = self.cache.entry( - CacheBucket::Wheels, - WheelCache::Index(&wheel.index) - .remote_wheel_dir(wheel_filename.name.as_ref()), - filename, // TODO should this be filename.stem() to match the other branch? - ); - fs::create_dir_all(&cache_entry.dir()).await?; - tokio::fs::rename(temp_file, &cache_entry.path()).await?; - - LocalWheel::Disk(DiskWheel { - dist: dist.clone(), - target: cache_entry - .with_file(wheel_filename.stem()) - .path() - .to_path_buf(), - path: cache_entry.into_path_buf(), - filename: wheel_filename, - }) - }; - - Ok(local_wheel) + Ok(LocalWheel::Unzipped(UnzippedWheel { + dist: dist.clone(), + target, + filename: wheel_filename, + })) } Dist::Built(BuiltDist::DirectUrl(wheel)) => { diff --git a/crates/puffin-distribution/src/download.rs b/crates/puffin-distribution/src/download.rs index b3ca4a98d..30f7f12c7 100644 --- a/crates/puffin-distribution/src/download.rs +++ b/crates/puffin-distribution/src/download.rs @@ -21,19 +21,6 @@ pub struct UnzippedWheel { pub(crate) target: PathBuf, } -/// A downloaded wheel that's stored in-memory. -#[derive(Debug, Clone)] -pub struct InMemoryWheel { - /// The remote distribution from which this wheel was downloaded. - pub(crate) dist: Dist, - /// The parsed filename. - pub(crate) filename: WheelFilename, - /// The contents of the wheel. - pub(crate) buffer: Vec, - /// The expected path to the downloaded wheel's entry in the cache. - pub(crate) target: PathBuf, -} - /// A downloaded wheel that's stored on-disk. #[derive(Debug, Clone)] pub struct DiskWheel { @@ -64,7 +51,6 @@ pub struct BuiltWheel { #[derive(Debug, Clone)] pub enum LocalWheel { Unzipped(UnzippedWheel), - InMemory(InMemoryWheel), Disk(DiskWheel), Built(BuiltWheel), } @@ -74,7 +60,6 @@ impl LocalWheel { pub fn target(&self) -> &Path { match self { LocalWheel::Unzipped(wheel) => &wheel.target, - LocalWheel::InMemory(wheel) => &wheel.target, LocalWheel::Disk(wheel) => &wheel.target, LocalWheel::Built(wheel) => &wheel.target, } @@ -84,7 +69,6 @@ impl LocalWheel { pub fn remote(&self) -> &Dist { match self { LocalWheel::Unzipped(wheel) => wheel.remote(), - LocalWheel::InMemory(wheel) => wheel.remote(), LocalWheel::Disk(wheel) => wheel.remote(), LocalWheel::Built(wheel) => wheel.remote(), } @@ -94,7 +78,6 @@ impl LocalWheel { pub fn filename(&self) -> &WheelFilename { match self { LocalWheel::Unzipped(wheel) => &wheel.filename, - LocalWheel::InMemory(wheel) => &wheel.filename, LocalWheel::Disk(wheel) => &wheel.filename, LocalWheel::Built(wheel) => &wheel.filename, } @@ -115,13 +98,6 @@ impl DiskWheel { } } -impl InMemoryWheel { - /// Return the [`Dist`] from which this wheel was downloaded. - pub fn remote(&self) -> &Dist { - &self.dist - } -} - impl BuiltWheel { /// Return the [`Dist`] from which this source distribution that this wheel was built from was /// downloaded. diff --git a/crates/puffin-distribution/src/lib.rs b/crates/puffin-distribution/src/lib.rs index 25507c235..134b9d19b 100644 --- a/crates/puffin-distribution/src/lib.rs +++ b/crates/puffin-distribution/src/lib.rs @@ -1,5 +1,5 @@ pub use distribution_database::{DistributionDatabase, DistributionDatabaseError}; -pub use download::{BuiltWheel, DiskWheel, InMemoryWheel, LocalWheel}; +pub use download::{BuiltWheel, DiskWheel, LocalWheel}; pub use index::{BuiltWheelIndex, RegistryWheelIndex}; pub use reporter::Reporter; pub use source::{SourceDistCachedBuilder, SourceDistError}; diff --git a/crates/puffin-distribution/src/unzip.rs b/crates/puffin-distribution/src/unzip.rs index 6708cc576..37a8f6843 100644 --- a/crates/puffin-distribution/src/unzip.rs +++ b/crates/puffin-distribution/src/unzip.rs @@ -3,19 +3,13 @@ use std::path::Path; use puffin_extract::{unzip_archive, Error}; use crate::download::BuiltWheel; -use crate::{DiskWheel, InMemoryWheel, LocalWheel}; +use crate::{DiskWheel, LocalWheel}; pub trait Unzip { /// Unzip a wheel into the target directory. fn unzip(&self, target: &Path) -> Result<(), Error>; } -impl Unzip for InMemoryWheel { - fn unzip(&self, target: &Path) -> Result<(), Error> { - unzip_archive(std::io::Cursor::new(&self.buffer), target) - } -} - impl Unzip for DiskWheel { fn unzip(&self, target: &Path) -> Result<(), Error> { unzip_archive(fs_err::File::open(&self.path)?, target) @@ -32,7 +26,6 @@ impl Unzip for LocalWheel { fn unzip(&self, target: &Path) -> Result<(), Error> { match self { LocalWheel::Unzipped(_) => Ok(()), - LocalWheel::InMemory(wheel) => wheel.unzip(target), LocalWheel::Disk(wheel) => wheel.unzip(target), LocalWheel::Built(wheel) => wheel.unzip(target), }