Add critical section API wrappers

newsfragments/4587.added.md

@@ -0,0 +1,2 @@
+* Added `with_critical_section`, a safe wrapper around the Python Critical
+  Section API added in Python 3.13 for the free-threaded build.

pyo3-ffi/src/cpython/critical_section.rs

@@ -9,13 +9,47 @@ pub struct PyCriticalSection {
     _cs_mutex: *mut PyMutex,
 }
 
+#[cfg(Py_GIL_DISABLED)]
+impl PyCriticalSection {
+    pub const fn new() -> PyCriticalSection {
+        PyCriticalSection {
+            _cs_prev: 0,
+            _cs_mutex: std::ptr::null_mut(),
+        }
+    }
+}
+
+#[cfg(Py_GIL_DISABLED)]
+impl Default for PyCriticalSection {
+    fn default() -> Self {
+        PyCriticalSection::new()
+    }
+}
+
 #[repr(C)]
 #[cfg(Py_GIL_DISABLED)]
 pub struct PyCriticalSection2 {
     _cs_base: PyCriticalSection,
     _cs_mutex2: *mut PyMutex,
 }
 
+#[cfg(Py_GIL_DISABLED)]
+impl PyCriticalSection2 {
+    pub const fn new() -> PyCriticalSection2 {
+        PyCriticalSection2 {
+            _cs_base: PyCriticalSection::new(),
+            _cs_mutex2: std::ptr::null_mut(),
+        }
+    }
+}
+
+#[cfg(Py_GIL_DISABLED)]
+impl Default for PyCriticalSection2 {
+    fn default() -> Self {
+        PyCriticalSection2::new()
+    }
+}
+
 #[cfg(not(Py_GIL_DISABLED))]
 opaque_struct!(PyCriticalSection);
 

src/sync.rs

@@ -5,7 +5,7 @@
 //!
 //! [PEP 703]: https://peps.python.org/pep-703/
 use crate::{
-    types::{any::PyAnyMethods, PyString},
+    types::{any::PyAnyMethods, PyAny, PyString},
     Bound, Py, PyResult, PyTypeCheck, Python,
 };
 use std::cell::UnsafeCell;
@@ -330,11 +330,67 @@ impl Interned {
     }
 }
 
+/// Executes a closure with a Python critical section held on an object.
+///
+/// Acquires the per-object lock for the object `op` that is held
+/// until the closure `f` is finished.
+///
+/// This is structurally equivalent to the use of the paired
+/// Py_BEGIN_CRITICAL_SECTION and Py_END_CRITICAL_SECTION macros.
+///
+/// A no-op on GIL-enabled builds, where the critical section API is exposed as
+/// a no-op by the Python C API.
+///
+/// Provides weaker locking guarantees than traditional locks, but can in some
+/// cases be used to provide guarantees similar to the GIL without the risk of
+/// deadlocks associated with traditional locks.
+///
+/// Many CPython C API functions do not acquire the per-object lock on objects
+/// passed to Python. You should not expect critical sections applied to
+/// built-in types to prevent concurrent modification. This API is most useful
+/// for user-defined types with full control over how the internal state for the
+/// type is managed.
+#[cfg_attr(not(Py_GIL_DISABLED), allow(unused_variables))]
+pub fn with_critical_section<F, R>(object: &Bound<'_, PyAny>, f: F) -> R
+where
+    F: FnOnce() -> R,
+{
+    #[cfg(Py_GIL_DISABLED)]
+    {
+        struct Guard(crate::ffi::PyCriticalSection);
+
+        impl Drop for Guard {
+            fn drop(&mut self) {
+                unsafe {
+                    crate::ffi::PyCriticalSection_End(&mut self.0);
+                }
+            }
+        }
+
+        let _guard = unsafe {
+            let mut section = std::mem::zeroed();
+            crate::ffi::PyCriticalSection_Begin(&mut section, object.as_ptr());
+            Guard(section)
+        };
+        f()
+    }
+    #[cfg(not(Py_GIL_DISABLED))]
+    {
+        f()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
-    use crate::types::{dict::PyDictMethods, PyDict};
+    use crate::types::{PyDict, PyDictMethods};
+
+    #[cfg(feature = "macros")]
+    use std::sync::{
+        atomic::{AtomicBool, Ordering},
+        Barrier,
+    };
 
     #[test]
     fn test_intern() {
@@ -381,4 +437,40 @@ mod tests {
             assert!(cell_py.clone_ref(py).get(py).unwrap().is_none(py));
         })
     }
+
+    #[cfg(feature = "macros")]
+    #[test]
+    fn test_critical_section() {
+        let barrier = Barrier::new(2);
+
+        #[crate::pyclass(crate = "crate")]
+        struct BoolWrapper(AtomicBool);
+
+        let bool_wrapper = Python::with_gil(|py| -> Py<BoolWrapper> {
+            Py::new(py, BoolWrapper(AtomicBool::new(false))).unwrap()
+        });
+
+        std::thread::scope(|s| {
+            s.spawn(|| {
+                Python::with_gil(|py| {
+                    let b = bool_wrapper.bind(py);
+                    with_critical_section(b, || {
+                        barrier.wait();
+                        std::thread::sleep(std::time::Duration::from_millis(10));
+                        b.borrow().0.store(true, Ordering::Release);
+                    })
+                });
+            });
+            s.spawn(|| {
+                barrier.wait();
+                Python::with_gil(|py| {
+                    let b = bool_wrapper.bind(py);
+                    // this blocks until the other thread's critical section finishes
+                    with_critical_section(b, || {
+                        assert!(b.borrow().0.load(Ordering::Acquire));
+                    });
+                });
+            });
+        });
+    }
 }