From 13ddb72e5acc718c87a25553b029e3b0ffe9a683 Mon Sep 17 00:00:00 2001
From: Alex Gaynor <alex.gaynor@gmail.com>
Date: Sat, 15 May 2021 10:05:58 -0400
Subject: [PATCH] Reduce function calling overhead by ~25% in
 ReferencePool::update_counts:

1) Place both increfs and decrefs behind a single mutex, rather than two. Even uncontended mutexes aren't free to acquire.
2) Keep a dirty tracking bool to avoid acquiring any mutexes at all in the common case of no modifications (because everything happened with the GIL held)
---
 src/gil.rs | 46 ++++++++++++++++++++++++++++------------------
 1 file changed, 28 insertions(+), 18 deletions(-)

diff --git a/src/gil.rs b/src/gil.rs
index 15a7488c048..8245241e540 100644
--- a/src/gil.rs
+++ b/src/gil.rs
@@ -5,7 +5,7 @@
 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::ManuallyDrop, ptr::NonNull, sync::atomic};
 
 static START: Once = Once::new();
 
@@ -287,50 +287,60 @@ 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 {
+        macro_rules! swap_vec {
             // 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);
+                if !$cell.is_empty() {
+                    std::mem::swap(&mut out, &mut $cell);
                 }
-                drop(locked);
                 out
             }};
         }
 
+        let prev = self.dirty.swap(false, atomic::Ordering::Acquire);
+        if !prev {
+            return;
+        }
+
+        let mut ops = self.pointer_ops.lock();
         // 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 swap_vec!(ops.0) {
             unsafe { ffi::Py_INCREF(ptr.as_ptr()) };
         }
 
-        for ptr in swap_vec_with_lock!(self.pointers_to_decref) {
+        for ptr in swap_vec!(ops.1) {
             unsafe { ffi::Py_DECREF(ptr.as_ptr()) };
         }
     }
@@ -753,15 +763,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()));