add sync::OnceExt trait

Author: davidhewitt

Diff for: src/sealed.rs

@@ -53,3 +53,5 @@ impl Sealed for ModuleDef {}
 
 impl<T: crate::type_object::PyTypeInfo> Sealed for PyNativeTypeInitializer<T> {}
 impl<T: crate::pyclass::PyClass> Sealed for PyClassInitializer<T> {}
+
+impl Sealed for std::sync::Once {}

Diff for: src/sync.rs

@@ -5,10 +5,17 @@
 //!
 //! [PEP 703]: https://peps.python.org/pep-703/
 use crate::{
+    ffi,
+    sealed::Sealed,
     types::{any::PyAnyMethods, PyAny, PyString},
     Bound, Py, PyResult, PyTypeCheck, Python,
 };
-use std::{cell::UnsafeCell, marker::PhantomData, mem::MaybeUninit, sync::Once};
+use std::{
+    cell::UnsafeCell,
+    marker::PhantomData,
+    mem::MaybeUninit,
+    sync::{Once, OnceState},
+};
 
 #[cfg(not(Py_GIL_DISABLED))]
 use crate::PyVisit;
@@ -473,6 +480,58 @@ where
     }
 }
 
+/// Helper trait for `Once` to help avoid deadlocking when using a `Once` when attached to a
+/// Python thread.
+pub trait OnceExt: Sealed {
+    /// Similar to [`call_once`][Once::call_once], but releases the Python GIL temporarily
+    /// if blocking on another thread currently calling this `Once`.
+    fn call_once_py_attached(&self, py: Python<'_>, f: impl FnOnce());
+
+    /// Similar to [`call_once_force`][Once::call_once_force], but releases the Python GIL
+    /// temporarily if blocking on another thread currently calling this `Once`.
+    fn call_once_force_py_attached(&self, py: Python<'_>, f: impl FnOnce(&OnceState));
+}
+
+impl OnceExt for Once {
+    fn call_once_py_attached(&self, _py: Python<'_>, f: impl FnOnce()) {
+        if self.is_completed() {
+            return;
+        }
+
+        // Safety: we are currently attached to the GIL, and we expect to block. We will save
+        // the current thread state and restore it as soon as we are done blocking.
+        let mut ts = Some(unsafe { ffi::PyEval_SaveThread() });
+
+        self.call_once(|| {
+            unsafe { ffi::PyEval_RestoreThread(ts.take().unwrap()) };
+            f();
+        });
+        if let Some(ts) = ts {
+            // Some other thread filled this Once, so we need to restore the GIL state.
+            unsafe { ffi::PyEval_RestoreThread(ts) };
+        }
+    }
+
+    fn call_once_force_py_attached(&self, _py: Python<'_>, f: impl FnOnce(&OnceState)) {
+        if self.is_completed() {
+            return;
+        }
+
+        // Safety: we are currently attached to the GIL, and we expect to block. We will save
+        // the current thread state and restore it as soon as we are done blocking.
+        let mut ts = Some(unsafe { ffi::PyEval_SaveThread() });
+
+        self.call_once_force(|state| {
+            unsafe { ffi::PyEval_RestoreThread(ts.take().unwrap()) };
+            f(state);
+        });
+        if let Some(ts) = ts {
+            // Some other thread filled this Once, so we need to restore the GIL state.
+            unsafe { ffi::PyEval_RestoreThread(ts) };
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

Diff for: tests/test_declarative_module.rs

@@ -1,9 +1,11 @@
 #![cfg(feature = "macros")]
 
+use std::sync::Once;
+
 use pyo3::create_exception;
 use pyo3::exceptions::PyException;
 use pyo3::prelude::*;
-use pyo3::sync::GILOnceCell;
+use pyo3::sync::{GILOnceCell, OnceExt};
 
 #[path = "../src/tests/common.rs"]
 mod common;
@@ -149,9 +151,17 @@ mod declarative_module2 {
 
 fn declarative_module(py: Python<'_>) -> &Bound<'_, PyModule> {
     static MODULE: GILOnceCell<Py<PyModule>> = GILOnceCell::new();
-    MODULE
-        .get_or_init(py, || pyo3::wrap_pymodule!(declarative_module)(py))
-        .bind(py)
+    static ONCE: Once = Once::new();
+
+    // Guarantee that the module is only ever initialized once; GILOnceCell can race.
+    // TODO: use OnceLock when MSRV >= 1.70
+    ONCE.call_once_py_attached(py, || {
+        MODULE
+            .set(py, pyo3::wrap_pymodule!(declarative_module)(py))
+            .expect("only ever set once");
+    });
+
+    MODULE.get(py).expect("once is completed").bind(py)
 }
 
 #[test]