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, } + +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 OnceMap { { 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 Default for OnceMap {