add sync::OnceExt and sync::OnceLockExt traits

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`] and an extension trait [`OnceExt`] that
+enable functionality similar to these types but 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
 
 ## 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,58 @@ 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` trait into scope. This 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. These functions are analogous to `Once::call_once` and
+`Once::call_once_force` except they both accept a `Python<'py>` token in
+addition to an `FnOnce`. Both functions release the GIL and re-acquire it before
+executing the function, avoiding deadlocks with the GIL that are possible
+without using these functions. Here is an example of how to use this function to
+enable single-initialization of a runtime cache:
+
+```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` trait.

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,76 @@ 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));
+}
+
+struct Guard(Option<*mut crate::ffi::PyThreadState>);
+
+impl Drop for Guard {
+    fn drop(&mut self) {
+        if let Some(ts) = self.0 {
+            unsafe { ffi::PyEval_RestoreThread(ts) };
+        }
+    }
+}
+
+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);
+    }
+}
+
+#[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 mut ts = Guard(Some(unsafe { ffi::PyEval_SaveThread() }));
+
+    once.call_once(|| {
+        unsafe { ffi::PyEval_RestoreThread(ts.0.take().unwrap()) };
+        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 mut ts = Guard(Some(unsafe { ffi::PyEval_SaveThread() }));
+
+    once.call_once_force(|state| {
+        unsafe { ffi::PyEval_RestoreThread(ts.0.take().unwrap()) };
+        f(state);
+    });
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -589,4 +666,38 @@ mod tests {
             });
         });
     }
+
+    #[test]
+    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, || {});
+            });
+        });
+    }
 }

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]