Reuse distribution structs in Resolver's source_distribution.rs (#264)

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())
+    }
+}

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;

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<Requirement>,
@@ -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))
             }
         }
     }

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<Path>) -> 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<PathBuf> {
-        let Ok(read_dir) = fs_err::read_dir(self.version(package_name, version)) else {
+    ) -> Result<Option<Metadata21>> {
+        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<Metadata21> {
+        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<PathBuf> {
+        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<Metadata21> {
-    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<Metadata21> {
-    let dist_info = spawn_blocking(move || -> Result<String> {
-        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<Path>) -> Result<Metadata21> {
+    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())?)
 }