gh-112075: Add critical sections for most dict APIs (#114508)

Starts adding thread safety to dict objects.


Use @critical_section for APIs which are exposed via argument clinic and don't directly correlate with a public C API which needs to acquire the lock
Use a _lock_held suffix for keeping changes to complicated functions simple and just wrapping them with a critical section
Acquire and release the lock in an existing function where it won't be overly disruptive to the existing logic

1 files changed, 46 insertions, 0 deletions

diff --git a/Include/internal/pycore_critical_section.h b/Include/internal/pycore_critical_section.h
index bf2bbff..38ed8cd 100644
--- a/Include/internal/pycore_critical_section.h
+++ b/Include/internal/pycore_critical_section.h
@@ -104,12 +104,37 @@ extern "C" {
 # define Py_END_CRITICAL_SECTION2()                                     \
         _PyCriticalSection2_End(&_cs2);                                 \
     }
+
+// Asserts that the mutex is locked.  The mutex must be held by the
+// top-most critical section otherwise there's the possibility
+// that the mutex would be swalled out in some code paths.
+#define  _Py_CRITICAL_SECTION_ASSERT_MUTEX_LOCKED(mutex) \
+    _PyCriticalSection_AssertHeld(mutex)
+
+// Asserts that the mutex for the given object is locked. The mutex must
+// be held by the top-most critical section otherwise there's the
+// possibility that the mutex would be swalled out in some code paths.
+#ifdef Py_DEBUG
+
+#define  _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(op)                           \
+    if (Py_REFCNT(op) != 1) {                                                    \
+        _Py_CRITICAL_SECTION_ASSERT_MUTEX_LOCKED(&_PyObject_CAST(op)->ob_mutex); \
+    }
+
+#else   /* Py_DEBUG */
+
+#define  _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(op)
+
+#endif  /* Py_DEBUG */
+
 #else  /* !Py_GIL_DISABLED */
 // The critical section APIs are no-ops with the GIL.
 # define Py_BEGIN_CRITICAL_SECTION(op)
 # define Py_END_CRITICAL_SECTION()
 # define Py_BEGIN_CRITICAL_SECTION2(a, b)
 # define Py_END_CRITICAL_SECTION2()
+# define _Py_CRITICAL_SECTION_ASSERT_MUTEX_LOCKED(mutex)
+# define _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(op)
 #endif  /* !Py_GIL_DISABLED */
 
 typedef struct {
@@ -236,6 +261,27 @@ _PyCriticalSection2_End(_PyCriticalSection2 *c)
 PyAPI_FUNC(void)
 _PyCriticalSection_SuspendAll(PyThreadState *tstate);
 
+#ifdef Py_GIL_DISABLED
+
+static inline void
+_PyCriticalSection_AssertHeld(PyMutex *mutex) {
+#ifdef Py_DEBUG
+    PyThreadState *tstate = _PyThreadState_GET();
+    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
+}
+
+#endif
+
 #ifdef __cplusplus
 }
 #endif