Move GIL counting from GILPool to GILGuard (#4188)

src/gil.rs

@@ -167,6 +167,7 @@ impl GILGuard {
     /// `GILGuard::Ensured` will be returned.
     pub(crate) fn acquire() -> Self {
         if gil_is_acquired() {
+            increment_gil_count();
             return GILGuard::Assumed;
         }
 
@@ -215,15 +216,29 @@ impl GILGuard {
     /// as part of multi-phase interpreter initialization.
     pub(crate) fn acquire_unchecked() -> Self {
         if gil_is_acquired() {
+            increment_gil_count();
             return GILGuard::Assumed;
         }
 
         let gstate = unsafe { ffi::PyGILState_Ensure() }; // acquire GIL
         #[allow(deprecated)]
         let pool = unsafe { mem::ManuallyDrop::new(GILPool::new()) };
 
+        increment_gil_count();
+
         GILGuard::Ensured { gstate, pool }
     }
+    /// Acquires the `GILGuard` while assuming that the GIL is already held.
+    pub(crate) unsafe fn assume() -> Self {
+        increment_gil_count();
+        GILGuard::Assumed
+    }
+
+    /// 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.
@@ -238,6 +253,7 @@ impl Drop for GILGuard {
                 ffi::PyGILState_Release(*gstate);
             },
         }
+        decrement_gil_count();
     }
 }
 
@@ -378,7 +394,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 +442,6 @@ impl Drop for GILPool {
                 }
             }
         }
-        decrement_gil_count();
     }
 }
 
@@ -687,19 +701,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(), 2);
+            assert_eq!(get_gil_count(), 1);
 
             drop(pool2);
             assert_eq!(get_gil_count(), 1);

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,
@@ -171,17 +170,19 @@ trampoline!(
 ///
 /// Panics during execution are trapped so that they don't propagate through any
 /// outer FFI boundary.
+///
+/// The GIL must already be held when this is called.
 #[inline]
-pub(crate) fn trampoline<F, R>(body: F) -> R
+pub(crate) unsafe fn trampoline<F, R>(body: F) -> R
 where
     F: for<'py> FnOnce(Python<'py>) -> PyResult<R> + UnwindSafe,
     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();
+
+    // SAFETY: This function requires the GIL to already be held.
+    let guard = unsafe { GILGuard::assume() };
+    let py = guard.python();
     let out = panic_result_into_callback_output(
         py,
         panic::catch_unwind(move || -> PyResult<_> { body(py) }),
@@ -218,17 +219,19 @@ where
 ///
 /// # Safety
 ///
-/// ctx must be either a valid ffi::PyObject or NULL
+/// - ctx must be either a valid ffi::PyObject or NULL
+/// - The GIL must already be held when this is called.
 #[inline]
 unsafe fn trampoline_unraisable<F>(body: F, ctx: *mut ffi::PyObject)
 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();
+
+    // SAFETY: The GIL is already held.
+    let guard = unsafe { GILGuard::assume() };
+    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)))
     {

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())
     }
 }