Merge pull request #3250 from PyO3/pool-opts

Two minor pool optimizations

src/gil.rs

@@ -3,8 +3,12 @@
 use crate::impl_::not_send::{NotSend, NOT_SEND};
 use crate::{ffi, Python};
 use parking_lot::{const_mutex, Mutex, Once};
-use std::cell::{Cell, RefCell};
-use std::{mem, ptr::NonNull, sync::atomic};
+use std::cell::Cell;
+#[cfg(debug_assertions)]
+use std::cell::RefCell;
+#[cfg(not(debug_assertions))]
+use std::cell::UnsafeCell;
+use std::{mem, ptr::NonNull};
 
 static START: Once = Once::new();
 
@@ -33,7 +37,10 @@ thread_local_const_init! {
     static GIL_COUNT: Cell<isize> = const { Cell::new(0) };
 
     /// Temporarily hold objects that will be released when the GILPool drops.
-    static OWNED_OBJECTS: RefCell<Vec<NonNull<ffi::PyObject>>> = const { RefCell::new(Vec::new()) };
+    #[cfg(debug_assertions)]
+    static OWNED_OBJECTS: RefCell<PyObjVec> = const { RefCell::new(Vec::new()) };
+    #[cfg(not(debug_assertions))]
+    static OWNED_OBJECTS: UnsafeCell<PyObjVec> = const { UnsafeCell::new(Vec::new()) };
 }
 
 const GIL_LOCKED_DURING_TRAVERSE: isize = -1;
@@ -238,38 +245,34 @@ 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 {
-    dirty: atomic::AtomicBool,
     // .0 is INCREFs, .1 is DECREFs
     pointer_ops: Mutex<(PyObjVec, PyObjVec)>,
 }
 
 impl ReferencePool {
     const fn new() -> Self {
         Self {
-            dirty: atomic::AtomicBool::new(false),
             pointer_ops: const_mutex((Vec::new(), Vec::new())),
         }
     }
 
     fn register_incref(&self, obj: NonNull<ffi::PyObject>) {
         self.pointer_ops.lock().0.push(obj);
-        self.dirty.store(true, atomic::Ordering::Release);
     }
 
     fn register_decref(&self, obj: NonNull<ffi::PyObject>) {
         self.pointer_ops.lock().1.push(obj);
-        self.dirty.store(true, atomic::Ordering::Release);
     }
 
     fn update_counts(&self, _py: Python<'_>) {
-        let prev = self.dirty.swap(false, atomic::Ordering::Acquire);
-        if !prev {
+        let mut ops = self.pointer_ops.lock();
+        if ops.0.is_empty() && ops.1.is_empty() {
             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.
@@ -377,7 +380,16 @@ impl GILPool {
         // Update counts of PyObjects / Py that have been cloned or dropped since last acquisition
         POOL.update_counts(Python::assume_gil_acquired());
         GILPool {
-            start: OWNED_OBJECTS.try_with(|o| o.borrow().len()).ok(),
+            start: OWNED_OBJECTS
+                .try_with(|owned_objects| {
+                    #[cfg(debug_assertions)]
+                    let len = owned_objects.borrow().len();
+                    #[cfg(not(debug_assertions))]
+                    // SAFETY: This is not re-entrant.
+                    let len = unsafe { (*owned_objects.get()).len() };
+                    len
+                })
+                .ok(),
             _not_send: NOT_SEND,
         }
     }
@@ -391,18 +403,21 @@ impl GILPool {
 
 impl Drop for GILPool {
     fn drop(&mut self) {
-        if let Some(obj_len_start) = self.start {
-            let dropping_obj = OWNED_OBJECTS.with(|holder| {
-                // `holder` must be dropped before calling Py_DECREF, or Py_DECREF may call
-                // `GILPool::drop` recursively, resulting in invalid borrowing.
-                let mut holder = holder.borrow_mut();
-                if obj_len_start < holder.len() {
-                    holder.split_off(obj_len_start)
+        if let Some(start) = self.start {
+            let owned_objects = OWNED_OBJECTS.with(|owned_objects| {
+                #[cfg(debug_assertions)]
+                let mut owned_objects = owned_objects.borrow_mut();
+                #[cfg(not(debug_assertions))]
+                // SAFETY: `OWNED_OBJECTS` is released before calling Py_DECREF,
+                // or Py_DECREF may call `GILPool::drop` recursively, resulting in invalid borrowing.
+                let owned_objects = unsafe { &mut *owned_objects.get() };
+                if start < owned_objects.len() {
+                    owned_objects.split_off(start)
                 } else {
                     Vec::new()
                 }
             });
-            for obj in dropping_obj {
+            for obj in owned_objects {
                 unsafe {
                     ffi::Py_DECREF(obj.as_ptr());
                 }
@@ -451,7 +466,15 @@ pub unsafe fn register_decref(obj: NonNull<ffi::PyObject>) {
 pub unsafe fn register_owned(_py: Python<'_>, obj: NonNull<ffi::PyObject>) {
     debug_assert!(gil_is_acquired());
     // Ignores the error in case this function called from `atexit`.
-    let _ = OWNED_OBJECTS.try_with(|holder| holder.borrow_mut().push(obj));
+    let _ = OWNED_OBJECTS.try_with(|owned_objects| {
+        #[cfg(debug_assertions)]
+        owned_objects.borrow_mut().push(obj);
+        #[cfg(not(debug_assertions))]
+        // SAFETY: This is not re-entrant.
+        unsafe {
+            (*owned_objects.get()).push(obj);
+        }
+    });
 }
 
 /// Increments pyo3's internal GIL count - to be called whenever GILPool or GILGuard is created.
@@ -500,7 +523,11 @@ mod tests {
     }
 
     fn owned_object_count() -> usize {
-        OWNED_OBJECTS.with(|holder| holder.borrow().len())
+        #[cfg(debug_assertions)]
+        let len = OWNED_OBJECTS.with(|owned_objects| owned_objects.borrow().len());
+        #[cfg(not(debug_assertions))]
+        let len = OWNED_OBJECTS.with(|owned_objects| unsafe { (*owned_objects.get()).len() });
+        len
     }
 
     fn pool_inc_refs_does_not_contain(obj: &PyObject) -> bool {