From a0de83001c0ec6ea85067b811c573d55bc907435 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 20 Nov 2024 11:45:14 -0500 Subject: [PATCH] Parallelize network requests in `uv tree --outdated` (#9280) ## Summary Closes https://github.com/astral-sh/uv/issues/9266. --- crates/uv-resolver/src/lock/map.rs | 5 +++ crates/uv/src/commands/pip/latest.rs | 5 ++- crates/uv/src/commands/pip/list.rs | 19 +++++----- crates/uv/src/commands/project/tree.rs | 48 ++++++++++++++++---------- crates/uv/src/lib.rs | 1 + 5 files changed, 51 insertions(+), 27 deletions(-) diff --git a/crates/uv-resolver/src/lock/map.rs b/crates/uv-resolver/src/lock/map.rs index e9cfdd042..e97377dfe 100644 --- a/crates/uv-resolver/src/lock/map.rs +++ b/crates/uv-resolver/src/lock/map.rs @@ -13,6 +13,11 @@ impl Default for PackageMap { } impl PackageMap { + /// Insert a value by [`PackageId`]. + pub fn insert(&mut self, package: Package, value: T) -> Option { + self.0.insert(package.id, value) + } + /// Get a value by [`PackageId`]. pub(crate) fn get(&self, package_id: &PackageId) -> Option<&T> { self.0.get(package_id) diff --git a/crates/uv/src/commands/pip/latest.rs b/crates/uv/src/commands/pip/latest.rs index 552ea7cfd..1331c29f5 100644 --- a/crates/uv/src/commands/pip/latest.rs +++ b/crates/uv/src/commands/pip/latest.rs @@ -1,3 +1,4 @@ +use tracing::debug; use uv_client::{RegistryClient, VersionFiles}; use uv_distribution_filename::DistFilename; use uv_distribution_types::{IndexCapabilities, IndexUrl}; @@ -10,7 +11,7 @@ use uv_warnings::warn_user_once; /// /// The returned distribution is guaranteed to be compatible with the provided tags and Python /// requirement. -#[derive(Debug)] +#[derive(Debug, Copy, Clone)] pub(crate) struct LatestClient<'env> { pub(crate) client: &'env RegistryClient, pub(crate) capabilities: &'env IndexCapabilities, @@ -27,6 +28,8 @@ impl<'env> LatestClient<'env> { package: &PackageName, index: Option<&IndexUrl>, ) -> anyhow::Result, uv_client::Error> { + debug!("Fetching latest version of: `{package}`"); + let archives = match self.client.simple(package, index, self.capabilities).await { Ok(archives) => archives, Err(err) => { diff --git a/crates/uv/src/commands/pip/list.rs b/crates/uv/src/commands/pip/list.rs index 46937be0e..9ec38157b 100644 --- a/crates/uv/src/commands/pip/list.rs +++ b/crates/uv/src/commands/pip/list.rs @@ -3,8 +3,7 @@ use std::fmt::Write; use anstream::println; use anyhow::Result; -use futures::stream::FuturesUnordered; -use futures::TryStreamExt; +use futures::StreamExt; use itertools::Itertools; use owo_colors::OwoColorize; use rustc_hash::FxHashMap; @@ -15,7 +14,7 @@ use uv_cache::{Cache, Refresh}; use uv_cache_info::Timestamp; use uv_cli::ListFormat; use uv_client::{Connectivity, RegistryClientBuilder}; -use uv_configuration::{IndexStrategy, KeyringProviderType, TrustedHost}; +use uv_configuration::{Concurrency, IndexStrategy, KeyringProviderType, TrustedHost}; use uv_distribution_filename::DistFilename; use uv_distribution_types::{Diagnostic, IndexCapabilities, IndexLocations, InstalledDist, Name}; use uv_fs::Simplified; @@ -44,6 +43,7 @@ pub(crate) async fn pip_list( keyring_provider: KeyringProviderType, allow_insecure_host: Vec, connectivity: Connectivity, + concurrency: Concurrency, strict: bool, exclude_newer: Option, python: Option<&str>, @@ -111,15 +111,18 @@ pub(crate) async fn pip_list( }; // Fetch the latest version for each package. - results - .iter() + let mut fetches = futures::stream::iter(&results) .map(|dist| async { let latest = client.find_latest(dist.name(), None).await?; Ok::<(&PackageName, Option), uv_client::Error>((dist.name(), latest)) }) - .collect::>() - .try_collect::>() - .await? + .buffer_unordered(concurrency.downloads); + + let mut map = FxHashMap::default(); + while let Some((package, version)) = fetches.next().await.transpose()? { + map.insert(package, version); + } + map } else { FxHashMap::default() }; diff --git a/crates/uv/src/commands/project/tree.rs b/crates/uv/src/commands/project/tree.rs index 48ca9d13c..599b30a69 100644 --- a/crates/uv/src/commands/project/tree.rs +++ b/crates/uv/src/commands/project/tree.rs @@ -1,8 +1,8 @@ use std::path::Path; use anstream::print; -use anyhow::Result; -use futures::{stream, StreamExt}; +use anyhow::{Error, Result}; +use futures::StreamExt; use uv_cache::{Cache, Refresh}; use uv_cache_info::Timestamp; @@ -11,7 +11,6 @@ use uv_configuration::{ Concurrency, DevGroupsSpecification, LowerBound, TargetTriple, TrustedHost, }; use uv_distribution_types::IndexCapabilities; -use uv_pep440::Version; use uv_pep508::PackageName; use uv_python::{PythonDownloads, PythonPreference, PythonRequest, PythonVersion}; use uv_resolver::{PackageMap, TreeDisplay}; @@ -182,21 +181,34 @@ pub(crate) async fn tree( }; // Fetch the latest version for each package. - stream::iter(lock.packages()) - .filter_map(|package| async { - let index = package.index(workspace.install_path()).ok()??; - let filename = client - .find_latest(package.name(), Some(&index)) - .await - .ok()??; - if filename.version() == package.version() { - None - } else { - Some((package.clone(), filename.into_version())) - } - }) - .collect::>() - .await + let mut fetches = futures::stream::iter(lock.packages().iter().filter_map(|package| { + // Filter to packages that are derived from a registry. + let index = match package.index(workspace.install_path()) { + Ok(Some(index)) => index, + Ok(None) => return None, + Err(err) => return Some(Err(err)), + }; + Some(Ok((package, index))) + })) + .map(|result| async move { + let (package, index) = result?; + let Some(filename) = client.find_latest(package.name(), Some(&index)).await? else { + return Ok(None); + }; + if filename.version() == package.version() { + return Ok(None); + } + Ok::, Error>(Some((package, filename.into_version()))) + }) + .buffer_unordered(concurrency.downloads); + + let mut map = PackageMap::default(); + while let Some(entry) = fetches.next().await.transpose()? { + if let Some((package, version)) = entry { + map.insert(package.clone(), version); + } + } + map } else { PackageMap::default() }; diff --git a/crates/uv/src/lib.rs b/crates/uv/src/lib.rs index cc32ad5e9..8ddecaa4a 100644 --- a/crates/uv/src/lib.rs +++ b/crates/uv/src/lib.rs @@ -615,6 +615,7 @@ async fn run(mut cli: Cli) -> Result { args.settings.keyring_provider, globals.allow_insecure_host, globals.connectivity, + globals.concurrency, args.settings.strict, args.settings.exclude_newer, args.settings.python.as_deref(),