Implement OnceExt & MutexExt for parking_lot & lock_api (#5044)

* Implement `OnceExt` & `MutexExt` for `parking_lot` & `lock_api`

* Implement `MutexExt` for `Arc<lock_api::Mutex>`

* Move `MutexExt` lock result into type parameter

* Remove lifetime from associated type of `OnceExt`

* Add docs

* add changelog

* Add new features to `full`

* mention `parking_lot` in guide

* Add tests

* Update docs to mention all cargo features

Co-authored-by: Nathan Goldbaum <nathan.goldbaum@gmail.com>

* Also test `parking_lot::Mutex::lock_py_attached` on `lock_api` feature

* Update newsfragments/5044.added.md

Co-authored-by: Nathan Goldbaum <nathan.goldbaum@gmail.com>

* Add additional case to `once_ext` test

* enable `arc_lock` feature on `parking_lot` during testing

---------

Co-authored-by: Nathan Goldbaum <nathan.goldbaum@gmail.com>

Author: bschoenmaeckers

Cargo.toml

@@ -49,6 +49,8 @@ time = { version = "0.3.38", default-features = false, optional = true }
 serde = { version = "1.0", optional = true }
 smallvec = { version = "1.0", optional = true }
 uuid = { version = "1.11.0", optional = true  }
+lock_api = { version = "0.4", optional = true }
+parking_lot = { version = "0.12", optional = true}
 
 [target.'cfg(not(target_has_atomic = "64"))'.dependencies]
 portable-atomic = "1.0"
@@ -68,6 +70,7 @@ futures = "0.3.28"
 tempfile = "3.12.0"
 static_assertions = "1.1.0"
 uuid = { version = "1.10.0", features = ["v4"] }
+parking_lot = { version = "0.12.3", features = ["arc_lock"]}
 
 [build-dependencies]
 pyo3-build-config = { path = "pyo3-build-config", version = "=0.25.0-dev", features = ["resolve-config"] }
@@ -115,6 +118,9 @@ auto-initialize = []
 # Enables `Clone`ing references to Python objects `Py<T>` which panics if the GIL is not held.
 py-clone = []
 
+parking_lot = ["dep:parking_lot", "lock_api"]
+arc_lock = ["lock_api", "lock_api/arc_lock", "parking_lot?/arc_lock"]
+
 # Optimizes PyObject to Vec conversion and so on.
 nightly = []
 
@@ -124,6 +130,7 @@ full = [
     "macros",
     # "multiple-pymethods", # Not supported by wasm
     "anyhow",
+    "arc_lock",
     "bigdecimal",
     "chrono",
     "chrono-tz",
@@ -133,10 +140,12 @@ full = [
     "eyre",
     "hashbrown",
     "indexmap",
+    "lock_api",
     "num-bigint",
     "num-complex",
     "num-rational",
     "ordered-float",
+    "parking_lot",
     "py-clone",
     "rust_decimal",
     "serde",

guide/src/class/thread-safety.md

@@ -101,7 +101,7 @@ impl MyClass {
 }
 ```
 
-If you need to lock around state stored in the Python interpreter or otherwise call into the Python C API while a lock is held, you might find the `MutexExt` trait useful. It provides a `lock_py_attached` method for `std::sync::Mutex` that avoids deadlocks with the GIL or other global synchronization events in the interpreter.
+If you need to lock around state stored in the Python interpreter or otherwise call into the Python C API while a lock is held, you might find the `MutexExt` trait useful. It provides a `lock_py_attached` method for `std::sync::Mutex` that avoids deadlocks with the GIL or other global synchronization events in the interpreter. Additionally, support for the `parking_lot` and `lock_api` synchronization libraries is gated behind the `parking_lot` and `lock_api` features. You can also enable the `arc_lock` feature if you need the `arc_lock` features of either library.
 
 ### Wrapping unsynchronized data
 

guide/src/features.md

@@ -121,6 +121,10 @@ These features enable conversions between Python types and types from other Rust
 
 Adds a dependency on [anyhow](https://docs.rs/anyhow). Enables a conversion from [anyhow](https://docs.rs/anyhow)’s [`Error`](https://docs.rs/anyhow/latest/anyhow/struct.Error.html) type to [`PyErr`]({{#PYO3_DOCS_URL}}/pyo3/struct.PyErr.html), for easy error handling.
 
+### `arc_lock`
+
+Enables Pyo3's `MutexExt` trait for all Mutexes that extend on [`lock_api::Mutex`](https://docs.rs/lock_api/latest/lock_api/struct.Mutex.html) and are wrapped in an [`Arc`](https://doc.rust-lang.org/std/sync/struct.Arc.html) type. Like [`Arc<parking_lot::Mutex>`](https://docs.rs/parking_lot/latest/parking_lot/type.Mutex.html#method.lock_arc)
+
 ### `bigdecimal`
 Adds a dependency on [bigdecimal](https://docs.rs/bigdecimal) and enables conversions into its [`BigDecimal`](https://docs.rs/bigdecimal/latest/bigdecimal/struct.BigDecimal.html) type.
 
@@ -168,6 +172,10 @@ Adds a dependency on [jiff@0.2](https://docs.rs/jiff/0.2) and requires MSRV 1.70
 - [Zoned](https://docs.rs/jiff/0.2/jiff/struct.Zoned.html) -> [`PyDateTime`]({{#PYO3_DOCS_URL}}/pyo3/types/struct.PyDateTime.html)
 - [Timestamp](https://docs.rs/jiff/0.2/jiff/struct.Timestamp.html) -> [`PyDateTime`]({{#PYO3_DOCS_URL}}/pyo3/types/struct.PyDateTime.html)
 
+### `lock_api`
+
+Adds a dependency on [lock_api](https://docs.rs/lock_api) and enables Pyo3's `MutexExt` trait for all mutexes that extend on [`lock_api::Mutex`](https://docs.rs/lock_api/latest/lock_api/struct.Mutex.html) (like `parking_lot` or `spin`).
+
 ### `num-bigint`
 
 Adds a dependency on [num-bigint](https://docs.rs/num-bigint) and enables conversions into its [`BigInt`](https://docs.rs/num-bigint/latest/num_bigint/struct.BigInt.html) and [`BigUint`](https://docs.rs/num-bigint/latest/num_bigint/struct.BigUint.html) types.
@@ -186,6 +194,10 @@ Adds a dependency on [ordered-float](https://docs.rs/ordered-float) and enables
 - [NotNan](https://docs.rs/ordered-float/latest/ordered_float/struct.NotNan.html) -> [`PyFloat`]({{#PYO3_DOCS_URL}}/pyo3/types/struct.PyFloat.html)
 - [OrderedFloat](https://docs.rs/ordered-float/latest/ordered_float/struct.OrderedFloat.html) -> [`PyFloat`]({{#PYO3_DOCS_URL}}/pyo3/types/struct.PyFloat.html)
 
+### `parking-lot`
+
+Adds a dependency on [parking_lot](https://docs.rs/parking_lot) and enables Pyo3's `OnceExt` & `MutexExt` traits for [`parking_lot::Once`](https://docs.rs/parking_lot/latest/parking_lot/struct.Once.html) and [`parking_lot::Mutex`](https://docs.rs/parking_lot/latest/parking_lot/type.Mutex.html) types.
+
 ### `rust_decimal`
 
 Adds a dependency on [rust_decimal](https://docs.rs/rust_decimal) and enables conversions into its [`Decimal`](https://docs.rs/rust_decimal/latest/rust_decimal/struct.Decimal.html) type.

newsfragments/5044.added.md

@@ -0,0 +1 @@
+Implement `OnceExt` & `MutexExt` for `parking_lot` & `lock_api`. Use the new extension traits by enabling the `arc_lock`, `lock_api`, or `parking_lot` cargo features.

src/sealed.rs

@@ -57,3 +57,9 @@ impl<T: crate::pyclass::PyClass> Sealed for PyClassInitializer<T> {}
 
 impl Sealed for std::sync::Once {}
 impl<T> Sealed for std::sync::Mutex<T> {}
+#[cfg(feature = "lock_api")]
+impl<R, T> Sealed for lock_api::Mutex<R, T> {}
+#[cfg(feature = "parking_lot")]
+impl Sealed for parking_lot::Once {}
+#[cfg(feature = "arc_lock")]
+impl<R: lock_api::RawMutex, T> Sealed for std::sync::Arc<lock_api::Mutex<R, T>> {}

src/sync.rs

@@ -536,13 +536,16 @@ mod once_lock_ext_sealed {
 /// Helper trait for `Once` to help avoid deadlocking when using a `Once` when attached to a
 /// Python thread.
 pub trait OnceExt: Sealed {
+    ///The state of `Once`
+    type OnceState;
+
     /// 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));
+    fn call_once_force_py_attached(&self, py: Python<'_>, f: impl FnOnce(&Self::OnceState));
 }
 
 /// Extension trait for [`std::sync::OnceLock`] which helps avoid deadlocks between the Python
@@ -565,21 +568,23 @@ pub trait OnceLockExt<T>: once_lock_ext_sealed::Sealed {
 
 /// Extension trait for [`std::sync::Mutex`] which helps avoid deadlocks between
 /// the Python interpreter and acquiring the `Mutex`.
-pub trait MutexExt<T>: Sealed {
+pub trait MutexExt<'a, T, R>: Sealed
+where
+    Self: 'a,
+{
     /// Lock this `Mutex` in a manner that cannot deadlock with the Python interpreter.
     ///
     /// Before attempting to lock the mutex, this function detaches from the
     /// Python runtime. When the lock is acquired, it re-attaches to the Python
     /// runtime before returning the `LockResult`. This avoids deadlocks between
     /// the GIL and other global synchronization events triggered by the Python
     /// interpreter.
-    fn lock_py_attached(
-        &self,
-        py: Python<'_>,
-    ) -> std::sync::LockResult<std::sync::MutexGuard<'_, T>>;
+    fn lock_py_attached(&'a self, py: Python<'_>) -> R;
 }
 
 impl OnceExt for Once {
+    type OnceState = OnceState;
+
     fn call_once_py_attached(&self, py: Python<'_>, f: impl FnOnce()) {
         if self.is_completed() {
             return;
@@ -597,6 +602,41 @@ impl OnceExt for Once {
     }
 }
 
+#[cfg(feature = "parking_lot")]
+impl OnceExt for parking_lot::Once {
+    type OnceState = parking_lot::OnceState;
+
+    fn call_once_py_attached(&self, _py: Python<'_>, f: impl FnOnce()) {
+        if self.state().done() {
+            return;
+        }
+
+        let ts_guard = unsafe { SuspendGIL::new() };
+
+        self.call_once(move || {
+            drop(ts_guard);
+            f();
+        });
+    }
+
+    fn call_once_force_py_attached(
+        &self,
+        _py: Python<'_>,
+        f: impl FnOnce(&parking_lot::OnceState),
+    ) {
+        if self.state().done() {
+            return;
+        }
+
+        let ts_guard = unsafe { SuspendGIL::new() };
+
+        self.call_once_force(move |state| {
+            drop(ts_guard);
+            f(&state);
+        });
+    }
+}
+
 #[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
@@ -612,11 +652,13 @@ impl<T> OnceLockExt<T> for std::sync::OnceLock<T> {
     }
 }
 
-impl<T> MutexExt<T> for std::sync::Mutex<T> {
+impl<'a, T> MutexExt<'a, T, std::sync::LockResult<std::sync::MutexGuard<'a, T>>>
+    for std::sync::Mutex<T>
+{
     fn lock_py_attached(
-        &self,
+        &'a self,
         _py: Python<'_>,
-    ) -> std::sync::LockResult<std::sync::MutexGuard<'_, T>> {
+    ) -> std::sync::LockResult<std::sync::MutexGuard<'a, T>> {
         // If try_lock is successful or returns a poisoned mutex, return them so
         // the caller can deal with them. Otherwise we need to use blocking
         // lock, which requires detaching from the Python runtime to avoid
@@ -638,6 +680,41 @@ impl<T> MutexExt<T> for std::sync::Mutex<T> {
     }
 }
 
+#[cfg(feature = "lock_api")]
+impl<'a, R: lock_api::RawMutex, T> MutexExt<'a, T, lock_api::MutexGuard<'a, R, T>>
+    for lock_api::Mutex<R, T>
+{
+    fn lock_py_attached(&'a self, _py: Python<'_>) -> lock_api::MutexGuard<'a, R, T> {
+        if let Some(guard) = self.try_lock() {
+            return guard;
+        }
+
+        let ts_guard = unsafe { SuspendGIL::new() };
+        let res = self.lock();
+        drop(ts_guard);
+        res
+    }
+}
+
+#[cfg(feature = "arc_lock")]
+impl<'a, R, T> MutexExt<'a, T, lock_api::ArcMutexGuard<R, T>>
+    for std::sync::Arc<lock_api::Mutex<R, T>>
+where
+    R: lock_api::RawMutex,
+    Self: 'a,
+{
+    fn lock_py_attached(&self, _py: Python<'_>) -> lock_api::ArcMutexGuard<R, T> {
+        if let Some(guard) = self.try_lock_arc() {
+            return guard;
+        }
+
+        let ts_guard = unsafe { SuspendGIL::new() };
+        let res = self.lock_arc();
+        drop(ts_guard);
+        res
+    }
+}
+
 #[cold]
 fn init_once_py_attached<F, T>(once: &Once, _py: Python<'_>, f: F)
 where
@@ -966,36 +1043,47 @@ 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());
+        macro_rules! test_once {
+            ($once:expr, $is_poisoned:expr) => {{
+                // adapted from the example in the docs for Once::try_once_force
+                let init = $once;
+                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, || {});
-                });
-            });
+                    // poisoning propagates
+                    let handle = s.spawn(|| {
+                        Python::with_gil(|py| {
+                            init.call_once_py_attached(py, || {});
+                        });
+                    });
+
+                    assert!(handle.join().is_err());
 
-            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!($is_poisoned(state.clone()));
+                        });
 
-            // 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, || {});
+                    });
+
+                    // calling call_once_force should return immediately without calling the closure
+                    Python::with_gil(|py| init.call_once_force_py_attached(py, |_| panic!()));
                 });
+            }};
+        }
 
-                // once any success happens, we stop propagating the poison
-                init.call_once_py_attached(py, || {});
-            });
-        });
+        test_once!(Once::new(), OnceState::is_poisoned);
+        #[cfg(feature = "parking_lot")]
+        test_once!(parking_lot::Once::new(), parking_lot::OnceState::poisoned);
     }
 
     #[cfg(rustc_has_once_lock)]
@@ -1047,6 +1135,54 @@ mod tests {
         });
     }
 
+    #[cfg(feature = "macros")]
+    #[cfg(all(
+        any(feature = "parking_lot", feature = "lock_api"),
+        not(target_arch = "wasm32") // We are building wasm Python with pthreads disabled
+    ))]
+    #[test]
+    fn test_parking_lot_mutex_ext() {
+        macro_rules! test_mutex {
+            ($guard:ty ,$mutex:stmt) => {{
+                let barrier = Barrier::new(2);
+
+                let mutex = Python::with_gil({ $mutex });
+
+                std::thread::scope(|s| {
+                    s.spawn(|| {
+                        Python::with_gil(|py| {
+                            let b: $guard = mutex.lock_py_attached(py);
+                            barrier.wait();
+                            // sleep to ensure the other thread actually blocks
+                            std::thread::sleep(std::time::Duration::from_millis(10));
+                            (*b).bind(py).borrow().0.store(true, Ordering::Release);
+                            drop(b);
+                        });
+                    });
+                    s.spawn(|| {
+                        barrier.wait();
+                        Python::with_gil(|py| {
+                            // blocks until the other thread releases the lock
+                            let b: $guard = mutex.lock_py_attached(py);
+                            assert!((*b).bind(py).borrow().0.load(Ordering::Acquire));
+                        });
+                    });
+                });
+            }};
+        }
+
+        test_mutex!(parking_lot::MutexGuard<'_, _>, |py| {
+            parking_lot::Mutex::new(Py::new(py, BoolWrapper(AtomicBool::new(false))).unwrap())
+        });
+
+        #[cfg(feature = "arc_lock")]
+        test_mutex!(parking_lot::ArcMutexGuard<_, _>, |py| {
+            let mutex =
+                parking_lot::Mutex::new(Py::new(py, BoolWrapper(AtomicBool::new(false))).unwrap());
+            std::sync::Arc::new(mutex)
+        });
+    }
+
     #[cfg(not(target_arch = "wasm32"))] // We are building wasm Python with pthreads disabled
     #[test]
     fn test_mutex_ext_poison() {