From d90dc420351083ca0df1e5ee122813d892ce54aa Mon Sep 17 00:00:00 2001
From: Bas Schoenmaeckers <b.schoenmaeckers@zanders.eu>
Date: Fri, 9 Aug 2024 16:23:22 +0200
Subject: [PATCH] Add PyCriticalSection lock to Dict iterator

---
 pyo3-ffi/src/cpython/criticalsection.rs | 10 ++++++
 pyo3-ffi/src/cpython/mod.rs             |  4 +++
 src/types/dict.rs                       | 48 ++++++++++++++++++++++++-
 3 files changed, 61 insertions(+), 1 deletion(-)
 create mode 100644 pyo3-ffi/src/cpython/criticalsection.rs

diff --git a/pyo3-ffi/src/cpython/criticalsection.rs b/pyo3-ffi/src/cpython/criticalsection.rs
new file mode 100644
index 00000000000..62b5f966fc6
--- /dev/null
+++ b/pyo3-ffi/src/cpython/criticalsection.rs
@@ -0,0 +1,10 @@
+use crate::PyObject;
+use std::os::raw::c_void;
+
+#[repr(C)]
+pub struct PyCriticalSection(*mut c_void, *mut c_void);
+
+extern "C" {
+    pub fn PyCriticalSection_Begin(m: *mut PyCriticalSection, op: *mut PyObject);
+    pub fn PyCriticalSection_End(m: *mut PyCriticalSection);
+}
diff --git a/pyo3-ffi/src/cpython/mod.rs b/pyo3-ffi/src/cpython/mod.rs
index 1710dbc4122..bd53de8b698 100644
--- a/pyo3-ffi/src/cpython/mod.rs
+++ b/pyo3-ffi/src/cpython/mod.rs
@@ -31,6 +31,8 @@ pub(crate) mod pymem;
 pub(crate) mod pystate;
 pub(crate) mod pythonrun;
 // skipped sysmodule.h
+#[cfg(Py_GIL_DISABLED)]
+pub(crate) mod criticalsection;
 pub(crate) mod floatobject;
 pub(crate) mod pyframe;
 pub(crate) mod tupleobject;
@@ -43,6 +45,8 @@ pub use self::bytesobject::*;
 pub use self::ceval::*;
 pub use self::code::*;
 pub use self::compile::*;
+#[cfg(Py_GIL_DISABLED)]
+pub use self::criticalsection::*;
 pub use self::descrobject::*;
 #[cfg(not(PyPy))]
 pub use self::dictobject::*;
diff --git a/src/types/dict.rs b/src/types/dict.rs
index 0fb6a711013..de3b63d203d 100644
--- a/src/types/dict.rs
+++ b/src/types/dict.rs
@@ -375,6 +375,8 @@ pub struct BoundDictIterator<'py> {
     ppos: ffi::Py_ssize_t,
     di_used: ffi::Py_ssize_t,
     len: ffi::Py_ssize_t,
+    #[cfg(Py_GIL_DISABLED)]
+    cs: ffi::PyCriticalSection,
 }
 
 impl<'py> Iterator for BoundDictIterator<'py> {
@@ -441,14 +443,33 @@ impl<'py> ExactSizeIterator for BoundDictIterator<'py> {
     }
 }
 
+#[cfg(Py_GIL_DISABLED)]
+impl Drop for BoundDictIterator<'_> {
+    fn drop(&mut self) {
+        unsafe {
+            ffi::PyCriticalSection_End(& mut self.cs);
+        }
+    }
+}
+
 impl<'py> BoundDictIterator<'py> {
     fn new(dict: Bound<'py, PyDict>) -> Self {
         let len = dict_len(&dict);
+
+        #[cfg(Py_GIL_DISABLED)]
+        let cs = unsafe {
+            let mut cs = std::mem::MaybeUninit::zeroed();
+            ffi::PyCriticalSection_Begin(cs.as_mut_ptr(), dict.as_ptr());
+            cs.assume_init()
+        };
+
         BoundDictIterator {
             dict,
             ppos: 0,
             di_used: len,
             len,
+            #[cfg(Py_GIL_DISABLED)]
+            cs,
         }
     }
 }
@@ -481,6 +502,8 @@ mod borrowed_iter {
         dict: Borrowed<'a, 'py, PyDict>,
         ppos: ffi::Py_ssize_t,
         len: ffi::Py_ssize_t,
+        #[cfg(Py_GIL_DISABLED)]
+        cs: ffi::PyCriticalSection,
     }
 
     impl<'a, 'py> Iterator for BorrowedDictIter<'a, 'py> {
@@ -519,10 +542,33 @@ mod borrowed_iter {
         }
     }
 
+    #[cfg(Py_GIL_DISABLED)]
+    impl Drop for BorrowedDictIter<'_, '_> {
+        fn drop(&mut self) {
+            unsafe {
+                ffi::PyCriticalSection_End(& mut self.cs);
+            }
+        }
+    }
+
     impl<'a, 'py> BorrowedDictIter<'a, 'py> {
         pub(super) fn new(dict: Borrowed<'a, 'py, PyDict>) -> Self {
             let len = dict_len(&dict);
-            BorrowedDictIter { dict, ppos: 0, len }
+
+            #[cfg(Py_GIL_DISABLED)]
+            let cs = unsafe {
+                let mut cs = std::mem::MaybeUninit::zeroed();
+                ffi::PyCriticalSection_Begin(cs.as_mut_ptr(), dict.as_ptr());
+                cs.assume_init()
+            };
+
+            BorrowedDictIter {
+                dict,
+                ppos: 0,
+                len,
+                #[cfg(Py_GIL_DISABLED)]
+                cs,
+            }
         }
     }
 }