Change GILState to be monitored by linked list instead of count

src/gil.rs

@@ -2,13 +2,79 @@
 
 //! Interaction with python's global interpreter lock
 
-use crate::{ffi, internal_tricks::Unsendable, Python};
+use crate::{ffi, Python};
 use std::cell::{Cell, RefCell, UnsafeCell};
 use std::sync::atomic::{spin_loop_hint, AtomicBool, Ordering};
 use std::{any, mem::ManuallyDrop, ptr::NonNull, sync};
 
 static START: sync::Once = sync::Once::new();
 
+#[derive(PartialEq, Eq, Clone, Copy)]
+pub(crate) enum GILStateType {
+    // A pool is currently managing pointer lifetimes
+    Held,
+
+    // The GIL is temporary released for allow_threads
+    Released
+}
+
+/// An internal marker for pyo3's GIL state tracking. Functions as a doubly-linked list.
+pub(crate) struct GILState {
+    ty: GILStateType,
+    previous: *mut GILState,
+    next: *mut GILState
+}
+
+impl GILState {
+    // Add this GILState to the end of the linked-list, and return it.
+    fn new(ty: GILStateType) -> Box<Self> {
+        GIL_STATE.with(|state| {
+            let previous = state.get();
+            let mut slf = Box::new(Self {
+                ty,
+                previous,
+                next: std::ptr::null_mut()
+            });
+
+            // Updating the state requires manipulating raw pointers
+            unsafe {
+                let ptr = &mut (*slf) as *mut _;
+
+                // Update doubly linked list
+                if !previous.is_null() {
+                    (*previous).next = ptr;
+                }
+
+                // Set current state
+                state.set(ptr);
+            }
+
+            slf
+        })
+    }
+}
+
+impl Drop for GILState {
+    // Remove this pointer from the doubly-linked list
+    fn drop(&mut self) {
+        let previous = self.previous;
+        let next = self.next;
+
+        unsafe {
+            if !previous.is_null() {
+                (*previous).next = next;
+            }
+
+            if !next.is_null() {
+                (*next).previous = previous;
+            } else {
+                // This was the end of the linked list; update the thread state
+                GIL_STATE.with(|state| state.set(previous));
+            }
+        }
+    }
+}
+
 thread_local! {
     /// This is a internal counter in pyo3 monitoring whether this thread has the GIL.
     ///
@@ -18,6 +84,13 @@ thread_local! {
     /// As a result, if this thread has the GIL, GIL_COUNT is greater than zero.
     static GIL_COUNT: Cell<u32> = Cell::new(0);
 
+    /// This is pyo3's internal tracking of the GIL for this thread.
+    ///
+    /// It is updated by set_gil_state, which is called by GILPool::new() and by Python::allow_threads().
+    ///
+    /// It is also updated when a GILState drops.
+    static GIL_STATE: Cell<*mut GILState> = Cell::new(std::ptr::null_mut());
+
     /// These are objects owned by the current thread, to be released when the GILPool drops.
     static OWNED_OBJECTS: RefCell<Vec<NonNull<ffi::PyObject>>> = RefCell::new(Vec::with_capacity(256));
 
@@ -33,7 +106,10 @@ thread_local! {
 ///  2) PyGILState_Check always returns 1 if the sub-interpreter APIs have ever been called,
 ///     which could lead to incorrect conclusions that the GIL is held.
 fn gil_is_acquired() -> bool {
-    GIL_COUNT.with(|c| c.get() > 0)
+    GIL_STATE.with(|s| {
+        let state_ptr = s.get();
+        !state_ptr.is_null() && unsafe { (*state_ptr).ty == GILStateType::Held }
+    })
 }
 
 /// Prepares the use of Python in a free-threaded context.
@@ -210,22 +286,22 @@ static POOL: ReleasePool = ReleasePool::new();
 pub struct GILPool {
     owned_objects_start: usize,
     owned_anys_start: usize,
-    // Stable solution for impl !Send
-    no_send: Unsendable,
+    state: Box<GILState>,
 }
 
 impl GILPool {
     /// # Safety
     /// This function requires that GIL is already acquired.
     #[inline]
     pub unsafe fn new() -> GILPool {
+        let state = set_gil_state(GILStateType::Held);
         increment_gil_count();
         // Release objects that were dropped since last GIL acquisition
         POOL.release_pointers(Python::assume_gil_acquired());
         GILPool {
             owned_objects_start: OWNED_OBJECTS.with(|o| o.borrow().len()),
             owned_anys_start: OWNED_ANYS.with(|o| o.borrow().len()),
-            no_send: Unsendable::default(),
+            state,
         }
     }
     pub unsafe fn python(&self) -> Python {
@@ -298,6 +374,10 @@ pub unsafe fn register_any<'p, T: 'static>(obj: T) -> &'p T {
     })
 }
 
+pub(crate) fn set_gil_state(ty: GILStateType) -> Box<GILState> {
+    GILState::new(ty)
+}
+
 /// Increment pyo3's internal GIL count - to be called whenever GILPool or GILGuard is created.
 #[inline(always)]
 fn increment_gil_count() {
@@ -319,7 +399,7 @@ fn decrement_gil_count() {
 
 #[cfg(test)]
 mod test {
-    use super::{GILPool, GIL_COUNT, OWNED_OBJECTS};
+    use super::{GILPool, GIL_COUNT, OWNED_OBJECTS, POOL};
     use crate::{ffi, gil, AsPyPointer, IntoPyPointer, PyObject, Python, ToPyObject};
     use std::ptr::NonNull;
 
@@ -475,4 +555,33 @@ mod test {
         drop(gil);
         assert_eq!(get_gil_count(), 0);
     }
+
+    #[test]
+    fn test_allow_threads() {
+        let gil = Python::acquire_gil();
+        let py = gil.python();
+        let object = get_object();
+
+        py.allow_threads(move || {
+            // Should be no pointers to drop
+            assert!(unsafe { (*POOL.pointers_to_drop.get()).is_empty() });
+
+            // Dropping object without the GIL should put the pointer in the pool
+            drop(object);
+            let obj_count = unsafe { (*POOL.pointers_to_drop.get()).len() };
+            assert_eq!(obj_count, 1);
+
+            // Now repeat dropping an object, with the GIL.
+            let gil = Python::acquire_gil();
+
+            // (Acquiring the GIL should have cleared the pool).
+            assert!(unsafe { (*POOL.pointers_to_drop.get()).is_empty() });
+
+            let object = get_object();
+            drop(object);
+
+            // Previous drop should have decreased count immediately instead of put in pool
+            assert!(unsafe { (*POOL.pointers_to_drop.get()).is_empty() });
+        })
+    }
 }

src/python.rs

@@ -135,9 +135,11 @@ impl<'p> Python<'p> {
         // The `Send` bound on the closure prevents the user from
         // transferring the `Python` token into the closure.
         unsafe {
+            let state = gil::set_gil_state(gil::GILStateType::Released);
             let save = ffi::PyEval_SaveThread();
             let result = f();
             ffi::PyEval_RestoreThread(save);
+            drop(state);
             result
         }
     }