critical_section: helpers for fine-grained locking

Critical sections are helpers to replace the global interpreter lock with finer grained locking. They provide similar guarantees to the GIL and avoid the deadlock risk that plain locking involves. Critical sections are implicitly ended whenever the GIL would be released. They are resumed when the GIL would be acquired. Nested critical sections behave as-if they're interleaved.
colesbury · Apr 23, 2023 · 6845b13 · 6845b13
1 parent 4584be5
commit 6845b13
Show file tree

Hide file tree

Showing 9 changed files with 160 additions and 0 deletions.
diff --git a/Include/cpython/pystate.h b/Include/cpython/pystate.h
@@ -172,6 +172,8 @@ struct _ts {
     int trash_delete_nesting;
     PyObject *trash_delete_later;
 
+    uintptr_t critical_section;
+
     /* Called when a thread state is deleted normally, but not when it
      * is destroyed after fork().
      * Pain:  to prevent rare but fatal shutdown errors (issue 18808),

diff --git a/Makefile.pre.in b/Makefile.pre.in
@@ -379,6 +379,7 @@ PYTHON_OBJS=	\
 		Python/codecs.o \
 		Python/compile.o \
 		Python/context.o \
+		Python/critical_section.o \
 		Python/dynamic_annotations.o \
 		Python/errors.o \
 		Python/frame.o \
@@ -1643,6 +1644,7 @@ PYTHON_HEADERS= \
 		$(srcdir)/Include/internal/pycore_compile.h \
 		$(srcdir)/Include/internal/pycore_condvar.h \
 		$(srcdir)/Include/internal/pycore_context.h \
+		$(srcdir)/Include/internal/pycore_critical_section.h \
 		$(srcdir)/Include/internal/pycore_dict.h \
 		$(srcdir)/Include/internal/pycore_dict_state.h \
 		$(srcdir)/Include/internal/pycore_descrobject.h \

diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
@@ -15,6 +15,7 @@
 #include "pycore_atomic_funcs.h" // _Py_atomic_int_get()
 #include "pycore_bitutils.h"     // _Py_bswap32()
 #include "pycore_compile.h"      // _PyCompile_CodeGen, _PyCompile_OptimizeCfg
+#include "pycore_critical_section.h"
 #include "pycore_fileutils.h"    // _Py_normpath
 #include "pycore_frame.h"        // _PyInterpreterFrame
 #include "pycore_gc.h"           // PyGC_Head
@@ -244,6 +245,34 @@ test_hashtable(PyObject *self, PyObject *Py_UNUSED(args))
     Py_RETURN_NONE;
 }
 
+static PyObject *
+test_critical_sections(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    PyThreadState *tstate = PyThreadState_GET();
+    _PyMutex m1, m2;
+    memset(&m1, 0, sizeof(m1));
+    memset(&m2, 0, sizeof(m2));
+
+    struct _Py_critical_section c;
+    _Py_critical_section_begin(&c, &m1);
+    assert(_PyMutex_is_locked(&m1));
+
+    /* nested critical section re-using lock */
+    struct _Py_critical_section c2;
+    _Py_critical_section_begin(&c2, &m1);
+    assert(_PyMutex_is_locked(&m1));
+    assert(_Py_critical_section_is_active(tstate->critical_section));
+    assert(!_Py_critical_section_is_active(c2.prev));
+    _Py_critical_section_end(&c2);
+
+    /* mutex is re-locked */
+    assert(_PyMutex_is_locked(&m1));
+
+    _Py_critical_section_end(&c);
+    assert(!_PyMutex_is_locked(&m1));
+
+    Py_RETURN_NONE;
+}
 
 static PyObject *
 test_get_config(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(args))
@@ -634,6 +663,7 @@ static PyMethodDef TestMethods[] = {
     _TESTINTERNALCAPI_COMPILER_CODEGEN_METHODDEF
     _TESTINTERNALCAPI_OPTIMIZE_CFG_METHODDEF
     {"get_interp_settings", get_interp_settings, METH_VARARGS, NULL},
+    {"test_critical_sections", test_critical_sections, METH_NOARGS},
     {NULL, NULL} /* sentinel */
 };
 

