Merge pull request #970 from davidhewitt/py-gil-safe

Cleanups to ensure GIL-safety of Py<T> and PyObject methods
PyO3 · Jun 15, 2020 · 3e905eb · 3e905eb
2 parents 090204c + a85d157
commit 3e905eb
Show file tree

Hide file tree

Showing 17 changed files with 153 additions and 102 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -14,6 +14,9 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 - Call `Py_Finalize` at exit to flush buffers, etc. [#943](https://github.com/PyO3/pyo3/pull/943)
 - Add type parameter to PyBuffer. #[951](https://github.com/PyO3/pyo3/pull/951)
 - Require `Send` bound for `#[pyclass]`. [#966](https://github.com/PyO3/pyo3/pull/966)
+- Add `Python` argument to most methods on `PyObject` and `Py<T>` to ensure GIL safety. [#970](https://github.com/PyO3/pyo3/pull/970)
+- Change signature of `PyTypeObject::type_object()` - now takes `Python` argument and returns `&PyType`. [#970](https://github.com/PyO3/pyo3/pull/970)
+- Change return type of `PyTuple::slice()` and `PyTuple::split_from()` from `Py<PyTuple>` to `&PyTuple`. [#970](https://github.com/PyO3/pyo3/pull/970)
 - Change return type of `PyTuple::as_slice` to `&[&PyAny]`. [#971](https://github.com/PyO3/pyo3/pull/971)
 
 ### Removed

diff --git a/src/conversion.rs b/src/conversion.rs
@@ -419,7 +419,7 @@ pub unsafe trait FromPyPointer<'p>: Sized {
     unsafe fn from_owned_ptr_or_panic(py: Python<'p>, ptr: *mut ffi::PyObject) -> &'p Self {
         match Self::from_owned_ptr_or_opt(py, ptr) {
             Some(s) => s,
-            None => err::panic_after_error(),
+            None => err::panic_after_error(py),
         }
     }
     unsafe fn from_owned_ptr(py: Python<'p>, ptr: *mut ffi::PyObject) -> &'p Self {
@@ -436,7 +436,7 @@ pub unsafe trait FromPyPointer<'p>: Sized {
     unsafe fn from_borrowed_ptr_or_panic(py: Python<'p>, ptr: *mut ffi::PyObject) -> &'p Self {
         match Self::from_borrowed_ptr_or_opt(py, ptr) {
             Some(s) => s,
-            None => err::panic_after_error(),
+            None => err::panic_after_error(py),
         }
     }
     unsafe fn from_borrowed_ptr(py: Python<'p>, ptr: *mut ffi::PyObject) -> &'p Self {

diff --git a/src/err.rs b/src/err.rs
@@ -1,12 +1,13 @@
 // Copyright (c) 2017-present PyO3 Project and Contributors
 
+use crate::gil::ensure_gil;
 use crate::panic::PanicException;
 use crate::type_object::PyTypeObject;
 use crate::types::PyType;
 use crate::{exceptions, ffi};
 use crate::{
-    AsPyPointer, FromPy, FromPyPointer, IntoPy, IntoPyPointer, Py, PyAny, PyObject, Python,
-    ToBorrowedObject, ToPyObject,
+    AsPyPointer, FromPy, FromPyPointer, IntoPy, IntoPyPointer, Py, PyAny, PyNativeType, PyObject,
+    Python, ToBorrowedObject, ToPyObject,
 };
 use libc::c_int;
 use std::ffi::CString;
@@ -88,11 +89,14 @@ impl PyErr {
         T: PyTypeObject,
         V: ToPyObject + 'static,
     {
-        let ty = T::type_object();
+        let gil = ensure_gil();
+        let py = unsafe { gil.python() };
+
+        let ty = T::type_object(py);
         assert_ne!(unsafe { ffi::PyExceptionClass_Check(ty.as_ptr()) }, 0);
 
         PyErr {
-            ptype: ty,
+            ptype: ty.into(),
             pvalue: PyErrValue::ToObject(Box::new(value)),
             ptraceback: None,
         }
@@ -103,12 +107,12 @@ impl PyErr {
     /// `exc` is the exception type; usually one of the standard exceptions
     /// like `exceptions::RuntimeError`.
     /// `args` is the a tuple of arguments to pass to the exception constructor.
-    pub fn from_type<A>(exc: Py<PyType>, args: A) -> PyErr
+    pub fn from_type<A>(exc: &PyType, args: A) -> PyErr
     where
         A: ToPyObject + 'static,
     {
         PyErr {
-            ptype: exc,
+            ptype: exc.into(),
             pvalue: PyErrValue::ToObject(Box::new(args)),
             ptraceback: None,
         }
@@ -119,11 +123,14 @@ impl PyErr {
     where
         T: PyTypeObject,
     {
-        let ty = T::type_object();
+        let gil = ensure_gil();
+        let py = unsafe { gil.python() };
+
+        let ty = T::type_object(py);
         assert_ne!(unsafe { ffi::PyExceptionClass_Check(ty.as_ptr()) }, 0);
 
         PyErr {
-            ptype: ty,
+            ptype: ty.into(),
             pvalue: value,
             ptraceback: None,
         }
@@ -140,19 +147,21 @@ impl PyErr {
 
         if unsafe { ffi::PyExceptionInstance_Check(ptr) } != 0 {
             PyErr {
-                ptype: unsafe { Py::from_borrowed_ptr(ffi::PyExceptionInstance_Class(ptr)) },
+                ptype: unsafe {
+                    Py::from_borrowed_ptr(obj.py(), ffi::PyExceptionInstance_Class(ptr))
+                },
                 pvalue: PyErrValue::Value(obj.into()),
                 ptraceback: None,
             }
         } else if unsafe { ffi::PyExceptionClass_Check(obj.as_ptr()) } != 0 {
             PyErr {
-                ptype: unsafe { Py::from_borrowed_ptr(ptr) },
+                ptype: unsafe { Py::from_borrowed_ptr(obj.py(), ptr) },
                 pvalue: PyErrValue::None,
                 ptraceback: None,
             }
         } else {
             PyErr {
-                ptype: exceptions::TypeError::type_object(),
+                ptype: exceptions::TypeError::type_object(obj.py()).into(),
                 pvalue: PyErrValue::ToObject(Box::new("exceptions must derive from BaseException")),
                 ptraceback: None,
             }
@@ -179,9 +188,9 @@ impl PyErr {
             let mut ptraceback: *mut ffi::PyObject = std::ptr::null_mut();
             ffi::PyErr_Fetch(&mut ptype, &mut pvalue, &mut ptraceback);
 
-            let err = PyErr::new_from_ffi_tuple(ptype, pvalue, ptraceback);
+            let err = PyErr::new_from_ffi_tuple(py, ptype, pvalue, ptraceback);
 
-            if ptype == PanicException::type_object().as_ptr() {
+            if ptype == PanicException::type_object(py).as_ptr() {
                 let msg: String = PyAny::from_borrowed_ptr_or_opt(py, pvalue)
                     .and_then(|obj| obj.extract().ok())
                     .unwrap_or_else(|| String::from("Unwrapped panic from Python code"));
@@ -233,31 +242,30 @@ impl PyErr {
     }
 
     unsafe fn new_from_ffi_tuple(
+        py: Python,
         ptype: *mut ffi::PyObject,
         pvalue: *mut ffi::PyObject,
         ptraceback: *mut ffi::PyObject,
     ) -> PyErr {
         // Note: must not panic to ensure all owned pointers get acquired correctly,
         // and because we mustn't panic in normalize().
 
-        let pvalue = if let Some(obj) =
-            PyObject::from_owned_ptr_or_opt(Python::assume_gil_acquired(), pvalue)
-        {
+        let pvalue = if let Some(obj) = PyObject::from_owned_ptr_or_opt(py, pvalue) {
             PyErrValue::Value(obj)
         } else {
             PyErrValue::None
         };
 
         let ptype = if ptype.is_null() {
-            <exceptions::SystemError as PyTypeObject>::type_object()
+            <exceptions::SystemError as PyTypeObject>::type_object(py).into()
         } else {
-            Py::from_owned_ptr(ptype)
+            Py::from_owned_ptr(py, ptype)
         };
 
         PyErr {
             ptype,
             pvalue,
-            ptraceback: PyObject::from_owned_ptr_or_opt(Python::assume_gil_acquired(), ptraceback),
+            ptraceback: PyObject::from_owned_ptr_or_opt(py, ptraceback),
         }
     }
 
@@ -288,12 +296,12 @@ impl PyErr {
     }
 
     /// Returns true if the current exception is instance of `T`.
-    pub fn is_instance<T>(&self, _py: Python) -> bool
+    pub fn is_instance<T>(&self, py: Python) -> bool
     where
         T: PyTypeObject,
     {
         unsafe {
-            ffi::PyErr_GivenExceptionMatches(self.ptype.as_ptr(), T::type_object().as_ptr()) != 0
+            ffi::PyErr_GivenExceptionMatches(self.ptype.as_ptr(), T::type_object(py).as_ptr()) != 0
         }
     }
 
@@ -328,7 +336,7 @@ impl PyErr {
         let mut ptraceback = ptraceback.into_ptr();
         unsafe {
             ffi::PyErr_NormalizeException(&mut ptype, &mut pvalue, &mut ptraceback);
-            PyErr::new_from_ffi_tuple(ptype, pvalue, ptraceback)
+            PyErr::new_from_ffi_tuple(py, ptype, pvalue, ptraceback)
         }
     }
 
@@ -565,7 +573,7 @@ impl_to_pyerr!(std::string::FromUtf16Error, exceptions::UnicodeDecodeError);
 impl_to_pyerr!(std::char::DecodeUtf16Error, exceptions::UnicodeDecodeError);
 impl_to_pyerr!(std::net::AddrParseError, exceptions::ValueError);
 
-pub fn panic_after_error() -> ! {
+pub fn panic_after_error(_py: Python) -> ! {
     unsafe {
         ffi::PyErr_Print();
     }

diff --git a/src/exceptions.rs b/src/exceptions.rs
@@ -89,7 +89,7 @@ macro_rules! import_exception {
 macro_rules! import_exception_type_object {
     ($module: expr, $name: ident) => {
         unsafe impl $crate::type_object::PyTypeObject for $name {
-            fn type_object() -> $crate::Py<$crate::types::PyType> {
+            fn type_object(py: $crate::Python) -> &$crate::types::PyType {
                 use $crate::type_object::LazyHeapType;
                 static TYPE_OBJECT: LazyHeapType = LazyHeapType::new();
 
@@ -111,7 +111,7 @@ macro_rules! import_exception_type_object {
                     }
                 });
 
-                unsafe { $crate::Py::from_borrowed_ptr(ptr.as_ptr() as *mut $crate::ffi::PyObject) }
+                unsafe { py.from_borrowed_ptr(ptr.as_ptr() as *mut $crate::ffi::PyObject) }
             }
         }
     };
@@ -173,7 +173,7 @@ macro_rules! create_exception {
 macro_rules! create_exception_type_object {
     ($module: ident, $name: ident, $base: ty) => {
         unsafe impl $crate::type_object::PyTypeObject for $name {
-            fn type_object() -> $crate::Py<$crate::types::PyType> {
+            fn type_object(py: $crate::Python) -> &$crate::types::PyType {
                 use $crate::type_object::LazyHeapType;
                 static TYPE_OBJECT: LazyHeapType = LazyHeapType::new();
 
@@ -186,7 +186,7 @@ macro_rules! create_exception_type_object {
                     )
                 });
 
-                unsafe { $crate::Py::from_borrowed_ptr(ptr.as_ptr() as *mut $crate::ffi::PyObject) }
+                unsafe { py.from_borrowed_ptr(ptr.as_ptr() as *mut $crate::ffi::PyObject) }
             }
         }
     };
@@ -215,8 +215,8 @@ macro_rules! impl_native_exception (
             }
         }
         unsafe impl PyTypeObject for $name {
-            fn type_object() -> $crate::Py<$crate::types::PyType> {
-                unsafe { $crate::Py::from_borrowed_ptr(ffi::$exc_name) }
+            fn type_object(py: $crate::Python) -> &$crate::types::PyType {
+                unsafe { py.from_borrowed_ptr(ffi::$exc_name) }
             }
         }
     );

diff --git a/src/gil.rs b/src/gil.rs
@@ -370,6 +370,35 @@ extern "C" fn finalize() {
     }
 }
 
+/// Ensure the GIL is held, useful in implementation of APIs like PyErr::new where it's
+/// inconvenient to force the user to acquire the GIL.
+pub(crate) fn ensure_gil() -> EnsureGIL {
+    if gil_is_acquired() {
+        EnsureGIL(None)
+    } else {
+        EnsureGIL(Some(GILGuard::acquire()))
+    }
+}
+
+/// Struct used internally which avoids acquiring the GIL where it's not necessary.
+pub(crate) struct EnsureGIL(Option<GILGuard>);
+
+impl EnsureGIL {
+    /// Get the GIL token.
+    ///
+    /// # Safety
+    /// If `self.0` is `None`, then this calls [Python::assume_gil_acquired].
+    /// Thus this method could be used to get access to a GIL token while the GIL is not held.
+    /// Care should be taken to only use the returned Python in contexts where it is certain the
+    /// GIL continues to be held.
+    pub unsafe fn python(&self) -> Python {
+        match &self.0 {
+            Some(gil) => gil.python(),
+            None => Python::assume_gil_acquired(),
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::{gil_is_acquired, GILPool, GIL_COUNT, OWNED_OBJECTS, POOL};
@@ -569,53 +598,62 @@ mod test {
     #[test]
     fn test_clone_with_gil() {
         let gil = Python::acquire_gil();
-        let obj = get_object(gil.python());
-        let count = obj.get_refcnt();
+        let py = gil.python();
+
+        let obj = get_object(py);
+        let count = obj.get_refcnt(py);
 
         // Cloning with the GIL should increase reference count immediately
         #[allow(clippy::redundant_clone)]
         let c = obj.clone();
-        assert_eq!(count + 1, c.get_refcnt());
+        assert_eq!(count + 1, c.get_refcnt(py));
     }
 
     #[test]
     fn test_clone_without_gil() {
         let gil = Python::acquire_gil();
         let py = gil.python();
         let obj = get_object(py);
-        let count = obj.get_refcnt();
+        let count = obj.get_refcnt(py);
 
         // Cloning without GIL should not update reference count
         drop(gil);
         let c = obj.clone();
-        assert_eq!(count, obj.get_refcnt());
+        assert_eq!(
+            count,
+            obj.get_refcnt(unsafe { Python::assume_gil_acquired() })
+        );
 
         // Acquring GIL will clear this pending change
         let gil = Python::acquire_gil();
+        let py = gil.python();
 
         // Total reference count should be one higher
-        assert_eq!(count + 1, obj.get_refcnt());
+        assert_eq!(count + 1, obj.get_refcnt(py));
 
-        // Clone dropped then GIL released
+        // Clone dropped
         drop(c);
-        drop(gil);
 
         // Overall count is now back to the original, and should be no pending change
-        assert_eq!(count, obj.get_refcnt());
+        assert_eq!(count, obj.get_refcnt(py));
     }
 
     #[test]
     fn test_clone_in_other_thread() {
         let gil = Python::acquire_gil();
-        let obj = get_object(gil.python());
-        let count = obj.get_refcnt();
+        let py = gil.python();
+        let obj = get_object(py);
+        let count = obj.get_refcnt(py);
 
         // Move obj to a thread which does not have the GIL, and clone it
         let t = std::thread::spawn(move || {
             // Cloning without GIL should not update reference count
             #[allow(clippy::redundant_clone)]
             let _ = obj.clone();
-            assert_eq!(count, obj.get_refcnt());
+            assert_eq!(
+                count,
+                obj.get_refcnt(unsafe { Python::assume_gil_acquired() })
+            );
 
             // Return obj so original thread can continue to use
             obj
@@ -631,13 +669,13 @@ mod test {
 
         // Re-acquring GIL will clear these pending changes
         drop(gil);
-        let _gil = Python::acquire_gil();
+        let gil = Python::acquire_gil();
 
         assert!(POOL.pointers_to_incref.lock().is_empty());
         assert!(POOL.pointers_to_decref.lock().is_empty());
 
         // Overall count is still unchanged
-        assert_eq!(count, obj.get_refcnt());
+        assert_eq!(count, obj.get_refcnt(gil.python()));
     }
 
     #[test]