fix: reclaim entries when duplicate key destructors panic

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
 

asyncband/src/once/once_map/mod.rs

@@ -38,9 +38,13 @@ struct Entry<K, V> {
     cell: OnceCell<V>,
 }
 
-enum Lookup<K, V> {
+enum Lookup<'a, K, V, S>
+where
+    K: Eq + Hash,
+    S: BuildHasher,
+{
     Ready(V),
-    Pending(Arc<Entry<K, V>>),
+    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<K, V>
+    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<Entry<K, V>>) -> Lookup<K, V>
+    fn classify(&self, entry: Arc<Entry<K, V>>) -> 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<V, E>,
     {
-        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)

asyncband/src/once/once_map/tests.rs

@@ -16,16 +16,19 @@
 // under the License.
 
 use std::cell::Cell;
+use std::cell::RefCell;
 use std::future::pending;
 use std::future::poll_fn;
 use std::hash::BuildHasherDefault;
 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,101 @@ 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::<KeyWithDrop, i32>::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_eq!(
+                result.unwrap_err().downcast_ref::<&str>(),
+                Some(&"duplicate key destructor")
+            );
+            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_eq!(
+            result.unwrap_err().downcast_ref::<&str>(),
+            Some(&"duplicate key destructor")
+        );
+        assert_eq!(map.len(), 0);
+        assert_eq!(Arc::strong_count(&value), 1);
+    }
+}

asyncband/src/singleflight/mod.rs

@@ -71,12 +71,16 @@ where
     K: Eq + Hash,
     S: BuildHasher,
 {
-    fn get_or_insert(&self, key: K) -> Arc<Entry<K, V>> {
+    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<Q>(&self, key: &Q)
@@ -157,9 +162,7 @@ where
     K: Eq + Hash,
     S: BuildHasher,
 {
-    fn new(group: &'a Group<K, V, S>, key: K) -> Self {
-        let entry = group.get_or_insert(key);
-
+    fn new(group: &'a Group<K, V, S>, entry: Arc<Entry<K, V>>) -> 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<V, E>,
     {
-        let guard = WorkCleanupGuard::new(self, key);
+        let guard = self.get_or_insert(key);
         let entry = guard.entry();
         let result = entry
             .cell

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,51 @@ 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::<KeyWithDrop, i32>::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_eq!(
+                result.unwrap_err().downcast_ref::<&str>(),
+                Some(&"duplicate key destructor")
+            );
+            drop(first.borrow_mut().take());
+            assert_eq!(entry_count(&group), 0);
+        }
+    }
+}