Add pyo3::once_cell::OnceCell

src/exceptions.rs

@@ -90,28 +90,26 @@ macro_rules! import_exception_type_object {
     ($module: expr, $name: ident) => {
         unsafe impl $crate::type_object::PyTypeObject for $name {
             fn type_object(py: $crate::Python) -> &$crate::types::PyType {
-                use $crate::type_object::LazyHeapType;
-                static TYPE_OBJECT: LazyHeapType = LazyHeapType::new();
-
-                let ptr = TYPE_OBJECT.get_or_init(|py| {
-                    let imp = py
-                        .import(stringify!($module))
-                        .expect(concat!("Can not import module: ", stringify!($module)));
-                    let cls = imp.get(stringify!($name)).expect(concat!(
-                        "Can not load exception class: {}.{}",
-                        stringify!($module),
-                        ".",
-                        stringify!($name)
-                    ));
-
-                    unsafe {
-                        std::ptr::NonNull::new_unchecked(
-                            $crate::IntoPyPointer::into_ptr(cls) as *mut _
-                        )
-                    }
-                });
-
-                unsafe { py.from_borrowed_ptr(ptr.as_ptr() as *mut $crate::ffi::PyObject) }
+                use $crate::once_cell::OnceCell;
+                use $crate::AsPyRef;
+                static TYPE_OBJECT: OnceCell<$crate::Py<$crate::types::PyType>> = OnceCell::new();
+
+                TYPE_OBJECT
+                    .get_or_init(py, || {
+                        let imp = py
+                            .import(stringify!($module))
+                            .expect(concat!("Can not import module: ", stringify!($module)));
+                        let cls = imp.get(stringify!($name)).expect(concat!(
+                            "Can not load exception class: {}.{}",
+                            stringify!($module),
+                            ".",
+                            stringify!($name)
+                        ));
+
+                        cls.extract()
+                            .expect("Imported exception should be a type object")
+                    })
+                    .as_ref(py)
             }
         }
     };
@@ -174,19 +172,24 @@ macro_rules! create_exception_type_object {
     ($module: ident, $name: ident, $base: ty) => {
         unsafe impl $crate::type_object::PyTypeObject for $name {
             fn type_object(py: $crate::Python) -> &$crate::types::PyType {
-                use $crate::type_object::LazyHeapType;
-                static TYPE_OBJECT: LazyHeapType = LazyHeapType::new();
-
-                let ptr = TYPE_OBJECT.get_or_init(|py| {
-                    $crate::PyErr::new_type(
-                        py,
-                        concat!(stringify!($module), ".", stringify!($name)),
-                        Some(py.get_type::<$base>()),
-                        None,
-                    )
-                });
-
-                unsafe { py.from_borrowed_ptr(ptr.as_ptr() as *mut $crate::ffi::PyObject) }
+                use $crate::once_cell::OnceCell;
+                use $crate::AsPyRef;
+                static TYPE_OBJECT: OnceCell<$crate::Py<$crate::types::PyType>> = OnceCell::new();
+
+                TYPE_OBJECT
+                    .get_or_init(py, || unsafe {
+                        $crate::Py::from_owned_ptr(
+                            py,
+                            $crate::PyErr::new_type(
+                                py,
+                                concat!(stringify!($module), ".", stringify!($name)),
+                                Some(py.get_type::<$base>()),
+                                None,
+                            )
+                            .as_ptr() as *mut $crate::ffi::PyObject,
+                        )
+                    })
+                    .as_ref(py)
             }
         }
     };

src/ffi/datetime.rs

@@ -14,10 +14,10 @@
 use crate::ffi::Py_hash_t;
 use crate::ffi::{PyObject, PyTypeObject};
 use crate::ffi::{PyObject_TypeCheck, Py_TYPE};
+use crate::once_cell::OnceCell;
+use crate::Python;
 use std::ops::Deref;
 use std::os::raw::{c_char, c_int, c_uchar};
-use std::ptr;
-use std::sync::Once;
 #[cfg(not(PyPy))]
 use {crate::ffi::PyCapsule_Import, std::ffi::CString};
 
@@ -196,29 +196,9 @@ pub struct PyDateTime_Delta {
     pub microseconds: c_int,
 }
 
-// C API Capsule
-// Note: This is "roll-your-own" lazy-static implementation is necessary because
-// of the interaction between the call_once locks and the GIL. It turns out that
-// calling PyCapsule_Import releases and re-acquires the GIL during the import,
-// so if you have two threads attempting to use the PyDateTimeAPI singleton
-// under the GIL, it causes a deadlock; what happens is:
-//
-// Thread 1 acquires GIL
-// Thread 1 acquires the call_once lock
-// Thread 1 calls PyCapsule_Import, thus releasing the GIL
-// Thread 2 acquires the GIL
-// Thread 2 blocks waiting for the call_once lock
-// Thread 1 blocks waiting for the GIL
-//
-// However, Python's import mechanism acquires a module-specific lock around
-// each import call, so all call importing datetime will return the same
-// object, making the call_once lock superfluous. As such, we can weaken
-// the guarantees of the cache, such that PyDateTime_IMPORT can be called
-// until __PY_DATETIME_API_UNSAFE_CACHE is populated, which will happen exactly
-// one time. So long as PyDateTime_IMPORT has no side effects (it should not),
-// this will be at most a slight waste of resources.
-static PY_DATETIME_API_ONCE: Once = Once::new();
-static mut PY_DATETIME_API_UNSAFE_CACHE: *const PyDateTime_CAPI = ptr::null();
+// Python already shares this object between threads, so it's no more evil for us to do it too!
+unsafe impl Sync for PyDateTime_CAPI {}
+static PY_DATETIME_API: OnceCell<&'static PyDateTime_CAPI> = OnceCell::new();
 
 #[derive(Debug)]
 pub struct PyDateTimeAPI {
@@ -233,13 +213,7 @@ impl Deref for PyDateTimeAPI {
     type Target = PyDateTime_CAPI;
 
     fn deref(&self) -> &'static PyDateTime_CAPI {
-        unsafe {
-            if !PY_DATETIME_API_UNSAFE_CACHE.is_null() {
-                &(*PY_DATETIME_API_UNSAFE_CACHE)
-            } else {
-                PyDateTime_IMPORT()
-            }
-        }
+        unsafe { PyDateTime_IMPORT() }
     }
 }
 
@@ -251,25 +225,27 @@ impl Deref for PyDateTimeAPI {
 /// Use this function only if you want to eagerly load the datetime module,
 /// such as if you do not want the first call to a datetime function to be
 /// slightly slower than subsequent calls.
+///
+/// # Safety
+/// The Python GIL must be held.
 pub unsafe fn PyDateTime_IMPORT() -> &'static PyDateTime_CAPI {
-    // PyPy expects the C-API to be initialized via PyDateTime_Import, so trying to use
-    // `PyCapsule_Import` will behave unexpectedly in pypy.
-    #[cfg(PyPy)]
-    let py_datetime_c_api = PyDateTime_Import();
-
-    #[cfg(not(PyPy))]
-    let py_datetime_c_api = {
-        // PyDateTime_CAPSULE_NAME is a macro in C
-        let PyDateTime_CAPSULE_NAME = CString::new("datetime.datetime_CAPI").unwrap();
-
-        PyCapsule_Import(PyDateTime_CAPSULE_NAME.as_ptr(), 1) as *const PyDateTime_CAPI
-    };
-
-    PY_DATETIME_API_ONCE.call_once(move || {
-        PY_DATETIME_API_UNSAFE_CACHE = py_datetime_c_api;
-    });
-
-    &(*PY_DATETIME_API_UNSAFE_CACHE)
+    let py = Python::assume_gil_acquired();
+    PY_DATETIME_API.get_or_init(py, || {
+        // PyPy expects the C-API to be initialized via PyDateTime_Import, so trying to use
+        // `PyCapsule_Import` will behave unexpectedly in pypy.
+        #[cfg(PyPy)]
+        let py_datetime_c_api = PyDateTime_Import();
+
+        #[cfg(not(PyPy))]
+        let py_datetime_c_api = {
+            // PyDateTime_CAPSULE_NAME is a macro in C
+            let PyDateTime_CAPSULE_NAME = CString::new("datetime.datetime_CAPI").unwrap();
+
+            &*(PyCapsule_Import(PyDateTime_CAPSULE_NAME.as_ptr(), 1) as *const PyDateTime_CAPI)
+        };
+
+        py_datetime_c_api
+    })
 }
 
 /// Type Check macros

src/lib.rs

@@ -183,6 +183,7 @@ mod instance;
 mod internal_tricks;
 pub mod marshal;
 mod object;
+pub mod once_cell;
 pub mod panic;
 pub mod prelude;
 pub mod pycell;

src/once_cell.rs

@@ -0,0 +1,58 @@
+use crate::Python;
+use std::cell::UnsafeCell;
+
+/// A write-once cell similar to once_cell::OnceCell.
+///
+/// Unlike once_cell::sync which uses locks to achieve thread safety, this implementation uses
+/// the Python GIL to mediate concurrent access.
+pub struct OnceCell<T>(UnsafeCell<Option<T>>);
+
+// T: Send is needed for Sync because the thread which drops the OnceCell can be different
+// to the thread which fills it.
+unsafe impl<T: Send + Sync> Sync for OnceCell<T> {}
+unsafe impl<T: Send> Send for OnceCell<T> {}
+
+impl<T> OnceCell<T> {
+    pub const fn new() -> Self {
+        Self(UnsafeCell::new(None))
+    }
+
+    pub fn get(&self, _py: Python) -> Option<&T> {
+        // Safe because if the cell has not yet been written, None is returned.
+        unsafe { &*self.0.get() }.as_ref()
+    }
+
+    pub fn get_or_init<F>(&self, py: Python, f: F) -> &T
+    where
+        F: FnOnce() -> T,
+    {
+        let inner = unsafe { &*self.0.get() }.as_ref();
+        if let Some(value) = inner {
+            return value;
+        }
+
+        // Note that f() could release the GIL, so it's possible that another thread writes to this
+        // OnceCell before f() finishes. That's fine; we'll just have to discard the value
+        // computed here and accept a bit of wasted computation.
+        let value = f();
+        let _ = self.set(py, value);
+
+        self.get(py).unwrap()
+    }
+
+    pub fn get_mut(&mut self) -> Option<&mut T> {
+        // Safe because we have &mut self
+        unsafe { &mut *self.0.get() }.as_mut()
+    }
+
+    pub fn set(&self, _py: Python, value: T) -> Result<(), T> {
+        // Safe because GIL is held, so no other thread can be writing to this cell concurrently.
+        let inner = unsafe { &mut *self.0.get() };
+        if inner.is_some() {
+            return Err(value);
+        }
+
+        *inner = Some(value);
+        Ok(())
+    }
+}

src/type_object.rs

@@ -6,7 +6,6 @@ use crate::pyclass_init::PyObjectInit;
 use crate::types::{PyAny, PyType};
 use crate::{ffi, AsPyPointer, Python};
 use std::cell::UnsafeCell;
-use std::ptr::NonNull;
 use std::sync::atomic::{AtomicBool, Ordering};
 
 /// `T: PyLayout<U>` represents that `T` is a concrete representaion of `U` in Python heap.
@@ -136,45 +135,6 @@ where
     }
 }
 
