add &Collector argument to reclaimers

Author: ibraheemdev

docs/GUIDE.md

@@ -209,14 +209,17 @@ impl<T> Stack<T> {
 If you need to run custom reclamation code, you can write a custom reclaimer.
 
 ```rust,ignore
-collector.retire(value, |value: *mut Node<T>| unsafe {
+collector.retire(value, |value: *mut Node<T>, _collector: &Collector| unsafe {
     // Safety: The value was allocated with `Box::new`.
     let value = Box::from_raw(ptr);
     println!("Dropping {value}");
     drop(value);
 });
 ```
 
+Note that the reclaimer receives a reference to the collector as its second
+argument, allowing for recursive reclamation.
+
 [`defer_retire`]:
   https://docs.rs/seize/latest/seize/trait.Guard.html#tymethod.defer_retire
 [`Guard::protect`]:

src/collector.rs

@@ -2,7 +2,6 @@ use crate::raw::{self, membarrier, Thread};
 use crate::{LocalGuard, OwnedGuard};
 
 use std::fmt;
-use std::sync::atomic::{AtomicUsize, Ordering};
 
 /// A concurrent garbage collector.
 ///
@@ -14,10 +13,8 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 /// Every instance of a concurrent data structure should typically own its
 /// `Collector`. This allows the garbage collection of non-`'static` values, as
 /// memory reclamation is guaranteed to run when the `Collector` is dropped.
+#[repr(transparent)]
 pub struct Collector {
-    /// A unique identifier for a collector.
-    id: usize,
-
     /// The underlying raw collector instance.
     pub(crate) raw: raw::Collector,
 }
@@ -34,9 +31,6 @@ impl Collector {
 
     /// Creates a new collector.
     pub fn new() -> Self {
-        // A counter for collector IDs.
-        static ID: AtomicUsize = AtomicUsize::new(0);
-
         // Initialize the `membarrier` module, detecting the presence of
         // operating-system strong barrier APIs.
         membarrier::detect();
@@ -50,7 +44,6 @@ impl Collector {
         let batch_size = cpus.max(Self::DEFAULT_BATCH_SIZE);
 
         Self {
-            id: ID.fetch_add(1, Ordering::Relaxed),
             raw: raw::Collector::new(cpus, batch_size),
         }
     }
@@ -187,14 +180,16 @@ impl Collector {
     /// Alternative, a custom reclaimer function can be used.
     ///
     /// ```
-    /// let collector = seize::Collector::new();
+    /// use seize::Collector;
+    ///
+    /// let collector = Collector::new();
     ///
     /// // Allocate a value and immediately retire it.
     /// let value: *mut usize = Box::into_raw(Box::new(1_usize));
     ///
     /// // Safety: The value was never shared.
     /// unsafe {
-    ///     collector.retire(value, |ptr: *mut usize| unsafe {
+    ///     collector.retire(value, |ptr: *mut usize, _collector: &Collector| unsafe {
     ///         // Safety: The value was allocated with `Box::new`.
     ///         let value = Box::from_raw(ptr);
     ///         println!("Dropping {value}");
@@ -203,7 +198,7 @@ impl Collector {
     /// }
     /// ```
     #[inline]
-    pub unsafe fn retire<T>(&self, ptr: *mut T, reclaim: unsafe fn(*mut T)) {
+    pub unsafe fn retire<T>(&self, ptr: *mut T, reclaim: unsafe fn(*mut T, &Collector)) {
         debug_assert!(!ptr.is_null(), "attempted to retire a null pointer");
 
         // Note that `add` doesn't ever actually reclaim the pointer immediately if
@@ -236,6 +231,11 @@ impl Collector {
     pub unsafe fn reclaim_all(&self) {
         unsafe { self.raw.reclaim_all() };
     }
+
+    // Create a reference to `Collector` from an underlying `raw::Collector`.
+    pub(crate) fn from_raw(raw: &raw::Collector) -> &Collector {
+        unsafe { &*(raw as *const raw::Collector as *const Collector) }
+    }
 }
 
 impl Eq for Collector {}
@@ -244,7 +244,7 @@ impl PartialEq for Collector {
     /// Checks if both references point to the same collector.
     #[inline]
     fn eq(&self, other: &Self) -> bool {
-        self.id == other.id
+        self.raw.id == other.raw.id
     }
 }
 

src/guard.rs

@@ -76,7 +76,7 @@ pub trait Guard {
     /// The retired pointer must no longer be accessible to any thread that
     /// enters after it is removed. Additionally, the pointer must be valid
     /// to pass to the provided reclaimer, once it is safe to reclaim.
-    unsafe fn defer_retire<T>(&self, ptr: *mut T, reclaim: unsafe fn(*mut T));
+    unsafe fn defer_retire<T>(&self, ptr: *mut T, reclaim: unsafe fn(*mut T, &Collector));
 }
 
 /// A guard that keeps the current thread marked as active.
@@ -169,7 +169,7 @@ impl Guard for LocalGuard<'_> {
     /// Retires a value, running `reclaim` when no threads hold a reference to
     /// it.
     #[inline]
-    unsafe fn defer_retire<T>(&self, ptr: *mut T, reclaim: unsafe fn(*mut T)) {
+    unsafe fn defer_retire<T>(&self, ptr: *mut T, reclaim: unsafe fn(*mut T, &Collector)) {
         // Safety:
         // - `self.thread` is the current thread.
         // - The validity of the pointer is guaranteed by the caller.
@@ -292,7 +292,7 @@ impl Guard for OwnedGuard<'_> {
     /// Retires a value, running `reclaim` when no threads hold a reference to
     /// it.
     #[inline]
-    unsafe fn defer_retire<T>(&self, ptr: *mut T, reclaim: unsafe fn(*mut T)) {
+    unsafe fn defer_retire<T>(&self, ptr: *mut T, reclaim: unsafe fn(*mut T, &Collector)) {
         // Safety: `self.reservation` is owned by the current thread.
         let reservation = unsafe { &*self.reservation };
         let _lock = reservation.lock.lock().unwrap();

src/raw/collector.rs

@@ -27,6 +27,9 @@ pub struct Collector {
     /// exiting.
     reservations: ThreadLocal<CachePadded<Reservation>>,
 
+    /// A unique identifier for a collector.
+    pub(crate) id: usize,
+
     /// The minimum number of nodes required in a batch before attempting
     /// retirement.
     pub(crate) batch_size: usize,
@@ -36,7 +39,11 @@ impl Collector {
     /// Create a collector with the provided batch size and initial thread
     /// count.
     pub fn new(threads: usize, batch_size: usize) -> Self {
+        // A counter for collector IDs.
+        static ID: AtomicUsize = AtomicUsize::new(0);
+
         Self {
+            id: ID.fetch_add(1, Ordering::Relaxed),
             reservations: ThreadLocal::with_capacity(threads),
             batches: ThreadLocal::with_capacity(threads),
             batch_size: batch_size.next_power_of_two(),
@@ -120,7 +127,7 @@ impl Collector {
             atomic::fence(Ordering::Acquire);
 
             // Decrement the reference counts of any batches that were retired.
-            unsafe { Collector::traverse(head) }
+            unsafe { self.traverse(head) }
         }
     }
 
@@ -138,7 +145,7 @@ impl Collector {
 
         if head != Entry::INACTIVE {
             // Decrement the reference counts of any batches that were retired.
-            unsafe { Collector::traverse(head) }
+            unsafe { self.traverse(head) }
         }
     }
 
@@ -157,7 +164,12 @@ impl Collector {
     /// Additionally, current thread must have unique access to the batch for
     /// the provided `thread`.
     #[inline]
-    pub unsafe fn add<T>(&self, ptr: *mut T, reclaim: unsafe fn(*mut T), thread: Thread) {
+    pub unsafe fn add<T>(
+        &self,
+        ptr: *mut T,
+        reclaim: unsafe fn(*mut T, &crate::Collector),
+        thread: Thread,
+    ) {
         // Safety: The caller guarantees we have unique access to the batch.
         let local_batch = unsafe { self.batches.load(thread).get() };
 
@@ -170,12 +182,13 @@ impl Collector {
             // Safety: `LocalBatch::DROP` means we have unique access to the collector.
             // Additionally, the caller guarantees that the pointer is valid for the
             // provided reclaimer.
-            unsafe { reclaim(ptr) }
+            unsafe { reclaim(ptr, crate::Collector::from_raw(self)) }
             return;
         }
 
         // Safety: `fn(*mut T) and fn(*mut U)` are ABI compatible if `T, U: Sized`.
-        let reclaim: unsafe fn(*mut ()) = unsafe { std::mem::transmute(reclaim) };
+        let reclaim: unsafe fn(*mut (), &crate::Collector) =
+            unsafe { std::mem::transmute(reclaim) };
 
         // Safety: The caller guarantees we have unique access to the batch.
         let len = unsafe {
@@ -364,7 +377,7 @@ impl Collector {
             // Additionally, the local batch has been reset and we are not holding on to any
             // mutable references, so any recursive calls to `retire` during
             // reclamation are valid.
-            unsafe { Collector::free_batch(batch) }
+            unsafe { self.free_batch(batch) }
         }
     }
 
@@ -376,7 +389,7 @@ impl Collector {
     /// `list` must be a valid reservation list.
     #[cold]
     #[inline(never)]
-    unsafe fn traverse(mut list: *mut Entry) {
+    unsafe fn traverse(&self, mut list: *mut Entry) {
         while !list.is_null() {
             let curr = list;
 
@@ -396,7 +409,7 @@ impl Collector {
 
                     // Safety: We have the last reference to the batch and it has been removed from
                     // our reservation list.
-                    Collector::free_batch(batch)
+                    self.free_batch(batch)
                 }
             }
         }
@@ -433,7 +446,7 @@ impl Collector {
             // ensured it is non-null. Additionally, the local batch was reset
             // above, so the batch is inaccessible through recursive calls to
             // `retire`.
-            unsafe { Collector::free_batch(batch) };
+            unsafe { self.free_batch(batch) };
 
             // Reset the batch.
             //
@@ -453,10 +466,10 @@ impl Collector {
     /// access the local batch; the batch being retired must be unreachable
     /// through any recursive calls.
     #[inline]
-    unsafe fn free_batch(batch: *mut Batch) {
+    unsafe fn free_batch(&self, batch: *mut Batch) {
         // Safety: We have a unique reference to the batch.
         for entry in unsafe { (*batch).entries.iter_mut() } {
-            unsafe { (entry.reclaim)(entry.ptr.cast()) };
+            unsafe { (entry.reclaim)(entry.ptr.cast(), crate::Collector::from_raw(self)) };
         }
 
         unsafe { LocalBatch::free(batch) };
@@ -530,7 +543,7 @@ struct Entry {
     ptr: *mut (),
 
     /// The function used to reclaim the object.
-    reclaim: unsafe fn(*mut ()),
+    reclaim: unsafe fn(*mut (), &crate::Collector),
 
     /// The state of the retired object.
     state: EntryState,

src/reclaim.rs

@@ -6,12 +6,14 @@
 
 use std::ptr;
 
+use crate::Collector;
+
 /// Reclaims memory allocated with [`Box`].
 ///
 /// # Safety
 ///
 /// The safety requirements of [`Box::from_raw`] apply.
-pub unsafe fn boxed<T>(ptr: *mut T) {
+pub unsafe fn boxed<T>(ptr: *mut T, _collector: &Collector) {
     unsafe { drop(Box::from_raw(ptr)) }
 }
 
@@ -20,6 +22,6 @@ pub unsafe fn boxed<T>(ptr: *mut T) {
 /// # Safety
 ///
 /// The safety requirements of [`ptr::drop_in_place`] apply.
-pub unsafe fn in_place<T>(ptr: *mut T) {
+pub unsafe fn in_place<T>(ptr: *mut T, _collector: &Collector) {
     unsafe { ptr::drop_in_place::<T>(ptr) }
 }

tests/lib.rs

@@ -3,7 +3,7 @@ use seize::{reclaim, Collector, Guard};
 use std::mem::ManuallyDrop;
 use std::ptr;
 use std::sync::atomic::{AtomicPtr, AtomicUsize, Ordering};
-use std::sync::{mpsc, Arc, Barrier, OnceLock};
+use std::sync::{mpsc, Arc, Barrier};
 use std::thread;
 
 #[test]
@@ -154,34 +154,33 @@ fn refresh() {
 
 #[test]
 fn recursive_retire() {
-    fn collector() -> &'static Collector {
-        static COLLECTOR: OnceLock<Collector> = OnceLock::new();
-        COLLECTOR.get_or_init(|| Collector::new().batch_size(1))
-    }
-
     struct Recursive {
         _value: usize,
         pointers: Vec<*mut usize>,
     }
 
+    let collector = Collector::new().batch_size(1);
+
     let ptr = boxed(Recursive {
         _value: 0,
         pointers: (0..cfg::ITEMS).map(boxed).collect(),
     });
 
     unsafe {
-        collector().retire(ptr, |link| {
-            let value = Box::from_raw(link.cast::<Recursive>());
+        collector.retire(ptr, |ptr: *mut Recursive, collector| {
+            let value = Box::from_raw(ptr);
+
             for pointer in value.pointers {
-                collector().retire(pointer, reclaim::boxed);
-                let mut guard = collector().enter();
+                collector.retire(pointer, reclaim::boxed);
+
+                let mut guard = collector.enter();
                 guard.flush();
                 guard.refresh();
                 drop(guard);
             }
         });
 
-        collector().enter().flush();
+        collector.enter().flush();
     }
 }
 
@@ -209,34 +208,29 @@ fn reclaim_all() {
 fn recursive_retire_reclaim_all() {
     struct Recursive {
         _value: usize,
-        collector: *mut Collector,
         pointers: Vec<*mut 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 collector = Collector::new().batch_size(cfg::ITEMS * 2);
         let dropped = Arc::new(AtomicUsize::new(0));
 
         let ptr = boxed(Recursive {
             _value: 0,
-            collector,
             pointers: (0..cfg::ITEMS)
                 .map(|_| boxed(DropTrack(dropped.clone())))
                 .collect(),
         });
 
-        (*collector).retire(ptr, |link| {
-            let value = Box::from_raw(link.cast::<Recursive>());
-            let collector = value.collector;
+        collector.retire(ptr, |ptr: *mut Recursive, collector| {
+            let value = Box::from_raw(ptr);
             for pointer in value.pointers {
                 (*collector).retire(pointer, reclaim::boxed);
             }
         });
 
-        (*collector).reclaim_all();
+        collector.reclaim_all();
         assert_eq!(dropped.load(Ordering::Relaxed), cfg::ITEMS);
-        let _ = Box::from_raw(collector);
     }
 }