diff --git a/PCbuild/_freeze_module.vcxproj b/PCbuild/_freeze_module.vcxproj
@@ -186,6 +186,7 @@
     <ClCompile Include="..\Python\codecs.c" />
     <ClCompile Include="..\Python\compile.c" />
     <ClCompile Include="..\Python\context.c" />
+    <ClCompile Include="..\Python\critical_section.c" />
     <ClCompile Include="..\Python\dtoa.c" />
     <ClCompile Include="..\Python\dynamic_annotations.c" />
     <ClCompile Include="..\Python\dynload_win.c" />

diff --git a/PCbuild/_freeze_module.vcxproj.filters b/PCbuild/_freeze_module.vcxproj.filters
@@ -100,6 +100,9 @@
     <ClCompile Include="..\Python\context.c">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="..\Python\critical_section.c">
+      <Filter>Source Files</Filter>
+    </ClCompile>
     <ClCompile Include="..\Objects\descrobject.c">
       <Filter>Source Files</Filter>
     </ClCompile>

diff --git a/PCbuild/pythoncore.vcxproj b/PCbuild/pythoncore.vcxproj
@@ -210,6 +210,7 @@
     <ClInclude Include="..\Include\internal\pycore_compile.h" />
     <ClInclude Include="..\Include\internal\pycore_condvar.h" />
     <ClInclude Include="..\Include\internal\pycore_context.h" />
+    <ClInclude Include="..\Include\internal\pycore_critical_section.h" />
     <ClInclude Include="..\Include\internal\pycore_descrobject.h" />
     <ClInclude Include="..\Include\internal\pycore_dict.h" />
     <ClInclude Include="..\Include\internal\pycore_dict_state.h" />
@@ -508,6 +509,7 @@
     <ClCompile Include="..\Python\codecs.c" />
     <ClCompile Include="..\Python\compile.c" />
     <ClCompile Include="..\Python\context.c" />
+    <ClCompile Include="..\Python\critical_section.c" />
     <ClCompile Include="..\Python\dynamic_annotations.c" />
     <ClCompile Include="..\Python\dynload_win.c" />
     <ClCompile Include="..\Python\errors.c" />

diff --git a/PCbuild/pythoncore.vcxproj.filters b/PCbuild/pythoncore.vcxproj.filters
@@ -537,6 +537,9 @@
     <ClInclude Include="..\Include\internal\pycore_context.h">
       <Filter>Include\internal</Filter>
     </ClInclude>
+    <ClInclude Include="..\Include\internal\pycore_critical_section.h">
+      <Filter>Include\internal</Filter>
+    </ClInclude>
     <ClInclude Include="..\Include\internal\pycore_descrobject.h">
       <Filter>Include\internal</Filter>
     </ClInclude>
@@ -1118,6 +1121,9 @@
     <ClCompile Include="..\Python\compile.c">
       <Filter>Python</Filter>
     </ClCompile>
+    <ClCompile Include="..\Python\critical_section.c">
+      <Filter>Python</Filter>
+    </ClCompile>
     <ClCompile Include="..\Python\dynamic_annotations.c">
       <Filter>Python</Filter>
     </ClCompile>

