avoid leaking recursive calls to retire when dropping collector

Author: ibraheemdev

src/collector.rs

@@ -262,6 +262,28 @@ impl Collector {
         unsafe { self.raw.add(ptr, reclaim, Thread::current()) }
     }
 
+    /// Reclaim any values that have been retired.
+    ///
+    /// This method reclaims any objects that have been retired across *all* threads.
+    /// After calling this method, any values that were previous retired, or retired
+    /// recursively on the current thread during this call, will have been reclaimed.
+    ///
+    /// # Safety
+    ///
+    /// This function is **extremely unsafe** to call. It is only sound when no threads are
+    /// currently active, whether accessing values that have been retired or accessing the
+    /// collector through any type of guard. This is akin to having a unique reference to the
+    /// collector. However, this method takes a shared reference, as reclaimers to be run by this
+    /// thread are allowed to access the collector recursively.
+    ///
+    /// # Notes
+    ///
+    /// Note that if reclaimers initialize guards across threads, or initialize owned guards,
+    /// objects retired through those guards may not be reclaimed.
+    pub unsafe fn reclaim_all(&self) {
+        unsafe { self.raw.reclaim_all() };
+    }
+
     pub(crate) fn ptr_eq(this: &Collector, other: &Collector) -> bool {
         ptr::eq(this.unique, other.unique)
     }

src/deferred.rs

@@ -17,7 +17,7 @@ use crate::{AsLink, Collector, Link};
 /// concurrently. It is not meant to be used to amortize the cost of retirement, which is done
 /// through thread-local batches controlled with [`Collector::batch_size`], as access from a single-thread
 /// can be more expensive than is required. Deferred batches are useful when you need to control when
-/// a batch of nodes is retired directly, a relatively rare use case.
+/// a batch of objects is retired directly, a relatively rare use case.
 ///
 /// # Examples
 ///
@@ -31,13 +31,13 @@ use crate::{AsLink, Collector, Link};
 ///     .map(|i| AtomicPtr::new(collector.link_boxed(i)))
 ///     .collect::<Arc<[_]>>();
 ///
-/// // create a batch of nodes to retire
+/// // create a batch of objects to retire
 /// let mut batch = Deferred::new();
 ///
 /// for item in items.iter() {
 ///     // make the item unreachable with an atomic swap
 ///     let old = item.swap(std::ptr::null_mut(), Ordering::AcqRel);
-///     // don't retire just yet, add the node to the batch
+///     // don't retire just yet, add the object to the batch
 ///     unsafe { batch.defer(old) };
 /// }
 ///
