From 509dc83fd35bed54826839cd6a1ac2296fba976c Mon Sep 17 00:00:00 2001 From: konsti Date: Wed, 11 Dec 2024 15:46:36 +0100 Subject: [PATCH] Resolver module improvements (#9773) Further small refactorings for the resolver. --- .../src/resolver/batch_prefetch.rs | 33 +++-- crates/uv-resolver/src/resolver/mod.rs | 115 +++++++++--------- 2 files changed, 84 insertions(+), 64 deletions(-) diff --git a/crates/uv-resolver/src/resolver/batch_prefetch.rs b/crates/uv-resolver/src/resolver/batch_prefetch.rs index 957f53cb3840..bfc9012ec3ab 100644 --- a/crates/uv-resolver/src/resolver/batch_prefetch.rs +++ b/crates/uv-resolver/src/resolver/batch_prefetch.rs @@ -39,13 +39,31 @@ enum BatchPrefetchStrategy { /// have to fetch the metadata for a lot of versions. /// /// Note that these all heuristics that could totally prefetch lots of irrelevant versions. -#[derive(Default)] pub(crate) struct BatchPrefetcher { + // Internal types. tried_versions: FxHashMap, last_prefetch: FxHashMap, + // Shared (e.g., `Arc`) types. + capabilities: IndexCapabilities, + index: InMemoryIndex, + request_sink: Sender, } impl BatchPrefetcher { + pub(crate) fn new( + capabilities: IndexCapabilities, + index: InMemoryIndex, + request_sink: Sender, + ) -> Self { + Self { + tried_versions: FxHashMap::default(), + last_prefetch: FxHashMap::default(), + capabilities, + index, + request_sink, + } + } + /// Prefetch a large number of versions if we already unsuccessfully tried many versions. pub(crate) fn prefetch_batches( &mut self, @@ -55,9 +73,6 @@ impl BatchPrefetcher { current_range: &Range, unchangeable_constraints: Option<&Term>>, python_requirement: &PythonRequirement, - request_sink: &Sender, - in_memory: &InMemoryIndex, - capabilities: &IndexCapabilities, selector: &CandidateSelector, env: &ResolverEnvironment, ) -> anyhow::Result<(), ResolveError> { @@ -79,12 +94,12 @@ impl BatchPrefetcher { // This is immediate, we already fetched the version map. let versions_response = if let Some(index) = index { - in_memory + self.index .explicit() .wait_blocking(&(name.clone(), index.clone())) .ok_or_else(|| ResolveError::UnregisteredTask(name.to_string()))? } else { - in_memory + self.index .implicit() .wait_blocking(name) .ok_or_else(|| ResolveError::UnregisteredTask(name.to_string()))? @@ -166,7 +181,7 @@ impl BatchPrefetcher { // Avoid prefetching built distributions that don't support _either_ PEP 658 (`.metadata`) // or range requests. if !(wheel.file.dist_info_metadata - || capabilities.supports_range_requests(&wheel.index)) + || self.capabilities.supports_range_requests(&wheel.index)) { debug!("Abandoning prefetch for {wheel} due to missing registry capabilities"); return Ok(()); @@ -214,9 +229,9 @@ impl BatchPrefetcher { ); prefetch_count += 1; - if in_memory.distributions().register(candidate.version_id()) { + if self.index.distributions().register(candidate.version_id()) { let request = Request::from(dist); - request_sink.blocking_send(request)?; + self.request_sink.blocking_send(request)?; } } diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs index fe87350e1cf1..91ef84257055 100644 --- a/crates/uv-resolver/src/resolver/mod.rs +++ b/crates/uv-resolver/src/resolver/mod.rs @@ -306,7 +306,11 @@ impl ResolverState ResolverState ResolverState None, - CompatibleDist::SourceDist { sdist, .. } - | CompatibleDist::IncompatibleWheel { sdist, .. } => { - sdist.file.requires_python.as_ref() - } - CompatibleDist::CompatibleWheel { wheel, .. } => wheel.file.requires_python.as_ref(), - }; - let incompatibility = requires_python.and_then(|requires_python| { - if python_requirement.installed() == python_requirement.target() { - if !python_requirement - .installed() - .is_contained_by(requires_python) - { - return if matches!(dist, CompatibleDist::CompatibleWheel { .. }) { - Some(IncompatibleDist::Wheel(IncompatibleWheel::RequiresPython( - requires_python.clone(), - PythonRequirementKind::Installed, - ))) - } else { - Some(IncompatibleDist::Source( - IncompatibleSource::RequiresPython( - requires_python.clone(), - PythonRequirementKind::Installed, - ), - )) - }; - } - } else { - if !python_requirement.target().is_contained_by(requires_python) { - return if matches!(dist, CompatibleDist::CompatibleWheel { .. }) { - Some(IncompatibleDist::Wheel(IncompatibleWheel::RequiresPython( - requires_python.clone(), - PythonRequirementKind::Target, - ))) - } else { - Some(IncompatibleDist::Source( - IncompatibleSource::RequiresPython( - requires_python.clone(), - PythonRequirementKind::Target, - ), - )) - }; - } - } - None - }); - - // The version is incompatible due to its Python requirement. - if let Some(incompatibility) = incompatibility { + // Check whether the version is incompatible due to its Python requirement. + if let Some(incompatibility) = Self::check_requires_python(dist, python_requirement) { return Ok(Some(ResolverVersion::Unavailable( candidate.version().clone(), UnavailableVersion::IncompatibleDist(incompatibility), @@ -1180,6 +1132,59 @@ impl ResolverState Option { + let requires_python = match dist { + CompatibleDist::InstalledDist(_) => None, + CompatibleDist::SourceDist { sdist, .. } + | CompatibleDist::IncompatibleWheel { sdist, .. } => { + sdist.file.requires_python.as_ref() + } + CompatibleDist::CompatibleWheel { wheel, .. } => wheel.file.requires_python.as_ref(), + }?; + if python_requirement.installed() == python_requirement.target() { + if !python_requirement + .installed() + .is_contained_by(requires_python) + { + return if matches!(dist, CompatibleDist::CompatibleWheel { .. }) { + Some(IncompatibleDist::Wheel(IncompatibleWheel::RequiresPython( + requires_python.clone(), + PythonRequirementKind::Installed, + ))) + } else { + Some(IncompatibleDist::Source( + IncompatibleSource::RequiresPython( + requires_python.clone(), + PythonRequirementKind::Installed, + ), + )) + }; + } + } else { + if !python_requirement.target().is_contained_by(requires_python) { + return if matches!(dist, CompatibleDist::CompatibleWheel { .. }) { + Some(IncompatibleDist::Wheel(IncompatibleWheel::RequiresPython( + requires_python.clone(), + PythonRequirementKind::Target, + ))) + } else { + Some(IncompatibleDist::Source( + IncompatibleSource::RequiresPython( + requires_python.clone(), + PythonRequirementKind::Target, + ), + )) + }; + } + } + None + } + /// Given a candidate package and version, return its dependencies. #[instrument(skip_all, fields(%package, %version))] fn get_dependencies_forking(