diff --git a/Python/critical_section.c b/Python/critical_section.c
@@ -0,0 +1,103 @@
+#include "Python.h"
+#include "pycore_critical_section.h"
+
+void
+_Py_critical_section_begin_slow(struct _Py_critical_section *c, _PyMutex *m)
+{
+    PyThreadState *tstate = PyThreadState_GET();
+    c->mutex = NULL;
+    c->prev = (uintptr_t)tstate->critical_section;
+    tstate->critical_section = (uintptr_t)c;
+
+    _PyMutex_lock(m);
+    c->mutex = m;
+}
+
+void
+_Py_critical_section2_begin_slow(struct _Py_critical_section2 *c,
+                                 _PyMutex *m1, _PyMutex *m2, int flag)
+{
+    PyThreadState *tstate = PyThreadState_GET();
+    c->base.mutex = NULL;
+    c->mutex2 = NULL;
+    c->base.prev = tstate->critical_section;
+    tstate->critical_section = (uintptr_t)c | _Py_CRITICAL_SECTION_TWO_MUTEXES;
+
+    if (!flag) {
+        _PyMutex_lock(m1);
+    }
+    _PyMutex_lock(m2);
+    c->base.mutex = m1;
+    c->mutex2 = m2;
+}
+
+struct _Py_critical_section *
+_Py_critical_section_untag(uintptr_t tag)
+{
+    tag &= ~_Py_CRITICAL_SECTION_MASK;
+    return (struct _Py_critical_section *)tag;
+}
+
+// Release all locks held by critical sections. This is called by
+// _PyThreadState_Detach.
+void
+_Py_critical_section_end_all(PyThreadState *tstate)
+{
+    uintptr_t *tagptr;
+    struct _Py_critical_section *c;
+    struct _Py_critical_section2 *c2;
+
+    tagptr = &tstate->critical_section;
+    while (*tagptr && _Py_critical_section_is_active(*tagptr)) {
+        c = _Py_critical_section_untag(*tagptr);
+
+        if (c->mutex) {
+            _PyMutex_unlock(c->mutex);
+            if ((*tagptr & _Py_CRITICAL_SECTION_TWO_MUTEXES)) {
+                c2 = (struct _Py_critical_section2 *)c;
+                if (c2->mutex2) {
+                    _PyMutex_unlock(c2->mutex2);
+                }
+            }
+        }
+
+        *tagptr |= _Py_CRITICAL_SECTION_INACTIVE;
+        tagptr = &c->prev;
+    }
+}
+
+void
+_Py_critical_section_resume(PyThreadState *tstate)
+{
+    uintptr_t p;
+    struct _Py_critical_section *c;
+    struct _Py_critical_section2 *c2;
+
+    p = tstate->critical_section;
+    c = _Py_critical_section_untag(p);
+    assert(!_Py_critical_section_is_active(p));
+
+    _PyMutex *m1 = NULL, *m2 = NULL;
+
+    m1 = c->mutex;
+    c->mutex = NULL;
+    if ((p & _Py_CRITICAL_SECTION_TWO_MUTEXES)) {
+        c2 = (struct _Py_critical_section2 *)c;
+        m2 = c2->mutex2;
+        c2->mutex2 = NULL;
+    }
+
+    if (m1) {
+        _PyMutex_lock(m1);
+    }
+    if (m2) {
+        _PyMutex_lock(m2);
+    }
+
+    c->mutex = m1;
+    if (m2) {
+        c2->mutex2 = m2;
+    }
+
+    tstate->critical_section &= ~_Py_CRITICAL_SECTION_INACTIVE;
+}
diff --git a/Python/pystate.c b/Python/pystate.c
@@ -4,6 +4,7 @@
 #include "Python.h"
 #include "pycore_ceval.h"
 #include "pycore_code.h"           // stats
+#include "pycore_critical_section.h"
 #include "pycore_frame.h"
 #include "pycore_initconfig.h"
 #include "pycore_lock.h"          // _PyRawEvent
@@ -257,6 +258,12 @@ _PyThreadState_Attach(PyThreadState *tstate)
             &tstate->status,
             _Py_THREAD_DETACHED,
             _Py_THREAD_ATTACHED)) {
+
+        // resume previous critical section
+        if (tstate->critical_section != 0) {
+            _Py_critical_section_resume(tstate);
+        }
+
         return 1;
     }
     return 0;
@@ -265,6 +272,10 @@ _PyThreadState_Attach(PyThreadState *tstate)
 static void
 _PyThreadState_Detach(PyThreadState *tstate)
 {
+    if (tstate->critical_section != 0) {
+        _Py_critical_section_end_all(tstate);
+    }
+
     _Py_atomic_store_int(&tstate->status, _Py_THREAD_DETACHED);
 }