mirror of
https://github.com/astral-sh/uv.git
synced 2025-07-07 13:25:00 +00:00
Avoid race condition in OnceMap
(#3987)
## 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:
parent
29ea5d5d9a
commit
1ffe18d861
1 changed files with 43 additions and 24 deletions
|
@ -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,20 +47,29 @@ 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 entry = self.items.get(key)?;
|
||||
let notify = {
|
||||
let entry = self.items.get(key)?;
|
||||
match entry.value() {
|
||||
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() {
|
||||
Value::Filled(value) => Some(value.clone()),
|
||||
Value::Waiting(notify) => {
|
||||
let notify = notify.clone();
|
||||
drop(entry);
|
||||
notify.notified().await;
|
||||
|
||||
let entry = self.items.get(key).expect("map is append-only");
|
||||
match entry.value() {
|
||||
Value::Filled(value) => Some(value.clone()),
|
||||
Value::Waiting(_) => unreachable!("notify was called"),
|
||||
}
|
||||
}
|
||||
Value::Waiting(_) => unreachable!("notify was called"),
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -67,20 +77,29 @@ 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 fn wait_blocking(&self, key: &K) -> Option<V> {
|
||||
let entry = self.items.get(key)?;
|
||||
let notify = {
|
||||
let entry = self.items.get(key)?;
|
||||
match entry.value() {
|
||||
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() {
|
||||
Value::Filled(value) => Some(value.clone()),
|
||||
Value::Waiting(notify) => {
|
||||
let notify = notify.clone();
|
||||
drop(entry);
|
||||
futures::executor::block_on(notify.notified());
|
||||
|
||||
let entry = self.items.get(key).expect("map is append-only");
|
||||
match entry.value() {
|
||||
Value::Filled(value) => Some(value.clone()),
|
||||
Value::Waiting(_) => unreachable!("notify was called"),
|
||||
}
|
||||
}
|
||||
Value::Waiting(_) => unreachable!("notify was called"),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue