Move GIL counting from GILPool to GILGuard

PyO3 · May 16, 2024 · ba5ff95 · ba5ff95
1 parent 8de1787
commit ba5ff95
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 22 deletions.
diff --git a/src/gil.rs b/src/gil.rs
@@ -222,8 +222,16 @@ impl GILGuard {
         #[allow(deprecated)]
         let pool = unsafe { mem::ManuallyDrop::new(GILPool::new()) };
 
+        increment_gil_count();
+
         GILGuard::Ensured { gstate, pool }
     }
+
+    /// Gets the Python token associated with this [`GILGuard`].
+    #[inline]
+    pub fn python(&self) -> Python<'_> {
+        unsafe { Python::assume_gil_acquired() }
+    }
 }
 
 /// The Drop implementation for `GILGuard` will release the GIL.
@@ -236,6 +244,8 @@ impl Drop for GILGuard {
                 mem::ManuallyDrop::drop(pool);
 
                 ffi::PyGILState_Release(*gstate);
+
+                decrement_gil_count();
             },
         }
     }
@@ -378,7 +388,6 @@ impl GILPool {
     /// As well as requiring the GIL, see the safety notes on `Python::new_pool`.
     #[inline]
     pub unsafe fn new() -> GILPool {
-        increment_gil_count();
         // Update counts of PyObjects / Py that have been cloned or dropped since last acquisition
         #[cfg(not(pyo3_disable_reference_pool))]
         POOL.update_counts(Python::assume_gil_acquired());
@@ -427,7 +436,6 @@ impl Drop for GILPool {
                 }
             }
         }
-        decrement_gil_count();
     }
 }
 
@@ -687,19 +695,19 @@ mod tests {
             assert_eq!(get_gil_count(), 1);
 
             let pool = unsafe { GILPool::new() };
-            assert_eq!(get_gil_count(), 2);
+            assert_eq!(get_gil_count(), 1);
 
             let pool2 = unsafe { GILPool::new() };
-            assert_eq!(get_gil_count(), 3);
+            assert_eq!(get_gil_count(), 1);
 
             drop(pool);
-            assert_eq!(get_gil_count(), 2);
+            assert_eq!(get_gil_count(), 1);
 
             Python::with_gil(|_| {
                 // nested with_gil doesn't update gil count
-                assert_eq!(get_gil_count(), 2);
+                assert_eq!(get_gil_count(), 1);
             });
-            assert_eq!(get_gil_count(), 2);
+            assert_eq!(get_gil_count(), 1);
 
             drop(pool2);
             assert_eq!(get_gil_count(), 1);

diff --git a/src/impl_/trampoline.rs b/src/impl_/trampoline.rs
@@ -9,8 +9,7 @@ use std::{
     panic::{self, UnwindSafe},
 };
 
-#[allow(deprecated)]
-use crate::gil::GILPool;
+use crate::gil::GILGuard;
 use crate::{
     callback::PyCallbackOutput, ffi, ffi_ptr_ext::FfiPtrExt, impl_::panic::PanicTrap,
     methods::IPowModulo, panic::PanicException, types::PyModule, Py, PyResult, Python,
@@ -178,10 +177,9 @@ where
     R: PyCallbackOutput,
 {
     let trap = PanicTrap::new("uncaught panic at ffi boundary");
-    // Necessary to construct a pool until PyO3 0.22 when the GIL Refs API is fully disabled
-    #[allow(deprecated)]
-    let pool = unsafe { GILPool::new() };
-    let py = pool.python();
+
+    let guard = GILGuard::acquire();
+    let py = guard.python();
     let out = panic_result_into_callback_output(
         py,
         panic::catch_unwind(move || -> PyResult<_> { body(py) }),
@@ -225,10 +223,10 @@ where
     F: for<'py> FnOnce(Python<'py>) -> PyResult<()> + UnwindSafe,
 {
     let trap = PanicTrap::new("uncaught panic at ffi boundary");
-    // Necessary to construct a pool until PyO3 0.22 when the GIL Refs API is fully disabled
-    #[allow(deprecated)]
-    let pool = GILPool::new();
-    let py = pool.python();
+
+    let guard = GILGuard::acquire();
+    let py = guard.python();
+
     if let Err(py_err) = panic::catch_unwind(move || body(py))
         .unwrap_or_else(|payload| Err(PanicException::from_panic_payload(payload)))
     {

diff --git a/src/marker.rs b/src/marker.rs
@@ -411,10 +411,10 @@ impl Python<'_> {
     where
         F: for<'py> FnOnce(Python<'py>) -> R,
     {
-        let _guard = GILGuard::acquire();
+        let guard = GILGuard::acquire();
 
         // SAFETY: Either the GIL was already acquired or we just created a new `GILGuard`.
-        f(unsafe { Python::assume_gil_acquired() })
+        f(guard.python())
     }
 
     /// Like [`Python::with_gil`] except Python interpreter state checking is skipped.
@@ -445,10 +445,9 @@ impl Python<'_> {
     where
         F: for<'py> FnOnce(Python<'py>) -> R,
     {
-        let _guard = GILGuard::acquire_unchecked();
+        let guard = GILGuard::acquire_unchecked();
 
-        // SAFETY: Either the GIL was already acquired or we just created a new `GILGuard`.
-        f(Python::assume_gil_acquired())
+        f(guard.python())
     }
 }