@@ -51,7 +51,7 @@ pub struct Deferred {
 }
 
 impl Deferred {
-    /// Create a new batch of deferred nodes.
+    /// Create a new batch of deferred objects.
     pub const fn new() -> Deferred {
         Deferred {
             head: AtomicPtr::new(ptr::null_mut()),
@@ -98,14 +98,14 @@ impl Deferred {
     }
 
     /// Retires a batch of values, running `reclaim` when no threads hold a reference to any
-    /// nodes in the batch.
+    /// objects in the batch.
     ///
     /// Note that this method is disconnected from any guards on the current thread,
     /// so the pointers may be reclaimed immediately.
     ///
     /// # Safety
     ///
-    /// The safety requirements of [`Collector::retire`] apply to each node in the batch.
+    /// The safety requirements of [`Collector::retire`] apply to each object in the batch.
     ///
     /// [`Collector::retire`]: crate::Collector::retire
     pub unsafe fn retire_all(&mut self, collector: &Collector, reclaim: unsafe fn(*mut Link)) {
@@ -114,7 +114,7 @@ impl Deferred {
         unsafe { collector.raw.add_batch(self, reclaim, Thread::current()) }
     }
 
-    /// Run a function for each node in the batch.
+    /// Run a function for each object in the batch.
     ///
     /// This function does not consume the batch and can be called multiple
     /// times, **before retirement**.

src/raw.rs

@@ -255,6 +255,12 @@ impl Collector {
         // safety: local batches are only accessed by the current thread
         let batch = unsafe { (*local_batch).get_or_init(self.batch_size) };
 
+        // if we are in a recursive call during drop, reclaim immediately
+        if batch == LocalBatch::DROP {
+            unsafe { reclaim(ptr.cast::<Link>()) }
+            return;
+        }
+
         // `ptr` is guaranteed to be a valid pointer that can be cast to a node (`T: AsLink`)
         //
         // any other thread with a reference to the pointer only has a shared
@@ -301,6 +307,12 @@ impl Collector {
         // safety: local batches are only accessed by the current thread
         let batch = unsafe { (*local_batch).get_or_init(self.batch_size) };
 
+        // if we are in a recursive call during drop, reclaim immediately
+        if batch == LocalBatch::DROP {
+            deferred.for_each(|ptr| unsafe { reclaim(ptr.cast::<Link>()) });
+            return;
+        }
+
         // keep track of the oldest node in the batch
         let min_epoch = *deferred.min_epoch.get_mut();
 
@@ -368,7 +380,7 @@ impl Collector {
         let batch = unsafe { (*local_batch).batch };
 
         // there is nothing to retire
-        if batch.is_null() {
+        if batch.is_null() || batch == LocalBatch::DROP {
             return;
         }
 
@@ -536,6 +548,35 @@ impl Collector {
         }
     }
 
+    // Reclaim all values in the collector, including recursive calls to retire.
+    //
+    // # Safety
+    //
+    // No threads may be accessing the collector or any values that have been retired.
+    pub unsafe fn reclaim_all(&self) {
+        for local_batch in self.batches.iter() {
+            let local_batch = local_batch.value.get();
+
+            unsafe {
+                let batch = (*local_batch).batch;
+
+                // there is nothing to reclaim
+                if batch.is_null() {
+                    continue;
+                }
+
+                // tell any recursive calls to `retire` to reclaim immediately
+                (*local_batch).batch = LocalBatch::DROP;
+
+                // safety: we have &mut self and the batch is non-null
+                Collector::free_batch(batch);
+
+                // reset the batch
+                (*local_batch).batch = ptr::null_mut();
+            }
+        }
+    }
+
     // Free a reservation list.
     //
     // # Safety
@@ -556,20 +597,8 @@ impl Collector {
 
 impl Drop for Collector {
     fn drop(&mut self) {
-        for batch in self.batches.iter() {
-            unsafe {
-                // safety: we have &mut self
-                let batch = (*batch.value.get()).batch;
-
-                // there is nothing to reclaim
-                if batch.is_null() {
-                    continue;
-                }
-
-                // safety: we have &mut self and the batch is non-null
-                Collector::free_batch(batch);
-            }
-        }
+        // safety: we have &mut self
+        unsafe { self.reclaim_all() };
     }
 }
 
@@ -651,7 +680,7 @@ impl Entry {
     //
     // While null indicates an empty list, INACTIVE indicates the thread has no active
     // guards and is not accessing any objects.
-    pub const INACTIVE: *mut Entry = -1_isize as usize as _;
+    pub const INACTIVE: *mut Entry = usize::MAX as _;
 }
 
 // A pointer to a batch, unique to the current thread.
@@ -668,6 +697,10 @@ impl Default for LocalBatch {
 }
 
 impl LocalBatch {
+    // This is set during a call to reclaim_all, signalling recursive calls to retire
+    // to reclaim immediately.
+    const DROP: *mut Batch = usize::MAX as _;
+
     // Return a pointer to the batch, initializing it if the batch was null.
     fn get_or_init(&mut self, capacity: usize) -> *mut Batch {
         if self.batch.is_null() {

tests/lib.rs

@@ -177,6 +177,66 @@ fn recursive_retire() {
     }
 }
 
+#[test]
+fn reclaim_all() {
+    let collector = Collector::new().batch_size(2);
+
+    for _ in 0..cfg::ITER {
+        let dropped = Arc::new(AtomicUsize::new(0));
+
+        let items = (0..cfg::ITEMS)
+            .map(|_| AtomicPtr::new(collector.link_boxed(DropTrack(dropped.clone()))))
+            .collect::<Vec<_>>();
+
+        for item in items {
+            unsafe {
+                collector.retire(
+                    item.load(Ordering::Relaxed),
+                    reclaim::boxed::<Linked<DropTrack>>,
+                )
+            };
+        }
+
+        unsafe { collector.reclaim_all() };
+        assert_eq!(dropped.load(Ordering::Relaxed), cfg::ITEMS);
+    }
+}
+
+#[test]
+fn recursive_retire_reclaim_all() {
+    struct Recursive {
+        _value: usize,
+        collector: *mut Collector,
+        pointers: Vec<*mut Linked<DropTrack>>,
+    }
+
+    unsafe {
+        // make sure retire runs in drop, not immediately
+        let collector = Box::into_raw(Box::new(Collector::new().batch_size(cfg::ITEMS * 2)));
+        let dropped = Arc::new(AtomicUsize::new(0));
+
+        let ptr = (*collector).link_boxed(Recursive {
+            _value: 0,
+            collector,
+            pointers: (0..cfg::ITEMS)
+                .map(|_| (*collector).link_boxed(DropTrack(dropped.clone())))
+                .collect(),
+        });
+
+        (*collector).retire(ptr, |link| {
+            let value = Box::from_raw(link.cast::<Linked<Recursive>>());
+            let collector = value.value.collector;
+            for pointer in value.value.pointers {
+                (*collector).retire(pointer, reclaim::boxed::<Linked<DropTrack>>);
+            }
+        });
+
+        (*collector).reclaim_all();
+        assert_eq!(dropped.load(Ordering::Relaxed), cfg::ITEMS);
+        let _ = Box::from_raw(collector);
+    }
+}
+
 #[test]
 fn deferred() {
     let collector = Arc::new(Collector::new().batch_size(2));