diff --git a/src/gil.rs b/src/gil.rs index abcb160595d..3d8bc731679 100644 --- a/src/gil.rs +++ b/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 { + 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 = 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>> = 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,8 +286,7 @@ 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, } impl GILPool { @@ -219,13 +294,14 @@ impl GILPool { /// 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::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() }); + }) + } } diff --git a/src/python.rs b/src/python.rs index 14af58597db..a149b02ea37 100644 --- a/src/python.rs +++ b/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 } }