From cfffcbb2697abe36f1841c73fafbb00f584da99b Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Wed, 3 Jan 2024 16:18:27 -0400
Subject: [PATCH] Cancel waiting tasks on resolution error (#753)

## Summary

I don't understand why this works (because I don't understand why it's
erroring) but it does. See:
https://github.com/astral-sh/puffin/pull/746#issuecomment-1875722454.

## Test Plan

```
cargo run --bin puffin pip-install requires-transitive-incompatible-with-transitive-8329cfc0 --extra-index-url https://test.pypi.org/simple -n
```
---
 crates/puffin-resolver/src/resolver/index.rs | 12 ++++++++++++
 crates/puffin-resolver/src/resolver/mod.rs   |  6 +++++-
 crates/puffin-traits/src/once_map.rs         |  7 +++++++
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/crates/puffin-resolver/src/resolver/index.rs b/crates/puffin-resolver/src/resolver/index.rs
index 5bb42d811..ef7116b2f 100644
--- a/crates/puffin-resolver/src/resolver/index.rs
+++ b/crates/puffin-resolver/src/resolver/index.rs
@@ -24,3 +24,15 @@ pub(crate) struct Index {
     /// A map from source URL to precise URL.
     pub(crate) redirects: OnceMap<Url, Url>,
 }
+
+impl Index {
+    /// Cancel all waiting tasks.
+    ///
+    /// Warning: waiting on tasks that have been canceled will cause the index to hang.
+    pub(crate) fn cancel_all(&self) {
+        self.packages.cancel_all();
+        self.distributions.cancel_all();
+        self.incompatibilities.cancel_all();
+        self.redirects.cancel_all();
+    }
+}
diff --git a/crates/puffin-resolver/src/resolver/mod.rs b/crates/puffin-resolver/src/resolver/mod.rs
index b47c33b8c..55ed080e2 100644
--- a/crates/puffin-resolver/src/resolver/mod.rs
+++ b/crates/puffin-resolver/src/resolver/mod.rs
@@ -202,7 +202,11 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
             }
             resolution = resolve_fut => {
                 resolution.map_err(|err| {
-                    // Add version information to improve unsat error messages
+                    // Ensure that any waiting tasks are cancelled prior to accessing any of the
+                    // index entries.
+                    self.index.cancel_all();
+
+                    // Add version information to improve unsat error messages.
                     if let ResolveError::NoSolution(err) = err {
                         ResolveError::NoSolution(err.with_available_versions(&self.python_requirement, &self.index.packages).with_selector(self.selector.clone()))
                     } else {
diff --git a/crates/puffin-traits/src/once_map.rs b/crates/puffin-traits/src/once_map.rs
index a90436098..3b56722e7 100644
--- a/crates/puffin-traits/src/once_map.rs
+++ b/crates/puffin-traits/src/once_map.rs
@@ -71,6 +71,13 @@ impl<K: Eq + Hash, V> OnceMap<K, V> {
     {
         self.wait_map.get(key)
     }
+
+    /// Cancel all waiting tasks.
+    ///
+    /// Warning: waiting on tasks that have been canceled will cause the map to hang.
+    pub fn cancel_all(&self) {
+        self.wait_map.cancel_all();
+    }
 }
 
 impl<K: Eq + Hash + Clone, V> Default for OnceMap<K, V> {