mirrors/uv
1
0
Fork 0
mirror of https://github.com/astral-sh/uv.git synced 2025-07-07 21:35:00 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

Avoid race condition in OnceMap (#3987)

Browse source 
## Summary

Fixes a race condition in `OnceMap::wait_blocking` where the inserted
value could potentially be missed, leading to a deadlock. Fairly certain
this will resolve https://github.com/astral-sh/uv/issues/3724.
This commit is contained in:
Ibraheem Ahmed 2024-06-03 12:25:58 -04:00 • committed by GitHub
parent 29ea5d5d9a
commit 1ffe18d861
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 43 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

47
crates/once-map/src/lib.rs
View file

@ -1,5 +1,6 @@
use std::borrow::Borrow;
use std::hash::{BuildHasher, Hash, RandomState};
use std::pin::pin;
use std::sync::Arc;
use dashmap::DashMap;
@ -46,13 +47,24 @@ impl<K: Eq + Hash, V: Clone, H: BuildHasher + Clone> OnceMap<K, V, H> {
///
/// Will hang if [`OnceMap::done`] isn't called for this key.
pub async fn wait(&self, key: &K) -> Option<V> {
let notify = {
let entry = self.items.get(key)?;
match entry.value() {
Value::Filled(value) => Some(value.clone()),
Value::Waiting(notify) => {
let notify = notify.clone();
drop(entry);
notify.notified().await;
Value::Filled(value) => return Some(value.clone()),
Value::Waiting(notify) => notify.clone(),
}
};
// Register the waiter for calls to `notify_waiters`.
let notification = pin!(notify.notified());
// Make sure the value wasn't inserted in-between us checking the map and registering the waiter.
if let Value::Filled(value) = self.items.get(key).expect("map is append-only").value() {
return Some(value.clone());
};
// Wait until the value is inserted.
notification.await;
let entry = self.items.get(key).expect("map is append-only");
match entry.value() {
@ -60,20 +72,29 @@ impl<K: Eq + Hash, V: Clone, H: BuildHasher + Clone> OnceMap<K, V, H> {
Value::Waiting(_) => unreachable!("notify was called"),
}
}
}
}
/// Wait for the result of a job that is running, in a blocking context.
///
/// Will hang if [`OnceMap::done`] isn't called for this key.
pub fn wait_blocking(&self, key: &K) -> Option<V> {
let notify = {
let entry = self.items.get(key)?;
match entry.value() {
Value::Filled(value) => Some(value.clone()),
Value::Waiting(notify) => {
let notify = notify.clone();
drop(entry);
futures::executor::block_on(notify.notified());
Value::Filled(value) => return Some(value.clone()),
Value::Waiting(notify) => notify.clone(),
}
};
// Register the waiter for calls to `notify_waiters`.
let notification = pin!(notify.notified());
// Make sure the value wasn't inserted in-between us checking the map and registering the waiter.
if let Value::Filled(value) = self.items.get(key).expect("map is append-only").value() {
return Some(value.clone());
};
// Wait until the value is inserted.
futures::executor::block_on(notification);
let entry = self.items.get(key).expect("map is append-only");
match entry.value() {
@ -81,8 +102,6 @@ impl<K: Eq + Hash, V: Clone, H: BuildHasher + Clone> OnceMap<K, V, H> {
Value::Waiting(_) => unreachable!("notify was called"),
}
}
}
}
/// Return the result of a previous job, if any.
pub fn get<Q: ?Sized + Hash + Eq>(&self, key: &Q) -> Option<V>