Cache the setup.py resolution (#327)

Cache the resolution for the setup.py requirements (`pip`, `setuptools`,
`wheels`) across builds.

crates/puffin-build/src/lib.rs

@@ -7,6 +7,7 @@ use std::io::BufRead;
 use std::path::{Path, PathBuf};
 use std::process::{Command, Output};
 use std::str::FromStr;
+use std::sync::Arc;
 
 use flate2::read::GzDecoder;
 use fs_err as fs;
@@ -17,6 +18,7 @@ use serde::{Deserialize, Serialize};
 use tar::Archive;
 use tempfile::{tempdir, TempDir};
 use thiserror::Error;
+use tokio::sync::Mutex;
 use tracing::{debug, instrument};
 use zip::ZipArchive;
 
@@ -93,12 +95,20 @@ impl Pep517Backend {
     }
 }
 
+/// Uses an [`Arc`] internally, clone freely
+#[derive(Debug, Default, Clone)]
+pub struct SourceDistributionBuildContext {
+    /// Cache the first resolution of `pip`, `setuptools` and `wheel` we made for setup.py (and
+    /// some PEP 517) builds so we can reuse it
+    setup_py_requirements: Arc<Mutex<Option<Vec<Requirement>>>>,
+}
+
 /// Holds the state through a series of PEP 517 frontend to backend calls or a single setup.py
 /// invocation.
 ///
 /// This keeps both the temp dir and the result of a potential `prepare_metadata_for_build_wheel`
 /// call which changes how we call `build_wheel`.
-pub struct SourceDistributionBuilder {
+pub struct SourceDistributionBuild {
     temp_dir: TempDir,
     source_tree: PathBuf,
     /// `Some` if this is a PEP 517 build
@@ -116,7 +126,7 @@ pub struct SourceDistributionBuilder {
     metadata_directory: Option<PathBuf>,
 }
 
-impl SourceDistributionBuilder {
+impl SourceDistributionBuild {
     /// Create a virtual environment in which to build a source distribution, extracting the
     /// contents from an archive if necessary.
     pub async fn setup(
@@ -124,7 +134,8 @@ impl SourceDistributionBuilder {
         subdirectory: Option<&Path>,
         interpreter_info: &InterpreterInfo,
         build_context: &impl BuildContext,
-    ) -> Result<SourceDistributionBuilder, Error> {
+        setup_py_requirements: SourceDistributionBuildContext,
+    ) -> Result<SourceDistributionBuild, Error> {
         let temp_dir = tempdir()?;
 
         // TODO(konstin): Parse and verify filenames
@@ -180,16 +191,22 @@ impl SourceDistributionBuilder {
                 })?;
             build_system.requires.clone()
         } else {
-            // TODO(konstin): Resolve those once globally and cache per puffin invocation
             let requirements = vec![
                 Requirement::from_str("wheel").unwrap(),
                 Requirement::from_str("setuptools").unwrap(),
                 Requirement::from_str("pip").unwrap(),
             ];
-            let resolved_requirements = build_context
+            let mut resolution = setup_py_requirements.setup_py_requirements.lock().await;
-                .resolve(&requirements)
+            let resolved_requirements = if let Some(resolved_requirements) = &*resolution {
-                .await
+                resolved_requirements.clone()
-                .map_err(|err| Error::RequirementsInstall("setup.py build (resolve)", err))?;
+            } else {
+                let resolved_requirements = build_context
+                    .resolve(&requirements)
+                    .await
+                    .map_err(|err| Error::RequirementsInstall("setup.py build (resolve)", err))?;
+                *resolution = Some(resolved_requirements.clone());
+                resolved_requirements
+            };
             build_context
                 .install(&resolved_requirements, &venv)
                 .await

crates/puffin-dev/src/build.rs

@@ -7,10 +7,10 @@ use directories::ProjectDirs;
 use fs_err as fs;
 
 use platform_host::Platform;
-use puffin_build::SourceDistributionBuilder;
 use puffin_client::RegistryClientBuilder;
 use puffin_dispatch::BuildDispatch;
 use puffin_interpreter::Virtualenv;
+use puffin_traits::BuildContext;
 
 #[derive(Parser)]
 pub(crate) struct BuildArgs {
@@ -52,13 +52,8 @@ pub(crate) async fn build(args: BuildArgs) -> Result<PathBuf> {
         venv.interpreter_info().clone(),
         fs::canonicalize(venv.python_executable())?,
     );
-    let builder = SourceDistributionBuilder::setup(
+    let wheel = build_dispatch
-        &args.sdist,
+        .build_source_distribution(&args.sdist, args.subdirectory.as_deref(), &wheel_dir)
-        args.subdirectory.as_deref(),
+        .await?;
-        venv.interpreter_info(),
-        &build_dispatch,
-    )
-    .await?;
-    let wheel = builder.build(&wheel_dir)?;
     Ok(wheel_dir.join(wheel))
 }

crates/puffin-dispatch/src/lib.rs

@@ -13,7 +13,7 @@ use tracing::{debug, instrument};
 
 use pep508_rs::Requirement;
 use platform_tags::Tags;
-use puffin_build::SourceDistributionBuilder;
+use puffin_build::{SourceDistributionBuild, SourceDistributionBuildContext};
 use puffin_client::RegistryClient;
 use puffin_installer::{Builder, Downloader, Installer, PartitionedRequirements, Unzipper};
 use puffin_interpreter::{InterpreterInfo, Virtualenv};
@@ -27,6 +27,7 @@ pub struct BuildDispatch {
     cache: PathBuf,
     interpreter_info: InterpreterInfo,
     base_python: PathBuf,
+    source_distribution_builder: SourceDistributionBuildContext,
 }
 
 impl BuildDispatch {
@@ -41,6 +42,7 @@ impl BuildDispatch {
             cache,
             interpreter_info,
             base_python,
+            source_distribution_builder: SourceDistributionBuildContext::default(),
         }
     }
 }
@@ -228,9 +230,14 @@ impl BuildContext for BuildDispatch {
         wheel_dir: &'a Path,
     ) -> Pin<Box<dyn Future<Output = Result<String>> + Send + 'a>> {
         Box::pin(async move {
-            let builder =
+            let builder = SourceDistributionBuild::setup(
-                SourceDistributionBuilder::setup(sdist, subdirectory, &self.interpreter_info, self)
+                sdist,
-                    .await?;
+                subdirectory,
+                &self.interpreter_info,
+                self,
+                self.source_distribution_builder.clone(),
+            )
+            .await?;
             Ok(builder.build(wheel_dir)?)
         })
     }

crates/puffin-traits/src/lib.rs

@@ -53,6 +53,7 @@ pub trait BuildContext {
     /// All (potentially nested) source distribution builds use the same base python and can reuse
     /// it's metadata (e.g. wheel compatibility tags).
     fn interpreter_info(&self) -> &InterpreterInfo;
+
     /// The system (or conda) python interpreter to create venvs.
     fn base_python(&self) -> &Path;