diff --git a/src/sealed.rs b/src/sealed.rs index cc835bee3b8..0a2846b134a 100644 --- a/src/sealed.rs +++ b/src/sealed.rs @@ -53,3 +53,5 @@ impl Sealed for ModuleDef {} impl Sealed for PyNativeTypeInitializer {} impl Sealed for PyClassInitializer {} + +impl Sealed for std::sync::Once {} diff --git a/src/sync.rs b/src/sync.rs index 65a81d06bd5..161327470cc 100644 --- a/src/sync.rs +++ b/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 --git a/tests/test_declarative_module.rs b/tests/test_declarative_module.rs index a911702ce20..93e0e1366f0 100644 --- a/tests/test_declarative_module.rs +++ b/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> = 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]