Use borrowed data in BuildDispatch (#679)

This PR uses borrowed data in `BuildDispatch` which makes creating a
`BuildDispatch` extremely cheap (only one allocation, for the Python
executable). I can be talked out of this, it will have no measurable
impact.

crates/puffin-dev/src/build.rs

@@ -53,14 +53,16 @@ pub(crate) async fn build(args: BuildArgs) -> Result<PathBuf> {
 
     let platform = Platform::current()?;
     let venv = Virtualenv::from_env(platform, &cache)?;
+    let client = RegistryClientBuilder::new(cache.clone()).build();
+    let index_urls = IndexUrls::default();
 
     let build_dispatch = BuildDispatch::new(
-        RegistryClientBuilder::new(cache.clone()).build(),
-        cache,
-        venv.interpreter().clone(),
+        &client,
+        &cache,
+        venv.interpreter(),
+        &index_urls,
         venv.python_executable(),
         false,
-        IndexUrls::default(),
     );
 
     let builder = SourceBuild::setup(

crates/puffin-dev/src/resolve_cli.rs

@@ -51,13 +51,15 @@ pub(crate) async fn resolve_cli(args: ResolveCliArgs) -> Result<()> {
     let platform = Platform::current()?;
     let venv = Virtualenv::from_env(platform, &cache)?;
     let client = RegistryClientBuilder::new(cache.clone()).build();
+    let index_urls = IndexUrls::default();
+
     let build_dispatch = BuildDispatch::new(
-        client.clone(),
-        cache.clone(),
-        venv.interpreter().clone(),
+        &client,
+        &cache,
+        venv.interpreter(),
+        &index_urls,
         venv.python_executable(),
         args.no_build,
-        IndexUrls::default(),
     );
 
     // Copied from `BuildDispatch`

crates/puffin-dev/src/resolve_many.rs

@@ -4,11 +4,8 @@ use std::sync::Arc;
 
 use anyhow::Result;
 use clap::Parser;
-use fs_err as fs;
-use futures::stream::FuturesUnordered;
 use futures::StreamExt;
 use indicatif::ProgressStyle;
-use tokio::sync::Semaphore;
 use tokio::time::Instant;
 use tracing::{info, info_span, span, Level, Span};
 use tracing_indicatif::span_ext::IndicatifSpanExt;
@@ -19,13 +16,13 @@ use puffin_cache::{Cache, CacheArgs};
 use puffin_client::RegistryClientBuilder;
 use puffin_dispatch::BuildDispatch;
 use puffin_interpreter::Virtualenv;
-use puffin_normalize::PackageName;
 use puffin_traits::BuildContext;
 use pypi_types::IndexUrls;
 
 #[derive(Parser)]
 pub(crate) struct ResolveManyArgs {
-    list: PathBuf,
+    /// Path to a file containing one requirement per line.
+    requirements: PathBuf,
     #[clap(long)]
     limit: Option<usize>,
     /// Don't build source distributions. This means resolving will not run arbitrary code. The
@@ -42,71 +39,62 @@ pub(crate) struct ResolveManyArgs {
 pub(crate) async fn resolve_many(args: ResolveManyArgs) -> Result<()> {
     let cache = Cache::try_from(args.cache_args)?;
 
-    let data = fs::read_to_string(&args.list)?;
+    let data = fs_err::read_to_string(&args.requirements)?;
     let lines = data.lines().map(Requirement::from_str);
     let requirements: Vec<Requirement> = if let Some(limit) = args.limit {
         lines.take(limit).collect::<Result<_, _>>()?
     } else {
         lines.collect::<Result<_, _>>()?
     };
+    let total = requirements.len();
 
     let platform = Platform::current()?;
     let venv = Virtualenv::from_env(platform, &cache)?;
+    let client = RegistryClientBuilder::new(cache.clone()).build();
+    let index_urls = IndexUrls::default();
+
     let build_dispatch = BuildDispatch::new(
-        RegistryClientBuilder::new(cache.clone()).build(),
-        cache.clone(),
-        venv.interpreter().clone(),
+        &client,
+        &cache,
+        venv.interpreter(),
+        &index_urls,
         venv.python_executable(),
         args.no_build,
-        IndexUrls::default(),
     );
-
-    let build_dispatch_arc = Arc::new(build_dispatch);
-    let mut tasks = FuturesUnordered::new();
-    let semaphore = Arc::new(Semaphore::new(args.num_tasks));
+    let build_dispatch = Arc::new(build_dispatch);
 
     let header_span = info_span!("resolve many");
     header_span.pb_set_style(&ProgressStyle::default_bar());
-    let total = requirements.len();
     header_span.pb_set_length(total as u64);
-
     let _header_span_enter = header_span.enter();
 
-    let tf_models_nightly = PackageName::from_str("tf-models-nightly").unwrap();
-    for requirement in requirements {
-        if requirement.name == tf_models_nightly {
-            continue;
-        }
-        let build_dispatch_arc = build_dispatch_arc.clone();
-        let semaphore = semaphore.clone();
-        tasks.push(tokio::spawn(async move {
-            let permit = semaphore.clone().acquire_owned().await.unwrap();
-            let span = span!(Level::TRACE, "resolving");
-            let _enter = span.enter();
-            let start = Instant::now();
+    let mut tasks = futures::stream::iter(requirements)
+        .map(|requirement| {
+            let build_dispatch = build_dispatch.clone();
+            async move {
+                let span = span!(Level::TRACE, "fetching");
+                let _enter = span.enter();
+                let start = Instant::now();
 
-            let result = build_dispatch_arc.resolve(&[requirement.clone()]).await;
+                let result = build_dispatch.resolve(&[requirement.clone()]).await;
 
-            drop(permit);
-            (requirement.to_string(), start.elapsed(), result)
-        }));
-    }
+                (requirement.to_string(), start.elapsed(), result)
+            }
+        })
+        .buffer_unordered(args.num_tasks);
 
     let mut success = 0usize;
     let mut errors = Vec::new();
-
     while let Some(result) = tasks.next().await {
-        let (package, duration, result) = result.unwrap();
-
+        let (package, duration, result) = result;
         match result {
-            Ok(resolution) => {
+            Ok(_) => {
                 info!(
-                    "Success ({}/{}, {} ms): {} ({} package(s))",
+                    "Success ({}/{}, {} ms): {}",
                     success + errors.len(),
                     total,
                     duration.as_millis(),
                     package,
-                    resolution.len(),
                 );
                 success += 1;
             }
@@ -124,6 +112,7 @@ pub(crate) async fn resolve_many(args: ResolveManyArgs) -> Result<()> {
         }
         Span::current().pb_inc(1);
     }
+
     info!("Errors: {}", errors.join(", "));
     info!("Success: {}, Error: {}", success, errors.len());
     Ok(())