Share in-flight map across resolutions (#932)

## Summary

This PR fixes a subtle bug in `pip install` when using `--reinstall`. If
a package depends on a build system directly (e.g., `waitress` depends
on `setuptools`), and then you have other packages that also need the
build system to build a source distribution, right now, we don't share
the `OnceMap` between those cases.

This lifts the `InFlight` tracking up a level, so that it's initialized
once per command, then shared everywhere.

## Test Plan

I'm having trouble coming up with an identical test-case and hesitant to
add this slow test to the suite... But if you run `pip install
--reinstall` with:

```
waitress @ git+https://github.com/zanieb/waitress
devpi-server @ git+https://github.com/zanieb/devpi#subdirectory=server
```

It fails consistently on `main` and passes here.

crates/puffin-cli/src/commands/pip_compile.rs

@@ -26,7 +26,7 @@ use puffin_normalize::ExtraName;
 use puffin_resolver::{
     DisplayResolutionGraph, Manifest, PreReleaseMode, ResolutionMode, ResolutionOptions, Resolver,
 };
-use puffin_traits::SetupPyStrategy;
+use puffin_traits::{InFlight, SetupPyStrategy};
 use requirements_txt::EditableRequirement;
 
 use crate::commands::reporters::{DownloadReporter, ResolverReporter};
@@ -154,6 +154,9 @@ pub(crate) async fn pip_compile(
         FlatIndex::from_entries(entries, &tags)
     };
 
+    // Track in-flight downloads, builds, etc., across resolutions.
+    let in_flight = InFlight::default();
+
     let options = ResolutionOptions::new(resolution_mode, prerelease_mode, exclude_newer);
     let build_dispatch = BuildDispatch::new(
         &client,
@@ -161,6 +164,7 @@ pub(crate) async fn pip_compile(
         &interpreter,
         &index_locations,
         &flat_index,
+        &in_flight,
         interpreter.sys_executable().to_path_buf(),
         setup_py,
         no_build,
@@ -227,13 +231,6 @@ pub(crate) async fn pip_compile(
         editable_metadata,
     );
 
-    // Resolve the flat indexes from `--find-links`.
-    let flat_index = {
-        let client = FlatIndexClient::new(&client, &cache);
-        let entries = client.fetch(index_locations.flat_indexes()).await?;
-        FlatIndex::from_entries(entries, &tags)
-    };
-
     // Resolve the dependencies.
     let resolver = Resolver::new(
         manifest,

crates/puffin-cli/src/commands/pip_install.rs

@@ -28,7 +28,7 @@ use puffin_normalize::PackageName;
 use puffin_resolver::{
     Manifest, PreReleaseMode, ResolutionGraph, ResolutionMode, ResolutionOptions, Resolver,
 };
-use puffin_traits::{OnceMap, SetupPyStrategy};
+use puffin_traits::{InFlight, SetupPyStrategy};
 use requirements_txt::EditableRequirement;
 
 use crate::commands::reporters::{DownloadReporter, InstallReporter, ResolverReporter};
@@ -144,6 +144,9 @@ pub(crate) async fn pip_install(
         FlatIndex::from_entries(entries, tags)
     };
 
+    // Track in-flight downloads, builds, etc., across resolutions.
+    let in_flight = InFlight::default();
+
     let options = ResolutionOptions::new(resolution_mode, prerelease_mode, exclude_newer);
 
     let build_dispatch = BuildDispatch::new(
@@ -152,6 +155,7 @@ pub(crate) async fn pip_install(
         &interpreter,
         &index_locations,
         &flat_index,
+        &in_flight,
         venv.python_executable(),
         setup_py,
         no_build,
@@ -221,6 +225,7 @@ pub(crate) async fn pip_install(
         &index_locations,
         tags,
         &client,
+        &in_flight,
         &build_dispatch,
         &cache,
         &venv,
@@ -420,6 +425,7 @@ async fn install(
     index_urls: &IndexLocations,
     tags: &Tags,
     client: &RegistryClient,
+    in_flight: &InFlight,
     build_dispatch: &BuildDispatch<'_>,
     cache: &Cache,
     venv: &Virtualenv,
@@ -488,8 +494,9 @@ async fn install(
         let downloader = Downloader::new(cache, tags, client, build_dispatch)
             .with_reporter(DownloadReporter::from(printer).with_length(remote.len() as u64));
 
+        // STOPSHIP(charlie): This needs to be shared!
         let wheels = downloader
-            .download(remote, &OnceMap::default())
+            .download(remote, in_flight)
             .await
             .context("Failed to download distributions")?;
 

crates/puffin-cli/src/commands/pip_sync.rs

@@ -14,7 +14,7 @@ use puffin_client::{FlatIndex, FlatIndexClient, RegistryClient, RegistryClientBu
 use puffin_dispatch::BuildDispatch;
 use puffin_installer::{Downloader, InstallPlan, Reinstall, ResolvedEditable, SitePackages};
 use puffin_interpreter::Virtualenv;
-use puffin_traits::{OnceMap, SetupPyStrategy};
+use puffin_traits::{InFlight, SetupPyStrategy};
 use pypi_types::Yanked;
 use requirements_txt::EditableRequirement;
 
@@ -70,6 +70,9 @@ pub(crate) async fn pip_sync(
         FlatIndex::from_entries(entries, tags)
     };
 
+    // Track in-flight downloads, builds, etc., across resolutions.
+    let in_flight = InFlight::default();
+
     // Prep the build context.
     let build_dispatch = BuildDispatch::new(
         &client,
@@ -77,6 +80,7 @@ pub(crate) async fn pip_sync(
         venv.interpreter(),
         &index_locations,
         &flat_index,
+        &in_flight,
         venv.python_executable(),
         setup_py,
         no_build,
@@ -211,7 +215,7 @@ pub(crate) async fn pip_sync(
             .with_reporter(DownloadReporter::from(printer).with_length(remote.len() as u64));
 
         let wheels = downloader
-            .download(remote, &OnceMap::default())
+            .download(remote, &in_flight)
             .await
             .context("Failed to download distributions")?;
 

crates/puffin-cli/src/commands/venv.rs

@@ -15,7 +15,7 @@ use puffin_cache::Cache;
 use puffin_client::{FlatIndex, FlatIndexClient, RegistryClientBuilder};
 use puffin_dispatch::BuildDispatch;
 use puffin_interpreter::Interpreter;
-use puffin_traits::{BuildContext, SetupPyStrategy};
+use puffin_traits::{BuildContext, InFlight, SetupPyStrategy};
 
 use crate::commands::ExitStatus;
 use crate::printer::Printer;
@@ -139,6 +139,9 @@ async fn venv_impl(
             FlatIndex::from_entries(entries, tags)
         };
 
+        // Track in-flight downloads, builds, etc., across resolutions.
+        let in_flight = InFlight::default();
+
         // Prep the build context.
         let build_dispatch = BuildDispatch::new(
             &client,
@@ -146,6 +149,7 @@ async fn venv_impl(
             interpreter,
             index_locations,
             &flat_index,
+            &in_flight,
             venv.python_executable(),
             SetupPyStrategy::default(),
             true,

crates/puffin-dev/src/build.rs

@@ -12,7 +12,7 @@ use puffin_cache::{Cache, CacheArgs};
 use puffin_client::{FlatIndex, RegistryClientBuilder};
 use puffin_dispatch::BuildDispatch;
 use puffin_interpreter::Virtualenv;
-use puffin_traits::{BuildContext, BuildKind, SetupPyStrategy};
+use puffin_traits::{BuildContext, BuildKind, InFlight, SetupPyStrategy};
 
 #[derive(Parser)]
 pub(crate) struct BuildArgs {
@@ -57,6 +57,7 @@ pub(crate) async fn build(args: BuildArgs) -> Result<PathBuf> {
     let index_urls = IndexLocations::default();
     let flat_index = FlatIndex::default();
     let setup_py = SetupPyStrategy::default();
+    let in_flight = InFlight::default();
 
     let build_dispatch = BuildDispatch::new(
         &client,
@@ -64,6 +65,7 @@ pub(crate) async fn build(args: BuildArgs) -> Result<PathBuf> {
         venv.interpreter(),
         &index_urls,
         &flat_index,
+        &in_flight,
         venv.python_executable(),
         setup_py,
         false,

crates/puffin-dev/src/install_many.rs

@@ -25,7 +25,7 @@ use puffin_installer::Downloader;
 use puffin_interpreter::Virtualenv;
 use puffin_normalize::PackageName;
 use puffin_resolver::DistFinder;
-use puffin_traits::{BuildContext, OnceMap, SetupPyStrategy};
+use puffin_traits::{BuildContext, InFlight, SetupPyStrategy};
 
 #[derive(Parser)]
 pub(crate) struct InstallManyArgs {
@@ -62,6 +62,7 @@ pub(crate) async fn install_many(args: InstallManyArgs) -> Result<()> {
     let index_locations = IndexLocations::default();
     let flat_index = FlatIndex::default();
     let setup_py = SetupPyStrategy::default();
+    let in_flight = InFlight::default();
     let tags = venv.interpreter().tags()?;
 
     let build_dispatch = BuildDispatch::new(
@@ -70,6 +71,7 @@ pub(crate) async fn install_many(args: InstallManyArgs) -> Result<()> {
         venv.interpreter(),
         &index_locations,
         &flat_index,
+        &in_flight,
         venv.python_executable(),
         setup_py,
         args.no_build,
@@ -141,7 +143,7 @@ async fn install_chunk(
 
     let mut registry_index = RegistryWheelIndex::new(build_dispatch.cache(), tags, index_locations);
     let (cached, uncached): (Vec<_>, Vec<_>) = dists.into_iter().partition_map(|dist| {
-        // We always want the wheel for the latest version not whatever matching is in cache
+        // We always want the wheel for the latest version not whatever matching is in cache.
         let VersionOrUrl::Version(version) = dist.version_or_url() else {
             unreachable!();
         };
@@ -155,7 +157,7 @@ async fn install_chunk(
     info!("Cached: {}, Uncached {}", cached.len(), uncached.len());
 
     let downloader = Downloader::new(build_dispatch.cache(), tags, client, build_dispatch);
-    let in_flight = OnceMap::default();
+    let in_flight = InFlight::default();
     let fetches: Vec<_> = futures::stream::iter(uncached)
         .map(|dist| downloader.get_wheel(dist, &in_flight))
         .buffer_unordered(50)

crates/puffin-dev/src/resolve_cli.rs

@@ -17,7 +17,7 @@ use puffin_client::{FlatIndex, FlatIndexClient, RegistryClientBuilder};
 use puffin_dispatch::BuildDispatch;
 use puffin_interpreter::Virtualenv;
 use puffin_resolver::{Manifest, ResolutionOptions, Resolver};
-use puffin_traits::SetupPyStrategy;
+use puffin_traits::{InFlight, SetupPyStrategy};
 
 #[derive(ValueEnum, Default, Clone)]
 pub(crate) enum ResolveCliFormat {
@@ -65,6 +65,7 @@ pub(crate) async fn resolve_cli(args: ResolveCliArgs) -> Result<()> {
         let entries = client.fetch(index_locations.flat_indexes()).await?;
         FlatIndex::from_entries(entries, venv.interpreter().tags()?)
     };
+    let in_flight = InFlight::default();
 
     let build_dispatch = BuildDispatch::new(
         &client,
@@ -72,6 +73,7 @@ pub(crate) async fn resolve_cli(args: ResolveCliArgs) -> Result<()> {
         venv.interpreter(),
         &index_locations,
         &flat_index,
+        &in_flight,
         venv.python_executable(),
         SetupPyStrategy::default(),
         args.no_build,

crates/puffin-dev/src/resolve_many.rs

@@ -20,7 +20,7 @@ use puffin_client::{FlatIndex, RegistryClient, RegistryClientBuilder};
 use puffin_dispatch::BuildDispatch;
 use puffin_interpreter::Virtualenv;
 use puffin_normalize::PackageName;
-use puffin_traits::{BuildContext, SetupPyStrategy};
+use puffin_traits::{BuildContext, InFlight, SetupPyStrategy};
 
 #[derive(Parser)]
 pub(crate) struct ResolveManyArgs {
@@ -76,6 +76,7 @@ pub(crate) async fn resolve_many(args: ResolveManyArgs) -> Result<()> {
     let index_locations = IndexLocations::default();
     let flat_index = FlatIndex::default();
     let setup_py = SetupPyStrategy::default();
+    let in_flight = InFlight::default();
 
     let build_dispatch = BuildDispatch::new(
         &client,
@@ -83,6 +84,7 @@ pub(crate) async fn resolve_many(args: ResolveManyArgs) -> Result<()> {
         venv.interpreter(),
         &index_locations,
         &flat_index,
+        &in_flight,
         venv.python_executable(),
         setup_py,
         args.no_build,

crates/puffin-dispatch/src/lib.rs

@@ -9,7 +9,7 @@ use anyhow::{bail, Context, Result};
 use itertools::Itertools;
 use tracing::{debug, instrument};
 
-use distribution_types::{CachedDist, DistributionId, IndexLocations, Name, Resolution};
+use distribution_types::{IndexLocations, Name, Resolution};
 use pep508_rs::Requirement;
 use puffin_build::{SourceBuild, SourceBuildContext};
 use puffin_cache::Cache;
@@ -17,7 +17,7 @@ use puffin_client::{FlatIndex, RegistryClient};
 use puffin_installer::{Downloader, InstallPlan, Installer, Reinstall, SitePackages};
 use puffin_interpreter::{Interpreter, Virtualenv};
 use puffin_resolver::{Manifest, ResolutionOptions, Resolver};
-use puffin_traits::{BuildContext, BuildKind, OnceMap, SetupPyStrategy};
+use puffin_traits::{BuildContext, BuildKind, InFlight, SetupPyStrategy};
 
 /// The main implementation of [`BuildContext`], used by the CLI, see [`BuildContext`]
 /// documentation.
@@ -27,12 +27,12 @@ pub struct BuildDispatch<'a> {
     interpreter: &'a Interpreter,
     index_locations: &'a IndexLocations,
     flat_index: &'a FlatIndex,
+    in_flight: &'a InFlight,
     base_python: PathBuf,
     setup_py: SetupPyStrategy,
     no_build: bool,
     source_build_context: SourceBuildContext,
     options: ResolutionOptions,
-    in_flight: OnceMap<DistributionId, Result<CachedDist, String>>,
 }
 
 impl<'a> BuildDispatch<'a> {
@@ -43,6 +43,7 @@ impl<'a> BuildDispatch<'a> {
         interpreter: &'a Interpreter,
         index_locations: &'a IndexLocations,
         flat_index: &'a FlatIndex,
+        in_flight: &'a InFlight,
         base_python: PathBuf,
         setup_py: SetupPyStrategy,
         no_build: bool,
@@ -53,12 +54,12 @@ impl<'a> BuildDispatch<'a> {
             interpreter,
             index_locations,
             flat_index,
+            in_flight,
             base_python,
             setup_py,
             no_build,
             source_build_context: SourceBuildContext::default(),
             options: ResolutionOptions::default(),
-            in_flight: OnceMap::default(),
         }
     }
 
@@ -184,7 +185,7 @@ impl<'a> BuildContext for BuildDispatch<'a> {
                 );
 
                 downloader
-                    .download(remote, &self.in_flight)
+                    .download(remote, self.in_flight)
                     .await
                     .context("Failed to download and build distributions")?
             };

crates/puffin-installer/src/downloader.rs

@@ -7,14 +7,12 @@ use tokio::task::JoinError;
 use tracing::{instrument, warn};
 use url::Url;
 
-use distribution_types::{
+use distribution_types::{CachedDist, Dist, Identifier, LocalEditable, RemoteSource, SourceDist};
-    CachedDist, Dist, DistributionId, Identifier, LocalEditable, RemoteSource, SourceDist,
-};
 use platform_tags::Tags;
 use puffin_cache::Cache;
 use puffin_client::RegistryClient;
 use puffin_distribution::{DistributionDatabase, DistributionDatabaseError, LocalWheel, Unzip};
-use puffin_traits::{BuildContext, OnceMap};
+use puffin_traits::{BuildContext, InFlight};
 
 use crate::editable::BuiltEditable;
 
@@ -66,7 +64,7 @@ impl<'a, Context: BuildContext + Send + Sync> Downloader<'a, Context> {
     pub fn download_stream<'stream>(
         &'stream self,
         distributions: Vec<Dist>,
-        in_flight: &'stream OnceMap<DistributionId, Result<CachedDist, String>>,
+        in_flight: &'stream InFlight,
     ) -> impl Stream<Item = Result<CachedDist, Error>> + 'stream {
         futures::stream::iter(distributions)
             .map(|dist| async {
@@ -86,7 +84,7 @@ impl<'a, Context: BuildContext + Send + Sync> Downloader<'a, Context> {
     pub async fn download(
         &self,
         mut distributions: Vec<Dist>,
-        in_flight: &OnceMap<DistributionId, Result<CachedDist, String>>,
+        in_flight: &InFlight,
     ) -> Result<Vec<CachedDist>, Error> {
         // Sort the distributions by size.
         distributions
@@ -154,13 +152,9 @@ impl<'a, Context: BuildContext + Send + Sync> Downloader<'a, Context> {
 
     /// Download, build, and unzip a single wheel.
     #[instrument(skip_all, fields(name = % dist, size = ? dist.size(), url = dist.file().map(|file| file.url.to_string()).unwrap_or_default()))]
-    pub async fn get_wheel(
+    pub async fn get_wheel(&self, dist: Dist, in_flight: &InFlight) -> Result<CachedDist, Error> {
-        &self,
-        dist: Dist,
-        in_flight: &OnceMap<DistributionId, Result<CachedDist, String>>,
-    ) -> Result<CachedDist, Error> {
         let id = dist.distribution_id();
-        let wheel = if in_flight.register(&id) {
+        let wheel = if in_flight.downloads.register(&id) {
             let download: LocalWheel = self
                 .database
                 .get_or_build_wheel(dist.clone())
@@ -169,16 +163,17 @@ impl<'a, Context: BuildContext + Send + Sync> Downloader<'a, Context> {
             let result = Self::unzip_wheel(download).await;
             match result {
                 Ok(cached) => {
-                    in_flight.done(id, Ok(cached.clone()));
+                    in_flight.downloads.done(id, Ok(cached.clone()));
                     cached
                 }
                 Err(err) => {
-                    in_flight.done(id, Err(err.to_string()));
+                    in_flight.downloads.done(id, Err(err.to_string()));
                     return Err(err);
                 }
             }
         } else {
             in_flight
+                .downloads
                 .wait(&id)
                 .await
                 .value()

crates/puffin-traits/src/lib.rs

@@ -6,7 +6,7 @@ use std::path::{Path, PathBuf};
 
 use anyhow::Result;
 
-use distribution_types::Resolution;
+use distribution_types::{CachedDist, DistributionId, Resolution};
 pub use once_map::OnceMap;
 use pep508_rs::Requirement;
 use puffin_cache::Cache;
@@ -124,6 +124,12 @@ pub trait SourceBuildTrait {
         -> impl Future<Output = Result<String>> + Send + 'a;
 }
 
+#[derive(Default)]
+pub struct InFlight {
+    /// The in-flight distribution downloads.
+    pub downloads: OnceMap<DistributionId, Result<CachedDist, String>>,
+}
+
 /// The strategy to use when building source distributions that lack a `pyproject.toml`.
 #[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
 pub enum SetupPyStrategy {