use MutexExt instead of ThreadStateGuard

Author: bschoenmaeckers

src/datetime.rs

@@ -60,14 +60,14 @@ use std::hash::Hash;
 use std::marker::PhantomData;
 use std::sync::Mutex;
 
+use pyo3::sync::MutexExt;
 use pyo3::{Bound, Py, Python};
 use rustc_hash::FxHashMap;
 
 use crate::dtype::{clone_methods_impl, Element, PyArrayDescr, PyArrayDescrMethods};
 use crate::npyffi::{
     PyArray_DatetimeDTypeMetaData, PyDataType_C_METADATA, NPY_DATETIMEUNIT, NPY_TYPES,
 };
-use crate::ThreadStateGuard;
 
 /// Represents the [datetime units][datetime-units] supported by NumPy
 ///
@@ -224,13 +224,10 @@ impl TypeDescriptors {
 
     #[allow(clippy::wrong_self_convention)]
     fn from_unit<'py>(&self, py: Python<'py>, unit: NPY_DATETIMEUNIT) -> Bound<'py, PyArrayDescr> {
-        // Detach from the runtime to avoid deadlocking on acquiring the mutex.
-        let ts_guard = ThreadStateGuard::new();
-
-        let mut dtypes = self.dtypes.lock().expect("dtype cache poisoned");
-
-        // Now we hold the mutex so it's safe to re-attach to the runtime.
-        drop(ts_guard);
+        let mut dtypes = self
+            .dtypes
+            .lock_py_attached(py)
+            .expect("dtype cache poisoned");
 
         let dtype = match dtypes.get_or_insert_with(Default::default).entry(unit) {
             Entry::Occupied(entry) => entry.into_mut(),

src/lib.rs

@@ -149,23 +149,6 @@ mod doctest {
 #[inline(always)]
 fn cold() {}
 
-/// An RAII guard for avoiding deadlocks with the GIL or other global
-/// synchronization events in the Python runtime
-// FIXME create a proper MutexExt trait that handles poisoning and upstream to PyO3
-struct ThreadStateGuard(*mut pyo3::ffi::PyThreadState);
-
-impl ThreadStateGuard {
-    fn new() -> ThreadStateGuard {
-        ThreadStateGuard(unsafe { pyo3::ffi::PyEval_SaveThread() })
-    }
-}
-
-impl Drop for ThreadStateGuard {
-    fn drop(&mut self) {
-        unsafe { pyo3::ffi::PyEval_RestoreThread(self.0) };
-    }
-}
-
 /// Create a [`PyArray`] with one, two or three dimensions.
 ///
 /// This macro is backed by [`ndarray::array`].

src/strings.rs

@@ -10,6 +10,7 @@ use std::os::raw::c_char;
 use std::str;
 use std::sync::Mutex;
 
+use pyo3::sync::MutexExt;
 use pyo3::{
     ffi::{Py_UCS1, Py_UCS4},
     Bound, Py, Python,
@@ -19,7 +20,6 @@ use rustc_hash::FxHashMap;
 use crate::dtype::{clone_methods_impl, Element, PyArrayDescr, PyArrayDescrMethods};
 use crate::npyffi::PyDataType_SET_ELSIZE;
 use crate::npyffi::NPY_TYPES;
-use crate::ThreadStateGuard;
 
 /// A newtype wrapper around [`[u8; N]`][Py_UCS1] to handle [`byte` scalars][numpy-bytes] while satisfying coherence.
 ///
@@ -179,13 +179,10 @@ impl TypeDescriptors {
         byteorder: c_char,
         size: usize,
     ) -> Bound<'py, PyArrayDescr> {
-        // Detach from the runtime to avoid deadlocking on acquiring the mutex.
-        let ts_guard = ThreadStateGuard::new();
-
-        let mut dtypes = self.dtypes.lock().expect("dtype cache poisoned");
-
-        // Now we hold the mutex so it's safe to re-attach to the runtime.
-        drop(ts_guard);
+        let mut dtypes = self
+            .dtypes
+            .lock_py_attached(py)
+            .expect("dtype cache poisoned");
 
         let dtype = match dtypes.get_or_insert_with(Default::default).entry(size) {
             Entry::Occupied(entry) => entry.into_mut(),