-/// Lazy type object for Exceptions
-#[doc(hidden)]
-pub struct LazyHeapType {
-    value: UnsafeCell<Option<NonNull<ffi::PyTypeObject>>>,
-    initialized: AtomicBool,
-}
-
-impl LazyHeapType {
-    pub const fn new() -> Self {
-        LazyHeapType {
-            value: UnsafeCell::new(None),
-            initialized: AtomicBool::new(false),
-        }
-    }
-
-    pub fn get_or_init<F>(&self, constructor: F) -> NonNull<ffi::PyTypeObject>
-    where
-        F: Fn(Python) -> NonNull<ffi::PyTypeObject>,
-    {
-        if !self
-            .initialized
-            .compare_and_swap(false, true, Ordering::Acquire)
-        {
-            // We have to get the GIL before setting the value to the global!!!
-            let gil = Python::acquire_gil();
-            unsafe {
-                *self.value.get() = Some(constructor(gil.python()));
-            }
-        }
-        unsafe { (*self.value.get()).unwrap() }
-    }
-}
-
-// This is necessary for making static `LazyHeapType`s
-//
-// Type objects are shared between threads by the Python interpreter anyway, so it is no worse
-// to allow sharing on the Rust side too.
-unsafe impl Sync for LazyHeapType {}
-
 /// Lazy type object for PyClass
 #[doc(hidden)]
 pub struct LazyStaticType {