Use a blocking approach to reduce allocation overhead of the lock-free reference pool

Author: adamreichold

src/gil.rs

@@ -282,7 +282,8 @@ impl Drop for GILGuard {
 #[cfg(not(pyo3_disable_reference_pool))]
 /// Thread-safe storage for objects which were dec_ref while the GIL was not held.
 struct PendingDecRef {
-    obj: NonNull<ffi::PyObject>,
+    objs: [NonNull<ffi::PyObject>; 32],
+    cnt: usize,
     next: *mut PendingDecRef,
 }
 
@@ -291,17 +292,57 @@ static POOL: AtomicPtr<PendingDecRef> = AtomicPtr::new(null_mut());
 
 #[cfg(not(pyo3_disable_reference_pool))]
 fn enqueue_decref(obj: NonNull<ffi::PyObject>) {
+    let old_top = POOL.swap(null_mut(), Ordering::AcqRel);
+
+    let mut curr = old_top;
+
+    while !curr.is_null() {
+        // SAFETY: We have exclusive ownership of the full stack starting at `old_top`.
+        let val = unsafe { &mut *curr };
+
+        if val.cnt < val.objs.len() {
+            val.objs[val.cnt] = obj;
+            val.cnt += 1;
+
+            let mut last_next = &mut val.next;
+
+            while !last_next.is_null() {
+                last_next = unsafe { &mut (**last_next).next };
+            }
+
+            POOL.fetch_update(Ordering::AcqRel, Ordering::Relaxed, |new_top| {
+                *last_next = new_top;
+                Some(old_top)
+            })
+            .unwrap();
+            return;
+        }
+
+        curr = val.next;
+    }
+
+    // There is no existing stack or it has no unused capacity remaining,
+    // hence we allocate a new block and prepend it locally before publishing.
+    let mut objs = [NonNull::dangling(); 32];
+    objs[0] = obj;
+
     let val = PendingDecRef {
-        obj,
-        next: null_mut(),
+        objs,
+        cnt: 1,
+        next: old_top,
     };
 
     let top = Box::into_raw(Box::new(val));
 
-    let next = unsafe { &mut (*top).next };
+    // SAFETY: We just allocated `top` using `Box::new`.
+    let mut last_next = unsafe { &mut (*top).next };
+
+    while !last_next.is_null() {
+        last_next = unsafe { &mut (**last_next).next };
+    }
 
     POOL.fetch_update(Ordering::AcqRel, Ordering::Relaxed, |new_top| {
-        *next = new_top;
+        *last_next = new_top;
         Some(top)
     })
     .unwrap();
@@ -314,7 +355,9 @@ fn update_counts(_py: Python<'_>) {
         // SAFETY: Was enqueued using `Box::into_raw`.
         let val = unsafe { Box::from_raw(top) };
 
-        unsafe { ffi::Py_DECREF(val.obj.as_ptr()) };
+        for obj in &val.objs[..val.cnt] {
+            unsafe { ffi::Py_DECREF(obj.as_ptr()) };
+        }
 
         top = val.next;
     }
@@ -581,11 +624,13 @@ mod tests {
             // SAFETY: Was enqueued using `Box::into_raw`.
             let val = unsafe { Box::from_raw(top) };
 
-            if val.obj.as_ptr() == obj.as_ptr() {
-                found = true;
-            }
+            for obj1 in &val.objs[..val.cnt] {
+                if obj1.as_ptr() == obj.as_ptr() {
+                    found = true;
+                }
 
-            unsafe { ffi::Py_DECREF(val.obj.as_ptr()) };
+                unsafe { ffi::Py_DECREF(obj1.as_ptr()) };
+            }
 
             top = val.next;
         }