Add a fast path to (single-mutex) critical section locking _iff_ the mutex

is already held by the currently active, top-most critical section of this
thread. This can matter a lot for indirectly recursive critical sections
without intervening critical sections.

Author: Yhg1s

Include/internal/pycore_critical_section.h

@@ -109,8 +109,18 @@ _PyCriticalSection_IsActive(uintptr_t tag)
 static inline void
 _PyCriticalSection_BeginMutex(PyCriticalSection *c, PyMutex *m)
 {
+    // As an optimisation for locking the same object recursively, skip
+    // locking if the mutex is currently locked and the topmost, active,
+    // single critical section.
+    PyThreadState *tstate = _PyThreadState_GET();
+    if (tstate->critical_section &&
+        !(tstate->critical_section & _Py_CRITICAL_SECTION_MASK) &&
+        ((PyCriticalSection *)tstate->critical_section)->_cs_mutex == m) {
+        c->_cs_mutex = NULL;
+        c->_cs_prev = 0;
+        return;
+    }
     if (PyMutex_LockFast(m)) {
-        PyThreadState *tstate = _PyThreadState_GET();
         c->_cs_mutex = m;
         c->_cs_prev = tstate->critical_section;
         tstate->critical_section = (uintptr_t)c;
@@ -145,6 +155,12 @@ _PyCriticalSection_Pop(PyCriticalSection *c)
 static inline void
 _PyCriticalSection_End(PyCriticalSection *c)
 {
+    // If the mutex is NULL, we used the fast path in
+    // _PyCriticalSection_BeginMutex for locks already held in the top-most
+    // critical section, and we shouldn't unlock or pop this critical section.
+    if (c->_cs_mutex == NULL) {
+        return;
+    }
     PyMutex_Unlock(c->_cs_mutex);
     _PyCriticalSection_Pop(c);
 }
@@ -199,6 +215,14 @@ _PyCriticalSection2_Begin(PyCriticalSection2 *c, PyObject *a, PyObject *b)
 static inline void
 _PyCriticalSection2_End(PyCriticalSection2 *c)
 {
+    // if mutex1 is NULL, we used the fast path in
+    // _PyCriticalSection_BeginMutex for mutexes that are already held,
+    // which should only happen when mutex1 and mutex2 were the same mutex,
+    // and mutex2 should also be NULL.
+    if (c->_cs_base._cs_mutex == NULL) {
+        assert(c->_cs_mutex2 == NULL);
+        return;
+    }
     if (c->_cs_mutex2) {
         PyMutex_Unlock(c->_cs_mutex2);
     }