add sync::OnceExt and sync::OnceLockExt traits (#4676)

* add `sync::OnceExt` trait

* newsfragment

* refactor to use RAII guard

* Add docs for single-initialization

* mark init logic with #[cold]

* attempt to include OnceLockExt as well

* Add OnceLockExt

* ignore clippy MSRV lint

* simplify

* fix wasm tests

---------

Co-authored-by: Nathan Goldbaum <[email protected]>

guide/src/faq.md

@@ -13,9 +13,11 @@ Sorry that you're having trouble using PyO3. If you can't find the answer to you
 5. Thread A is blocked, because it waits to re-acquire the GIL which thread B still holds.
 6. Deadlock.
 
-PyO3 provides a struct [`GILOnceCell`] which works similarly to these types but avoids risk of deadlocking with the Python GIL. This means it can be used in place of other choices when you are experiencing the deadlock described above. See the documentation for [`GILOnceCell`] for further details and an example how to use it.
+PyO3 provides a struct [`GILOnceCell`] which implements a single-initialization API based on these types that relies on the GIL for locking. If the GIL is released or there is no GIL, then this type allows the initialization function to race but ensures that the data is only ever initialized once. If you need to ensure that the initialization function is called once and only once, you can make use of the [`OnceExt`] and [`OnceLockExt`] extension traits that enable using the standard library types for this purpose but provide new methods for these types that avoid the risk of deadlocking with the Python GIL. This means they can be used in place of other choices when you are experiencing the deadlock described above. See the documentation for [`GILOnceCell`] and [`OnceExt`] for further details and an example how to use them.
 
 [`GILOnceCell`]: {{#PYO3_DOCS_URL}}/pyo3/sync/struct.GILOnceCell.html
+[`OnceExt`]: {{#PYO3_DOCS_URL}}/pyo3/sync/trait.OnceExt.html
+[`OnceLockExt`]: {{#PYO3_DOCS_URL}}/pyo3/sync/trait.OnceLockExt.html
 
 ## I can't run `cargo test`; or I can't build in a Cargo workspace: I'm having linker issues like "Symbol not found" or "Undefined reference to _PyExc_SystemError"!
 

guide/src/free-threading.md

@@ -152,6 +152,60 @@ We plan to allow user-selectable semantics for mutable pyclass definitions in
 PyO3 0.24, allowing some form of opt-in locking to emulate the GIL if that is
 needed.
 
+## Thread-safe single initialization
+
+Until version 0.23, PyO3 provided only `GILOnceCell` to enable deadlock-free
+single initialization of data in contexts that might execute arbitrary Python
+code. While we have updated `GILOnceCell` to avoid thread safety issues
+triggered only under the free-threaded build, the design of `GILOnceCell` is
+inherently thread-unsafe, in a manner that can be problematic even in the
+GIL-enabled build.
+
+If, for example, the function executed by `GILOnceCell` releases the GIL or
+calls code that releases the GIL, then it is possible for multiple threads to
+try to race to initialize the cell. While the cell will only ever be intialized
+once, it can be problematic in some contexts that `GILOnceCell` does not block
+like the standard library `OnceLock`.
+
+In cases where the initialization function must run exactly once, you can bring
+the `OnceExt` or `OnceLockExt` traits into scope. The `OnceExt` trait adds
+`OnceExt::call_once_py_attached` and `OnceExt::call_once_force_py_attached`
+functions to the api of `std::sync::Once`, enabling use of `Once` in contexts
+where the GIL is held. Similarly, `OnceLockExt` adds
+`OnceLockExt::get_or_init_py_attached`. These functions are analogous to
+`Once::call_once`, `Once::call_once_force`, and `OnceLock::get_or_init` except
+they accept a `Python<'py>` token in addition to an `FnOnce`. All of these
+functions release the GIL and re-acquire it before executing the function,
+avoiding deadlocks with the GIL that are possible without using the PyO3
+extension traits. Here is an example of how to use `OnceExt` to
+enable single-initialization of a runtime cache holding a `Py<PyDict>`.
+
+```rust
+# fn main() {
+# use pyo3::prelude::*;
+use std::sync::Once;
+use pyo3::sync::OnceExt;
+use pyo3::types::PyDict;
+
+struct RuntimeCache {
+    once: Once,
+    cache: Option<Py<PyDict>>
+}
+
+let mut cache = RuntimeCache {
+    once: Once::new(),
+    cache: None
+};
+
+Python::with_gil(|py| {
+    // guaranteed to be called once and only once
+    cache.once.call_once_py_attached(py, || {
+        cache.cache = Some(PyDict::new(py).unbind());
+    });
+});
+# }
+```
+
 ## `GILProtected` is not exposed
 
 `GILProtected` is a PyO3 type that allows mutable access to static data by

guide/src/migration.md

@@ -230,7 +230,12 @@ PyO3 0.23 introduces preliminary support for the new free-threaded build of
 CPython 3.13. PyO3 features that implicitly assumed the existence of the GIL are
 not exposed in the free-threaded build, since they are no longer safe.  Other
 features, such as `GILOnceCell`, have been internally rewritten to be threadsafe
-without the GIL.
+without the GIL, although note that `GILOnceCell` is inherently racey. You can
+also use `OnceExt::call_once_py_attached` or
+`OnceExt::call_once_force_py_attached` to enable use of `std::sync::Once` in
+code that has the GIL acquired without risking a dealock with the GIL. We plan
+We plan to expose more extension traits in the future that make it easier to
+write code for the GIL-enabled and free-threaded builds of Python.
 
 If you make use of these features then you will need to account for the
 unavailability of this API in the free-threaded build. One way to handle it is

newsfragments/4676.added.md

@@ -0,0 +1 @@
+Add `pyo3::sync::OnceExt` and `pyo3::sync::OnceLockExt` traits.

pyo3-build-config/src/lib.rs

@@ -138,6 +138,10 @@ fn resolve_cross_compile_config_path() -> Option<PathBuf> {
 pub fn print_feature_cfgs() {
     let rustc_minor_version = rustc_minor_version().unwrap_or(0);
 
+    if rustc_minor_version >= 70 {
+        println!("cargo:rustc-cfg=rustc_has_once_lock");
+    }
+
     // invalid_from_utf8 lint was added in Rust 1.74
     if rustc_minor_version >= 74 {
         println!("cargo:rustc-cfg=invalid_from_utf8_lint");
@@ -175,6 +179,7 @@ pub fn print_expected_cfgs() {
     println!("cargo:rustc-check-cfg=cfg(pyo3_leak_on_drop_without_reference_pool)");
     println!("cargo:rustc-check-cfg=cfg(diagnostic_namespace)");
     println!("cargo:rustc-check-cfg=cfg(c_str_lit)");
+    println!("cargo:rustc-check-cfg=cfg(rustc_has_once_lock)");
 
     // allow `Py_3_*` cfgs from the minimum supported version up to the
     // maximum minor version (+1 for development for the next)

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 {}

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,139 @@ where
     }
 }
 
+#[cfg(rustc_has_once_lock)]
+mod once_lock_ext_sealed {
+    pub trait Sealed {}
+    impl<T> Sealed for std::sync::OnceLock<T> {}
+}
+
+/// 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));
+}
+
+// Extension trait for [`std::sync::OnceLock`] which helps avoid deadlocks between the Python
+/// interpreter and initialization with the `OnceLock`.
+#[cfg(rustc_has_once_lock)]
+pub trait OnceLockExt<T>: once_lock_ext_sealed::Sealed {
+    /// Initializes this `OnceLock` with the given closure if it has not been initialized yet.
+    ///
+    /// If this function would block, this function detaches from the Python interpreter and
+    /// reattaches before calling `f`. This avoids deadlocks between the Python interpreter and
+    /// the `OnceLock` in cases where `f` can call arbitrary Python code, as calling arbitrary
+    /// Python code can lead to `f` itself blocking on the Python interpreter.
+    ///
+    /// By detaching from the Python interpreter before blocking, this ensures that if `f` blocks
+    /// then the Python interpreter cannot be blocked by `f` itself.
+    fn get_or_init_py_attached<F>(&self, py: Python<'_>, f: F) -> &T
+    where
+        F: FnOnce() -> T;
+}
+
+struct Guard(*mut crate::ffi::PyThreadState);
+
+impl Drop for Guard {
+    fn drop(&mut self) {
+        unsafe { ffi::PyEval_RestoreThread(self.0) };
+    }
+}
+
+impl OnceExt for Once {
+    fn call_once_py_attached(&self, py: Python<'_>, f: impl FnOnce()) {
+        if self.is_completed() {
+            return;
+        }
+
+        init_once_py_attached(self, py, f)
+    }
+
+    fn call_once_force_py_attached(&self, py: Python<'_>, f: impl FnOnce(&OnceState)) {
+        if self.is_completed() {
+            return;
+        }
+
+        init_once_force_py_attached(self, py, f);
+    }
+}
+
+#[cfg(rustc_has_once_lock)]
+impl<T> OnceLockExt<T> for std::sync::OnceLock<T> {
+    fn get_or_init_py_attached<F>(&self, py: Python<'_>, f: F) -> &T
+    where
+        F: FnOnce() -> T,
+    {
+        // this trait is guarded by a rustc version config
+        // so clippy's MSRV check is wrong
+        #[allow(clippy::incompatible_msrv)]
+        // Use self.get() first to create a fast path when initialized
+        self.get()
+            .unwrap_or_else(|| init_once_lock_py_attached(self, py, f))
+    }
+}
+
+#[cold]
+fn init_once_py_attached<F, T>(once: &Once, _py: Python<'_>, f: F)
+where
+    F: FnOnce() -> T,
+{
+    // 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 ts_guard = Guard(unsafe { ffi::PyEval_SaveThread() });
+
+    once.call_once(move || {
+        drop(ts_guard);
+        f();
+    });
+}
+
+#[cold]
+fn init_once_force_py_attached<F, T>(once: &Once, _py: Python<'_>, f: F)
+where
+    F: FnOnce(&OnceState) -> T,
+{
+    // 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 ts_guard = Guard(unsafe { ffi::PyEval_SaveThread() });
+
+    once.call_once_force(move |state| {
+        drop(ts_guard);
+        f(state);
+    });
+}
+
+#[cfg(rustc_has_once_lock)]
+#[cold]
+fn init_once_lock_py_attached<'a, F, T>(
+    lock: &'a std::sync::OnceLock<T>,
+    _py: Python<'_>,
+    f: F,
+) -> &'a T
+where
+    F: FnOnce() -> T,
+{
+    // SAFETY: we are currently attached to a Python thread
+    let ts_guard = Guard(unsafe { ffi::PyEval_SaveThread() });
+
+    // this trait is guarded by a rustc version config
+    // so clippy's MSRV check is wrong
+    #[allow(clippy::incompatible_msrv)]
+    // By having detached here, we guarantee that `.get_or_init` cannot deadlock with
+    // the Python interpreter
+    let value = lock.get_or_init(move || {
+        drop(ts_guard);
+        f()
+    });
+
+    value
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -589,4 +729,56 @@ mod tests {
             });
         });
     }
+
+    #[test]
+    #[cfg(not(target_arch = "wasm32"))] // We are building wasm Python with pthreads disabled
+    fn test_once_ext() {
+        // adapted from the example in the docs for Once::try_once_force
+        let init = Once::new();
+        std::thread::scope(|s| {
+            // poison the once
+            let handle = s.spawn(|| {
+                Python::with_gil(|py| {
+                    init.call_once_py_attached(py, || panic!());
+                })
+            });
+            assert!(handle.join().is_err());
+
+            // poisoning propagates
+            let handle = s.spawn(|| {
+                Python::with_gil(|py| {
+                    init.call_once_py_attached(py, || {});
+                });
+            });
+
+            assert!(handle.join().is_err());
+
+            // call_once_force will still run and reset the poisoned state
+            Python::with_gil(|py| {
+                init.call_once_force_py_attached(py, |state| {
+                    assert!(state.is_poisoned());
+                });
+
+                // once any success happens, we stop propagating the poison
+                init.call_once_py_attached(py, || {});
+            });
+        });
+    }
+
+    #[cfg(rustc_has_once_lock)]
+    #[cfg(not(target_arch = "wasm32"))] // We are building wasm Python with pthreads disabled
+    #[test]
+    fn test_once_lock_ext() {
+        let cell = std::sync::OnceLock::new();
+        std::thread::scope(|s| {
+            assert!(cell.get().is_none());
+
+            s.spawn(|| {
+                Python::with_gil(|py| {
+                    assert_eq!(*cell.get_or_init_py_attached(py, || 12345), 12345);
+                });
+            });
+        });
+        assert_eq!(cell.get(), Some(&12345));
+    }
 }

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]