Propagate cancellation errors in `OnceMap` (#1032)

## Summary

Ensures that if an operation is cancelled in one thread, we propagate it
to others rather than panicking.

Related to https://github.com/astral-sh/puffin/issues/1005.

Cargo.lock

@@ -1895,6 +1895,7 @@ name = "once-map"
 version = "0.0.1"
 dependencies = [
  "rustc-hash",
+ "thiserror",
  "waitmap",
 ]
 
@@ -2572,6 +2573,7 @@ dependencies = [
  "fs-err",
  "futures",
  "install-wheel-rs",
+ "once-map",
  "pep440_rs 0.3.12",
  "pep508_rs",
  "platform-tags",

crates/once-map/Cargo.toml

@@ -14,4 +14,5 @@ workspace = true
 
 [dependencies]
 rustc-hash = { workspace = true }
+thiserror = { workspace = true }
 waitmap = { workspace = true }

crates/once-map/src/lib.rs

@@ -54,14 +54,14 @@ 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<Q: ?Sized + Hash + Eq>(&self, key: &Q) -> Ref<'_, K, V, RandomState>
+    pub async fn wait<Q: ?Sized + Hash + Eq>(
+        &self,
+        key: &Q,
+    ) -> Result<Ref<'_, K, V, RandomState>, Error>
     where
         K: Borrow<Q> + for<'a> From<&'a Q>,
     {
-        self.wait_map
+        self.wait_map.wait(key).await.ok_or(Error::Canceled)
-            .wait(key)
-            .await
-            .expect("This operation is never cancelled")
     }
 
     /// Return the result of a previous job, if any.
@@ -88,3 +88,9 @@ impl<K: Eq + Hash + Clone, V> Default for OnceMap<K, V> {
         }
     }
 }
+
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[error("The operation was canceled")]
+    Canceled,
+}

crates/puffin-installer/Cargo.toml

@@ -16,6 +16,7 @@ workspace = true
 distribution-filename = { path = "../distribution-filename" }
 distribution-types = { path = "../distribution-types" }
 install-wheel-rs = { path = "../install-wheel-rs", default-features = false }
+once-map = { path = "../once-map" }
 pep440_rs = { path = "../pep440-rs" }
 pep508_rs = { path = "../pep508-rs" }
 platform-tags = { path = "../platform-tags" }

crates/puffin-installer/src/downloader.rs

@@ -29,6 +29,8 @@ pub enum Error {
     Editable(#[from] DistributionDatabaseError),
     #[error("Unzip failed in another thread: {0}")]
     Thread(String),
+    #[error(transparent)]
+    OnceMap(#[from] once_map::Error),
 }
 
 /// Download, build, and unzip a set of distributions.
@@ -176,7 +178,7 @@ impl<'a, Context: BuildContext + Send + Sync> Downloader<'a, Context> {
             in_flight
                 .downloads
                 .wait(&id)
-                .await
+                .await?
                 .value()
                 .clone()
                 .map_err(Error::Thread)?

crates/puffin-resolver/src/error.rs

@@ -36,6 +36,9 @@ pub enum ResolveError {
     #[error(transparent)]
     Join(#[from] tokio::task::JoinError),
 
+    #[error(transparent)]
+    OnceMap(#[from] once_map::Error),
+
     #[error("Package metadata name `{metadata}` does not match given name `{given}`")]
     NameMismatch {
         given: PackageName,

crates/puffin-resolver/src/resolver/mod.rs

@@ -476,7 +476,7 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
                 } else {
                     // Otherwise, assume this is a source distribution.
                     let dist = PubGrubDistribution::from_url(package_name, url);
-                    let entry = self.index.distributions.wait(&dist.package_id()).await;
+                    let entry = self.index.distributions.wait(&dist.package_id()).await?;
                     let metadata = entry.value();
                     let version = &metadata.version;
                     if range.contains(version) {
@@ -489,7 +489,7 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
 
             PubGrubPackage::Package(package_name, extra, None) => {
                 // Wait for the metadata to be available.
-                let entry = self.index.packages.wait(package_name).await;
+                let entry = self.index.packages.wait(package_name).await?;
                 let version_map = entry.value();
 
                 if let Some(extra) = extra {
@@ -638,7 +638,7 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
                     return Ok(Dependencies::Known(constraints));
                 }
 
-                let entry = self.index.distributions.wait(&package_id).await;
+                let entry = self.index.distributions.wait(&package_id).await?;
                 let metadata = entry.value();
 
                 let mut constraints = PubGrubDependencies::from_requirements(
@@ -763,7 +763,7 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
             // Pre-fetch the package and distribution metadata.
             Request::Prefetch(package_name, range) => {
                 // Wait for the package metadata to become available.
-                let entry = self.index.packages.wait(&package_name).await;
+                let entry = self.index.packages.wait(&package_name).await?;
                 let version_map = entry.value();
 
                 // Try to find a compatible version. If there aren't any compatible versions,