memorypool: add optional support for the system allocator in Rc (#48305)

Author: jkarneges

benches/memorypool.rs

@@ -43,6 +43,28 @@ fn bench_memorypool_rc_new<const N: usize>(c: &mut Criterion, op_kcount: usize)
     });
 }
 
+fn bench_system_rc_new<const N: usize>(c: &mut Criterion, op_kcount: usize) {
+    let op_count = op_kcount * 1000;
+    let bytes = mem::size_of::<[u64; N]>();
+    let instances = RefCell::new(Vec::with_capacity(op_count));
+
+    c.bench_function(&format!("sys-rc-new-{bytes}b-x{op_kcount}k"), |b| {
+        b.iter_batched_ref(
+            || instances.borrow_mut().clear(),
+            |_| {
+                let instances = &mut *instances.borrow_mut();
+                let mut next_value: [u64; N] = [0; N];
+                while next_value[0] < op_count as u64 {
+                    let n = memorypool::Rc::new(next_value);
+                    instances.push(n);
+                    next_value[0] += 1;
+                }
+            },
+            BatchSize::PerIteration,
+        )
+    });
+}
+
 fn bench_std_rc_new<const N: usize>(c: &mut Criterion, op_kcount: usize) {
     let op_count = op_kcount * 1000;
     let bytes = mem::size_of::<[u64; N]>();
@@ -88,6 +110,28 @@ fn bench_memorypool_rc_drop<const N: usize>(c: &mut Criterion, op_kcount: usize)
     });
 }
 
+fn bench_system_rc_drop<const N: usize>(c: &mut Criterion, op_kcount: usize) {
+    let op_count = op_kcount * 1000;
+    let bytes = mem::size_of::<[u64; N]>();
+    let instances = RefCell::new(Vec::with_capacity(op_count));
+
+    c.bench_function(&format!("sys-rc-drp-{bytes}b-x{op_kcount}k"), |b| {
+        b.iter_batched_ref(
+            || {
+                let instances = &mut *instances.borrow_mut();
+                let mut next_value: [u64; N] = [0; N];
+                while next_value[0] < op_count as u64 {
+                    let n = memorypool::Rc::new(next_value);
+                    instances.push(n);
+                    next_value[0] += 1;
+                }
+            },
+            |_| instances.borrow_mut().clear(),
+            BatchSize::PerIteration,
+        )
+    });
+}
+
 fn bench_std_rc_drop<const N: usize>(c: &mut Criterion, op_kcount: usize) {
     let op_count = op_kcount * 1000;
     let bytes = mem::size_of::<[u64; N]>();
@@ -260,9 +304,11 @@ fn criterion_benchmark(c: &mut Criterion) {
     }
 
     bench_memorypool_rc_new::<1>(c, OP_KCOUNT);
+    bench_system_rc_new::<1>(c, OP_KCOUNT);
     bench_std_rc_new::<1>(c, OP_KCOUNT);
 
     bench_memorypool_rc_drop::<1>(c, OP_KCOUNT);
+    bench_system_rc_drop::<1>(c, OP_KCOUNT);
     bench_std_rc_drop::<1>(c, OP_KCOUNT);
 
     bench_memorypool_rc_new::<80>(c, OP_KCOUNT);

src/core/memorypool.rs

@@ -159,6 +159,20 @@ pub struct Rc<T> {
 }
 
 impl<T> Rc<T> {
+    pub fn new(v: T) -> Self {
+        let ptr = Box::leak(Box::new(RcInner {
+            refs: Cell::new(1),
+            memory: Cell::new(None),
+            key: Cell::new(0),
+            value: v,
+        }));
+
+        Self {
+            ptr: ptr.into(),
+            phantom: PhantomData,
+        }
+    }
+
     pub fn try_new_in(v: T, memory: &std::rc::Rc<RcMemory<T>>) -> Result<Self, AllocError> {
         let (key, ptr) = memory
             .insert(RcInner {
@@ -194,16 +208,24 @@ impl<T> Rc<T> {
 
     #[inline(never)]
     fn drop_slow(&mut self) {
-        let key = self.inner().key.get();
-
-        // Entries contain a std::rc::Rc to the Memory they are contained in,
-        // and we need to be careful the Memory is not dropped while an entry
-        // is being removed. To ensure this, the Rc is moved out of the entry
-        // and dropped only after the entry is removed.
-
-        let memory = self.inner().memory.take().unwrap();
-
-        memory.remove(key);
+        if let Some(memory) = self.inner().memory.take() {
+            // Slab allocations contain a std::rc::Rc to the Memory they are
+            // contained in, and we need to be careful the Memory is not
+            // dropped while an entry is being removed. To ensure this, the
+            // Rc is moved out of the entry above, and it is dropped only
+            // after the entry is removed.
+
+            let key = self.inner().key.get();
+
+            memory.remove(key);
+        } else {
+            // If there is no reference to slab memory, then the memory is
+            // managed by the system allocator. To free it, we simply convert
+            // the pointer back to a Box and drop it.
+
+            // SAFETY: While this Rc is alive, ptr is always valid
+            unsafe { drop(Box::from_raw(self.ptr.as_mut())) };
+        }
     }
 }
 
@@ -333,8 +355,40 @@ impl ReusableVec {
 mod tests {
     use super::*;
 
+    struct FlagOnDrop {
+        dropped: std::rc::Rc<Cell<bool>>,
+    }
+
+    impl Drop for FlagOnDrop {
+        fn drop(&mut self) {
+            self.dropped.set(true);
+        }
+    }
+
     #[test]
     fn test_rc() {
+        let dropped = std::rc::Rc::new(Cell::new(false));
+
+        let f = FlagOnDrop {
+            dropped: std::rc::Rc::clone(&dropped),
+        };
+
+        let a = Rc::new(f);
+        assert_eq!(a.inner().refs.get(), 1);
+
+        let b = Rc::clone(&a);
+        assert_eq!(a.inner().refs.get(), 2);
+
+        drop(a);
+        assert_eq!(b.inner().refs.get(), 1);
+        assert!(!dropped.get());
+
+        drop(b);
+        assert!(dropped.get());
+    }
+
+    #[test]
+    fn test_rc_in() {
         let memory = std::rc::Rc::new(RcMemory::new(2));
         assert_eq!(memory.len(), 0);