From 9e443ca37d762f5eea3ce8af1c9a7ff56a262eae Mon Sep 17 00:00:00 2001 From: tison Date: Thu, 5 Nov 2026 21:03:39 +0800 Subject: [PATCH 1/2] fix: reclaim entries when duplicate key destructors panic --- CHANGELOG.md | 1 + asyncband/src/once/once_map/mod.rs | 59 ++++++++--------- asyncband/src/once/once_map/tests.rs | 95 ++++++++++++++++++++++++++-- asyncband/src/singleflight/mod.rs | 21 +++--- asyncband/src/singleflight/tests.rs | 48 ++++++++++++++ asyncband/src/test_support.rs | 43 +++++++++++++ 6 files changed, 225 insertions(+), 42 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5d5452e4..91c00408 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -37,6 +37,7 @@ All notable changes to this project will be documented in this file. ### Bug fixes * Prevent `OnceMap` from keeping an empty entry when a computation fails, panics, or is cancelled while a concurrent `get` looks up the same key. +* Prevent `OnceMap` and `singleflight::Group` from leaking entries or values when a duplicate key's destructor panics, including when the original computation is cancelled during that destructor. ### Improvements diff --git a/asyncband/src/once/once_map/mod.rs b/asyncband/src/once/once_map/mod.rs index 9bf434ae..8e660279 100644 --- a/asyncband/src/once/once_map/mod.rs +++ b/asyncband/src/once/once_map/mod.rs @@ -38,9 +38,13 @@ struct Entry { cell: OnceCell, } -enum Lookup { +enum Lookup<'a, K, V, S> +where + K: Eq + Hash, + S: BuildHasher, +{ Ready(V), - Pending(Arc>), + Pending(ComputeCleanupGuard<'a, K, V, S>), } /// A hash map that runs computation only once for each key and stores the result. @@ -75,39 +79,38 @@ where K: Eq + Hash, S: BuildHasher, { - fn get_or_insert(&self, key: K) -> Lookup + fn get_or_insert(&self, key: K) -> Lookup<'_, K, V, S> where V: Clone, { let hash = self.hasher.hash_one(&key); - let entry = { - let mut entries = self.entries.lock(); - if let Some(entry) = entries - .find(hash, |entry| entry.key.eq(&key)) - .map(Arc::clone) - { - entry - } else { - let entry = Arc::new(Entry { - hash, - key, - cell: OnceCell::new(), - }); - entries.insert_unique(hash, Arc::clone(&entry), |entry| entry.hash); - entry - } - }; + let mut entries = self.entries.lock(); + if let Some(entry) = entries.find(hash, |entry| entry.key.eq(&key)) { + let entry = Arc::clone(entry); + drop(entries); + let lookup = self.classify(entry); + // Own the value or cleanup guard before a duplicate key can unwind or reenter. + drop(key); + return lookup; + } - Self::classify(entry) + let entry = Arc::new(Entry { + hash, + key, + cell: OnceCell::new(), + }); + entries.insert_unique(hash, Arc::clone(&entry), |entry| entry.hash); + drop(entries); + self.classify(entry) } - fn classify(entry: Arc>) -> Lookup + fn classify(&self, entry: Arc>) -> Lookup<'_, K, V, S> where V: Clone, { match entry.cell.get().cloned() { Some(value) => Lookup::Ready(value), - None => Lookup::Pending(entry), + None => Lookup::Pending(ComputeCleanupGuard::new(self, entry)), } } @@ -309,12 +312,11 @@ where where F: AsyncFnOnce() -> V, { - let entry = match self.get_or_insert(key) { + let guard = match self.get_or_insert(key) { Lookup::Ready(value) => return value, - Lookup::Pending(entry) => entry, + Lookup::Pending(guard) => guard, }; - let guard = ComputeCleanupGuard::new(self, entry); let result = guard.entry().cell.get_or_init(func).await.clone(); guard.dismiss(); result @@ -336,12 +338,11 @@ where where F: AsyncFnOnce() -> Result, { - let entry = match self.get_or_insert(key) { + let guard = match self.get_or_insert(key) { Lookup::Ready(value) => return Ok(value), - Lookup::Pending(entry) => entry, + Lookup::Pending(guard) => guard, }; - let guard = ComputeCleanupGuard::new(self, entry); let result = guard.entry().cell.get_or_try_init(func).await?.clone(); guard.dismiss(); Ok(result) diff --git a/asyncband/src/once/once_map/tests.rs b/asyncband/src/once/once_map/tests.rs index 2aee5af9..236523c6 100644 --- a/asyncband/src/once/once_map/tests.rs +++ b/asyncband/src/once/once_map/tests.rs @@ -16,6 +16,7 @@ // under the License. use std::cell::Cell; +use std::cell::RefCell; use std::future::pending; use std::future::poll_fn; use std::hash::BuildHasherDefault; @@ -23,9 +24,11 @@ use std::hash::Hasher; use std::panic::AssertUnwindSafe; use std::panic::catch_unwind; use std::pin::pin; +use std::rc::Rc; use std::task::Poll; use super::*; +use crate::test_support::KeyWithDrop; use crate::test_support::poll_once; // These tests stay next to the implementation because they inspect private state. @@ -157,11 +160,95 @@ fn colliding_pending_entries_are_tracked_independently() { unreachable!() }; - assert!(Arc::ptr_eq(&first, &first_waiter)); - assert!(!Arc::ptr_eq(&first, &second)); + assert!(Arc::ptr_eq(first.entry(), first_waiter.entry())); + assert!(!Arc::ptr_eq(first.entry(), second.entry())); drop(first_waiter); - map.cleanup_abandoned_entry(first); - map.cleanup_abandoned_entry(second); + drop(first); + drop(second); assert_eq!(map.len(), 0); } + +#[test] +fn duplicate_key_panic_cleans_up_pending_compute() { + for fallible in [false, true] { + for cancel_during_drop in [false, true] { + let map = Rc::new(OnceMap::::new()); + let first_map = Rc::clone(&map); + let mut first = Box::pin(async move { + first_map + .compute( + KeyWithDrop { + key: "key", + on_drop: None, + }, + async || pending().await, + ) + .await + }); + assert!(poll_once(first.as_mut()).is_pending()); + let first = Rc::new(RefCell::new(Some(first))); + let dropping_first = Rc::clone(&first); + let dropping_map = Rc::clone(&map); + let key = KeyWithDrop { + key: "key", + on_drop: Some(Box::new(move || { + // Reenter the map and optionally cancel its last other caller before unwinding. + assert_eq!(dropping_map.len(), 1); + if cancel_during_drop { + drop(dropping_first.borrow_mut().take()); + } + panic!("duplicate key destructor"); + })), + }; + let result = catch_unwind(AssertUnwindSafe(|| { + if fallible { + let _ = poll_once(pin!(map.try_compute(key, async || Ok::<_, ()>(1)))); + } else { + let _ = poll_once(pin!(map.compute(key, async || 1))); + } + })); + assert!(result.is_err()); + drop(first.borrow_mut().take()); + assert_eq!(map.len(), 0); + } + } +} + +#[test] +fn duplicate_key_panic_drops_cached_value_clone() { + for fallible in [false, true] { + let map = Rc::new(OnceMap::new()); + let value = Arc::new(()); + assert!( + poll_once(pin!(map.compute( + KeyWithDrop { + key: "key", + on_drop: None + }, + async || Arc::clone(&value), + ))) + .is_ready() + ); + let dropping_map = Rc::clone(&map); + let key = KeyWithDrop { + key: "key", + on_drop: Some(Box::new(move || { + dropping_map.discard("key"); + panic!("duplicate key destructor"); + })), + }; + let result = catch_unwind(AssertUnwindSafe(|| { + if fallible { + let _ = poll_once(pin!( + map.try_compute(key, async || Ok::<_, ()>(Arc::clone(&value))) + )); + } else { + let _ = poll_once(pin!(map.compute(key, async || Arc::clone(&value)))); + } + })); + assert!(result.is_err()); + assert_eq!(map.len(), 0); + assert_eq!(Arc::strong_count(&value), 1); + } +} diff --git a/asyncband/src/singleflight/mod.rs b/asyncband/src/singleflight/mod.rs index 104193ac..91763df9 100644 --- a/asyncband/src/singleflight/mod.rs +++ b/asyncband/src/singleflight/mod.rs @@ -71,12 +71,16 @@ where K: Eq + Hash, S: BuildHasher, { - fn get_or_insert(&self, key: K) -> Arc> { + fn get_or_insert(&self, key: K) -> WorkCleanupGuard<'_, K, V, S> { let hash = self.hasher.hash_one(&key); let mut entries = self.entries.lock(); - // Drop duplicate keys after unlocking: their destructors may reenter the group. if let Some(entry) = entries.find(hash, |entry| entry.key.eq(&key)) { - return entry.clone(); + let entry = Arc::clone(entry); + drop(entries); + let guard = WorkCleanupGuard::new(self, entry); + // Establish cleanup before a duplicate key can unwind or reenter the group. + drop(key); + return guard; } let entry = Arc::new(Entry { @@ -85,7 +89,8 @@ where cell: OnceCell::new(), }); entries.insert_unique(hash, entry.clone(), |entry| entry.hash); - entry + drop(entries); + WorkCleanupGuard::new(self, entry) } fn remove(&self, key: &Q) @@ -157,9 +162,7 @@ where K: Eq + Hash, S: BuildHasher, { - fn new(group: &'a Group, key: K) -> Self { - let entry = group.get_or_insert(key); - + fn new(group: &'a Group, entry: Arc>) -> Self { Self { group, entry: Some(entry), @@ -285,7 +288,7 @@ where where F: AsyncFnOnce() -> V, { - let guard = WorkCleanupGuard::new(self, key); + let guard = self.get_or_insert(key); let entry = guard.entry(); let result = entry .cell @@ -355,7 +358,7 @@ where where F: AsyncFnOnce() -> Result, { - let guard = WorkCleanupGuard::new(self, key); + let guard = self.get_or_insert(key); let entry = guard.entry(); let result = entry .cell diff --git a/asyncband/src/singleflight/tests.rs b/asyncband/src/singleflight/tests.rs index d97ad5b6..f12540c6 100644 --- a/asyncband/src/singleflight/tests.rs +++ b/asyncband/src/singleflight/tests.rs @@ -16,14 +16,17 @@ // under the License. use std::cell::Cell; +use std::cell::RefCell; use std::future::pending; use std::future::poll_fn; use std::panic::AssertUnwindSafe; use std::panic::catch_unwind; use std::pin::pin; +use std::rc::Rc; use std::task::Poll; use super::*; +use crate::test_support::KeyWithDrop; use crate::test_support::poll_once; // These tests stay next to the implementation because they inspect private state. @@ -104,3 +107,48 @@ fn failed_try_work_preserves_entry_for_waiter_retry() { assert_eq!(poll_once(retry.as_mut()), Poll::Ready(Ok("success"))); assert_eq!(entry_count(&group), 0); } + +#[test] +fn duplicate_key_panic_cleans_up_pending_work() { + for fallible in [false, true] { + for cancel_during_drop in [false, true] { + let group = Rc::new(Group::::new()); + let first_group = Rc::clone(&group); + let mut first = Box::pin(async move { + first_group + .work( + KeyWithDrop { + key: "key", + on_drop: None, + }, + async || pending().await, + ) + .await + }); + assert!(poll_once(first.as_mut()).is_pending()); + let first = Rc::new(RefCell::new(Some(first))); + let dropping_first = Rc::clone(&first); + let dropping_group = Rc::clone(&group); + let key = KeyWithDrop { + key: "key", + on_drop: Some(Box::new(move || { + assert_eq!(entry_count(&dropping_group), 1); + if cancel_during_drop { + drop(dropping_first.borrow_mut().take()); + } + panic!("duplicate key destructor"); + })), + }; + let result = catch_unwind(AssertUnwindSafe(|| { + if fallible { + let _ = poll_once(pin!(group.try_work(key, async || Ok::<_, ()>(1)))); + } else { + let _ = poll_once(pin!(group.work(key, async || 1))); + } + })); + assert!(result.is_err()); + drop(first.borrow_mut().take()); + assert_eq!(entry_count(&group), 0); + } + } +} diff --git a/asyncband/src/test_support.rs b/asyncband/src/test_support.rs index 4bf84413..fd873299 100644 --- a/asyncband/src/test_support.rs +++ b/asyncband/src/test_support.rs @@ -24,3 +24,46 @@ use std::task::Waker; pub fn poll_once(future: Pin<&mut F>) -> Poll { future.poll(&mut Context::from_waker(Waker::noop())) } + +#[cfg(any(feature = "once-map", feature = "singleflight"))] +pub use self::key_with_drop::KeyWithDrop; + +#[cfg(any(feature = "once-map", feature = "singleflight"))] +mod key_with_drop { + use std::borrow::Borrow; + use std::hash::Hash; + use std::hash::Hasher; + + pub struct KeyWithDrop { + pub key: &'static str, + pub on_drop: Option>, + } + + impl PartialEq for KeyWithDrop { + fn eq(&self, other: &Self) -> bool { + self.key == other.key + } + } + + impl Eq for KeyWithDrop {} + + impl Hash for KeyWithDrop { + fn hash(&self, state: &mut H) { + self.key.hash(state); + } + } + + impl Borrow for KeyWithDrop { + fn borrow(&self) -> &str { + self.key + } + } + + impl Drop for KeyWithDrop { + fn drop(&mut self) { + if let Some(on_drop) = self.on_drop.take() { + on_drop(); + } + } + } +} From fcb5debb01a87827225d42144eb09206994a3f6e Mon Sep 17 00:00:00 2001 From: tison Date: Fri, 6 Nov 2026 23:11:57 +0800 Subject: [PATCH 2/2] test: verify the duplicate key destructor panic payload --- asyncband/src/once/once_map/tests.rs | 10 ++++++++-- asyncband/src/singleflight/tests.rs | 5 ++++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/asyncband/src/once/once_map/tests.rs b/asyncband/src/once/once_map/tests.rs index 236523c6..54532e9c 100644 --- a/asyncband/src/once/once_map/tests.rs +++ b/asyncband/src/once/once_map/tests.rs @@ -208,7 +208,10 @@ fn duplicate_key_panic_cleans_up_pending_compute() { let _ = poll_once(pin!(map.compute(key, async || 1))); } })); - assert!(result.is_err()); + assert_eq!( + result.unwrap_err().downcast_ref::<&str>(), + Some(&"duplicate key destructor") + ); drop(first.borrow_mut().take()); assert_eq!(map.len(), 0); } @@ -247,7 +250,10 @@ fn duplicate_key_panic_drops_cached_value_clone() { let _ = poll_once(pin!(map.compute(key, async || Arc::clone(&value)))); } })); - assert!(result.is_err()); + assert_eq!( + result.unwrap_err().downcast_ref::<&str>(), + Some(&"duplicate key destructor") + ); assert_eq!(map.len(), 0); assert_eq!(Arc::strong_count(&value), 1); } diff --git a/asyncband/src/singleflight/tests.rs b/asyncband/src/singleflight/tests.rs index f12540c6..aa91ddc2 100644 --- a/asyncband/src/singleflight/tests.rs +++ b/asyncband/src/singleflight/tests.rs @@ -146,7 +146,10 @@ fn duplicate_key_panic_cleans_up_pending_work() { let _ = poll_once(pin!(group.work(key, async || 1))); } })); - assert!(result.is_err()); + assert_eq!( + result.unwrap_err().downcast_ref::<&str>(), + Some(&"duplicate key destructor") + ); drop(first.borrow_mut().take()); assert_eq!(entry_count(&group), 0); }