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

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

@ -1,5 +1,6 @@
use std::borrow::Borrow; use std::borrow::Borrow;
use std::hash::{BuildHasher, Hash, RandomState}; use std::hash::{BuildHasher, Hash, RandomState};
use std::pin::pin;
use std::sync::Arc; use std::sync::Arc;
use dashmap::DashMap; 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. /// Will hang if [`OnceMap::done`] isn't called for this key.
pub async fn wait(&self, key: &K) -> Option<V> { 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() { match entry.value() {
Value::Filled(value) => Some(value.clone()), Value::Filled(value) => Some(value.clone()),
Value::Waiting(notify) => { Value::Waiting(_) => unreachable!("notify was called"),
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"),
}
}
} }
} }
@ -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. /// Will hang if [`OnceMap::done`] isn't called for this key.
pub fn wait_blocking(&self, key: &K) -> Option<V> { 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() { match entry.value() {
Value::Filled(value) => Some(value.clone()), Value::Filled(value) => Some(value.clone()),
Value::Waiting(notify) => { Value::Waiting(_) => unreachable!("notify was called"),
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"),
}
}
} }
} }