Absorb to critical section

DinoV · Jan 19, 2024 · c731ed6 · c731ed6
1 parent e38e8c8
commit c731ed6
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/Include/internal/pycore_critical_section.h b/Include/internal/pycore_critical_section.h
@@ -260,8 +260,15 @@ _PyCriticalSection_AssertHeld(PyMutex *mutex)
 {
 #ifdef Py_DEBUG
     PyThreadState *tstate = _PyThreadState_GET();
-    _PyCriticalSection *cs = (_PyCriticalSection *)tstate->critical_section;
-    assert(cs != NULL && cs->mutex == mutex);
+    uintptr_t prev = tstate->critical_section;
+    if (prev & _Py_CRITICAL_SECTION_TWO_MUTEXES) {
+        _PyCriticalSection2 *cs = (_PyCriticalSection2 *)(prev & ~_Py_CRITICAL_SECTION_MASK);
+        assert(cs != NULL && (cs->base.mutex == mutex || cs->mutex2 == mutex));
+    } else {
+        _PyCriticalSection *cs = (_PyCriticalSection *)(tstate->critical_section & ~_Py_CRITICAL_SECTION_MASK);
+        assert(cs != NULL && cs->mutex == mutex);
+    }
+
 #endif
 }