diff --git a/crates/once-map/src/lib.rs b/crates/once-map/src/lib.rs index 0ac543981..990626f6e 100644 --- a/crates/once-map/src/lib.rs +++ b/crates/once-map/src/lib.rs @@ -11,11 +11,14 @@ use tokio::sync::Notify; /// requests for metadata. When multiple tasks start the same query in parallel, e.g. through source /// dist builds, we want to wait until the other task is done and get a reference to the same /// result. -pub struct OnceMap { +/// +/// 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>, } -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 @@ -34,8 +37,7 @@ impl OnceMap { /// Submit the result of a job you registered. pub fn done(&self, key: K, value: V) { - if let Some(Value::Waiting(notify)) = self.items.insert(key, Value::Filled(Arc::new(value))) - { + if let Some(Value::Waiting(notify)) = self.items.insert(key, Value::Filled(value)) { notify.notify_waiters(); } } @@ -43,7 +45,7 @@ impl OnceMap { /// Wait for the result of a job that is running. /// /// Will hang if [`OnceMap::done`] isn't called for this key. - pub async fn wait(&self, key: &K) -> Option> { + pub async fn wait(&self, key: &K) -> Option { let entry = self.items.get(key)?; match entry.value() { Value::Filled(value) => Some(value.clone()), @@ -62,7 +64,7 @@ impl OnceMap { } /// Return the result of a previous job, if any. - pub fn get(&self, key: &Q) -> Option> + pub fn get(&self, key: &Q) -> Option where K: Borrow, { @@ -84,5 +86,5 @@ impl Default for OnceMap { enum Value { Waiting(Arc), - Filled(Arc), + Filled(V), } diff --git a/crates/uv-auth/src/middleware.rs b/crates/uv-auth/src/middleware.rs index facc0dd41..da8a625d7 100644 --- a/crates/uv-auth/src/middleware.rs +++ b/crates/uv-auth/src/middleware.rs @@ -312,13 +312,12 @@ impl AuthMiddleware { ); if !self.cache().fetches.register(key.clone()) { - let credentials = Arc::<_>::unwrap_or_clone( - self.cache() - .fetches - .wait(&key) - .await - .expect("The key must exist after register is called"), - ); + let credentials = self + .cache() + .fetches + .wait(&key) + .await + .expect("The key must exist after register is called"); if credentials.is_some() { trace!("Using credentials from previous fetch for {url}"); diff --git a/crates/uv-resolver/src/error.rs b/crates/uv-resolver/src/error.rs index ec519a390..6bc327f78 100644 --- a/crates/uv-resolver/src/error.rs +++ b/crates/uv-resolver/src/error.rs @@ -232,7 +232,7 @@ impl NoSolutionError { mut self, python_requirement: &PythonRequirement, visited: &DashSet, - package_versions: &OnceMap, + package_versions: &OnceMap>, ) -> Self { let mut available_versions = IndexMap::default(); for package in self.derivation_tree.packages() { diff --git a/crates/uv-resolver/src/resolution.rs b/crates/uv-resolver/src/resolution.rs index 535417a85..614697b8d 100644 --- a/crates/uv-resolver/src/resolution.rs +++ b/crates/uv-resolver/src/resolution.rs @@ -1,5 +1,6 @@ use std::borrow::Cow; use std::hash::BuildHasherDefault; +use std::sync::Arc; use anyhow::Result; use itertools::Itertools; @@ -71,8 +72,8 @@ impl ResolutionGraph { pub(crate) fn from_state( selection: &SelectedDependencies, pins: &FilePins, - packages: &OnceMap, - distributions: &OnceMap, + packages: &OnceMap>, + distributions: &OnceMap>, 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 614baf206..aae0a3129 100644 --- a/crates/uv-resolver/src/resolver/index.rs +++ b/crates/uv-resolver/src/resolver/index.rs @@ -11,21 +11,21 @@ use crate::resolver::provider::{MetadataResponse, VersionsResponse}; pub struct InMemoryIndex { /// A map from package name to the metadata for that package and the index where the metadata /// came from. - pub(crate) packages: OnceMap, + pub(crate) packages: OnceMap>, /// A map from package ID to metadata for that distribution. - pub(crate) distributions: OnceMap, + pub(crate) distributions: OnceMap>, } impl InMemoryIndex { /// Insert a [`VersionsResponse`] into the index. pub fn insert_package(&self, package_name: PackageName, response: VersionsResponse) { - self.packages.done(package_name, response); + self.packages.done(package_name, Arc::new(response)); } /// Insert a [`Metadata23`] into the index. pub fn insert_metadata(&self, version_id: VersionId, response: MetadataResponse) { - self.distributions.done(version_id, response); + self.distributions.done(version_id, Arc::new(response)); } /// Get the [`VersionsResponse`] for a given package name, without waiting. diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs index 2a8c52f54..fc5055d3e 100644 --- a/crates/uv-resolver/src/resolver/mod.rs +++ b/crates/uv-resolver/src/resolver/mod.rs @@ -1046,13 +1046,15 @@ impl< match response? { Some(Response::Package(package_name, version_map)) => { trace!("Received package metadata for: {package_name}"); - self.index.packages.done(package_name, version_map); + self.index + .packages + .done(package_name, Arc::new(version_map)); } Some(Response::Installed { dist, metadata }) => { trace!("Received installed distribution metadata for: {dist}"); self.index.distributions.done( dist.version_id(), - MetadataResponse::Found(ArchiveMetadata::from(metadata)), + Arc::new(MetadataResponse::Found(ArchiveMetadata::from(metadata))), ); } Some(Response::Dist { @@ -1069,7 +1071,9 @@ impl< } _ => {} } - self.index.distributions.done(dist.version_id(), metadata); + self.index + .distributions + .done(dist.version_id(), Arc::new(metadata)); } Some(Response::Dist { dist: Dist::Source(dist), @@ -1085,7 +1089,9 @@ impl< } _ => {} } - self.index.distributions.done(dist.version_id(), metadata); + self.index + .distributions + .done(dist.version_id(), Arc::new(metadata)); } None => {} }