diff --git a/crates/puffin-distribution/src/lib.rs b/crates/puffin-distribution/src/lib.rs index da8cf41db..103c85029 100644 --- a/crates/puffin-distribution/src/lib.rs +++ b/crates/puffin-distribution/src/lib.rs @@ -251,3 +251,43 @@ impl std::fmt::Display for InstalledDistribution { write!(f, "{}@{}", self.name(), self.version()) } } + +/// Unowned reference to a [`RemoteDistribution`]. +#[derive(Debug, Clone)] +pub struct RemoteDistributionRef<'a> { + name: &'a PackageName, + version: &'a Version, + file: &'a File, +} + +impl<'a> RemoteDistributionRef<'a> { + pub fn new(name: &'a PackageName, version: &'a Version, file: &'a File) -> Self { + Self { + name, + version, + file, + } + } + + pub fn name(&self) -> &PackageName { + self.name + } + + pub fn version(&self) -> &Version { + self.version + } + + pub fn file(&self) -> &File { + self.file + } + + pub fn id(&self) -> String { + format!("{}-{}", DistInfoName::from(self.name()), self.version()) + } +} + +impl std::fmt::Display for RemoteDistributionRef<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}@{}", self.name(), self.version()) + } +} diff --git a/crates/puffin-resolver/src/lib.rs b/crates/puffin-resolver/src/lib.rs index 9086af10a..23b12369a 100644 --- a/crates/puffin-resolver/src/lib.rs +++ b/crates/puffin-resolver/src/lib.rs @@ -4,7 +4,6 @@ pub use prerelease_mode::PreReleaseMode; pub use resolution::Graph; pub use resolution_mode::ResolutionMode; pub use resolver::{Reporter as ResolverReporter, Resolver}; -pub use source_distribution::BuiltSourceDistributionCache; pub use wheel_finder::{Reporter as WheelFinderReporter, WheelFinder}; mod candidate_selector; diff --git a/crates/puffin-resolver/src/resolver.rs b/crates/puffin-resolver/src/resolver.rs index 3df2ad6a2..d7c5636cb 100644 --- a/crates/puffin-resolver/src/resolver.rs +++ b/crates/puffin-resolver/src/resolver.rs @@ -14,13 +14,14 @@ use pubgrub::range::Range; use pubgrub::solver::{Incompatibility, State}; use pubgrub::type_aliases::DependencyConstraints; use tokio::select; -use tracing::{debug, trace}; +use tracing::{debug, error, trace}; use waitmap::WaitMap; use distribution_filename::{SourceDistributionFilename, WheelFilename}; use pep508_rs::{MarkerEnvironment, Requirement}; use platform_tags::Tags; use puffin_client::RegistryClient; +use puffin_distribution::RemoteDistributionRef; use puffin_package::dist_info_name::DistInfoName; use puffin_package::package_name::PackageName; use puffin_package::pypi_types::{File, Metadata21, SimpleJson}; @@ -33,8 +34,7 @@ use crate::manifest::Manifest; use crate::pubgrub::{iter_requirements, version_range}; use crate::pubgrub::{PubGrubPackage, PubGrubPriorities, PubGrubVersion, MIN_VERSION}; use crate::resolution::Graph; -use crate::source_distribution::{download_and_build_sdist, read_dist_info}; -use crate::BuiltSourceDistributionCache; +use crate::source_distribution::SourceDistributionBuildTree; pub struct Resolver<'a, Context: BuildContext + Sync> { requirements: Vec, @@ -581,30 +581,32 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> { .await } Request::Sdist(file, package_name, version) => { - let cached_wheel = self.build_context.cache().and_then(|cache| { - BuiltSourceDistributionCache::new(cache).find_wheel( - &package_name, - &version, - self.tags, - ) - }); - let metadata21 = if let Some(cached_wheel) = cached_wheel { - read_dist_info(cached_wheel).await - } else { - download_and_build_sdist( - &file, - &package_name, - &version, - self.client, - self.build_context, - ) - .await - } - .map_err(|err| ResolveError::SourceDistribution { - filename: file.filename.clone(), - err, - })?; - Ok(Response::Sdist(file, metadata21)) + let build_tree = SourceDistributionBuildTree::new(self.build_context); + let distribution = RemoteDistributionRef::new(&package_name, &version, &file); + let metadata = match build_tree.find_dist_info(&distribution, self.tags) { + Ok(Some(metadata)) => metadata, + Ok(None) => build_tree + .download_and_build_sdist(&distribution, self.client) + .await + .map_err(|err| ResolveError::SourceDistribution { + filename: file.filename.clone(), + err, + })?, + Err(err) => { + error!( + "Failed to read source distribution {} from cache: {}", + file.filename, err + ); + build_tree + .download_and_build_sdist(&distribution, self.client) + .await + .map_err(|err| ResolveError::SourceDistribution { + filename: file.filename.clone(), + err, + })? + } + }; + Ok(Response::Sdist(file, metadata)) } } } diff --git a/crates/puffin-resolver/src/source_distribution.rs b/crates/puffin-resolver/src/source_distribution.rs index dd6cc9782..413bd61f9 100644 --- a/crates/puffin-resolver/src/source_distribution.rs +++ b/crates/puffin-resolver/src/source_distribution.rs @@ -4,49 +4,87 @@ use std::str::FromStr; use anyhow::Result; use fs_err::tokio as fs; use tempfile::tempdir; -use tokio::task::spawn_blocking; use tokio_util::compat::FuturesAsyncReadCompatExt; use tracing::debug; use url::Url; use zip::ZipArchive; use distribution_filename::WheelFilename; -use pep440_rs::Version; use platform_tags::Tags; use puffin_client::RegistryClient; -use puffin_package::package_name::PackageName; -use puffin_package::pypi_types::{File, Metadata21}; +use puffin_distribution::RemoteDistributionRef; +use puffin_package::pypi_types::Metadata21; use puffin_traits::BuildContext; const BUILT_WHEELS_CACHE: &str = "built-wheels-v0"; -/// TODO(konstin): Find a better home for me? -/// /// Stores wheels built from source distributions. We need to keep those separate from the regular /// wheel cache since a wheel with the same name may be uploaded after we made our build and in that /// case the hashes would clash. -pub struct BuiltSourceDistributionCache(PathBuf); +pub(crate) struct SourceDistributionBuildTree<'a, T: BuildContext>(&'a T); -impl BuiltSourceDistributionCache { - pub fn new(path: impl AsRef) -> Self { - Self(path.as_ref().join(BUILT_WHEELS_CACHE)) +impl<'a, T: BuildContext> SourceDistributionBuildTree<'a, T> { + /// Initialize a [`SourceDistributionBuildTree`] from a [`BuildContext`]. + pub(crate) fn new(build_context: &'a T) -> Self { + Self(build_context) } - pub fn version(&self, name: &PackageName, version: &Version) -> PathBuf { - self.0.join(name.to_string()).join(version.to_string()) - } - - /// Search for a wheel matching the tags that was built from the given source distribution. - pub fn find_wheel( + /// Read the [`Metadata21`] from a built source distribution, if it exists in the cache. + pub(crate) fn find_dist_info( &self, - package_name: &PackageName, - version: &Version, + distribution: &RemoteDistributionRef<'_>, tags: &Tags, - ) -> Option { - let Ok(read_dir) = fs_err::read_dir(self.version(package_name, version)) else { + ) -> Result> { + self.find_wheel(distribution, tags) + .map(read_dist_info) + .transpose() + } + + /// Download and build a source distribution, storing the built wheel in the cache. + pub(crate) async fn download_and_build_sdist( + &self, + distribution: &RemoteDistributionRef<'_>, + client: &RegistryClient, + ) -> Result { + debug!("Building: {}", distribution.file().filename); + let url = Url::parse(&distribution.file().url)?; + let reader = client.stream_external(&url).await?; + let mut reader = tokio::io::BufReader::new(reader.compat()); + let temp_dir = tempdir()?; + + let sdist_dir = temp_dir.path().join("sdist"); + tokio::fs::create_dir(&sdist_dir).await?; + let sdist_file = sdist_dir.join(&distribution.file().filename); + let mut writer = tokio::fs::File::create(&sdist_file).await?; + tokio::io::copy(&mut reader, &mut writer).await?; + + let wheel_dir = self.0.cache().map_or_else( + || temp_dir.path().join(BUILT_WHEELS_CACHE), + |cache| cache.join(BUILT_WHEELS_CACHE).join(distribution.id()), + ); + fs::create_dir_all(&wheel_dir).await?; + + let disk_filename = self + .0 + .build_source_distribution(&sdist_file, &wheel_dir) + .await?; + + let metadata21 = read_dist_info(wheel_dir.join(disk_filename))?; + + debug!("Finished building: {}", distribution.file().filename); + Ok(metadata21) + } + + /// Search for a wheel matching the tags that was built from the given source distribution. + fn find_wheel(&self, distribution: &RemoteDistributionRef<'_>, tags: &Tags) -> Option { + let wheel_dir = self + .0 + .cache()? + .join(BUILT_WHEELS_CACHE) + .join(distribution.id()); + let Ok(read_dir) = fs_err::read_dir(wheel_dir) else { return None; }; - for entry in read_dir { let Ok(entry) = entry else { continue; @@ -64,56 +102,22 @@ impl BuiltSourceDistributionCache { } } -pub(crate) async fn download_and_build_sdist( - file: &File, - package_name: &PackageName, - version: &Version, - client: &RegistryClient, - build_context: &impl BuildContext, -) -> Result { - debug!("Building: {}", &file.filename); - let url = Url::parse(&file.url)?; - let reader = client.stream_external(&url).await?; - let mut reader = tokio::io::BufReader::new(reader.compat()); - let temp_dir = tempdir()?; - - let sdist_dir = temp_dir.path().join("sdist"); - tokio::fs::create_dir(&sdist_dir).await?; - let sdist_file = sdist_dir.join(&file.filename); - let mut writer = tokio::fs::File::create(&sdist_file).await?; - tokio::io::copy(&mut reader, &mut writer).await?; - - let wheel_dir = if let Some(cache) = &build_context.cache() { - BuiltSourceDistributionCache::new(cache).version(package_name, version) - } else { - temp_dir.path().join("wheels") - }; - - fs::create_dir_all(&wheel_dir).await?; - - let disk_filename = build_context - .build_source_distribution(&sdist_file, &wheel_dir) - .await?; - - let metadata21 = read_dist_info(wheel_dir.join(disk_filename)).await?; - - debug!("Finished building: {}", &file.filename); - Ok(metadata21) -} - -pub(crate) async fn read_dist_info(wheel: PathBuf) -> Result { - let dist_info = spawn_blocking(move || -> Result { - let mut archive = ZipArchive::new(std::fs::File::open(&wheel)?)?; - let dist_info_prefix = install_wheel_rs::find_dist_info( - &WheelFilename::from_str(wheel.file_name().unwrap().to_string_lossy().as_ref())?, - &mut archive, - )?; - let dist_info = std::io::read_to_string( - archive.by_name(&format!("{dist_info_prefix}.dist-info/METADATA"))?, - )?; - Ok(dist_info) - }) - .await - .unwrap()?; +/// Read the [`Metadata21`] from a wheel. +fn read_dist_info(wheel: impl AsRef) -> Result { + let mut archive = ZipArchive::new(std::fs::File::open(&wheel)?)?; + let dist_info_prefix = install_wheel_rs::find_dist_info( + &WheelFilename::from_str( + wheel + .as_ref() + .file_name() + .unwrap() + .to_string_lossy() + .as_ref(), + )?, + &mut archive, + )?; + let dist_info = std::io::read_to_string( + archive.by_name(&format!("{dist_info_prefix}.dist-info/METADATA"))?, + )?; Ok(Metadata21::parse(dist_info.as_bytes())?) }