Add OnceLockExt

guide/src/faq.md

@@ -13,8 +13,7 @@ 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`] 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.
+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 ot 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

guide/src/free-threading.md

@@ -168,15 +168,17 @@ 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:
+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() {

newsfragments/4676.added.md

@@ -1 +1 @@
-Add `pyo3::sync::OnceExt` trait.
+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/sync.rs

@@ -480,6 +480,7 @@ where
     }
 }
 
+#[cfg(rustc_has_once_lock)]
 mod once_lock_ext_sealed {
     pub trait Sealed {}
     impl<T> Sealed for std::sync::OnceLock<T> {}
@@ -499,6 +500,7 @@ pub trait OnceExt: Sealed {
 
 // 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.
     ///
@@ -542,6 +544,7 @@ impl OnceExt for Once {
     }
 }
 
+#[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
@@ -560,10 +563,10 @@ where
 {
     // 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() }));
+    let ts_guard = Guard(Some(unsafe { ffi::PyEval_SaveThread() }));
 
     once.call_once(|| {
-        unsafe { ffi::PyEval_RestoreThread(ts.0.take().unwrap()) };
+        drop(ts_guard);
         f();
     });
 }
@@ -575,14 +578,15 @@ where
 {
     // 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() }));
+    let ts_guard = Guard(Some(unsafe { ffi::PyEval_SaveThread() }));
 
     once.call_once_force(|state| {
-        unsafe { ffi::PyEval_RestoreThread(ts.0.take().unwrap()) };
+        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>,
@@ -593,14 +597,12 @@ where
     F: FnOnce() -> T,
 {
     // SAFETY: we are currently attached to a Python thread
-    let mut ts = Guard(Some(unsafe { ffi::PyEval_SaveThread() }));
+    let ts_guard = Guard(Some(unsafe { ffi::PyEval_SaveThread() }));
 
     // By having detached here, we guarantee that `.get_or_init` cannot deadlock with
     // the Python interpreter
     let value = lock.get_or_init(move || {
-        let ts = ts.0.take().expect("ts is set to Some above");
-        // SAFETY: ts is a valid thread state and needs restoring
-        unsafe { ffi::PyEval_RestoreThread(ts) };
+        drop(ts_guard);
         f()
     });
 
@@ -757,4 +759,20 @@ mod tests {
             });
         });
     }
+
+    #[cfg(rustc_has_once_lock)]
+    #[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));
+    }
 }