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.

Author: colesbury

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),

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 \

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 */
 };
 

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" />

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>

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" />

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>

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;
+}

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);
 }