feat: scoped_critical_section

Signed-off-by: Henry Schreiner <[email protected]>

Author: henryiii

docs/advanced/misc.rst

@@ -295,8 +295,10 @@ This module is sub-interpreter safe, for both ``shared_gil`` ("legacy") and
 function concurrently from different threads. This is safe because each sub-interpreter's GIL
 protects it's own Python objects from concurrent access.
 
-However, the module is no longer free-threading safe, for the same reason as before, because the
-calculation is not synchronized. We can synchronize it using a Python critical section.
+However, the module is no longer free-threading safe, for the same reason as
+before, because the calculation is not synchronized. We can synchronize it
+using a Python critical section. This will do nothing if not in free-threaded
+Python. You can have it lock one or two Python objects. You cannot nest it.
 
 .. code-block:: cpp
     :emphasize-lines: 1,5,10
@@ -305,12 +307,11 @@ calculation is not synchronized. We can synchronize it using a Python critical s
         m.def("calc_next", []() {
             size_t old;
             py::dict g = py::globals();
-            Py_BEGIN_CRITICAL_SECTION(g);
+            py::scoped_critical_section guard(g);
             if (!g.contains("myseed"))
                 g["myseed"] = 0;
             old = g["myseed"];
             g["myseed"] = (old + 1) * 10;
-            Py_END_CRITICAL_SECTION();
             return old;
         });
     }

include/pybind11/gil.h

@@ -199,3 +199,43 @@ class gil_scoped_release {
 PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)
 
 #endif // !PYBIND11_SIMPLE_GIL_MANAGEMENT
+
+PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
+
+class scoped_critical_section {
+public:
+    /// This does not do anything if there's a GIL. On free-threaded Python,
+    /// it locks an object. This uses the CPython API, which has limits
+    scoped_critical_section(handle obj) : single(true) {
+#ifdef Py_GIL_DISABLED
+        PyCriticalSection_Begin(&section, obj.ptr());
+#endif
+    }
+
+    scoped_critical_section(handle obj1, handle obj2) : single(true) {
+#ifdef Py_GIL_DISABLED
+        PyCriticalSection2_Begin(&section2, obj1.ptr(), obj2.ptr());
+#endif
+    }
+
+    ~scoped_critical_section() {
+#ifdef Py_GIL_DISABLED
+        if(single) {
+            PyCriticalSection_End(&section);
+        } else {
+            PyCriticalSection2_End(&section2);
+        }
+#endif
+    }
+
+private:
+    bool single;
+#ifdef Py_GIL_DISABLED
+    union {
+        PyCriticalSection section;
+        PyCriticalSection2 section2;
+    };
+#endif
+};
+
+PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)

tests/test_embed/test_interpreter.cpp

@@ -365,13 +365,9 @@ TEST_CASE("Threads") {
 #ifdef Py_GIL_DISABLED
 #    if PY_VERSION_HEX < 0x030E0000
                 std::lock_guard<std::mutex> lock(mutex);
-                locals["count"] = locals["count"].cast<int>() + 1;
 #    else
-                Py_BEGIN_CRITICAL_SECTION(locals.ptr());
-                locals["count"] = locals["count"].cast<int>() + 1;
-                Py_END_CRITICAL_SECTION();
+                py::scoped_critical_section lock(locals);
 #    endif
-#else
                 locals["count"] = locals["count"].cast<int>() + 1;
 #endif
             });