diff --git a/crates/once-map/src/lib.rs b/crates/once-map/src/lib.rs index 3fcedba49..00286d227 100644 --- a/crates/once-map/src/lib.rs +++ b/crates/once-map/src/lib.rs @@ -1,5 +1,5 @@ use std::borrow::Borrow; -use std::hash::Hash; +use std::hash::{BuildHasher, Hash, RandomState}; use std::sync::Arc; use dashmap::DashMap; @@ -14,11 +14,11 @@ use tokio::sync::Notify; /// /// Note that this always clones the value out of the underlying map. Because /// of this, it's common to wrap the `V` in an `Arc` to make cloning cheap. -pub struct OnceMap { - items: DashMap>, +pub struct OnceMap { + items: DashMap, H>, } -impl OnceMap { +impl OnceMap { /// Register that you want to start a job. /// /// If this method returns `true`, you need to start a job and call [`OnceMap::done`] eventually @@ -97,10 +97,10 @@ impl OnceMap { } } -impl Default for OnceMap { +impl Default for OnceMap { fn default() -> Self { Self { - items: DashMap::new(), + items: DashMap::with_hasher(H::default()), } } } diff --git a/crates/uv-resolver/src/error.rs b/crates/uv-resolver/src/error.rs index 5132d61e2..b2f01832c 100644 --- a/crates/uv-resolver/src/error.rs +++ b/crates/uv-resolver/src/error.rs @@ -6,11 +6,10 @@ use std::sync::Arc; use indexmap::IndexMap; use pubgrub::range::Range; use pubgrub::report::{DefaultStringReporter, DerivationTree, External, Reporter}; -use rustc_hash::FxHashMap; +use rustc_hash::{FxHashMap, FxHashSet}; -use dashmap::{DashMap, DashSet}; +use dashmap::DashMap; use distribution_types::{BuiltDist, IndexLocations, InstalledDist, ParsedUrlError, SourceDist}; -use once_map::OnceMap; use pep440_rs::Version; use pep508_rs::Requirement; use uv_normalize::PackageName; @@ -19,7 +18,9 @@ use crate::candidate_selector::CandidateSelector; use crate::dependency_provider::UvDependencyProvider; use crate::pubgrub::{PubGrubPackage, PubGrubPython, PubGrubReportFormatter}; use crate::python_requirement::PythonRequirement; -use crate::resolver::{IncompletePackage, UnavailablePackage, UnavailableReason, VersionsResponse}; +use crate::resolver::{ + FxOnceMap, IncompletePackage, UnavailablePackage, UnavailableReason, VersionsResponse, +}; #[derive(Debug, thiserror::Error)] pub enum ResolveError { @@ -235,8 +236,8 @@ impl NoSolutionError { pub(crate) fn with_available_versions( mut self, python_requirement: &PythonRequirement, - visited: &DashSet, - package_versions: &OnceMap>, + visited: &FxHashSet, + package_versions: &FxOnceMap>, ) -> Self { let mut available_versions = IndexMap::default(); for package in self.derivation_tree.packages() { diff --git a/crates/uv-resolver/src/resolution/graph.rs b/crates/uv-resolver/src/resolution/graph.rs index a3853deba..69646c18f 100644 --- a/crates/uv-resolver/src/resolution/graph.rs +++ b/crates/uv-resolver/src/resolution/graph.rs @@ -10,7 +10,6 @@ use distribution_types::{ Dist, DistributionMetadata, Name, ParsedUrlError, Requirement, ResolvedDist, VersionId, VersionOrUrlRef, }; -use once_map::OnceMap; use pep440_rs::{Version, VersionSpecifier}; use pep508_rs::MarkerEnvironment; use uv_normalize::{ExtraName, PackageName}; @@ -22,6 +21,7 @@ use crate::preferences::Preferences; use crate::pubgrub::{PubGrubDistribution, PubGrubPackage}; use crate::redirect::url_to_precise; use crate::resolution::AnnotatedDist; +use crate::resolver::FxOnceMap; use crate::{ lock, InMemoryIndex, Lock, LockError, Manifest, MetadataResponse, ResolveError, VersionsResponse, @@ -45,8 +45,8 @@ impl ResolutionGraph { pub(crate) fn from_state( selection: &SelectedDependencies, pins: &FilePins, - packages: &OnceMap>, - distributions: &OnceMap>, + packages: &FxOnceMap>, + distributions: &FxOnceMap>, state: &State, preferences: &Preferences, editables: Editables, diff --git a/crates/uv-resolver/src/resolver/index.rs b/crates/uv-resolver/src/resolver/index.rs index fedcbcffe..24e52c477 100644 --- a/crates/uv-resolver/src/resolver/index.rs +++ b/crates/uv-resolver/src/resolver/index.rs @@ -1,7 +1,9 @@ +use std::hash::BuildHasherDefault; use std::sync::Arc; use distribution_types::VersionId; use once_map::OnceMap; +use rustc_hash::FxHasher; use uv_normalize::PackageName; use crate::resolver::provider::{MetadataResponse, VersionsResponse}; @@ -14,20 +16,22 @@ pub struct InMemoryIndex(Arc); struct SharedInMemoryIndex { /// A map from package name to the metadata for that package and the index where the metadata /// came from. - packages: OnceMap>, + packages: FxOnceMap>, /// A map from package ID to metadata for that distribution. - distributions: OnceMap>, + distributions: FxOnceMap>, } +pub(crate) type FxOnceMap = OnceMap>; + impl InMemoryIndex { /// Returns a reference to the package metadata map. - pub fn packages(&self) -> &OnceMap> { + pub fn packages(&self) -> &FxOnceMap> { &self.0.packages } /// Returns a reference to the distribution metadata map. - pub fn distributions(&self) -> &OnceMap> { + pub fn distributions(&self) -> &FxOnceMap> { &self.0.distributions } } diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs index 87f9f636c..7c194a0ae 100644 --- a/crates/uv-resolver/src/resolver/mod.rs +++ b/crates/uv-resolver/src/resolver/mod.rs @@ -7,8 +7,8 @@ use std::sync::Arc; use std::thread; use anyhow::Result; -use dashmap::{DashMap, DashSet}; -use futures::{FutureExt, StreamExt}; +use dashmap::DashMap; +use futures::{FutureExt, StreamExt, TryFutureExt}; use itertools::Itertools; use pubgrub::error::PubGrubError; use pubgrub::range::Range; @@ -49,6 +49,7 @@ use crate::pubgrub::{ use crate::python_requirement::PythonRequirement; use crate::resolution::ResolutionGraph; use crate::resolver::batch_prefetch::BatchPrefetcher; +pub(crate) use crate::resolver::index::FxOnceMap; pub use crate::resolver::index::InMemoryIndex; pub use crate::resolver::provider::{ DefaultResolverProvider, MetadataResponse, PackageVersionsResult, ResolverProvider, @@ -202,8 +203,6 @@ struct ResolverState { unavailable_packages: DashMap, /// Incompatibilities for packages that are unavailable at specific versions. incomplete_packages: DashMap>, - /// The set of all registry-based packages visited during resolution. - visited: DashSet, reporter: Option>, } @@ -286,7 +285,6 @@ impl index: index.clone(), unavailable_packages: DashMap::default(), incomplete_packages: DashMap::default(), - visited: DashSet::default(), selector: CandidateSelector::for_resolution(options, &manifest, markers), dependency_mode: options.dependency_mode, urls: Urls::from_manifest(&manifest, markers, options.dependency_mode)?, @@ -332,7 +330,11 @@ impl let (request_sink, request_stream) = mpsc::channel(300); // Run the fetcher. - let requests_fut = state.clone().fetch(provider.clone(), request_stream).fuse(); + let requests_fut = state + .clone() + .fetch(provider.clone(), request_stream) + .map_err(|err| (err, FxHashSet::default())) + .fuse(); // Spawn the PubGrub solver on a dedicated thread. let solver = state.clone(); @@ -347,7 +349,7 @@ impl let resolve_fut = async move { rx.await - .map_err(|_| ResolveError::ChannelClosed) + .map_err(|_| (ResolveError::ChannelClosed, FxHashSet::default())) .and_then(|result| result) }; @@ -357,13 +359,13 @@ impl state.on_complete(); Ok(resolution) } - Err(err) => { + Err((err, visited)) => { // Add version information to improve unsat error messages. Err(if let ResolveError::NoSolution(err) = err { ResolveError::NoSolution( err.with_available_versions( &state.python_requirement, - &state.visited, + &visited, state.index.packages(), ) .with_selector(state.selector.clone()) @@ -381,11 +383,23 @@ impl } impl ResolverState { - /// Run the PubGrub solver. #[instrument(skip_all)] fn solve( self: Arc, request_sink: Sender, + ) -> Result)> { + let mut visited = FxHashSet::default(); + self.solve_tracked(&mut visited, request_sink) + .map_err(|err| (err, visited)) + } + + /// Run the PubGrub solver, updating the `visited` set for each package visited during + /// resolution. + #[instrument(skip_all)] + fn solve_tracked( + self: Arc, + visited: &mut FxHashSet, + request_sink: Sender, ) -> Result { let root = PubGrubPackage::Root(self.project.clone()); let mut prefetcher = BatchPrefetcher::default(); @@ -450,6 +464,7 @@ impl ResolverState ResolverState, pins: &mut FilePins, + visited: &mut FxHashSet, request_sink: &Sender, ) -> Result, ResolveError> { match package { @@ -805,7 +821,7 @@ impl ResolverState version_maps.as_slice(),