Fix unsafe reclamation when dropping maps with a shared collector

src/raw/mod.rs

@@ -18,7 +18,7 @@ use self::utils::{untagged, AtomicPtrFetchOps, Counter, Parker, StrictProvenance
 use crate::map::{Compute, Operation, ResizeMode};
 use crate::Equivalent;
 
-use seize::{Collector, LocalGuard, OwnedGuard};
+use seize::{Collector, Guard as _, LocalGuard, OwnedGuard};
 use utils::{MapGuard, Stack, VerifiedGuard};
 
 /// A lock-free hash-table.
@@ -2724,36 +2724,55 @@ impl<K, V, G> Clone for Iter<'_, K, V, G> {
 
 impl<K, V, S> Drop for HashMap<K, V, S> {
     fn drop(&mut self) {
-        let mut raw = *self.table.get_mut();
+        let raw = *self.table.get_mut();
 
-        // Make sure all objects are reclaimed before the collector is dropped.
-        //
-        // Dropping a table depends on accessing the collector for deferred retirement,
-        // using the shared collector pointer that is invalidated by drop.
-        //
-        // Safety: We have a unique reference to the collector.
-        unsafe { self.collector.reclaim_all() };
+        if Arc::get_mut(&mut self.collector).is_some() {
+            // Make sure all objects are reclaimed before the collector is dropped.
+            //
+            // Safety: `Arc::get_mut` proves that no other strong or weak references
+            // to the collector exist, so no guards can still be active.
+            unsafe { self.collector.reclaim_all() };
 
-        // Drop all nested tables and entries.
-        while !raw.is_null() {
-            // Safety: The root and next tables are always valid pointers to a
-            // table allocation, or null.
-            let mut table = unsafe { Table::from_raw(raw) };
+            // Safety: We have exclusive access to the map and its collector.
+            unsafe { drop_tables(raw, &self.collector) };
+        } else {
+            let guard = self.collector.enter();
 
-            // Read the next table pointer before dropping the current one.
-            let next = *table.state_mut().next.get_mut();
+            // Defer destruction of this map's table chain through the shared
+            // collector so active guards from any map sharing it remain protected.
+            //
+            // Safety: Dropping the map makes its root table unreachable. `raw` is
+            // the valid root of the complete table chain owned by this map.
+            unsafe {
+                guard.defer_retire(raw, |raw, collector| {
+                    drop_tables(raw, collector);
+                });
+            }
+        }
+    }
+}
 
-            // Safety: We have unique access to the table and do
-            // not access the entries after this call.
-            unsafe { drop_entries(table) };
+// Drop all tables and entries owned by a map.
+//
+// # Safety
+//
+// `raw` must be the root of a valid table chain that is no longer accessible.
+unsafe fn drop_tables<K, V>(mut raw: *mut RawTable<Entry<K, V>>, collector: &Collector) {
+    while !raw.is_null() {
+        // Safety: The root and next tables are always valid pointers to a
+        // table allocation, or null.
+        let mut table = unsafe { Table::from_raw(raw) };
 
-            // Safety: We have unique access to the table and do
-            // not access it after this call.
-            unsafe { drop_table(table, &self.collector) };
+        // Read the next table pointer before dropping the current one.
+        let next = *table.state_mut().next.get_mut();
 
-            // Continue for all nested tables.
-            raw = next;
-        }
+        // Safety: The table is inaccessible and its entries are not used afterward.
+        unsafe { drop_entries(table) };
+
+        // Safety: The table is inaccessible and is not used afterward.
+        unsafe { drop_table(table, collector) };
+
+        raw = next;
     }
 }
 

tests/basic.rs

@@ -1,9 +1,13 @@
 // Adapted from: https://github.com/jonhoo/flurry/blob/main/tests/basic.rs
 
-use papaya::{Compute, HashMap, OccupiedError, Operation};
+use papaya::{Compute, HashMap, OccupiedError, Operation, ResizeMode};
+use seize::Collector;
 
 use std::hash::{BuildHasher, BuildHasherDefault, Hasher};
-use std::sync::Arc;
+use std::sync::{
+    atomic::{AtomicUsize, Ordering},
+    Arc,
+};
 
 mod common;
 use common::with_map;
@@ -580,6 +584,51 @@ fn current_kv_dropped() {
     });
 }
 
+#[test]
+fn shared_collector_drop_respects_active_guards() {
+    struct DropProbe(Arc<AtomicUsize>);
+
+    impl Drop for DropProbe {
+        fn drop(&mut self) {
+            self.0.fetch_add(1, Ordering::SeqCst);
+        }
+    }
+
+    let collector = Arc::new(Collector::new());
+    let drops = Arc::new(AtomicUsize::new(0));
+    let protected = HashMap::builder()
+        .shared_collector(Arc::clone(&collector))
+        .build();
+    let growing = HashMap::builder()
+        .capacity(1)
+        .resize_mode(ResizeMode::Blocking)
+        .shared_collector(Arc::clone(&collector))
+        .build();
+    let trigger = HashMap::<usize, usize>::builder()
+        .shared_collector(Arc::clone(&collector))
+        .build();
+
+    let guard = protected.guard();
+    protected.insert(0, DropProbe(Arc::clone(&drops)), &guard);
+    let reference = protected.get(&0, &guard).unwrap();
+    std::hint::black_box(reference);
+    protected.remove(&0, &guard);
+
+    for key in 0..1_024 {
+        growing.pin().insert(key, key);
+    }
+    assert_eq!(drops.load(Ordering::SeqCst), 0);
+
+    drop(trigger);
+    assert_eq!(drops.load(Ordering::SeqCst), 0);
+
+    drop(guard);
+    drop(growing);
+    drop(protected);
+    drop(collector);
+    assert_eq!(drops.load(Ordering::SeqCst), 1);
+}
+
 #[test]
 fn empty_maps_equal() {
     with_map::<usize, usize>(|map1| {