once-map: avoid hard-coding `Arc` (#3242)

The only thing a `OnceMap` really needs to be able to do with the value
is to clone it. All extant uses benefited from having this done for them
by automatically wrapping values in an `Arc`. But this isn't necessarily
true for all things. For example, a value might have an `Arc` internally
to making cloning cheap in other contexts, and it doesn't make sense to
re-wrap it in an `Arc` just to use it with a `OnceMap`. Or
alternatively, cloning might just be cheap enough on its own that an
`Arc` isn't worth it.

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<K: Eq + Hash, V> {
+///
+/// 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<V>` to make cloning cheap.
+pub struct OnceMap<K, V> {
     items: DashMap<K, Value<V>>,
 }
 
-impl<K: Eq + Hash, V> OnceMap<K, V> {
+impl<K: Eq + Hash, V: Clone> OnceMap<K, V> {
     /// 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<K: Eq + Hash, V> OnceMap<K, V> {
 
     /// 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<K: Eq + Hash, V> OnceMap<K, V> {
     /// 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<Arc<V>> {
+    pub async fn wait(&self, key: &K) -> Option<V> {
         let entry = self.items.get(key)?;
         match entry.value() {
             Value::Filled(value) => Some(value.clone()),
@@ -62,7 +64,7 @@ impl<K: Eq + Hash, V> OnceMap<K, V> {
     }
 
     /// Return the result of a previous job, if any.
-    pub fn get<Q: ?Sized + Hash + Eq>(&self, key: &Q) -> Option<Arc<V>>
+    pub fn get<Q: ?Sized + Hash + Eq>(&self, key: &Q) -> Option<V>
     where
         K: Borrow<Q>,
     {
@@ -84,5 +86,5 @@ impl<K: Eq + Hash + Clone, V> Default for OnceMap<K, V> {
 
 enum Value<V> {
     Waiting(Arc<Notify>),
-    Filled(Arc<V>),
+    Filled(V),
 }

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}");

crates/uv-resolver/src/error.rs

@@ -232,7 +232,7 @@ impl NoSolutionError {
         mut self,
         python_requirement: &PythonRequirement,
         visited: &DashSet<PackageName>,
-        package_versions: &OnceMap<PackageName, VersionsResponse>,
+        package_versions: &OnceMap<PackageName, Arc<VersionsResponse>>,
     ) -> Self {
         let mut available_versions = IndexMap::default();
         for package in self.derivation_tree.packages() {

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<UvDependencyProvider>,
         pins: &FilePins,
-        packages: &OnceMap<PackageName, VersionsResponse>,
-        distributions: &OnceMap<VersionId, MetadataResponse>,
+        packages: &OnceMap<PackageName, Arc<VersionsResponse>>,
+        distributions: &OnceMap<VersionId, Arc<MetadataResponse>>,
         state: &State<UvDependencyProvider>,
         preferences: &Preferences,
         editables: Editables,

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<PackageName, VersionsResponse>,
+    pub(crate) packages: OnceMap<PackageName, Arc<VersionsResponse>>,
 
     /// A map from package ID to metadata for that distribution.
-    pub(crate) distributions: OnceMap<VersionId, MetadataResponse>,
+    pub(crate) distributions: OnceMap<VersionId, Arc<MetadataResponse>>,
 }
 
 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.

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 => {}
             }