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:
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

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,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"),
}
}