Reduce function calling overhead by ~25% in ReferencePool::update_counts:

src/gil.rs

@@ -5,7 +5,11 @@
 use crate::{ffi, internal_tricks::Unsendable, Python};
 use parking_lot::{const_mutex, Mutex, Once};
 use std::cell::{Cell, RefCell};
-use std::{mem::ManuallyDrop, ptr::NonNull};
+use std::{
+    mem::{self, ManuallyDrop},
+    ptr::NonNull,
+    sync::atomic,
+};
 
 static START: Once = Once::new();
 
@@ -287,50 +291,51 @@ impl Drop for GILGuard {
     }
 }
 
+// Vector of PyObject(
+type PyObjVec = Vec<NonNull<ffi::PyObject>>;
+
 /// Thread-safe storage for objects which were inc_ref / dec_ref while the GIL was not held.
 struct ReferencePool {
-    pointers_to_incref: Mutex<Vec<NonNull<ffi::PyObject>>>,
-    pointers_to_decref: Mutex<Vec<NonNull<ffi::PyObject>>>,
+    dirty: atomic::AtomicBool,
+    // .0 is INCREFs, .1 is DECREFs
+    pointer_ops: Mutex<(PyObjVec, PyObjVec)>,
 }
 
 impl ReferencePool {
     const fn new() -> Self {
         Self {
-            pointers_to_incref: const_mutex(Vec::new()),
-            pointers_to_decref: const_mutex(Vec::new()),
+            dirty: atomic::AtomicBool::new(false),
+            pointer_ops: const_mutex((Vec::new(), Vec::new())),
         }
     }
 
     fn register_incref(&self, obj: NonNull<ffi::PyObject>) {
-        self.pointers_to_incref.lock().push(obj)
+        self.pointer_ops.lock().0.push(obj);
+        self.dirty.store(true, atomic::Ordering::Release);
     }
 
     fn register_decref(&self, obj: NonNull<ffi::PyObject>) {
-        self.pointers_to_decref.lock().push(obj)
+        self.pointer_ops.lock().1.push(obj);
+        self.dirty.store(true, atomic::Ordering::Release);
     }
 
     fn update_counts(&self, _py: Python) {
-        macro_rules! swap_vec_with_lock {
-            // Get vec from one of ReferencePool's mutexes via lock, swap vec if needed, unlock.
-            ($cell:expr) => {{
-                let mut locked = $cell.lock();
-                let mut out = Vec::new();
-                if !locked.is_empty() {
-                    std::mem::swap(&mut out, &mut *locked);
-                }
-                drop(locked);
-                out
-            }};
+        let prev = self.dirty.swap(false, atomic::Ordering::Acquire);
+        if !prev {
+            return;
         }
 
+        let mut ops = self.pointer_ops.lock();
+        let (increfs, decrefs) = mem::take(&mut *ops);
+        drop(ops);
         // Always increase reference counts first - as otherwise objects which have a
         // nonzero total reference count might be incorrectly dropped by Python during
         // this update.
-        for ptr in swap_vec_with_lock!(self.pointers_to_incref) {
+        for ptr in increfs {
             unsafe { ffi::Py_INCREF(ptr.as_ptr()) };
         }
 
-        for ptr in swap_vec_with_lock!(self.pointers_to_decref) {
+        for ptr in decrefs {
             unsafe { ffi::Py_DECREF(ptr.as_ptr()) };
         }
     }
@@ -753,15 +758,15 @@ mod test {
 
         // The pointer should appear once in the incref pool, and once in the
         // decref pool (for the clone being created and also dropped)
-        assert_eq!(&*POOL.pointers_to_incref.lock(), &vec![ptr]);
-        assert_eq!(&*POOL.pointers_to_decref.lock(), &vec![ptr]);
+        assert_eq!(&POOL.pointer_ops.lock().0, &vec![ptr]);
+        assert_eq!(&POOL.pointer_ops.lock().1, &vec![ptr]);
 
         // Re-acquring GIL will clear these pending changes
         drop(gil);
         let gil = Python::acquire_gil();
 
-        assert!(POOL.pointers_to_incref.lock().is_empty());
-        assert!(POOL.pointers_to_decref.lock().is_empty());
+        assert!(POOL.pointer_ops.lock().0.is_empty());
+        assert!(POOL.pointer_ops.lock().1.is_empty());
 
         // Overall count is still unchanged
         assert_eq!(count, obj.get_refcnt(gil.python()));