From ad4f909e0e7890e027c4ae7fea74586667242ad3 Mon Sep 17 00:00:00 2001
From: Sam Gross <colesbury@gmail.com>
Date: Thu, 15 Feb 2024 08:37:54 -0500
Subject: [PATCH] gh-115432: Add critical section variant that handles a NULL
 object (#115433)

This adds `Py_XBEGIN_CRITICAL_SECTION` and
`Py_XEND_CRITICAL_SECTION`, which accept a possibly NULL object as an
argument. If the argument is NULL, then nothing is locked or unlocked.
Otherwise, they behave like `Py_BEGIN/END_CRITICAL_SECTION`.
---
 Include/internal/pycore_critical_section.h    | 29 +++++++++++++++++++
 .../test_critical_sections.c                  |  9 ++++++
 2 files changed, 38 insertions(+)

diff --git a/Include/internal/pycore_critical_section.h b/Include/internal/pycore_critical_section.h
index 38ed8cd69804ba..30820a24f5bb64 100644
--- a/Include/internal/pycore_critical_section.h
+++ b/Include/internal/pycore_critical_section.h
@@ -96,6 +96,15 @@ extern "C" {
         _PyCriticalSection_End(&_cs);                                   \
     }
 
+# define Py_XBEGIN_CRITICAL_SECTION(op)                                 \
+    {                                                                   \
+        _PyCriticalSection _cs_opt = {0};                               \
+        _PyCriticalSection_XBegin(&_cs_opt, _PyObject_CAST(op))
+
+# define Py_XEND_CRITICAL_SECTION()                                     \
+        _PyCriticalSection_XEnd(&_cs_opt);                              \
+    }
+
 # define Py_BEGIN_CRITICAL_SECTION2(a, b)                               \
     {                                                                   \
         _PyCriticalSection2 _cs2;                                       \
@@ -131,6 +140,8 @@ extern "C" {
 // The critical section APIs are no-ops with the GIL.
 # define Py_BEGIN_CRITICAL_SECTION(op)
 # define Py_END_CRITICAL_SECTION()
+# define Py_XBEGIN_CRITICAL_SECTION(op)
+# define Py_XEND_CRITICAL_SECTION()
 # define Py_BEGIN_CRITICAL_SECTION2(a, b)
 # define Py_END_CRITICAL_SECTION2()
 # define _Py_CRITICAL_SECTION_ASSERT_MUTEX_LOCKED(mutex)
@@ -187,6 +198,16 @@ _PyCriticalSection_Begin(_PyCriticalSection *c, PyMutex *m)
     }
 }
 
+static inline void
+_PyCriticalSection_XBegin(_PyCriticalSection *c, PyObject *op)
+{
+#ifdef Py_GIL_DISABLED
+    if (op != NULL) {
+        _PyCriticalSection_Begin(c, &_PyObject_CAST(op)->ob_mutex);
+    }
+#endif
+}
+
 // Removes the top-most critical section from the thread's stack of critical
 // sections. If the new top-most critical section is inactive, then it is
 // resumed.
@@ -209,6 +230,14 @@ _PyCriticalSection_End(_PyCriticalSection *c)
     _PyCriticalSection_Pop(c);
 }
 
+static inline void
+_PyCriticalSection_XEnd(_PyCriticalSection *c)
+{
+    if (c->mutex) {
+        _PyCriticalSection_End(c);
+    }
+}
+
 static inline void
 _PyCriticalSection2_Begin(_PyCriticalSection2 *c, PyMutex *m1, PyMutex *m2)
 {
diff --git a/Modules/_testinternalcapi/test_critical_sections.c b/Modules/_testinternalcapi/test_critical_sections.c
index 1f7e311558b27c..94da0468fcf149 100644
--- a/Modules/_testinternalcapi/test_critical_sections.c
+++ b/Modules/_testinternalcapi/test_critical_sections.c
@@ -49,6 +49,15 @@ test_critical_sections(PyObject *self, PyObject *Py_UNUSED(args))
     Py_END_CRITICAL_SECTION2();
     assert_nogil(!PyMutex_IsLocked(&d2->ob_mutex));
 
+    // Optional variant behaves the same if the object is non-NULL
+    Py_XBEGIN_CRITICAL_SECTION(d1);
+    assert_nogil(PyMutex_IsLocked(&d1->ob_mutex));
+    Py_XEND_CRITICAL_SECTION();
+
+    // No-op
+    Py_XBEGIN_CRITICAL_SECTION(NULL);
+    Py_XEND_CRITICAL_SECTION();
+
     Py_DECREF(d2);
     Py_DECREF(d1);
     Py_RETURN_NONE;